A differentiable software suite for accelerated simulation of turbulent flows

O artigo apresenta o IncompressibleNavierStokes.jl, um pacote Julia de código aberto que utiliza kernels diferenciáveis e otimizações de memória para simular eficientemente escoamentos turbulentos em GPUs, permitindo o treinamento de modelos de fechamento por redes neurais diretamente dentro de simulações de grandes vórtices.

O essencial

Imagine que você está tentando prever o clima, projetar um carro mais aerodinâmico ou entender como o sangue flui em uma veia. Para fazer isso, os cientistas usam equações matemáticas complexas chamadas Equações de Navier-Stokes. Elas descrevem como os fluidos (como água e ar) se movem.

O problema é que resolver essas equações para cenários turbulentos (como um furacão ou o ar passando por um carro em alta velocidade) é como tentar adivinhar o futuro de cada gota de chuva em uma tempestade. É computacionalmente caro e difícil.

Este artigo apresenta uma nova ferramenta chamada IncompressibleNavierStokes.jl (ou apenas INS.jl). Pense nela como um "super-gerenciador de tráfego" para fluidos, escrito em uma linguagem de programação moderna chamada Julia.

Aqui está uma explicação simples do que ela faz e por que é especial, usando analogias do dia a dia:

1. O "Motor" que roda em qualquer lugar (Hardware-Agnóstico)

Imagine que você tem um motor de carro. Normalmente, você precisa de um motor diferente para uma pista de terra (CPU) e outro para uma pista de asfalto (GPU).
O INS.jl é como um motor universal. O mesmo código de programação funciona perfeitamente tanto em computadores comuns quanto em supercomputadores com placas gráficas (GPUs) poderosas. Isso significa que os cientistas não precisam reescrever o software para cada máquina nova; eles apenas dizem "rodar" e ele se adapta.

2. O "Espelho Mágico" (Diferenciabilidade)

Esta é a parte mais inovadora. Normalmente, simular um fluido é como dirigir um carro: você vira o volante e o carro se move. Se você quiser saber como virar o volante para chegar mais rápido a um destino, você teria que tentar muitas vezes, errar e corrigir.

O INS.jl tem um espelho mágico (chamado de "adjunto" ou "diferenciável").

• Como funciona: Se você simular o fluxo de ar e o resultado não for perfeito, o software pode olhar para trás (como um filme em reverso) e dizer exatamente: "Se você tivesse mudado aquela pequena coisa no início, o resultado teria sido 10% melhor".
• Por que importa: Isso permite treinar Inteligência Artificial (Redes Neurais) diretamente dentro da simulação. Em vez de apenas simular o vento, o software aprende a prever o vento de forma mais inteligente, ajustando-se sozinho enquanto roda. É como ter um piloto de teste que aprende a dirigir melhor a cada volta, sem sair do carro.

3. O "Detetive de Memória" (Otimização)

Simular turbulência exige muita memória (RAM), como tentar guardar milhões de fotos em um celular pequeno.
O INS.jl é um mestre em economia de espaço. Em vez de criar uma nova pasta para cada cálculo intermediário (o que enche a memória), ele reutiliza as mesmas pastas, apagando o que não serve mais e escrevendo o novo ali mesmo.

• O Resultado: Com essa técnica, eles conseguiram rodar simulações extremamente detalhadas (chamadas de DNS - Simulação Numérica Direta) em uma única placa de vídeo de consumo (como uma RTX 4090), algo que antes exigia supercomputadores gigantescos. É como conseguir fazer uma festa para 1.000 pessoas em uma sala pequena, apenas organizando a mobília de forma inteligente.

4. A "Caixa de Ferramentas" (Flexibilidade)

O software foi feito para ser fácil de modificar. Se um cientista quiser adicionar um novo fenômeno físico (como temperatura, que faz o ar subir e descer, ou forças aleatórias), ele não precisa reescrever o motor inteiro. Basta adicionar uma "peça" nova na caixa de ferramentas.
É como montar um LEGO: você pode trocar uma peça por outra para mudar a forma do castelo, sem precisar construir tudo do zero.

5. O "Laboratório de Testes" (Validação)

Para garantir que o software não está inventando coisas, os autores o testaram contra dados reais de um "canal de turbulência" (um fluxo de água em um tubo).

• Eles compararam os resultados do software com dados de referência de outros cientistas.
• O Veredito: O software acertou em cheio na maioria dos pontos, mostrando que é preciso e confiável. Eles também testaram modelos de "fechamento" (fórmulas que ajudam a prever turbulência em escalas menores) e viram que ele consegue rodar modelos clássicos e novos modelos baseados em IA com a mesma facilidade.

Resumo da Ópera

O INS.jl é uma ferramenta de código aberto que torna a simulação de fluidos:

1. Mais rápida: Roda em qualquer hardware.
2. Mais inteligente: Permite que a IA aprenda com a simulação em tempo real.
3. Mais acessível: Roda em computadores pessoais potentes, não apenas em supercomputadores.
4. Confiável: Foi rigorosamente testado e documentado.

Em suma, é como dar aos cientistas um "super-poder" para entender e prever o comportamento da água e do ar com uma precisão e velocidade que antes eram impossíveis, tudo isso em um único pacote de software fácil de usar.

Resumo técnico

1. Problema e Motivação

A simulação de escoamentos turbulentos, especialmente através de Simulação Numérica Direta (DNS) e Simulação de Grandes Vórtices (LES), enfrenta dois desafios principais:

1. Custo Computacional e Memória: DNS de alta resolução requerem armazenamento massivo de matrizes esparsas e intermediários, limitando a resolução alcançável em hardware de consumo (como GPUs de desktop).
2. Integração de Modelos de Aprendizado de Máquina: O treinamento de modelos de fechamento (closure models) baseados em redes neurais para LES tradicionalmente exige acoplamento externo entre solvers de CFD e bibliotecas de aprendizado de máquina (ex: Python e C++), introduzindo sobrecarga de comunicação e complexidade. Além disso, a diferenciação automática (AD) eficiente através de solvers CFD é difícil de implementar devido a operações de mutação de memória e kernels personalizados.

O objetivo do trabalho é apresentar uma solução unificada que permita simulações de alta performance, otimizadas para hardware acelerador (GPU/CPU), e que seja totalmente diferenciável para permitir o treinamento "a posteriori" de modelos de fechamento neural diretamente dentro do solver.

2. Metodologia

O artigo descreve o desenvolvimento e a implementação do pacote IncompressibleNavierStokes.jl (INS.jl), escrito na linguagem de programação Julia.

• Discretização Numérica:

  • Utiliza o método de volumes finitos de segunda ordem em grades cartesianas escalonadas (staggered grids), seguindo a abordagem de Harlow e Welch.
  • A pressão é armazenada nos centros dos volumes, enquanto as componentes de velocidade são armazenadas nas faces. Isso acopla naturalmente os campos de pressão e velocidade, evitando instabilidades artificiais.
  • Suporta grades uniformes e não uniformes (com funções de estiramento como cosseno, tangente hiperbólica e geométrica).
  • A integração temporal utiliza métodos de Runge-Kutta (incluindo variantes de baixo armazenamento como Wray3) com projeção de pressão via equação de Poisson para garantir a condição de divergência nula.

• Arquitetura de Software e Diferenciabilidade:

  • Kernels Agnósticos de Hardware: Todos os operadores diferenciais (divergência, gradiente de pressão, difusão, convecção) são implementados como kernels usando KernelAbstractions.jl. O mesmo código fonte é compilado para CPUs multithread ou GPUs, dependendo do tipo de array de entrada.
  • Otimização de Memória: Para permitir DNS de precisão dupla em GPUs de 24GB (ex: RTX 4090) até resoluções de $840^3$, o software emprega:
    1. Reutilização de arrays (kernels mutantes que sobrescrevem buffers pré-alocados).
    2. Métodos de Runge-Kutta de baixo armazenamento.
    3. Computação preguiçosa (lazy) de componentes de tensores convectivos.
  • Diferenciação Automática Reversa: O solver é totalmente diferenciável. Os autores implementaram kernels adjuntos (pullbacks) escritos à mão para todos os operadores discretos, registrados através da interface ChainRules.jl. Isso permite que o Zygote.jl realize diferenciação reversa através de todo o solver, incluindo a solução da equação de Poisson.
  • Separação Mutante/Puro: Para manter a eficiência na simulação direta, usam-se funções mutantes (que alteram arrays in-place). Para a diferenciação (necessária no treinamento de redes neurais), usa-se variantes puras (que retornam novos arrays), garantindo compatibilidade com o AD.

• Integração de Redes Neurais:

  • O pacote integra nativamente modelos de fechamento de redes neurais da biblioteca Lux.jl.
  • Permite o treinamento "a posteriori": o modelo neural é embutido no solver LES, e os gradientes da função de perda são calculados propagando o erro através de múltiplos passos de tempo (unrolling), sem necessidade de frameworks externos de acoplamento.

3. Contribuições Principais

1. INS.jl: Um pacote de código aberto em Julia para DNS e LES de escoamentos incompressíveis, destacando-se por ser totalmente diferenciável.
2. Kernels Adjointos Manuais: A implementação de regras de pullback (adjuntos) para todos os operadores discretos, permitindo diferenciação reversa eficiente e precisa através de kernels de GPU personalizados.
3. Pipeline Unificado: Capacidade de realizar todo o fluxo de trabalho (geração de dados DNS, treinamento de modelo de fechamento neural e avaliação LES) em um único ambiente e linguagem, eliminando a necessidade de interfaces entre linguagens (ex: Python-C++).
4. Otimização de Hardware: Demonstração de DNS de precisão dupla em resoluções extremas ($840^3$) em uma única GPU de datacenter (H100) ou de consumo (RTX 4090), graças a técnicas agressivas de gerenciamento de memória.
5. Práticas de Engenharia de Software: O projeto adota rigorosas práticas de CI/CD, testes automatizados, documentação gerada via literate programming e versionamento semântico, garantindo reprodutibilidade e manutenibilidade.

4. Resultados

Os resultados foram validados através de dois casos de teste principais:

• Vórtice de Taylor-Green 2D (Verificação Numérica):

  • Confirmou-se a convergência de segunda ordem para precisão dupla (Float64) e simples (Float32) em grades uniformes.
  • A precisão meia (Float16) mostrou-se inadequada para a discretização espacial direta devido a erros de arredondamento e instabilidade, embora possa ser útil para inferência em redes neurais.

• Escoamento Turbulento em Canal (Re$\tau$ = 180):

  • DNS: Os resultados de DNS em grade uniforme ($512 \times 1024 \times 256$) mostraram excelente concordância com dados de referência de Vreman e Kuerten para perfis de velocidade média e flutuações RMS. Pequenas discrepâncias em momentos de ordem superior (3ª e 4ª) foram atribuídas à ordem de precisão do esquema numérico (2ª ordem vs. 4ª ordem no referência).
  • LES: Simulações em grades mais grossas ($128 \times 64 \times 64$) com modelos de viscosidade turbulenta clássicos (Smagorinsky, WALE, QR, Vreman) reproduziram comportamentos esperados na literatura. O modelo Smagorinsky mostrou-se excessivamente dissipativo perto da parede, enquanto os modelos WALE, QR e Vreman apresentaram melhor desempenho, com perfis de velocidade e flutuações mais próximos dos dados DNS de referência.

5. Significado e Impacto

O trabalho representa um avanço significativo na interseção entre Dinâmica dos Fluidos Computacional (CFD) e Aprendizado de Máquina (ML):

• Aceleração do Desenvolvimento de Modelos: Ao permitir o treinamento de modelos de fechamento neural diretamente dentro do solver diferenciável, remove-se a barreira de acoplamento externo, acelerando a pesquisa em "Physics-Informed Machine Learning".
• Acesso a Alta Performance: A otimização de memória permite que pesquisadores realizem simulações de DNS de alta fidelidade em hardware de laboratório padrão (uma única GPU), democratizando o acesso a dados de alta qualidade para treinamento de IA.
• Reprodutibilidade: A ênfase em práticas de software robustas e o uso de Julia (com seu sistema de tipos e múltiplo despacho) estabelecem um novo padrão para a transparência e reprodutibilidade em códigos científicos de alto desempenho.
• Futuro: O trabalho abre caminho para o uso de estratégias de precisão mista (mixed-precision) e a expansão para computação distribuída (multi-GPU/nós), além da exploração de diferenciadores de baixo nível como Enzyme.jl para suportar código mutante diretamente.

Em resumo, o INS.jl fornece uma infraestrutura moderna, eficiente e flexível para a próxima geração de simulações de turbulência assistidas por inteligência artificial.