教程 | 如何在数值模式中丝滑耦合AI模型？以WRF-Chem为例

#include <torch/script.h> // LibTorch 核心头文件
#include <torch/torch.h>
#include <iostream>
#include <memory>
#include <unistd.h>
    
    // 使用 extern "C" 确保 C++ 编译器以 C 语言标准编译函数名，以便 Fortran 能正确识别和链接
extern"C"void saprc99_DNN(double *c_cin, double *c_cout, int batch_size);

void saprc99_DNN(double *c_cin, double *c_cout, int batch_size) {
    // 1. 定义并加载序列化的 PyTorch 模型
    //    模型文件 "AIPC_MHSA.pt" 需要与最终的可执行文件位于同一目录
    torch::jit::script::Module module;
    try {
        module = torch::jit::load("AIPC_MHSA.pt");
    } catch (const c10::Error& e) {
        std::cerr << "error loading the model\n";
        return;
    }

    // 2. 将 Fortran 传入的 double* 数组转换为 LibTorch 张量
    //    a. 使用 torch::from_blob 将原始数据指针封装为张量，不发生内存拷贝
    auto cinput = torch::from_blob(c_cin, {batch_size, 91}, torch::kDouble).unsqueeze(1); // 形状: [batch_size, channel=1, 91 features]
    //    b. 根据模型要求，将数据类型转换为 Float
    auto cinput_ = cinput.to(torch::kFloat);

    // 3. 准备模型输入
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(cinput_);

    // 4. 执行模型前向推理
    at::Tensor output = module.forward(inputs).toTensor();
    
    // 5. 将输出张量转换回 Fortran 需要的 double 类型
    output = output.to(torch::kDouble);
    auto res = output.accessor<double, 2>(); // 使用 accessor 获取数据访问器

    // 6. 将结果从输出张量内存中拷贝回 Fortran 的数组指针
    std::memcpy(c_cout, res.data(), 74 * batch_size * sizeof(double));
}

# 指定最低 CMake 版本要求
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)

# 定义项目名称和支持的语言
project(Demo LANGUAGES C CXX Fortran)

# 查找已安装的 LibTorch 包，这是关键依赖
find_package(Torch REQUIRED)

# 从源文件 saprc99DNN.cpp 创建一个名为 saprc99DNN 的静态库
add_library(saprc99DNN STATIC saprc99DNN.cpp)

# 将 LibTorch 的库链接到我们的静态库上
target_link_libraries(saprc99DNN "${TORCH_LIBRARIES}")

# 设置 C++ 标准为 C++14 或更高，以兼容 LibTorch
set_property(TARGET saprc99DNN PROPERTY CXX_STANDARD 14)

MODULE module_saprc99_DNN

! 定义 Fortran 与 C 语言之间数据类型的对应关系
USE, INTRINSIC :: iso_c_binding

IMPLICITNONE

! 接口定义
INTERFACE
    ! 定义一个与 C++ 函数匹配的 Fortran 子程序
    ! bind(c, name="saprc99_DNN") 是关键，它告诉 Fortran 编译器：
    ! 1. 使用 C 语言的调用约定 (bind(c))
    ! 2. 这个子程序在链接时对应的外部符号名为 "saprc99_DNN"
    SUBROUTINE saprc99_DNN(c_cin, c_cout, batch_size) BIND(C, name='saprc99_DNN')
      USEiso_c_binding
      IMPLICITNONE

      ! 参数类型必须与 C 语言类型精确对应
      REAL(c_double), INTENT(IN)  :: c_cin(91, *)
      REAL(c_double), INTENT(OUT) :: c_cout(74, *)
      INTEGER(c_int), VALUE       :: batch_size
    ENDSUBROUTINE saprc99_DNN
ENDINTERFACE

ENDMODULE module_saprc99_DNN

! 在程序的声明部分添加 USE 语句
USE module_saprc99_DNN

! ... 省略其他代码 ...

! 在需要进行化学计算的地方
CALL saprc99_DNN(input_array, output_array, grid_cell_count)

! ... 省略其他代码 ...

# 在 LIB_EXTERNAL 变量中添加 AI 静态库和 LibTorch 的路径
# L$ 指示链接器搜索库的路径
# l 指示链接具体的库名

LIB_EXTERNAL = \
               -L$(WRF_SRC_ROOT_DIR)/external/io_netcdf -lwrfio_nf ... \
               # 添加我们自己的静态库
               $(WRF_SRC_ROOT_DIR)/external/libsaprc99DNN.a \
               # 添加 LibTorch 及其依赖的路径和库
               /opt/rh/devtoolset-11/root/lib/gcc/x86_64-redhat-linux/11/libstdc++.a \
               -L/home/zhxia/env/libtorch_cpu/lib -ltorch_cpu -lc10 ...

$ ln -s /path/to/your/model/AIPC_MHSA.pt wrf/run/AIPC_MHSA.pt