# 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 bash pip install -e ".[complete]" pip install dask jupyter #### Dockerfile（方法二） docker build -t : . 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 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