def test_basic_raise_codegen(): """Test code generation of the Raise statement.""" cbuild = RawCodeBuilder() class TimeStepUnderflow(RuntimeError): pass cbuild.add_and_get_ids(Raise(TimeStepUnderflow, "underflow", id="raise")) cbuild.commit() code = create_DAGCode_with_init_and_main_phases( init_statements=[], main_statements=cbuild.statements) codegen = PythonCodeGenerator(class_name="Method") Method = codegen.get_class(code) # noqa method = Method({}) method.set_up(t_start=0, dt_start=0, context={}) try: # initialization for _result in method.run_single_step(): pass # first primary step for _result in method.run_single_step(): raise AssertionError() except method.TimeStepUnderflow: pass except Method.StepError as e: assert e.condition == "TimeStepUnderflow" except Exception as e: assert not e, e
def test_global_name_distinctness(): """Test whether the code generator gives globals distinct names.""" cbuild = RawCodeBuilder() cbuild.add_and_get_ids( Assign(id="assign_y^", assignee="<p>y^", assignee_subscript=(), expression=1), Assign(id="assign_y*", assignee="<p>y*", assignee_subscript=(), expression=0), YieldState(id="return", time=0, time_id="final", expression=var("<p>y^") + var("<p>y*"), component_id="y", depends_on=["assign_y^", "assign_y*"])) cbuild.commit() code = create_DAGCode_with_init_and_main_phases( init_statements=[], main_statements=cbuild.statements) codegen = PythonCodeGenerator(class_name="Method") Method = codegen.get_class(code) # noqa method = Method({}) method.set_up(t_start=0, dt_start=0, context={}) hist = list(method.run(max_steps=2)) assert len(hist) == 3 assert isinstance(hist[1], method.StateComputed) assert hist[1].state_component == 1
def main(): def rhs(t, y): u, v = y return np.array([v, -u / t**2], dtype=np.float64) def soln(t): inner = np.sqrt(3) / 2 * np.log(t) return np.sqrt(t) * (5 * np.sqrt(3) / 3 * np.sin(inner) + np.cos(inner)) from dagrt.codegen import PythonCodeGenerator codegen = PythonCodeGenerator("AdaptiveRK") tolerances = [1.0e-1, 1.0e-2, 1.0e-3, 1.0e-5] errors = [] for tol in tolerances: method = adaptive_rk_method(tol) AdaptiveRK = codegen.get_class(method) # noqa: N806 solver = AdaptiveRK({"<func>g": rhs}) solver.set_up(t_start=1.0, dt_start=0.1, context={"y": np.array([1., 3.])}) for evt in solver.run(t_end=10.0): final_time = evt.t errors.append(np.abs(solver.global_state_y[0] - soln(final_time))) print("Tolerance\tError") print("-" * 25) for tol, error in zip(tolerances, errors): print(f"{tol:.2e}\t{error:.2e}")
def test_basic_codegen(): """Test whether the code generator returns a working method. The generated method always returns 0.""" cbuild = RawCodeBuilder() cbuild.add_and_get_ids( YieldState(id="return", time=0, time_id="final", expression=0, component_id="<state>", depends_on=[])) cbuild.commit() code = create_DAGCode_with_init_and_main_phases( init_statements=[], main_statements=cbuild.statements) codegen = PythonCodeGenerator(class_name="Method") print(codegen(code)) Method = codegen.get_class(code) # noqa method = Method({}) method.set_up(t_start=0, dt_start=0, context={}) hist = [s for s in method.run(max_steps=2)] assert len(hist) == 3 assert isinstance(hist[0], method.StepCompleted) assert hist[0].current_phase == "init" assert isinstance(hist[1], method.StateComputed) assert hist[1].state_component == 0 assert isinstance(hist[2], method.StepCompleted) assert hist[2].current_phase == "main"
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_basic_conditional_codegen(): """Test whether the code generator generates branches properly.""" cbuild = RawCodeBuilder() cbuild.add_and_get_ids( Assign(id="then_branch", assignee="<state>y", assignee_subscript=(), expression=1, condition=True), Assign(id="else_branch", assignee="<state>y", assignee_subscript=(), expression=0, condition=False), Nop(id="branch", depends_on=["then_branch", "else_branch"]), YieldState(id="return", time=0, time_id="final", expression=var("<state>y"), component_id="<state>", depends_on=["branch"])) cbuild.commit() code = create_DAGCode_with_init_and_main_phases( init_statements=[], main_statements=cbuild.statements) codegen = PythonCodeGenerator(class_name="Method") Method = codegen.get_class(code) # noqa method = Method({}) method.set_up(t_start=0, dt_start=0, context={"y": 6}) hist = [s for s in method.run(max_steps=2)] assert len(hist) == 3 assert isinstance(hist[1], method.StateComputed) assert hist[1].state_component == 1 assert isinstance(hist[2], method.StepCompleted)
def test_basic_fail_step_codegen(): """Test code generation of the FailStep statement.""" cbuild = RawCodeBuilder() cbuild.add_and_get_ids(FailStep(id="fail")) cbuild.commit() code = create_DAGCode_with_init_and_main_phases( init_statements=[], main_statements=cbuild.statements) codegen = PythonCodeGenerator(class_name="Method") Method = codegen.get_class(code) # noqa method = Method({}) method.set_up(t_start=0, dt_start=0, context={}) print(codegen(code)) for _evt in method.run_single_step(): pass with pytest.raises(method.FailStepException): for evt in method.run_single_step(): print(evt)
def test_function_name_distinctness(): """Test whether the code generator gives functions distinct names.""" cbuild = RawCodeBuilder() cbuild.add_and_get_ids( YieldState(id="return", time=0, time_id="final", expression=var("<func>y^")() + var("<func>y*")(), component_id="y")) cbuild.commit() code = create_DAGCode_with_init_and_main_phases( init_statements=[], main_statements=cbuild.statements) codegen = PythonCodeGenerator(class_name="Method") Method = codegen.get_class(code) # noqa method = Method({"<func>y^": lambda: 0, "<func>y*": lambda: 1}) method.set_up(t_start=0, dt_start=0, context={}) hist = list(method.run(max_steps=2)) assert len(hist) == 3 assert isinstance(hist[1], method.StateComputed) assert hist[1].state_component == 1
def find_stability_bdry(code, prec, make_k, angle): # Generated code doesn't pickle well->generate here. from dagrt.codegen import PythonCodeGenerator integrator_cls = PythonCodeGenerator("Integrator").get_class(code) def predicate(amag): return is_stable(integrator_cls, make_k(angle, amag)) mag = find_truth_bdry(predicate, prec=prec) return make_k(angle, mag)
def make_multirate_method(f2f, s2f, f2s, s2s, ratio=2, order=3): """Return the object that drives the multirate method for the given parameters.""" FastestFirst = "Fq" code = TwoRateAdamsBashforthMethod(FastestFirst, order, ratio).generate() MRABMethod = PythonCodeGenerator(class_name='MRABMethod').get_class(code) rhs_map = {'<func>f2f': f2f, '<func>s2f': s2f, '<func>f2s': f2s, '<func>s2s': s2s} return MRABMethod(rhs_map)
def test_basic_assign_rhs_codegen(): """Test whether the code generator generates RHS evaluation code properly.""" cbuild = RawCodeBuilder() cbuild.add_and_get_ids( Assign(id="assign_rhs1", assignee="<state>y", assignee_subscript=(), expression=var("y")(t=var("<t>")), depends_on=[]), Assign(id="assign_rhs2", assignee="<state>y", assignee_subscript=(), expression=var("yy")(t=var("<t>"), y=var("<state>y")), depends_on=["assign_rhs1"]), YieldState(id="return", time=0, time_id="final", expression=var("<state>y"), component_id="<state>", depends_on=["assign_rhs2"])) cbuild.commit() code = create_DAGCode_with_init_and_main_phases( init_statements=[], main_statements=cbuild.statements) codegen = PythonCodeGenerator(class_name="Method") Method = codegen.get_class(code) # noqa def y(t): return 6 def yy(t, y): return y + 6 method = Method({"y": y, "yy": yy}) method.set_up(t_start=0, dt_start=0, context={"y": 0}) hist = [s for s in method.run(max_steps=2)] assert len(hist) == 3 assert isinstance(hist[1], method.StateComputed) assert hist[1].state_component == 12 assert isinstance(hist[2], method.StepCompleted)
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 set_up_rk4(field_var_name, dt, fields, rhs, t_start=0): from leap.rk import LSRK4MethodBuilder from dagrt.codegen import PythonCodeGenerator dt_method = LSRK4MethodBuilder(component_id=field_var_name) dt_code = dt_method.generate() dt_stepper_class = PythonCodeGenerator("TimeStep").get_class(dt_code) dt_stepper = dt_stepper_class({"<func>" + dt_method.component_id: rhs}) dt_stepper.set_up(t_start=t_start, dt_start=dt, context={dt_method.component_id: fields}) return dt_stepper
def generate_singlerate_leap_advancer(timestepper, component_id, rhs, t, dt, state): """Generate Leap code to advance all state at the same timestep, without substepping. Parameters ---------- timestepper An instance of :class:`leap.MethodBuilder` that advances the state from t=time to t=(time+dt), and returns the advanced state. component_id State id (required input for leap method generation) rhs Function that should return the time derivative of the state. This function should take time and state as arguments, with a call looking like rhs(t, state). t: float Time at which to start dt: float Initial timestep to be set by leap method state: numpy.ndarray Agglomerated object array containing at least the state variables that will be advanced by this stepper Returns ------- dagrt.codegen.python.StepperInterface Python class implementing leap method, and generated by dagrt """ code = timestepper.generate() from dagrt.codegen import PythonCodeGenerator codegen = PythonCodeGenerator(class_name="Method") stepper_cls = codegen.get_class(code)(function_map={ "<func>" + component_id: rhs, }) stepper_cls.set_up(t_start=t, dt_start=dt, context={component_id: state}) return stepper_cls
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 run(): from leap.rk.imex import KennedyCarpenterIMEXARK4MethodBuilder from dagrt.codegen import PythonCodeGenerator # Construct the method generator. mgen = KennedyCarpenterIMEXARK4MethodBuilder("y", atol=_atol) # Generate the code for the method. code = mgen.generate() from leap.implicit import replace_AssignImplicit code = replace_AssignImplicit(code, {"solve": solver_hook}) IMEXIntegrator = PythonCodeGenerator("IMEXIntegrator").get_class(code) # Set up the problem and run the method. from functools import partial problem = KapsProblem(epsilon=0.001) integrator = IMEXIntegrator( function_map={ "<func>expl_y": problem.nonstiff, "<func>impl_y": problem.stiff, "<func>solver": partial(solver, problem.stiff, problem.jacobian), "<func>j": problem.jacobian }) integrator.set_up(t_start=problem.t_start, dt_start=1.0e-1, context={"y": problem.initial()}) t = None y = None for event in integrator.run(t_end=problem.t_end): if isinstance(event, integrator.StateComputed): t = event.t y = event.state_component print("Error: " + str(np.linalg.norm(y - problem.exact(t))))
def python_method_impl_codegen(code, **kwargs): from dagrt.codegen import PythonCodeGenerator codegen = PythonCodeGenerator(class_name="Method") #with open("outf.py", "w") as outf: # outf.write(codegen(code)) return codegen.get_class(code)(**kwargs)
def python_method_impl_codegen(code, **kwargs): from dagrt.codegen import PythonCodeGenerator codegen = PythonCodeGenerator(class_name="Method") return codegen.get_class(code)(**kwargs)
def make_3_component_multirate_method(problem, ratios, order=3, code_only=False, return_rhs_map=False): """Return the object that drives the multirate method for the given parameters. """ from leap.multistep.multirate import (MultiRateMultiStepMethodBuilder, MultiRateHistory) ncomponents = 3 assert problem.ncomponents == ncomponents code = MultiRateMultiStepMethodBuilder( default_order=3, system_description=(("dt", "comp0", "=", MultiRateHistory(ratios[0], "<func>rhs0to0", ("comp0", )), MultiRateHistory(ratios[0], "<func>rhs1to0", ("comp1", ))), ("dt", "comp1", "=", MultiRateHistory(ratios[1], "<func>rhs0to1", ("comp0", )), MultiRateHistory(ratios[1], "<func>rhs1to1", ("comp1", )), MultiRateHistory(ratios[1], "<func>rhs2to1", ("comp2", ))), ("dt", "comp2", "=", MultiRateHistory(ratios[2], "<func>rhs1to2", ("comp1", )), MultiRateHistory(ratios[2], "<func>rhs2to2", ("comp2", )))), static_dt=True).generate() from functools import partial if code_only and not return_rhs_map: return code rhs_map = {} for i in range(3): rhs_map["<func>rhs%dto%d" % (i, i)] = partial(problem.rhs, i, i) if i > 0: rhs_map["<func>rhs%dto%d" % (i, i - 1)] = (partial( problem.rhs, i, i - 1)) if i < ncomponents - 1: rhs_map["<func>rhs%dto%d" % (i, i + 1)] = (partial( problem.rhs, i, i + 1)) if code_only: if return_rhs_map: return code, rhs_map else: return code from dagrt.codegen import PythonCodeGenerator MRABMethod = PythonCodeGenerator(class_name="MRABMethod").get_class( code) # noqa if return_rhs_map: return MRABMethod(rhs_map), rhs_map else: return MRABMethod(rhs_map)
def test_dependent_state(order=3, step_ratio=3): # Solve # f' = f+s # s' = -f+s def true_f(t): return np.exp(t) * np.sin(t) def true_s(t): return np.exp(t) * np.cos(t) method = MultiRateMultiStepMethodBuilder(order, ( ( "dt", "fast", "=", MRHistory(1, "<func>f", ( "two_fast", "slow", )), ), ("dt", "slow", "=", MRHistory(step_ratio, "<func>s", ("fast", "slow"))), ( "two_fast", "=", MRHistory(step_ratio, "<func>twice", ("fast", )), ), ), static_dt=True) code = method.generate() print(code) from pytools.convergence import EOCRecorder eocrec = EOCRecorder() from dagrt.codegen import PythonCodeGenerator codegen = PythonCodeGenerator(class_name="Method") stepper_cls = codegen.get_class(code) for n in range(4, 7): t = 0 dt = 2**(-n) final_t = 10 stepper = stepper_cls( function_map={ "<func>f": lambda t, two_fast, slow: 0.5 * two_fast + slow, "<func>s": lambda t, fast, slow: -fast + slow, "<func>twice": lambda t, fast: 2 * fast, }) stepper.set_up(t_start=t, dt_start=dt, context={ "fast": true_f(t), "slow": true_s(t), }) f_times = [] f_values = [] s_times = [] s_values = [] for event in stepper.run(t_end=final_t): if isinstance(event, stepper_cls.StateComputed): if event.component_id == "fast": f_times.append(event.t) f_values.append(event.state_component) elif event.component_id == "slow": s_times.append(event.t) s_values.append(event.state_component) else: assert False, event.component_id f_times = np.array(f_times) s_times = np.array(s_times) f_values_true = true_f(f_times) s_values_true = true_s(s_times) f_err = f_values - f_values_true s_err = s_values - s_values_true error = ( la.norm(f_err) / la.norm(f_values_true) + # noqa: W504 la.norm(s_err) / la.norm(s_values_true)) eocrec.add_data_point(dt, error) print(eocrec.pretty_print()) orderest = eocrec.estimate_order_of_convergence()[0, 1] assert orderest > 3 * 0.95