jax-cfd

.github/workflows/ci-build.yml

+name: CI
+
+on:
+  # Triggers the workflow on push or pull request events but only for the main branch
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+  # Allows you to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [3.7, 3.9]
+    steps:
+    - name: Cancel previous
+      uses: styfle/cancel-workflow-action@0.7.0
+      with:
+        access_token: ${{ github.token }}
+      if: ${{github.ref != 'refs/head/main'}}
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Get pip cache dir
+      id: pip-cache
+      run: |
+        python -m pip install --upgrade pip wheel
+        echo "::set-output name=dir::$(pip cache dir)"
+    - name: pip cache
+      uses: actions/cache@v2
+      with:
+        path: ${{ steps.pip-cache.outputs.dir }}
+        key: ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}
+    - name: Install dependencies
+      run: |
+        pip install jax jaxlib
+    - name: Test with pytest
+      # validation_test is too slow (it needs to run on TPUs)
+      # time_stepping_test requires reducing tolerance too much in GitHub actions (maybe a JAX bug?)
+      run: |
+        pip install -e ".[complete]"
+        pytest -n auto --forked jax_cfd --ignore=jax_cfd/base/validation_test.py --ignore=jax_cfd/spectral/time_stepping_test.py

.gitignore

+content
+.pytest_cache
+test/
+*.egg-info

CONTRIBUTING.md

+# How to Contribute
+
+We'd love to accept your patches and contributions to this project. There are
+just a few small guidelines you need to follow.
+
+## Contributor License Agreement
+
+Contributions to this project must be accompanied by a Contributor License
+Agreement (CLA). You (or your employer) retain the copyright to your
+contribution; this simply gives us permission to use and redistribute your
+contributions as part of the project. Head over to
+<https://cla.developers.google.com/> to see your current agreements on file or
+to sign a new one.
+
+You generally only need to submit a CLA once, so if you've already submitted one
+(even if it was for a different project), you probably don't need to do it
+again.
+
+## Code Reviews
+
+All submissions, including submissions by project members, require review. We
+use GitHub pull requests for this purpose. Consult
+[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
+information on using pull requests.
+
+## Community Guidelines
+
+This project follows
+[Google's Open Source Community Guidelines](https://opensource.google/conduct/).

Dockerfile

+FROM image.sourcefind.cn:5000/dcu/admin/base/jax:0.4.23-ubuntu20.04-dtk24.04-py310
\ No newline at end of file

LICENSE

+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

+# jax-cfd_jax
+
+## 仓库介绍
+
+JAX-CFD是一个探索机器学习、自动微分和硬件加速器（GPU/TPU）在计算流体动力学中潜在应用的实验研究项目。它是用JAX实现的。
+
+### 组织方式
+
+JAX-CFD是围绕子模块组织的：
+
+- jax_cfd.base：用于CFD的核心有限体积/差分方法，用JAX编写。
+
+- jax_cfd.spectral：用于CFD的核心伪谱方法，用JAX编写。
+
+- jax_cfd.ml：用于CFD的机器学习增强模型，用JAX和Haiku编写。
+
+- jax_cfd.data：用于准备、评估和后处理使用JAX-CFD创建的数据的数据处理工具，用Xarray和Pillow编写。
+
+### 数值方法
+
+JAX-CFD目前专注于非稳态湍流流动：
+- 空间离散化：
+
+    - 在交错网格（"Arakawa C"或"MAC"网格）上使用有限体积/差分方法，每个单元格的中心定义压力，速度分量定义在相应的面上。
+    - 对于涡度使用伪谱方法，采用抗混叠滤波技术处理非线性项以保持稳定性。
+
+- 时间离散化：目前仅支持一阶时间离散化，对对流使用显式时间步进，对扩散使用隐式或显式时间步进。
+- 压力求解：使用CG（共轭梯度法）或实值FFT（快速傅里叶变换）的快速对角化（适用于周期性边界条件）。
+- 边界条件：目前仅支持周期性边界条件。
+- 对流：我们实现了二阶精度的"Van Leer"方案。
+- 封闭模型：我们目前实现了Smagorinsky涡粘模型。
+
+从长远来看，我们有兴趣扩展JAX-CFD以实现与相关研究相关的方法，例如：
+- 共位网格
+- 替代边界条件（例如，非周期性边界和浸没边界方法）
+- 更高阶的时间步进
+- 几何多重网格
+- 稳态模拟（例如，雷诺平均N-S方程RANS）
+- 在多个TPU/GPU上的分布式模拟
+
+我们欢迎在这些方面的任何合作！在开始重要工作之前，请通过GitHub或电子邮件与我们联系以协调。
+
+
+## 应用案例
+### 论文
+
+**Machine learning–accelerated computationalfluid dynamics**
+
+* https://www.pnas.org/doi/epdf/10.1073/pnas.2101784118
+
+### 模型结构
+
+该模型包括一个卷积神经网络，该网络控制在标准数值求解器的对流计算中的学习近似。ψ和u分别指的是被输送和输送速度分量。对于空间维度，有d*d个对流通量模块的副本，对应于每个空间方向中每个速度分量的通量。
+
+![alt text](asset/model_structure.png)
+
+### 算法原理
+CNN网络具有平移不变性。这意味着网络对输入数据的平移不敏感，即使图像中的对象位置发生了变化，网络也能识别出对象。在流体动力学的上下文中，网络能够识别出流动中相同的模式，无论它们在空间中如何移动。此外，卷积神经网络在空间上是局部化的，网络的操作集中在输入数据的局部区域。在流体动力学的背景下，网络专注于网格点周围的局部流速场，而不是整个流速场。这种局部化允许网络捕捉局部流动特征，如涡旋、边界层等，这对于准确模拟流体行为至关重要。
+
+![alt text](asset/alg.png)
+
+### 环境配置
+
+#### Docker（方法一）
+
+    docker pull image.sourcefind.cn:5000/dcu/admin/base/jax:0.4.23-ubuntu20.04-dtk24.04-py310
+
+    docker run --shm-size 10g --network=host --name=jax_cfd --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash
+
+    pip install -e ".[complete]"
+
+    pip install dask jupyter
+
+#### Dockerfile（方法二）
+
+    docker build -t <IMAGE_NAME>:<TAG> .
+
+    docker run --shm-size 10g --network=host --name=jax_cfd --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash
+
+    pip install -e ".[complete]"
+
+    pip install dask jupyter
+
+#### Anaconda（方法三）
+
+    DTK驱动：dtk24.04
+    python：python3.10
+    jax: 0.4.23
+
+Tips：以上dtk驱动、python、jax等DCU相关工具版本需要严格一一对应
+
+2、其他非特殊库
+
+    pip install dask jupyter
+
+### 数据集
+
+1、安装gsutil
+
+    curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-367.0.0-linux-x86_64.tar.gz
+    tar zxvf google-cloud-sdk-367.0.0-linux-x86_64.tar.gz
+    cd /安装目录/google-cloud-sdk
+    bash install.sh
+
+    # 添加环境变量
+    vim ~/.bashrc
+    export PATH=$PATH:/安装目录/google-cloud-sdk/bin
+    source ~/.bashrc
+
+2、下载数据
+
+    cd jax-cfd && mkdir content
+
+    gsutil -m cp gs://gresearch/jax-cfd/public_eval_datasets/kolmogorov_re_1000/eval_*.nc ./content
+
+    gsutil -m cp -r gs://gresearch/jax-cfd/public_eval_datasets/kolmogorov_re_1000_fig1 ./content
+
+    gsutil -m cp gs://gresearch/jax-cfd/public_eval_datasets/kolmogorov_re_1000/learned*.nc ./content
+
+    gsutil -m cp gs://gresearch/jax-cfd/public_eval_datasets/kolmogorov_re_1000/long_eval*.nc ./content
+
+
+### 推理
+
+1、下载模型
+
+    gsutil -m cp gs://gresearch/jax-cfd/public_models/*.pkl ./content
+
+2、执行
+
+    # 执行该目录下的ipynb中的代码
+    cd notebooks
+
+    # 采用交错网格的二维有限体积法模拟
+    demo.ipynb
+
+    # 使用伪谱法解算器的二维模拟
+    spectral_forced_turbulence.ipynb
+
+    # 二维通道流模拟
+    channel_flow_demo.ipynb
+
+    # 在共位网格上使用有限体积法（FVM）的二维模拟
+    collocated_demo.ipynb
+
+    # PNAS论文
+    # 数据分析及评估
+    ml_accelerated_cfd_data_analysis.ipynb
+
+    # 预训练模型
+    ml_model_inference_demo.ipynb
+    
+
+注意：跳过代码中需要下载数据的部分。
+
+### result
+
+![alt text](asset/result1.png)
+
+
+#### 精度
+
+无
+
+## 应用场景
+
+### 算法类别
+
+`流体动力学`
+
+### 热点应用行业
+
+`化工,气象,能源`  
+
+## 源码仓库及问题反馈
+
+* https://developer.hpccube.com/codes/modelzoo/jax-cfd_jax
+
+## 参考资料
+
+* https://github.com/google/jax-cfd/tree/main
+
+
+
+<!-- jupyter notebook --no-browser --ip=0.0.0.0 --allow-root -->
\ No newline at end of file

README_official.md

+# JAX-CFD: Computational Fluid Dynamics in JAX
+
+Authors: Dmitrii Kochkov, Jamie A. Smith, Peter Norgaard, Gideon Dresdner, Ayya Alieva, Stephan Hoyer
+
+JAX-CFD is an experimental research project for exploring the potential of
+machine learning, automatic differentiation and hardware accelerators (GPU/TPU)
+for computational fluid dynamics. It is implemented in
+[JAX](https://github.com/google/jax).
+
+To learn more about our general approach, read our paper [Machine learning accelerated computational fluid dynamics](https://www.pnas.org/content/118/21/e2101784118)
+(PNAS 2021).
+
+## Getting started
+
+The "notebooks" directory contains several demonstrations of using the JAX-CFD
+code.
+
+- Demos of different simulation setups:
+  - [2D simulation with FVM on a staggered grid](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/demo.ipynb)
+  - [2D simulation with a psuedo-spectral solver](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/spectral_forced_turbulence.ipynb)
+  - [2D simulation of channel flow](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/channel_flow_demo.ipynb)
+  - [2D simulation with FVM on a collocated grid](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/collocated_demo.ipynb) (experimental)
+
+- Reproduce results from our PNAS paper:
+  - [Data analysis and evaluation](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/ml_accelerated_cfd_data_analysis.ipynb)
+  - [Running our pre-trained models](https://colab.research.google.com/github/google/jax-cfd/blob/main/notebooks/ml_model_inference_demo.ipynb)
+
+## Organization
+
+JAX-CFD is organized around sub-modules:
+
+- `jax_cfd.base`: core finite volume/difference methods for CFD, written in JAX.
+- `jax_cfd.spectral`: core pseudospectral methods for CFD, written in JAX.
+- `jax_cfd.ml`: machine learning augmented models for CFD,
+  written in JAX and [Haiku](https://dm-haiku.readthedocs.io/en/latest/).
+- `jax_cfd.data`: data processing utilities for preparing, evaluating and
+  post-processing data created with JAX-CFD, written in
+  [Xarray](http://xarray.pydata.org/) and
+  [Pillow](https://pillow.readthedocs.io/).
+
+A base install with `pip install jax-cfd` only requires NumPy, SciPy and JAX.
+To install dependencies for the other submodules, use `pip install jax-cfd[ml]`,
+`pip install jax-cfd[data]` or `pip install jax-cfd[complete]`.
+
+## Numerics
+
+JAX-CFD is currently focused on unsteady turbulent flows:
+
+- *Spatial discretization*:
+  - *Finite volume/difference* methods on a staggered grid (the "Arakawa C" or
+  "MAC" grid) with pressure at the center of each cell and velocity components
+  defined on corresponding faces.
+  - *Pseudospectral* methods for vorticity which use anti-aliasing filtering
+  techniques for non-linear terms to maintain stability.
+- *Temporal discretization*: Currently only first-order temporal
+  discretization, using explicit time-stepping for advection and either implicit
+  or explicit time-stepping for diffusion.
+- *Pressure solves*: Either CG or fast diagonalization with real-valued FFTs
+  (suitable for periodic boundary conditions).
+- *Boundary conditions*: Currently only periodic boundary conditions are
+  supported.
+- *Advection*: We implement 2nd order accurate "Van Leer" schemes.
+- *Closures*: We currently implement Smagorinsky eddy-viscosity models.
+
+TODO: add a notebook explaining our numerical models in more depth.
+
+In the long term, we're interested in expanding JAX-CFD to implement methods
+relevant for related research, e.g.,
+
+- Colocated grids
+- Alternative boundary conditions (e.g., non-periodic boundaries and immersed
+  boundary methods)
+- Higher order time-stepping
+- Geometric multigrid
+- Steady state simulation (e.g., RANS)
+- Distributed simulations across multiple TPUs/GPUs
+
+We would welcome collaboration on any of these! Please reach out (either on
+GitHub or by email) to coordinate before starting significant work.
+
+## Projects using JAX-CFD
+
+- [Variational Data Assimilation with a Learned Inverse Observation Operator](https://github.com/googleinterns/invobs-data-assimilation)
+
+## Other awesome projects
+
+Other differentiable CFD codes compatible with deep learning:
+
+- [PhiFlow](https://github.com/tum-pbs/PhiFlow/) supports TensorFlow, PyTorch and JAX
+- [Fluid simulation in Autograd](https://github.com/HIPS/autograd#end-to-end-examples)
+
+JAX for science:
+
+- [JAX-MD](https://github.com/google/jax-md)
+- [JAX-DFT](https://github.com/google-research/google-research/tree/master/jax_dft)
+- [jax-cosmo](https://github.com/DifferentiableUniverseInitiative/jax_cosmo)
+- [Veros](https://github.com/team-ocean/veros)
+
+Did we miss something? Please let us know!
+
+## Citation
+
+If you use our finite volume method (FVM) or ML models, please cite:
+
+```
+@article{Kochkov2021-ML-CFD,
+  author = {Kochkov, Dmitrii and Smith, Jamie A. and Alieva, Ayya and Wang, Qing and Brenner, Michael P. and Hoyer, Stephan},
+  title = {Machine learning{\textendash}accelerated computational fluid dynamics},
+  volume = {118},
+  number = {21},
+  elocation-id = {e2101784118},
+  year = {2021},
+  doi = {10.1073/pnas.2101784118},
+  publisher = {National Academy of Sciences},
+  issn = {0027-8424},
+  URL = {https://www.pnas.org/content/118/21/e2101784118},
+  eprint = {https://www.pnas.org/content/118/21/e2101784118.full.pdf},
+  journal = {Proceedings of the National Academy of Sciences}
+}
+```
+
+If you use our spectral code, please cite:
+
+```
+@article{Dresdner2022-Spectral-ML,
+  doi = {10.48550/ARXIV.2207.00556},
+  url = {https://arxiv.org/abs/2207.00556},
+  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},
+  title = {Learning to correct spectral methods for simulating turbulent flows},
+  publisher = {arXiv},
+  year = {2022},
+  copyright = {arXiv.org perpetual, non-exclusive license}
+}
+```
+
+## Local development
+
+To locally install for development:
+```
+git clone https://github.com/google/jax-cfd.git
+cd jax-cfd
+pip install jaxlib
+pip install -e ".[complete]"
+```
+
+Then to manually run the test suite:
+```
+pytest -n auto jax_cfd --dist=loadfile --ignore=jax_cfd/base/validation_test.py
+```

jax_cfd/__init__.py

+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Defines the JAX-CFD module for computational fluid dynamics."""
+
+__version__ = '0.2.1'
+
+import jax_cfd.base

jax_cfd/base/__init__.py

+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Non-learned "base" physics routines for JAX-CFD."""
+
+import jax_cfd.base.advection
+import jax_cfd.base.array_utils
+import jax_cfd.base.boundaries
+import jax_cfd.base.diffusion
+import jax_cfd.base.equations
+import jax_cfd.base.fast_diagonalization
+import jax_cfd.base.finite_differences
+import jax_cfd.base.forcings
+import jax_cfd.base.funcutils
+import jax_cfd.base.grids
+import jax_cfd.base.initial_conditions
+import jax_cfd.base.interpolation
+import jax_cfd.base.pressure
+import jax_cfd.base.resize
+import jax_cfd.base.subgrid_models
+import jax_cfd.base.time_stepping
+import jax_cfd.base.validation_problems