TimeProjector

demos/structure_preservation/nse_helicity_energy.py

@@ -0,0 +1,221 @@
+from firedrake import *
+from irksome import Dt, GalerkinTimeStepper
+from irksome.galerkin_stepper import TimeProjector
+from irksome.labeling import TimeQuadratureLabel
+from FIAT import ufc_simplex
+from FIAT.quadrature_schemes import create_quadrature
+from scipy import special
+from firedrake.petsc import PETSc
+
+print = PETSc.Sys.Print
+
+ufcline = ufc_simplex(1)
+
+'''
+Thanks to Boris Andrews for providing the Hill vortex functions
+'''
+'''
+Hill vortex functions
+'''
+# Bessel function parameters
+bessel_J_root = 5.7634591968945506
+bessel_J_root_threehalves = bessel_J(3/2, bessel_J_root)
+
+
+# (r, theta, phi) components of Hill vortex
+def hill_r(r, theta, radius):
+    rho = r / radius
+    return 2 * (
+        bessel_J(3/2, bessel_J_root*rho) / rho**(3/2)
+        - bessel_J_root_threehalves
+    ) * cos(theta)
+
+
+def hill_theta(r, theta, radius):
+    rho = r / radius
+    return (
+        bessel_J_root * bessel_J(5/2, bessel_J_root*rho) / rho**(1/2)
+        + 2 * bessel_J_root_threehalves
+        - 2 * bessel_J(3/2, bessel_J_root*rho) / rho**(3/2)
+    ) * sin(theta)
+
+
+def hill_phi(r, theta, radius):
+    rho = r / radius
+    return bessel_J_root * (
+        bessel_J(3/2, bessel_J_root*rho) / rho**(3/2)
+        - bessel_J_root_threehalves
+    ) * rho * sin(theta)
+
+
+# Hill vortex (Cartesian)
+def hill(vec, radius):
+    (x, y, z) = vec
+    rho = sqrt(x*x + y*y)
+
+    r = sqrt(dot(vec, vec))
+    theta = pi/2-atan2(z, rho)
+    phi = atan2(y, x)
+
+    r_dir = as_vector([cos(phi)*sin(theta), sin(phi)*sin(theta), cos(theta)])
+    theta_dir = as_vector([cos(phi)*cos(theta), sin(phi)*cos(theta), -sin(theta)])
+    phi_dir = as_vector([-sin(phi), cos(phi), 0])
+
+    return conditional(  # If we're outside the vortex...
+        ge(r, radius),
+        zero((3,)),
+        conditional(  # If we're at the origin...
+            le(r, 1e-13),
+            as_vector([0, 0, 2*((bessel_J_root/2)**(3/2)/special.gamma(5/2) - bessel_J_root_threehalves)]),
+            (hill_r(r, theta, radius) * r_dir
+             + hill_theta(r, theta, radius) * theta_dir
+             + hill_phi(r, theta, radius) * phi_dir)
+        )
+    )
+
+
+def stokes_pair(msh):
+    V = VectorFunctionSpace(msh, "CG", 2)
+    Q = FunctionSpace(msh, "CG", 1)
+    return V, Q
+
+
+def nse_naive(msh, order, t, dt, Re, solver_parameters=None):
+    hill_expr = hill(SpatialCoordinate(msh), 0.25)
+    V, Q = stokes_pair(msh)
+    Z = V * Q
+    up = Function(Z)
+    up.subfunctions[0].interpolate(hill_expr)
+    # VTKFile("output/hill.pvd").write(up.subfunctions[0])
+
+    v, q = TestFunctions(Z)
+    u, p = split(up)
+
+    Qhigh = create_quadrature(ufcline, 3*order-1)
+    Qlow = create_quadrature(ufcline, 2*(order-1))
+    Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights())
+    Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights())
+
+    F = (Llow(inner(Dt(u), v) * dx)
+         + Lhigh(-inner(cross(u, curl(u)), v) * dx)
+         + 1/Re * inner(grad(u), grad(v)) * dx
+         - inner(p, div(v)) * dx
+         - inner(div(u), q) * dx)
+
+    bcs = DirichletBC(Z.sub(0), 0, "on_boundary")
+
+    stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs)
+    Q1 = inner(u, u) * dx
+    Q2 = inner(u, curl(u)) * dx
+    invariants = [Q1, Q2]
+    return stepper, invariants
+
+
+def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None):
+    hill_expr = hill(SpatialCoordinate(msh), 0.25)
+    V, Q = stokes_pair(msh)
+    Z = V * V * V * Q * Q
+    up = Function(Z)
+    up.subfunctions[0].interpolate(hill_expr)
+
+    aux_indices = (1, 2, 4)
+    v, v1, v2, q, q1 = TestFunctions(Z)
+    u, w1, w2, p, r1 = split(up)
+
+    Qhigh = create_quadrature(ufcline, 3*(order-1))
+    Qlow = create_quadrature(ufcline, 2*(order-1))
+    Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights())
+    Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights())
+
+    F = (Llow(inner(Dt(u), v) * dx)
+         + Lhigh(-inner(cross(w1, w2), v) * dx)
+         + Llow(1/Re * inner(grad(w1), grad(v)) * dx)
+         - inner(p, div(v)) * dx
+         - inner(div(u), q) * dx
+         + inner(w1 - u, v1) * dx
+         + inner(w2 - curl(u), v2) * dx
+         - inner(r1, div(v2)) * dx
+         - inner(div(w2), q1) * dx
+         )
+
+    bcs = [DirichletBC(Z.sub(i), 0, "on_boundary") for i in range(3)]
+
+    stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs, aux_indices=aux_indices)
+    Q1 = inner(u, u) * dx
+    Q2 = inner(u, curl(u)) * dx
+    invariants = [Q1, Q2]
+    return stepper, invariants
+
+
+def nse_project(msh, order, t, dt, Re, solver_parameters=None):
+    hill_expr = hill(SpatialCoordinate(msh), 0.25)
+    V, Q = stokes_pair(msh)
+    Z = V * V * Q * Q
+    up = Function(Z)
+    up.subfunctions[0].interpolate(hill_expr)
+
+    aux_indices = (1, 3)
+    v, v2, q, q1 = TestFunctions(Z)
+    u, w2, p, r1 = split(up)
+
+    Qhigh = create_quadrature(ufcline, 3*(order-1))
+    Qlow = create_quadrature(ufcline, 2*(order-1))
+    Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights())
+    Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights())
+
+    # Eliminate w1 only
+    Qproj = create_quadrature(ufcline, 2*order-1)
+    w1 = TimeProjector(u, order-1, Qproj)
+
+    F = (Llow(inner(Dt(u), v) * dx)
+         + Lhigh(-inner(cross(w1, w2), v) * dx)
+         + Llow(1/Re * inner(grad(w1), grad(v)) * dx)
+         - inner(p, div(v)) * dx
+         - inner(div(u), q) * dx
+         + inner(w2 - curl(u), v2)*dx
+         - inner(r1, div(v2)) * dx
+         - inner(div(w2), q1) * dx
+         )
+
+    bcs = [DirichletBC(Z.sub(i), 0, "on_boundary") for i in range(2)]
+
+    stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs, aux_indices=aux_indices)
+    Q1 = inner(u, u) * dx
+    Q2 = inner(u, curl(u)) * dx
+    invariants = [Q1, Q2]
+    return stepper, invariants
+
+
+def run_nse(stepper, invariants):
+    t = stepper.t
+    dt = stepper.dt
+    row = [float(t), *map(assemble, invariants)]
+    print(*(f"{r:.8e}" for r in row))
+    # while float(t) < 3 * 2**(-6):
+    for k in range(2):
+        stepper.advance()
+        t.assign(t + dt)
+
+        row = [float(t), *map(assemble, invariants)]
+        print(*(f"{r:.8e}" for r in row))
+
+
+order = 2
+N = 8
+msh = UnitCubeMesh(N, N, N)
+msh.coordinates.dat.data[:, :] -= 0.5
+
+t = Constant(0)
+dt = Constant(2**-10)
+Re = Constant(2**16)
+
+solvers = {
+    "naive": nse_naive,
+    "project": nse_project,
+    "aux": nse_aux_variable,
+}
+
+for name, solver in solvers.items():
+    print(name)
+    t.assign(0)
+    run_nse(*solver(msh, order, t, dt, Re))

irksome/__init__.py

@@ -10,6 +10,7 @@
 )
 from .tableaux.pep_explicit_rk import PEPRK
 from .ufl.deriv import Dt, expand_time_derivatives
+from .ufl.time_projector import TimeProjector
 from .tableaux.dirk_imex_tableaux import DIRK_IMEX
 from .tableaux.ars_dirk_imex_tableaux import ARS_DIRK_IMEX
 from .tableaux.sspk_tableau import SSPK_DIRK_IMEX, SSPButcherTableau
@@ -42,6 +43,7 @@
     "RadauIIA",
     "SSPButcherTableau",
     "SSPK_DIRK_IMEX",
+    "TimeProjector",
     "WSODIRK",
 ]
 

irksome/galerkin_stepper.py

@@ -1,12 +1,14 @@
 from FIAT import (Bernstein, DiscontinuousLagrange,
                   GaussRadau, IntegratedLegendre, Lagrange,
-                  NodalEnrichedElement, RestrictedElement)
+                  Legendre, NodalEnrichedElement, RestrictedElement)
 from ufl.classes import Zero
-from ufl import as_ufl, as_tensor
+from ufl import as_ufl, as_tensor, Coefficient
+from ufl.domain import as_domain
 
 from .base_time_stepper import StageCoupledTimeStepper
 from .bcs import bc2space, extract_bcs, stage2spaces4bc
 from .ufl.deriv import TimeDerivative, expand_time_derivatives
+from .ufl.time_projector import expand_time_projectors
 from .ufl.estimate_degrees import TimeDegreeEstimator, get_degree_mapping
 from .labeling import split_quadrature, as_form
 from .scheme import GalerkinCollocationScheme, create_time_quadrature, ufc_line
@@ -20,7 +22,7 @@
 from .stage_value import getFormStage
 
 import numpy as np
-from firedrake import TestFunction, Constant
+from firedrake import TestFunction, Constant, Function, VectorFunctionSpace
 from firedrake.bcs import EquationBCSplit
 
 
@@ -65,17 +67,34 @@ def getElements(basis_type, order):
     return L_trial, L_test
 
 
-def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices):
+def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices=None):
+    qpts = Q.get_points()
+    qwts = Q.get_weights()
+
+    # internal state to be used inside projected expressions
+    u1 = Function(u0)
+    # symbolic Coefficient with the temporal test function
+    mesh = as_domain(u0.function_space().mesh())
+    R = VectorFunctionSpace(mesh, "Real", 0, dim=L_test.space_dimension())
+    phi = Coefficient(R)
+    # apply time projectors
+    F = expand_time_projectors(F, L_trial, t, dt, u0, u1, stages, phi)
+    # tabulate the temporal test function
+    ref_el = L_test.get_reference_element()
+    phisub = vecconst(Legendre(ref_el, L_test.degree()).tabulate(0, qpts)[(0,)].T)
+
     # preprocess time derivatives
     F = expand_time_derivatives(F, t=t, timedep_coeffs=(u0,))
     v, = F.arguments()
     V = v.function_space()
-    assert V == u0.function_space()
+
+    # assert V == u0.function_space()
     i0, = L_trial.entity_dofs()[0][0]
 
     qpts = Q.get_points()
     qwts = Q.get_weights()
     assert qpts.size >= L_test.space_dimension()
+
     tabulate_trials = L_trial.tabulate(1, qpts)
     trial_vals = tabulate_trials[(0,)]
     trial_dvals = tabulate_trials[(1,)]
@@ -109,7 +128,10 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices)
         repl[q] = {t: t + qpts[q] * dt,
                    v: vsub[q] * dt,
                    u0: usub[q],
-                   dtu0: dtu0sub[q] / dt}
+                   dtu0: dtu0sub[q] / dt,
+                   u1: u0,
+                   phi: phisub[q]}
+
     Fnew = sum(replace(F, repl[q]) for q in repl)
     return Fnew
 

irksome/ufl/estimate_degrees.py

@@ -15,6 +15,7 @@
 )
 
 from .deriv import TimeDerivative
+from .time_projector import TimeProjector
 
 
 class TimeDegreeEstimator(DAGTraverser):
@@ -64,6 +65,10 @@ def terminal(self, o):
     def time_derivative(self, o, degree):
         return max(degree - 1, 0)
 
+    @process.register(TimeProjector)
+    def time_projector(self, o):
+        return o.order
+
     @process.register(Abs)
     @process.register(Conj)
     @process.register(Curl)