diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d2e24ec..73b98fb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -68,6 +68,9 @@ jobs: - name: "Main Script" run: | CONDA_ENVIRONMENT=.test-conda-env-py3.yml + #echo "git+https://github.com/ecisneros8/pyrometheus.git#egg=pyrometheus" >> requirements.txt + echo "git+https://github.com/cmikida2/pyrometheus.git@jax#egg=pyrometheus" >> requirements.txt + curl -L -O -k https://tiker.net/ci-support-v0 . ./ci-support-v0 build_py_project_in_conda_env @@ -109,6 +112,9 @@ jobs: - name: "Main Script" run: | CONDA_ENVIRONMENT=.test-conda-env-py3.yml + #echo "git+https://github.com/ecisneros8/pyrometheus.git#egg=pyrometheus" >> requirements.txt + echo "git+https://github.com/cmikida2/pyrometheus.git@jax#egg=pyrometheus" >> requirements.txt + curl -L -O -k https://tiker.net/ci-support-v0 . ./ci-support-v0 build_py_project_in_conda_env diff --git a/.test-conda-env-py3.yml b/.test-conda-env-py3.yml index 2331cb1..792018f 100644 --- a/.test-conda-env-py3.yml +++ b/.test-conda-env-py3.yml @@ -8,6 +8,9 @@ dependencies: - git - numpy - scipy +- matplotlib +- jax +- cantera - matplotlib-base - pip diff --git a/leap/__init__.py b/leap/__init__.py index 8d4c75f..569bd45 100644 --- a/leap/__init__.py +++ b/leap/__init__.py @@ -163,6 +163,136 @@ def norm(expr): # }}} +# {{{ adaptivity w/order changing + +class AdaptiveOrderMethodBuilderMixin(MethodBuilder): + """ + This class expected the following members to be defined: state, t, dt. + """ + + def __init__(self, atol=0, rtol=0, max_dt_growth=None, + min_dt_shrinkage=None, error_const=None): + self.adaptive = bool(atol or rtol) + self.atol = atol + self.rtol = rtol + + if max_dt_growth is None: + max_dt_growth = 5 + + if min_dt_shrinkage is None: + min_dt_shrinkage = 0.1 + + self.max_dt_growth = max_dt_growth + self.min_dt_shrinkage = min_dt_shrinkage + + self.error_const = error_const + + def finish_nonadaptive(self, cb, high_order_estimate, low_order_estimate): + raise NotImplementedError() + + def rotate_history(self, cb, high_order_estimate, low_order_estimate): + raise NotImplementedError() + + def rescale_interval(self, cb): + raise NotImplementedError() + + def finish_adaptive(self, cb, high_order_estimate, low_order_estimate, + hist, order, factor, equal_steps): + from pymbolic import var + from pymbolic.primitives import Comparison, LogicalOr, Max, Min + from dagrt.expression import IfThenElse + + rel_error_raw = var("rel_error_raw") + rel_error = var("rel_error") + rel_error_p1 = var("rel_error_p1") + rel_error_m1 = var("rel_error_m1") + factor_p1 = var("factor_p1") + factor_m1 = var("factor_m1") + + def norm(expr): + return var("norm_2")(expr) + + def weighted_norm(x, y): + # Weighted RMS norm. + # FIXME: need to support array of relative tolerances? + denom = self.rtol * var("elementwise_abs")(y) + self.atol + length = var("len")(y) + return norm(x/denom) / length ** 0.5 + + cb(rel_error_raw, weighted_norm( + self.error_const[order]*(high_order_estimate - low_order_estimate), + high_order_estimate)) + + cb(rel_error, IfThenElse(Comparison(rel_error_raw, "==", 0), + 1.0e-14, rel_error_raw)) + + with cb.if_(LogicalOr((Comparison(rel_error, ">", 1), + var("isnan")(rel_error)))): + + with cb.if_(var("isnan")(rel_error)): + cb(self.dt, self.min_dt_shrinkage * self.dt) + with cb.else_(): + # Default Scipy BDF safety factor. + cb(self.dt, Max((0.81 * self.dt + * rel_error ** (-1.0 / (order + 1)), + self.min_dt_shrinkage * self.dt))) + cb(factor, Max((0.81 + * rel_error ** (-1.0 / (order + 1)), + self.min_dt_shrinkage))) + self.rescale_interval(cb) + cb(self.dt_monitor, self.dt) + cb(equal_steps, 0) + with cb.if_(self.t + self.dt, "==", self.t): + cb.raise_(TimeStepUnderflow) + with cb.else_(): + cb.fail_step() + + with cb.else_(): + # Rotate state history. + self.rotate_history(cb, high_order_estimate, low_order_estimate) + cb(equal_steps, equal_steps+1) + cb(factor, 1) + cb(factor_p1, 0) + cb(factor_m1, 0) + with cb.if_(Comparison(equal_steps, ">=", order+1)): + + # Now we need to use the history to determine if we should + # update order. + with cb.if_(Comparison(order, ">", 1)): + cb(rel_error_m1, weighted_norm( + self.error_const[order-1]*hist[order], + high_order_estimate)) + cb(factor_m1, 0.81 + * rel_error_m1 ** (-1.0 / (order))) + with cb.else_(): + cb(factor_m1, 0) + with cb.if_(Comparison(order, "<", 5)): + cb(rel_error_p1, weighted_norm( + self.error_const[order+1]*hist[order+2], + high_order_estimate)) + cb(factor_p1, 0.81 + * rel_error_p1 ** (-1.0 / (order + 2))) + with cb.else_(): + cb(rel_error_p1, 0) + # Based on these new error norms, do we modify the order? + cb(factor, 0.81 * rel_error ** (-1.0 / (order + 1))) + cb(factor, Max((factor, factor_p1, factor_m1))) + with cb.if_(Comparison(factor, "==", factor_p1)): + cb(order, order + 1) + with cb.if_(Comparison(factor, "==", factor_m1)): + cb(order, order - 1) + cb(factor, Min((self.max_dt_growth, factor))) + self.rescale_interval(cb) + cb(equal_steps, 0) + # This updates :
should not be set before this is called. + self.finish_nonadaptive(cb, high_order_estimate, + low_order_estimate) + cb(self.dt, self.dt * factor) + cb(self.dt_monitor, self.dt) + +# }}} + + # {{{ diagnostics class TimeStepUnderflow(RuntimeError): diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index 660dbbc..3541680 100644 --- a/leap/multistep/__init__.py +++ b/leap/multistep/__init__.py @@ -4,7 +4,7 @@ __copyright__ = """ Copyright (C) 2007 Andreas Kloeckner Copyright (C) 2014, 2015 Matt Wala -Copyright (C) 2015 Cory Mikida +Copyright (C) 2015, 2021 Cory Mikida """ __license__ = """ @@ -29,7 +29,7 @@ import numpy as np import numpy.linalg as la -from leap import MethodBuilder +from leap import MethodBuilder, AdaptiveOrderMethodBuilderMixin from pymbolic import var @@ -37,6 +37,7 @@ .. autoclass:: AdamsIntegrationFunctionFamily .. autoclass:: AdamsMonomialIntegrationFunctionFamily .. autoclass:: AdamsBashforthMethodBuilder +.. autoclass:: AdaptiveBDFMethodBuilder """ @@ -205,6 +206,73 @@ def emit_adams_extrapolation(cb, name_gen, # }}} +# {{{ BDF/NDF helper functions + + +def emit_R_computation(cb, order, factor, name_gen): + """Compute the matrix for changing the differences array. + + Source: Shampine, Lawrence F., and Mark W. Reichelt. + "The matlab ode suite." Section 2.2""" + array = var("array") + k = var(name_gen("m_k")) + i = var(name_gen("mat_I")) + cb(i, array(order)) + cb(i[k-1], k, loops=[(k.name, 1, order+1)]) + # I = np.arange(1, order + 1)[:, None] + j = var(name_gen("mat_J")) + cb(j, array(order)) + cb(j[k-1], k, loops=[(k.name, 1, order+1)]) + # J = np.arange(1, order + 1) + m = var(name_gen("M")) + cb(m, array((order + 1)*(order + 1))) + # M = np.zeros((order + 1, order + 1)) + l_ind = var(name_gen("m_l")) + cb(m[k + l_ind*(order+1)], 0, + loops=[(l_ind.name, 0, order+1), (k.name, 0, order+1)]) + cb(m[k+1 + (l_ind+1)*(order+1)], (i[l_ind] - 1 - factor * j[k]) / i[l_ind], + loops=[(l_ind.name, 0, order), (k.name, 0, order)]) + cb(m[l_ind], 1, loops=[(l_ind.name, 0, order+1)]) + # M[1:, 1:] = (I - 1 - factor * J) / I + transpose = var("transpose") + r = var(name_gen("R")) + cb(r, array((order + 1)*(order + 1))) + cb(r[k], m[k], loops=[(k.name, 0, order+1)]) + cb(r[k + (l_ind)*(order+1)], m[k + (l_ind)*(order+1)] + * r[k + (l_ind - 1)*(order+1)], + loops=[(l_ind.name, 1, order+1), (k.name, 0, order+1)]) + cb(r, transpose(r, order+1)) + # return np.cumprod(M, axis=0) + return r + + +def change_D(cb, d, order, factor): + """Change differences array in-place when step size is changed. + Source: Shampine, Lawrence F., and Mark W. Reichelt. + "The matlab ode suite." Section 2.2""" + from pytools import UniqueNameGenerator + name_gen = UniqueNameGenerator() + r = emit_R_computation(cb, order, factor, name_gen) + u = emit_R_computation(cb, order, 1, name_gen) + ru = var(name_gen("RU")) + matmul = var("matmul") + # RU = R.dot(U) + cb(ru, matmul(r, u, order+1, order+1)) + i = var(name_gen("m_i")) + j = var(name_gen("m_j")) + array_utype = var("array_utype") + d_out = var(name_gen("D_out")) + cb(d_out, array_utype(order+1, d[0])) + # D[:order + 1] = np.dot(RU.T, D[:order + 1]) + cb(d_out[j], ru[j*(order+1)] * d[0], + loops=[(j.name, 0, order+1)]) + cb(d_out[j], d_out[j] + ru[i + j*(order+1)] * d[i], + loops=[(i.name, 1, order+1), (j.name, 0, order+1)]) + cb(d[j], d_out[j], loops=[(j.name, 0, order+1)]) + + +# }}} + # {{{ ab method @@ -416,4 +484,326 @@ def rk_bootstrap(self, cb): # }}} +# {{{ Scipy-esque adaptive BDF/NDF + + +class AdaptiveBDFMethodBuilder(AdaptiveOrderMethodBuilderMixin): + """ + Source: Shampine, Lawrence F., and Mark W. Reichelt. + "The matlab ode suite." SIAM journal on scientific + computing 18.1 (1997): 1-22. + + User-supplied context: + + component_id: The value that is integrated + + component_id: The right hand side function + """ + + def __init__(self, component_id, state_filter_name=None, fixed_order=None, + use_high_order=True, atol=0, rtol=0, max_dt_growth=None, + min_dt_shrinkage=None, ndf=False, bootstrap_factor=None): + + # For adaptive BDF/NDF, we fix maximum order of 5. + self.max_order = 5 + + self.component_id = component_id + + # Declare variables + if fixed_order: + self.fixed_order = fixed_order + else: + self.fixed_order = 0 + if self.fixed_order and atol != 0: + raise ValueError("Cannot run adaptively with fixed order") + if self.fixed_order and rtol != 0: + raise ValueError("Cannot run adaptively with fixed order") + # Take smaller faux-bootstrap steps + if self.fixed_order > 1: + if bootstrap_factor: + self.bootstrap_factor = bootstrap_factor + else: + self.bootstrap_factor = 1 + self.step = var("

step") + self.step_factor = var("

step_factor") + self.equal_steps = var("

equal_steps") + self.order = var("

order") + self.function = var("" + component_id) + self.history = var("

hist") + self.kappa = var("

kappa") + self.gamma = var("

gamma") + self.alpha = var("

alpha") + self.dt_monitor = var("

dt_monitor") + self.error_const = var("

error_const") + + self.state = var("" + component_id) + self.t = var("") + self.dt = var("

") + self.dt_save = var("

dt_save") + + self.ndf = ndf + + self.state_filter_name = state_filter_name + if state_filter_name is not None: + self.state_filter = var("" + state_filter_name) + else: + self.state_filter = None + + AdaptiveOrderMethodBuilderMixin.__init__( + self, + atol=atol, + rtol=rtol, + max_dt_growth=max_dt_growth, + min_dt_shrinkage=min_dt_shrinkage, + error_const=self.error_const) + + self.atol = atol + self.rtol = rtol + self.use_high_order = use_high_order + + def eval_rhs(self, t, y): + """Return a node that evaluates the RHS at the given time and + component value.""" + from pymbolic.primitives import CallWithKwargs + return CallWithKwargs(function=self.function, + parameters=(), + kw_parameters={"t": t, self.component_id: y}) + + def generate(self): + """ + :returns: :class:`dagrt.language.DAGCode` + """ + + from dagrt.language import DAGCode, CodeBuilder + + # Initialization + with CodeBuilder(name="initialization") as cb_init: + j = var("j") + rhs_dummy = var("rhs_dummy") + array = var("array") + array_utype = var("array_utype") + cb_init(self.kappa, array(6)) + cb_init(self.gamma, array(6)) + cb_init(self.alpha, array(6)) + cb_init(self.error_const, array(6)) + cb_init(self.history, array_utype(8, self.state)) + cb_init(self.order, 1) + cb_init(self.equal_steps, 0) + # Faux-bootstrap steps needed + if self.fixed_order > 1: + cb_init(self.dt_save, self.dt) + cb_init(self.dt, self.dt/self.bootstrap_factor) + # NDF vs. BDF: + if self.ndf: + # Optimized for minimal truncation error. + cb_init(self.kappa[0], 0) + cb_init(self.kappa[1], -0.1850) + cb_init(self.kappa[2], -1/9) + cb_init(self.kappa[3], -0.0823) + cb_init(self.kappa[4], -0.0415) + cb_init(self.kappa[5], 0) + cb_init(self.error_const[0], 1) + cb_init(self.error_const[1], 0.315) + cb_init(self.error_const[2], 0.16666667) + cb_init(self.error_const[3], 0.09911667) + cb_init(self.error_const[4], 0.11354167) + cb_init(self.error_const[5], 0.16666667) + else: + cb_init(self.kappa[j], 0, loops=[(j.name, 0, 6)]) + cb_init(self.error_const[0], 1) + cb_init(self.error_const[1], 0.5) + cb_init(self.error_const[2], 0.33333333) + cb_init(self.error_const[3], 0.25) + cb_init(self.error_const[4], 0.2) + cb_init(self.error_const[5], 0.16666667) + cb_init(self.gamma[0], 0) + cb_init(self.gamma[1], 1) + cb_init(self.gamma[2], 1.5) + cb_init(self.gamma[3], 1.83333333) + cb_init(self.gamma[4], 2.08333333) + cb_init(self.gamma[5], 2.28333333) + cb_init(self.alpha, (1 - self.kappa) * self.gamma) + cb_init(self.history[0], self.state) + cb_init(rhs_dummy, self.dt * self.eval_rhs(self.t, self.state)) + cb_init(self.history[1], rhs_dummy) + cb_init(self.step, 1) + cb_init(self.dt_monitor, self.dt) + + # Primary + with CodeBuilder(name="primary") as cb_primary: + self.generate_primary(cb_primary) + + # Based on an initialization at first order, and with + # a predictor-corrector formulation, we require 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") + + def rotate_history(self, cb, state_est_high, state_est_low): + # Rotate the history. + i = var("i") + cb(self.history[self.order + 2], (state_est_high - state_est_low) + - self.history[self.order + 1]) + cb(self.history[self.order + 1], (state_est_high - state_est_low)) + cb(self.history[self.order-i], self.history[self.order-i] + + self.history[self.order+1-i], + loops=[(i.name, 0, self.order+1)]) + + def rescale_interval(self, cb): + # Incorporate the scaling change + # required when the timestep changes. + change_D(cb, self.history, self.order, self.step_factor) + + def yield_state(self, cb): + # Update state and time. + cb(self.t, self.t + self.dt) + # If fixed order specified, increase order + # after each step until we hit it ("quasi-bootstrap") + if self.fixed_order > 1: + from pymbolic.primitives import Comparison + with cb.if_(Comparison(self.order, "<", self.fixed_order)): + cb(self.order, self.order + 1) + with cb.if_(Comparison(self.dt, "!=", self.dt_save)): + cb(self.step_factor, 1.1) + with cb.if_(Comparison(self.t, "==", self.dt_save)): + cb(self.step_factor, self.dt_save / self.dt) + cb(self.dt, self.dt_save) + cb(self.dt_monitor, self.dt) + self.rescale_interval(cb) + with cb.else_(): + with cb.if_(Comparison(self.t + 1.1*self.dt, ">", self.dt_save)): + cb(self.step_factor, (self.dt_save - self.t)/self.dt) + cb(self.dt, self.dt_save - self.t) + cb(self.dt_monitor, self.dt) + self.rescale_interval(cb) + with cb.else_(): + cb(self.dt, self.step_factor*self.dt) + cb(self.dt_monitor, self.dt) + self.rescale_interval(cb) + cb.yield_state(expression=self.state, + component_id=self.component_id, + time_id="", time=self.t) + + def generate_primary(self, cb): + + from pytools import UniqueNameGenerator + name_gen = UniqueNameGenerator() + rhs_next_var = var("rhs_next_var") + rhs_var_to_unknown = {} + unkvar = cb.fresh_var("unk") + rhs_var_to_unknown[rhs_next_var] = unkvar + equations = [] + unknowns = set() + knowns = set() + + unknowns.add(rhs_next_var) + + # First, check to see if dt got modified out from + # under us in order to adhere to some end time. + from pymbolic.primitives import Comparison + with cb.if_(Comparison(self.dt, "!=", self.dt_monitor)): + cb(self.step_factor, self.dt/self.dt_monitor) + cb(self.dt_monitor, self.dt) + self.rescale_interval(cb) + + # Use existing scaled histories to form prediction. + j = var(name_gen("m_j")) + y_predict = var("y_predict") + cb(y_predict, self.history[0]) + cb(y_predict, y_predict + self.history[j], + loops=[(j.name, 1, self.order+1)]) + + state_est_low = y_predict + + # Form expression(s) for correction. + psi = var("psi") + cb(psi, 0) + cb(psi, psi + (1 / self.alpha[self.order]) + * self.gamma[j]*self.history[j], + loops=[(j.name, 1, self.order+1)]) + state_est_high = (self.dt / self.alpha[self.order]) * \ + rhs_next_var - psi + state_est_low + + # Build the implicit solve expression. + from dagrt.expression import collapse_constants + from pymbolic.mapper.distributor import DistributeMapper as DistMap + solve_expression = collapse_constants( + rhs_next_var - self.eval_rhs(self.t + self.dt, + DistMap()(state_est_high)), + list(unknowns) + [self.state], + cb.assign, cb.fresh_var) + equations.append(solve_expression) + + # {{{ 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_est_low}, + solver_id="solve") + + del equations[:] + knowns.update(unknowns) + unknowns.clear() + elif len(unknowns) > len(equations): + raise ValueError("BDF/NDF implicit timestep has more " + "unknowns than equations") + elif len(unknowns) < len(equations): + raise ValueError("BDF/NDF implicit timestep has more " + "equations than unknowns") + + # }}} + + # Update the state now that we've solved. + if self.state_filter is not None: + state_est_low = self.state_filter(state_est_low) + state_est_high = self.state_filter(state_est_high) + + # this is needed in case the order changes + high_est = var("high_est") + low_est = var("low_est") + cb(high_est, state_est_high) + cb(low_est, state_est_low) + + self.finish(cb, high_est, low_est) + + def finish(self, cb, high_est, low_est): + if not self.adaptive: + self.rotate_history(cb, high_est, low_est) + self.finish_nonadaptive(cb, high_est, low_est) + else: + self.finish_adaptive(cb, high_est, low_est, self.history, + self.order, self.step_factor, + self.equal_steps) + + def finish_nonadaptive(self, cb, high_est, low_est): + if self.use_high_order: + est = high_est + else: + est = low_est + + cb(self.state, est) + self.yield_state(cb) + +# }}} + # vim: fdm=marker diff --git a/requirements.txt b/requirements.txt index de5cd4e..d5dbf72 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ git+https://github.com/inducer/pytools.git#egg=pytools git+https://github.com/inducer/pymbolic.git#egg=pymbolic -git+https://github.com/inducer/dagrt.git#egg=dagrt +git+https://github.com/cmikida2/dagrt.git@bdf-builtins#=egg=dagrt diff --git a/test/reactor_system.py b/test/reactor_system.py new file mode 100644 index 0000000..4643d93 --- /dev/null +++ b/test/reactor_system.py @@ -0,0 +1,158 @@ +def make_jax_pyro_class(ptk_base_cls, usr_np): + + class PyroJaxNumpy(ptk_base_cls): + + def _pyro_make_array(self, res_list): + """This works around (e.g.) numpy.exp not working with object arrays of numpy + scalars. It defaults to making object arrays, however if an array + consists of all scalars, it makes a "plain old" :class:`numpy.ndarray`. + + See ``this numpy bug `__ + for more context. + """ + + from numbers import Number + # Needed to play nicely with Jax, which frequently creates + # arrays of shape () when handed numbers + all_numbers = all( + isinstance(e, Number) + or (isinstance(e, self.usr_np.ndarray) and e.shape == ()) + for e in res_list) + + if all_numbers: + return self.usr_np.array(res_list, dtype=self.usr_np.float64) + + result = self.usr_np.empty_like(res_list, dtype=object, + shape=(len(res_list),)) + + # 'result[:] = res_list' may look tempting, however: + # https://github.com/numpy/numpy/issues/16564 + for idx in range(len(res_list)): + result[idx] = res_list[idx] + + return result + + def _pyro_norm(self, argument, normord): + """This works around numpy.linalg norm not working with scalars. + + If the argument is a regular ole number, it uses :func:`numpy.abs`, + otherwise it uses ``usr_np.linalg.norm``. + """ + # Wrap norm for scalars + from numbers import Number + if isinstance(argument, Number): + return self.usr_np.abs(argument) + # Needed to play nicely with Jax, which frequently creates + # arrays of shape () when handed numbers + if isinstance(argument, self.usr_np.ndarray) and argument.shape == (): + return self.usr_np.abs(argument) + return self.usr_np.linalg.norm(argument, normord) + + return PyroJaxNumpy(usr_np=usr_np) + + +class ReactorSystemOde(object): + """ + Cantera example problem: reactor2.py + + Two reactors connected with a piston, with heat loss to the environment + + This script simulates the following situation. A closed cylinder with volume 2 + m^3 is divided into two equal parts by a massless piston that moves with speed + proportional to the pressure difference between the two sides. It is + initially held in place in the middle. One side is filled with 1000 K argon at + 20 atm, and the other with a combustible 500 K methane/air mixture at 0.1 atm + (phi = 1.1). At t = 0 the piston is released and begins to move due to the + large pressure difference, compressing and heating the methane/air mixture, + which eventually explodes. At the same time, the argon cools as it expands. + The piston is adiabatic, but some heat is lost through the outer cylinder + walls to the environment. + + Note that this simulation, being zero-dimensional, takes no account of shock + wave propagation. It is somewhat artifical, but nevertheless instructive. + """ + def __init__(self, gas1, gas2, np): + import cantera as ct + import pyrometheus as pyro + ptk_base_cls_1 = pyro.get_thermochem_class(gas1) + ptk_base_cls_2 = pyro.get_thermochem_class(gas2) + self.gas1 = make_jax_pyro_class(ptk_base_cls_1, np) + self.gas2 = make_jax_pyro_class(ptk_base_cls_2, np) + self.gas2_ct = gas2 + self.env = ct.Reservoir(ct.Solution("air.xml")) + # Initial volume of each reactor is 1.0, so... + self.np = np + self.gas1_mass = gas1.density + self.gas2_mass = gas2.density + self.gas1_mw = self.np.asarray(gas1.molecular_weights) + self.gas2_mw = self.np.asarray(gas2.molecular_weights) + + def __call__(self, t, y): + + # State vector is [U1, V1, Y_1, Y_2, ... Y_K, U2, V2, Y_1, .... Y_K] + # Set gases. + temp1 = y[0] + rho1 = self.gas1_mass * (1.0/y[1]) + p1 = self.gas1.get_pressure(rho1, temp1, self.np.asarray([y[2]])) + + temp2 = y[3] + rho2 = self.gas2_mass * (1.0/y[4]) + p2 = self.gas2.get_pressure(rho2, temp2, self.np.asarray(y[5:])) + + # Volume rate of change (move piston based on pressure + # difference between reactors) + area = 1.0 + k = 0.5e-4 + u1 = 100.0 + dvdt_1 = k*area*(p1 - p2) + + # Mass fraction rate of change (via production + # rates as is typical) + # Reactor 1 is pure argon + wdot_1 = 0.0 + dydt_1 = wdot_1 * self.gas1_mw * (1.0 / rho1) + + # Internal energy rate of change + # Pressure work first. + dtempdt_1 = -p1*dvdt_1 + # Include heat transfer via the piston wall + dtempdt_1 += -area*u1*(temp1 - temp2) + # Partial molar internal energies + e0_rt1 = (self.gas1.get_species_enthalpies_rt(temp1) - 1.0) * \ + self.gas1.gas_constant * temp1 + dtempdt_1 += -self.np.dot(e0_rt1, wdot_1*(1.0/rho1)*self.gas1_mass) + dtempdt_1 = dtempdt_1 * (1.0 + / (self.gas1_mass + * self.gas1.get_mixture_specific_heat_cv_mass(temp1, + self.np.asarray([y[2]])))) + + # Volume rate of change (move piston based on pressure + # difference between reactors) + dvdt_2 = k*area*(p2 - p1) + + # Mass fraction rate of change (via production + # rates as is typical) + wdot_2 = self.gas2.get_net_production_rates(rho2, temp2, + self.np.asarray(y[5:])) + dydt_2 = wdot_2 * self.gas2_mw * (1.0 / rho2) + + # Internal energy rate of change + # Pressure work first. + dtempdt_2 = -p2*dvdt_2 + # Include heat transfer via the piston wall + dtempdt_2 += -area*u1*(temp2 - temp1) + # Include heat loss to the environment (air reservoir) + # via specified wall + area2 = 1.0 + u2 = 500.0 + dtempdt_2 += -area2*u2*(temp2 - self.env.T) + e0_rt2 = (self.gas2.get_species_enthalpies_rt(temp2) - 1.0) \ + * self.gas2.gas_constant * temp2 + dtempdt_2 += -self.np.dot(e0_rt2, wdot_2*(1.0/rho2)*self.gas2_mass) + dtempdt_2 = dtempdt_2 * (1.0 + / (self.gas2_mass + * self.gas2.get_mixture_specific_heat_cv_mass(temp2, + self.np.asarray(y[5:])))) + + return self.np.hstack((self.np.hstack((dtempdt_1, dvdt_1, dydt_1)), + self.np.hstack((dtempdt_2, dvdt_2, dydt_2)))) diff --git a/test/test_bdf.py b/test/test_bdf.py new file mode 100755 index 0000000..3a0936c --- /dev/null +++ b/test/test_bdf.py @@ -0,0 +1,181 @@ +#! /usr/bin/env python + +__copyright__ = "Copyright (C) 2014 Andreas Kloeckner" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +# avoid spurious: pytest.mark.parametrize is not callable +# pylint: disable=not-callable + +import sys +import pytest + +from leap.multistep import AdaptiveBDFMethodBuilder +from stiff_test_systems import VanDerPolProblem, KapsProblem +import numpy as np + +import logging + +from utils import ( # noqa + python_method_impl_interpreter as pmi_int, + python_method_impl_codegen as pmi_cg) + +logger = logging.getLogger(__name__) + + +def kaps_solver(f, t, sub_y, coeff, guess): + from scipy.optimize import root + return root(lambda unk: unk - f(t=t, y=sub_y + coeff*unk), guess).x + + +def VDPJac(t, y): + jac = np.zeros((2, 2)) + mu = 30 + jac[0, 0] = 0 + jac[0, 1] = 1 + jac[1, 0] = -2*mu*y[0]*y[1] - 1 + jac[1, 1] = -mu*(y[0]*y[0] - 1) + return jac + + +def newton_solver(t, sub_y, coeff, guess): + + vdp = VanDerPolProblem() + d = 0 + corr_norm = 1.0 + reltol = 1e-6 + abstol = 0 + y_old = guess.copy() + y_guess = guess.copy() + corr_weights = np.zeros(2) + # Match Scipy BDF + newton_tol = 0.01 + newton_maxiter = 4 + corr_norm_old = None + converged = False + # Try pulling this out of the loop. + jac = VDPJac(t, y_old) + import scipy.linalg as la + lu = la.lu_factor(np.eye(2) - coeff*jac, overwrite_a=True) + # Check convergence w/weighted norm... + for j in range(0, 2): + corr_weights[j] = (reltol * np.abs(y_old[j]) + abstol) + psi = -(sub_y - guess) + for _i in range(newton_maxiter): + rhs = vdp(t, y_guess) + corr = la.lu_solve(lu, coeff*rhs - psi - d, overwrite_b=True) + y_guess += corr + d += corr + # RMS norm: + corr_norm = np.linalg.norm(corr / corr_weights) / 2 ** 0.5 + if corr_norm_old is not None: + rate = corr_norm / corr_norm_old + else: + rate = None + if rate is not None and rate / (1 - rate) * corr_norm < newton_tol: + converged = True + break + corr_norm_old = corr_norm + + if converged is False: + raise ValueError("Newton failed to converge") + + # Return RHS that gives y_n+1 = y_guess + return (y_guess + psi - y_old)/coeff + + +def solver_hook(solve_expr, solve_var, solver_id, guess): + from dagrt.expression import match, substitute + + pieces = match("unk - rhs(t=t, y=sub_y + coeff*unk)", solve_expr, + pre_match={"unk": solve_var}) + pieces["guess"] = guess + return substitute("solver(t, sub_y, coeff, guess)", pieces) + + +@pytest.mark.parametrize("problem, method, expected_order", [ + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=1, max_dt_growth=10), 1], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=1, max_dt_growth=10, ndf=True), 1], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=2, max_dt_growth=10, bootstrap_factor=10), 2], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=2, max_dt_growth=10, bootstrap_factor=10, ndf=True), 2], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=3, max_dt_growth=10, bootstrap_factor=100), 3], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=3, max_dt_growth=10, bootstrap_factor=100, ndf=True), 3], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=4, max_dt_growth=10, bootstrap_factor=1000), 4], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=4, max_dt_growth=10, bootstrap_factor=1000, ndf=True), 4], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=5, max_dt_growth=10, bootstrap_factor=20000), 5], + [KapsProblem(epsilon=0.9), AdaptiveBDFMethodBuilder( + "y", fixed_order=5, max_dt_growth=10, bootstrap_factor=20000, ndf=True), 5], + ]) +def test_convergence(python_method_impl, problem, method, expected_order, + show_dag=False, plot_solution=False): + pytest.importorskip("scipy") + + from functools import partial + dts = 2**-np.array(range(3, 8), dtype=np.float64) + function_map = {"y": problem, + "solver": partial(kaps_solver, problem)} + from utils import check_simple_convergence + check_simple_convergence(method=method, method_impl=python_method_impl, + expected_order=expected_order, show_dag=show_dag, + plot_solution=plot_solution, + function_map=function_map, dts=dts, + implicit=True, problem=problem, + solver_hook=solver_hook) + + +# {{{ adaptive test + +@pytest.mark.parametrize("method", [ + AdaptiveBDFMethodBuilder( + "y", rtol=1e-6), + AdaptiveBDFMethodBuilder( + "y", rtol=1e-6, ndf=True), + ]) +def test_adaptive_timestep(python_method_impl, method, show_dag=False, + plot=False): + pytest.importorskip("scipy") + from utils import check_adaptive_timestep + check_adaptive_timestep(python_method_impl=python_method_impl, method=method, + ss_frac=0.7, bs_frac=0.2, ss_targ=0.01, + bs_targ=0.05, show_dag=show_dag, plot=plot, + implicit=True, solver=newton_solver, + solver_hook=solver_hook) + +# }}} + + +if __name__ == "__main__": + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: filetype=pyopencl:fdm=marker diff --git a/test/test_bdf_reactors.py b/test/test_bdf_reactors.py new file mode 100644 index 0000000..3f5dd4c --- /dev/null +++ b/test/test_bdf_reactors.py @@ -0,0 +1,314 @@ +""" +Cantera example problem: reactor2.py + +Two reactors connected with a piston, with heat loss to the environment + +This script simulates the following situation. A closed cylinder with volume 2 +m^3 is divided into two equal parts by a massless piston that moves with speed +proportional to the pressure difference between the two sides. It is +initially held in place in the middle. One side is filled with 1000 K argon at +20 atm, and the other with a combustible 500 K methane/air mixture at 0.1 atm +(phi = 1.1). At t = 0 the piston is released and begins to move due to the +large pressure difference, compressing and heating the methane/air mixture, +which eventually explodes. At the same time, the argon cools as it expands. +The piston is adiabatic, but some heat is lost through the outer cylinder +walls to the environment. + +Note that this simulation, being zero-dimensional, takes no account of shock +wave propagation. It is somewhat artifical, but nevertheless instructive. +""" + +import numpy as np +import pytest + + +def test_vs_scipy(): + pytest.importorskip("scipy") + pytest.importorskip("cantera") + pytest.importorskip("pyrometheus") + pytest.importorskip("jax") + + import scipy.linalg as la + import scipy.integrate + import cantera as ct + from reactor_system import ReactorSystemOde + import jax + from jax import jit, jacfwd + import jax.numpy as jnp + + jax.config.update("jax_enable_x64", 1) + + #------------------------------ + # First create each gas needed. + #------------------------------ + + # create an argon gas object and set its state + ar = ct.Solution("argon.xml") + ar.TP = 1000.0, 20.0 * ct.one_atm + + # use GRI-Mech 3.0 for the methane/air mixture, and set its initial state + gas = ct.Solution("gri30.xml") + gas.TP = 500.0, 0.2 * ct.one_atm + gas.set_equivalence_ratio(1.1, "CH4:1.0", "O2:2, N2:7.52") + + ar_mass = ar.density + gas_mass = gas.density + + # Now the problem is set up, and we're ready to solve it. + #print('finished setup, begin solution...') + + n_steps = 300 + + # Initialize reactor states for Leap. + state = ar.T + state = np.append(state, 1.0) + state = np.append(state, ar.Y) + state = np.append(state, gas.T) + state = np.append(state, 1.0) + state = np.append(state, gas.Y) + + from leap.multistep import AdaptiveBDFMethodBuilder + + rtol = 1e-4 + atol = 1e-16 + + method = AdaptiveBDFMethodBuilder("y", use_high_order=True, ndf=True, + atol=atol, rtol=rtol, max_dt_growth=10, + min_dt_shrinkage=0.2) + + code = method.generate() + + ode = ReactorSystemOde(ar, gas, np) + ode_jax = ReactorSystemOde(ar, gas, jnp) + + #f = lambda y: ode_jax(t, y) + def f(y): + return ode_jax(t, y) + + jacobian = jit(jacfwd(f)) + + def ReactorJac(t, y): + return jacobian(y) + + def newton_solver(t, sub_y, coeff, guess): + + d = 0 + corr_norm = 1.0 + reltol = 1e-4 + abstol = 1e-16 + y_old = guess.copy() + y_guess = guess.copy() + corr_weights = np.zeros(58) + # Match Scipy BDF + newton_tol = 0.01 + newton_maxiter = 4 + corr_norm_old = None + converged = False + jac = ReactorJac(t, y_old) + lu = la.lu_factor(np.eye(58) - coeff*jac, overwrite_a=True) + # Check convergence w/weighted norm... + for j in range(0, 58): + corr_weights[j] = (reltol * np.abs(y_old[j]) + abstol) + psi = -(sub_y - guess) + for _i in range(newton_maxiter): + rhs = ode(t, y_guess) + corr = la.lu_solve(lu, coeff*rhs - psi - d, overwrite_b=True) + y_guess += corr + d += corr + # RMS norm: + corr_norm = np.linalg.norm(corr / corr_weights) / 58 ** 0.5 + if corr_norm_old is not None: + rate = corr_norm / corr_norm_old + else: + rate = None + if rate is not None and rate / (1 - rate) * corr_norm < newton_tol: + converged = True + break + corr_norm_old = corr_norm + + if converged is False: + raise ValueError("Newton failed to converge") + + # Return RHS that gives y_n+1 = y_guess + return (y_guess + psi - y_old)/coeff + + def solver_hook(solve_expr, solve_var, solver_id, guess): + from dagrt.expression import match, substitute + + pieces = match("unk - rhs(t=t, y=sub_y + coeff*unk)", solve_expr, + pre_match={"unk": solve_var}) + pieces["guess"] = guess + return substitute("solver(t, sub_y, coeff, guess)", pieces) + + from leap.implicit import replace_AssignImplicit + code = replace_AssignImplicit(code, {"solve": solver_hook}) + + from dagrt.codegen import PythonCodeGenerator + codegen = PythonCodeGenerator(class_name="Method") + + stepper_cls = codegen.get_class(code) + + t = 0.0 + # Match Scipy's initial timestep + dt = 1.535031777168217e-05 + final_t = 4.e-4*n_steps + + stepper = stepper_cls( + function_map={ + "y": lambda t, y: ode(t, y), + "solver": newton_solver, + }) + + stepper.set_up( + t_start=t, dt_start=dt, + context={ + "y": state, + }) + + times = [] + values = [] + new_times = [] + new_values = [] + step_sizes = [] + last_t = 0.0 + istep = 0 + + for event in stepper.run(t_end=final_t): + if isinstance(event, stepper_cls.StateComputed): + assert event.component_id == "y" + # print("Reactor 2 Temp: T = ", event.state_component[3]) + new_times.append(event.t) + new_values.append(event.state_component) + elif isinstance(event, stepper_cls.StepCompleted): + # print("Step completed: t = ", event.t) + # print("Step completed: dt = ", event.t - last_t) + istep += 1 + if not new_times: + continue + times.extend(new_times) + values.extend(new_values) + step_sizes.append(event.t - last_t) + last_t = event.t + # Get Leap to match Scipy's end behavior + if times[-1] + stepper.dt > final_t: + stepper.dt = final_t - times[-1] + del new_times[:] + del new_values[:] + elif isinstance(event, stepper_cls.StepFailed): + del new_times[:] + del new_values[:] + + times = np.array(times) + + # Need to obtain the pressures and temperatures for both + # the fast and slow reactors. + + # Slow reactor + slow_temps = [] + slow_press = [] + slow_vol = [] + slow_energy = [] + for i in range(0, len(times)): + ar.TDY = values[i][0], ar_mass/values[i][1], values[i][2] + slow_temps.append(ar.T) + slow_press.append(ar.P/1e5) + slow_vol.append(values[i][1]) + slow_energy.append(ar.u) + + # Fast reactor + fast_temps = [] + fast_press = [] + fast_vol = [] + fast_energy = [] + fast_neg = [] + for i in range(0, len(times)): + gas.TDY = values[i][3], gas_mass/values[i][4], values[i][5:] + fast_temps.append(gas.T) + fast_press.append(gas.P/1e5) + fast_vol.append(values[i][4]) + fast_energy.append(gas.u) + fast_neg.append(0.0) + for j in range(0, gas.n_species): + # Track both magnitude and number of negative species + if values[i][5+j] < 0: + fast_neg[-1] -= values[i][5+j] + + step_sizes = np.array(step_sizes) + + # Now, replicate the process using Scipy and compare the two + + t = 0.0 + dt = 4.e-3 + final_t = 4.e-4*n_steps + + # New interface + solver = scipy.integrate.BDF(ode, 0, state, final_t, rtol=rtol, + atol=atol, jac=ReactorJac) + + # Need to obtain the pressures and temperatures for both + # the fast and slow reactors. + + # Slow reactor + sslow_temps = [] + sslow_press = [] + sslow_vol = [] + sslow_mass = [] + sslow_energy = [] + sfast_temps = [] + sfast_press = [] + sfast_vol = [] + sfast_mass = [] + sfast_energy = [] + sfast_species = [] + sfast_neg_sum = [] + sfast_neg_num = [] + for _j in range(0, gas.n_species): + sfast_species.append([]) + stimes = [] + ssteps = [] + + # New interface + while solver.status == "running" and solver.t < final_t: + solver.step() + ar.TDY = solver.y[0], ar_mass/solver.y[1], solver.y[2] + sslow_vol.append(solver.y[1]) + sslow_temps.append(ar.T) + sslow_mass.append(solver.y[1]*ar.density) + sslow_energy.append(ar.u) + sslow_press.append(ar.P/1e5) + gas.TDY = solver.y[3], gas_mass/solver.y[4], solver.y[5:] + frac_sum = 0 + num = 0 + for j in range(0, gas.n_species): + sfast_species[j].append(solver.y[5+j]) + if solver.y[5+j] < 0: + frac_sum -= solver.y[5+j] + num += 1 + sfast_neg_sum.append(frac_sum) + sfast_neg_num.append(num) + sfast_vol.append(solver.y[4]) + sfast_temps.append(gas.T) + sfast_press.append(gas.P/1e5) + sfast_mass.append(solver.y[4]*gas.density) + sfast_energy.append(gas.u) + stimes.append(solver.t) + ssteps.append(solver.t - solver.t_old) + # print("Reactor 2 Temp: ", solver.y[3]) + # print("Time: ", solver.t) + # print("Timestep: ", solver.t - solver.t_old) + + # Testing criterion + # 1.) Ensure Leap solution doesn't contain negative mass fractions + assert max(fast_neg) < 1e-16 + # 2.) Ensure final Reactor 2 temperatures of Scipy and Leap are the same + assert abs(fast_temps[-1] - sfast_temps[-1]) < 0.5 + # 3.) As a proxy for similar (but not exact) adaptive timestepping, + # ensure that Scipy and Leap took a similar number of steps + # to reach the same end time. + assert abs((len(times) - len(stimes))/len(stimes)) < 0.05 + + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) diff --git a/test/test_fortran_bdf.py b/test/test_fortran_bdf.py new file mode 100644 index 0000000..87ce5ca --- /dev/null +++ b/test/test_fortran_bdf.py @@ -0,0 +1,73 @@ +""" +Fortran version of the reactors. All we need for now is to +try running a Fortran code generator with the method. +""" + + +def test_vs_scipy(): + + from leap.multistep import AdaptiveBDFMethodBuilder + + rtol = 1e-4 + atol = 1e-16 + + method = AdaptiveBDFMethodBuilder("y", use_high_order=True, ndf=False, + atol=atol, rtol=rtol, max_dt_growth=10, + min_dt_shrinkage=0.2) + + code = method.generate() + + import dagrt.codegen.fortran as f + + # For now, we will register simple dummy rhs and solver functions. + from dagrt.function_registry import ( + base_function_registry, register_ode_rhs, + register_function, UserType) + rhs_function = "y" + solver_function = "solver" + freg = register_ode_rhs(base_function_registry, "y", + identifier=rhs_function) + freg = freg.register_codegen(rhs_function, "fortran", + f.CallCode(""" + ${result} = -2*${y} + """)) + freg = register_function(freg, solver_function, + ("t", "sub_y", "coeff", "guess",), + result_names=("result",), result_kinds=(UserType("y"),)) + freg = freg.register_codegen(solver_function, "fortran", + f.CallCode(""" + ${result} = -2*${sub_y} + """)) + + def solver_hook(solve_expr, solve_var, solver_id, guess): + from dagrt.expression import match, substitute + + pieces = match("unk - rhs(t=t, y=sub_y + coeff*unk)", solve_expr, + pre_match={"unk": solve_var}) + pieces["guess"] = guess + return substitute("solver(t, sub_y, coeff, guess)", pieces) + + from leap.implicit import replace_AssignImplicit + code = replace_AssignImplicit(code, {"solve": solver_hook}) + + # Make some dummy thing to use to troubleshoot...user array + # of size one to start with... + codegen = f.CodeGenerator("BDFMethod", + user_type_map={"y": f.ArrayType((2,), + f.BuiltinType("real (kind=8)"),)}, + function_registry=freg, + module_preamble=""" + ! lines copied to the start of the module, e.g. to say: + ! use ModStuff + """) + + code_str = codegen(code) + text_file = open("test.f90", "w") + text_file.write(code_str) + text_file.close() + + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) diff --git a/test/test_imex.py b/test/test_imex.py index 5aa25d7..f6ea99c 100755 --- a/test/test_imex.py +++ b/test/test_imex.py @@ -57,58 +57,21 @@ def solver_hook(solve_expr, solve_var, solver_id, guess): "y", use_high_order=False), 3], [KapsProblem(epsilon=0.9), KennedyCarpenterIMEXARK4MethodBuilder("y"), 4], ]) -def test_convergence(python_method_impl, problem, method, expected_order): +def test_convergence(python_method_impl, problem, method, expected_order, + show_dag=False, plot_solution=False): pytest.importorskip("scipy") - - code = method.generate() - - from leap.implicit import replace_AssignImplicit - code = replace_AssignImplicit(code, {"solve": solver_hook}) - - from pytools.convergence import EOCRecorder - eocrec = EOCRecorder() - - for n in range(2, 7): - dt = 2**(-n) - - y_0 = problem.initial() - t_start = problem.t_start - t_end = problem.t_end - - from functools import partial - interp = python_method_impl(code, function_map={ - "expl_y": problem.nonstiff, - "impl_y": problem.stiff, - "solver": partial(solver, problem.stiff), - }) - - interp.set_up(t_start=t_start, dt_start=dt, context={"y": y_0}) - - times = [] - values = [] - - for event in interp.run(t_end=t_end): - if isinstance(event, interp.StateComputed): - values.append(event.state_component) - times.append(event.t) - - times = np.array(times) - values = np.array(values) - - assert abs(times[-1] - t_end) < 1e-10 - - times = np.array(times) - - error = np.linalg.norm(values[-1] - problem.exact(t_end)) - eocrec.add_data_point(dt, error) - - print("------------------------------------------------------") - print("expected order %d" % expected_order) - print("------------------------------------------------------") - print(eocrec.pretty_print()) - - orderest = eocrec.estimate_order_of_convergence()[0, 1] - assert orderest > 0.9 * expected_order + from functools import partial + dts = 2**-np.array(range(3, 8), dtype=np.float64) + function_map = {"expl_y": problem.nonstiff, + "impl_y": problem.stiff, + "solver": partial(solver, problem.stiff)} + from utils import check_simple_convergence + check_simple_convergence(method=method, method_impl=python_method_impl, + expected_order=expected_order, show_dag=show_dag, + plot_solution=plot_solution, + function_map=function_map, dts=dts, + implicit=True, problem=problem, + solver_hook=solver_hook) @pytest.mark.parametrize("problem, method", [ diff --git a/test/test_rk.py b/test/test_rk.py index 6786476..e90489f 100755 --- a/test/test_rk.py +++ b/test/test_rk.py @@ -37,16 +37,11 @@ SSPRK22MethodBuilder, SSPRK33MethodBuilder, ) from leap.rk.imex import KennedyCarpenterIMEXARK4MethodBuilder -import numpy as np - -import logging from utils import ( # noqa python_method_impl_interpreter as pmi_int, python_method_impl_codegen as pmi_cg) -logger = logging.getLogger(__name__) - # {{{ non-adaptive test @@ -90,77 +85,11 @@ def test_rk_accuracy(python_method_impl, method, expected_order, ]) def test_adaptive_timestep(python_method_impl, method, show_dag=False, plot=False): - # Use "DEBUG" to trace execution - logging.basicConfig(level=logging.INFO) - - component_id = method.component_id - code = method.generate() - print(code) - #1/0 - - if show_dag: - from dagrt.language import show_dependency_graph - show_dependency_graph(code) - - from stiff_test_systems import VanDerPolProblem - example = VanDerPolProblem() - y = example.initial() - - interp = python_method_impl(code, - function_map={"" + component_id: example}) - interp.set_up(t_start=example.t_start, dt_start=1e-5, context={component_id: y}) - - times = [] - values = [] - - new_times = [] - new_values = [] - - last_t = 0 - step_sizes = [] - - for event in interp.run(t_end=example.t_end): - if isinstance(event, interp.StateComputed): - assert event.component_id == component_id - - new_values.append(event.state_component) - new_times.append(event.t) - elif isinstance(event, interp.StepCompleted): - if not new_times: - continue - - step_sizes.append(event.t - last_t) - last_t = event.t - - times.extend(new_times) - values.extend(new_values) - del new_times[:] - del new_values[:] - elif isinstance(event, interp.StepFailed): - del new_times[:] - del new_values[:] - - logger.info("failed step at t=%s" % event.t) - - times = np.array(times) - values = np.array(values) - step_sizes = np.array(step_sizes) - - if plot: - import matplotlib.pyplot as pt - pt.plot(times, values[:, 1], "x-") - pt.show() - pt.plot(times, step_sizes, "x-") - pt.show() - - step_sizes = np.array(step_sizes) - small_step_frac = len(np.nonzero(step_sizes < 0.01)[0]) / len(step_sizes) - big_step_frac = len(np.nonzero(step_sizes > 0.05)[0]) / len(step_sizes) - - print("small_step_frac (<0.01): %g - big_step_frac (>.05): %g" - % (small_step_frac, big_step_frac)) - assert small_step_frac <= 0.35, small_step_frac - assert big_step_frac >= 0.16, big_step_frac + from utils import check_adaptive_timestep + check_adaptive_timestep(python_method_impl=python_method_impl, method=method, + ss_frac=0.35, bs_frac=0.16, ss_targ=0.01, + bs_targ=0.05, show_dag=show_dag, plot=plot, + implicit=False) # }}} diff --git a/test/utils.py b/test/utils.py index 60d3ac3..bf9dce0 100644 --- a/test/utils.py +++ b/test/utils.py @@ -21,6 +21,9 @@ """ import numpy as np +import logging + +logger = logging.getLogger(__name__) # {{{ things to pass for python_method_impl @@ -82,10 +85,10 @@ def initial(self): def exact(self, t): inner = np.sqrt(3) / 2 * np.log(t) - return np.sqrt(t) * ( + return np.array([np.sqrt(t) * ( 5 * np.sqrt(3) / 3 * np.sin(inner) - + np.cos(inner) - ) + + np.cos(inner)), (3*np.cos(inner) + + np.sqrt(3) / 3 * np.sin(inner))/np.sqrt(t)], dtype=np.float64) def __call__(self, t, y): u, v = y @@ -97,13 +100,14 @@ def __call__(self, t, y): def check_simple_convergence(method, method_impl, expected_order, problem=None, dts=_default_dts, - show_dag=False, plot_solution=False): + show_dag=False, plot_solution=False, + function_map=None, implicit=False, + solver_hook=None): if problem is None: problem = DefaultProblem() component_id = method.component_id code = method.generate() - print(code) if show_dag: from dagrt.language import show_dependency_graph @@ -112,14 +116,19 @@ def check_simple_convergence(method, method_impl, expected_order, from pytools.convergence import EOCRecorder eocrec = EOCRecorder() + if function_map is None: + function_map = {"" + component_id: problem} + + if implicit: + from leap.implicit import replace_AssignImplicit + code = replace_AssignImplicit(code, {"solve": solver_hook}) + for dt in dts: t = problem.t_start y = problem.initial() final_t = problem.t_end - interp = method_impl(code, function_map={ - "" + component_id: problem, - }) + interp = method_impl(code, function_map=function_map) interp.set_up(t_start=t, dt_start=dt, context={component_id: y}) times = [] @@ -127,12 +136,13 @@ def check_simple_convergence(method, method_impl, expected_order, for event in interp.run(t_end=final_t): if isinstance(event, interp.StateComputed): assert event.component_id == component_id - values.append(event.state_component[0]) + values.append(event.state_component) times.append(event.t) assert abs(times[-1] - final_t) / final_t < 0.1 times = np.array(times) + values = np.array(values) if plot_solution: import matplotlib.pyplot as pt @@ -140,7 +150,7 @@ def check_simple_convergence(method, method_impl, expected_order, pt.plot(times, problem.exact(times), label="true") pt.show() - error = abs(values[-1] - problem.exact(final_t)) + error = np.linalg.norm(values[-1] - problem.exact(final_t)) eocrec.add_data_point(dt, error) print("------------------------------------------------------") @@ -153,4 +163,89 @@ def check_simple_convergence(method, method_impl, expected_order, assert orderest > expected_order * 0.9 +def check_adaptive_timestep(python_method_impl, method, ss_frac, bs_frac, + ss_targ, bs_targ, show_dag=False, plot=False, + implicit=False, solver=None, + solver_hook=None): + # Use "DEBUG" to trace execution + logging.basicConfig(level=logging.INFO) + + component_id = method.component_id + code = method.generate() + + if implicit: + from leap.implicit import replace_AssignImplicit + code = replace_AssignImplicit(code, {"solve": solver_hook}) + + if show_dag: + from dagrt.language import show_dependency_graph + show_dependency_graph(code) + + from stiff_test_systems import VanDerPolProblem + example = VanDerPolProblem() + y = example.initial() + + if implicit: + interp = python_method_impl(code, + function_map={"" + component_id: example, + "solver": solver}) + else: + interp = python_method_impl(code, + function_map={"" + component_id: example}) + interp.set_up(t_start=example.t_start, dt_start=1e-5, context={component_id: y}) + + times = [] + values = [] + + new_times = [] + new_values = [] + + last_t = 0 + step_sizes = [] + + for event in interp.run(t_end=example.t_end): + if isinstance(event, interp.StateComputed): + assert event.component_id == component_id + + new_values.append(event.state_component) + new_times.append(event.t) + elif isinstance(event, interp.StepCompleted): + if not new_times: + continue + + step_sizes.append(event.t - last_t) + last_t = event.t + + times.extend(new_times) + values.extend(new_values) + del new_times[:] + del new_values[:] + elif isinstance(event, interp.StepFailed): + del new_times[:] + del new_values[:] + + logger.info("failed step at t=%s" % event.t) + + times = np.array(times) + values = np.array(values) + step_sizes = np.array(step_sizes) + + if plot: + import matplotlib.pyplot as pt + pt.clf() + pt.plot(times, values[:, 1], "x-") + pt.show() + pt.plot(times, step_sizes, "x-") + pt.show() + + step_sizes = np.array(step_sizes) + small_step_frac = len(np.nonzero(step_sizes < ss_targ)[0]) / len(step_sizes) + big_step_frac = len(np.nonzero(step_sizes > bs_targ)[0]) / len(step_sizes) + + print("small_step_frac (<%g): %g - big_step_frac (>%g): %g" + % (ss_targ, small_step_frac, bs_targ, big_step_frac)) + assert small_step_frac <= ss_frac, small_step_frac + assert big_step_frac >= bs_frac, big_step_frac + + # vim: foldmethod=marker