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 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 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 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 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 fuse_two_dags(dag1,
dag2,
phase_correspondences=None,
should_disambiguate_name=None):
from dagrt.language import DAGCode
new_phases = {}
for phase_name in frozenset(dag1.phases) | frozenset(dag2.phases):
phase1 = dag1.phases.get(phase_name)
phase2 = dag2.phases.get(phase_name)
new_phases[phase_name] = fuse_two_phases(phase_name, phase1, phase2)
if dag1.initial_phase != dag2.initial_phase:
raise ValueError("DAGs don't agree on initial phase")
return DAGCode(new_phases, dag1.initial_phase)
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 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 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_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")
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 strang_splitting(dag1, dag2, stepping_phase):
"""Given two time advancement routines (in *dag1* and *dag2*), returns a
single second-order accurate time advancement routine representing the sum
of both of those advancements.
:arg dag1: a :class:`dagrt.language.DAGCode`
:arg dag2: a :class:`dagrt.language.DAGCode`
:arg stepping_phase: the name of the phase in *dag1* and *dag2* that carries
out time stepping to which Strang splitting is to be applied.
:returns: a :class:`dagrt.language.DAGCode`
"""
from pymbolic.mapper.substitutor import make_subst_func, SubstitutionMapper
# {{{ disambiguate
id1 = dag1.existing_var_names()
id2 = dag1.existing_var_names()
from pytools import UniqueNameGenerator
vng = UniqueNameGenerator(id1 | id2)
from pymbolic import var
subst2 = {}
for clash in id1 & id2:
if not clash.startswith("<") or clash.startswith("<p>"):
unclash = vng(clash)
subst2[clash] = var(unclash)
subst2_mapper = SubstitutionMapper(make_subst_func(subst2))
# }}}
all_phases = frozenset(dag1.phases) | frozenset(dag2.phases)
from dagrt.language import DAGCode, ExecutionPhase
new_phases = {}
for phase_name in all_phases:
phase1 = dag1.phases.get(phase_name)
phase2 = dag2.phases.get(phase_name)
substed_s2_stmts = [
stmt.map_expressions(subst2_mapper) for stmt in phase2.statements
]
if phase_name == stepping_phase:
assert phase1 is not None
assert phase2 is not None
from pymbolic import var
dt_half = SubstitutionMapper(
make_subst_func({"<dt>": var("<dt>") / 2}))
phase1_half_dt = [
stmt.map_expressions(dt_half) for stmt in phase1.statements
]
if phase1.next_phase != phase2.next_phase:
raise ValueError(
"DAGs don't agree on default "
f"phase transition out of phase '{phase_name}'")
s2_name = phase_name + "_s2"
s3_name = phase_name + "_s3"
assert s2_name not in all_phases
assert s3_name not in all_phases
"""
du/dt = A + B
Time interval is [0,1]
1. Starting with u0, solve du / dt = A from t = 0 to 1/2, get u1
2. Starting with u1, solve du / dt = B from t = 0 to 1, get u2
3. Starting with u2, solve du / dt = A from t = 1/2 to 1, get u3
4. Return u3
"""
new_phases[phase_name] = ExecutionPhase(
name=phase_name,
next_phase=s2_name,
statements=(_update_t_by_dt_factor(
0, _elide_yield_state(phase1_half_dt))))
new_phases[s2_name] = ExecutionPhase(
name=s2_name,
next_phase=s3_name,
statements=(_update_t_by_dt_factor(
1 / 2, _elide_yield_state(substed_s2_stmts))))
new_phases[s3_name] = ExecutionPhase(name=s3_name,
next_phase=phase1.next_phase,
statements=phase1_half_dt)
else:
from dagrt.transform import fuse_two_phases
new_phases[phase_name] = fuse_two_phases(
phase_name, phase1, phase2.copy(statements=substed_s2_stmts))
if dag1.initial_phase != dag2.initial_phase:
raise ValueError("DAGs don't agree on initial phase")
return DAGCode(new_phases, dag1.initial_phase)