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 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_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 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_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 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 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 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
def python_method_impl_codegen(code, **kwargs):
from dagrt.codegen import PythonCodeGenerator
codegen = PythonCodeGenerator(class_name="Method")
return codegen.get_class(code)(**kwargs)
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)
