{"id":13409120,"url":"https://github.com/google/jax-cfd","last_synced_at":"2025-10-21T19:31:47.980Z","repository":{"id":37850677,"uuid":"350436461","full_name":"google/jax-cfd","owner":"google","description":"Computational Fluid Dynamics in JAX","archived":false,"fork":false,"pushed_at":"2024-08-08T18:19:28.000Z","size":6044,"stargazers_count":773,"open_issues_count":30,"forks_count":110,"subscribers_count":14,"default_branch":"main","last_synced_at":"2025-01-07T16:08:20.123Z","etag":null,"topics":["cfd","jax"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/google.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-03-22T17:46:12.000Z","updated_at":"2025-01-03T17:08:54.000Z","dependencies_parsed_at":"2023-02-10T11:15:29.501Z","dependency_job_id":"38f1303e-7d0d-40c1-910a-1c59817d44f3","html_url":"https://github.com/google/jax-cfd","commit_stats":{"total_commits":166,"total_committers":18,"mean_commits":9.222222222222221,"dds":0.6506024096385542,"last_synced_commit":"d215f13282bd63045fb3455f8fac061653428040"},"previous_names":[],"tags_count":3,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Fjax-cfd","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Fjax-cfd/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Fjax-cfd/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Fjax-cfd/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google","download_url":"https://codeload.github.com/google/jax-cfd/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":243593289,"owners_count":20316160,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cfd","jax"],"created_at":"2024-07-30T20:00:58.124Z","updated_at":"2025-10-21T19:31:42.952Z","avatar_url":"https://github.com/google.png","language":"Jupyter Notebook","funding_links":[],"categories":["AI/ML for Simulation","Computational Fluid Dynamics","💻 Computing Frameworks"],"sub_categories":["Verified vs declared","Differential programming","Specialized Frameworks"],"readme":"# JAX-CFD: Computational Fluid Dynamics in JAX\n\nAuthors: Dmitrii Kochkov, Jamie A. Smith, Peter Norgaard, Gideon Dresdner, Ayya Alieva, Stephan Hoyer\n\nJAX-CFD is an experimental research project for exploring the potential of\nmachine learning, automatic differentiation and hardware accelerators (GPU/TPU)\nfor computational fluid dynamics. It is implemented in\n[JAX](https://github.com/google/jax).\n\nTo learn more about our general approach, read our paper [Machine learning accelerated computational fluid dynamics](https://www.pnas.org/content/118/21/e2101784118)\n(PNAS 2021).\n\n## Getting started\n\nThe \"notebooks\" directory contains several demonstrations of using the JAX-CFD\ncode.\n\n- Demos of different simulation setups:\n  - [2D simulation with FVM on a staggered grid](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/demo.ipynb)\n  - [2D simulation with a psuedo-spectral solver](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/spectral_forced_turbulence.ipynb)\n  - [2D simulation of channel flow](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/channel_flow_demo.ipynb)\n  - [2D simulation with FVM on a collocated grid](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/collocated_demo.ipynb) (experimental)\n\n- Reproduce results from our PNAS paper:\n  - [Data analysis and evaluation](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/ml_accelerated_cfd_data_analysis.ipynb)\n  - [Running our pre-trained models](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/ml_model_inference_demo.ipynb)\n\n## Organization\n\nJAX-CFD is organized around sub-modules:\n\n- `jax_cfd.base`: core finite volume/difference methods for CFD, written in JAX.\n- `jax_cfd.spectral`: core pseudospectral methods for CFD, written in JAX.\n- `jax_cfd.ml`: machine learning augmented models for CFD,\n  written in JAX and [Haiku](https://dm-haiku.readthedocs.io/en/latest/).\n- `jax_cfd.data`: data processing utilities for preparing, evaluating and\n  post-processing data created with JAX-CFD, written in\n  [Xarray](http://xarray.pydata.org/) and\n  [Pillow](https://pillow.readthedocs.io/).\n\nA base install with `pip install jax-cfd` only requires NumPy, SciPy and JAX.\nTo install dependencies for the other submodules, use `pip install jax-cfd[ml]`,\n`pip install jax-cfd[data]` or `pip install jax-cfd[complete]`.\n\n## Numerics\n\nJAX-CFD is currently focused on unsteady turbulent flows:\n\n- *Spatial discretization*:\n  - *Finite volume/difference* methods on a staggered grid (the \"Arakawa C\" or\n  \"MAC\" grid) with pressure at the center of each cell and velocity components\n  defined on corresponding faces.\n  - *Pseudospectral* methods for vorticity which use anti-aliasing filtering\n  techniques for non-linear terms to maintain stability.\n- *Temporal discretization*: Currently only first-order temporal\n  discretization, using explicit time-stepping for advection and either implicit\n  or explicit time-stepping for diffusion.\n- *Pressure solves*: Either CG or fast diagonalization with real-valued FFTs\n  (suitable for periodic boundary conditions).\n- *Boundary conditions*: Currently only periodic boundary conditions are\n  supported.\n- *Advection*: We implement 2nd order accurate \"Van Leer\" schemes.\n- *Closures*: We currently implement Smagorinsky eddy-viscosity models.\n\nTODO: add a notebook explaining our numerical models in more depth.\n\nIn the long term, we're interested in expanding JAX-CFD to implement methods\nrelevant for related research, e.g.,\n\n- Colocated grids\n- Alternative boundary conditions (e.g., non-periodic boundaries and immersed\n  boundary methods)\n- Higher order time-stepping\n- Geometric multigrid\n- Steady state simulation (e.g., RANS)\n- Distributed simulations across multiple TPUs/GPUs\n\nWe would welcome collaboration on any of these! Please reach out (either on\nGitHub or by email) to coordinate before starting significant work.\n\n## Projects using JAX-CFD\n\n- [Variational Data Assimilation with a Learned Inverse Observation Operator](https://github.com/googleinterns/invobs-data-assimilation)\n\n## Other awesome projects\n\nOther differentiable CFD codes compatible with deep learning:\n\n- [PhiFlow](https://github.com/tum-pbs/PhiFlow/) supports TensorFlow, PyTorch and JAX\n- [Fluid simulation in Autograd](https://github.com/HIPS/autograd#end-to-end-examples)\n\nJAX for science:\n\n- [JAX-MD](https://github.com/google/jax-md)\n- [JAX-DFT](https://github.com/google-research/google-research/tree/master/jax_dft)\n- [jax-cosmo](https://github.com/DifferentiableUniverseInitiative/jax_cosmo)\n- [Veros](https://github.com/team-ocean/veros)\n\nDid we miss something? Please let us know!\n\n## Citation\n\nIf you use our finite volume method (FVM) or ML models, please cite:\n\n```\n@article{Kochkov2021-ML-CFD,\n  author = {Kochkov, Dmitrii and Smith, Jamie A. and Alieva, Ayya and Wang, Qing and Brenner, Michael P. and Hoyer, Stephan},\n  title = {Machine learning{\\textendash}accelerated computational fluid dynamics},\n  volume = {118},\n  number = {21},\n  elocation-id = {e2101784118},\n  year = {2021},\n  doi = {10.1073/pnas.2101784118},\n  publisher = {National Academy of Sciences},\n  issn = {0027-8424},\n  URL = {https://www.pnas.org/content/118/21/e2101784118},\n  eprint = {https://www.pnas.org/content/118/21/e2101784118.full.pdf},\n  journal = {Proceedings of the National Academy of Sciences}\n}\n```\n\nIf you use our spectral code, please cite:\n\n```\n@article{Dresdner2022-Spectral-ML,\n  doi = {10.48550/ARXIV.2207.00556},\n  url = {https://arxiv.org/abs/2207.00556},\n  author = {Dresdner, Gideon and Kochkov, Dmitrii and Norgaard, Peter and Zepeda-Núñez, Leonardo and Smith, Jamie A. and Brenner, Michael P. and Hoyer, Stephan},\n  title = {Learning to correct spectral methods for simulating turbulent flows},\n  publisher = {arXiv},\n  year = {2022},\n  copyright = {arXiv.org perpetual, non-exclusive license}\n}\n```\n\n## Local development\n\nTo locally install for development:\n```\ngit clone https://github.com/google/jax-cfd.git\ncd jax-cfd\npip install jaxlib\npip install -e \".[complete]\"\n```\n\nThen to manually run the test suite:\n```\npytest -n auto jax_cfd --dist=loadfile --ignore=jax_cfd/base/validation_test.py\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle%2Fjax-cfd","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle%2Fjax-cfd","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle%2Fjax-cfd/lists"}