二维声波问题

概述

声波方程求解是医疗超声、地质勘探等领域中的核心技术，大规模声波方程求解面临算力和存储的挑战。声波方程求解器一般采用频域求解算法和时域求解算法，时域求解算法的代表是时域有限差分法 (TDFD)，频域求解算法包括频域有限差分法 (FDFD)、有限元法 (FEM) 和 CBS (Convergent Born series) 迭代法。CBS 方法由于不引入频散误差，且求解的内存需求低，因此受到工程和学术界的广泛关注。尤其是 Osnabrugge et al. (2016) 解决了该方法的收敛性问题，使得 CBS 方法的应用具有更广阔的前景。

本案例将演示如何调用 MindFlow 提供的 CBS API 实现二维声波方程的求解。

问题描述

声波方程求解中，波速场和震源信息是输入参数，求解输出的是时空分布的波场。

二维声波方程的表达式如下

时域表达式	频域表达式
$\frac{\partial^{2} u}{\partial t^{2}} - c^{2} Δ u = f$	$- ω^{2} \hat{u} - c^{2} Δ \hat{u} = \hat{f}$

其中

$u (x, t) [L]$ 变形位移 (压强除以密度)，标量
$c (x) [L / T]$ 波速，标量
$f (x, t) [L / T^{2}]$ 震源激励 (体积分布力)，标量

实际求解中，为了降低参数维度，一般先将参数无量纲化，然后针对无量纲方程和参数进行求解，最后恢复解的量纲。选取 $ω$ 、 $\hat{f}$ 和 $d$ （网格间距，本案例要求网格在各方向间距相等）对频域方程做无量纲化，可得频域无量纲方程：

u^{*} + c^{* 2} \tilde{Δ} + f^{*} = 0

其中

$u^{*} = \hat{u} ω^{2} / \hat{f}$ 为无量纲变形位移
$c^{*} = c / (ω d)$ 为无量纲波速
$\tilde{Δ}$ 为归一化 Laplace 算子，即网格间距均为 1 时的 Laplace 算子
$f^{*}$ 为标记震源位置的 mask，即在震源作用点为 1，其余位置为 0

本案例中 src 包可以在 src 下载。

[1]:

import os
import numpy as np
import scipy
import pandas as pd
import mindspore as ms
from mindspore import Tensor

from mindflow.utils import load_yaml_config

from cbs.cbs import CBS
from src import visual
from solve_acoustic import solve_cbs

定义输入参数和输出采样方式

本案例所需的输入为有量纲 2D 速度场、震源位置列表、震源波形，在文件 config.yaml 中指定输入文件名。为了方便用户直接验证，本案例在本链接中提供了预置的输入数据，请下载所需要的数据集，并保存在 ./dataset 目录下。数据集包括速度场 velocity.npy、震源位置列表 srclocs.csv、震源波形 srcwaves.csv。用户可仿照输入文件格式自行修改输入参数。

输出为时空分布的波场，为了明确输出如何在时间和频率上采样，需在 config.yaml 文件中指定 dt, nt 等参数。

由于输入的震源波形在时间上的采样率可能与输出所要求的不一致，因此需对其进行插值。

[2]:

ms.set_context(device_target='Ascend', device_id=0, mode=ms.GRAPH_MODE)

config = load_yaml_config('config.yaml')

data_config = config['data']
solve_config = config['solve']
summary_config = config['summary']

# 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)

# 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)

# read & interp source wave
df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_wave']))
inter_func = scipy.interpolate.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)

# 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']

选取待求频点

确定了输出采样方式即确定了所有待求频点。但为了减少计算量，也可以只选择部分频点进行求解，其余频点通过插值获得。具体的频点降采样方式由 config.yaml 文件中的 downsample_mode, downsample_rate 指定。默认为不做降采样，即求解除 $ω = 0$ 之外的所有频点。

[3]:

# select omegas
no = len(omegas_all) // solve_config['downsample_rate']

if solve_config['downsample_mode'] == 'exp':
    omegas_sel = np.exp(np.linspace(np.log(omegas_all[1]), np.log(omegas_all[-1]), no))
elif solve_config['downsample_mode'] == 'square':
    omegas_sel = np.linspace(omegas_all[1]**.5, omegas_all[-1]**.5, no)**2
else:
    omegas_sel = np.linspace(omegas_all[1], omegas_all[-1], no)

执行仿真

将相关数组定义为 Tensor，调用 solve_cbs()，在 NPU 执行求解。由于显存限制，求解过程在频点维度分批执行，batch 数量由用户在 config.yaml 中指定，不要求整除频点数（允许最后一个 batch 的 size 与其他 batch 不一致）。求解完成后，会保存频域求解结果到文件 u_star.npy。

[ ]:

# 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)

ur, ui = solve_cbs(cbs, velo, slocs, omegas_sel, dx=dx, n_batches=solve_config['n_batches']) # shape (ns, no, len(receiver_zs), nx)

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)

仿真结果后处理

CBS 求解的是无量纲化的频域方程，但下游任务通常希望在时域观察有量纲波场的演化过程，因此最后将求解结果恢复量纲化并转回时域。恢复量纲的方式为 $\hat{u} = u^{*} \hat{f} / ω^{2}$ ，若在前述的“选取待求频点”步骤中对频点做了降采样，则在此处需沿频率方向插值恢复所有频点的解。然后对有量纲的频域波场 $\hat{u}$ 做 Fourier 反变换得到时域波场 $u$ 。将时域波场保存至文件 u_time.npy。

[ ]:

# recover dimension and interpolate to full frequency domain
u_star /= omegas_sel.reshape(-1, 1, 1)**2
u_star = scipy.interpolate.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)

# 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
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'))