JAX-Fluids: A Differentiable Fluid Dynamics Package for Machine-Learning Accelerated CFD

Strömungen sind allgegenwärtig in der Natur und ingenieurwissenschaftlichen Anwendungen. Die numerische Vorhersage kompressibler Strömungen ist jedoch herausfordernd, da sie durch intrinsisch nichtlineare Mechanismen sowie komplexe Strömungsphänomene wie Turbulenz, Stoßwellen, Materialgrenzflächen und chemische Reaktionen geprägt sind. Präzise Simulationen in der numerischen Strömungsmechanik (Computational Fluid Dynamics, CFD) erfordern daher hochentwickelte numerische Methoden, die speziell auf diese Probleme zugeschnitten sind. Der jüngste Erfolg von maschinellem Lernen (ML) zur Lösung partieller Differentialgleichungen (PDEs) zeigt zudem, dass ML das Potenzial besitzt, konventionelle CFD-Methoden zu unterstützen und zu beschleunigen. Die Symbiose aus ML und CFD erfordert neue Solver und methodische Paradigmen. Im Rahmen unserer Arbeiten zur Unterstützung von ML-basierter CFD haben wir JAX-Fluids entwickelt. JAX-Fluids ist ein umfassender, vollständig differenzierbarer CFD-Solver für kompressible Strömungen, implementiert in JAX. Das Framework löst die kompressiblen Einphasen-, Zweiphasen- und Mehrkomponenten-Navier-Stokes-Gleichungen. Jede einzelne Routine in JAX-Fluids ist automatisch differenzierbar, was eine End-to-End-Optimierung numerischer Verfahren ermöglicht. 

Zu den zentralen Eigenschaften von JAX-Fluids gehören unter anderem: 

• Nahtlose Integration von Modellen für maschinelles Lernen in einen multiphysikalischen CFD-Solver
• Differenzierbare Algorithmen zur End-to-End-Optimierung datengetriebener Modelle
• Schnelle Prototypentwicklung durch die Implementierung in Python
• Effiziente Skalierung auf CPU-, GPU- und TPU-HPC-Systemen 

Aktuelle Entwicklungsarbeiten umfassen: 

• Weiterentwicklung hin zu einem differenzierbaren Multiphysik-Solver
• Entwicklung numerischer Methoden für Zweiphasenströmungen
• End-to-End-Optimierung von Surrogatmodellen
• Inverse Probleme
• Quantifizierung von Simulationsunsicherheiten  

Integration maschinellen Lernens

JAX-Fluids wurde mit dem Ziel entwickelt, Forschung an der Schnittstelle von numerischer Strömungsmechanik und maschinellem Lernen zu erleichtern. Die Integration und Hybridisierung von Machine-Learning Modellen mit etablierten CFD-Algorithmen ist ein zentrales Merkmal von JAX-Fluids. Aus diesem Grund wurde JAX als Backend für den CFD-Solver gewählt. Dadurch wird eine nahtlose Einbindung datengetriebener Modelle in den Solver sowie eine End-to-End-Optimierung neuronaler Netze ermöglicht. 

Integration neuronaler Netze 

JAX ist eine Python-Bibliothek für hochperformantes numerisches Rechnen und maschinelles Lernen. Die Integration neuronaler Subroutinen in JAX-Fluids ist entsprechend unkompliziert. So wurde beispielsweise eine neuronale WENO3-NN-Zellflächenrekonstruktion entwickelt, die die klassische WENO3-JS-Rekonstruktion ersetzen kann. Das Architekturkonzept zeigt, dass Eingabe und Ausgabe von WENO3-NN identisch zur klassischen WENO3-JS-Methode sind. Der Unterschied besteht darin, dass beim WENO3-NN eine neuronale Netzwerkstruktur das analytische Glattheitsmaß ersetzt. 

End-to-End-Optimierung 

Der gesamte CFD-Algorithmus ist vollständig differenzierbar, wodurch das Training neuronaler Netze im End-to-End-Ansatz möglich wird. Der Optimierungsprozess von JAX-Fluids lässt sich anhand des Trainings eines maschinell gelernten ILES-Modells (ML-ILES) veranschaulichen: Ausgehend von einer Anfangsbedingung wird das ML-CFD-Modell mit trainierbaren Parametern über mehrere Zeitschritte in JAX-Fluids zeitlich entwickelt und ergibt eine räumlich-zeitliche Trajektorie (JXF). Diese wird mit einer Referenzlösung (hier einer gefilterten DNS-Trajektorie) verglichen, um eine Verlustfunktion zu berechnen. Mit der JAX-Routine jax.grad wird anschließend der Gradient dieser Verlustfunktion bestimmt, der zur Aktualisierung der Modellparameter verwendet wird. 

JAX-Fluids zeichnet sich als vielseitiger Löser zur Simulation komplexer Mehrphasen-, reaktiver und turbulenter Strömungen aus. Die Mehrphasenmodelle zeigen ihre Leistungsfähigkeit insbesondere in Szenarien, in denen starke Stoßwellen mit Tropfen oder Blasen interagieren. Darüber hinaus ermöglicht JAX-Fluids detaillierte Simulationen reaktiver Strömungen unter Berücksichtigung komplexer chemischer Prozesse. Eine weitere Stärke liegt in der effektiven numerischen Untersuchung turbulenter Strömungen. Unter Nutzung der scharfen Level-Set-Methode erlaubt JAX-Fluids zudem die Einbindung komplexer Geometrien als eingebettete Randbedingungen (Immersed Boundaries), wodurch die Flexibilität und Anwendbarkeit des Frameworks weiter erhöht wird.

Mehrphasen-Strömungen

JAX-Fluids implementiert zwei unterschiedliche Zweiphasenmodelle: eine auf Level-Set basierende scharfe Grenzflächenmethode (LSM) sowie eine diffuse Methode (DIM). Die LSM erhält während der gesamten Simulation eine scharfe Grenzfläche, während die DIM eine kontrollierte künstliche Durchmischung in einem schmalen Bereich um die Grenzfläche zulässt. Sowohl die LSM- als auch die DIM-Algorithmen sind vollständig automatisch differenzierbar.

Reaktiven Mehrkomponenten Strömungen

JAX-Fluids bietet Funktionalitäten zur Simulation reaktiver Mehrkomponentenströmungen. Dabei lösen wir die kompressiblen reaktiven Navier-Stokes-Gleichungen für Mehrkomponentensysteme unter Berücksichtigung detaillierter Chemie- und Transportmodelle. Materialeigenschaften und Reaktionsmechanismen können bequem über Cantera-Inputdateien konfiguriert werden. Die Implementierung wurde anhand kanonischer Testfälle validiert, beispielsweise durch den Vergleich der JAX-Fluids-Ergebnisse mit Cantera für Wasserstoff-Sauerstoff-Verbrennung (San-Diego-Mechanismus) unter stöchiometrischen Bedingungen sowie für Methan-Sauerstoff-Verbrennung (GRI-Mech-3.0-Mechanismus). Darüber hinaus wurde eine reaktive Schock-Blasen-Simulation untersucht, bei der eine Stoßwelle auf eine Wasserstoff-Sauerstoff-Xenon-Blase trifft. Aufgrund der Stoßfokussierung entzündet sich das Gemisch, sobald Druck und Temperatur ausreichende Werte erreichen.

Kompressible Turbulenz

JAX-Fluids unterstützt die Simulation kompressibler turbulenter Strömungen mit Stoßwellen. Insbesondere verfügt JAX-Fluids über Fähigkeiten zur Durchführung von Direct Numerical Simulations (DNS) sowie impliziten Large-Eddy Simulations (ILES). Dabei kommen moderne nichtlineare Diskretisierungsschemata mit gezielt gesteuertem Subgrid-Scale-Verhalten zum Einsatz, wie etwa die ALDM-Methode oder TENO-basierte Verfahren mit adaptiver Dissipationskontrolle. In jüngerer Zeit wird zudem die Entwicklung implizieter Large-Eddy Simulationen mit Hilfe von maschinellen Lernes (ML-ILES) untersucht.

Immersed-Boundaries

Veröffentlichungen

• Bezgin, D., Buhendwa, A. & Adams, N. A. (2023): Jax-Fuids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows. Computer Physics Communications, 282, 108527.