def test_switch_phases(python_method_impl):
from dagrt.language import CodeBuilder, ExecutionPhase
with CodeBuilder(name="state_1") as builder_1:
builder_1(var("<state>x"), 1)
builder_1.switch_phase("state_2")
with CodeBuilder(name="state_2") as builder_2:
builder_2.yield_state(var("<state>x"), "x", 0, "final")
code = DAGCode(phases={
"state_1":
ExecutionPhase(name="state_1",
next_phase="state_1",
statements=builder_1.statements),
"state_2":
ExecutionPhase(name="state_2",
next_phase="state_2",
statements=builder_2.statements)
},
initial_phase="state_1")
from utils import execute_and_return_single_result
result = execute_and_return_single_result(python_method_impl,
code,
initial_context={"x": 0},
max_steps=2)
assert result == 1
def generate(self, explainer=None):
"""
:arg explainer: a subclass of :class:`SchemeExplainerBase`, possibly
:class:`TextualSchemeExplainer`, or *None*.
:returns: :class:`dagrt.language.DAGCode`
"""
if explainer is None:
explainer = SchemeExplainerBase()
from dagrt.language import DAGCode, CodeBuilder
with CodeBuilder(name="initialization") as cb_init:
self.emit_initialization(cb_init)
with CodeBuilder(name="primary") as cb_primary:
self.emit_adams_method(cb_primary, explainer)
with CodeBuilder(name="bootstrap") as cb_bootstrap:
self.emit_rk_bootstrap(cb_bootstrap)
return DAGCode(phases={
"initialization":
cb_init.as_execution_phase("bootstrap"),
"bootstrap":
cb_bootstrap.as_execution_phase("bootstrap"),
"primary":
cb_primary.as_execution_phase("primary"),
},
initial_phase="initialization")
def generate(self):
"""
:returns: :class:`dagrt.language.DAGCode`
"""
comp_id = self.component_id
from pymbolic import var
dt = var("<dt>")
t = var("<t>")
residual = var("<p>residual_" + comp_id)
state = var("<state>" + comp_id)
rhs_func = var(self.rhs_func_name)
with CodeBuilder("initialization") as cb:
cb(residual, 0)
cb_init = cb
# Primary.
rhs_val = var("rhs_val")
with CodeBuilder("primary") as cb:
# https://github.com/PyCQA/pylint/issues/3387
for a, b, c in self.coeffs: # pylint: disable=not-an-iterable
cb(rhs_val, rhs_func(t=t + c * dt, **{comp_id: state}))
cb(residual, a * residual + dt * rhs_val)
new_state_expr = state + b * residual
if self.state_filter is not None:
new_state_expr = self.state_filter(
**{comp_id: new_state_expr})
cb(state, new_state_expr)
cb.yield_state(state, comp_id, t + dt, "final")
cb(t, t + dt)
cb_primary = cb
from dagrt.language import DAGCode
return DAGCode(phases={
"initial":
cb_init.as_execution_phase(next_phase="primary"),
"primary":
cb_primary.as_execution_phase(next_phase="primary")
},
initial_phase="initial")
def test_arrays_and_linalg(python_method_impl):
with CodeBuilder(name="primary") as cb:
cb("n", "4")
cb("nodes", "`<builtin>array`(n)")
cb("vdm", "`<builtin>array`(n*n)")
cb("identity", "`<builtin>array`(n*n)")
cb("nodes[i]", "i/n", loops=[("i", 0, "n")])
cb("identity[i]", "0", loops=[("i", 0, "n*n")])
cb("identity[i*n + i]", "1", loops=[("i", 0, "n")])
cb("vdm[j*n + i]", "nodes[i]**j", loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_inverse", "`<builtin>linear_solve`(vdm, identity, n, n)")
cb("myarray", "`<builtin>matmul`(vdm, vdm_inverse, n, n)")
cb((), "`<builtin>print`(myarray)")
cb.yield_state("myarray", "result", 0, "final")
from utils import execute_and_return_single_result
code = create_DAGCode_with_steady_phase(cb.statements)
result = execute_and_return_single_result(python_method_impl, code)
result = result.reshape(4, 4, order="F")
assert la.norm(result - np.eye(4)) < 1e-10
def test_svd(python_method_impl):
with CodeBuilder(name="primary") as cb:
cb("n", 3)
cb("nodes", "`<builtin>array`(n)")
cb("vdm", "`<builtin>array`(n*n)")
cb("identity", "`<builtin>array`(n*n)")
cb("nodes[i]", "i/n", loops=[("i", 0, "n")])
cb("vdm[j*n + i]", "nodes[i]**j", loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_u, vdm_sigma, vdm_vt", "`<builtin>svd`(vdm, n)")
cb("vdm_usigma", "`<builtin>array`(n*n)")
cb("vdm_v", "`<builtin>array`(n*n)")
cb("vdm_usigma[i + j*n]",
"vdm_u[i + j*n] * vdm_sigma[j]",
loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_v[i + j*n]",
"vdm_vt[j + i*n]",
loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_2", "`<builtin>matmul`(vdm_usigma, vdm_vt, n, n)")
cb("diff", "vdm-vdm_2")
cb((), "`<builtin>print`(diff)")
cb.yield_state("diff", "result", 0, "final")
from utils import execute_and_return_single_result
code = create_DAGCode_with_steady_phase(cb.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert la.norm(result) < 1e-10
def test_CodeBuilder_assign(python_method_impl):
with CodeBuilder("phase") as builder:
builder(var("x"), 1)
builder.yield_state(var("x"), "x", 0, "final")
code = create_DAGCode_with_steady_phase(builder.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert result == 1
def test_class_preamble():
from dagrt.language import CodeBuilder
with CodeBuilder(name="primary") as cb:
cb.assign("<t>", "<t> + <dt>")
cb.yield_state("f()", "f", 0, "final")
code = create_DAGCode_with_steady_phase(cb.statements)
from dagrt.codegen import PythonCodeGenerator
import dagrt.function_registry as freg
preamble = """
@staticmethod
def f():
return 1
"""
f = freg.Function(identifier="f",
language_to_codegen={"python": lambda *args: "self.f()"})
generator = PythonCodeGenerator(
"PythonMethod",
class_preamble=preamble,
function_registry=freg.base_function_registry.register(f))
class_ = generator.get_class(code)
method = class_(function_map={})
method.set_up(t_start=0, dt_start=1, context={})
events = list(method.run(t_end=1))
assert events
assert isinstance(events[0], class_.StateComputed)
assert events[0].state_component == 1
def test_missing_state_detection():
"""Check that the code generator detects there is a missing state."""
from dagrt.language import CodeBuilder
with CodeBuilder(name="state_1") as cb:
cb.switch_phase("state_2")
code = create_DAGCode_with_steady_phase(statements=cb.statements)
with pytest.raises(CodeGenerationError):
verify_code(code)
def test_arrays_and_looping(python_method_impl):
with CodeBuilder(name="primary") as cb:
cb("myarray", "`<builtin>array`(20)")
cb("myarray[i]", "i", loops=[("i", 0, 20)])
cb.yield_state("myarray[15]", "result", 0, "final")
from utils import execute_and_return_single_result
code = create_DAGCode_with_steady_phase(cb.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert result == 15
def test_CodeBuilder_condition_with_else_not_taken(python_method_impl):
with CodeBuilder("phase") as builder:
builder(var("x"), 1)
with builder.if_(var("x"), "==", 1):
builder(var("x"), 2)
with builder.else_():
builder(var("x"), 3)
builder.yield_state(var("x"), "x", 0, "final")
code = create_DAGCode_with_steady_phase(builder.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert result == 2
def adaptive_rk_method(tol):
"""
The difference between the results of Euler's method
y_e = y_n + h f(t_n, y_n)
and Heun's method
y_h = y_n + h / 2 (f(t_n, y_n), f(t_n + dt, y_e))
can be used as an estimate of the high order error term in Euler's method.
This allows us to adapt the time step so that the local error falls within a
specified tolerance.
"""
# Set up variables
y = var("<state>y")
g = var("<func>g")
y_e = var("y_e")
y_h = var("y_h")
dt = var("<dt>")
t = var("<t>")
dt_old = var("dt_old")
# Helpers for expression fragments
def norm(val):
return var("<builtin>norm_2")(val)
def dt_scaling(tol, err):
# Return a suitable scaling factor for dt.
# Ensure to guard against excessive increase or divide-by-zero.
from pymbolic.primitives import Max, Min
return Min(((tol / Max((1.0e-16, norm(err))))**(1 / 2), 2.0))
# Code for the main state
with CodeBuilder("adaptrk") as cb:
# Euler
cb(y_e, y + dt * g(t, y))
# Heun
cb(y_h, y + dt / 2 * (g(t, y) + g(t + dt, y_e)))
# Adaptation
cb(dt_old, dt)
cb(dt, dt * dt_scaling(tol, y_h - y_e))
# Accept or reject step
with cb.if_(norm(y_h - y_e), ">=", tol):
cb.fail_step()
with cb.else_():
cb(y, y_h)
cb(t, t + dt_old)
return DAGCode.from_phases_list([cb.as_execution_phase("adaptrk")],
"adaptrk")
def test_elementwise_abs():
with CodeBuilder(name="primary") as cb:
cb("i", "<builtin>array(20)")
cb("i[j]", "-j", loops=(("j", 0, 20), ))
# Test new builtin on an array type.
cb("k", "<builtin>elementwise_abs(i)")
with cb.if_("k[20] > 19"):
cb.raise_(AbsFailure)
with cb.if_("k[20] < 19"):
cb.raise_(AbsFailure)
# Test new builtin on a scalar.
cb("l", "<builtin>elementwise_abs(-20)")
with cb.if_("l > 20"):
cb.raise_(AbsFailure)
with cb.if_("l < 20"):
cb.raise_(AbsFailure)
cb("y", "<func>f(0, <state>ytype)")
cb("<state>ytype", "y")
# Test new builtin on a usertype.
cb("<state>ytype", "<builtin>elementwise_abs(<state>ytype)")
# (We check this in the outer test code)
code = create_DAGCode_with_steady_phase(cb.statements)
rhs_function = "<func>f"
from dagrt.function_registry import (base_function_registry,
register_ode_rhs)
freg = register_ode_rhs(base_function_registry,
"ytype",
identifier=rhs_function,
input_names=("y", ))
freg = freg.register_codegen(
rhs_function, "fortran",
f.CallCode("""
${result} = -2*${y}
"""))
codegen = f.CodeGenerator(
"element_abs_test",
function_registry=freg,
user_type_map={"ytype": f.ArrayType((100, ), f.BuiltinType("real*8"))},
timing_function="second")
code_str = codegen(code)
run_fortran([
("element_abs.f90", code_str),
("test_element_abs.f90", read_file("test_element_abs.f90")),
],
fortran_libraries=["lapack", "blas"])
def demo_rk_adaptive():
from dagrt.language import CodeBuilder
from pymbolic import var
# Set tolerance
tol = 1e-3
# Bulk declare symbolic names
k1, k2, k3, k4, t, dt, dt_old, y, y_hi, y_lo, f, norm = (
var(name) for name in
"k1 k2 k3 k4 <t> <dt> dt_old <state>y y_hi y_lo <func>f <builtin>norm_inf"
.split()) # noqa
with CodeBuilder("primary") as cb:
# Calculate the RK stage values
cb(k1, f(t, y))
cb(k2, f(t + 1 / 2 * dt, y + dt * (1 / 2 * k1)))
cb(k3, f(t + 3 / 4 * dt, y + dt * (3 / 4 * k2)))
cb(k4, f(t + dt, y + dt * (2 / 9 * k1 + 1 / 3 * k2 + 4 / 9 * k3)))
# Compute the low and high order solutions
cb(y_lo, y + dt * (7 / 24 * k1 + 1 / 4 * k2 + 1 / 3 * k3 + 1 / 8 * k4))
cb(y_hi, y + dt * (2 / 9 * k1 + 1 / 3 * k2 + 4 / 9 * k3))
# Save the value of dt
cb(dt_old, dt)
# Update dt based on the error estimate
err_est = norm(y_lo - y_hi)
order = 3
cb(dt, 0.9 * dt * (tol / err_est)**(1 / order))
# Adapt the step size
with cb.if_(err_est, "<=", tol):
# Update t and y
cb(y, y_hi)
cb(t, t + dt_old)
cb.yield_state(expression=y,
component_id="y",
time=t,
time_id=None) # noqa
with cb.else_():
cb.fail_step()
from dagrt.language import DAGCode
code = DAGCode(
phases={"primary": cb.as_execution_phase(next_phase="primary")},
initial_phase="primary")
print(code)
# Generate and run the method.
from dagrt.codegen import PythonCodeGenerator
cls = PythonCodeGenerator("RKAdaptiveMethod").get_class(code)
eocrec = get_convergence_data(cls, problem=KapsProblem(0.001))
print(eocrec.pretty_print())
def test_CodeBuilder_restart_step(python_method_impl):
with CodeBuilder("init") as builder_init:
builder_init("<p>x", "0")
with CodeBuilder("state1") as builder1:
builder1("<p>x", "<p>x + 1")
with builder1.if_("<p>x == 1"):
builder1.restart_step()
with CodeBuilder("state2") as builder2:
builder2.yield_state(var("<p>x"), "x", 0, "final")
phases = [
builder_init.as_execution_phase(next_phase="state1"),
builder1.as_execution_phase(next_phase="state2"),
builder2.as_execution_phase(next_phase="state2")
]
code = DAGCode.from_phases_list(phases, "init")
result = execute_and_return_single_result(python_method_impl,
code,
max_steps=4)
assert result == 2
def generate(self, solver_hook):
"""Return code that implements the implicit Euler method for the single
state component supported."""
with CodeBuilder(name="primary") as cb:
self._make_primary(cb)
code = DAGCode.from_phases_list(
[cb.as_execution_phase(next_phase="primary")],
initial_phase="primary")
from leap.implicit import replace_AssignImplicit
return replace_AssignImplicit(code,
{self.SOLVER_EXPRESSION_ID: solver_hook})
def demo_rk_implicit():
from dagrt.language import CodeBuilder
from pymbolic import var
k1, k2, t, dt, y, f = (
var(name) for name in
"k1 k2 <t> <dt> <state>y <func>f".split())
gamma = (2 - 2**0.5) / 2
with CodeBuilder("primary") as cb:
cb.assign_implicit_1(
k1,
solve_component=k1,
expression=k1 - f(t + gamma * dt, y + dt * gamma * k1),
guess=f(t, y))
cb.assign_implicit_1(
k2,
solve_component=k2,
expression=k2 - f(t + dt, y + dt * ((1 - gamma) * k1 + gamma * k2)), # noqa
guess=k1)
cb(y, y + dt * ((1 - gamma) * k1 + gamma * k2))
cb(t, t + dt)
cb.yield_state(y, "y", t, None)
from dagrt.language import DAGCode
code = DAGCode(phases={
"primary": cb.as_execution_phase(next_phase="primary")
},
initial_phase="primary")
def solver_hook(solve_expr, unknown, solver_id, guess):
from dagrt.expression import match, substitute
pieces = match(
"k - <func>rhs(time, y + dt * (c0 + c1 * k))",
solve_expr,
pre_match={"k": unknown})
return substitute("-10 * (dt * c0 + y) / (10 * dt * c1 + 1)", pieces)
from leap.implicit import replace_AssignImplicit
code = replace_AssignImplicit(code, solver_hook)
print(code)
from dagrt.codegen import PythonCodeGenerator
IRKMethodBuilder = PythonCodeGenerator("IRKMethodBuilder").get_class(code)
eocrec = get_convergence_data(IRKMethodBuilder, SimpleDecayProblem())
print(eocrec.pretty_print())
def generate(self):
"""
:returns: a :class:`~dagrt.language.DAGCode` that can be used to
generate code for this method.
"""
# {{{ check coefficients are explicit
nstages = len(self.alpha)
for n in range(1, nstages + 1):
if len(self.alpha[n - 1]) > n or len(self.beta[n - 1]) > n:
raise ValueError("only explicit SSP schemes are supported")
# }}}
# {{{ primary phase
comp_id = self.component_id
with CodeBuilder(name="primary") as cb:
stages = [self.state
] + [cb.fresh_var(f"s{i}") for i in range(nstages)]
for i in range(0, nstages):
states = sum(alpha * stages[j]
for j, alpha in enumerate(self.alpha[i]))
rhss = sum(beta * self.rhs_func(t=self.t + self.c[i] * self.dt,
**{comp_id: stages[j]})
for j, beta in enumerate(self.beta[i]))
expr = states + self.dt * rhss
if self.state_filter is not None:
expr = self.state_filter(**{comp_id: expr})
cb(stages[i + 1], expr)
# finish
cb(self.state, stages[-1])
cb.yield_state(self.state, comp_id, self.t + self.dt, "final")
cb(self.t, self.t + self.dt)
# }}}
return DAGCode(phases={
"primary": cb.as_execution_phase(next_phase="primary"),
},
initial_phase="primary")
def test_arrays_and_linalg():
from dagrt.function_registry import base_function_registry as freg
with CodeBuilder(name="primary") as cb:
cb("n", "4")
cb("nodes", "`<builtin>array`(n)")
cb("vdm", "`<builtin>array`(n*n)")
cb("identity", "`<builtin>array`(n*n)")
cb("nodes[i]", "i/n", loops=[("i", 0, "n")])
cb("identity[i]", "0", loops=[("i", 0, "n*n")])
cb("identity[i*n + i]", "1", loops=[("i", 0, "n")])
cb("vdm[j*n + i]", "nodes[i]**j", loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_inverse", "`<builtin>linear_solve`(vdm, identity, n, n)")
cb("myarray", "`<builtin>matmul`(vdm, vdm_inverse, n, n)")
cb("myzero", "myarray - identity")
cb((), "`<builtin>print`(myzero)")
with cb.if_("`<builtin>norm_2`(myzero) > 10**(-8)"):
cb.raise_(MatrixInversionFailure)
code = create_DAGCode_with_steady_phase(cb.statements)
codegen = f.CodeGenerator("arrays",
function_registry=freg,
user_type_map={},
emit_instrumentation=True,
timing_function="second")
code_str = codegen(code)
if 0:
with open("arrays.f90", "wt") as outf:
outf.write(code_str)
run_fortran([
("arrays.f90", code_str),
("test_arrays_and_linalg.f90",
read_file("test_arrays_and_linalg.f90")),
],
fortran_libraries=["lapack", "blas"])
def get_IfThenElse_test_code_and_expected_result():
from dagrt.expression import IfThenElse
with CodeBuilder(name="primary") as cb:
cb(var("c1"), IfThenElse(True, 0, 1))
cb(var("c2"), IfThenElse(False, 0, 1))
cb(var("c3"), IfThenElse(IfThenElse(True, True, False), 0, 1))
cb(var("c4"), IfThenElse(IfThenElse(False, True, False), 0, 1))
cb(var("c5"), IfThenElse(True, IfThenElse(True, 0, 1), 2))
cb(var("c6"), IfThenElse(True, IfThenElse(False, 0, 1), 2))
cb(var("c7"), IfThenElse(False, 0, IfThenElse(True, 1, 2)))
cb(var("c8"), IfThenElse(False, 0, IfThenElse(False, 1, 2)))
cb(var("c9"), 1 + IfThenElse(True, 0, 1))
cb(var("c10"), 1 + IfThenElse(False, 0, 1))
cb.yield_state(tuple(var("c" + str(i)) for i in range(1, 11)),
"result", 0, "final")
code = create_DAGCode_with_steady_phase(cb.statements)
return (code, (0, 1, 0, 1, 0, 1, 1, 2, 1, 2))
def test_self_dep_in_loop():
with CodeBuilder(name="primary") as cb:
cb("y", "<state>y")
cb("y",
"<func>f(0, 2*i*<func>f(0, y if i > 2 else 2*y))",
loops=(("i", 0, 5), ))
cb("<state>y", "y")
code = create_DAGCode_with_steady_phase(cb.statements)
rhs_function = "<func>f"
from dagrt.function_registry import (base_function_registry,
register_ode_rhs)
freg = register_ode_rhs(base_function_registry,
"ytype",
identifier=rhs_function,
input_names=("y", ))
freg = freg.register_codegen(
rhs_function, "fortran",
f.CallCode("""
${result} = -2*${y}
"""))
codegen = f.CodeGenerator(
"selfdep",
function_registry=freg,
user_type_map={"ytype": f.ArrayType((100, ), f.BuiltinType("real*8"))},
timing_function="second")
code_str = codegen(code)
run_fortran([
("selfdep.f90", code_str),
("test_selfdep.f90", read_file("test_selfdep.f90")),
],
fortran_libraries=["lapack", "blas"])
def generate(self):
"""
:returns: :class:`dagrt.language.DAGCode`
"""
from pytools import UniqueNameGenerator
name_gen = UniqueNameGenerator()
from dagrt.language import DAGCode, CodeBuilder
array = var("<builtin>array")
rhs_var = var("rhs_var")
# Initialization
with CodeBuilder(name="initialization") as cb_init:
cb_init(self.step, 1)
# Primary
with CodeBuilder(name="primary") as cb_primary:
if not self.static_dt:
time_history_data = self.time_history + [self.t]
time_hist_var = var(name_gen("time_history"))
cb_primary(time_hist_var, array(self.hist_length))
for i in range(self.hist_length):
cb_primary(time_hist_var[i], time_history_data[i] - self.t)
time_hist = time_hist_var
t_end = self.dt
dt_factor = 1
else:
time_hist = list(range(-self.hist_length+1, 0+1)) # noqa pylint:disable=invalid-unary-operand-type
dt_factor = self.dt
t_end = 1
cb_primary(rhs_var, self.eval_rhs(self.t, self.state))
history = self.history + [rhs_var]
ab_sum = emit_ab_integration(
cb_primary, name_gen,
self.function_family,
time_hist, history,
0, t_end)
state_est = self.state + dt_factor * ab_sum
if self.state_filter is not None:
state_est = self.state_filter(state_est)
cb_primary(self.state, state_est)
# Rotate history and time history.
for i in range(self.hist_length - 1):
cb_primary(self.history[i], history[i + 1])
if not self.static_dt:
cb_primary(self.time_history[i], time_history_data[i + 1])
cb_primary(self.t, self.t + self.dt)
cb_primary.yield_state(expression=self.state,
component_id=self.component_id,
time_id='', time=self.t)
if self.hist_length == 1:
# The first order method requires no bootstrapping.
return DAGCode(
phases={
"initial": cb_init.as_execution_phase(next_phase="primary"),
"primary": cb_primary.as_execution_phase(next_phase="primary")
},
initial_phase="initial")
# Bootstrap
with CodeBuilder(name="bootstrap") as cb_bootstrap:
self.rk_bootstrap(cb_bootstrap)
cb_bootstrap(self.t, self.t + self.dt)
cb_bootstrap.yield_state(expression=self.state,
component_id=self.component_id,
time_id='', time=self.t)
cb_bootstrap(self.step, self.step + 1)
with cb_bootstrap.if_(self.step, "==", self.hist_length):
cb_bootstrap.switch_phase("primary")
return DAGCode(
phases={
"initialization": cb_init.as_execution_phase("bootstrap"),
"bootstrap": cb_bootstrap.as_execution_phase("bootstrap"),
"primary": cb_primary.as_execution_phase("primary"),
},
initial_phase="initialization")
def generate_butcher(self, stage_coeff_set_names, stage_coeff_sets,
rhs_funcs, estimate_coeff_set_names,
estimate_coeff_sets):
"""
:arg stage_coeff_set_names: a list of names/string identifiers
for stage coefficient sets
:arg stage_coeff_sets: a mapping from set names to stage coefficients
:arg rhs_funcs: a mapping from set names to right-hand-side
functions
:arg estimate_coeffs_set_names: a list of names/string identifiers
for estimate coefficient sets
:arg estimate_coeffs_sets: a mapping from estimate coefficient set
names to cofficients.
"""
from pymbolic import var
comp = self.component_id
dt = self.dt
t = self.t
state = self.state
nstages = len(self.c)
# {{{ check coefficients for plausibility
for name in stage_coeff_set_names:
for istage in range(nstages):
coeff_sum = sum(stage_coeff_sets[name][istage])
assert abs(coeff_sum - self.c[istage]) < 1e-12, (
name, istage, coeff_sum, self.c[istage])
# }}}
# {{{ initialization
last_rhss = {}
with CodeBuilder(name="initialization") as cb:
for name in stage_coeff_set_names:
if (name in self.recycle_last_stage_coeff_set_names
and _is_first_stage_same_as_last_stage(
self.c, stage_coeff_sets[name])):
last_rhss[name] = var("<p>last_rhs_" + name)
cb(last_rhss[name], rhs_funcs[name](t=t, **{comp: state}))
cb_init = cb
# }}}
stage_rhs_vars = {}
rhs_var_to_unknown = {}
for name in stage_coeff_set_names:
stage_rhs_vars[name] = [
cb.fresh_var(f"rhs_{name}_s{i}") for i in range(nstages)
]
# These are rhss if they are not yet known and pending an implicit solve.
for i, rhsvar in enumerate(stage_rhs_vars[name]):
unkvar = cb.fresh_var(f"unk_{name}_s{i}")
rhs_var_to_unknown[rhsvar] = unkvar
knowns = set()
# {{{ stage loop
last_state_est_var = cb.fresh_var("last_state_est")
last_state_est_var_valid = False
with CodeBuilder(name="primary") as cb:
equations = []
unknowns = set()
def make_known(v):
unknowns.discard(v)
knowns.add(v)
for istage in range(nstages):
for name in stage_coeff_set_names:
c = self.c[istage]
my_rhs = stage_rhs_vars[name][istage]
if (name in self.recycle_last_stage_coeff_set_names
and istage == 0
and _is_first_stage_same_as_last_stage(
self.c, stage_coeff_sets[name])):
cb(my_rhs, last_rhss[name])
make_known(my_rhs)
else:
is_implicit = False
state_increment = 0
for src_name in stage_coeff_set_names:
coeffs = stage_coeff_sets[src_name][istage]
for src_istage, coeff in enumerate(coeffs):
rhsval = stage_rhs_vars[src_name][src_istage]
if rhsval not in knowns:
unknowns.add(rhsval)
is_implicit = True
state_increment += dt * coeff * rhsval
state_est = state + state_increment
if (self.state_filter is not None and not (
# reusing last output state
c == 0 and all(
len(stage_coeff_sets[src_name][istage]) ==
0 for src_name in stage_coeff_set_names))):
state_est = self.state_filter(state_est)
if is_implicit:
rhs_expr = rhs_funcs[name](t=t + c * dt,
**{
comp: state_est
})
from dagrt.expression import collapse_constants
solve_expression = collapse_constants(
my_rhs - rhs_expr,
list(unknowns) + [self.state], cb.assign,
cb.fresh_var)
equations.append(solve_expression)
if istage + 1 == nstages:
last_state_est_var_valid = False
else:
if istage + 1 == nstages:
cb(last_state_est_var, state_est)
state_est = last_state_est_var
last_state_est_var_valid = True
rhs_expr = rhs_funcs[name](t=t + c * dt,
**{
comp: state_est
})
cb(my_rhs, rhs_expr)
make_known(my_rhs)
# {{{ emit solve if possible
if unknowns and len(unknowns) == len(equations):
# got a square system, let's solve
assignees = [unk.name for unk in unknowns]
from pymbolic import substitute
subst_dict = {
rhs_var.name: rhs_var_to_unknown[rhs_var]
for rhs_var in unknowns
}
cb.assign_implicit(
assignees=assignees,
solve_components=[
rhs_var_to_unknown[unk].name
for unk in unknowns
],
expressions=[
substitute(eq, subst_dict) for eq in equations
],
# TODO: Could supply a starting guess
other_params={"guess": state},
solver_id="solve")
del equations[:]
knowns.update(unknowns)
unknowns.clear()
# }}}
# Compute solution estimates.
estimate_vars = [
cb.fresh_var("est_" + name)
for name in estimate_coeff_set_names
]
for iest, name in enumerate(estimate_coeff_set_names):
out_coeffs = estimate_coeff_sets[name]
if (last_state_est_var_valid and # noqa: W504
_is_last_stage_same_as_output(self.c, stage_coeff_sets,
out_coeffs)):
state_est = last_state_est_var
else:
state_increment = 0
for src_name in stage_coeff_set_names:
state_increment += sum(
coeff * stage_rhs_vars[src_name][src_istage]
for src_istage, coeff in enumerate(out_coeffs))
state_est = state + dt * state_increment
if self.state_filter is not None:
state_est = self.state_filter(state_est)
cb(estimate_vars[iest], state_est)
# This updates <t>.
self.finish(cb, estimate_coeff_set_names, estimate_vars)
# These updates have to happen *after* finish because before we
# don't yet know whether finish will accept the new state.
for name in stage_coeff_set_names:
if (name in self.recycle_last_stage_coeff_set_names
and _is_first_stage_same_as_last_stage(
self.c, stage_coeff_sets[name])):
cb(last_rhss[name], stage_rhs_vars[name][-1])
cb_primary = cb
# }}}
return DAGCode(phases={
"initial":
cb_init.as_execution_phase(next_phase="primary"),
"primary":
cb_primary.as_execution_phase(next_phase="primary")
},
initial_phase="initial")