Clean up UFL manipulation (#130)

* Clean up UFL manipulation

* unfiy replace

Author: pbrubeck

irksome/deriv.py

@@ -5,9 +5,6 @@
 from ufl.algorithms.map_integrands import map_integrand_dags, map_expr_dag
 from ufl.algorithms.apply_derivatives import GenericDerivativeRuleset
 from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering
-from ufl.tensors import ListTensor
-from ufl.indexed import Indexed
-from ufl.core.multiindex import FixedIndex
 
 
 @ufl_type(num_ops=1,
@@ -21,9 +18,6 @@ class TimeDerivative(Derivative):
     __slots__ = ()
 
     def __new__(cls, f):
-        if isinstance(f, ListTensor):
-            # Push TimeDerivative inside ListTensor
-            return ListTensor(*map(TimeDerivative, f.ufl_operands))
         return Derivative.__new__(cls)
 
     def __init__(self, f):
@@ -32,14 +26,6 @@ def __init__(self, f):
     def __str__(self):
         return "d{%s}/dt" % (self.ufl_operands[0],)
 
-    def _simplify_indexed(self, multiindex):
-        """Return a simplified Expr used in the constructor of Indexed(self, multiindex)."""
-        # Push Indexed inside TimeDerivative
-        if all(isinstance(i, FixedIndex) for i in multiindex):
-            f, = self.ufl_operands
-            return TimeDerivative(Indexed(f, multiindex))
-        return Derivative._simplify_indexed(self, multiindex)
-
 
 def Dt(f, order=1):
     """Short-hand function to produce a :class:`TimeDerivative` of a given order."""
@@ -72,9 +58,10 @@ def time_derivative(self, o, f):
         else:
             return map_expr_dag(self, f)
 
-    def _linear_op(self, o, f):
-        return o._ufl_expr_reconstruct_(f)
+    def _linear_op(self, o, *operands):
+        return o._ufl_expr_reconstruct_(*operands)
 
+    indexed = _linear_op
     derivative = _linear_op
     grad = _linear_op
     curl = _linear_op

irksome/dirk_stepper.py

@@ -5,7 +5,7 @@
 from ufl.constantvalue import as_ufl
 
 from .deriv import TimeDerivative, expand_time_derivatives
-from .tools import component_replace, replace, MeshConstant, vecconst
+from .tools import replace, MeshConstant, vecconst
 from .bcs import bc2space
 
 
@@ -36,7 +36,7 @@ def getFormDIRK(F, ks, butch, t, dt, u0, bcs=None):
     repl = {t: t + c * dt,
             u0: g + k * (a * dt),
             TimeDerivative(u0): k}
-    stage_F = component_replace(F, repl)
+    stage_F = replace(F, repl)
 
     bcnew = []
 

irksome/discontinuous_galerkin_stepper.py

@@ -7,7 +7,7 @@
 from .bcs import stage2spaces4bc
 from .deriv import expand_time_derivatives
 from .manipulation import extract_terms, strip_dt_form
-from .tools import component_replace, replace, vecconst
+from .tools import replace, vecconst
 import numpy as np
 from firedrake import TestFunction
 
@@ -97,21 +97,21 @@ def getFormDiscGalerkin(F, L, Q, t, dt, u0, stages, bcs=None):
     # Jump terms
     repl = {u0: u_np[0] - u0,
             v: v_np[0]}
-    Fnew = component_replace(F_dtless, repl)
+    Fnew = replace(F_dtless, repl)
 
     # Terms with time derivatives
     for q in range(len(qpts)):
         repl = {t: t + qpts[q] * dt,
                 v: vsub[q] * dt,
                 u0: dtu0sub[q] / dt}
-        Fnew += component_replace(F_dtless, repl)
+        Fnew += replace(F_dtless, repl)
 
     # Handle the rest of the terms
     for q in range(len(qpts)):
         repl = {t: t + qpts[q] * dt,
                 v: vsub[q] * dt,
                 u0: usub[q]}
-        Fnew += component_replace(F_remainder, repl)
+        Fnew += replace(F_remainder, repl)
 
     # Oh, honey, is it the boundary conditions?
     if bcs is None:

irksome/galerkin_stepper.py

@@ -6,7 +6,7 @@
 from .base_time_stepper import StageCoupledTimeStepper
 from .bcs import bc2space, stage2spaces4bc
 from .deriv import TimeDerivative, expand_time_derivatives
-from .tools import component_replace, replace, vecconst
+from .tools import replace, vecconst
 import numpy as np
 from firedrake import TestFunction
 
@@ -91,7 +91,7 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, bcs=None):
                 v: vsub[q] * dt,
                 u0: usub[q],
                 dtu0: dtu0sub[q] / dt}
-        Fnew += component_replace(F, repl)
+        Fnew += replace(F, repl)
 
     # Oh, honey, is it the boundary conditions?
     if bcs is None:

irksome/imex.py

@@ -9,7 +9,7 @@
 from .ButcherTableaux import RadauIIA
 from .deriv import TimeDerivative, expand_time_derivatives
 from .stage_value import getFormStage
-from .tools import AI, ConstantOrZero, IA, MeshConstant, replace, component_replace, getNullspace, get_stage_space
+from .tools import AI, ConstantOrZero, IA, MeshConstant, replace, getNullspace, get_stage_space
 from .bcs import bc2space
 
 
@@ -66,15 +66,15 @@ def getFormExplicit(Fexp, butch, u0, UU, t, dt, splitting=None):
         for i in range(num_stages):
             # replace test function
             repl = {v: v_np[i]}
-            Ftmp = component_replace(Fexp, repl)
+            Ftmp = replace(Fexp, repl)
 
             # replace the solution with stage values
             for j in range(num_stages):
                 repl = {t: t + C[j] * dt,
                         u0: u_np[j]}
 
                 # and sum the contribution
-                replF = component_replace(Ftmp, repl)
+                replF = replace(Ftmp, repl)
                 Fit += Ait[i, j] * dt * replF
                 Fprop += Aprop[i, j] * dt * replF
     elif splitting == IA:
@@ -84,23 +84,23 @@ def getFormExplicit(Fexp, butch, u0, UU, t, dt, splitting=None):
                     u0: u_np[i],
                     v: v_np[i]}
 
-            Fit += dt * component_replace(Fexp, repl)
+            Fit += dt * replace(Fexp, repl)
 
         # dense contribution to propagator
         AinvAexp = vecconst(np.linalg.solve(butch.A, Aexp))
 
         for i in range(num_stages):
             # replace test function
             repl = {v: v_np[i]}
-            Ftmp = component_replace(Fexp, repl)
+            Ftmp = replace(Fexp, repl)
 
             # replace the solution with stage values
             for j in range(num_stages):
                 repl = {t: t + C[j] * dt,
                         u0: u_np[j]}
 
                 # and sum the contribution
-                Fprop += AinvAexp[i, j] * dt * component_replace(Ftmp, repl)
+                Fprop += AinvAexp[i, j] * dt * replace(Ftmp, repl)
     else:
         raise NotImplementedError(
             "Must specify splitting to either IA or AI")
@@ -329,13 +329,13 @@ def getFormsDIRKIMEX(F, Fexp, ks, khats, butch, t, dt, u0, bcs=None):
     repl = {t: t + c * dt,
             u0: g + dt * a * k,
             TimeDerivative(u0): k}
-    stage_F = component_replace(F, repl)
+    stage_F = replace(F, repl)
 
     # Explicit replacement, solve at time t + chat * dt, for khat
     replhat = {t: t + chat * dt,
                u0: ghat}
 
-    Fhat = inner(khat, vhat)*dx + component_replace(Fexp, replhat)
+    Fhat = inner(khat, vhat)*dx + replace(Fexp, replhat)
 
     bcnew = []
 

irksome/nystrom_stepper.py

@@ -1,7 +1,7 @@
 from .base_time_stepper import StageCoupledTimeStepper
 from .bcs import BCStageData, bc2space
 from .deriv import Dt, TimeDerivative, expand_time_derivatives
-from .tools import component_replace, replace, vecconst
+from .tools import replace, vecconst
 from firedrake import TestFunction, as_ufl
 import numpy
 from ufl import zero
@@ -106,7 +106,7 @@ def getFormNystrom(F, tableau, t, dt, u0, ut0, stages,
                 u0: u0 + ut0 * (c[i] * dt) + Abark[i] * dt**2,
                 dtu: ut0 + Ak[i] * dt,
                 dt2u: k_np[i]}
-        Fnew += component_replace(F, repl)
+        Fnew += replace(F, repl)
 
     if bcs is None:
         bcs = []

irksome/stage_derivative.py

@@ -5,7 +5,7 @@
 from firedrake import assemble, dx, inner, norm
 
 from ufl.constantvalue import as_ufl, zero
-from .tools import component_replace, replace, AI, vecconst
+from .tools import AI, replace, vecconst
 from .deriv import Dt, TimeDerivative, expand_time_derivatives
 from .bcs import EmbeddedBCData, BCStageData, bc2space
 from .manipulation import extract_terms
@@ -84,7 +84,7 @@ def getForm(F, butch, t, dt, u0, stages, bcs=None, bc_type=None, splitting=AI):
                 v: v_np[i],
                 u0: u0 + A1w[i] * dt,
                 dtu: A2invw[i]}
-        Fnew += component_replace(F, repl)
+        Fnew += replace(F, repl)
 
     if bcs is None:
         bcs = []

irksome/stage_value.py

@@ -11,7 +11,7 @@
 from .ButcherTableaux import CollocationButcherTableau
 from .deriv import expand_time_derivatives
 from .manipulation import extract_terms, strip_dt_form
-from .tools import AI, is_ode, replace, component_replace, vecconst
+from .tools import AI, is_ode, replace, vecconst
 from .base_time_stepper import StageCoupledTimeStepper
 
 
@@ -117,15 +117,15 @@ def getFormStage(F, butch, t, dt, u0, stages, bcs=None, splitting=None, vandermo
         repl = {t: t + c[i] * dt,
                 v: A2invTv[i],
                 u0: w_np[i] - u0}
-        Fnew += component_replace(F_dtless, repl)
+        Fnew += replace(F_dtless, repl)
 
     # Handle the rest of the terms
     for i in range(num_stages):
         # replace the solution with stage values
         repl = {t: t + c[i] * dt,
                 v: A1Tv[i] * dt,
                 u0: w_np[i]}
-        Fnew += component_replace(F_remainder, repl)
+        Fnew += replace(F_remainder, repl)
 
     if bcs is None:
         bcs = []
@@ -212,7 +212,7 @@ def get_update_solver(self, update_solver_parameters):
         for i in range(self.num_stages):
             repl = {t: t + C[i] * dt,
                     u0: u_np[i]}
-            Fupdate += dt * B[i] * component_replace(split_form.remainder, repl)
+            Fupdate += dt * B[i] * replace(split_form.remainder, repl)
 
         # And the BC's for the update -- just the original BC at t+dt
         update_bcs = []

irksome/tools.py

@@ -2,12 +2,9 @@
 from functools import reduce
 import numpy
 from firedrake import Function, FunctionSpace, MixedVectorSpaceBasis, Constant
-from ufl.algorithms.analysis import extract_type, has_exact_type
-from ufl.algorithms.map_integrands import map_integrand_dags
-from ufl.classes import CoefficientDerivative, Zero
-from ufl.constantvalue import as_ufl
-from ufl.corealg.multifunction import MultiFunction
-from ufl.tensors import as_tensor
+from ufl.algorithms.analysis import extract_type
+from ufl import as_tensor, zero
+from ufl import replace as ufl_replace
 
 from irksome.deriv import TimeDerivative
 
@@ -54,58 +51,10 @@ def getNullspace(V, Vbig, num_stages, nullspace):
     return nspnew
 
 
-# Update for UFL's replace that performs post-order traversal and hence replaces
-# more complicated expressions first.
-class MyReplacer(MultiFunction):
-    def __init__(self, mapping):
-        super().__init__()
-        self.replacements = mapping
-        if not all(k.ufl_shape == v.ufl_shape for k, v in mapping.items()):
-            raise ValueError("Replacement expressions must have the same shape as what they replace.")
-
-    def expr(self, o):
-        if o in self.replacements:
-            return self.replacements[o]
-        else:
-            return self.reuse_if_untouched(o, *map(self, o.ufl_operands))
-
-
 def replace(e, mapping):
-    """Replace subexpressions in expression.
-
-    @param e:
-        An Expr or Form.
-    @param mapping:
-        A dict with from:to replacements to perform.
-    """
-    mapping2 = dict((k, as_ufl(v)) for (k, v) in mapping.items())
-
-    # Workaround for problem with delayed derivative evaluation
-    # The problem is that J = derivative(f(g, h), g) does not evaluate immediately
-    # So if we subsequently do replace(J, {g: h}) we end up with an expression:
-    # derivative(f(h, h), h)
-    # rather than what were were probably thinking of:
-    # replace(derivative(f(g, h), g), {g: h})
-    #
-    # To fix this would require one to expand derivatives early (which
-    # is not attractive), or make replace lazy too.
-    if has_exact_type(e, CoefficientDerivative):
-        # Hack to avoid circular dependencies
-        from ufl.algorithms.ad import expand_derivatives
-        e = expand_derivatives(e)
-
-    return map_integrand_dags(MyReplacer(mapping2), e)
-
-
-def component_replace(e, mapping):
-    """Replace, recurring on components"""
-    cmapping = {}
-    for key, value in mapping.items():
-        cmapping[key] = as_tensor(value)
-        if key.ufl_shape:
-            for j in numpy.ndindex(key.ufl_shape):
-                cmapping[key[j]] = value[j]
-    return replace(e, cmapping)
+    """A wrapper for ufl.replace that allows numpy arrays."""
+    cmapping = {k: as_tensor(v) for k, v in mapping.items()}
+    return ufl_replace(e, cmapping)
 
 
 # Utility functions that help us refactor
@@ -120,11 +69,13 @@ def IA(A):
 def is_ode(f, u):
     """Given a form defined over a function `u`, checks if
     (each bit of) u appears under a time derivative."""
-    blah = extract_type(f, TimeDerivative)
-
-    Dtbits = set(b.ufl_operands[0] for b in blah)
-    ubits = set(u[i] for i in numpy.ndindex(u.ufl_shape))
-    return Dtbits == ubits
+    derivs = extract_type(f, TimeDerivative)
+    Dtbits = []
+    for k in derivs:
+        op, = k.ufl_operands
+        Dtbits.extend(op[i] for i in numpy.ndindex(op.ufl_shape))
+    ubits = [u[i] for i in numpy.ndindex(u.ufl_shape)]
+    return set(Dtbits) == set(ubits)
 
 
 # Utility class for constants on a mesh
@@ -139,7 +90,7 @@ def Constant(self, val=0.0):
 
 def ConstantOrZero(x, MC=None):
     const = MC.Constant if MC else Constant
-    return Zero() if abs(complex(x)) < 1.e-10 else const(x)
+    return zero() if abs(complex(x)) < 1.e-10 else const(x)
 
 
 vecconst = numpy.vectorize(ConstantOrZero)