examples/adaptive-rk/adaptive-rk.py

@@ -14,8 +14,10 @@ The same functionality is supported using leap.rk.ODE23MethodBuilder.
 """
 
 import logging
+
 import numpy as np
 
+
 logging.basicConfig(level=logging.INFO)
 
 

examples/imex/imex.py

@@ -84,9 +84,10 @@ def solver_hook(solve_expr, solve_var, solver_id, guess):
 
 
 def run():
-    from leap.rk.imex import KennedyCarpenterIMEXARK4MethodBuilder
     from dagrt.codegen import PythonCodeGenerator
 
+    from leap.rk.imex import KennedyCarpenterIMEXARK4MethodBuilder
+
     # Construct the method generator.
     mgen = KennedyCarpenterIMEXARK4MethodBuilder("y", atol=_atol)
 

examples/implicit_euler/implicit_euler.py

 #!/usr/bin/env python
 """Implicit Euler timestepper"""
 
-from leap import MethodBuilder
-from dagrt.language import DAGCode, CodeBuilder
+from dagrt.language import CodeBuilder, DAGCode
 from pymbolic import var
 from pymbolic.primitives import CallWithKwargs
 
+from leap import MethodBuilder
+
+
 __copyright__ = "Copyright (C) 2014 Matt Wala"
 
 __license__ = """

examples/implicit_euler/test_implicit_euler.py

 #! /usr/bin/env python
 
-import pytest
 import sys
+
 import numpy as np
+import pytest
 
 
 __copyright__ = "Copyright (C) 2014 Andreas Kloeckner, Matt Wala"

examples/rk/run_rk_method.py

@@ -26,9 +26,10 @@ THE SOFTWARE.
 """
 
 
-from leap.rk import ODE23MethodBuilder
 import numpy as np
 
+from leap.rk import ODE23MethodBuilder
+
 
 def main(show_dag=False, plot_solution=False):
     component_id = "y"

examples/semi-implicit-rk/semi-implicit-rk.py

@@ -11,9 +11,10 @@ Source:
 """
 
 
-import numpy as np
 import logging
 
+import numpy as np
+
 
 logging.basicConfig(level=logging.INFO)
 

examples/stability/mr-max-eigvals.py

 #!/usr/bin/env python
 
 import matplotlib
+
+
 matplotlib.use("Agg")  # noqa
 
+from functools import partial
+
+import matplotlib.pyplot as pt
 import numpy as np
 import numpy.linalg as la
+
 from leap.multistep.multirate import TwoRateAdamsBashforthMethodBuilder
-import matplotlib.pyplot as pt
-from functools import partial
 
 
 def process_eigval(evaluate_mat, speed_factor, prec, major_eigval):
@@ -26,10 +30,10 @@ def process_eigval(evaluate_mat, speed_factor, prec, major_eigval):
 
 
 def main():
-    from leap.step_matrix import StepMatrixFinder
-
     from pymbolic import var
 
+    from leap.step_matrix import StepMatrixFinder
+
     speed_factor = 10
     method_name = "Fq"
     order = 3

examples/stability/mr-stability-diagram.py

 #!/usr/bin/env python
 
 import matplotlib
+
+
 matplotlib.use("Agg")  # noqa
 
+import matplotlib.pyplot as pt
 import numpy as np
 import numpy.linalg as la
+
 from leap.multistep.multirate import TwoRateAdamsBashforthMethodBuilder
-import matplotlib.pyplot as pt
 
 
 def main():
-    from leap.step_matrix import StepMatrixFinder
-
     from pymbolic import var
 
+    from leap.step_matrix import StepMatrixFinder
+
     speed_factor = 10
     method_name = "Fq"
     order = 3
@@ -65,9 +68,10 @@ def main():
 
             return (np.abs(eigvals) <= 1 + tol).all()
 
-        from leap.stability import find_truth_bdry
         from functools import partial
 
+        from leap.stability import find_truth_bdry
+
         points = []
 
         for angle in angles:

examples/stability/multirate-step-matrix.py

 #!/usr/bin/env python
 import numpy as np
 import numpy.linalg as la
+
 from leap.multistep.multirate import TwoRateAdamsBashforthMethodBuilder
 
 
 def main():
-    from leap.step_matrix import StepMatrixFinder
-
     from pymbolic import var
 
+    from leap.step_matrix import StepMatrixFinder
+
     speed_factor = 10
     step_ratio = 7
     method_name = "Fq"
@@ -58,9 +59,10 @@ def main():
 
         return (np.abs(eigvals) <= 1 + tol).all()
 
-    from leap.stability import find_truth_bdry
     from functools import partial
 
+    from leap.stability import find_truth_bdry
+
     prec = 1e-5
     print("stable imaginary timestep:",
             find_truth_bdry(partial(is_stable, 1j), prec=prec))

examples/stability/plot-stability-regions.py

 #!/usr/bin/env python
 
 import matplotlib
+
+
 matplotlib.use("Agg")
 
 
@@ -24,8 +26,8 @@ def main(save_pdfs=True):
     pt.ylabel(r"Im $\lambda$ / RHS calls")
     pt.grid()
 
-    import leap.rk as rk
     import leap.multistep as multistep
+    import leap.rk as rk
 
     for label, method, factor in [
             #("ode23", rk.ODE23MethodBuilder("y", use_high_order=True), 1),

examples/stability/step-matrix.py

 #!/usr/bin/env python
 import numpy as np
 import numpy.linalg as la
-from leap.rk import RK4MethodBuilder  # noqa
+
 from leap.multistep import AdamsBashforthMethodBuilder  # noqa
+from leap.rk import RK4MethodBuilder  # noqa
 
 
 def main():
-    from leap.step_matrix import StepMatrixFinder
-
     from pymbolic import var
 
+    from leap.step_matrix import StepMatrixFinder
+
     #method = RK4MethodBuilder("y")
     method = AdamsBashforthMethodBuilder("y", order=3, static_dt=True)
 
@@ -42,9 +43,10 @@ def main():
 
         return (np.abs(eigvals) <= 1 + tol).all()
 
-    from leap.stability import find_truth_bdry
     from functools import partial
 
+    from leap.stability import find_truth_bdry
+
     prec = 1e-5
     print("stable imaginary timestep:",
             find_truth_bdry(partial(is_stable, 1j), prec=prec))

examples/variable-coeff-wave-equation/wave-equation.py

@@ -12,13 +12,12 @@ with piecewise constant coefficients c(x) using a multirate multistep method.
 import argparse
 import fnmatch
 import logging
+import os
+from contextlib import contextmanager
+
 import matplotlib
 import numpy as np
 import numpy.linalg as la
-import os
-
-
-from contextlib import contextmanager
 
 
 logging.basicConfig(level=logging.INFO)
@@ -261,7 +260,7 @@ def make_3_component_multirate_method(
 
     """
     from leap.multistep.multirate import (
-            MultiRateMultiStepMethodBuilder, MultiRateHistory)
+        MultiRateHistory, MultiRateMultiStepMethodBuilder)
 
     ncomponents = 3
     assert problem.ncomponents == ncomponents

leap/__init__.py

@@ -113,9 +113,9 @@ class TwoOrderAdaptiveMethodBuilderMixin(MethodBuilder):
         raise NotImplementedError()
 
     def finish_adaptive(self, cb, high_order_estimate, low_order_estimate):
+        from dagrt.expression import IfThenElse
         from pymbolic import var
         from pymbolic.primitives import Comparison, LogicalOr, Max, Min
-        from dagrt.expression import IfThenElse
 
         norm_start_state = var("norm_start_state")
         norm_end_state = var("norm_end_state")

leap/multistep/__init__.py

@@ -29,9 +29,11 @@ THE SOFTWARE.
 
 import numpy as np
 import numpy.linalg as la
-from leap import MethodBuilder
+
 from pymbolic import var
 
+from leap import MethodBuilder
+
 
 __doc__ = """
 .. autoclass:: AdamsIntegrationFunctionFamily
@@ -43,8 +45,8 @@ __doc__ = """
 # {{{ Adams-Bashforth integration (with and without dynamic time steps)
 
 def _linear_comb(coefficients, vectors):
-    from operator import add
     from functools import reduce
+    from operator import add
     return reduce(add,
             (coeff * v for coeff, v in zip(coefficients, vectors)))
 
@@ -277,7 +279,7 @@ class AdamsBashforthMethodBuilder(MethodBuilder):
         from pytools import UniqueNameGenerator
         name_gen = UniqueNameGenerator()
 
-        from dagrt.language import DAGCode, CodeBuilder
+        from dagrt.language import CodeBuilder, DAGCode
 
         array = var("<builtin>array")
         rhs_var = var("rhs_var")

leap/multistep/multirate/__init__.py

@@ -27,17 +27,19 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
 
-from pytools import Record
-from leap import MethodBuilder
+from dataclasses import dataclass, replace
+from typing import Any, Optional, Tuple
+
 from pymbolic import var
 
+from leap import MethodBuilder
 from leap.multistep import _linear_comb
 
 
 __doc__ = """
-
 .. autoclass:: rhs_policy
 .. autoclass:: MultiRateHistory
+    :members:
 .. autoclass:: MultiRateMultiStepMethodBuilder
 
 .. autoclass:: SchemeExplainerBase
@@ -47,7 +49,7 @@ __doc__ = """
 
 # {{{ system description
 
-class rhs_policy:  # noqa
+class rhs_policy:  # noqa: N801
     """
     .. attribute:: late
     .. attribute:: early
@@ -58,44 +60,47 @@ class rhs_policy:  # noqa
     early_and_late = 2
 
 
-class MultiRateHistory(Record):
+@dataclass(frozen=True)
+class MultiRateHistory:
+    interval: int
+    """An integer indicating the interval (relative to the smallest available
+    time step) at which the history is to be update, where an update will typically
+    involve a call to :attr:`func_name`.
     """
-    .. automethod:: __init__
+    func_name: str
+    """The name of the function to call on history updates."""
+    arguments: Tuple[str, ...]
+    """A tuple of component names which are passed to :attr:`func_name`."""
+    order: Optional[int] = None
+    """The approximation order of the Adams method to be used for this history.
+    If *None*, the default method is assumed.
+    """
+    rhs_policy: int = rhs_policy.late
+    """One of the constants in :class:`rhs_policy`."""
+    invalidate_computed_state: bool = False
+    """A flag that denotes whether evaluating the right-hand side :attr:`func_name`
+    should force a recomputation of any state that depended upon now-superseded
+    state.
+    """
+    hist_length: Optional[int] = None
+    """The length of the history. If the length is greater than :attr:`order`,
+    we use a least-squares solve rather than a linear sole to obtain the
+    Adams coefficients for this history.
     """
-    def __init__(self, interval, func_name, arguments, order=None,
-            rhs_policy=rhs_policy.late, invalidate_computed_state=False,
-            hist_length=None):
-        """
-        :arg interval: An integer indicating the interval (relative to the
-            smallest available timestep) at which this history is to be
-            updated.
-            (where each update will typically involve a call to *func_name*)
-        :arg arguments: A tuple of component names
-            (see :class:`MultiRateMultiStepMethodBuilder`)
-            which are passed to this right-hand side function.
-        :arg order: The Adams approximation order to be used for this
-            history, or None if the method default is to be used.
-        :arg rhs_policy: One of the constants in :class:`rhs_policy`
-        :arg invalidate_dependent_state: Whether evaluating this
-            right-hand side should force a recomputation of any
-            state that depended upon now-superseded state.
-        :arg hist_length: history length.  If greater than order, we use a
-            least-squares solve rather than a linear solve to obtain the Adams
-            coefficients for this history
-        """
-        super().__init__(
-                interval=interval,
-                func_name=func_name,
-                arguments=arguments,
-                order=order,
-                rhs_policy=rhs_policy,
-                invalidate_computed_state=invalidate_computed_state,
-                hist_length=hist_length)
 
     @property
     def history_length(self):
         return self.hist_length
 
+    def copy(self, **kwargs: Any) -> "MultiRateHistory":
+        from warnings import warn
+
+        warn(f"{type(self).__name__}.copy is deprecated and will be removed in "
+             "2025. This is now a dataclass and should use the replace function",
+             DeprecationWarning, stacklevel=2)
+
+        return replace(self, **kwargs)
+
 
 class RHS(MultiRateHistory):
     def __init__(self, *args, **kwargs):
@@ -323,8 +328,9 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder):
                 if hist_length is None:
                     hist_length = order
 
-                new_component_rhss.append(rhs.copy(order=order,
-                    hist_length=hist_length))
+                new_component_rhss.append(
+                    replace(rhs, order=order, hist_length=hist_length)
+                )
 
             new_rhss.append(tuple(new_component_rhss))
 
@@ -711,8 +717,13 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder):
 
     # }}}
 
-    class StateContribExplanation(Record):
-        pass
+    @dataclass(frozen=True)
+    class StateContribExplanation:
+        rhs: str
+        """Name of the right-hand side function to call."""
+        from_substeps: int
+        using: Tuple[str, ...]
+        """A tuple of variables used."""
 
     # {{{ main method generation
 
@@ -820,9 +831,8 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder):
                     dt_factor = self.dt
 
                 from leap.multistep import (
-                        AdamsMonomialIntegrationFunctionFamily,
-                        emit_adams_integration,
-                        emit_adams_extrapolation)
+                    AdamsMonomialIntegrationFunctionFamily, emit_adams_extrapolation,
+                    emit_adams_integration)
 
                 if self.is_ode_component[comp_name]:
                     contrib = dt_factor*emit_adams_integration(
@@ -1113,7 +1123,7 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder):
         if explainer is None:
             explainer = SchemeExplainerBase()
 
-        from dagrt.language import DAGCode, CodeBuilder
+        from dagrt.language import CodeBuilder, DAGCode
 
         with CodeBuilder(name="initialization") as cb_init:
             self.emit_initialization(cb_init)

leap/rk/__init__.py

@@ -27,11 +27,12 @@ THE SOFTWARE.
 """
 
 import numpy as np
-from leap import MethodBuilder, TwoOrderAdaptiveMethodBuilderMixin
-from dagrt.language import CodeBuilder, DAGCode
 
+from dagrt.language import CodeBuilder, DAGCode
 from pymbolic import var
 
+from leap import MethodBuilder, TwoOrderAdaptiveMethodBuilderMixin
+
 
 __doc__ = """
 .. data:: ORDER_TO_RK_METHOD_BUILDER

leap/rk/imex.py

@@ -2,6 +2,7 @@
 
 
 from pymbolic import var
+
 from leap import TwoOrderAdaptiveMethodBuilderMixin
 from leap.rk import ButcherTableauMethodBuilder
 

leap/stability.py

@@ -23,9 +23,10 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
 
-import numpy as np
 from cmath import pi
 
+import numpy as np
+
 
 __doc__ = """
 .. autofunction:: find_stability_region

leap/step_matrix.py

@@ -21,23 +21,33 @@ THE SOFTWARE.
 """
 
 from collections import namedtuple
-from dagrt.expression import EvaluationMapper
+from dataclasses import dataclass, replace
+from typing import List, Tuple
+
 import numpy as np
+
 from dagrt.exec_numpy import FailStepException
+from dagrt.expression import EvaluationMapper
 from pymbolic.interop.maxima import MaximaStringifyMapper
-from pytools import Record
+from pymbolic.primitives import Expression
 
-__doc__ = """
 
+__doc__ = """
 .. autoclass:: StepMatrixFinder
 .. autofunction:: fast_evaluator
 """
 
 
-class SparseStepMatrix(Record):
-
-    def __init__(self, shape, indices, data):
-        Record.__init__(self, shape=shape, indices=indices, data=data)
+@dataclass(frozen=True)
+class SparseStepMatrix:
+    shape: Tuple[int, int]
+    """The shape of the step matrix."""
+    indices: List[Tuple[int, int]]
+    """A list of ``(i, j)`` tuples given the index in the matrix for each element
+    in :attr:`data`.
+    """
+    data: List[Expression]
+    """A list of matrix entries."""
 
 
 class LeapMaximaStringifyMapper(MaximaStringifyMapper):
@@ -197,8 +207,8 @@ class StepMatrixFinder:
 
         components, initial_vals = self.run_symbolic_step(phase_name, shapes)
 
-        from pymbolic.mapper.differentiator import DifferentiationMapper
         from pymbolic.mapper.dependency import DependencyMapper
+        from pymbolic.mapper.differentiator import DifferentiationMapper
         dependencies = DependencyMapper()
 
         nv = len(components)
@@ -306,15 +316,16 @@ def fast_evaluator(matrix, sparse=False):
 
     substitutor = SubstitutionMapper(make_identifier)
 
-    from pymbolic import compile
     # functools.partial ensures the resulting object is picklable.
     from functools import partial
 
+    from pymbolic import compile
+
     if sparse:
         data = [substitutor(entry) for entry in matrix.data]
         var_order, renamed_vars = get_var_order_from_name_map()
         compiled_entries = [compile(entry, renamed_vars) for entry in data]
-        compiled_matrix = matrix.copy(data=compiled_entries)
+        compiled_matrix = replace(matrix, data=compiled_entries)
     else:
         matrix = substitutor(matrix)
         var_order, renamed_vars = get_var_order_from_name_map()
@@ -333,7 +344,7 @@ def _eval_compiled_matrix(compiled_matrix, var_order, var_assignments):
     arguments = [var_assignments[name] for name in var_order]
     if isinstance(compiled_matrix, SparseStepMatrix):
         evaluated_data = [entry(*arguments) for entry in compiled_matrix.data]
-        return compiled_matrix.copy(data=evaluated_data)
+        return replace(compiled_matrix, data=evaluated_data)
     else:
         return compiled_matrix(*arguments)
 

leap/transform.py

@@ -30,7 +30,7 @@ __doc__ = """
 
 
 def _elide_yield_state(statements):
-    from dagrt.language import YieldState, Nop
+    from dagrt.language import Nop, YieldState
     return [stmt
             if not isinstance(stmt, YieldState)
             else Nop(id=stmt.id, depends_on=stmt.depends_on)
@@ -40,7 +40,7 @@ def _elide_yield_state(statements):
 def _update_t_by_dt_factor(factor, statements):
     from dagrt.language import Assign, Nop
     from pymbolic import var
-    from pymbolic.mapper.substitutor import make_subst_func, SubstitutionMapper
+    from pymbolic.mapper.substitutor import SubstitutionMapper, make_subst_func
 
     mapper = SubstitutionMapper(
         make_subst_func({"<dt>": factor * var("<dt>")}))
@@ -69,7 +69,7 @@ def strang_splitting(dag1, dag2, stepping_phase):
     :returns: a :class:`dagrt.language.DAGCode`
     """
 
-    from pymbolic.mapper.substitutor import make_subst_func, SubstitutionMapper
+    from pymbolic.mapper.substitutor import SubstitutionMapper, make_subst_func
 
     # {{{ disambiguate