I would like to share a short Python snippet I’ve been using very often to generate Fortran code through symbolic simplification and common sub-expression elimination:
# Copyright (c) 2020 Ivan Pribec. All rights reserved.
#
# This work is licensed under the terms of the MIT license.
# For a copy, see <https://opensource.org/licenses/MIT>.
import sympy as sp
from sympy.printing.fcode import print_fcode
def apply_cse_and_print_fcode(expr,assign_to,source_format='free',standard=95):
"""Apply common subexpression simplification and print Fortran code
Input:
expr - a list of expressions we would like to evaluate
assign_to - the symbols we would like to assign them to (either a tuple of strings or sp.symbols)
Output:
The expressions will be printed to the standard output.
"""
rvar, rexpr = sp.cse(sp.simplify(expr))
print("! Reduced variables")
for var, var_expr in rvar:
print_fcode(var_expr,assign_to=var,source_format=source_format,standard=standard)
print("! Reduced expressions")
for var, var_expr in zip(assign_to,rexpr):
print_fcode(var_expr,assign_to=var,source_format=source_format,standard=standard)
As an example I will be computing the water activity of a wet material, a_w, which is a function of moisture content X and temperature T, according to the Oswin model equation:
where the parameters (a_1,a_2,b_1,b_2) are empiricial fitting coefficients. In some applications we might also be interested in the derivative of a_w with respect to the variables X and T. Using SymPy we can easily find these using the following snippet (you can run this in a Jupyter notebook):
# Create symbols
X, T, a1, a2, b1, b2 = sp.symbols('X T a1 a2 b1 b2')
# Define symbolic expression
a = a1 + a2*T
b = b1 + b2*T
aw = (X/a)**(1/b)/(1 + (X/a)**(1/b))
# Find symbolic derivatives
daw_dX = sp.diff(aw,X)
daw_dT = sp.diff(aw,T)
To generate Fortran code for the expression and derivatives all I need to do now is run the following code:
output = (aw,daw_dX,daw_dT)
output_var = ('aw','daw_dX','daw_dT')
apply_cse_and_print_fcode(output,output_var)
which will print the reduced variables and expressions to standard output:
! Reduced variables
x0 = T*a2 + a1
x1 = 1d0/x0
x2 = X*x1
x3 = T*b2 + b1
x4 = 1d0/x3
x5 = x2**x4
x6 = x5 + 1
x7 = x5/x6**2
! Reduced expressions
aw = x5/x6
daw_dX = x4*x7/X
daw_dT = -x1*x7*(a2*x3 + b2*x0*log(x2))/x3**2
I can now wrap this output into a subroutine block, specifying the precise interface and applying any further changes (i.e. replace the d0 format specifier to my desired precision):
subroutine oswin_model(X,T,a1,a2,b1,b2,aw,daw_dX,daw_dT)
real(wp), intent(in) :: X, T
real(wp), intent(in) :: a1, a2, b1, b2,
real(wp), intent(out) :: aw, daw_dX, daw_dT
real(wp) :: x0,x1,x2,x3,x4,x5,x6,x7
! Reduced variables
x0 = T*a2 + a1
x1 = 1.0_wp/x0
x2 = X*x1
x3 = T*b2 + b1
x4 = 1.0_wp/x3
x5 = x2**x4
x6 = x5 + 1
x7 = x5/x6**2
! Reduced expressions
aw = x5/x6
daw_dX = x4*x7/X
daw_dT = -x1*x7*(a2*x3 + b2*x0*log(x2))/x3**2
end subroutine
As a second example, say you want to evaluate the Jacobian of a system of ordinary differential equations. I will use the Lorenz attractor as a trivial example:
First we define the symbolic system of differential equations:
x,y,z = sp.symbols('x y z')
s,r,b = sp.symbols('sigma rho beta')
dx_dt = s*(y-x)
dy_dt = x*(r - z) - y
dz_dt = x*y - b*z
system = sp.Matrix([dx_dt,dy_dt,dz_dt])
The SymPy Matrix object has a built-in Jacobian function, that can be called as system.jacobian((x,y,z)) to produce the output:
To output the Fortran code we again use the function defined above,
output = (system,system.jacobian((x,y,z)))
output_vars = ('F','J')
apply_cse_and_print_fcode(output,output_vars)
which generates the output:
! Reduced variables
x0 = -x
x1 = rho - z
! Reduced expressions
F(1, 1) = sigma*(x0 + y)
F(2, 1) = x*x1 - y
F(3, 1) = -beta*z + x*y
J(1, 1) = -sigma
J(2, 1) = x1
J(3, 1) = y
J(1, 2) = sigma
J(2, 2) = -1
J(3, 2) = x
J(1, 3) = 0
J(2, 3) = x0
J(3, 3) = -beta
We can see in this case, there is not much need to introduce temporary variables. Nevertheless, the symbolic route is much easier than writing the Jacobian by hand. In principle, it is possible to fine-tune the format of the output arrays by passing additional arguments to the Fortran printer in SymPy. Personally, for short functions, I prefer to do it manually using a good text editor.
If you are interested in learning more about code generation using Sympy, I suggest reading the tutorial “Automatic Code Generation with SymPy”. For those who have a MATLAB license, a similar tool is available in the Symbolic Math Toolbox.
Hopefully, someone will find this useful.
