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
