CBS模型昇腾NPU部署指导

chmod +x Ascend-hdk-<chip_type>-npu-driver_<version>_linux-<arch>.run
chmod +x Ascend-hdk-<chip_type>-npu-firmware_<version>.run
./Ascend-hdk-<chip_type>-npu-driver_<version>_linux-<arch>.run --full --force
./Ascend-hdk-<chip_type>-npu-firmware_<version>.run --full

#基于PyTorch框架设置环境变量
source /usr/local/Ascend/cann/set_env.sh # 修改为实际安装的Toolkit包路径
source /usr/local/Ascend/nnal/atb/set_env.sh # 修改为实际安装的nnal包路径

# 安装torch和torch_npu构建参考 https://gitcode.com/ascend/pytorch/releases
pip3 install torch-2.7.1-cp310-cp310-manylinux_2_28_aarch64.whl 
pip3 install torch_npu-2.7.1rc1-cp310-cp310-manylinux_2_28_aarch64.whl

pip install mindspore==2.8.0 -i https://repo.mindspore.cn/pypi/simple --trusted-host repo.mindspore.cn --extra-index-url https://repo.huaweicloud.com/repository/pypi/simple

git clone https://gitcode.com/ascend/MindSpeed.git
cd MindSpeed
git checkout master  # checkout commit from MindSpeed master
pip3 install -r requirements.txt 
pip3 install -e .
cd ..

git clone https://gitcode.com/ascend/MindSpeed-LLM.git 
git clone https://github.com/NVIDIA/Megatron-LM.git  # 从github下载Megatron-LM，请确保网络能访问
cd Megatron-LM
git checkout core_v0.12.1
cp -r megatron ../MindSpeed-LLM/
cd ../MindSpeed-LLM
git checkout master

pip3 install -r requirements.txt  # 安装其余依赖库

# 行15-22：导入昇腾兼容的MindSpore核心模块
import mindspore as ms
from mindspore import Tensor, nn, ops, mint, numpy as mnp, lazy_inline
from mindscience import DFTn, IDFTn, DST, IDST
import mindspore as ms
from mindspore import nn, ops, dtype as mstype, mint
from mindspore.ops import DataType, CustomRegOp, CustomOpBuilder

# 行31/71/109：所有计算单元继承MindSpore的nn.Cell
class MixedDSTDFTn(nn.Cell):
class MixedIDSTDFTn(nn.Cell):
class CBSBlock(nn.Cell):
class CBS(nn.Cell):

# 行50-51：昇腾维度交换算子（替代numpy.swapaxes）
dft_ar = ops.swapdims(dft_ar, 2, -1) 
dft_ai = ops.swapdims(dft_ai, 2, -1)

# 行82-83：昇腾转置算子（替代numpy.transpose）
ur = mint.transpose(ur, -2, -1)
ui = mint.transpose(ui, -2, -1)

# 行343：昇腾concat算子（替代numpy.concatenate）
return ops.concat((left, original, right))

# 行156-157：ops.select替代原生if-else（昇腾图模式兼容的条件分支）
dur = ops.select(cond, -vgi / (eps + 1e-8), gvr - ur) 
dui = ops.select(cond, vgr / (eps + 1e-8), gvi - ui)

# 行413：construct内的for循环（MindSpore昇腾图模式支持的循环规范）
for _ in range(self.n_iter):
    dur, dui = self.cbs_block(ur, ui, vr, vi, gr, gi, rhs, eps)
    ur = ur + dur
    ui = ui + dui
    errs = self.compute_error(dur, dui, ur, ui)
    errs_list.append(errs)

# 行138-150：复数张量拆分为实虚部单独计算（昇腾兼容）
if ops.is_complex(f_star):
    f_star_real = get_real(f_star)
    f_star_imag = get_imag(f_star)
    c_star_squared = c_star**2
    rhs_real = f_star_real / c_star_squared
    rhs_imag = f_star_imag / c_star_squared
    rhs_real_padded = ops.pad(rhs_real, self.pml_size_xyz)
    rhs_imag_padded = ops.pad(rhs_imag, self.pml_size_xyz)
    rhs = make_complex(rhs_real_padded, rhs_imag_padded)
else:
    rhs = mint.nn.functional.pad(f_star / c_star**2, self.pml_size_xyz)

# 行124-130：op_v/op_g算子均拆分实虚部计算，无直接复数操作
def op_v(self, ur, ui, vr, vi):
    wr = ur * vr - ui * vi
    wi = ur * vi + ui * vr
    return wr, wi
def op_g(self, ur, ui, gr, gi):
    fur, fui = self.dft_cell(ur, ui)
    gur = gr * fur - gi * fui
    gui = gi * fur + gr * fui
    wr, wi = self.idft_cell(gur, gui)
    return wr, wi

"""
CBS 算法推理主脚本
功能：基于给定的速度场和源项，求解频域声波方程的波场（实部/虚部）
适配：2D 场景（3D 可修改 shape/dxs 等参数扩展）
"""
import os
import numpy as np
import mindspore as ms
import matplotlib.pyplot as plt
from cbs import CBS

# ====================== 1. 环境配置 ======================
ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend") 
ms.set_seed(42)  # 固定随机种子
save_dir = "./cbs_inference_results"
os.makedirs(save_dir, exist_ok=True)

# ====================== 2. 推理参数配置 ======================
# 2D 空间网格形状 (nz, nx)，3D 改为 (nz, ny, nx)
shape = (128, 128)
# 网格步长 (dz, dx)，3D 改为 (dz, dy, dx)
dxs = (0.01, 0.01)
# 迭代配置
n_iter_per_call = 20  # 单次 construct 调用的迭代次数
tol = 1e-3            # 收敛误差阈值
max_iter = 10000      # 最大迭代次数
# 边界配置
btype = "pml"         # 边界类型：pml / freesurface
pml_size = 16         # PML 层厚度
alpha = 1.0           # PML 衰减强度
rampup = 12           # PML 过渡平滑度
remove_pml = True     # 输出是否移除 PML 层

# ====================== 3. 生成测试数据（速度场 + 源项） ======================
batch_size = 1  # 批次大小
channel = 1     # 通道数

# 3.1 速度场 c_star (非维度化，物理含义：介质波速)
# 构造分层速度场（示例）：上层低速，下层高速
c_star = np.ones((batch_size, channel) + shape, dtype=np.float32) * 1.5
z_split = shape[0] // 2
c_star[:, :, z_split:, :] = 2.5  # 下半部分速度更高

# 3.2 源项 f_star (源位置掩码，仅源点位置为1，其余为0)
f_star = np.zeros((batch_size, channel) + shape, dtype=np.float32)
# 源点位置：网格中心
src_z, src_x = shape[0] // 2, shape[1] // 2
f_star[:, :, src_z, src_x] = 1.0  # 点源

# 转换为 MindSpore Tensor
c_star = ms.Tensor(c_star, dtype=ms.float32)
f_star = ms.Tensor(f_star, dtype=ms.float32)

# ====================== 4. 初始化 CBS 求解器 ======================
cbs_solver = CBS(
    shape=shape,
    dxs=dxs,
    n_iter=n_iter_per_call,
    pml_size=pml_size,
    alpha=alpha,
    rampup=rampup,
    btype=btype,
    remove_pml=False,  # 迭代过程中保留 PML，最终输出再移除
    epsilon=None       # 自动计算稳定系数
)

# ====================== 5. 执行推理（求解波场） ======================
print("开始 CBS 推理求解...")
# 初始波场设为 0（ur_init/ui_init=None 时自动初始化为 0）
ur, ui, errs_list = cbs_solver.solve(
    c_star=c_star,
    f_star=f_star,
    tol=tol,
    max_iter=max_iter,
    remove_pml=remove_pml,
    print_info=True
)

# ====================== 6. 结果分析与可视化 ======================
# 6.1 迭代误差分析
errs_np = [err.asnumpy() for err in errs_list]
errs_max = [np.max(err) for err in errs_np]
errs_mean = [np.mean(err) for err in errs_np]

# 绘制迭代误差曲线
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.plot(errs_max, label="Max Error")
plt.plot(errs_mean, label="Mean Error")
plt.xlabel(f"Iteration Step (×{n_iter_per_call})")
plt.ylabel("Relative Error")
plt.yscale("log")
plt.title("迭代误差变化")
plt.legend()
plt.grid(True, alpha=0.3)

# 6.2 波场可视化（实部 + 虚部）
ur_np = ur.asnumpy()[0, 0]  # 取 batch=0, channel=0 的 2D 波场实部
ui_np = ui.asnumpy()[0, 0]  # 取 batch=0, channel=0 的 2D 波场虚部

plt.subplot(1, 2, 2)
plt.imshow(ur_np, cmap="seismic", aspect="auto")
plt.colorbar(label="波场实部")
plt.title("波场实部分布")
plt.xlabel("X 网格")
plt.ylabel("Z 网格")

plt.tight_layout()
plt.savefig(os.path.join(save_dir, "cbs_wavefield_real.png"), dpi=300, bbox_inches="tight")

# 绘制虚部分布
plt.figure(figsize=(8, 6))
plt.imshow(ui_np, cmap="seismic", aspect="auto")
plt.colorbar(label="波场虚部")
plt.title("波场虚部分布")
plt.xlabel("X 网格")
plt.ylabel("Z 网格")
plt.savefig(os.path.join(save_dir, "cbs_wavefield_imag.png"), dpi=300, bbox_inches="tight")

# ====================== 7. 结果保存 ======================
# 保存波场数据（实部/虚部）
np.save(os.path.join(save_dir, "wavefield_real.npy"), ur_np)
np.save(os.path.join(save_dir, "wavefield_imag.npy"), ui_np)
# 保存迭代误差
np.save(os.path.join(save_dir, "iteration_errors.npy"), errs_np)

print(f"推理完成！结果已保存至 {save_dir}")
print(f"波场形状（移除 PML 后）：{ur_np.shape}")
print(f"最终迭代最大误差：{errs_max[-1]:.6f}")
print(f"最终迭代平均误差：{errs_mean[-1]:.6f}")