Resumo Técnico: JAX-FVM

Definição do Problema

A Dinâmica de Fluidos Computacional (CFD) tem sido tradicionalmente utilizada para simulações diretas, mas muitas aplicações críticas — como otimização de forma, identificação de parâmetros e modelagem de turbulência — exigem a resolução de problemas inversos. Essas tarefas necessitam do cálculo de gradientes das saídas do solver em relação aos inputs (ex: condições de contorno, parâmetros físicos ou estados iniciais). Embora os métodos adjuntos tenham abordado isso historicamente, eles são codificados de forma intrusiva (code-intrusive) e frequentemente desconectados dos ecossistemas modernos de aprendizado de máquina (ML).

Desenvolvimentos recentes em programação diferenciável (ex: JAX, PyTorch) permitiram a diferenciabilidade de ponta a ponta, possibilitando que gradientes sejam computados via diferenciação automática (AD) sem a necessidade de adjuntos derivados manualmente. No entanto, os frameworks de CFD diferenciável existentes (como JAX-CFD, JAX-Fluids e ΦFlow) são predominantemente restritos a grades cartesianas estruturadas ou discretizações espectrais. Isso cria uma lacuna significativa: malhas não estruturadas, que são essenciais para lidar com domínios geometricamente complexos na CFD industrial, permanecem amplamente ausentes do ecossistema de simulação diferenciável devido à sua conectividade irregular, que conflita com os modelos de execução vetorizada de aceleradores. Além disso, simular robustamente fluxos compressíveis em malhas não estruturadas requer esquemas com estabilidade de entropia para evitar instabilidades não lineares e oscilações espúrias próximas a choques, um recurso que é frequentemente difícil de implementar dentro de um framework totalmente diferenciável.

Metodologia

O artigo apresenta o JAX-FVM, um solver de método de volumes finitos (FVM) totalmente diferenciável e de código aberto para as equações de Euler e Navier-Stokes compressíveis bidimensionais em malhas triangulares não estruturadas. O solver é escrito inteiramente em JAX, garantindo que cada operação — desde a conectividade da malha e avaliação de fluxo até o limitador de inclinação e integração temporal — seja compilada via just-in-time (JIT), vetorizada e diferenciável.

Discretização Numérica

• Equações Governantes: O solver aborda a forma conservativa das equações de Navier-Stokes compressíveis 2D, fechadas pela lei dos gases ideais.
• Representação da Malha: O domínio é triangulado usando o gerador MeshPy/Triangle. A estrutura de dados da malha armazena vértices, baricentros de células, áreas e conectividade de faces, suportando fronteiras periódicas e vários marcadores de condições de contorno.
• Formulação Semidiscreta: A atualização centrada na célula é derivada via teorema da divergência. O resíduo é computado como a soma dos fluxos numéricos sobre as faces das células.
• Reconstrução de Gradiente: A precisão de segunda ordem é alcançada através da reconstrução MUSCL de variáveis primitivas. Os gradientes celulares são computados usando um ajuste de mínimos quadrados ponderados (LSQ), resolvido via um solver linear vetorizado (jax.vmap). O limitador de inclinação é aplicado utilizando o limitador de Venkatakrishnan ou o clássico limitador minmod.
• Fluxos Numéricos: O núcleo do esquema é uma formulação de estabilidade de entropia.
  • O fluxo convectivo base é um fluxo de dois pontos conservativo de entropia (Tadmor/Ismail-Roe).
  • A dissipação é adicionada via um termo de dissipação de Rusanov (Lax-Friedrichs local) ou Roe baseado em variáveis de entropia.
  • Crucialmente, a ação do Jacobiano necessária para mapear saltos de variáveis de entropia para variáveis conservativas é computada exatamente usando AD de modo direto (jax.jvp), em vez de derivação analítica manual.
  • Os fluxos viscosos são avaliados usando gradientes LSQ e estados primitivos médios por face.
• Integração Temporal: O solver suporta integradores explícitos (RK2, RK3, RK4) e integradores implícitos livres de matriz (matrix-free) (Backward Euler, SDIRK2). As etapas implícitas resolvem sistemas não lineares via iterações de Newton, onde os produtos Jacobiano-vetor são computados de forma matrix-free usando AD de modo direto e passados para um solver GMRES.

Arquitetura de Software

• Diferenciabilidade: Toda a lógica de resíduo e de passos de tempo é encapsulada em jax.jit. A AD de modo direto é usada internamente para a dissipação de entropia e Jacobianos de solvers implícitos, enquanto a AD de modo reverso permite a propagação de gradientes para qualquer função escalar de volta aos inputs.
• Agnosticismo de Hardware: O código roda transparentemente em CPUs ou GPUs (via JAX habilitado para CUDA).
• Integração: A biblioteca depende de flax e optax, facilitando a construção de loops de otimização e design inverso diretamente sobre o solver.

Principais Contribuições

1. Primeiro FVM Não Estruturado Diferenciável de Fluxo Compressível Disponível Aberto: O JAX-FVM preenche a lacuna deixada pelas ferramentas de CFD diferenciáveis existentes que são restritas a grades estruturadas. É, tanto quanto sabem os autores, um dos primeiros solvers diferenciáveis para equações de Euler/Navier-Stokes compressíveis operando nativamente em malhas triangulares não estruturadas.
2. Estabilidade de Entropia via AD: A implementação integra a teoria de fluxo conservativo de entropia com dissipação de variáveis de entropia. Ao usar AD para computar as ações de Jacobiano necessárias para os termos de dissipação, o solver alcança estabilidade de entropia provável sem a derivação manual de Jacobianos complexos.
3. Diferenciabilidade de Ponta a Ponta: O solver permite o cálculo direto de gradientes de funcionais de fluxo em relação a condições iniciais, dados de contorno, parâmetros físicos (viscosidade, condutividade, $\gamma$) ou modelos neurais incorporados.
4. Plataforma de Pesquisa: O código fornece uma implementação compacta e orientada à pesquisa que combina geometria não estruturada, numéricos consistentes com entropia e AD baseado em JAX, servindo como uma fundação para problemas inversos baseados em gradiente e modelagem híbrida física/ML.

Resultados e Verificação

O artigo valida o JAX-FVM através de uma suíte de casos de verificação padrão:

• Degrau Frontal Mach 3: Um problema de choque não estacionário canônico envolvendo um choque de proa destacado, um choque de Mach, choques refletidos e uma linha de deslizamento (slip line). O solver captura com sucesso essas características em uma triangulação totalmente não estruturada usando o fluxo conservativo de entropia, reconstrução MUSCL com limitação de Venkatakrishnan e integração temporal RK2 explícita.
• Benchmarks Adicionais: A distribuição inclui 19 configurações Riemann bidimensionais (Lax e Liu), reflexão de Mach dupla, instabilidade de Kelvin-Helmholtz, vórtice de Taylor-Green e vários problemas de vórtice de baixo Mach.
• Diagnósticos: O solver monitora a entropia total, energia cinética, enstrofia, palinstrofia, vorticidade e número de Mach para verificar propriedades de preservação de estrutura.

Significância e Alegações

O artigo posiciona o JAX-FVM como uma ponte entre a CFD de alta fidelidade em geometrias complexas e a programação diferenciável moderna. Sua significância reside em permitir o design inverso baseado em gradiente e a otimização com restrição de PDE para fluxos compressíveis em malhas não estruturadas, uma capacidade anteriormente indisponível em frameworks diferenciáveis de código aberto.

Os autores alegam que, ao combinar capacidades de malha não estruturada com numéricos estáveis de entropia e diferenciabilidade de ponta a ponta, o JAX-FVM cria uma plataforma natural para:

• Recuperar parâmetros a partir de dados esparsos (ex: resistência vascular).
• Otimização de forma baseada em pressões de superfície.
• Aprender fechamentos de turbulência ou discretizações baseadas em dados de maneira consistente com o solver.

O trabalho é modesto em seu escopo, reconhecendo limitações atuais: é restrito a malhas triangulares 2D (com extensão para 3D planejada), depende do gerador MeshPy (com leitor VTK planejado) e atualmente carece de modelos de turbulência (fechamentos RANS ou LES são trabalhos futuros). O solver é apresentado primariamente como uma ferramenta de pesquisa para facilitar o desenvolvimento de modelos híbridos física/ML e fluxos de trabalho de design inverso.