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")