Sympy Codegen生成Fortran矩阵代码的两类技术问题咨询

我用Sympy处理了一组含符号表达式的矩阵，通过codegen成功生成了.f90文件并能在其他代码中调用，但还有两个问题需要解决：

1. 如何在codegen中指定生成双精度实数类型的代码？
2. 当前每个矩阵对应的Fortran子程序只包含参与计算的参数，怎么让每个子程序接收所有变量作为参数（哪怕部分变量没参与计算），实现统一调用格式：CALL matA11(all_the_variables_in_the_f90_file,name_matrix)和CALL matA21(all_the_variables_in_the_f90_file,name_matrix)，避免用户逐个指定参数。

生成的测试用.f90文件示例 #

subroutine matA11(Gamma1, g, p0, rho0, yml, out_8665343569832042533)
implicit none
REAL*8, intent(in) :: Gamma1
REAL*8, intent(in) :: g
REAL*8, intent(in) :: p0
REAL*8, intent(in) :: rho0
REAL*8, intent(in) :: yml
REAL*8, intent(out), dimension(1:4, 1:4) :: out_8665343569832042533

out_8665343569832042533(1, 1) = 0
out_8665343569832042533(2, 1) = 0
out_8665343569832042533(3, 1) = 0
out_8665343569832042533(4, 1) = rho0*yml/(Gamma1*p0)
out_8665343569832042533(1, 2) = 0
out_8665343569832042533(2, 2) = 0
out_8665343569832042533(3, 2) = 0
out_8665343569832042533(4, 2) = 0
out_8665343569832042533(1, 3) = 0
out_8665343569832042533(2, 3) = 0
out_8665343569832042533(3, 3) = 0
out_8665343569832042533(4, 3) = 0
out_8665343569832042533(1, 4) = g*yml
out_8665343569832042533(2, 4) = 0
out_8665343569832042533(3, 4) = 0
out_8665343569832042533(4, 4) = 0

end subroutine

subroutine matA21(drdmu, dyml, g, m, mu, r, ymlp, &
      out_825899070853517470)
implicit none
REAL*8, intent(in) :: drdmu
REAL*8, intent(in) :: dyml
REAL*8, intent(in) :: g
REAL*8, intent(in) :: m
REAL*8, intent(in) :: mu
REAL*8, intent(in) :: r
REAL*8, intent(in) :: ymlp
REAL*8, intent(out), dimension(1:4, 1:3) :: out_825899070853517470

out_825899070853517470(1, 1) = -drdmu*dyml*g*mu**2/r + drdmu*dyml*g/r
out_825899070853517470(2, 1) = 0
out_825899070853517470(3, 1) = 0
out_825899070853517470(4, 1) = 0
out_825899070853517470(1, 2) = -drdmu*g*m*ymlp/r
out_825899070853517470(2, 2) = 0
out_825899070853517470(3, 2) = 0
out_825899070853517470(4, 2) = 0
out_825899070853517470(1, 3) = 0
out_825899070853517470(2, 3) = 0
out_825899070853517470(3, 3) = 0
out_825899070853517470(4, 3) = 0
end subroutine

1. 指定生成双精度实数类型 #

Sympy的codegen可以通过自定义类型映射强制生成双精度代码，有两种常用实现方式：

方法一：单个表达式生成时指定类型 #

使用sympy.printing.fcode的type_aliases参数，将默认的real映射为real*8：

from sympy import symbols, fcode

Gamma1, g, p0, rho0, yml = symbols('Gamma1 g p0 rho0 yml')
expr = rho0*yml/(Gamma1*p0)

# 生成双精度代码
print(fcode(expr, type_aliases={'real': 'real*8'}))

方法二：批量生成子程序时配置 #

使用sympy.utilities.codegen.codegen函数，通过type_aliases参数统一设置类型：

from sympy import symbols, Matrix
from sympy.utilities.codegen import codegen

# 定义变量
Gamma1, g, p0, rho0, yml = symbols('Gamma1 g p0 rho0 yml')
matA11 = Matrix([
    [0, 0, 0, g*yml],
    [0, 0, 0, 0],
    [0, 0, 0, 0],
    [rho0*yml/(Gamma1*p0), 0, 0, 0]
])

# 生成双精度Fortran代码
codegen(
    ('matA11', matA11),
    language='fortran',
    filename='matrices',
    # 映射为REAL*8，也可替换为'double precision'适配F95标准
    type_aliases={'real': 'real*8'}
)

2. 让子程序接收所有变量作为参数 #

Sympy默认只将表达式中用到的变量设为参数，要强制传入所有变量，可通过添加哑变量引用实现：在矩阵表达式中加入未使用变量乘以0的项（不影响计算结果，但会被Sympy识别为必要参数）。

示例代码 #

from sympy import symbols, Matrix, ZeroMatrix
from sympy.utilities.codegen import codegen

# 定义所有需要传入的全局变量
all_vars = symbols('Gamma1 g p0 rho0 yml drdmu dyml m mu r ymlp')

# 构造matA11原始表达式
matA11_raw = Matrix([
    [0, 0, 0, g*yml],
    [0, 0, 0, 0],
    [0, 0, 0, 0],
    [rho0*yml/(Gamma1*p0), 0, 0, 0]
])
# 找出未用到的变量，添加哑引用
unused_vars_A11 = [var for var in all_vars if var not in matA11_raw.free_symbols]
dummy_A11 = ZeroMatrix(4,4)
for var in unused_vars_A11:
    dummy_A11 += var * ZeroMatrix(4,4)
matA11_full = matA11_raw + dummy_A11

# 同理处理matA21
matA21_raw = Matrix([
    [-drdmu*dyml*g*mu**2/r + drdmu*dyml*g/r, -drdmu*g*m*ymlp/r, 0],
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0]
])
unused_vars_A21 = [var for var in all_vars if var not in matA21_raw.free_symbols]
dummy_A21 = ZeroMatrix(4,3)
for var in unused_vars_A21:
    dummy_A21 += var * ZeroMatrix(4,3)
matA21_full = matA21_raw + dummy_A21

# 批量生成统一参数格式的代码
codegen(
    [('matA11', matA11_full), ('matA21', matA21_full)],
    language='fortran',
    filename='matrices_unified',
    type_aliases={'real': 'real*8'}
)

效果说明 #

生成的两个子程序都会包含所有全局变量作为输入参数，调用时可统一传入全部变量，无需区分每个子程序的参数列表。

内容的提问来源于stack exchange，提问作者arkhose u