diff --git a/MindFlow/acoustic/README.md b/MindFlow/acoustic/README.md index beebfcfe809ef9c1bb0b0af2eaaf95034f75baf4..178134fc47794719f9b9818e551b6afc7778ccc2 100644 --- a/MindFlow/acoustic/README.md +++ b/MindFlow/acoustic/README.md @@ -1,10 +1,10 @@ -# 2D Acoustic Wave Equation CBS Solver +# 2D/3D Acoustic Wave Equation CBS Solver ## Overview The solution of the acoustic wave equation is a core technology in fields such as medical ultrasound and geological exploration. Large-scale acoustic wave equation solving faces challenges in computing power and storage. Acoustic wave equation solvers generally use frequency domain solving algorithms and time domain solving algorithms. The representative of time domain solving algorithms is the Time Domain Finite Difference (TDFD) method, and the frequency domain solving algorithms include Frequency Domain Finite Difference (FDFD), Finite Element Method (FEM), and Convergent Born Series (CBS) iterative method. The CBS method is widely recognized in the engineering and academic communities due to its low memory requirements and the absence of dispersion errors. In particular, [Osnabrugge et al. (2016)](https://linkinghub.elsevier.com/retrieve/pii/S0021999116302595) solved the convergence problem of this method, making the application of the CBS method have broader prospects. The AI model based on the CBS computational structure is a typical representative of the dual-driven paradigm of physics and AI, including [Stanziola et al. (2022)](http://arxiv.org/abs/2212.04948), [Zeng et al. (2023)](http://arxiv.org/abs/2312.15575), etc. -This case will demonstrate how to call the CBS API provided by MindFlow to solve the 2D acoustic wave equation. +This case will demonstrate how to call the CBS API provided by MindFlow to solve the 2D/3D acoustic wave equation. ## Theoretical Background @@ -12,7 +12,7 @@ This case will demonstrate how to call the CBS API provided by MindFlow to solve In the solution of the acoustic wave equation, the velocity field and source information are input parameters, and the output is the spatiotemporal distribution of the wavefield. -The expression of the 2D acoustic wave equation is as follows: +The expression of the 2D/3D acoustic wave equation is as follows: | Time Domain Expression | Frequency Domain Expression | | ----------------------------------------------------- | ------------------------------------------------- | @@ -87,29 +87,31 @@ $$ The content is translated into English as follows: - Non-dimensionalization of input parameters; -- Non-dimensionalization of frequency domain 2D acoustic wave equation CBS solver; +- Non-dimensionalization of frequency domain 2D/3D acoustic wave equation CBS solver; - Dimensional restoration of the solution; - Time-frequency transformation of the solution. -The core solving process is parallelized for different source locations and different frequency points. Due to the large number of frequency points, it is divided into `n_batches` batches to solve sequentially along the frequency direction. The required input for the case is placed in the `dataset/` directory in the form of files, and the file names are passed in through `config.yaml`. The output results include the solution of the non-dimensionalized equation in the frequency domain `u_star.npy`, the dimensional final solution converted to the time domain `u_time.npy`, and the visualization animation of the time domain solution `wave.gif`. +The core solving process is parallelized for different source locations and different frequency points. Due to the large number of frequency points, it is divided into `n_batches` batches to solve sequentially along the frequency direction. The required input for the case is placed in the `dataset/` directory in the form of files, and the file names are passed in through `config_2d.yaml` (2D) or `config_3d.yaml` (3D). The output results include the solution of the non-dimensionalized equation in the frequency domain `u_star.npy`, the dimensional final solution converted to the time domain `u_time.npy`, and the visualization animation of the time domain solution `wave.gif`. + +The case introduces two new user-configurable parameters: `pml_size` and `dxs`. The `pml_size` parameter defines the thickness of the absorbing boundary (which may vary by direction), while `dxs` specifies the spatial discretization step size (which may also vary by direction). ## Quick Start -To facilitate direct verification by users, preset inputs are provided [here](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic). Please download the data and put them in `./dataset` in the case directory. The data include the velocity field `velocity.npy`, source location list `srclocs.csv`, and source waveform `srcwaves.csv`. Users can modify the input parameters according to the file format. +To facilitate direct verification by users, preset inputs are provided [here](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic). Please download the data and put them in `./dataset` in the case directory. The data include the velocity field `velocity_2d.npy`, source location list `srclocs_2d.csv`, and source waveform `srcwaves_2d.csv` in 2D and `velocity_3d.npy`, `srclocs_3d.csv`, and `srcwaves_3d.csv` in 3D. Users can modify the input parameters according to the file format. ### Method 1: Running the `solve_acoustic.py` script ```shell -python solve_acoustic.py --config_file_path ./config.yaml --device_id 0 --mode GRAPH +python solve_acoustic.py --dim 2 --device_id 0 --mode GRAPH ``` Where -`--config_file_path` represents the path of the configuration file, with a default value of `./config.yaml`; +`--dim` represents the dimension of space; `--device_id` represents the ID of the computing card used, which can be filled in according to the actual situation, and the most idle one will be automatically selected from all the computing cards by default; -`--mode` represents the running mode, 'GRAPH' represents static graph mode, and 'PYNATIVE' represents dynamic graph mode. +`--mode` represents the running mode, `GRAPH` represents static graph mode, and `PYNATIVE` represents dynamic graph mode. ### Method 2: Running Jupyter Notebook @@ -117,28 +119,46 @@ You can use the [Chinese version](./acoustic_CN.ipynb) and the [English version] ## Result Display +### 2D Model + The evolution of the wave field excited by different source locations for the same velocity model over time is shown in the following figure. -![wave.gif](images/wave.gif) +![wave_2d.gif](images/wave_2d.gif) The iterative convergence process of the equation residual is shown in the following figure, with each line representing a frequency point. The number of iterations required to reach the convergence threshold varies for different frequency points, and the number of iterations in the same batch depends on the slowest converging frequency point. -![errors.png](images/errors.png) +![errors_2d.png](images/errors_2d.png) + +### 3D Model + +Due to the large scale of the 3D model, we manually partitioned the frequency points across NPU cards for the purpose of demonstrating the correctness of numerical results: they were divided into 5 parts and distributed across NPU cards 0 to 4. Each card's frequency points were further split into 50 batches (automatically handled by the program). If users run the program directly, they will obtain numerical results based on a small-scale (coarse grid points) execution on a single NPU card. + +Distribute the frequency points across 5 NPU cards. The CBS iteration error convergence curves for each card are as follows + +![errors_3d](images/errors_3d.png) + +From left to right are the wave speed distribution, X-T acoustic pressure distribution, and Y-T acoustic pressure distribution. The X-T plot depicts the propagation of acoustic waves along the X-axis over time (with the vertical coordinate increasing downward) on the Z=0 plane, while the Y-T plot describes the propagation along the Y-axis under the same conditions. + +![velocity_pressure_3d](images/velocity_pressure_3d.png) + +Below is an animated visualization of wave propagation over time at the X-Z cross-section of the 3D environment (a single source is shown for clarity of wave propagation and reflection patterns). The animation demonstrates both direct wave propagation through subsurface layers and reflected waves from geological interfaces. + +sound_pressure_3d ## Performance -| Parameter | Ascend | -|:----------------------:|:--------------------------:| -| Hardware | Ascend NPU | -| MindSpore Version | >=2.3.0 | -| Dataset | [Marmousi velocity model](https://en.wikipedia.org/wiki/Marmousi_model) slices, included in the `dataset/` path of the case | -| Number of Parameters | No trainable parameters | -| Solver Parameters | batch_size=300, tol=1e-3, max_iter=10000 | -| Convergence Iterations | batch 0: 1320, batch 1: 560, batch 2: 620, batch 3: 680| -| Solver Speed (ms/iteration) | 500 | +| Parameter | 2D | 3D | +|:----------------------:|:--------------------------:|:----------------------:| +| Hardware | Ascend NPU | Ascend NPU | +| MindSpore Version | >=2.4.0 | >=2.4.0 | +| Dataset | [Marmousi velocity model](https://en.wikipedia.org/wiki/Marmousi_model) slices, included in the `dataset/` path of the case | 3D SEAM Barrett model, included in the `dataset/` path of the case | +| Number of Parameters | No trainable parameters | No trainable parameters | +| Solver Parameters | batch_size=300, tol=1e-3, max_iter=10000 | batch_size=437, tol=1e-3, max_iter=10000 | +| Convergence Iterations | batch 0: 1320, batch 1: 560, batch 2: 620, batch 3: 680| batch 0: 1180, batch 1: 540, batch 2: 460, batch 3: 460 | +| Solver Speed (ms/iteration) | 500 | 2800 | ## Contributors -gitee id: [WhFanatic](https://gitee.com/WhFanatic) +gitee id: [WhFanatic](https://gitee.com/WhFanatic), [zhaog6](https://gitee.com/zhaog6) -email: hainingwang1995@gmail.com \ No newline at end of file +email: hainingwang1995@gmail.com, zhaog6@lsec.cc.ac.cn diff --git a/MindFlow/acoustic/README_CN.md b/MindFlow/acoustic/README_CN.md index 7512de08d5b83c1f75b1887e9cb1397f1f10e015..71bdeff32acef8a4da90bfc2d2cde921b9c21d78 100644 --- a/MindFlow/acoustic/README_CN.md +++ b/MindFlow/acoustic/README_CN.md @@ -1,10 +1,10 @@ -# 2D 声波方程 CBS 求解 +# 2D/3D 声波方程 CBS 求解 ## 概述 声波方程求解是医疗超声、地质勘探等领域中的核心技术，大规模声波方程求解面临算力和存储的挑战。声波方程求解器一般采用频域求解算法和时域求解算法，时域求解算法的代表是时域有限差分法 (TDFD)，频域求解算法包括频域有限差分法 (FDFD)、有限元法 (FEM) 和 CBS (Convergent Born series) 迭代法。CBS 方法由于不引入频散误差，且求解的内存需求低，因此受到工程和学术界的广泛关注。尤其是 [Osnabrugge et al. (2016)](https://linkinghub.elsevier.com/retrieve/pii/S0021999116302595) 解决了该方法的收敛性问题，使得 CBS 方法的应用具有更广阔的前景。基于 CBS 的计算结构所提出的 AI 模型也是物理与 AI 双驱动范式的典型代表，包括 [Stanziola et al. (2022)](http://arxiv.org/abs/2212.04948)，[Zeng et al. (2023)](http://arxiv.org/abs/2312.15575) 等。 -本案例将演示如何调用 MindFlow 提供的 CBS API 实现二维声波方程的求解。 +本案例将演示如何调用 MindFlow 提供的 CBS API 实现二维/三维声波方程的求解。 ## 理论背景 @@ -12,7 +12,7 @@ 声波方程求解中，波速场和震源信息是输入参数，求解输出的是时空分布的波场。 -二维声波方程的表达式如下 +二维/三维声波方程的表达式如下 | 时域表达式 | 频域表达式 | | ----------------------------------------------------- | ------------------------------------------------- | @@ -87,31 +87,33 @@ $$ 具体包含以下步骤 - 输入参数无量纲化； -- 频域无量纲化 2D 声波方程 CBS 求解； +- 频域无量纲化 2D/3D 声波方程 CBS 求解； - 求解结果恢复量纲化； - 求解结果时频转换。 其中核心求解的过程针对不同震源位置和不同频点同时并行求解，由于频点数可能较多，因此沿频率方向分为 `n_batches` 个批次依次求解。 -案例所需的输入以文件的形式放置于 `dataset/` 中，文件名通过 `config.yaml` 传入。输出的结果为频域无量纲方程的解 `u_star.npy`、转换到时域的有量纲最终解 `u_time.npy`、针对时域解制作的可视化动图 `wave.gif`。 +案例所需的输入以文件的形式放置于 `dataset/` 中，文件名通过 `config_2d.yaml`（2D模型）和 `config_3d.yaml`（3D模型）传入。输出的结果为频域无量纲方程的解 `u_star.npy`、转换到时域的有量纲最终解 `u_time.npy`、针对时域解制作的可视化动图 `wave.gif`。 + +案例新增可供用户指定的参数 `pml_size` 和 `dxs`，前者用于描述吸收边界厚度（各方向可不同），后者用于描述空间离散步长（各方向可不同）。 ## 快速开始 -为了方便用户直接验证，本案例在本[链接](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic)中提供了预置的输入数据，请下载所需要的数据集，并保存在 `./dataset` 目录下。数据集包括速度场 `velocity.npy`、震源位置列表 `srclocs.csv`、震源波形 `srcwaves.csv`。用户可仿照输入文件格式自行修改输入参数。 +为了方便用户直接验证，本案例在本[链接](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic)中提供了预置的输入数据，请下载所需要的数据集，并保存在 `./dataset` 目录下。2D数据集包括速度场 `velocity_2d.npy`、震源位置列表 `srclocs_2d.csv`、震源波形 `srcwaves_2d.csv`，3D数据集包括速度场 `velocity_3d.npy`、震源位置列表 `srclocs_3d.csv`、震源波形 `srcwaves_3d.csv`。用户可仿照输入文件格式自行修改输入参数。 ### 运行方式一：`solve_acoustic.py` 脚本 ```shell -python solve_acoustic.py --config_file_path ./config.yaml --device_id 0 --mode GRAPH +python solve_acoustic.py --dim 2 --device_id 0 --mode GRAPH ``` 其中， -`--config_file_path`表示配置文件的路径，默认值'./config.yaml'； +`--dim`表示声波方程的维数，默认值为 `2`； `--device_id`表示使用的计算卡编号，可按照实际情况填写，默认从所有计算卡中自动选取最空闲的一张； -`--mode`表示运行的模式，'GRAPH'表示静态图模式, 'PYNATIVE'表示动态图模式。 +`--mode`表示运行的模式，`GRAPH`表示静态图模式, `PYNATIVE`表示动态图模式。 ### 运行方式二：运行 Jupyter Notebook @@ -119,28 +121,46 @@ python solve_acoustic.py --config_file_path ./config.yaml --device_id 0 --mode G ## 结果展示 +### 二维模型 + 针对同一个速度模型，不同震源位置激发的波场随时间演化过程如下图所示。 -![wave.gif](images/wave.gif) +![wave_2d.gif](images/wave_2d.gif) 方程残差的迭代收敛过程如下图所示，每根线代表一个频点。不同频点达到收敛阈值所需的迭代次数不同，同一批次的迭代次数取决于收敛最慢的频点。 -![errors.png](images/errors.png) +![errors_2d.png](images/errors_2d.png) + +### 三维模型 + +由于3D模型规模较大，这里仅为了展示数值结果的正确性，我们暂对频点手动划分到NPU卡：将其划分为5份，分别放入NPU 0号到4号这5张卡中，再将每一张卡中的频点分为50个批次（程序自动分）。如果用户直接运行程序，得到的将是单张NPU卡上跑的小规模（粗网格点）下的数值结果。 + +将频点分到5张NPU卡上，每张卡的CBS迭代误差收敛曲线图为 + +![errors_3d](images/errors_3d.png) + +如下从左至右分别为波速分布图、X-T声压分布图、Y-T声压分布图。其中X-T声压图描述了Z=0平面上声波随时间（纵坐标向下）沿X轴的传播过程，Y-T声压图描述了Z=0平面上声波随时间（纵坐标向下）沿Y轴的传播过程。 + +![velocity_pressure_3d](images/velocity_pressure_3d.png) + +如下为3D场景 X-Z 截面处的声波随时间演化动图（为看清声波的传播与反射，这里仅对一个震源进行演示），可以从中看到声波在地下的传播以及声波遇到地下介质的反射波。 + +sound_pressure_3d ## 性能 -| 参数 | Ascend | -|:----------------------:|:--------------------------:| -| 硬件资源 | 昇腾 NPU | -| MindSpore版本 | >=2.3.0 | -| 数据集 | [Marmousi 速度模型](https://en.wikipedia.org/wiki/Marmousi_model)切片，包含在案例 `dataset/` 路径中 | -| 参数量 | 无可学习参数 | -| 求解参数 | batch_size=300, tol=1e-3, max_iter=10000 | -| 收敛所需迭代数 | batch 0: 1320, batch 1: 560, batch 2: 620, batch 3: 680| -| 求解速度(ms/iteration) | 500 | +| 空间维数 | 2D | 3D | +|:----------------------:|:--------------------------:|:----------------------:| +| 硬件资源 | 昇腾 NPU | 昇腾 NPU | +| MindSpore版本 | >=2.4.0 | >=2.4.0 | +| 数据集 | [Marmousi 速度模型](https://en.wikipedia.org/wiki/Marmousi_model)切片，包含在案例 `dataset/` 路径中 | 3D SEAM Barrett模型，包含在案例 `dataset/` 路径中 | +| 参数量 | 无可学习参数 | 无可学习参数 | +| 求解参数 | batch_size=300, tol=1e-3, max_iter=10000 | batch_size=437, tol=1e-3, max_iter=10000 | +| 收敛所需迭代数 | batch 0: 1320, batch 1: 560, batch 2: 620, batch 3: 680| batch 0: 1180, batch 1: 540, batch 2: 460, batch 3: 460 | +| 求解速度(ms/iteration) | 500 | 2800 | ## 贡献者 -gitee id: [WhFanatic](https://gitee.com/WhFanatic) +gitee id: [WhFanatic](https://gitee.com/WhFanatic), [zhaog6](https://gitee.com/zhaog6) -email: hainingwang1995@gmail.com \ No newline at end of file +email: hainingwang1995@gmail.com, zhaog6@lsec.cc.ac.cn diff --git a/MindFlow/acoustic/acoustic.ipynb b/MindFlow/acoustic/acoustic.ipynb index 425cbeb80ebfa93098fe357f83ac7dd893f0887e..f8cdf44651d60f95187706bc854cfcc82ada95d7 100644 --- a/MindFlow/acoustic/acoustic.ipynb +++ b/MindFlow/acoustic/acoustic.ipynb @@ -5,15 +5,15 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# 2D acoustic problem\n", + "# 2D/3D acoustic problem\n", "\n", "[![Download Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_notebook.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/cfd_solver/mindspore_acoustic.ipynb) [![Download Sample Codes](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_download_code.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/cfd_solver/mindspore_acoustic.py) [![View Source Files](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_source.png)](https://gitee.com/mindspore/docs/blob/master/docs/mindflow/docs/source_zh_cn/cfd_solver/acoustic.ipynb)\n", "\n", "## Environment Installation\n", "\n", - "This case requires **MindSpore >= 2.3.0-rc2** version. Please refer to [MindSpore Installation](https://www.mindspore.cn/install) for details.\n", + "This case requires **MindSpore >= 2.4.0** version. Please refer to [MindSpore Installation](https://www.mindspore.cn/install) for details.\n", "\n", - "In addition, you need to install **MindFlow >=0.2.0** version. If it is not installed in the current environment, please follow the instructions below to choose the backend and version for installation." + "In addition, you need to install **MindFlow >=0.4.0** version. If it is not installed in the current environment, please follow the instructions below to choose the backend and version for installation." ] }, { @@ -22,7 +22,7 @@ "metadata": {}, "outputs": [], "source": [ - "mindflow_version = \"0.3.0\" # update if needed\n", + "mindflow_version = \"0.4.0\" # update if needed\n", "\n", "# Only NPU is supported.\n", "!pip uninstall -y mindflow-ascend\n", @@ -37,7 +37,7 @@ "\n", "Solving the acoustic wave equation is a core technology in fields such as medical ultrasound and geological exploration. Large-scale acoustic wave equation solvers face challenges in terms of computational power and storage. Solvers for the wave equation generally use either frequency domain algorithms or time domain algorithms. The representative time domain algorithm is the Time Domain Finite Difference (TDFD) method, while frequency domain algorithms include Frequency Domain Finite Difference (FDFD), Finite Element Method (FEM), and Convergent Born Series (CBS) iterative method. The CBS method, due to its low memory requirement and absence of dispersion error, has gained widespread attention in engineering and academia. In particular, [Osnabrugge et al. (2016)](https://linkinghub.elsevier.com/retrieve/pii/S0021999116302595) have addressed the convergence issue of this method, expanding the application prospects of the CBS method.\n", "\n", - "This case study will demonstrate how to invoke the CBS API provided by MindFlow to solve the two-dimensional acoustic wave equation." + "This case study will demonstrate how to invoke the CBS API provided by MindFlow to solve the two/three-dimensional acoustic wave equation." ] }, { @@ -49,7 +49,7 @@ "\n", "In solving the acoustic wave equation, the input parameters are the velocity field and source information, and the output is the spatiotemporal distribution of the wave field.\n", "\n", - "The expression of the two-dimensional acoustic wave equation is as follows\n", + "The expression of the two/three-dimensional acoustic wave equation is as follows\n", "\n", "| Time Domain Expression | Frequency Domain Expression |\n", "| ----------------------------------------------------- | ------------------------------------------------- |\n", @@ -85,7 +85,6 @@ "source": [ "import os\n", "import numpy as np\n", - "import scipy\n", "import pandas as pd\n", "import mindspore as ms\n", "from mindspore import Tensor\n", @@ -104,47 +103,73 @@ "source": [ "## Define input parameters and output sampling method\n", "\n", - "The required inputs for this case are dimensional 2D velocity field, source location list, and source waveform. The input file name is specified in the [config.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config.yaml) file. For user convenience, pre-set inputs are provided [here](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic). Please download the data and put them in `./dataset` in the case directory. The data include the velocity field `velocity.npy`, source location list `srclocs.csv`, and source waveform `srcwaves.csv`. Users can modify the input parameters based on the input file format.\n", + "The required inputs for this case are dimensional 2D/3D velocity field, source location list, and source waveform. The input file name is specified in the [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) (2D) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) (3D) file. For user convenience, pre-set inputs are provided [here](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic). Please download the data and put them in `./dataset` in the case directory. The data include the velocity field `velocity_2d.npy`, source location list `srclocs_2d.csv`, and source waveform `srcwaves_2d.csv` in 2D, and the velocity field `velocity_3d.npy`, source location list `srclocs_3d.csv`, and source waveform `srcwaves_3d.csv` in 3D. Users can modify the input parameters based on the input file format.\n", "\n", - "The output is a spatiotemporal distribution of the wavefield. To specify how the output is sampled in time and frequency, parameters such as `dt` and `nt` need to be specified in the [config.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config.yaml) file.\n", + "The output is a spatiotemporal distribution of the wavefield. To specify how the output is sampled in time and frequency, parameters such as `dt` and `nt` need to be specified in the [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) file.\n", "\n", "Since the sampling rate of the input source waveform in time may differ from the required output, interpolation needs to be performed." ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ms.set_context(device_target='Ascend', device_id=0, mode=ms.GRAPH_MODE)\n", "\n", - "config = load_yaml_config('config.yaml')\n", + "# Space dimension\n", + "dim = 2\n", + "\n", + "if dim == 2:\n", + " config = load_yaml_config(\"./config_2d.yaml\")\n", + "else:\n", + " config = load_yaml_config(\"./config_3d.yaml\")\n", "\n", "data_config = config['data']\n", "solve_config = config['solve']\n", "summary_config = config['summary']\n", "\n", - "# read time & frequency points\n", + "# Coarsen rate for reducing scale\n", + "rate = solve_config['coarsen_rate']\n", + "\n", + "# Read time & frequency points\n", "dt = solve_config['dt']\n", "nt = solve_config['nt']\n", - "ts = np.arange(nt) * dt\n", - "omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt)\n", + "dt *= rate # Increase the time iteraion step size\n", + "nt //= rate\n", + "receivers = None\n", + "\n", + "# Read velocity array\n", + "velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field']))\n", + "assert dim == len(velo.shape), \"dim and dimension of velocity should be equal\"\n", + "dx = data_config['velocity_dx']\n", + "dz = data_config['velocity_dz']\n", + "if dim == 2:\n", + " dxs = (rate*dz, rate*dx)\n", + " velo = velo[::rate, ::rate] # Coarsen velocity field\n", + "else:\n", + " dy = data_config['velocity_dy']\n", + " dxs = (rate*dz, rate*dy, rate*dx)\n", + " velo = velo[::rate, ::rate, ::rate] # Coarsen velocity field\n", "\n", - "# read source locations\n", + "# Read source locations\n", "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_locations']), index_col=0)\n", - "slocs = df[['y', 'x']].values # shape (ns, 2)\n", + "if dim == 2:\n", + " slocs = df[['y', 'x']].values # Shape (ns, 2)\n", + "else:\n", + " slocs = df[['z', 'y', 'x']].values # Shape (ns, 3)\n", "\n", - "# read & interp source wave\n", + "# Read & interp source wave\n", "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_wave']))\n", - "inter_func = scipy.interpolate.interp1d(df.t, df.f, bounds_error=False, fill_value=0)\n", - "src_waves = inter_func(ts) # shape (nt)\n", - "src_amplitudes = np.fft.rfft(src_waves) # shape (nt//2+1)\n", + "inter_func = interp1d(df.t, df.f, kind='cubic', bounds_error=False, fill_value=0)\n", "\n", - "# read velocity array\n", - "velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field']))\n", - "nz, nx = velo.shape\n", - "dx = data_config['velocity_dx']" + "ts = np.arange(nt) * dt\n", + "omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt)\n", + "\n", + "# Interpolation source wave\n", + "src_waves = inter_func(ts) # Shape (nt)\n", + "src_amplitudes = np.fft.rfft(src_waves) # Shape (nt//2+1)" ] }, { @@ -154,16 +179,16 @@ "source": [ "## Select desired frequency points\n", "\n", - "With the output sampling method determined, all the desired frequency points are in turn determined. However, in order to reduce computational load, it is also possible to select only a portion of the frequency points for calculation, while obtaining the remaining frequency points through interpolation. The specific frequency point downsampling method is specified by the `downsample_mode` and `downsample_rate` in the [config.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config.yaml) file. The default is no downsampling, which means solving all frequency points except $\\omega=0$." + "With the output sampling method determined, all the desired frequency points are in turn determined. However, in order to reduce computational load, it is also possible to select only a portion of the frequency points for calculation, while obtaining the remaining frequency points through interpolation. The specific frequency point downsampling method is specified by the `downsample_mode` and `downsample_rate` in the [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) (2D) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) (3D) file. The default is no downsampling, which means solving all frequency points except $\\omega=0$." ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ - "# select omegas\n", + "# Select omegas\n", "no = len(omegas_all) // solve_config['downsample_rate']\n", "\n", "if solve_config['downsample_mode'] == 'exp':\n", @@ -181,7 +206,7 @@ "source": [ "## Perform simulation\n", "\n", - "Define the relevant arrays as Tensors, call `solve_cbs()`, and execute the solution on the NPU. Due to memory limitations, the solution process is executed in batches in the frequency domain. The number of batches is specified by the user in `config.yaml` and does not need to be divisible by the number of frequency points (allowing the size of the last batch to be different from the other batches). After the solution is completed, the frequency domain solution results will be saved to the file `u_star.npy`." + "Define the relevant arrays as Tensors, call `solve_cbs()`, and execute the solution on the NPU. Due to memory limitations, the solution process is executed in batches in the frequency domain. The number of batches is specified by the user in `config_2d.yaml` or `config_3d.yaml` and does not need to be divisible by the number of frequency points (allowing the size of the last batch to be different from the other batches). The thickness and type of absorbing boundaries in each direction can be specified via parameters `pml_size` and `btype` in the file [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml). After the solution is completed, the frequency domain solution results will be saved to the file `u_star.npy`." ] }, { @@ -190,15 +215,23 @@ "metadata": {}, "outputs": [], "source": [ - "### send to NPU and perform computation\n", + "# Send to NPU and perform computation\n", "os.makedirs(summary_config['root_dir'], exist_ok=True)\n", "velo = Tensor(velo, dtype=ms.float32, const_arg=True)\n", - "cbs = CBS((nz, nx), remove_pml=False)\n", "\n", - "ur, ui = solve_cbs(cbs, velo, slocs, omegas_sel, dx=dx, n_batches=solve_config['n_batches']) # shape (ns, no, len(receiver_zs), nx)\n", + "dxs_nd = tuple(d / dxs[-1] for d in dxs) # Nondimensional dxs\n", + "pml_size = tuple(solve_config['pml_size'])\n", + "btype = solve_config['btype']\n", + "cbs = CBS(velo.shape, dxs=dxs_nd, pml_size=pml_size, alpha=1., rampup=12,\n", + " btype=btype, remove_pml=False)\n", + "\n", + "# Shape (ns, no, len(receiver_zs), nx) or (ns, no, len(receiver_zs), ny, nx)\n", + "ur, ui, errs = solve_cbs(\n", + " cbs, velo, slocs, omegas_sel, receivers=receivers, dxs=dxs, n_batches=solve_config['n_batches'])\n", "\n", - "u_star = np.squeeze(ur.numpy() + 1j * ui.numpy()) # shape (ns, no, len(krs), nx)\n", - "np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), u_star)" + "u_star = ur.numpy() + 1j * ui.numpy() # Shape (ns, no, len(krs), nx) or (ns, no, len(krs), ny, nx)\n", + "\n", + "np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), np.squeeze(u_star))" ] }, { @@ -208,7 +241,7 @@ "source": [ "## Post-processing\n", "\n", - "CBS solves the dimensionless frequency domain equation, but downstream tasks often require observing the evolution process of dimensional wavefields in the time domain. Therefore, the final solution is restored to dimensional and converted back to the time domain. The restoration method is given by $\\hat{u} = u^* hat{f} / \\omega^2$. If downsampling is performed on the frequency points in the \"Select desired frequency points\" step, interpolation along the frequency direction is required here to restore the solutions for all frequency points. Then, perform a Fourier inverse transform on the dimensional frequency domain wavefield $\\hat{u}$ to obtain the time domain wavefield $u$. Save the time domain wavefield to the file `u_time.npy`." + "CBS solves the dimensionless frequency domain equation, but downstream tasks often require observing the evolution process of dimensional wavefields in the time domain. Therefore, the final solution is restored to dimensional and converted back to the time domain. The restoration method is given by $\\hat{u} = u^* \\hat{f} / \\omega^2$. If downsampling is performed on the frequency points in the \"Select desired frequency points\" step, interpolation along the frequency direction is required here to restore the solutions for all frequency points. Then, perform a Fourier inverse transform on the dimensional frequency domain wavefield $\\hat{u}$ to obtain the time domain wavefield $u$. Save the time domain wavefield to the file `u_time.npy`." ] }, { @@ -217,12 +250,13 @@ "metadata": {}, "outputs": [], "source": [ - "# recover dimension and interpolate to full frequency domain\n", - "u_star /= omegas_sel.reshape(-1, 1, 1)**2\n", - "u_star = scipy.interpolate.interp1d(omegas_sel, u_star, axis=1, kind='cubic', bounds_error=False, fill_value=0)(omegas_all)\n", - "u_star *= src_amplitudes.reshape(-1, 1, 1)\n", + "# Recover dimension and interpolate to full frequency domain\n", + "ones = (1,) * dim\n", + "u_star /= omegas_sel.reshape(-1, *ones)**2\n", + "u_star = interp1d(omegas_sel, u_star, axis=1, kind='cubic', bounds_error=False, fill_value=0)(omegas_all)\n", + "u_star *= src_amplitudes.reshape(-1, *ones)\n", "\n", - "# transform to time domain\n", + "# Transform to time domain\n", "u_time = np.fft.irfft(u_star, axis=1)\n", "np.save(os.path.join(summary_config['root_dir'], 'u_time.npy'), u_time)" ] @@ -240,9 +274,13 @@ "metadata": {}, "outputs": [], "source": [ - "# visualize the result\n", + "# Visualize the result\n", "u_time = np.load(os.path.join(summary_config['root_dir'], 'u_time.npy'))\n", - "visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))" + "if dim == 2:\n", + " visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n", + "else:\n", + " visual.anim3d(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n", + "visual.plot_errs(errs, os.path.join(summary_config['root_dir'], 'errors.png'))" ] } ], diff --git a/MindFlow/acoustic/acoustic_CN.ipynb b/MindFlow/acoustic/acoustic_CN.ipynb index 1b631abc53ed49cabd3cc0113747c656da9a9dd2..ee90a19dff85cd03125e600b77dde47d0b99af2d 100644 --- a/MindFlow/acoustic/acoustic_CN.ipynb +++ b/MindFlow/acoustic/acoustic_CN.ipynb @@ -11,9 +11,9 @@ "\n", "## 环境安装\n", "\n", - "本案例要求 **MindSpore >= 2.3.0-rc2** 版本。具体请查看 [MindSpore安装](https://www.mindspore.cn/install)。\n", + "本案例要求 **MindSpore >= 2.4.0** 版本。具体请查看 [MindSpore安装](https://www.mindspore.cn/install)。\n", "\n", - "此外，你需要安装 **MindFlow >=0.2.0** 版本。如果当前环境还没有安装，请按照下列方式选择后端和版本进行安装。" + "此外，你需要安装 **MindFlow >=0.4.0** 版本。如果当前环境还没有安装，请按照下列方式选择后端和版本进行安装。" ] }, { @@ -22,7 +22,7 @@ "metadata": {}, "outputs": [], "source": [ - "mindflow_version = \"0.3.0\" # update if needed\n", + "mindflow_version = \"0.4.0\" # update if needed\n", "\n", "# Only NPU is supported.\n", "!pip uninstall -y mindflow-ascend\n", @@ -37,7 +37,7 @@ "\n", "声波方程求解是医疗超声、地质勘探等领域中的核心技术，大规模声波方程求解面临算力和存储的挑战。声波方程求解器一般采用频域求解算法和时域求解算法，时域求解算法的代表是时域有限差分法 (TDFD)，频域求解算法包括频域有限差分法 (FDFD)、有限元法 (FEM) 和 CBS (Convergent Born series) 迭代法。CBS 方法由于不引入频散误差，且求解的内存需求低，因此受到工程和学术界的广泛关注。尤其是 [Osnabrugge et al. (2016)](https://linkinghub.elsevier.com/retrieve/pii/S0021999116302595) 解决了该方法的收敛性问题，使得 CBS 方法的应用具有更广阔的前景。\n", "\n", - "本案例将演示如何调用 MindFlow 提供的 CBS API 实现二维声波方程的求解。" + "本案例将演示如何调用 MindFlow 提供的 CBS API 实现二维/三维声波方程的求解。" ] }, { @@ -49,7 +49,7 @@ "\n", "声波方程求解中，波速场和震源信息是输入参数，求解输出的是时空分布的波场。\n", "\n", - "二维声波方程的表达式如下\n", + "二维/三维声波方程的表达式如下\n", "\n", "| 时域表达式 | 频域表达式 |\n", "| ----------------------------------------------------- | ------------------------------------------------- |\n", @@ -85,7 +85,6 @@ "source": [ "import os\n", "import numpy as np\n", - "import scipy\n", "import pandas as pd\n", "import mindspore as ms\n", "from mindspore import Tensor\n", @@ -104,47 +103,73 @@ "source": [ "## 定义输入参数和输出采样方式\n", "\n", - "本案例所需的输入为有量纲 2D 速度场、震源位置列表、震源波形，在文件 [config.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config.yaml) 中指定输入文件名。为了方便用户直接验证，本案例在本[链接](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic)中提供了预置的输入数据，请下载所需要的数据集，并保存在 `./dataset` 目录下。数据集包括速度场 `velocity.npy`、震源位置列表 `srclocs.csv`、震源波形 `srcwaves.csv`。用户可仿照输入文件格式自行修改输入参数。\n", + "本案例所需的输入为有量纲 2D/3D 速度场、震源位置列表、震源波形，在文件 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) (2D) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) (3D) 中指定输入文件名。为了方便用户直接验证，本案例在本[链接](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic)中提供了预置的输入数据，请下载所需要的数据集，并保存在 `./dataset` 目录下。2D数据集包括速度场 `velocity_2d.npy`、震源位置列表 `srclocs_2d.csv`、震源波形 `srcwaves_2d.csv`，3D数据集包括速度场 `velocity_3d.npy`、震源位置列表 `srclocs_3d.csv`、震源波形 `srcwaves_3d.csv`。用户可仿照输入文件格式自行修改输入参数。\n", "\n", - "输出为时空分布的波场，为了明确输出如何在时间和频率上采样，需在 [config.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config.yaml) 文件中指定 `dt`, `nt` 等参数。\n", + "输出为时空分布的波场，为了明确输出如何在时间和频率上采样，需在 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) 或 [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) 文件中指定 `dt`, `nt` 等参数。\n", "\n", "由于输入的震源波形在时间上的采样率可能与输出所要求的不一致，因此需对其进行插值。" ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ms.set_context(device_target='Ascend', device_id=0, mode=ms.GRAPH_MODE)\n", "\n", - "config = load_yaml_config('config.yaml')\n", + "# Space dimension\n", + "dim = 2\n", + "\n", + "if dim == 2:\n", + " config = load_yaml_config(\"./config_2d.yaml\")\n", + "else:\n", + " config = load_yaml_config(\"./config_3d.yaml\")\n", "\n", "data_config = config['data']\n", "solve_config = config['solve']\n", "summary_config = config['summary']\n", "\n", - "# read time & frequency points\n", + "# Coarsen rate for reducing scale\n", + "rate = solve_config['coarsen_rate']\n", + "\n", + "# Read time & frequency points\n", "dt = solve_config['dt']\n", "nt = solve_config['nt']\n", - "ts = np.arange(nt) * dt\n", - "omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt)\n", + "dt *= rate # Increase the time iteraion step size\n", + "nt //= rate\n", + "receivers = None\n", + "\n", + "# Read velocity array\n", + "velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field']))\n", + "assert dim == len(velo.shape), \"dim and dimension of velocity should be equal\"\n", + "dx = data_config['velocity_dx']\n", + "dz = data_config['velocity_dz']\n", + "if dim == 2:\n", + " dxs = (rate*dz, rate*dx)\n", + " velo = velo[::rate, ::rate] # Coarsen velocity field\n", + "else:\n", + " dy = data_config['velocity_dy']\n", + " dxs = (rate*dz, rate*dy, rate*dx)\n", + " velo = velo[::rate, ::rate, ::rate] # Coarsen velocity field\n", "\n", - "# read source locations\n", + "# Read source locations\n", "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_locations']), index_col=0)\n", - "slocs = df[['y', 'x']].values # shape (ns, 2)\n", + "if dim == 2:\n", + " slocs = df[['y', 'x']].values # Shape (ns, 2)\n", + "else:\n", + " slocs = df[['z', 'y', 'x']].values # Shape (ns, 3)\n", "\n", - "# read & interp source wave\n", + "# Read & interp source wave\n", "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_wave']))\n", - "inter_func = scipy.interpolate.interp1d(df.t, df.f, bounds_error=False, fill_value=0)\n", - "src_waves = inter_func(ts) # shape (nt)\n", - "src_amplitudes = np.fft.rfft(src_waves) # shape (nt//2+1)\n", + "inter_func = interp1d(df.t, df.f, kind='cubic', bounds_error=False, fill_value=0)\n", "\n", - "# read velocity array\n", - "velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field']))\n", - "nz, nx = velo.shape\n", - "dx = data_config['velocity_dx']" + "ts = np.arange(nt) * dt\n", + "omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt)\n", + "\n", + "# Interpolation source wave\n", + "src_waves = inter_func(ts) # Shape (nt)\n", + "src_amplitudes = np.fft.rfft(src_waves) # Shape (nt//2+1)" ] }, { @@ -154,16 +179,16 @@ "source": [ "## 选取待求频点\n", "\n", - "确定了输出采样方式即确定了所有待求频点。但为了减少计算量，也可以只选择部分频点进行求解，其余频点通过插值获得。具体的频点降采样方式由 [config.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config.yaml) 文件中的 `downsample_mode`, `downsample_rate` 指定。默认为不做降采样，即求解除 $\\omega=0$ 之外的所有频点。" + "确定了输出采样方式即确定了所有待求频点。但为了减少计算量，也可以只选择部分频点进行求解，其余频点通过插值获得。具体的频点降采样方式由 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) 和 [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) 文件中的 `downsample_mode`, `downsample_rate` 指定。默认为不做降采样，即求解除 $\\omega=0$ 之外的所有频点。" ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ - "# select omegas\n", + "# Select omegas\n", "no = len(omegas_all) // solve_config['downsample_rate']\n", "\n", "if solve_config['downsample_mode'] == 'exp':\n", @@ -181,7 +206,7 @@ "source": [ "## 执行仿真\n", "\n", - "将相关数组定义为 Tensor，调用 `solve_cbs()`，在 NPU 执行求解。由于显存限制，求解过程在频点维度分批执行，batch 数量由用户在 `config.yaml` 中指定，不要求整除频点数（允许最后一个 batch 的 size 与其他 batch 不一致）。求解完成后，会保存频域求解结果到文件 `u_star.npy`。" + "将相关数组定义为 Tensor，调用 `solve_cbs()`，在 NPU 执行求解。由于显存限制，求解过程在频点维度分批执行，batch 数量由用户在 `config_2d.yaml` 或 `config_3d.yaml` 中指定，不要求整除频点数（允许最后一个 batch 的 size 与其他 batch 不一致）。各方向上的吸收边界厚度和边界类型可分别通过 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) 或 [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) 文件中的 `pml_size`, `btype` 来指定。求解完成后，会保存频域求解结果到文件 `u_star.npy`。" ] }, { @@ -190,15 +215,23 @@ "metadata": {}, "outputs": [], "source": [ - "### send to NPU and perform computation\n", + "# Send to NPU and perform computation\n", "os.makedirs(summary_config['root_dir'], exist_ok=True)\n", "velo = Tensor(velo, dtype=ms.float32, const_arg=True)\n", - "cbs = CBS((nz, nx), remove_pml=False)\n", "\n", - "ur, ui = solve_cbs(cbs, velo, slocs, omegas_sel, dx=dx, n_batches=solve_config['n_batches']) # shape (ns, no, len(receiver_zs), nx)\n", + "dxs_nd = tuple(d / dxs[-1] for d in dxs) # Nondimensional dxs\n", + "pml_size = tuple(solve_config['pml_size'])\n", + "btype = solve_config['btype']\n", + "cbs = CBS(velo.shape, dxs=dxs_nd, pml_size=pml_size, alpha=1., rampup=12,\n", + " btype=btype, remove_pml=False)\n", + "\n", + "# Shape (ns, no, len(receiver_zs), nx) or (ns, no, len(receiver_zs), ny, nx)\n", + "ur, ui, errs = solve_cbs(\n", + " cbs, velo, slocs, omegas_sel, receivers=receivers, dxs=dxs, n_batches=solve_config['n_batches'])\n", "\n", - "u_star = np.squeeze(ur.numpy() + 1j * ui.numpy()) # shape (ns, no, len(krs), nx)\n", - "np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), u_star)" + "u_star = ur.numpy() + 1j * ui.numpy() # Shape (ns, no, len(krs), nx) or (ns, no, len(krs), ny, nx)\n", + "\n", + "np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), np.squeeze(u_star))" ] }, { @@ -217,12 +250,13 @@ "metadata": {}, "outputs": [], "source": [ - "# recover dimension and interpolate to full frequency domain\n", - "u_star /= omegas_sel.reshape(-1, 1, 1)**2\n", - "u_star = scipy.interpolate.interp1d(omegas_sel, u_star, axis=1, kind='cubic', bounds_error=False, fill_value=0)(omegas_all)\n", - "u_star *= src_amplitudes.reshape(-1, 1, 1)\n", + "# Recover dimension and interpolate to full frequency domain\n", + "ones = (1,) * dim\n", + "u_star /= omegas_sel.reshape(-1, *ones)**2\n", + "u_star = interp1d(omegas_sel, u_star, axis=1, kind='cubic', bounds_error=False, fill_value=0)(omegas_all)\n", + "u_star *= src_amplitudes.reshape(-1, *ones)\n", "\n", - "# transform to time domain\n", + "# Transform to time domain\n", "u_time = np.fft.irfft(u_star, axis=1)\n", "np.save(os.path.join(summary_config['root_dir'], 'u_time.npy'), u_time)" ] @@ -240,9 +274,13 @@ "metadata": {}, "outputs": [], "source": [ - "# visualize the result\n", + "# Visualize the result\n", "u_time = np.load(os.path.join(summary_config['root_dir'], 'u_time.npy'))\n", - "visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))" + "if dim == 2:\n", + " visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n", + "else:\n", + " visual.anim3d(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n", + "visual.plot_errs(errs, os.path.join(summary_config['root_dir'], 'errors.png'))" ] } ], diff --git a/MindFlow/acoustic/config.yaml b/MindFlow/acoustic/config.yaml deleted file mode 100644 index 85b96730a7fe517351740aeddd41ed4d7e22bfce..0000000000000000000000000000000000000000 --- a/MindFlow/acoustic/config.yaml +++ /dev/null @@ -1,16 +0,0 @@ -data: - root_dir: 'dataset' - velocity_field: 'velocity.npy' - velocity_dx: 16. # grid interval of the velocity matrix - source_wave: 'srcwaves.csv' - source_locations: 'srclocs.csv' - -solve: - dt: 0.02 # time interval of the output - nt: 300 # number of time points of the output, must be even (required by rfft) - downsample_mode: 'linear' # way to downsample the frequency points, options: linear, exp, square - downsample_rate: 1 # only 1/downsample_rate frequency points will be solved - n_batches: 4 # the number of batches for frequencies to be diveded into - -summary: - root_dir: 'results' diff --git a/MindFlow/acoustic/config_2d.yaml b/MindFlow/acoustic/config_2d.yaml new file mode 100644 index 0000000000000000000000000000000000000000..74d356454be4a66e3a4a7c7c21c0a0c04ae92762 --- /dev/null +++ b/MindFlow/acoustic/config_2d.yaml @@ -0,0 +1,21 @@ +data: + root_dir: './dataset' + velocity_field: 'velocity_2d.npy' + velocity_dx: 16. # Grid interval of the velocity matrix on x axis, that is resolution on x axis + velocity_dz: 16. # Grid interval of the velocity matrix on z axis, that is resolution on z axis + source_wave: 'srcwaves_2d.csv' + source_locations: 'srclocs_2d.csv' + +solve: + dt: 0.02 # Time interval of the output + nt: 300 # Number of time points of the output, must be even (required by rfft) + coarsen_rate: 1 # Coarsen rate (<= space size) of velocity field and time point for NPU memory, + # satisfy (dz, dx) <- (coarsen_rate*dz, coarsen_rate*dx) + pml_size: [60, 60, 60, 60] # Size of absorbing boundary layer. The order is [z_up, z_down, x_left, x_right] + btype: 'pml' # Boundary type, choose one of pml, freesurface + downsample_mode: 'linear' # Way to downsample the frequency points, options: linear, exp, square + downsample_rate: 1 # Only 1/downsample_rate frequency points will be solved + n_batches: 4 # Number of batches for frequencies to be diveded into + +summary: + root_dir: 'results' diff --git a/MindFlow/acoustic/config_3d.yaml b/MindFlow/acoustic/config_3d.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f25b6d23b02e2e64e408793dbcddc99a0acb8f73 --- /dev/null +++ b/MindFlow/acoustic/config_3d.yaml @@ -0,0 +1,23 @@ +data: + root_dir: './dataset' + velocity_field: 'velocity_3d.npy' + velocity_dx: 12.5 # Grid interval of the velocity matrix on x axis, that is resolution on x axis + velocity_dy: 12.5 # Grid interval of the velocity matrix on y axis, that is resolution on y axis + velocity_dz: 6.25 # Grid interval of the velocity matrix on z axis, that is resolution on z axis + source_wave: 'srcwaves_3d.csv' + source_locations: 'srclocs_3d.csv' + +solve: + dt: 0.008 # Time interval of the output + nt: 875 # Number of time points of the output, must be even (required by rfft) + coarsen_rate: 4 # Coarsen rate (<= space size) of velocity field and time point for NPU memory, + # satisfy (dz, dy, dx) <- (coarsen_rate*dz, coarsen_rate*dy, coarsen_rate*dx) + pml_size: [100, 80, 50, 50, 50, 50] # Size of absorbing boundary layer. The order is + # [z_up, z_down, y_left, y_right, x_front, x_back] + btype: 'pml' # Boundary type, choose one of pml, freesurface + downsample_mode: 'linear' # Way to downsample the frequency points, options: linear, exp, square + downsample_rate: 1 # Only 1/downsample_rate frequency points will be solved + n_batches: 4 # Number of batches for frequencies to be diveded into + +summary: + root_dir: 'results' diff --git a/MindFlow/acoustic/solve_acoustic.py b/MindFlow/acoustic/solve_acoustic.py index 282f950228d9ebed8a59e608627e124f4010896d..219ae4157990f1d129264629d839f2c6db26abc4 100644 --- a/MindFlow/acoustic/solve_acoustic.py +++ b/MindFlow/acoustic/solve_acoustic.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -""""Solve 2D acoustic equation""""" +""""Solve 2D/3D acoustic equation""""" import os import argparse import numpy as np @@ -21,38 +21,59 @@ import pandas as pd import mindspore as ms from mindspore import ops, Tensor, numpy as mnp -from mindscience import CBS, load_yaml_config +from mindflow.utils import load_yaml_config + +from cbs.cbs import CBS from src import utils, visual -def solve_cbs(cbs, velo, slocs, omegas, receiver_zs=None, dx=1., n_batches=1): +def solve_cbs(cbs, velo, slocs, omegas, receivers=None, dxs=None, n_batches=1): ''' Solve for different source locations and frequencies using CBS (Convergent Born series) solver Args: - velo: 2d Tensor, the velocity field - slocs: (ns, 2) array, the source locations (z, x coordinates) to be solved - omegas: 1d array, the frequencies to be solved on - receiver_zs: 1d array, z coordinates of signal receivers. - Default is None, which means all signals will be received - dx: float, the grid interval along x & z directions - n_batches: int, the number of batches for frequencies to be diveded into + velo: 2d/3d Tensor, the velocity field. + slocs: (ns, 2) or (ns, 3) array, the source locations (z, x) or (z, y, x) coordinates + to be solved. + omegas: 1d array, the frequencies to be solved on. + receivers: Tuple of 1d array, (z, x) or (z, y, x) coordinates of signal receivers. + Default is None, which means all signals will be received. + dxs: Tuple of float, the grid interval along z & x (2D) or z & y & x (3D) directions. + Default is None. + n_batches: int, the number of batches for frequencies to be diveded into. Returns: u_real, u_imag: ''' no = len(omegas) ns = len(slocs) - nz, nx = velo.shape - - if receiver_zs is None: - receiver_zs = np.arange(nz) * dx + dim = len(velo.shape) + if dim == 2: + nz, nx = velo.shape + else: + nz, ny, nx = velo.shape + if dxs is None: + dxs = (1.0,) * dim + + if receivers is None: + if dim == 2: + receivers = (np.arange(nz) * dxs[0], np.arange(nx) * dxs[1]) + else: + receivers = (np.arange(nz) * dxs[0], np.arange(ny) * dxs[1], np.arange(nx) * dxs[2]) + + if dim == 2: + krzs = Tensor(np.rint(np.divide(receivers[0], dxs[0])), dtype=ms.int32, const_arg=False) + krxs = Tensor(np.rint(np.divide(receivers[1], dxs[1])), dtype=ms.int32, const_arg=False) + else: + krzs = Tensor(np.rint(np.divide(receivers[0], dxs[0])), dtype=ms.int32, const_arg=False) + krys = Tensor(np.rint(np.divide(receivers[1], dxs[1])), dtype=ms.int32, const_arg=False) + krxs = Tensor(np.rint(np.divide(receivers[2], dxs[2])), dtype=ms.int32, const_arg=False) - krs = Tensor(np.rint(np.divide(receiver_zs, dx)), dtype=ms.int32, const_arg=False) omegas = Tensor(omegas, dtype=ms.float32, const_arg=False) - masks = Tensor(utils.sloc2mask(slocs, (nz, nx), (dx, dx)), dtype=ms.float32, const_arg=False) # shape (ns, nz, nx) + # Shape (ns, nz, nx) for 2D and (ns, nz, ny, nx) for 3D + masks = Tensor(utils.sloc2mask(slocs, velo.shape, dxs), dtype=ms.float32, const_arg=False) - urs = [] # note: do hold the solution of each batch in list and cat to Tensor later - uis = [] # note: do not hold them by modifying Tensor slices, dynamic shape and error would be caused + urs = [] # Note: do hold the solution of each batch in list and cat to Tensor later + uis = [] # Note: do not hold them by modifying Tensor slices, dynamic shape and error would be caused errs = [] for n, i in enumerate(range(0, no, no // n_batches)): @@ -60,52 +81,97 @@ def solve_cbs(cbs, velo, slocs, omegas, receiver_zs=None, dx=1., n_batches=1): print(f'batch {n}, omega {float(omegas[i]):.4f} ~ {float(omegas[j-1]):.4f}') - c_star = velo / dx / omegas[i:j].reshape(-1, 1, 1) - f_star = masks.reshape(ns, 1, nz, nx) + if dim == 2: + c_star = velo / dxs[-1] / omegas[i:j].reshape(-1, 1, 1) + else: + c_star = velo / dxs[-1] / omegas[i:j].reshape(-1, 1, 1, 1) + + f_star = masks.reshape(ns, 1, *velo.shape) c_star, f_star = mnp.broadcast_arrays(c_star, f_star) - c_star = c_star.reshape(-1, 1, *c_star.shape[2:]) # shape (ns * no, 1, nz, nx) - f_star = f_star.reshape(-1, 1, *f_star.shape[2:]) # shape (ns * no, 1, nz, nx) + # Shape (ns * no, 1, nz, nx) for 2D and (ns * no, 1, nz, ny, nx) for 3D + c_star = c_star.reshape(-1, 1, *c_star.shape[2:]) + f_star = f_star.reshape(-1, 1, *f_star.shape[2:]) ur, ui, err = cbs.solve(c_star, f_star, tol=1e-3) - urs.append(ur[..., krs, :].reshape(ns, -1, len(krs), nx)) - uis.append(ui[..., krs, :].reshape(ns, -1, len(krs), nx)) + if dim == 2: + krzs_expand = krzs.reshape(-1, 1) + krxs_expand = krxs.reshape(1, -1) + urs.append(ur[..., krzs_expand, krxs_expand].reshape(ns, -1, len(krzs), len(krxs))) + uis.append(ui[..., krzs_expand, krxs_expand].reshape(ns, -1, len(krzs), len(krxs))) + else: + krzs_expand = krzs.reshape(-1, 1, 1) + krys_expand = krys.reshape(1, -1, 1) + krxs_expand = krxs.reshape(1, 1, -1) + urs.append(ur[..., krzs_expand, krys_expand, krxs_expand].reshape( + ns, -1, len(krzs), len(krys), len(krxs))) + uis.append(ui[..., krzs_expand, krys_expand, krxs_expand].reshape( + ns, -1, len(krzs), len(krys), len(krxs))) + errs.append(np.reshape(err, (-1, ns, j - i))) - u_real = ops.cat(urs, axis=1) # shape (ns, no, len(krs), nx) - u_imag = ops.cat(uis, axis=1) # shape (ns, no, len(krs), nx) + u_real = ops.cat(urs, axis=1) # Shape (ns, no, len(krs), nx) + u_imag = ops.cat(uis, axis=1) # Shape (ns, no, len(krs), nx) return u_real, u_imag, errs -def main(config): +def main(dim): + if dim == 2: + config = load_yaml_config("./config_2d.yaml") + elif dim == 3: + config = load_yaml_config("./config_3d.yaml") + else: + raise ValueError("The dim can only choose 2 or 3.") + data_config = config['data'] solve_config = config['solve'] summary_config = config['summary'] - # read time & frequency points + # Coarsen rate for reducing scale + rate = solve_config['coarsen_rate'] + + # Read time & frequency points dt = solve_config['dt'] nt = solve_config['nt'] - ts = np.arange(nt) * dt - omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt) + dt *= rate # Increase the time iteraion step size + nt //= rate + receivers = None - # read source locations + # Read velocity array + velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field'])) + if dim != len(velo.shape): + raise ValueError("The dim and dimension of velocity should be equal.") + dx = data_config['velocity_dx'] + dz = data_config['velocity_dz'] + if dim == 2: + dxs = (rate*dz, rate*dx) + velo = velo[::rate, ::rate] # Coarsen velocity field + else: + dy = data_config['velocity_dy'] + dxs = (rate*dz, rate*dy, rate*dx) + velo = velo[::rate, ::rate, ::rate] # Coarsen velocity field + + # Read source locations df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_locations']), index_col=0) - slocs = df[['y', 'x']].values # shape (ns, 2) + if dim == 2: + slocs = df[['y', 'x']].values # Shape (ns, 2) + else: + slocs = df[['z', 'y', 'x']].values # Shape (ns, 3) - # read & interp source wave + # Read & interp source wave df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_wave'])) - inter_func = interp1d(df.t, df.f, bounds_error=False, fill_value=0) - src_waves = inter_func(ts) # shape (nt) - src_amplitudes = np.fft.rfft(src_waves) # shape (nt//2+1) + inter_func = interp1d(df.t, df.f, kind='cubic', bounds_error=False, fill_value=0) - # read velocity array - velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field'])) - nz, nx = velo.shape - dx = data_config['velocity_dx'] + ts = np.arange(nt) * dt + omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt) + + # Interpolation source wave + src_waves = inter_func(ts) # Shape (nt) + src_amplitudes = np.fft.rfft(src_waves) # Shape (nt//2+1) - # select omegas + # Select omegas no = len(omegas_all) // solve_config['downsample_rate'] if solve_config['downsample_mode'] == 'exp': @@ -115,34 +181,45 @@ def main(config): else: omegas_sel = np.linspace(omegas_all[1], omegas_all[-1], no) - # send to NPU and perform computation + # Send to NPU and perform computation os.makedirs(summary_config['root_dir'], exist_ok=True) velo = Tensor(velo, dtype=ms.float32, const_arg=True) - cbs = CBS((nz, nx), remove_pml=False) + + # Solve Helmholtz equations with CBS + dxs_nd = tuple(d / dxs[-1] for d in dxs) # Nondimensional dxs + pml_size = tuple(solve_config['pml_size']) + btype = solve_config['btype'] + cbs = CBS(velo.shape, dxs=dxs_nd, pml_size=pml_size, alpha=1., rampup=12, + btype=btype, remove_pml=False) ur, ui, errs = solve_cbs( - cbs, velo, slocs, omegas_sel, dx=dx, n_batches=solve_config['n_batches']) # shape (ns, no, len(receiver_zs), nx) + cbs, velo, slocs, omegas_sel, receivers=receivers, dxs=dxs, n_batches=solve_config['n_batches']) + + u_star = ur.numpy() + 1j * ui.numpy() # Shape (ns, no, len(krzs), len(krys), len(krxs)) - u_star = np.squeeze(ur.numpy() + 1j * ui.numpy()) # shape (ns, no, len(krs), nx) - np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), u_star) + np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), np.squeeze(u_star)) - # recover dimension and interpolate to full frequency domain - u_star /= omegas_sel.reshape(-1, 1, 1)**2 + # Recover dimension and interpolate to full frequency domain + ones = (1,) * dim + u_star /= omegas_sel.reshape(-1, *ones)**2 u_star = interp1d(omegas_sel, u_star, axis=1, kind='cubic', bounds_error=False, fill_value=0)(omegas_all) - u_star *= src_amplitudes.reshape(-1, 1, 1) + u_star *= src_amplitudes.reshape(-1, *ones) - # transform to time domain + # Transform to time domain u_time = np.fft.irfft(u_star, axis=1) np.save(os.path.join(summary_config['root_dir'], 'u_time.npy'), u_time) - # visualize the result + # Visualize the result u_time = np.load(os.path.join(summary_config['root_dir'], 'u_time.npy')) - visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif')) + if dim == 2: + visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif')) + else: + visual.anim3d(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif')) visual.plot_errs(errs, os.path.join(summary_config['root_dir'], 'errors.png')) if __name__ == '__main__': - parser = argparse.ArgumentParser(description="Solve 2D acoustic equation with CBS") + parser = argparse.ArgumentParser(description="Solve 2D/3D acoustic equation with CBS") parser.add_argument( "--mode", type=str, @@ -156,7 +233,12 @@ if __name__ == '__main__': default=utils.choose_free_npu(), help="ID of the target device", ) - parser.add_argument("--config_file_path", type=str, default="./config.yaml") + parser.add_argument( + "--dim", + type=int, + default=3, + help="Dimension of acoustic equation" + ) args = parser.parse_args() ms.set_context( @@ -164,4 +246,4 @@ if __name__ == '__main__': device_id=args.device_id, mode=ms.GRAPH_MODE if args.mode.upper().startswith("GRAPH") else ms.PYNATIVE_MODE) - main(load_yaml_config(args.config_file_path)) + main(args.dim) diff --git a/mindscience/solvers/cbs.py b/mindscience/solvers/cbs.py index b6a2c17d4e14eeec135a521af924835ddc83d6da..1ad9108f54b6ae8ded4a67485e17e62479d2d546 100644 --- a/mindscience/solvers/cbs.py +++ b/mindscience/solvers/cbs.py @@ -13,26 +13,100 @@ # limitations under the License. # ============================================================================== """CBS solver for solving acoustic equation in frequency domain""" +import warnings from math import factorial +from time import time as toc import numpy as np import mindspore as ms from mindspore import Tensor, nn, ops, mint, numpy as mnp, lazy_inline +from mindscience import DFTn, IDFTn, DST, IDST -from ..sciops.fourier import DFTn, IDFTn + +class MixedDSTDFTn(nn.Cell): + ''' The mixed computation of DST and DFT in CBS iterations ''' + def __init__(self, shape): + """Configurations of mixed DST and DFT + + Args: + shape (tuple[int]): Only the spatial shape, not including the batch and channel dimensions. + """ + super().__init__() + self.dst_cell = DST(shape[:1]) + self.dft_cell = DFTn(shape[1:]) + + def construct(self, ar, ai): + ''' + Get mixed DST and DFT of the input + Args: + ar: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the real part of the input tensor. + ai: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the imaginary part of the input tensor. + ''' + dft_ar, dft_ai = self.dft_cell(ar, ai) + dft_ar = ops.swapdims(dft_ar, 2, -1) # Move z to last axis for DST + dft_ar = self.dst_cell(dft_ar) + dst_dft_ar = ops.swapdims(dft_ar, 2, -1) + dft_ai = ops.swapdims(dft_ai, 2, -1) + dft_ai = self.dst_cell(dft_ai) + dst_dft_ai = ops.swapdims(dft_ai, 2, -1) + + return dst_dft_ar, dst_dft_ai + + +class MixedIDSTDFTn(nn.Cell): + ''' The mixed computation of the inverse of DST and DFT in CBS iterations ''' + def __init__(self, shape): + """Configurations of the inverse of mixed DST and DFT + + Args: + shape (tuple[int]): Only the spatial shape, not including the batch and channel dimensions. + """ + super().__init__() + self.idst_cell = IDST(shape[:1]) + self.idft_cell = IDFTn(shape[1:]) + + def construct(self, ar, ai): + ''' + Get the inverse of mixed DST and DFT of the input + Args: + ar: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the real part of the input tensor. + ai: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the imaginary part of the input tensor. + ''' + ar = ops.swapdims(ar, 2, -1) + ar = self.idst_cell(ar) + idst_ar = ops.swapdims(ar, 2, -1) + ai = ops.swapdims(ai, 2, -1) + ai = self.idst_cell(ai) + idst_ai = ops.swapdims(ai, 2, -1) + idft_idst_ar, idft_idst_ai = self.idft_cell(idst_ar, idst_ai) + + return idft_idst_ar, idft_idst_ai class CBSBlock(nn.Cell): ''' The computation procedures for each iteration in CBS ''' @lazy_inline - def __init__(self, shape): + def __init__(self, shape, btype): ''' No trainable parameters, but the dft cells needs initialization Args: shape: Tuple of int, only the spatial shape, not including the batch and channel dimensions ''' super().__init__() - self.dft_cell = DFTn(shape) - self.idft_cell = IDFTn(shape) + self.btype = btype + if self.btype == "freesurface": + self.dft_cell = MixedDSTDFTn(shape) + self.idft_cell = MixedIDSTDFTn(shape) + else: + self.dft_cell = DFTn(shape) + self.idft_cell = IDFTn(shape) # Scattering potential calculation for real and imaginary parts def op_v(self, ur, ui, vr, vi): @@ -62,20 +136,34 @@ class CBSBlock(nn.Cell): return ops.stack([dur, dui]) - class CBS(nn.Cell): - ''' The CBS cell for solving 2D Helmholtz equation using the Convergente Born Series (CBS) + ''' The CBS cell for solving 2D/3D Helmholtz equation using the Convergent Born Series (CBS) method described in `Osnabrugge et al, 2016 `_ . - It is an iterative solver with spectral precision due to the Fourier approximation to the differential terms. - This cell can be used for acoustic wave simulation, - or be adapted to physics loss for self-supervised training of neural operators. + It is an iterative solver with spectral precision due to the Fourier approximation to + the differential terms. This cell can be used for acoustic wave simulation, or be adapted to + physics loss for self-supervised training of neural operators. ''' + def __init__(self, + shape, + dxs=None, + n_iter=20, + pml_size=None, + alpha=1.0, + rampup=12, + btype="pml", + remove_pml=True, + epsilon=None, + ): + """Configurations of the CBS solver Args: shape (tuple[int]): Only the spatial shape, not including the batch and channel dimensions. + dxs (tuple[float], optional): The grid interval along z & x (2D) or z & y & x (3D) directions + under nondimensionalization. Defaults to None, it represents assigning a default value + `(1,)*len(shape)`. n_iter (int, optional): Number of iterations in a single call. Defaults to 20. pml_size (int, tuple[int], optional): Number of grid layers to pad on each boundary for the wave - to attenuate. The order is (z_up, z_down, x_left, x_right) for 2D and (z_up, z_down, y_left, y_right, - x_front, x_back) for 3D. Defaults to None. + to attenuate. The order is (z_up, z_down, x_left, x_right) for 2D and (z_up, z_down, y_left, y_right, + x_front, x_back) for 3D. Defaults to None. 1) If None, assign a default value (If `btype` is `"pml"`, the size of absorption boundary in each direction is `shape[-1]//4`, if `btype` is `"freesurface"`, except for the upper surface absorption boundary size of 0, the rest of the absorption boundary sizes are @@ -85,82 +173,117 @@ class CBS(nn.Cell): individually. alpha (float, optional): The strength of wave attenuation in PML layers. Defaults to 1.0. rampup (int, optional): The smoothness of transition from interior domain to PML layers. Defaults to 12. + btype (string, optional): Boundary type (choose one of) pml freesurface. Defaults to 'pml'. remove_pml (bool, optional): Whether to remove the PML layers for the output. Defaults to True. epsilon (float, optional): The small value to stabilize the iteration. Defaults to None, calculating epsilon automatically. - ''' - def __init__(self, - shape, - n_iter=20, - pml_size=60, - alpha=1.0, - rampup=12, - remove_pml=True, - epsilon=None, - ): - """Configurations of the CBS solver""" + """ super().__init__() + if (dxs is not None) and (len(dxs) != len(shape)): + raise ValueError("Length of the shape and dxs should be equal.") + if btype not in ["pml", "freesurface"]: + raise ValueError("The boundary type should be one of `pml` and `freesurface`.") + + self.shape = shape self.n_iter = n_iter - self.pml_size = pml_size self.alpha = alpha self.rampup = rampup + self.btype = btype self.remove_pml = remove_pml self.epsilon = epsilon + self.dim = len(shape) + + if pml_size is None: + s = shape[-1] // 4 + self.pml_size = (s * int(btype == "pml"),) + (s,) * (2 * self.dim - 1) + elif isinstance(pml_size, int): + self.pml_size = (pml_size * int(btype == "pml"),) + (pml_size,) * (2 * self.dim - 1) + else: + self.pml_size = (int(btype == "pml") * pml_size[0],) + pml_size[1:] + + self.pml_size = list(self.pml_size) + for i, (a, b) in enumerate(zip(self.pml_size, [d for pair in zip(shape, shape) for d in pair])): + if a > b: + self.pml_size[i] = b + warnings.warn("The size of absorbing boundary should not be larger than the original " + f"grid ({a} > {b}), CBS.pml_size[{i}] will be assigned as {b}.", + category=UserWarning) + self.pml_size = tuple(self.pml_size) + + if dxs is None: + self.dxs = (1.0,) * self.dim + else: + self.dxs = dxs + + # Coordination reverse (x, y, z) -> (z, y, x) + self.pml_size_xyz = [ + num + for i in range(len(self.pml_size)-2, -2, -2) + for num in self.pml_size[i:i+2] + ] + + self.shape_padded = tuple(n + s1 + s2 for n, s1, s2 in zip(shape, self.pml_size[0::2], self.pml_size[1::2])) + + if self.btype == "freesurface": + # DST wave number for 1st-order derivative + fftfreq = [(np.pi / self.dxs[0] / self.shape_padded[0] * np.arange(1, self.shape_padded[0]+1))**2] + fftfreq += [4 * np.pi**2 * np.fft.fftfreq(n, d)**2 for n, d in zip(self.shape_padded[1:], self.dxs[1:])] + p_sq = sum(np.meshgrid(*fftfreq, indexing="ij")) + else: + p_sq = sum(np.meshgrid( + *[np.fft.fftfreq(n, d)**2 for n, d in zip(self.shape_padded, self.dxs)], + indexing="ij")) * (2 * np.pi)**2 - shape_padded = tuple(n + 2 * pml_size for n in shape) - - dxs = (1.0, 1.0) - p_sq = sum(np.meshgrid(*[np.fft.fftfreq(n, d)**2 - for n, d in zip(shape_padded, dxs)], indexing="ij")) * (2 * np.pi)**2 self.p_sq = Tensor(p_sq, dtype=ms.float32, const_arg=True) - pml_mask = 1 - np.pad(np.ones(shape), pml_size) - self.pml_mask = Tensor(pml_mask, dtype=ms.float32, const_arg=True) + self.pml_mask = 1 - ops.pad(ops.ones(shape), self.pml_size_xyz) - self.cbs_block = CBSBlock(shape_padded) + self.cbs_block = CBSBlock(self.shape_padded, self.btype) def cbs_params(self, c_star, f_star): - ''' compute constant variables for CBS iteration ''' - pml_size = self.pml_size - nz, nx = c_star.shape[-2:] - dxs = (1.0, 1.0) + ''' Compute constant variables for CBS iteration ''' omg = 1.0 - # source field, shape is (batch, 1, nz_padded, nx_padded) - rhs = ops.pad(f_star / c_star**2, [pml_size] * 4) + # Source field + # Shape (batch, 1, nz_padded, nx_padded) for 2D or (batch, 1, nz_padded, ny_padded, nx_padded) for 3D + rhs = ops.pad(f_star / c_star**2, self.pml_size_xyz) - # homogeneous k field - k_max = omg / ops.amin(c_star, axis=(-2, -1), keepdims=True) - k_min = omg / ops.amax(c_star, axis=(-2, -1), keepdims=True) - k_padded = self.mypadding(omg / c_star, [pml_size] * 4, values=k_max) - k0 = ops.sqrt(0.5 * (k_max**2 + k_min**2)) # shape is (batch, 1, 1, 1) + # Homogeneous k field + axis = tuple(range(-self.dim, 0, 1)) + k_max = omg / ops.amin(c_star, axis=axis, keepdims=True) + k_min = omg / ops.amax(c_star, axis=axis, keepdims=True) + k0 = ops.sqrt(0.5 * (k_max**2 + k_min**2)) # (batch, 1, 1, 1) + # Mirror padding gradually approaching `values` + k_padded = self.extrapolate(omg / c_star, self.pml_size_xyz, values=k_max) - # heterogeneous k field, shape is (batch, 1, nz_padded, nx_padded) - ksq_r, ksq_i = self.cbs_pml((nz, nx), dxs, k_padded, pml_size, self.alpha, self.rampup) + # Heterogeneous k field + # Shape (batch, 1, nz_padded, nx_padded) for 2D or (batch, 1, nz_padded, ny_padded, nx_padded) for 3D + ksq_r, ksq_i = self.cbs_pml(self.shape, self.dxs, k_padded, self.pml_size, + self.alpha, self.rampup) - ksq_r = ksq_r * self.pml_mask + ops.pad((omg / c_star)**2, [pml_size] * 4) * (1 - self.pml_mask) + ksq_r = ksq_r * self.pml_mask + ops.pad((omg / c_star)**2, self.pml_size_xyz) * (1 - self.pml_mask) ksq_i = ksq_i * self.pml_mask - eps = ops.amax((ksq_r - k0**2)**2 + ksq_i**2, axis=(-2, -1), keepdims=True)**.5 # shae is (batch, 1, 1, 1) + # Shape (batch, 1, 1, 1) for 2D, (batch, 1, 1, 1, 1) for 3D + eps = ops.amax((ksq_r - k0**2)**2 + ksq_i**2, axis=axis, keepdims=True)**.5 - # if epsilon given by user, use original BS instead of CBS + # If epsilon given by user, use original BS instead of CBS if isinstance(self.epsilon, (float, int)): eps = self.epsilon * ops.ones_like(eps) - # field variables needed by operator V & G - vr = ksq_r - k0**2 # shape is (batch, 1, nz_padded, nx_padded) - vi = ksq_i - eps # shape is (batch, 1, nz_padded, nx_padded) - gr = 1. / ((self.p_sq - k0**2)**2 + eps**2) * (self.p_sq - k0**2) # shape is (batch, 1, nz_padded, nx_padded) - gi = 1. / ((self.p_sq - k0**2)**2 + eps**2) * eps # shape is (batch, 1, nz_padded, nx_padded) + # Field variables needed by operator V & G + # Shape (batch, 1, nz_padded, nx_padded) for 2D or (batch, 1, nz_padded, ny_padded, nx_padded) for 3D + vr = ksq_r - k0**2 + vi = ksq_i - eps + gr = 1. / ((self.p_sq - k0**2)**2 + eps**2) * (self.p_sq - k0**2) + gi = 1. / ((self.p_sq - k0**2)**2 + eps**2) * eps return vr, vi, gr, gi, rhs, eps * (self.epsilon is None) @staticmethod def cbs_pml(shape, dxs, k0, pml_size, alpha, rampup): - ''' attenuate wavefield near boundary by modifying k ''' - shape_padded = tuple(n + 2 * pml_size for n in shape) - + ''' Construct the heterogeneous k field with PML BC embedded ''' def num(x): num_real = (alpha ** 2) * (rampup - alpha * x) * ((alpha * x) ** (rampup - 1)) num_imag = (alpha ** 2) * (2 * k0 * x) * ((alpha * x) ** (rampup - 1)) @@ -175,10 +298,18 @@ class CBS(nn.Cell): transform_real, transform_imag = num_real / den_x, num_imag / den_x return transform_real, transform_imag + def pml_padding(n, d, s1, s2): + original = (ops.abs(mnp.linspace(1 - n, n - 1, n)) - n) * (d / 2) + left = ops.arange((s1 - 0.5)*d, -d/2, -d, dtype=ms.float32) + right = ops.arange(d/2, (s2 + 0.5)*d, d, dtype=ms.float32) + return ops.concat((left, original, right)) + diff = ops.stack(mnp.meshgrid( - *[((ops.abs(mnp.linspace(1 - n, n - 1, n)) - n) / 2 + pml_size) * d for n, d in zip(shape_padded, dxs)], - indexing="ij"), axis=0) + *[pml_padding(n, d, s1, s2) + for n, d, s1, s2 in zip(shape, dxs, pml_size[0::2], pml_size[1::2])], indexing="ij"), axis=0) + # Reference: https://github.com/ucl-bug/jwave/blob/9d114ac6acade3c866e78a9984be47833119a9f9/ + # jwave/acoustics/time_harmonic.py#L153 diff *= (diff > 0).astype(ms.float32) / 4. dist = ops.norm(diff, dim=0) @@ -189,8 +320,8 @@ class CBS(nn.Cell): return ksq_r, ksq_i @staticmethod - def mypadding(data, padding, values): - ''' mirror padding gradually approaching `values` ''' + def extrapolate(data, padding, values): + ''' Mirror padding gradually approaching `values` ''' ndim = len(padding) // 2 for d in range(ndim): n0 = data.shape[-1 - d] # current dimension length @@ -201,10 +332,10 @@ class CBS(nn.Cell): # build matrix accessing mirror elements near boundary and applying attenuation weights = np.ones(n) weights[:n1] = np.arange(n1) / (n1 - 1) - weights[-n2:] = np.arange(n2)[::-1] / (n2 - 1) + weights[n1+n0:] = np.arange(n2)[::-1] / (n2 - 1) indices = n0 - 1 - np.abs(n0 - 1 - np.abs(np.arange(n) - n1)) indices[:n1] -= 1 - indices[-n2:] += 1 + indices[n1+n0:] += 1 flipmat = np.zeros([n, n0]) flipmat[range(n), indices] = weights @@ -218,38 +349,51 @@ class CBS(nn.Cell): return data + def compute_error(self, dur, dui, ur, ui): + ''' Compute the relative error of du ''' + dim = tuple(range(-self.dim, 0, 1)) + return (ops.sum(dur**2 + dui**2, dim=dim) / ops.sum(ur**2 + ui**2, dim=dim))**.5 + def construct(self, c_star, f_star, ur_init=None, ui_init=None): ''' Run the solver to solve non-dimensionalized 2D/3D acoustic equation for given c* and f* Args: - c_star: float (batch_size, 1, nz, nx) for 2D, the non-dimensionalized velocity field. - f_star: float (batch_size, 1, nz, nx) for 2D, the mask marking out the source locations. - ur_init, ui_init: float (batch_size, 1, NZ, NX) for 2D, initial wave field for iteration, - real & imag parts. - In 2D, if remove_pml is True, NZ = nz, NX = nx, otherwise - NZ = nz + pml_size[0] + pml_size[1], NX = nx + pml_size[2] + pml_size[3]. - In 3D, if remove_pml is True, NZ = nz, NY = ny, NX = nx, otherwise - NZ = nz + pml_size[0] + pml_size[1], NY = ny + pml_size[2] + pml_size[3], - NX = nx + pml_size[4] + pml_size[5]. - Default is None, which means initialize from 0. + c_star: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the non-dimensionalized velocity field. + f_star: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the mask marking out the source locations. + ur_init, ui_init: float (batch_size, 1, NZ, NX) for 2D, + float (batch_size, 1, NZ, NY, NX) for 3D, + initial wave field for iteration, real & imag parts. + In 2D, if remove_pml is True, NZ = nz, NX = nx, otherwise + NZ = nz + pml_size[0] + pml_size[1], NX = nx + pml_size[2] + pml_size[3]. + In 3D, if remove_pml is True, NZ = nz, NY = ny, NX = nx, otherwise + NZ = nz + pml_size[0] + pml_size[1], NY = ny + pml_size[2] + pml_size[3], + NX = nx + pml_size[4] + pml_size[5]. + Default is None, which means initialize from 0. ''' vr, vi, gr, gi, rhs, eps = self.cbs_params(c_star, f_star) - n0 = self.remove_pml * self.pml_size - n1 = (ur_init is None or self.remove_pml) * self.pml_size - n2 = (ui_init is None or self.remove_pml) * self.pml_size + # Padding under x-y-z coordination + n0 = [self.remove_pml * s for s in self.pml_size] + paddingr = [(ur_init is None or self.remove_pml) * s for s in self.pml_size_xyz] + paddingi = [(ui_init is None or self.remove_pml) * s for s in self.pml_size_xyz] - # construct initial field + # Construct initial field if ur_init is None: - ur_init = ops.zeros_like(c_star, dtype=ms.float32) # shape is (batch, 1, nz, nx) + ur_init = ops.zeros_like(c_star, dtype=ms.float32) # (batch, 1, nz, nx) if ui_init is None: - ui_init = ops.zeros_like(c_star, dtype=ms.float32) # shape is (batch, 1, nz, nx) + ui_init = ops.zeros_like(c_star, dtype=ms.float32) # (batch, 1, nz, nx) - # pad initial field - ur = ops.pad(ur_init, padding=[n1] * 4, value=0) # shape is (batch, 1, nz_padded, nx_padded) - ui = ops.pad(0 - ui_init, padding=[n2] * 4, value=0) # conjugate input as output + # Pad initial field + # Note: here u_init is conjugated, because the output is also conjugated + ur = ops.pad(ur_init, padding=paddingr, value=0) # Note: better padding (with gradual damping) can be applied + # Shape (batch, 1, nz_padded, nx_padded) for 2D, (batch, 1, nz_padded, ny_padded, nx_padded) for 3D + ui = ops.pad(0 - ui_init, padding=paddingi, value=0) - # start iteration + # Start iteration errs_list = [] for _ in range(self.n_iter): @@ -257,17 +401,22 @@ class CBS(nn.Cell): ur = ur + dur ui = ui + dui - # calculate iteration residual - errs = (ops.sum(dur**2 + dui**2, dim=(-2, -1)) / ops.sum(ur**2 + ui**2, dim=(-2, -1)))**.5 + # Calculate iteration residual + errs = self.compute_error(dur, dui, ur, ui) errs_list.append(errs) - # remove pml layer - nz, nx = ur.shape[-2:] - ur = ur[..., n0:nz - n0, n0:nx - n0] - ui = ui[..., n0:nz - n0, n0:nx - n0] + # Remove pml layer + if self.dim == 2: + nz, nx = ur.shape[-self.dim:] + ur = ur[..., n0[0]:nz - n0[1], n0[2]:nx - n0[3]] + ui = ui[..., n0[0]:nz - n0[1], n0[2]:nx - n0[3]] + else: + nz, ny, nx = ur.shape[-self.dim:] + ur = ur[..., n0[0]:nz - n0[1], n0[2]:ny - n0[3], n0[4]:nx - n0[5]] + ui = ui[..., n0[0]:nz - n0[1], n0[2]:ny - n0[3], n0[4]:nx - n0[5]] ui *= -1. - # note: conjugation here is because our definition of Fourier mode differs with that of jwave - # in that the frequency is opposite, leading to an opposite PML implementation + # Note: the conjugate here is because we define Fourier modes differently to JAX in that the frequencies + # are opposite, leading to opposite attenuation in PML, and finally the conjugation in results return ur, ui, errs_list @@ -284,23 +433,30 @@ class CBS(nn.Cell): """A convenient method for solving the equation to a given tolerance Args: - c_star: float (batch_size, 1, nz, nx) for 2D, the non-dimensionalized velocity field. - f_star: float (batch_size, 1, nz, nx) for 2D, the mask marking out the source locations. - ur_init, ui_init: float (batch_size, 1, NZ, NX) for 2D, initial wave field for iteration, - real & imag parts. - In 2D, if remove_pml is True, NZ = nz, NX = nx, otherwise - NZ = nz + pml_size[0] + pml_size[1], NX = nx + pml_size[2] + pml_size[3]. - In 3D, if remove_pml is True, NZ = nz, NY = ny, NX = nx, otherwise - NZ = nz + pml_size[0] + pml_size[1], NY = ny + pml_size[2] + pml_size[3], - NX = nx + pml_size[4] + pml_size[5]. - Default is None, which means initialize from 0. + c_star: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the non-dimensionalized velocity field. + f_star: float (batch_size, 1, nz, nx) for 2D, + float (batch_size, 1, nz, ny, nx) for 3D, + the mask marking out the source locations. + ur_init, ui_init: float (batch_size, 1, NZ, NX) for 2D, + float (batch_size, 1, NZ, NY, NX) for 3D, + initial wave field for iteration, real & imag parts. + In 2D, if remove_pml is True, NZ = nz, NX = nx, otherwise + NZ = nz + pml_size[0] + pml_size[1], NX = nx + pml_size[2] + pml_size[3]. + In 3D, if remove_pml is True, NZ = nz, NY = ny, NX = nx, otherwise + NZ = nz + pml_size[0] + pml_size[1], NY = ny + pml_size[2] + pml_size[3], + NX = nx + pml_size[4] + pml_size[5]. + Default is None, which means initialize from 0. tol (float, optional): The tolerance for the relative error. Defaults to 1e-3. max_iter (int, optional): The maximum iterations. Defaults to 10000. remove_pml (bool, optional): Whether to remove the PML layers for the output. Defaults to True. print_info (bool, optional): Whether to print the information of solution. Defaults to True. """ - assert not self.remove_pml, \ - 'PML layers cannot be removed during iteration, but can be removed for the final result' + msg = 'PML layers cannot be removed during iteration, but can be removed for the final result' + assert not self.remove_pml, msg + + tic = toc() ur, ui, errs_list = self(c_star, f_star, ur_init, ui_init) @@ -310,17 +466,30 @@ class CBS(nn.Cell): err_ave = float(errs_list[-1].mean()) if print_info: - print(f'step {(ep + 1) * self.n_iter}, max error {err_max:.6f}, \ - min error {err_min:.6f}, mean error {err_ave:.6f}') + print(f'step {(ep + 1) * self.n_iter}, max error {err_max:.6f}', end=', ') + print(f'min error {err_min:.6f}, mean error {err_ave:.6f}', end=', ') + print(f'mean step time {(toc() - tic) / self.n_iter:.4f}s') + tic = toc() if err_max < tol: + print(f"CBS solve converged within TOL {err_max}") break ur, ui, errs = self(c_star, f_star, ur, ui) errs_list += errs - if remove_pml and self.pml_size: - ur = ur[..., self.pml_size:-self.pml_size, self.pml_size:-self.pml_size] - ui = ui[..., self.pml_size:-self.pml_size, self.pml_size:-self.pml_size] + if remove_pml: + if self.dim == 2: + ur = ur[..., self.pml_size[0]:self.pml_size[0]+self.shape[0], + self.pml_size[2]:self.pml_size[2]+self.shape[1]] + ui = ui[..., self.pml_size[0]:self.pml_size[0]+self.shape[0], + self.pml_size[2]:self.pml_size[2]+self.shape[1]] + else: + ur = ur[..., self.pml_size[0]:self.pml_size[0]+self.shape[0], + self.pml_size[2]:self.pml_size[2]+self.shape[1], + self.pml_size[4]:self.pml_size[4]+self.shape[2]] + ui = ui[..., self.pml_size[0]:self.pml_size[0]+self.shape[0], + self.pml_size[2]:self.pml_size[2]+self.shape[1], + self.pml_size[4]:self.pml_size[4]+self.shape[2]] return ur, ui, errs_list diff --git a/tests/solver/test_cbs.py b/tests/solver/test_cbs.py index 6ef990afd11cfe5dde4c35b798f84e24d1c87efd..ca7e745ff6463639e44edf7efebe409c3bc04063 100644 --- a/tests/solver/test_cbs.py +++ b/tests/solver/test_cbs.py @@ -23,7 +23,7 @@ from mindspore import ops, Tensor from mindscience import CBS -def gen_input(): +def gen_input_2d(): ''' prepare c_star & f_star ''' resolution = 256 @@ -34,7 +34,7 @@ def gen_input(): velo[64:72, 64:72] *= 1.1 # add discontinuity on velocity field mask[64, 64] = 1 # add one source point - c_star = Tensor(velo / omgs.reshape(-1, 1, 1, 1), dtype=ms.float32) + c_star = Tensor(velo / omgs.reshape(-1, 1, 1), dtype=ms.float32) f_star = Tensor(np.broadcast_to(mask, c_star.shape), dtype=ms.float32, const_arg=True) return c_star, f_star @@ -44,7 +44,7 @@ def gen_input(): @pytest.mark.platform_arm_ascend910b_training @pytest.mark.env_onecard @pytest.mark.parametrize('mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE]) -def test_solve(mode): +def test_solve_2d(mode): """ Feature: Test CBS forward solving. Description: None. @@ -52,7 +52,7 @@ def test_solve(mode): """ ms.set_device('Ascend') ms.set_context(mode=mode) - c_star, f_star = gen_input() + c_star, f_star = gen_input_2d() warmup_steps = 1 cbs = CBS(c_star.shape[-2:], remove_pml=False) @@ -78,7 +78,7 @@ def test_solve(mode): @pytest.mark.platform_arm_ascend910b_training @pytest.mark.env_onecard @pytest.mark.parametrize('mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE]) -def test_grad(mode): +def test_grad_2d(mode): """ Feature: Test CBS solver backward propagation. Description: None. @@ -86,7 +86,7 @@ def test_grad(mode): """ ms.set_device('Ascend') ms.set_context(mode=mode) - c_star, f_star = gen_input() + c_star, f_star = gen_input_2d() warmup_steps = 1 cbs = CBS(c_star.shape[-2:], n_iter=5) @@ -106,6 +106,94 @@ def test_grad(mode): loss, grad = grd_func(c_star, f_star) time_spent = toc() - tic - assert isclose(loss, 8.737738, rel_tol=1e-3, abs_tol=1e-3) + assert isclose(loss, 8.4663, rel_tol=1e-3, abs_tol=1e-3) assert isclose(grad.std(), 0.01240166, rel_tol=1e-3, abs_tol=1e-3) assert time_spent <= 30 + + +def gen_input_3d(): + ''' prepare c_star & f_star ''' + resolution = 128 + + velo = np.ones([resolution] * 3) * 1500 / 20 + mask = np.zeros_like(velo) + omgs = np.arange(30, 40) * np.pi + + velo[32:40, 32:40, 32:40] *= 1.1 # add discontinuity on velocity field + mask[32, 32, 32] = 1 # add one source point + + c_star = Tensor(velo / omgs.reshape(-1, 1, 1, 1), dtype=ms.float32) + f_star = Tensor(np.broadcast_to(mask, c_star.shape), dtype=ms.float32, const_arg=True) + + return c_star, f_star + + +@pytest.mark.level0 +@pytest.mark.platform_arm_ascend910b_training +@pytest.mark.env_onecard +@pytest.mark.parametrize('mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE]) +def test_solve_3d(mode): + """ + Feature: Test CBS forward solving. + Description: None. + Expectation: Success or throw AssertionError. + """ + ms.set_device('Ascend') + ms.set_context(mode=mode) + c_star, f_star = gen_input_3d() + warmup_steps = 1 + + cbs = CBS(c_star.shape[-3:], remove_pml=False) + + # warmup runs to eliminate the initiation time + for _ in range(warmup_steps): + cbs(c_star, f_star) + + # run and time a complete solution process + tic = toc() + ur, ui, errs = cbs.solve(c_star, f_star) + time_spent = toc() - tic + n_steps = len(errs) + step_time = time_spent / n_steps + + assert n_steps <= 180 + assert isclose((ur - ui).std(), 0.00265927, rel_tol=1e-3, abs_tol=1e-3) + assert time_spent <= 60 + assert step_time <= 0.5 + + +@pytest.mark.level0 +@pytest.mark.platform_arm_ascend910b_training +@pytest.mark.env_onecard +@pytest.mark.parametrize('mode', [ms.GRAPH_MODE, ms.PYNATIVE_MODE]) +def test_grad_3d(mode): + """ + Feature: Test CBS solver backward propagation. + Description: None. + Expectation: Success or throw AssertionError. + """ + ms.set_device('Ascend') + ms.set_context(mode=mode) + c_star, f_star = gen_input_3d() + warmup_steps = 1 + + cbs = CBS(c_star.shape[-3:], n_iter=5) + + def loss_func(c, f): + ur, ui, _ = cbs(c, f) + return ops.norm(ur) + ops.norm(ui) + + grd_func = ms.value_and_grad(loss_func, 0, None) + + # warmup runs to eliminate the initiation time + for _ in range(warmup_steps): + grd_func(c_star, f_star) + + # run and time a complete forward & backward process + tic = toc() + loss, grad = grd_func(c_star, f_star) + time_spent = toc() - tic + + assert isclose(loss, 7.84422, rel_tol=1e-3, abs_tol=1e-3) + assert isclose(grad.std(), 0.00207235, rel_tol=1e-3, abs_tol=1e-3) + assert time_spent <= 30