README.md

# cfd

<!-- ## 仓库介绍

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

## 数据集

使用[SCNet](http://113.200.138.88:18080/aidatasets/project-dependency/jax_cfd) 高速下载通道下载，

或

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

注意：可以使用[SCNet](http://113.200.138.88:18080/aimodels/findsource-dependency/jax_cfd)高速下载通道下载。

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
    
注意：可以使用`jupyter notebook --no-browser --ip=0.0.0.0 --allow-root`启动运行环境。


## result

![alt text](asset/result1.png)


### 精度

无

## 应用场景

### 算法类别

`流体动力学`

### 热点应用行业

`化工,气象,能源`  

## 源码仓库及问题反馈

* https://developer.hpccube.com/codes/modelzoo/cfd_jax

## 参考资料

* https://github.com/google/jax-cfd/tree/main



<!-- jupyter notebook --no-browser --ip=0.0.0.0 --allow-root -->