devitocodes
Exponential complication time for code generation in C
Issues with the operator construction time for multi-stage methods. The ideia of the multi-stage methods implemented in Devito is to write a PDE system as first order in time:
where H is the spatial operator (with spatial derivatives and other spatial functions) and f is the source.
With this formulation, various time integrators can be implemented, including arbitrary-order Runge-Kutta methods with Strong Stability Preserving (SSP) properties:
where m is the number of stages of the method.
The problem resides in that when constructing the operator with the command
op = Operator(pde + [rec_term], subs=grid.spacing_map)
The C code generation requires exponentially increasing computational time with each additional stage of the method. The following figure illustrates the construction time for the 2D acoustic case using a second-order finite difference (FD) scheme.
As a strategy, besides the default optimization parameters, it was tested using
Operator(equations, opt=('advanced', {'expand': False}))
and
Operator(equations, opt='noop')
Nonetheless, the C code generation still required an exponential increasing computational time with each additional stage of the method.
Below is a minimal implementation for constructing the operator for a multi-stage Runge-Kutta method.
import devito as dv
from examples.seismic import Receiver, TimeAxis
dv.configuration['log-level'] = 'DEBUG'
import numpy as np
import sympy as sym
def acoustic_2d(stages,opt,space_order=2,fd_order=2):
# definition of the grid ----------------------------------------------------------------------------------------
grid = dv.Grid(origin=(0,0), extent=(1,1), shape=(3,3),dtype=np.float64)
(x, y) = grid.dimensions
dt = grid.stepping_dim.spacing
t = grid.time_dim
# definition of the medium velocity ------------------------------------------------------------------------------
param_fun = dv.Function(name="vel", grid=grid, space_order=space_order, dtype=np.float64)
param_fun.data[:] = 1
# definition of the wave displacement and velocity (the unknown functions) ---------------------------------------
fun_labels=['u','v']
U0=[dv.TimeFunction(name=fun_labels[i], grid=grid, space_order=space_order,time_order=1, dtype=np.float64) for i in range(2)]
# definition of the receiver ------------------------------------------------------------------------------------
time_range=TimeAxis(start=0, stop=1, num=101)
rec=Receiver(name='rec', grid=grid, npoint=3, time_range=time_range)
rec.coordinates.data[:, 0] = np.linspace(0,1,3)
rec.coordinates.data[:, 1] = 0.5
rec=rec.interpolate(expr=U0[0].forward)
# definition of the source ----------------------------------------------------------------------------------------
# spatial component
src_spat = dv.Function(name="src_spat", grid=grid, space_order=space_order, dtype=np.float64)
src_spat.data[1, 1] = 1
# time component function and its derivatives
f0=0.2
t_var=sym.Symbol('t_var')
src_time=[(1-2*(np.pi*f0*(t_var-1/f0))**2)*sym.exp(-(np.pi*f0*(t_var-1/f0))**2)]
for i in range(stages-1):
src_time=src_time+[src_time[-1].diff(t_var)]
# definition of the spatial operator for the 2D acoustic equation ------------------------------------------------
def sys_op_extended(U0,param_fun,stages,src_time,src_spat,x,y,t,dt,e_p):
system_op=[U0[1], (dv.Derivative(U0[0],(x,2),fd_order=fd_order)+
dv.Derivative(U0[0],(y,2),fd_order=fd_order))*param_fun**2]
for i in range(stages):
if e_p[stages-1-i]!=0:
system_op[1]+=src_spat*src_time[i].subs({t_var:t*dt})*e_p[stages-1-i]
e_p=[e_p[i]+dt*e_p[i+1] for i in range(stages-1)]+[e_p[-1]]
return system_op,e_p
# HORK (arbitrary order Runge-Kutta method) ----------------------------------------------------------------------
# High-order Runge-Kutta coefficients
alpha = np.zeros(stages)
alpha[0] = np.exp(0)
for i in range(1, stages):
alpha[i] = 1 / (i + 1) * alpha[i - 1]
alpha[1:i] = 1 / (np.array(range(1, i))) * alpha[:(i - 1)]
alpha[0] = 1 - np.sum(alpha[1:(i + 1)])
# High-order Runge-Kutta method ---------------------------------------------------------------------------------
e_p=[0]*stages
e_p[-1]=1
approx = [U0[i]*alpha[0] for i in range(2)]
aux=U0
for i in range(1,stages-1):
system_op,e_p=sys_op_extended(aux,param_fun,stages,src_time,src_spat,x,y,t,dt,e_p)
aux=[aux[j]+dt*system_op[j] for j in range(2)]
approx=[approx[j]+aux[j]*alpha[i] for j in range(2)]
system_op,e_p=sys_op_extended(aux,param_fun,stages,src_time,src_spat,x,y,t,dt,e_p)
aux=[aux[i]+dt*system_op[i] for i in range(2)]
system_op,e_p=sys_op_extended(aux,param_fun,stages,src_time,src_spat,x,y,t,dt,e_p)
aux=[aux[i]+dt*system_op[i] for i in range(2)]
approx=[approx[i]+aux[i]*alpha[stages-1] for i in range(2)]
pde=[dv.Eq(U0[i].forward,approx[i]) for i in range(2)]
# constructing the operator ---------------------------------------------------------------------------------------
if opt=='default':
op = dv.Operator(pde + [rec], subs=grid.spacing_map)
elif opt=='n_expand':
op = dv.Operator(pde + [rec], subs=grid.spacing_map,opt=('advanced', {'expand': False}))
else:
op = dv.Operator(pde + [rec], subs=grid.spacing_map,opt='noop')
# possibles values of opt:
# 'default' : with default Devito's optimization
# 'n_expand' : with expand==False
# 'noop' : with no optimization
# code for number of stages from 3 to 10 using opt=='default' (not to use a number of stages less tha 3)
opt='default'
for i in range(3,10):
print(i)
acoustic_2d(i,opt)
# code for number of stages from 3 to 10 using opt=='n_expand'
opt='n_expand'
for i in range(3,10):
print(i)
acoustic_2d(i,opt)
# code for number of stages from 3 to 10 using opt=='n_expand'
opt='noop'
for i in range(3,10):
print(i)
acoustic_2d(i,opt)
