Fix Constants

irksome/discontinuous_galerkin_stepper.py

@@ -9,9 +9,9 @@
 from .bcs import stage2spaces4bc
 from .manipulation import extract_terms, strip_dt_form
 from .stage import getBits
-from .tools import MeshConstant, getNullspace, replace
+from .tools import getNullspace, replace
 import numpy as np
-from firedrake import TestFunction, Function, NonlinearVariationalProblem as NLVP, NonlinearVariationalSolver as NLVS
+from firedrake import as_vector, Constant, dot, TestFunction, Function, NonlinearVariationalProblem as NLVP, NonlinearVariationalSolver as NLVS
 from firedrake.dmhooks import pop_parent, push_parent
 
 
@@ -55,8 +55,7 @@ def getFormDiscGalerkin(F, L, Q, t, dt, u0, bcs=None, nullspace=None):
     V = v.function_space()
     assert V == u0.function_space()
 
-    MC = MeshConstant(V.mesh())
-    vecconst = np.vectorize(MC.Constant)
+    vecconst = Constant
 
     num_fields = len(V)
     num_stages = L.space_dimension()
@@ -84,7 +83,9 @@ def getFormDiscGalerkin(F, L, Q, t, dt, u0, bcs=None, nullspace=None):
 
     # mass matrix later for BC
     mmat = np.multiply(basis_vals, qwts) @ basis_vals.T
-    mmat_inv = vecconst(np.linalg.inv(mmat))
+
+    # L2 projector
+    proj = Constant(np.linalg.solve(mmat, np.multiply(basis_vals, qwts)))
 
     u0bits, vbits, VVbits, UUbits = getBits(num_stages, num_fields,
                                             u0, UU, v, VV)
@@ -155,7 +156,6 @@ def getFormDiscGalerkin(F, L, Q, t, dt, u0, bcs=None, nullspace=None):
             Fnew += dt * qwts[q] * basis_vals[i, q] * replace(F_i, repl)
 
     # Oh, honey, is it the boundary conditions?
-    minv_test_vals = mmat_inv @ np.multiply(basis_vals, qwts)
     if bcs is None:
         bcs = []
     bcsnew = []
@@ -165,7 +165,7 @@ def getFormDiscGalerkin(F, L, Q, t, dt, u0, bcs=None, nullspace=None):
         for q in range(len(qpts)):
             tcur = t + qpts[q] * dt
             bcblah_at_qp[q] = replace(bcarg, {t: tcur})
-        bc_func_for_stages = minv_test_vals @ bcblah_at_qp
+        bc_func_for_stages = dot(proj, as_vector(bcblah_at_qp))
         for i in range(num_stages):
             Vbigi = stage2spaces4bc(bc, V, Vbig, i)
             bcsnew.append(bc.reconstruct(V=Vbigi, g=bc_func_for_stages[i]))

irksome/galerkin_stepper.py

@@ -8,9 +8,9 @@
 from .bcs import bc2space, stage2spaces4bc
 from .deriv import TimeDerivative
 from .stage import getBits
-from .tools import MeshConstant, getNullspace, replace
+from .tools import getNullspace, replace
 import numpy as np
-from firedrake import TestFunction, Function, NonlinearVariationalProblem as NLVP, NonlinearVariationalSolver as NLVS
+from firedrake import as_vector, dot, Constant, TestFunction, Function, NonlinearVariationalProblem as NLVP, NonlinearVariationalSolver as NLVS
 from firedrake.dmhooks import pop_parent, push_parent
 
 
@@ -57,9 +57,6 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, bcs=None, nullspace=None):
     V = v.function_space()
     assert V == u0.function_space()
 
-    MC = MeshConstant(V.mesh())
-    vecconst = np.vectorize(MC.Constant)
-
     num_fields = len(V)
     num_stages = L_test.space_dimension()
 
@@ -83,18 +80,20 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, bcs=None, nullspace=None):
 
     # mass-ish matrix later for BC
     mmat = np.multiply(test_vals, qwts) @ trial_vals[1:].T
-    mmat_inv = vecconst(np.linalg.inv(mmat))
+
+    # L2 projector
+    proj = Constant(np.linalg.solve(mmat, np.multiply(test_vals, qwts)))
 
     u0bits, vbits, VVbits, UUbits = getBits(num_stages, num_fields,
                                             u0, UU, v, VV)
 
     Fnew = Zero()
 
-    trial_vals = vecconst(trial_vals)
-    trial_dvals = vecconst(trial_dvals)
-    test_vals = vecconst(test_vals)
-    qpts = vecconst(qpts.reshape((-1,)))
-    qwts = vecconst(qwts)
+    trial_vals = Constant(trial_vals)
+    trial_dvals = Constant(trial_dvals)
+    test_vals = Constant(test_vals)
+    qpts = Constant(qpts.reshape((-1,)))
+    qwts = Constant(qwts)
 
     for i in range(num_stages):
         repl = {}
@@ -128,7 +127,6 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, bcs=None, nullspace=None):
             Fnew += dt * qwts[q] * test_vals[i, q] * replace(F_i, repl)
 
     # Oh, honey, is it the boundary conditions?
-    minv_test_vals = mmat_inv @ np.multiply(test_vals, qwts)
     if bcs is None:
         bcs = []
     bcsnew = []
@@ -139,7 +137,7 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, bcs=None, nullspace=None):
         for q in range(len(qpts)):
             tcur = t + qpts[q] * dt
             bcblah_at_qp[q] = replace(bcarg, {t: tcur}) - u0_sub * trial_vals[0, q]
-        bc_func_for_stages = minv_test_vals @ bcblah_at_qp
+        bc_func_for_stages = dot(proj, as_vector(bcblah_at_qp))
         for i in range(num_stages):
             Vbigi = stage2spaces4bc(bc, V, Vbig, i)
             bcsnew.append(bc.reconstruct(V=Vbigi, g=bc_func_for_stages[i]))

irksome/imex.py

@@ -1,13 +1,13 @@
 import FIAT
 import numpy as np
-from firedrake import (Function, NonlinearVariationalProblem,
+from firedrake import (Constant, Function, NonlinearVariationalProblem,
                        NonlinearVariationalSolver, TestFunction)
 from firedrake.dmhooks import pop_parent, push_parent
 from ufl.classes import Zero
 
 from .ButcherTableaux import RadauIIA
 from .stage import getBits, getFormStage
-from .tools import AI, IA, MeshConstant, replace
+from .tools import AI, IA, replace
 
 
 def riia_explicit_coeffs(k):
@@ -38,18 +38,15 @@ def getFormExplicit(Fexp, butch, u0, UU, t, dt, splitting=None):
     which really just differ by which constants are in them."""
     v = Fexp.arguments()[0]
     V = v.function_space()
-    msh = V.mesh()
     Vbig = UU.function_space()
     VV = TestFunction(Vbig)
 
     num_stages = butch.num_stages
     num_fields = len(V)
-    MC = MeshConstant(msh)
-    vc = np.vectorize(MC.Constant)
     Aexp = riia_explicit_coeffs(num_stages)
-    Aprop = vc(Aexp)
-    Ait = vc(butch.A)
-    C = vc(butch.c)
+    Aprop = Constant(Aexp)
+    Ait = Constant(butch.A)
+    C = Constant(butch.c)
 
     u0bits, vbits, VVbits, UUbits = getBits(num_stages, num_fields,
                                             u0, UU, v, VV)
@@ -99,7 +96,7 @@ def getFormExplicit(Fexp, butch, u0, UU, t, dt, splitting=None):
             Fit += dt * replace(Fexp, repl)
 
         # dense contribution to propagator
-        Ablah = vc(np.linalg.solve(butch.A, Aexp))
+        Ablah = Constant(np.linalg.solve(butch.A, Aexp))
 
         for i in range(num_stages):
             # replace test function

irksome/stage.py

@@ -5,7 +5,8 @@
 
 import numpy as np
 from FIAT import Bernstein, ufc_simplex
-from firedrake import (Function, NonlinearVariationalProblem,
+from firedrake import (Constant,
+                       Function, NonlinearVariationalProblem,
                        NonlinearVariationalSolver, TestFunction, dx,
                        inner, split)
 from firedrake.petsc import PETSc
@@ -24,20 +25,13 @@ def isiterable(x):
 
 
 def split_field(num_fields, u):
-    return np.array((u,) if num_fields == 1 else split(u), dtype="O")
+    ubits = np.array(split(u), dtype="O")
+    return ubits
 
 
 def split_stage_field(num_stages, num_fields, UU):
-    if num_fields == 1:
-        if num_stages == 1:   # single-stage method
-            UUbits = np.reshape(np.array((UU,), dtype='O'), (num_stages, num_fields))
-        else:  # multi-stage methods
-            UUbits = np.zeros((len(split(UU)),), dtype='O')
-            for (i, x) in enumerate(split(UU)):
-                UUbits[i] = np.zeros((1,), dtype='O')
-                UUbits[i][0] = x
-    else:
-        UUbits = np.reshape(np.asarray(split(UU), dtype="O"), (num_stages, num_fields))
+    UUbits = np.reshape(np.asarray(split(UU), dtype="O"),
+                        (num_stages, num_fields))
     return UUbits
 
 
@@ -128,7 +122,7 @@ def getFormStage(F, butch, u0, t, dt, bcs=None, splitting=None, vandermonde=None
                                             u0, ZZ, v, VV)
 
     MC = MeshConstant(V.mesh())
-    vecconst = np.vectorize(MC.Constant)
+    vecconst = Constant
 
     C = vecconst(butch.c)
     A = vecconst(butch.A)
@@ -145,13 +139,7 @@ def getFormStage(F, butch, u0, t, dt, bcs=None, splitting=None, vandermonde=None
     Vander_col = Vander[1:, 0]
     Vander0 = Vander[1:, 1:]
 
-    v0u0 = np.zeros((num_stages, num_fields), dtype="O")
-    for i in range(num_stages):
-        for j in range(num_fields):
-            v0u0[i, j] = Vander_col[i] * u0bits[j]
-
-    if num_fields == 1:
-        v0u0 = np.reshape(v0u0, (-1,))
+    v0u0 = np.outer(Vander_col, u0bits)
 
     UUbits = v0u0 + Vander0 @ ZZbits
 
@@ -173,9 +161,7 @@ def getFormStage(F, butch, u0, t, dt, bcs=None, splitting=None, vandermonde=None
             for j in range(num_fields):
                 repl[u0bits[j]] = UUbits[i][j] - u0bits[j]
                 repl[vbits[j]] = VVbits[i][j]
-
-            # Also get replacements right for indexing.
-            for j in range(num_fields):
+                # Also get replacements right for indexing.
                 for ii in np.ndindex(u0bits[j].ufl_shape):
                     repl[u0bits[j][ii]] = UUbits[i][j][ii] - u0bits[j][ii]
                     repl[vbits[j][ii]] = VVbits[i][j][ii]