diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py new file mode 100644 index 00000000..843c80e1 --- /dev/null +++ b/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)) diff --git a/irksome/__init__.py b/irksome/__init__.py index af848619..fcdca469 100644 --- a/irksome/__init__.py +++ b/irksome/__init__.py @@ -14,6 +14,7 @@ AdamsMoulton, ) from .tableaux.pep_explicit_rk import PEPRK +from .ufl.time_projector import TimeProjector from .ufl.deriv import Dt, expand_time_derivatives, check_irksome_import_order check_irksome_import_order() @@ -54,6 +55,7 @@ "RadauIIA", "SSPButcherTableau", "SSPK_DIRK_IMEX", + "TimeProjector", "WSODIRK", ] diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index d2b67f97..9fcb1321 100644 --- a/irksome/galerkin_stepper.py +++ b/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 diff --git a/irksome/ufl/estimate_degrees.py b/irksome/ufl/estimate_degrees.py index 2e7f98c7..b6dcd421 100644 --- a/irksome/ufl/estimate_degrees.py +++ b/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) diff --git a/irksome/ufl/time_projector.py b/irksome/ufl/time_projector.py new file mode 100644 index 00000000..b33bf7fe --- /dev/null +++ b/irksome/ufl/time_projector.py @@ -0,0 +1,94 @@ + +from functools import singledispatchmethod + +from ufl.core.operator import Operator +from ufl.core.ufl_type import ufl_type +from ufl.corealg.dag_traverser import DAGTraverser +from ufl.domain import as_domain +from ufl.algorithms.map_integrands import map_integrands +from ufl.algorithms import expand_derivatives, replace +from ufl import as_tensor, diff, dx, inner, outer +from ufl.classes import Expr, BaseForm +from ufl import Coefficient + +import numpy as np + + +@ufl_type( + num_ops=1, + inherit_shape_from_operand=0, + inherit_indices_from_operand=0, +) +class TimeProjector(Operator): + __slots__ = ("order", "quadrature") + + def __init__(self, expression, order, Q): + self.order = order + self.quadrature = Q + Operator.__init__(self, operands=(expression,)) + + def _ufl_expr_reconstruct_(self, *operands): + """Return a new object of the same type with new operands.""" + return TimeProjector(*operands, self.order, self.quadrature) + + +class TimeProjectorDispatcher(DAGTraverser): + def __init__(self, element, t, dt, u0, u1, stages, phi): + super().__init__() + self.L_trial = element + self.t = t + self.dt = dt + self.u0 = u0 + self.u1 = u1 + self.stages = stages + self.phi = np.reshape(phi, (-1,)) + + # Work around singledispatchmethod inheritance issue; + # see https://bugs.python.org/issue36457. + @singledispatchmethod + def process(self, o): + return super().process(o) + + @process.register(Expr) + @process.register(BaseForm) + def expr(self, o): + return self.reuse_if_untouched(o) + + @process.register(TimeProjector) + def time_projector(self, o): + from FIAT import Legendre + from firedrake import TestFunction, TensorFunctionSpace + from irksome.galerkin_stepper import getTermGalerkin + from irksome.constant import vecconst + # use the internal copy of the state, so it does not get updated again in the outer quadrature + order = o.order + Q = o.quadrature + assert order+1 <= len(self.phi) + f, = o.ufl_operands + mesh = as_domain(self.u0.function_space().mesh()) + R = TensorFunctionSpace(mesh, "DG", 0, shape=f.ufl_shape) + F = inner(f, TestFunction(R))*dx + F = replace(F, {self.u0: self.u1}) + + # compute the hierarchical mass matrix (always the identity) + ref_el = self.L_trial.get_reference_element() + L_test = Legendre(ref_el, order) + test_vals = L_test.tabulate(0, Q.get_points())[(0,)] + M = np.multiply(test_vals, Q.get_weights()) @ test_vals.T + Minv = vecconst(np.linalg.inv(M)) + + # compute modal expansion tested against c + c = Coefficient(R) + + test = outer(as_tensor(Minv[:order+1] @ self.phi) / self.dt, c) + Fc = getTermGalerkin(F, self.L_trial, L_test, Q, self.t, self.dt, self.u1, self.stages, test) + + # compute the L2-Riesz representation by undoing the integral against the test coefficient + fc = sum(it.integrand() for it in Fc.integrals()) + fproj = expand_derivatives(diff(fc, c)) + return fproj + + +def expand_time_projectors(expression, element, t, dt, u0, u1, stages, phi): + rules = TimeProjectorDispatcher(element, t, dt, u0, u1, stages, phi) + return map_integrands(rules, expression) diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index 65672006..45a159a3 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -1,8 +1,10 @@ import numpy as np import pytest from firedrake import * -from irksome import Dt, MeshConstant, ContinuousPetrovGalerkinScheme, GalerkinCollocationScheme, TimeStepper, GaussLegendre +from irksome import Dt, MeshConstant, ContinuousPetrovGalerkinScheme, GalerkinCollocationScheme, TimeStepper, TimeProjector, GaussLegendre from irksome.labeling import TimeQuadratureLabel +from irksome.scheme import create_time_quadrature +import numpy as np def run_1d_heat_dirichletbc(scheme, **kwargs): @@ -263,7 +265,7 @@ def test_wave_eq_galerkin(deg, N, order): assert np.allclose(energy[1:], energy[:-1]) -def kepler(V, order, t, dt, u0, solver_parameters): +def kepler_naive(V, order, t, dt, u0, solver_parameters): dim = V.value_size//2 u = Function(V) u.interpolate(u0) @@ -336,8 +338,46 @@ def kepler_aux_variable(V, order, t, dt, u0, solver_parameters): return stepper, invariants +def kepler_projector(V, order, t, dt, u0, solver_parameters): + dim = V.value_size//2 + u = Function(V) + u.interpolate(u0) + + uv = variable(u) + p = as_vector([uv[k] for k in range(dim)]) + q = as_vector([uv[k] for k in range(dim, 2*dim)]) + + T = 0.5*dot(p, p) + U = -1/sqrt(dot(q, q)) + + # Invariants + H = T + U + L = dot(p, perp(q)) + A1, A2 = U*q - L*perp(p) + invariants = [H*dx, L*dx, A1*dx, A2*dx] + + v = TestFunction(V) + dHdu = diff(H, uv) + dA1du = diff(A1, uv) + dA2du = diff(A2, uv) + + Qproj = create_time_quadrature(25) + w0 = TimeProjector(dHdu, order-1, Qproj) + w1 = TimeProjector(dA1du, order-1, Qproj) + w2 = TimeProjector(dA2du, order-1, Qproj) + determinant_forms = [v, w0, w1, w2] + tensor = as_tensor(determinant_forms) + + F = inner(Dt(u), v)*dx - (det(tensor) / (2*L*H))*dx + + scheme = ContinuousPetrovGalerkinScheme(order, quadrature_degree="auto") + stepper = TimeStepper(F, scheme, t, dt, u, + solver_parameters=solver_parameters) + return stepper, invariants + + @pytest.mark.parametrize('order', (1, 2)) -@pytest.mark.parametrize('problem', (kepler, kepler_aux_variable)) +@pytest.mark.parametrize('problem', (kepler_naive, kepler_aux_variable, kepler_projector)) def test_kepler(problem, order): msh = UnitIntervalMesh(1) MC = MeshConstant(msh)