From bab0a4fefef608250dabd59f6556297aea7827f2 Mon Sep 17 00:00:00 2001 From: "Robert C. Kirby" Date: Mon, 19 May 2025 13:47:46 +0100 Subject: [PATCH 01/24] Allow labeled forms with different quadrature degrees in galerkin --- irksome/galerkin_stepper.py | 112 ++++++++++++++++++++++++------------ irksome/labeling.py | 19 ++++++ tests/test_galerkin.py | 50 ++++++++++++++++ 3 files changed, 143 insertions(+), 38 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 6e919b66..a26bcb14 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -1,12 +1,14 @@ from FIAT import (Bernstein, DiscontinuousElement, DiscontinuousLagrange, - IntegratedLegendre, Lagrange, Legendre, - make_quadrature, ufc_simplex) + IntegratedLegendre, Lagrange, Legendre, ufc_simplex) +from FIAT.quadrature_schemes import create_quadrature from ufl import zero from ufl.constantvalue import as_ufl from .base_time_stepper import StageCoupledTimeStepper from .bcs import bc2space, stage2spaces4bc from .deriv import TimeDerivative, expand_time_derivatives +from .labeling import as_labelled_form from .tools import replace, vecconst +from .labeling import TimeQuadratureRule import numpy as np from firedrake import TestFunction @@ -34,7 +36,45 @@ def getElements(basis_type, order): return trial_el, test_el -def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, bcs=None): +def getTermGalerkin(Fcur, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test): + print(f"woohoo {len(qpts)}") + v, = Fcur.arguments() + Fcur = expand_time_derivatives(Fcur, t=t, timedep_coeffs=(u0,)) + tabulate_trials = L_trial.tabulate(1, qpts) + trial_vals = tabulate_trials[(0,)] + trial_dvals = tabulate_trials[(1,)] + test_vals = L_test.tabulate(0, qpts)[(0,)] + test_vals_w = np.multiply(test_vals, qwts) + + trial_vals = vecconst(trial_vals) + trial_dvals = vecconst(trial_dvals) + test_vals_w = vecconst(test_vals_w) + qpts = vecconst(np.reshape(qpts, (-1,))) + + # set up the pieces we need to work with to do our substitutions + v_np = np.reshape(test, (-1, *u0.ufl_shape)) + w_np = np.reshape(stages, (-1, *u0.ufl_shape)) + u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) + vsub = test_vals_w.T @ v_np + usub = trial_vals.T @ u_np + dtu0sub = trial_dvals.T @ u_np + + dtu0 = TimeDerivative(u0) + + Fnew = zero() + + # now loop over quadrature points + for q in range(len(qpts)): + repl = {t: t + qpts[q] * dt, + v: vsub[q] * dt, + u0: usub[q], + dtu0: dtu0sub[q] / dt} + Fnew += replace(Fcur, repl) + + return Fnew + + +def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None): """Given a time-dependent variational form, trial and test spaces, and a quadrature rule, produce UFL for the Galerkin-in-Time method. @@ -61,58 +101,54 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, bcs=None): - `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed on the Galerkin-in-time solution, """ - assert L_test.get_reference_element() == Q.ref_el - assert L_trial.get_reference_element() == Q.ref_el - assert Q.ref_el.get_spatial_dimension() == 1 + assert L_test.get_reference_element() == Qdefault.ref_el + assert L_trial.get_reference_element() == Qdefault.ref_el + assert Qdefault.ref_el.get_spatial_dimension() == 1 assert L_trial.get_order() == L_test.get_order() + 1 # preprocess time derivatives - F = expand_time_derivatives(F, t=t, timedep_coeffs=(u0,)) - v, = F.arguments() + F = as_labelled_form(F) + + v, = F.form.arguments() V = v.function_space() assert V == u0.function_space() num_stages = L_test.space_dimension() Vbig = stages.function_space() test = TestFunction(Vbig) - qpts = Q.get_points() - qwts = Q.get_weights() - tabulate_trials = L_trial.tabulate(1, qpts) + Fnew = zero() + for term in F.terms: + print(term) + labels = term.labels + print(labels) + quad_labels = [label for label in labels if isinstance(label, TimeQuadratureRule)] + if len(quad_labels) == 0: + qpts = Qdefault.get_points() + qwts = Qdefault.get_weights() + elif len(quad_labels) == 1: + qpts = quad_labels[0].x + qwts = quad_labels[0].w + else: + raise ValueError("Multiple quadrature labels on one form.") + + Fnew += getTermGalerkin(term.form, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test) + + # mass-ish matrix for BC, based on default quadrature rule + qpts = Qdefault.get_points() + qwts = Qdefault.get_weights() + tabulate_trials = L_trial.tabulate(0, qpts) trial_vals = tabulate_trials[(0,)] - trial_dvals = tabulate_trials[(1,)] test_vals = L_test.tabulate(0, qpts)[(0,)] test_vals_w = np.multiply(test_vals, qwts) - # mass-ish matrix later for BC mmat = test_vals_w @ trial_vals[1:].T - # L2 projector - proj = vecconst(np.linalg.solve(mmat, test_vals_w)) - trial_vals = vecconst(trial_vals) - trial_dvals = vecconst(trial_dvals) + proj = vecconst(np.linalg.solve(mmat, test_vals_w)) test_vals_w = vecconst(test_vals_w) + trial_vals = vecconst(trial_vals) qpts = vecconst(qpts.reshape((-1,))) - # set up the pieces we need to work with to do our substitutions - v_np = np.reshape(test, (num_stages, *u0.ufl_shape)) - w_np = np.reshape(stages, (num_stages, *u0.ufl_shape)) - u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) - vsub = test_vals_w.T @ v_np - usub = trial_vals.T @ u_np - dtu0sub = trial_dvals.T @ u_np - - dtu0 = TimeDerivative(u0) - - # now loop over quadrature points - Fnew = zero() - for q in range(len(qpts)): - repl = {t: t + qpts[q] * dt, - v: vsub[q] * dt, - u0: usub[q], - dtu0: dtu0sub[q] / dt} - Fnew += replace(F, repl) - # Oh, honey, is it the boundary conditions? if bcs is None: bcs = [] @@ -182,9 +218,9 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None, self.trial_el, self.test_el = getElements(basis_type, order) if quadrature is None: - quadrature = make_quadrature(ufc_line, order) + quadrature = create_quadrature(ufc_line, 2 * order) self.quadrature = quadrature - assert np.size(quadrature.get_points()) >= order + assert np.size(quadrature.get_points()) >= order + 1 super().__init__(F, t, dt, u0, order, bcs=bcs, **kwargs) diff --git a/irksome/labeling.py b/irksome/labeling.py index d744c415..ce154b09 100644 --- a/irksome/labeling.py +++ b/irksome/labeling.py @@ -2,6 +2,18 @@ explicit = Label("explicit") +empty_label = Label("") + + +class TimeQuadratureLabel(Label): + def __init__(self, x, w): + super().__init__(TimeQuadratureRule(x, w)) + + +class TimeQuadratureRule: + def __init__(self, x, w): + self.x = tuple(x) + self.w = tuple(w) def split_explicit(F): @@ -15,3 +27,10 @@ def split_explicit(F): map_if_true=keep, map_if_false=drop) return imp_part.form, exp_part.form + + +def as_labelled_form(F): + if not isinstance(F, LabelledForm): + return empty_label(F) + else: + return F diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index ba6d2c6c..266d3cfe 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -4,7 +4,9 @@ from firedrake import * from irksome import Dt, MeshConstant, GalerkinTimeStepper from irksome import TimeStepper, GaussLegendre +from irksome.labeling import TimeQuadratureLabel from FIAT import make_quadrature, ufc_simplex +from FIAT.quadrature_schemes import create_quadrature @pytest.mark.parametrize("order", [1, 2, 3]) @@ -165,6 +167,54 @@ def test_1d_heat_homogeneous_dirichletbc(order): assert (errornorm(u_GL, u) / norm(u)) < 1.e-10 +@pytest.mark.parametrize("order", [1, 2, 3]) +def test_1d_heat_homogeneous_dirichletbc_timequadlabels(order): + N = 20 + msh = UnitIntervalMesh(N) + V = FunctionSpace(msh, "CG", 1) + MC = MeshConstant(msh) + dt = MC.Constant(1.0 / N) + t = MC.Constant(0.0) + (x,) = SpatialCoordinate(msh) + + uexact = sin(pi*x)*exp(-(pi**2)*t) + rhs = Dt(uexact) - div(grad(uexact)) + bcs = DirichletBC(V, uexact, "on_boundary") + u = Function(V) + u.interpolate(uexact) + + v = TestFunction(V) + + ufc_line = ufc_simplex(1) + Qlow = create_quadrature(ufc_line, 2*order-2) + Qhigh = create_quadrature(ufc_line, 2*order+2) + Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) + Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) + + F = ( + Llow(inner(Dt(u), v) * dx) + + inner(grad(u), grad(v)) * dx + - Lhigh(inner(rhs, v) * dx) + ) + + luparams = {"mat_type": "aij", "ksp_type": "preonly", "pc_type": "lu"} + + stepper = GalerkinTimeStepper( + F, order, t, dt, u, bcs=bcs, + solver_parameters=luparams + ) + + t_end = 1.0 + while float(t) < t_end: + print(float(t)) + if float(t) + float(dt) > t_end: + dt.assign(t_end - float(t)) + stepper.advance() + t += dt + + assert errornorm(uexact, u) < 1.e-4 + + def galerkin_wave(n, deg, alpha, order): N = 2**n msh = UnitIntervalMesh(N) From 9bf231960379df6b71a65b0d1ceb588757be8622 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 19 May 2025 18:05:13 +0100 Subject: [PATCH 02/24] aux_indices --- irksome/galerkin_stepper.py | 55 +++++++++------ irksome/labeling.py | 4 +- tests/test_galerkin.py | 129 ++++++++++++++++++++++++++++++++++++ 3 files changed, 165 insertions(+), 23 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index a26bcb14..7e00288e 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -1,7 +1,7 @@ from FIAT import (Bernstein, DiscontinuousElement, DiscontinuousLagrange, IntegratedLegendre, Lagrange, Legendre, ufc_simplex) from FIAT.quadrature_schemes import create_quadrature -from ufl import zero +from ufl import Form from ufl.constantvalue import as_ufl from .base_time_stepper import StageCoupledTimeStepper from .bcs import bc2space, stage2spaces4bc @@ -11,6 +11,7 @@ from .labeling import TimeQuadratureRule import numpy as np from firedrake import TestFunction +from firedrake import assemble, derivative ufc_line = ufc_simplex(1) @@ -36,8 +37,7 @@ def getElements(basis_type, order): return trial_el, test_el -def getTermGalerkin(Fcur, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test): - print(f"woohoo {len(qpts)}") +def getTermGalerkin(Fcur, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test, aux_indices): v, = Fcur.arguments() Fcur = expand_time_derivatives(Fcur, t=t, timedep_coeffs=(u0,)) tabulate_trials = L_trial.tabulate(1, qpts) @@ -54,14 +54,19 @@ def getTermGalerkin(Fcur, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test): # set up the pieces we need to work with to do our substitutions v_np = np.reshape(test, (-1, *u0.ufl_shape)) w_np = np.reshape(stages, (-1, *u0.ufl_shape)) + + # Expand unknowns in the trial space, including the initial condition u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) - vsub = test_vals_w.T @ v_np usub = trial_vals.T @ u_np dtu0sub = trial_dvals.T @ u_np + if aux_indices is not None: + # Expand auxiliary variables in the test space rather than the trial space + usub[:, aux_indices] = test_vals.T @ w_np[:, aux_indices] - dtu0 = TimeDerivative(u0) + vsub = test_vals_w.T @ v_np - Fnew = zero() + dtu0 = TimeDerivative(u0) + Fnew = Form([]) # now loop over quadrature points for q in range(len(qpts)): @@ -74,7 +79,7 @@ def getTermGalerkin(Fcur, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test): return Fnew -def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None): +def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, aux_indices=None): """Given a time-dependent variational form, trial and test spaces, and a quadrature rule, produce UFL for the Galerkin-in-Time method. @@ -91,9 +96,11 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None): :arg u0: a :class:`Function` referring to the state of the PDE system at time `t` :arg stages: a :class:`Function` representing the stages to be solved for. - :arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof) + :kwarg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof) containing (possibly time-dependent) boundary conditions imposed on the system. + :kwarg aux_indices: a list of field indices to be discretized in the test space + rather than trial space. On output, we return a tuple consisting of four parts: @@ -117,7 +124,7 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None): Vbig = stages.function_space() test = TestFunction(Vbig) - Fnew = zero() + Fnew = Form([]) for term in F.terms: print(term) labels = term.labels @@ -127,12 +134,14 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None): qpts = Qdefault.get_points() qwts = Qdefault.get_weights() elif len(quad_labels) == 1: - qpts = quad_labels[0].x - qwts = quad_labels[0].w + qpts = np.asarray(quad_labels[0].x) + qwts = np.asarray(quad_labels[0].w) else: raise ValueError("Multiple quadrature labels on one form.") - Fnew += getTermGalerkin(term.form, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test) + Fnew += getTermGalerkin(term.form, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test, aux_indices) + + assemble(derivative(Fnew, stages), mat_type="aij").petscmat.convert("dense").view() # mass-ish matrix for BC, based on default quadrature rule qpts = Qdefault.get_points() @@ -183,23 +192,25 @@ class GalerkinTimeStepper(StageCoupledTimeStepper): The user may adjust this value between time steps. :arg u0: A :class:`firedrake.Function` containing the current state of the problem to be solved. - :arg bcs: An iterable of :class:`firedrake.DirichletBC` containing + :kwarg bcs: An iterable of :class:`firedrake.DirichletBC` containing the strongly-enforced boundary conditions. Irksome will manipulate these to obtain boundary conditions for each stage of the method. - :arg basis_type: A string indicating the finite element family (either + :kwarg basis_type: A string indicating the finite element family (either `'Lagrange'` or `'Bernstein'`) or the Lagrange variant for the test/trial spaces. Defaults to equispaced Lagrange elements. - :arg quadrature: A :class:`FIAT.QuadratureRule` indicating the quadrature + :kwarg quadrature: A :class:`FIAT.QuadratureRule` indicating the quadrature to be used in time, defaulting to GL with order points - :arg solver_parameters: A :class:`dict` of solver parameters that + :kwarg aux_indices: a list of field indices to be discretized in the test space + rather than trial space. + :kwarg solver_parameters: A :class:`dict` of solver parameters that will be used in solving the algebraic problem associated with each time step. - :arg appctx: An optional :class:`dict` containing application context. + :kwarg appctx: An optional :class:`dict` containing application context. This gets included with particular things that Irksome will pass into the nonlinear solver so that, say, user-defined preconditioners have access to it. - :arg nullspace: A list of tuples of the form (index, VSB) where + :kwarg nullspace: A list of tuples of the form (index, VSB) where index is an index into the function space associated with `u` and VSB is a :class: `firedrake.VectorSpaceBasis` instance to be passed to a @@ -207,10 +218,11 @@ class GalerkinTimeStepper(StageCoupledTimeStepper): associated with the Runge-Kutta method """ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None, - quadrature=None, **kwargs): + quadrature=None, aux_indices=None, **kwargs): assert order >= 1 self.order = order self.basis_type = basis_type + self.aux_indices = aux_indices V = u0.function_space() self.num_fields = len(V) @@ -224,7 +236,7 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None, super().__init__(F, t, dt, u0, order, bcs=bcs, **kwargs) - def get_form_and_bcs(self, stages, basis_type=None, order=None, quadrature=None, F=None): + def get_form_and_bcs(self, stages, basis_type=None, order=None, quadrature=None, aux_indices=None, F=None): if basis_type is None: basis_type = self.basis_type if order is None: @@ -237,7 +249,8 @@ def get_form_and_bcs(self, stages, basis_type=None, order=None, quadrature=None, return getFormGalerkin(F or self.F, trial_el, test_el, quadrature or self.quadrature, - self.t, self.dt, self.u0, stages, self.orig_bcs) + self.t, self.dt, self.u0, stages, self.orig_bcs, + aux_indices or self.aux_indices) def _update(self): k1, = self.trial_el.entity_dofs()[0][1] diff --git a/irksome/labeling.py b/irksome/labeling.py index ce154b09..ee0eccfe 100644 --- a/irksome/labeling.py +++ b/irksome/labeling.py @@ -12,8 +12,8 @@ def __init__(self, x, w): class TimeQuadratureRule: def __init__(self, x, w): - self.x = tuple(x) - self.w = tuple(w) + self.x = x + self.w = w def split_explicit(F): diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index 266d3cfe..b2078326 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -265,3 +265,132 @@ def galerkin_wave(n, deg, alpha, order): def test_wave_eq_galerkin(deg, N, order): energy = galerkin_wave(N, deg, 0.3, order) assert np.allclose(energy[1:], energy[:-1]) + + + + + +@pytest.mark.parametrize('order', (1,)) +def test_kepler(order): + msh = UnitIntervalMesh(1) + MC = MeshConstant(msh) + t = MC.Constant(0.0) + dt = MC.Constant(4/16) + + sd = 2 + nbody = 2 + dim = (nbody-1)*sd + R = FunctionSpace(msh, "R", 0) + V = MixedFunctionSpace([R]*dim) + Z = V * V + z = Function(Z) + + p0 = [0, 1] + q0 = [1, 0] + vals = p0 + q0 + for c, val in zip(z.subfunctions, vals): + c.assign(val) + + c = split(z) + p = as_vector(c[:dim]) + q = as_vector(c[dim:]) + J = as_matrix(np.kron([[0, -1], [1, 0]], np.eye(dim))) + + H = (0.5*dot(p, p) - 1/dot(q, q)**0.5)*dx + L = dot(p, perp(q))*dx + + test = TestFunction(Z) + Hz = derivative(H, z, test) + F = inner(Dt(z), test)*dx - replace(Hz, {test: dot(J.T, test)}) + + stepper = GalerkinTimeStepper(F, order, t, dt, z, solver_parameters={"mat_type": "nest", + "snes_rtol": 1E-14,}) + + energies = [] + + print(float(t), assemble(H), assemble(L)) + + T = 10 + while (float(t) < T): + if (float(t) + float(dt) > T): + dt.assign(T - float(t)) + stepper.advance() + t += dt + + print(float(t), assemble(H), assemble(L)) + energies.append(assemble(H)) + + +@pytest.mark.parametrize('order', (1,)) +def test_kepler_aux_variable(order): + msh = UnitIntervalMesh(1) + MC = MeshConstant(msh) + t = MC.Constant(0.0) + dt = MC.Constant(pi/10) + + sd = 2 + nbody = 2 + dim = (nbody-1)*sd + V = VectorFunctionSpace(msh, "DG", 0, dim=2*dim) + Z = V * V + z = Function(Z) + + pq, aux = split(z) + pq = variable(pq) + p = as_vector([pq[k] for k in range(dim)]) + q = as_vector([pq[k] for k in range(dim, 2*dim)]) + + J = as_matrix(np.kron([[0, -1], [1, 0]], np.eye(dim))) + + H = (0.5*dot(p, p) - 1/sqrt(dot(q, q))) + + L = dot(p, perp(q))*dx + + test = TestFunction(Z) + test_pq, test_aux = split(test) + + ufc_line = ufc_simplex(1) + Qhigh = create_quadrature(ufc_line, 4) + Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) + + F = inner(Dt(pq) - dot(J, aux), test_pq/dt)*dx + + Hpq = diff(H, pq) + F += Lhigh(inner(aux - Hpq, test_aux)*dx) + print(F.form) + H = H * dx + + # Initial condition + vals = (0, 2, 0.4, 0) + z.subfunctions[0].interpolate(as_vector(vals)) + + aux_indices = list(range(2*dim, 4*dim)) + #aux_indices = None + sp = { + #"snes_monitor": None, + "snes_linesearch_type": "l2", + #"snes_converged_reason": None, + "snes_atol": 1.0e-14, + "snes_rtol": 1.0e-14, + "mat_type": "dense", + "pc_type": "lu" + } + + stepper = GalerkinTimeStepper(F, order, t, dt, z, + solver_parameters=sp, + aux_indices=aux_indices) + + energies = [] + print(float(t), assemble(H), assemble(L)) + + T = 10 + #while (float(t) < T): + for _ in range(2): + if (float(t) + float(dt) > T): + dt.assign(T - float(t)) + print("F", assemble(stepper.solver._problem.F).dat.data) + stepper.advance() + t += dt + + print(float(t), assemble(H), assemble(L)) + energies.append(assemble(H)) From 7edacdf20e09bd82310b6cd2f45447ef7fa20f7a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 20 May 2025 15:29:04 +0100 Subject: [PATCH 03/24] Galerkin: set constant-in-time initial guess --- irksome/galerkin_stepper.py | 107 ++++++++++++++++++++---------------- irksome/labeling.py | 40 +++++++++++--- tests/test_galerkin.py | 92 +++++++++++++------------------ 3 files changed, 131 insertions(+), 108 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 7e00288e..71bd7506 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -1,17 +1,14 @@ from FIAT import (Bernstein, DiscontinuousElement, DiscontinuousLagrange, IntegratedLegendre, Lagrange, Legendre, ufc_simplex) from FIAT.quadrature_schemes import create_quadrature -from ufl import Form from ufl.constantvalue import as_ufl from .base_time_stepper import StageCoupledTimeStepper from .bcs import bc2space, stage2spaces4bc from .deriv import TimeDerivative, expand_time_derivatives -from .labeling import as_labelled_form +from .labeling import split_quadrature from .tools import replace, vecconst -from .labeling import TimeQuadratureRule import numpy as np -from firedrake import TestFunction -from firedrake import assemble, derivative +from firedrake import TestFunction, Constant ufc_line = ufc_simplex(1) @@ -37,9 +34,14 @@ def getElements(basis_type, order): return trial_el, test_el -def getTermGalerkin(Fcur, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test, aux_indices): - v, = Fcur.arguments() - Fcur = expand_time_derivatives(Fcur, t=t, timedep_coeffs=(u0,)) +def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices): + # preprocess time derivatives + F = expand_time_derivatives(F, t=t, timedep_coeffs=(u0,)) + v, = F.arguments() + assert v.function_space() == u0.function_space() + + qpts = Q.get_points() + qwts = Q.get_weights() tabulate_trials = L_trial.tabulate(1, qpts) trial_vals = tabulate_trials[(0,)] trial_dvals = tabulate_trials[(1,)] @@ -57,25 +59,31 @@ def getTermGalerkin(Fcur, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test, # Expand unknowns in the trial space, including the initial condition u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) + vsub = test_vals_w.T @ v_np usub = trial_vals.T @ u_np dtu0sub = trial_dvals.T @ u_np if aux_indices is not None: - # Expand auxiliary variables in the test space rather than the trial space - usub[:, aux_indices] = test_vals.T @ w_np[:, aux_indices] + cur = 0 + aux_comps = [] + V = v.function_space() + for i, Vi in enumerate(V): + if i in aux_indices: + aux_comps.extend(range(cur, cur+Vi.value_size)) + cur += Vi.value_size - vsub = test_vals_w.T @ v_np + # Expand auxiliary variables in the test space rather than the trial space + usub[:, aux_comps] = test_vals.T @ w_np[:, aux_comps] dtu0 = TimeDerivative(u0) - Fnew = Form([]) # now loop over quadrature points + repl = {} for q in range(len(qpts)): - repl = {t: t + qpts[q] * dt, - v: vsub[q] * dt, - u0: usub[q], - dtu0: dtu0sub[q] / dt} - Fnew += replace(Fcur, repl) - + repl[q] = {t: t + qpts[q] * dt, + v: vsub[q] * dt, + u0: usub[q], + dtu0: dtu0sub[q] / dt} + Fnew = sum(replace(F, repl[q]) for q in repl) return Fnew @@ -113,35 +121,13 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a assert Qdefault.ref_el.get_spatial_dimension() == 1 assert L_trial.get_order() == L_test.get_order() + 1 - # preprocess time derivatives - F = as_labelled_form(F) - - v, = F.form.arguments() - V = v.function_space() - assert V == u0.function_space() - num_stages = L_test.space_dimension() Vbig = stages.function_space() test = TestFunction(Vbig) - Fnew = Form([]) - for term in F.terms: - print(term) - labels = term.labels - print(labels) - quad_labels = [label for label in labels if isinstance(label, TimeQuadratureRule)] - if len(quad_labels) == 0: - qpts = Qdefault.get_points() - qwts = Qdefault.get_weights() - elif len(quad_labels) == 1: - qpts = np.asarray(quad_labels[0].x) - qwts = np.asarray(quad_labels[0].w) - else: - raise ValueError("Multiple quadrature labels on one form.") - - Fnew += getTermGalerkin(term.form, L_trial, L_test, qpts, qwts, t, dt, u0, stages, test, aux_indices) - - assemble(derivative(Fnew, stages), mat_type="aij").petscmat.convert("dense").view() + splitting = split_quadrature(F, Qdefault=Qdefault) + Fnew = sum(getTermGalerkin(Fcur, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices) + for Q, Fcur in splitting.items()) # mass-ish matrix for BC, based on default quadrature rule qpts = Qdefault.get_points() @@ -159,6 +145,7 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a qpts = vecconst(qpts.reshape((-1,))) # Oh, honey, is it the boundary conditions? + V = u0.function_space() if bcs is None: bcs = [] bcsnew = [] @@ -222,7 +209,6 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None, assert order >= 1 self.order = order self.basis_type = basis_type - self.aux_indices = aux_indices V = u0.function_space() self.num_fields = len(V) @@ -230,11 +216,14 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None, self.trial_el, self.test_el = getElements(basis_type, order) if quadrature is None: - quadrature = create_quadrature(ufc_line, 2 * order) + quad_degree = self.trial_el.degree() + self.test_el.degree() + quadrature = create_quadrature(ufc_line, quad_degree) self.quadrature = quadrature - assert np.size(quadrature.get_points()) >= order + 1 + assert np.size(quadrature.get_points()) >= order + self.aux_indices = aux_indices super().__init__(F, t, dt, u0, order, bcs=bcs, **kwargs) + self.set_initial_guess() def get_form_and_bcs(self, stages, basis_type=None, order=None, quadrature=None, aux_indices=None, F=None): if basis_type is None: @@ -256,3 +245,29 @@ def _update(self): k1, = self.trial_el.entity_dofs()[0][1] for i, u0bit in enumerate(self.u0.subfunctions): u0bit.assign(self.stages.subfunctions[self.num_fields*(k1-1)+i]) + + def set_initial_guess(self): + # Set a constant-in-time initial guess + ref_el = self.test_el.get_reference_element() + P0 = DiscontinuousLagrange(ref_el, 0) + P0 = P0.get_nodal_basis() + B = P0.get_coeffs() + + test_dual = self.test_el.get_dual_set() + test_dofs = np.dot(test_dual.to_riesz(P0), B) + + trial_dual = self.trial_el.get_dual_set() + trial_dofs = np.dot(trial_dual.to_riesz(P0), B) + + dof = Constant(0) + for k in range(self.num_stages): + for i, u0bit in enumerate(self.u0.subfunctions): + sbit = self.stages.subfunctions[self.num_fields*k+i] + if self.aux_indices and i in self.aux_indices: + dof.assign(test_dofs[k]) + else: + dof.assign(trial_dofs[k+1]) + if abs(float(dof)) < 1E-12: + sbit.zero() + else: + sbit.assign(u0bit * dof) diff --git a/irksome/labeling.py b/irksome/labeling.py index ee0eccfe..1711ee84 100644 --- a/irksome/labeling.py +++ b/irksome/labeling.py @@ -1,8 +1,7 @@ from firedrake.fml import Label, keep, drop, LabelledForm - +import numpy as np explicit = Label("explicit") -empty_label = Label("") class TimeQuadratureLabel(Label): @@ -15,6 +14,36 @@ def __init__(self, x, w): self.x = x self.w = w + def get_points(self): + return np.asarray(self.x) + + def get_weights(self): + return np.asarray(self.w) + + +def split_quadrature(F, Qdefault=None): + if not isinstance(F, LabelledForm): + return {Qdefault: F} + + quad_labels = set() + for term in F.terms: + cur_labels = [label for label in term.labels if isinstance(label, TimeQuadratureRule)] + if len(cur_labels) == 1: + quad_labels.update(cur_labels) + elif len(cur_labels) > 1: + raise ValueError("Multiple quadrature labels on one term.") + + splitting = {Q: F.label_map(lambda t: Q in t.labels, map_if_true=keep, map_if_false=drop) + for Q in quad_labels} + splitting[Qdefault] = F.label_map(lambda t: len(quad_labels.intersection(t.labels)) > 0, + map_if_true=drop, map_if_false=keep) + for Q in list(splitting): + try: + splitting[Q] = splitting[Q].form + except TypeError: + splitting.pop(Q) + return splitting + def split_explicit(F): if not isinstance(F, LabelledForm): @@ -27,10 +56,3 @@ def split_explicit(F): map_if_true=keep, map_if_false=drop) return imp_part.form, exp_part.form - - -def as_labelled_form(F): - if not isinstance(F, LabelledForm): - return empty_label(F) - else: - return F diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index b2078326..a7a4b829 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -267,66 +267,64 @@ def test_wave_eq_galerkin(deg, N, order): assert np.allclose(energy[1:], energy[:-1]) - - - -@pytest.mark.parametrize('order', (1,)) +@pytest.mark.parametrize('order', (1, 2)) def test_kepler(order): msh = UnitIntervalMesh(1) MC = MeshConstant(msh) t = MC.Constant(0.0) - dt = MC.Constant(4/16) + dt = MC.Constant(pi/100) + Nsteps = 10 sd = 2 nbody = 2 dim = (nbody-1)*sd - R = FunctionSpace(msh, "R", 0) - V = MixedFunctionSpace([R]*dim) - Z = V * V + Z = VectorFunctionSpace(msh, "DG", 0, dim=2*dim) z = Function(Z) - p0 = [0, 1] - q0 = [1, 0] - vals = p0 + q0 - for c, val in zip(z.subfunctions, vals): - c.assign(val) + pq = z + p = as_vector([pq[k] for k in range(dim)]) + q = as_vector([pq[k] for k in range(dim, 2*dim)]) - c = split(z) - p = as_vector(c[:dim]) - q = as_vector(c[dim:]) J = as_matrix(np.kron([[0, -1], [1, 0]], np.eye(dim))) H = (0.5*dot(p, p) - 1/dot(q, q)**0.5)*dx L = dot(p, perp(q))*dx + Qhigh = create_quadrature(ufc_simplex(1), 25) + Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) + test = TestFunction(Z) Hz = derivative(H, z, test) - F = inner(Dt(z), test)*dx - replace(Hz, {test: dot(J.T, test)}) + F = inner(Dt(z), test)*dx + Lhigh(-replace(Hz, {test: dot(J.T, test)})) - stepper = GalerkinTimeStepper(F, order, t, dt, z, solver_parameters={"mat_type": "nest", - "snes_rtol": 1E-14,}) + # Initial condition + vals = Constant((0, 2, 0.4, 0)) + z.subfunctions[0].interpolate(vals) + sp = { + "snes_linesearch_type": "l2", + "snes_atol": 1.0e-14, + "snes_rtol": 1.0e-14, + "mat_type": "dense", + "pc_type": "lu" + } + stepper = GalerkinTimeStepper(F, order, t, dt, z, solver_parameters=sp, basis_type="integral") energies = [] - print(float(t), assemble(H), assemble(L)) - - T = 10 - while (float(t) < T): - if (float(t) + float(dt) > T): - dt.assign(T - float(t)) + for _ in range(Nsteps): stepper.advance() t += dt - print(float(t), assemble(H), assemble(L)) energies.append(assemble(H)) -@pytest.mark.parametrize('order', (1,)) +@pytest.mark.parametrize('order', (1, 2)) def test_kepler_aux_variable(order): msh = UnitIntervalMesh(1) MC = MeshConstant(msh) t = MC.Constant(0.0) - dt = MC.Constant(pi/10) + dt = MC.Constant(pi/100) + Nsteps = 10 sd = 2 nbody = 2 @@ -342,55 +340,43 @@ def test_kepler_aux_variable(order): J = as_matrix(np.kron([[0, -1], [1, 0]], np.eye(dim))) - H = (0.5*dot(p, p) - 1/sqrt(dot(q, q))) + H = (0.5*dot(p, p) - 1/sqrt(dot(q, q)))*dx L = dot(p, perp(q))*dx test = TestFunction(Z) test_pq, test_aux = split(test) - ufc_line = ufc_simplex(1) - Qhigh = create_quadrature(ufc_line, 4) - Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) + Qlow = create_quadrature(ufc_simplex(1), 2*order-2) + Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) - F = inner(Dt(pq) - dot(J, aux), test_pq/dt)*dx + Qhigh = create_quadrature(ufc_simplex(1), 25) + Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) - Hpq = diff(H, pq) - F += Lhigh(inner(aux - Hpq, test_aux)*dx) - print(F.form) - H = H * dx + Hpq = diff(H.integrals()[0].integrand(), pq) + pq, aux = split(z) + F = Llow(inner(Dt(pq) - J*aux, test_pq)*dx) + F += Llow(inner(aux, test_aux)*dx) + Lhigh(-inner(Hpq, test_aux)*dx) # Initial condition - vals = (0, 2, 0.4, 0) - z.subfunctions[0].interpolate(as_vector(vals)) - - aux_indices = list(range(2*dim, 4*dim)) - #aux_indices = None + pq_init = Constant((0, 2, 0.4, 0)) + z.subfunctions[0].interpolate(pq_init) sp = { - #"snes_monitor": None, "snes_linesearch_type": "l2", - #"snes_converged_reason": None, "snes_atol": 1.0e-14, "snes_rtol": 1.0e-14, "mat_type": "dense", "pc_type": "lu" } - + aux_indices = list(range(2*dim, 4*dim)) stepper = GalerkinTimeStepper(F, order, t, dt, z, solver_parameters=sp, aux_indices=aux_indices) energies = [] print(float(t), assemble(H), assemble(L)) - - T = 10 - #while (float(t) < T): - for _ in range(2): - if (float(t) + float(dt) > T): - dt.assign(T - float(t)) - print("F", assemble(stepper.solver._problem.F).dat.data) + for _ in range(Nsteps): stepper.advance() t += dt - print(float(t), assemble(H), assemble(L)) energies.append(assemble(H)) From 004e3f5e5be21743ecc4ed39b0f6cf681a23f8db Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 20 May 2025 17:02:10 +0100 Subject: [PATCH 04/24] Conservation of all invariants --- tests/test_galerkin.py | 46 ++++++++++++++++++++++++++++-------------- 1 file changed, 31 insertions(+), 15 deletions(-) diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index a7a4b829..abcf4c69 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -330,22 +330,29 @@ def test_kepler_aux_variable(order): nbody = 2 dim = (nbody-1)*sd V = VectorFunctionSpace(msh, "DG", 0, dim=2*dim) - Z = V * V + Z = V * V * V * V z = Function(Z) - pq, aux = split(z) + pq, w0, w1, w2 = split(z) pq = variable(pq) p = as_vector([pq[k] for k in range(dim)]) q = as_vector([pq[k] for k in range(dim, 2*dim)]) - J = as_matrix(np.kron([[0, -1], [1, 0]], np.eye(dim))) + T = 0.5*dot(p, p) + U = -1/sqrt(dot(q, q)) - H = (0.5*dot(p, p) - 1/sqrt(dot(q, q)))*dx + # Invariants + H = T + U + L = dot(p, perp(q)) + A1, A2 = U*q - L*perp(p) - L = dot(p, perp(q))*dx + invariants = [H*dx, L*dx, A1*dx, A2*dx] + dHdu = diff(H, pq) + dA1du = diff(A1, pq) + dA2du = diff(A2, pq) test = TestFunction(Z) - test_pq, test_aux = split(test) + test_pq, v0, v1, v2 = split(test) Qlow = create_quadrature(ufc_simplex(1), 2*order-2) Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) @@ -353,30 +360,39 @@ def test_kepler_aux_variable(order): Qhigh = create_quadrature(ufc_simplex(1), 25) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) - Hpq = diff(H.integrals()[0].integrand(), pq) - pq, aux = split(z) - F = Llow(inner(Dt(pq) - J*aux, test_pq)*dx) - F += Llow(inner(aux, test_aux)*dx) + Lhigh(-inner(Hpq, test_aux)*dx) + # determinant_forms = [test_pq, dHdu, dA1du, dA2du] + determinant_forms = [test_pq, w0, w1, w2] + tensor = as_tensor(determinant_forms) + + F = Llow(inner(Dt(pq), test_pq)*dx) + Llow(-(det(tensor) / (2*L*H))*dx) + F += Llow(inner(w0, v0)*dx) + Lhigh(-inner(dHdu, v0)*dx) + F += Llow(inner(w1, v1)*dx) + Lhigh(-inner(dA1du, v1)*dx) + F += Llow(inner(w2, v2)*dx) + Lhigh(-inner(dA2du, v2)*dx) # Initial condition pq_init = Constant((0, 2, 0.4, 0)) z.subfunctions[0].interpolate(pq_init) + z.subfunctions[1].interpolate(dHdu) + z.subfunctions[2].interpolate(dA1du) + z.subfunctions[3].interpolate(dA2du) sp = { + "snes_converged_reason": None, "snes_linesearch_type": "l2", "snes_atol": 1.0e-14, "snes_rtol": 1.0e-14, "mat_type": "dense", "pc_type": "lu" } - aux_indices = list(range(2*dim, 4*dim)) + aux_indices = list(range(V.value_size, Z.value_size)) stepper = GalerkinTimeStepper(F, order, t, dt, z, solver_parameters=sp, + basis_type="integral", aux_indices=aux_indices) - energies = [] - print(float(t), assemble(H), assemble(L)) + print() + print(float(t), *map(assemble, invariants)) for _ in range(Nsteps): stepper.advance() t += dt - print(float(t), assemble(H), assemble(L)) - energies.append(assemble(H)) + print(f"u({float(t)}) = ", z.subfunctions[0].dat.data) + print(float(t), *map(assemble, invariants)) From 8f644520dea8e9e4bf8062fcb5d11c3c01469d51 Mon Sep 17 00:00:00 2001 From: "Robert C. Kirby" Date: Wed, 21 May 2025 11:13:06 +0100 Subject: [PATCH 05/24] mesh constant? --- tests/test_galerkin.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index 266d3cfe..9ec2e520 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -191,11 +191,11 @@ def test_1d_heat_homogeneous_dirichletbc_timequadlabels(order): Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) - F = ( - Llow(inner(Dt(u), v) * dx) - + inner(grad(u), grad(v)) * dx - - Lhigh(inner(rhs, v) * dx) - ) + F0 = inner(Dt(u), v) * dx + F1 = inner(grad(u), grad(v)) * dx + F2 = inner(rhs, v) * dx + F = Llow(F0) + F1 - Lhigh(F2) + luparams = {"mat_type": "aij", "ksp_type": "preonly", "pc_type": "lu"} From 094927af473762af27ae55e104a7d369f0522d44 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 21 May 2025 12:07:27 +0100 Subject: [PATCH 06/24] Fix tests --- irksome/galerkin_stepper.py | 11 ++- tests/test_galerkin.py | 134 +++++++++++++++--------------------- 2 files changed, 59 insertions(+), 86 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 71bd7506..a95daefa 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -65,13 +65,13 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices) if aux_indices is not None: cur = 0 aux_comps = [] - V = v.function_space() - for i, Vi in enumerate(V): - if i in aux_indices: - aux_comps.extend(range(cur, cur+Vi.value_size)) - cur += Vi.value_size + for field, Vsub in enumerate(v.function_space()): + if field in aux_indices: + aux_comps.extend(range(cur, cur+Vsub.value_size)) + cur += Vsub.value_size # Expand auxiliary variables in the test space rather than the trial space + test_vals = vecconst(test_vals) usub[:, aux_comps] = test_vals.T @ w_np[:, aux_comps] dtu0 = TimeDerivative(u0) @@ -140,7 +140,6 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a mmat = test_vals_w @ trial_vals[1:].T proj = vecconst(np.linalg.solve(mmat, test_vals_w)) - test_vals_w = vecconst(test_vals_w) trial_vals = vecconst(trial_vals) qpts = vecconst(qpts.reshape((-1,))) diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index abcf4c69..d399985e 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -267,76 +267,36 @@ def test_wave_eq_galerkin(deg, N, order): assert np.allclose(energy[1:], energy[:-1]) -@pytest.mark.parametrize('order', (1, 2)) -def test_kepler(order): - msh = UnitIntervalMesh(1) - MC = MeshConstant(msh) - t = MC.Constant(0.0) - dt = MC.Constant(pi/100) - Nsteps = 10 - - sd = 2 - nbody = 2 - dim = (nbody-1)*sd - Z = VectorFunctionSpace(msh, "DG", 0, dim=2*dim) - z = Function(Z) - - pq = z - p = as_vector([pq[k] for k in range(dim)]) - q = as_vector([pq[k] for k in range(dim, 2*dim)]) +def kepler(V, order, t, dt, u0, solver_parameters): + dim = V.value_size//2 + u = Function(V) + u.interpolate(u0) + p = as_vector([u[k] for k in range(dim)]) + q = as_vector([u[k] for k in range(dim, 2*dim)]) J = as_matrix(np.kron([[0, -1], [1, 0]], np.eye(dim))) - - H = (0.5*dot(p, p) - 1/dot(q, q)**0.5)*dx - L = dot(p, perp(q))*dx + H = (0.5*dot(p, p) - 1/sqrt(dot(q, q)))*dx Qhigh = create_quadrature(ufc_simplex(1), 25) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) - test = TestFunction(Z) - Hz = derivative(H, z, test) - F = inner(Dt(z), test)*dx + Lhigh(-replace(Hz, {test: dot(J.T, test)})) - - # Initial condition - vals = Constant((0, 2, 0.4, 0)) - z.subfunctions[0].interpolate(vals) - sp = { - "snes_linesearch_type": "l2", - "snes_atol": 1.0e-14, - "snes_rtol": 1.0e-14, - "mat_type": "dense", - "pc_type": "lu" - } - stepper = GalerkinTimeStepper(F, order, t, dt, z, solver_parameters=sp, basis_type="integral") - - energies = [] - print(float(t), assemble(H), assemble(L)) - for _ in range(Nsteps): - stepper.advance() - t += dt - print(float(t), assemble(H), assemble(L)) - energies.append(assemble(H)) - + test = TestFunction(V) + dHdu = derivative(H, u, test) + F = inner(Dt(u), test)*dx + Lhigh(-replace(dHdu, {test: dot(J.T, test)})) + stepper = GalerkinTimeStepper(F, order, t, dt, u, solver_parameters=solver_parameters) + return stepper, [H] -@pytest.mark.parametrize('order', (1, 2)) -def test_kepler_aux_variable(order): - msh = UnitIntervalMesh(1) - MC = MeshConstant(msh) - t = MC.Constant(0.0) - dt = MC.Constant(pi/100) - Nsteps = 10 - sd = 2 - nbody = 2 - dim = (nbody-1)*sd - V = VectorFunctionSpace(msh, "DG", 0, dim=2*dim) +def kepler_aux_variable(V, order, t, dt, u0, solver_parameters): + dim = V.value_size//2 Z = V * V * V * V z = Function(Z) + z.subfunctions[0].interpolate(u0) - pq, w0, w1, w2 = split(z) - pq = variable(pq) - p = as_vector([pq[k] for k in range(dim)]) - q = as_vector([pq[k] for k in range(dim, 2*dim)]) + u, w0, w1, w2 = split(z) + u = variable(u) + p = as_vector([u[k] for k in range(dim)]) + q = as_vector([u[k] for k in range(dim, 2*dim)]) T = 0.5*dot(p, p) U = -1/sqrt(dot(q, q)) @@ -347,12 +307,12 @@ def test_kepler_aux_variable(order): A1, A2 = U*q - L*perp(p) invariants = [H*dx, L*dx, A1*dx, A2*dx] - dHdu = diff(H, pq) - dA1du = diff(A1, pq) - dA2du = diff(A2, pq) + dHdu = diff(H, u) + dA1du = diff(A1, u) + dA2du = diff(A2, u) test = TestFunction(Z) - test_pq, v0, v1, v2 = split(test) + test_u, v0, v1, v2 = split(test) Qlow = create_quadrature(ufc_simplex(1), 2*order-2) Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) @@ -360,21 +320,34 @@ def test_kepler_aux_variable(order): Qhigh = create_quadrature(ufc_simplex(1), 25) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) - # determinant_forms = [test_pq, dHdu, dA1du, dA2du] - determinant_forms = [test_pq, w0, w1, w2] + # determinant_forms = [test_u, dHdu, dA1du, dA2du] + determinant_forms = [test_u, w0, w1, w2] tensor = as_tensor(determinant_forms) - F = Llow(inner(Dt(pq), test_pq)*dx) + Llow(-(det(tensor) / (2*L*H))*dx) + F = Llow(inner(Dt(u), test_u)*dx) + Lhigh(-(det(tensor) / (2*L*H))*dx) F += Llow(inner(w0, v0)*dx) + Lhigh(-inner(dHdu, v0)*dx) F += Llow(inner(w1, v1)*dx) + Lhigh(-inner(dA1du, v1)*dx) F += Llow(inner(w2, v2)*dx) + Lhigh(-inner(dA2du, v2)*dx) - # Initial condition - pq_init = Constant((0, 2, 0.4, 0)) - z.subfunctions[0].interpolate(pq_init) - z.subfunctions[1].interpolate(dHdu) - z.subfunctions[2].interpolate(dA1du) - z.subfunctions[3].interpolate(dA2du) + # Auxiliary variable subspaces + aux_indices = list(range(1, len(Z))) + stepper = GalerkinTimeStepper(F, order, t, dt, z, + solver_parameters=solver_parameters, + aux_indices=aux_indices) + return stepper, invariants + + +@pytest.mark.parametrize('order', (1, 2)) +@pytest.mark.parametrize('problem', (kepler, kepler_aux_variable)) +def test_kepler(problem, order): + msh = UnitIntervalMesh(1) + MC = MeshConstant(msh) + t = MC.Constant(0.0) + dt = MC.Constant(pi/100) + Nsteps = 2 + + dim = 2 + V = VectorFunctionSpace(msh, "DG", 0, dim=2*dim) sp = { "snes_converged_reason": None, "snes_linesearch_type": "l2", @@ -383,16 +356,17 @@ def test_kepler_aux_variable(order): "mat_type": "dense", "pc_type": "lu" } - aux_indices = list(range(V.value_size, Z.value_size)) - stepper = GalerkinTimeStepper(F, order, t, dt, z, - solver_parameters=sp, - basis_type="integral", - aux_indices=aux_indices) + + # Initial condition + u0 = Constant((0, 2, 0.4, 0)) + stepper, invariants = problem(V, order, t, dt, u0, sp) print() - print(float(t), *map(assemble, invariants)) + E0 = np.asarray(list(map(assemble, invariants))) + print(float(t), E0) for _ in range(Nsteps): stepper.advance() t += dt - print(f"u({float(t)}) = ", z.subfunctions[0].dat.data) - print(float(t), *map(assemble, invariants)) + Et = np.asarray(list(map(assemble, invariants))) + print(float(t), Et) + assert np.allclose(E0, Et) From eed8d45e58a55b349332ee678edeaf136813c7ff Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 21 May 2025 12:19:16 +0100 Subject: [PATCH 07/24] lint --- irksome/galerkin_stepper.py | 2 ++ tests/test_galerkin.py | 1 - 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index a95daefa..cd1ce3e1 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -202,6 +202,8 @@ class GalerkinTimeStepper(StageCoupledTimeStepper): instance to be passed to a `firedrake.MixedVectorSpaceBasis` over the larger space associated with the Runge-Kutta method + :kwarg aux_indices: a list of field indices to be discretized in the test space + rather than trial space. """ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None, quadrature=None, aux_indices=None, **kwargs): diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index 1802c3dd..a41c3bbb 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -196,7 +196,6 @@ def test_1d_heat_homogeneous_dirichletbc_timequadlabels(order): F2 = inner(rhs, v) * dx F = Llow(F0) + F1 - Lhigh(F2) - luparams = {"mat_type": "aij", "ksp_type": "preonly", "pc_type": "lu"} stepper = GalerkinTimeStepper( From ad23e4d53b7ff0af8151a01e35a646a0358917b1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 24 May 2025 13:01:52 +0100 Subject: [PATCH 08/24] Refactor auxiliary variables --- irksome/galerkin_stepper.py | 16 ++-------------- irksome/tools.py | 15 ++++++++++++++- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index cd1ce3e1..0701e537 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -6,7 +6,7 @@ from .bcs import bc2space, stage2spaces4bc from .deriv import TimeDerivative, expand_time_derivatives from .labeling import split_quadrature -from .tools import replace, vecconst +from .tools import replace, vecconst, replace_auxiliary_variables import numpy as np from firedrake import TestFunction, Constant @@ -36,6 +36,7 @@ def getElements(basis_type, order): def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices): # preprocess time derivatives + F = replace_auxiliary_variables(F, u0, aux_indices) F = expand_time_derivatives(F, t=t, timedep_coeffs=(u0,)) v, = F.arguments() assert v.function_space() == u0.function_space() @@ -57,23 +58,10 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices) v_np = np.reshape(test, (-1, *u0.ufl_shape)) w_np = np.reshape(stages, (-1, *u0.ufl_shape)) - # Expand unknowns in the trial space, including the initial condition u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) vsub = test_vals_w.T @ v_np usub = trial_vals.T @ u_np dtu0sub = trial_dvals.T @ u_np - if aux_indices is not None: - cur = 0 - aux_comps = [] - for field, Vsub in enumerate(v.function_space()): - if field in aux_indices: - aux_comps.extend(range(cur, cur+Vsub.value_size)) - cur += Vsub.value_size - - # Expand auxiliary variables in the test space rather than the trial space - test_vals = vecconst(test_vals) - usub[:, aux_comps] = test_vals.T @ w_np[:, aux_comps] - dtu0 = TimeDerivative(u0) # now loop over quadrature points diff --git a/irksome/tools.py b/irksome/tools.py index 197ef204..dcd62636 100644 --- a/irksome/tools.py +++ b/irksome/tools.py @@ -1,7 +1,7 @@ from operator import mul from functools import reduce import numpy -from firedrake import Function, FunctionSpace, MixedVectorSpaceBasis, Constant +from firedrake import Function, FunctionSpace, MixedVectorSpaceBasis, Constant, split from ufl.algorithms.analysis import extract_type from ufl import as_tensor, zero from ufl import replace as ufl_replace @@ -68,6 +68,19 @@ def replace(e, mapping): return ufl_replace(e, cmapping) +def replace_auxiliary_variables(F, u0, aux_indices): + """Discretize the fields corresponding to aux_indices in Dt(V).""" + if aux_indices is None: + return F + + components = [] + for i, usub in enumerate(split(u0)): + if i in aux_indices: + usub = TimeDerivative(usub) + components.extend(usub[i] for i in numpy.ndindex(usub.ufl_shape)) + return replace(F, {u0: numpy.reshape(components, u0.ufl_shape)}) + + # Utility functions that help us refactor def AI(A): return (A, numpy.eye(*A.shape, dtype=A.dtype)) From 232919119a0d9c8d884a165c280971fbedc965ca Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 28 May 2025 14:17:39 +0100 Subject: [PATCH 09/24] TimeProjector --- irksome/__init__.py | 2 +- irksome/galerkin_stepper.py | 100 ++++++++++++++++++++++++++++++++---- tests/test_galerkin.py | 63 +++++++++++++++++++---- 3 files changed, 144 insertions(+), 21 deletions(-) diff --git a/irksome/__init__.py b/irksome/__init__.py index 8b4ad5e2..73d4f000 100644 --- a/irksome/__init__.py +++ b/irksome/__init__.py @@ -27,5 +27,5 @@ from .stepper import TimeStepper # noqa: F401 from .tools import MeshConstant # noqa: F401 from .wso_dirk_tableaux import WSODIRK # noqa: F401 -from .galerkin_stepper import GalerkinTimeStepper # noqa: F401 +from .galerkin_stepper import GalerkinTimeStepper, TimeProjector # noqa: F401 from .discontinuous_galerkin_stepper import DiscontinuousGalerkinTimeStepper # noqa: F401 diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 0701e537..3e90c87d 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -8,7 +8,14 @@ from .labeling import split_quadrature from .tools import replace, vecconst, replace_auxiliary_variables import numpy as np -from firedrake import TestFunction, Constant + +from ufl import Coefficient +from ufl.core.operator import Operator +from ufl.core.ufl_type import ufl_type +from ufl.corealg.multifunction import MultiFunction +from ufl.algorithms.map_integrands import map_integrand_dags +from ufl.algorithms import expand_derivatives +from firedrake import Constant, Function, TestFunction, VectorFunctionSpace, diff ufc_line = ufc_simplex(1) @@ -34,15 +41,26 @@ def getElements(basis_type, order): return trial_el, test_el -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 = 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 = replace_auxiliary_variables(F, u0, aux_indices) F = expand_time_derivatives(F, t=t, timedep_coeffs=(u0,)) v, = F.arguments() - assert v.function_space() == u0.function_space() - - qpts = Q.get_points() - qwts = Q.get_weights() tabulate_trials = L_trial.tabulate(1, qpts) trial_vals = tabulate_trials[(0,)] trial_dvals = tabulate_trials[(1,)] @@ -55,10 +73,10 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices) qpts = vecconst(np.reshape(qpts, (-1,))) # set up the pieces we need to work with to do our substitutions - v_np = np.reshape(test, (-1, *u0.ufl_shape)) + v_np = np.reshape(test, (-1, *v.ufl_shape)) w_np = np.reshape(stages, (-1, *u0.ufl_shape)) - u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) + vsub = test_vals_w.T @ v_np usub = trial_vals.T @ u_np dtu0sub = trial_dvals.T @ u_np @@ -70,7 +88,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 @@ -82,7 +103,7 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a :arg F: UFL form for the semidiscrete ODE/DAE :arg L_trial: A :class:`FIAT.FiniteElement` for the trial functions in time - :arg L_test: A :class:`FIAT.FinteElement` for the test functions in time + :arg L_test: A :class:`FIAT.FiniteElement` for the test functions in time :arg Q: A :class:`FIAT.QuadratureRule` for the time integration :arg t: a :class:`Function` on the Real space over the same mesh as `u0`. This serves as a variable referring to the current time. @@ -260,3 +281,62 @@ def set_initial_guess(self): sbit.zero() else: sbit.assign(u0bit * dof) + + +@ufl_type( + num_ops=1, + inherit_shape_from_operand=0, + inherit_indices_from_operand=0, +) +class TimeProjector(Operator): + __slots__ = ("form", "order", "quadrature") + + def __init__(self, F, order, Q): + self.form = expand_derivatives(F) + self.order = order + self.quadrature = Q + c = Coefficient(F.arguments()[0].function_space()) + Operator.__init__(self, operands=(c,)) + + +class TimeProjectorDispatcher(MultiFunction): + def __init__(self, element, t, dt, u0, u1, stages, phi): + MultiFunction.__init__(self) + 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,)) + + def time_projector(self, o): + # use the internal copy of the state, so it does not get updated again in the outer quadrature + F = replace(o.form, {self.u0: self.u1}) + order = o.order + Q = o.quadrature + assert order+1 <= len(self.phi) + + # 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, = o.ufl_operands + test = np.outer(Minv[:order+1] @ self.phi, np.asarray(c, dtype=object)) / self.dt + 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 + + expr = MultiFunction.reuse_if_untouched + + +def expand_time_projectors(expression, element, t, dt, u0, u1, stages, phi): + rules = TimeProjectorDispatcher(element, t, dt, u0, u1, stages, phi) + return map_integrand_dags(rules, expression) diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index a41c3bbb..80c1c93d 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -2,7 +2,7 @@ import pytest from firedrake import * -from irksome import Dt, MeshConstant, GalerkinTimeStepper +from irksome import Dt, MeshConstant, GalerkinTimeStepper, TimeProjector from irksome import TimeStepper, GaussLegendre from irksome.labeling import TimeQuadratureLabel from FIAT import make_quadrature, ufc_simplex @@ -266,7 +266,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) @@ -274,16 +274,19 @@ def kepler(V, order, t, dt, u0, solver_parameters): p = as_vector([u[k] for k in range(dim)]) q = as_vector([u[k] for k in range(dim, 2*dim)]) J = as_matrix(np.kron([[0, -1], [1, 0]], np.eye(dim))) - H = (0.5*dot(p, p) - 1/sqrt(dot(q, q)))*dx + H = (0.5*dot(p, p) - 1/sqrt(dot(q, q))) - Qhigh = create_quadrature(ufc_simplex(1), 25) + Qlow = create_quadrature(ufc_simplex(1), 2*order-2) + Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) + + Qhigh = create_quadrature(ufc_simplex(1), 2*order+2) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) test = TestFunction(V) - dHdu = derivative(H, u, test) - F = inner(Dt(u), test)*dx + Lhigh(-replace(dHdu, {test: dot(J.T, test)})) + dHdu = derivative(H*dx, u, test) + F = Llow(inner(Dt(u), test)*dx) + Lhigh(-replace(dHdu, {test: dot(J.T, test)})) stepper = GalerkinTimeStepper(F, order, t, dt, u, solver_parameters=solver_parameters) - return stepper, [H] + return stepper, [H*dx] def kepler_aux_variable(V, order, t, dt, u0, solver_parameters): @@ -316,7 +319,7 @@ def kepler_aux_variable(V, order, t, dt, u0, solver_parameters): Qlow = create_quadrature(ufc_simplex(1), 2*order-2) Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) - Qhigh = create_quadrature(ufc_simplex(1), 25) + Qhigh = create_quadrature(ufc_simplex(1), 2*order+2) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) # determinant_forms = [test_u, dHdu, dA1du, dA2du] @@ -336,13 +339,53 @@ 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) + + p = as_vector([u[k] for k in range(dim)]) + q = as_vector([u[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 = derivative(H*dx, u, v) + dA1du = derivative(A1*dx, u, v) + dA2du = derivative(A2*dx, u, v) + + Qlow = create_quadrature(ufc_simplex(1), 2*order-2) + Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) + + Qhigh = create_quadrature(ufc_simplex(1), 2*order+2) + Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) + + w0 = TimeProjector(dHdu, order-1, Qhigh) + w1 = TimeProjector(dA1du, order-1, Qhigh) + w2 = TimeProjector(dA2du, order-1, Qhigh) + determinant_forms = [v, w0, w1, w2] + tensor = as_tensor(determinant_forms) + F = Llow(inner(Dt(u), v)*dx) + Lhigh(-(det(tensor) / (2*L*H))*dx) + + stepper = GalerkinTimeStepper(F, order, 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) t = MC.Constant(0.0) - dt = MC.Constant(pi/100) + dt = MC.Constant(0.01*pi) Nsteps = 2 dim = 2 From e9c3849adc8644fa3efb8f79a47eebef0bd4edee Mon Sep 17 00:00:00 2001 From: "Robert C. Kirby" Date: Thu, 29 May 2025 11:53:27 +0100 Subject: [PATCH 10/24] temp fix (?) on BC for Galerkin; WIP demo on structure preservation --- .../nse_helicity_energy.py | 186 ++++++++++++++++++ irksome/galerkin_stepper.py | 3 +- 2 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 demos/structure_preservation/nse_helicity_energy.py diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py new file mode 100644 index 00000000..37ccce3f --- /dev/null +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -0,0 +1,186 @@ +from firedrake import * +from irksome import Dt, GalerkinTimeStepper, TimeStepper, GaussLegendre +from irksome.galerkin_stepper import TimeProjector +from irksome.labeling import TimeQuadratureLabel +from FIAT import make_quadrature, ufc_simplex +import math +from scipy import special + +''' +Thanks to Boris Andrews for providing the Hill vortex functions +''' +''' +Hill vortex functions +''' +# UFL-compatible Bessel function +def besselJ(x, alpha, layers=10): + return sum([ + (-1)**m / math.factorial(m) / special.gamma(m + alpha + 1) + * (x/2)**(2*m+alpha) + for m in range(layers) + ]) + + + +# Bessel function parameters +besselJ_root = 5.7634591968945506 +besselJ_root_threehalves = besselJ(besselJ_root, 3/2) + + + +# (r, theta, phi) components of Hill vortex +def hill_r(r, theta, radius): + rho = r / radius + return 2 * ( + besselJ(besselJ_root*rho, 3/2) / rho**(3/2) + - besselJ_root_threehalves + ) * cos(theta) + +def hill_theta(r, theta, radius): + rho = r / radius + return ( + besselJ_root * besselJ(besselJ_root*rho, 5/2) / rho**(1/2) + + 2 * besselJ_root_threehalves + - 2 * besselJ(besselJ_root*rho, 3/2) / rho**(3/2) + ) * sin(theta) + +def hill_phi(r, theta, radius): + rho = r / radius + return besselJ_root * ( + besselJ(besselJ_root*rho, 3/2) / rho**(3/2) + - besselJ_root_threehalves + ) * rho * sin(theta) + + + +# Hill vortex (Cartesian) +def hill(vec, radius): + (x, y, z) = vec + + # Cylindrical/spherical coordinates + r_cyl = sqrt(x**2 + y**2) # Cylindrical radius + r_sph = sqrt(x**2 + y**2 + z**2) # Spherical radius + theta = conditional( # Spherical angle + le(r_cyl, 1e-13), + 0, + pi/2 - atan(z/r_cyl) + ) + + return conditional( # If we're outside the vortex... + ge(r_sph, radius), + as_vector([0, 0, 0]), + conditional( # If we're at the origin... + le(r_sph, 1e-13), + as_vector([0, 0, 2*((besselJ_root/2)**(3/2)/special.gamma(5/2) - besselJ_root_threehalves)]), + conditional( # If we're on the z axis... + le(r_cyl, 1e-13), + as_vector([0, 0, hill_r(r_sph, 0, radius)]), + as_vector( # Else... + hill_r(r_sph, theta, radius) * np.array([x, y, z]) / r_sph + + hill_theta(r_sph, theta, radius) * np.array([x*z, y*z, -r_cyl**2]) / r_sph / r_cyl + + hill_phi(r_sph, theta, radius) * np.array([-y, x, 0]) / r_cyl + ) + ) + ) + ) + + +def nse_naive(msh, order, t, dt, Re, solver_parameters=None): + hill_expr = hill(SpatialCoordinate(msh), 0.25) + V = VectorFunctionSpace(msh, "CG", 2) + W = FunctionSpace(msh, "CG", 1) + Z = V * W + up = Function(Z) + up.subfunctions[0].interpolate(hill_expr) + + v, w = TestFunctions(Z) + u, p = split(up) + + F = (inner(Dt(u), v) * dx + - inner(cross(u, curl(u)), v) * dx + + 1/Re * inner(grad(u), grad(v)) * dx + - inner(p, div(v)) * dx + + inner(div(u), w) * dx) + + bcs = DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary") + + stepper = GalerkinTimeStepper(F, 2, t, dt, up, bcs=bcs) + + Q1 = inner(u, u) * dx + Q2 = inner(u, curl(u)) * dx + + Q1s = [assemble(Q1)] + Q2s = [assemble(Q2)] + + print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") + while float(t) < 3 * 2**(-6): + stepper.advance() + t.assign(float(t) + float(dt)) + Q1s.append(assemble(Q1)) + Q2s.append(assemble(Q2)) + print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") + + return np.array(Q1s), np.array(Q2s) + + +def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): + hill_expr = hill(SpatialCoordinate(msh), 0.25) + V = VectorFunctionSpace(msh, "CG", 2) + W = FunctionSpace(msh, "CG", 1) + Z = V * W + up = Function(Z) + up.subfunctions[0].interpolate(hill_expr) + + v, w = TestFunctions(Z) + u, p = split(up) + + ufcline = ufc_simplex(1) + + Qhigh = make_quadrature(ufcline, 2*order+2) + Qlow = make_quadrature(ufcline, 2*order-2) + + Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) + Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) + + w1 = TimeProjector(u, order-1, Qhigh) + w2 = TimeProjector(curl(u), order-1, Qhigh) + + F = (inner(Dt(u), v) * dx + - inner(cross(w1, w2), v) * dx + + 1/Re * inner(grad(w1), grad(v)) * dx + - inner(p, div(v)) * dx + + inner(div(u), w) * dx) + + bcs = DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary") + + stepper = GalerkinTimeStepper(F, 2, t, dt, up, bcs=bcs) + + Q1 = inner(u, u) * dx + Q2 = inner(u, curl(u)) * dx + + Q1s = [assemble(Q1)] + Q2s = [assemble(Q2)] + + print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") + while float(t) < 3 * 2**(-6): + stepper.advance() + t.assign(float(t) + float(dt)) + Q1s.append(assemble(Q1)) + Q2s.append(assemble(Q2)) + print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") + + return np.array(Q1s), np.array(Q2s) + + + +N = 8 +msh = UnitCubeMesh(N, N, N) +t = Constant(0) +dt = Constant(2**-10) +msh.coordinates.dat.data[:, :] -= 0.5 +Re = Constant(2**8) +Q1s, Q2s = nse_project_both(msh, 2, t, dt, Re) +print(Q1s) +print(Q2s) + + diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 3e90c87d..300a74c6 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -161,7 +161,8 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a u0_sub = bc2space(bc, u0) g0 = as_ufl(bc._original_arg) Vg_np = np.array([replace(g0, {t: t + c * dt}) for c in qpts]) - Vg_np -= u0_sub * trial_vals[0] + for i in range(Vg_np.shape[0]): + Vg_np[i] -= u0_sub * trial_vals[0, i] g_np = proj @ Vg_np for i in range(num_stages): Vbigi = stage2spaces4bc(bc, V, Vbig, i) From e6c7750575662c542000d2d67f986d1dcd07bdfb Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 29 May 2025 13:49:41 +0100 Subject: [PATCH 11/24] TimeProjector(Expr) --- irksome/galerkin_stepper.py | 17 +++++++++-------- tests/test_galerkin.py | 11 ++++++----- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 3e90c87d..bfdb704b 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -15,7 +15,7 @@ from ufl.corealg.multifunction import MultiFunction from ufl.algorithms.map_integrands import map_integrand_dags from ufl.algorithms import expand_derivatives -from firedrake import Constant, Function, TestFunction, VectorFunctionSpace, diff +from firedrake import Constant, Function, TestFunction, TensorFunctionSpace, VectorFunctionSpace, diff, dx, inner ufc_line = ufc_simplex(1) @@ -289,14 +289,12 @@ def set_initial_guess(self): inherit_indices_from_operand=0, ) class TimeProjector(Operator): - __slots__ = ("form", "order", "quadrature") + __slots__ = ("order", "quadrature") - def __init__(self, F, order, Q): - self.form = expand_derivatives(F) + def __init__(self, expression, order, Q): self.order = order self.quadrature = Q - c = Coefficient(F.arguments()[0].function_space()) - Operator.__init__(self, operands=(c,)) + Operator.__init__(self, operands=(expression,)) class TimeProjectorDispatcher(MultiFunction): @@ -312,10 +310,13 @@ def __init__(self, element, t, dt, u0, u1, stages, phi): def time_projector(self, o): # use the internal copy of the state, so it does not get updated again in the outer quadrature - F = replace(o.form, {self.u0: self.u1}) order = o.order Q = o.quadrature assert order+1 <= len(self.phi) + f, = o.ufl_operands + R = TensorFunctionSpace(self.u0.function_space().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() @@ -325,7 +326,7 @@ def time_projector(self, o): Minv = vecconst(np.linalg.inv(M)) # compute modal expansion tested against c - c, = o.ufl_operands + c = Coefficient(R) test = np.outer(Minv[:order+1] @ self.phi, np.asarray(c, dtype=object)) / self.dt Fc = getTermGalerkin(F, self.L_trial, L_test, Q, self.t, self.dt, self.u1, self.stages, test) diff --git a/tests/test_galerkin.py b/tests/test_galerkin.py index 80c1c93d..b979adb8 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -344,8 +344,9 @@ def kepler_projector(V, order, t, dt, u0, solver_parameters): u = Function(V) u.interpolate(u0) - p = as_vector([u[k] for k in range(dim)]) - q = as_vector([u[k] for k in range(dim, 2*dim)]) + 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)) @@ -357,9 +358,9 @@ def kepler_projector(V, order, t, dt, u0, solver_parameters): invariants = [H*dx, L*dx, A1*dx, A2*dx] v = TestFunction(V) - dHdu = derivative(H*dx, u, v) - dA1du = derivative(A1*dx, u, v) - dA2du = derivative(A2*dx, u, v) + dHdu = diff(H, uv) + dA1du = diff(A1, uv) + dA2du = diff(A2, uv) Qlow = create_quadrature(ufc_simplex(1), 2*order-2) Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) From dee3e7f8d61478880adc181448c4dfdc3b133e63 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 29 May 2025 14:54:22 +0100 Subject: [PATCH 12/24] Add Labels --- .../nse_helicity_energy.py | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py index 37ccce3f..9cb6d5fc 100644 --- a/demos/structure_preservation/nse_helicity_energy.py +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -59,26 +59,25 @@ def hill(vec, radius): # Cylindrical/spherical coordinates r_cyl = sqrt(x**2 + y**2) # Cylindrical radius - r_sph = sqrt(x**2 + y**2 + z**2) # Spherical radius - theta = conditional( # Spherical angle - le(r_cyl, 1e-13), - 0, - pi/2 - atan(z/r_cyl) - ) + r_sph = sqrt(dot(vec, vec)) # Spherical radius + theta = pi/2 - atan2(z, r_cyl) + rvec = vec / r_sph + Rvec = as_vector([-y, x, 0]) / r_cyl + THvec = as_vector([x*z, y*z, -r_cyl**2]) / r_sph / r_cyl return conditional( # If we're outside the vortex... ge(r_sph, radius), - as_vector([0, 0, 0]), + zero((3,)), conditional( # If we're at the origin... le(r_sph, 1e-13), as_vector([0, 0, 2*((besselJ_root/2)**(3/2)/special.gamma(5/2) - besselJ_root_threehalves)]), conditional( # If we're on the z axis... le(r_cyl, 1e-13), as_vector([0, 0, hill_r(r_sph, 0, radius)]), - as_vector( # Else... - hill_r(r_sph, theta, radius) * np.array([x, y, z]) / r_sph - + hill_theta(r_sph, theta, radius) * np.array([x*z, y*z, -r_cyl**2]) / r_sph / r_cyl - + hill_phi(r_sph, theta, radius) * np.array([-y, x, 0]) / r_cyl + ( # Else... + hill_r(r_sph, theta, radius) * rvec + + hill_theta(r_sph, theta, radius) * THvec + + hill_phi(r_sph, theta, radius) * Rvec ) ) ) @@ -136,20 +135,21 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): ufcline = ufc_simplex(1) - Qhigh = make_quadrature(ufcline, 2*order+2) - Qlow = make_quadrature(ufcline, 2*order-2) + Qhigh = make_quadrature(ufcline, 3*order-1) + Qlow = make_quadrature(ufcline, 2*(order-1)) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) - w1 = TimeProjector(u, order-1, Qhigh) - w2 = TimeProjector(curl(u), order-1, Qhigh) + Qk = make_quadrature(ufcline, order) + w1 = TimeProjector(u, order-1, Qk) + w2 = TimeProjector(curl(u), order-1, Qk) - F = (inner(Dt(u), v) * dx - - inner(cross(w1, w2), v) * dx + F = (Llow(inner(Dt(u), v) * dx) + + Lhigh(- inner(cross(w1, w2), v) * dx) + ( + 1/Re * inner(grad(w1), grad(v)) * dx - inner(p, div(v)) * dx - + inner(div(u), w) * dx) + + inner(div(u), w) * dx)) bcs = DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary") @@ -179,7 +179,7 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): dt = Constant(2**-10) msh.coordinates.dat.data[:, :] -= 0.5 Re = Constant(2**8) -Q1s, Q2s = nse_project_both(msh, 2, t, dt, Re) +Q1s, Q2s = nse_naive(msh, 2, t, dt, Re) print(Q1s) print(Q2s) From eaa1474cb257de630cb55e06803357bb5e7092d7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 29 May 2025 15:04:08 +0100 Subject: [PATCH 13/24] minor change --- demos/structure_preservation/nse_helicity_energy.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py index 9cb6d5fc..6050be63 100644 --- a/demos/structure_preservation/nse_helicity_energy.py +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -103,7 +103,7 @@ def nse_naive(msh, order, t, dt, Re, solver_parameters=None): bcs = DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary") - stepper = GalerkinTimeStepper(F, 2, t, dt, up, bcs=bcs) + stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs) Q1 = inner(u, u) * dx Q2 = inner(u, curl(u)) * dx @@ -135,7 +135,7 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): ufcline = ufc_simplex(1) - Qhigh = make_quadrature(ufcline, 3*order-1) + Qhigh = make_quadrature(ufcline, 3*(order-1)) Qlow = make_quadrature(ufcline, 2*(order-1)) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) @@ -145,15 +145,14 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): w1 = TimeProjector(u, order-1, Qk) w2 = TimeProjector(curl(u), order-1, Qk) - F = (Llow(inner(Dt(u), v) * dx) + - Lhigh(- inner(cross(w1, w2), v) * dx) + ( - + 1/Re * inner(grad(w1), grad(v)) * dx + F = (Llow(inner(Dt(u), v) * dx + 1/Re * inner(grad(w1), grad(v)) * dx) + + Lhigh(-inner(cross(w1, w2), v) * dx) + ( - inner(p, div(v)) * dx + inner(div(u), w) * dx)) bcs = DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary") - stepper = GalerkinTimeStepper(F, 2, t, dt, up, bcs=bcs) + stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs) Q1 = inner(u, u) * dx Q2 = inner(u, curl(u)) * dx @@ -179,7 +178,7 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): dt = Constant(2**-10) msh.coordinates.dat.data[:, :] -= 0.5 Re = Constant(2**8) -Q1s, Q2s = nse_naive(msh, 2, t, dt, Re) +Q1s, Q2s = nse_project_both(msh, 1, t, dt, Re) print(Q1s) print(Q2s) From 3a70634956931302cba017151273f78d5c6315da Mon Sep 17 00:00:00 2001 From: "Robert C. Kirby" Date: Thu, 29 May 2025 15:07:44 +0100 Subject: [PATCH 14/24] fix quadrature degree --- demos/structure_preservation/nse_helicity_energy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py index 6050be63..202ec260 100644 --- a/demos/structure_preservation/nse_helicity_energy.py +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -135,8 +135,8 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): ufcline = ufc_simplex(1) - Qhigh = make_quadrature(ufcline, 3*(order-1)) - Qlow = make_quadrature(ufcline, 2*(order-1)) + Qhigh = make_quadrature(ufcline, 3*order-1) + Qlow = make_quadrature(ufcline, 2*order-1) Lhigh = TimeQuadratureLabel(Qhigh.get_points(), Qhigh.get_weights()) Llow = TimeQuadratureLabel(Qlow.get_points(), Qlow.get_weights()) From 3e85826525562b92e1e22c252b051be90f1253ea Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 30 May 2025 11:42:14 +0100 Subject: [PATCH 15/24] NSE + auxiliary variables --- .../nse_helicity_energy.py | 118 ++++++++++++++---- irksome/galerkin_stepper.py | 11 +- 2 files changed, 98 insertions(+), 31 deletions(-) diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py index 202ec260..4da3882b 100644 --- a/demos/structure_preservation/nse_helicity_energy.py +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -2,7 +2,8 @@ from irksome import Dt, GalerkinTimeStepper, TimeStepper, GaussLegendre from irksome.galerkin_stepper import TimeProjector from irksome.labeling import TimeQuadratureLabel -from FIAT import make_quadrature, ufc_simplex +from FIAT import ufc_simplex +from FIAT.quadrature_schemes import create_quadrature import math from scipy import special @@ -58,26 +59,27 @@ def hill(vec, radius): (x, y, z) = vec # Cylindrical/spherical coordinates - r_cyl = sqrt(x**2 + y**2) # Cylindrical radius - r_sph = sqrt(dot(vec, vec)) # Spherical radius - theta = pi/2 - atan2(z, r_cyl) - rvec = vec / r_sph - Rvec = as_vector([-y, x, 0]) / r_cyl - THvec = as_vector([x*z, y*z, -r_cyl**2]) / r_sph / r_cyl + rho = sqrt(x*x + y*y) # Cylindrical radius + r = sqrt(dot(vec, vec)) # Spherical radius + theta = pi/2 - atan2(z, rho) + + r_dir = vec / r + theta_dir = as_vector([x*z/rho, y*z/rho, -rho/r]) + phi_dir = as_vector([-y, x, 0]) / rho return conditional( # If we're outside the vortex... - ge(r_sph, radius), + ge(r, radius), zero((3,)), conditional( # If we're at the origin... - le(r_sph, 1e-13), + le(r, 1e-13), as_vector([0, 0, 2*((besselJ_root/2)**(3/2)/special.gamma(5/2) - besselJ_root_threehalves)]), conditional( # If we're on the z axis... - le(r_cyl, 1e-13), - as_vector([0, 0, hill_r(r_sph, 0, radius)]), + le(rho, 1e-13), + as_vector([0, 0, hill_r(r, 0, radius)]), ( # Else... - hill_r(r_sph, theta, radius) * rvec - + hill_theta(r_sph, theta, radius) * THvec - + hill_phi(r_sph, theta, radius) * Rvec + hill_r(r, theta, radius) * r_dir + + hill_theta(r, theta, radius) * theta_dir + + hill_phi(r, theta, radius) * phi_dir ) ) ) @@ -112,7 +114,8 @@ def nse_naive(msh, order, t, dt, Re, solver_parameters=None): Q2s = [assemble(Q2)] print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - while float(t) < 3 * 2**(-6): + #while float(t) < 3 * 2**(-6): + for k in range(2): stepper.advance() t.assign(float(t) + float(dt)) Q1s.append(assemble(Q1)) @@ -135,18 +138,19 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): ufcline = ufc_simplex(1) - Qhigh = make_quadrature(ufcline, 3*order-1) - Qlow = make_quadrature(ufcline, 2*order-1) + 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()) - Qk = make_quadrature(ufcline, order) - w1 = TimeProjector(u, order-1, Qk) - w2 = TimeProjector(curl(u), order-1, Qk) + Qproj = create_quadrature(ufcline, 2*order-1) + w1 = TimeProjector(u, order-1, Qproj) + w2 = TimeProjector(curl(u), order-1, Qproj) - F = (Llow(inner(Dt(u), v) * dx + 1/Re * inner(grad(w1), grad(v)) * dx) + + F = (Llow(inner(Dt(u), v) * dx) + Lhigh(-inner(cross(w1, w2), v) * dx) + ( + 1/Re * inner(grad(w1), grad(v)) * dx - inner(p, div(v)) * dx + inner(div(u), w) * dx)) @@ -161,7 +165,8 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): Q2s = [assemble(Q2)] print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - while float(t) < 3 * 2**(-6): + #while float(t) < 3 * 2**(-6): + for k in range(2): stepper.advance() t.assign(float(t) + float(dt)) Q1s.append(assemble(Q1)) @@ -171,15 +176,78 @@ def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): return np.array(Q1s), np.array(Q2s) +def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None): + hill_expr = hill(SpatialCoordinate(msh), 0.25) + V = VectorFunctionSpace(msh, "CG", 2) + W = FunctionSpace(msh, "CG", 1) + Z = V * V * V * W + up = Function(Z) + up.subfunctions[0].interpolate(hill_expr) + + v, v1, v2, q = TestFunctions(Z) + u, w1, w2, p = split(up) + + ufcline = ufc_simplex(1) + + 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) + ( + 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 + )) + + bcs = [DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary")] + + stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs, aux_indices=(1, 2)) + Q1 = inner(u, u) * dx + Q2 = inner(u, curl(u)) * dx + + Q1s = [assemble(Q1)] + Q2s = [assemble(Q2)] + + print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") + #while float(t) < 3 * 2**(-6): + for k in range(2): + stepper.advance() + t.assign(float(t) + float(dt)) + Q1s.append(assemble(Q1)) + Q2s.append(assemble(Q2)) + print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") + + return np.array(Q1s), np.array(Q2s) + + + +order = 2 N = 8 msh = UnitCubeMesh(N, N, N) t = Constant(0) dt = Constant(2**-10) msh.coordinates.dat.data[:, :] -= 0.5 Re = Constant(2**8) -Q1s, Q2s = nse_project_both(msh, 1, t, dt, Re) -print(Q1s) -print(Q2s) +print("naive") +t.assign(0) +Q1s, Q2s = nse_naive(msh, order, t, dt, Re) + +print("aux") +t.assign(0) +Q1s, Q2s = nse_aux_variable(msh, order, t, dt, Re) + + + +print("project") +t.assign(0) +Q1s, Q2s = nse_project_both(msh, order, t, dt, Re) +#print(Q1s) +#print(Q2s) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index be09e84b..341be1bd 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -1,7 +1,6 @@ from FIAT import (Bernstein, DiscontinuousElement, DiscontinuousLagrange, IntegratedLegendre, Lagrange, Legendre, ufc_simplex) from FIAT.quadrature_schemes import create_quadrature -from ufl.constantvalue import as_ufl from .base_time_stepper import StageCoupledTimeStepper from .bcs import bc2space, stage2spaces4bc from .deriv import TimeDerivative, expand_time_derivatives @@ -9,7 +8,7 @@ from .tools import replace, vecconst, replace_auxiliary_variables import numpy as np -from ufl import Coefficient +from ufl import as_tensor, as_ufl, Coefficient from ufl.core.operator import Operator from ufl.core.ufl_type import ufl_type from ufl.corealg.multifunction import MultiFunction @@ -160,13 +159,13 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a for bc in bcs: u0_sub = bc2space(bc, u0) g0 = as_ufl(bc._original_arg) - Vg_np = np.array([replace(g0, {t: t + c * dt}) for c in qpts]) - for i in range(Vg_np.shape[0]): - Vg_np[i] -= u0_sub * trial_vals[0, i] + Vg_np = np.array([replace(g0, {t: t + c * dt}) - u0_sub * phi + for c, phi in zip(qpts, trial_vals[0])]) g_np = proj @ Vg_np for i in range(num_stages): + gcur = as_tensor(g_np[i]) Vbigi = stage2spaces4bc(bc, V, Vbig, i) - bcsnew.append(bc.reconstruct(V=Vbigi, g=g_np[i])) + bcsnew.append(bc.reconstruct(V=Vbigi, g=gcur)) return Fnew, bcsnew From 569a4ead40b26572b57f3860c6d35de2ced617ee Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 30 May 2025 19:16:24 +0100 Subject: [PATCH 16/24] Galerkin solve for an additive update --- irksome/galerkin_stepper.py | 50 ++++++++++--------------------------- 1 file changed, 13 insertions(+), 37 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 341be1bd..0ce5256f 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -74,7 +74,7 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices= # set up the pieces we need to work with to do our substitutions v_np = np.reshape(test, (-1, *v.ufl_shape)) w_np = np.reshape(stages, (-1, *u0.ufl_shape)) - u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) + u_np = np.concatenate((vecconst(np.zeros((1, *u0.ufl_shape))), w_np)) vsub = test_vals_w.T @ v_np usub = trial_vals.T @ u_np @@ -86,7 +86,7 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices= for q in range(len(qpts)): repl[q] = {t: t + qpts[q] * dt, v: vsub[q] * dt, - u0: usub[q], + u0: u0 + usub[q], dtu0: dtu0sub[q] / dt, u1: u0, phi: phisub[q]} @@ -157,15 +157,18 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a bcs = [] bcsnew = [] for bc in bcs: - u0_sub = bc2space(bc, u0) - g0 = as_ufl(bc._original_arg) - Vg_np = np.array([replace(g0, {t: t + c * dt}) - u0_sub * phi - for c, phi in zip(qpts, trial_vals[0])]) - g_np = proj @ Vg_np + if bc._original_arg == 0: + gbig = [bc._original_arg] * num_stages + else: + g0 = as_ufl(bc._original_arg) + u0_sub = bc2space(bc, u0) + Vg_np = np.array([replace(g0, {t: t + q * dt}) - u0_sub * phi + for q, phi in zip(qpts, trial_vals[0])]) + gbig = as_tensor(proj @ Vg_np) + for i in range(num_stages): - gcur = as_tensor(g_np[i]) Vbigi = stage2spaces4bc(bc, V, Vbig, i) - bcsnew.append(bc.reconstruct(V=Vbigi, g=gcur)) + bcsnew.append(bc.reconstruct(V=Vbigi, g=gbig[i])) return Fnew, bcsnew @@ -233,7 +236,6 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None, self.aux_indices = aux_indices super().__init__(F, t, dt, u0, order, bcs=bcs, **kwargs) - self.set_initial_guess() def get_form_and_bcs(self, stages, basis_type=None, order=None, quadrature=None, aux_indices=None, F=None): if basis_type is None: @@ -254,33 +256,7 @@ def get_form_and_bcs(self, stages, basis_type=None, order=None, quadrature=None, def _update(self): k1, = self.trial_el.entity_dofs()[0][1] for i, u0bit in enumerate(self.u0.subfunctions): - u0bit.assign(self.stages.subfunctions[self.num_fields*(k1-1)+i]) - - def set_initial_guess(self): - # Set a constant-in-time initial guess - ref_el = self.test_el.get_reference_element() - P0 = DiscontinuousLagrange(ref_el, 0) - P0 = P0.get_nodal_basis() - B = P0.get_coeffs() - - test_dual = self.test_el.get_dual_set() - test_dofs = np.dot(test_dual.to_riesz(P0), B) - - trial_dual = self.trial_el.get_dual_set() - trial_dofs = np.dot(trial_dual.to_riesz(P0), B) - - dof = Constant(0) - for k in range(self.num_stages): - for i, u0bit in enumerate(self.u0.subfunctions): - sbit = self.stages.subfunctions[self.num_fields*k+i] - if self.aux_indices and i in self.aux_indices: - dof.assign(test_dofs[k]) - else: - dof.assign(trial_dofs[k+1]) - if abs(float(dof)) < 1E-12: - sbit.zero() - else: - sbit.assign(u0bit * dof) + u0bit.assign(u0bit + self.stages.subfunctions[self.num_fields*(k1-1)+i]) @ufl_type( From 0544bbad0fde56fcc1ee32ade9dc5575fd642aed Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 30 May 2025 19:16:32 +0100 Subject: [PATCH 17/24] Cleanup test --- .../nse_helicity_energy.py | 235 ++++++++---------- 1 file changed, 104 insertions(+), 131 deletions(-) diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py index 4da3882b..87dabe89 100644 --- a/demos/structure_preservation/nse_helicity_energy.py +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -6,6 +6,11 @@ from FIAT.quadrature_schemes import create_quadrature import math 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 @@ -14,7 +19,9 @@ Hill vortex functions ''' # UFL-compatible Bessel function -def besselJ(x, alpha, layers=10): +besselJ = bessel_J + +def besselJ(alpha, x, layers=10): return sum([ (-1)**m / math.factorial(m) / special.gamma(m + alpha + 1) * (x/2)**(2*m+alpha) @@ -22,50 +29,49 @@ def besselJ(x, alpha, layers=10): ]) - # Bessel function parameters besselJ_root = 5.7634591968945506 -besselJ_root_threehalves = besselJ(besselJ_root, 3/2) - +besselJ_root_threehalves = besselJ(3/2, besselJ_root) # (r, theta, phi) components of Hill vortex def hill_r(r, theta, radius): rho = r / radius return 2 * ( - besselJ(besselJ_root*rho, 3/2) / rho**(3/2) + besselJ(3/2, besselJ_root*rho) / rho**(3/2) - besselJ_root_threehalves ) * cos(theta) + def hill_theta(r, theta, radius): rho = r / radius return ( - besselJ_root * besselJ(besselJ_root*rho, 5/2) / rho**(1/2) + besselJ_root * besselJ(5/2, besselJ_root*rho) / rho**(1/2) + 2 * besselJ_root_threehalves - - 2 * besselJ(besselJ_root*rho, 3/2) / rho**(3/2) + - 2 * besselJ(3/2, besselJ_root*rho) / rho**(3/2) ) * sin(theta) + def hill_phi(r, theta, radius): rho = r / radius return besselJ_root * ( - besselJ(besselJ_root*rho, 3/2) / rho**(3/2) + besselJ(3/2, besselJ_root*rho) / rho**(3/2) - besselJ_root_threehalves ) * rho * sin(theta) - # Hill vortex (Cartesian) def hill(vec, radius): (x, y, z) = vec + rho = sqrt(x*x + y*y) - # Cylindrical/spherical coordinates - rho = sqrt(x*x + y*y) # Cylindrical radius - r = sqrt(dot(vec, vec)) # Spherical radius - theta = pi/2 - atan2(z, rho) + r = sqrt(dot(vec, vec)) + theta = pi/2-atan2(z, rho) + phi = atan2(y, x) - r_dir = vec / r - theta_dir = as_vector([x*z/rho, y*z/rho, -rho/r]) - phi_dir = as_vector([-y, x, 0]) / rho + 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), @@ -73,181 +79,148 @@ def hill(vec, radius): conditional( # If we're at the origin... le(r, 1e-13), as_vector([0, 0, 2*((besselJ_root/2)**(3/2)/special.gamma(5/2) - besselJ_root_threehalves)]), - conditional( # If we're on the z axis... - le(rho, 1e-13), - as_vector([0, 0, hill_r(r, 0, radius)]), - ( # Else... - hill_r(r, theta, radius) * r_dir - + hill_theta(r, theta, radius) * theta_dir - + hill_phi(r, theta, radius) * phi_dir - ) - ) + (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 = VectorFunctionSpace(msh, "CG", 2) - W = FunctionSpace(msh, "CG", 1) - Z = V * W + 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, w = TestFunctions(Z) + v, q = TestFunctions(Z) u, p = split(up) - F = (inner(Dt(u), v) * dx - - inner(cross(u, curl(u)), v) * dx + 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), w) * dx) + - inner(div(u), q) * dx) - bcs = DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary") + 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 - Q1s = [assemble(Q1)] - Q2s = [assemble(Q2)] - - print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - #while float(t) < 3 * 2**(-6): - for k in range(2): - stepper.advance() - t.assign(float(t) + float(dt)) - Q1s.append(assemble(Q1)) - Q2s.append(assemble(Q2)) - print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - - return np.array(Q1s), np.array(Q2s) - -def nse_project_both(msh, order, t, dt, Re, solver_parameters=None): +def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None): hill_expr = hill(SpatialCoordinate(msh), 0.25) - V = VectorFunctionSpace(msh, "CG", 2) - W = FunctionSpace(msh, "CG", 1) - Z = V * W + V, Q = stokes_pair(msh) + Z = V * V * V * Q up = Function(Z) up.subfunctions[0].interpolate(hill_expr) - v, w = TestFunctions(Z) - u, p = split(up) - - ufcline = ufc_simplex(1) + v, v1, v2, q = TestFunctions(Z) + u, w1, w2, 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()) - Qproj = create_quadrature(ufcline, 2*order-1) - w1 = TimeProjector(u, order-1, Qproj) - w2 = TimeProjector(curl(u), order-1, Qproj) - - F = (Llow(inner(Dt(u), v) * dx) + - Lhigh(-inner(cross(w1, w2), v) * dx) + ( - 1/Re * inner(grad(w1), grad(v)) * dx + 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), w) * dx)) + - inner(div(u), q) * dx + + inner(w1 - u, v1) * dx + + inner(w2 - curl(u), v2) * dx + ) - bcs = DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary") - - stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs) + bcs = [DirichletBC(Z.sub(i), 0, "on_boundary") for i in range(len(Z)-1)] + stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs, aux_indices=(1, 2)) Q1 = inner(u, u) * dx Q2 = inner(u, curl(u)) * dx + invariants = [Q1, Q2] + return stepper, invariants - Q1s = [assemble(Q1)] - Q2s = [assemble(Q2)] - - print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - #while float(t) < 3 * 2**(-6): - for k in range(2): - stepper.advance() - t.assign(float(t) + float(dt)) - Q1s.append(assemble(Q1)) - Q2s.append(assemble(Q2)) - print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - - return np.array(Q1s), np.array(Q2s) - -def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None): +def nse_project(msh, order, t, dt, Re, solver_parameters=None): hill_expr = hill(SpatialCoordinate(msh), 0.25) - V = VectorFunctionSpace(msh, "CG", 2) - W = FunctionSpace(msh, "CG", 1) - Z = V * V * V * W + V, Q = stokes_pair(msh) + Z = V * Q up = Function(Z) up.subfunctions[0].interpolate(hill_expr) - v, v1, v2, q = TestFunctions(Z) - u, w1, w2, p = split(up) - - ufcline = ufc_simplex(1) + 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(w1, w2), v) * dx) + ( - 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 - )) + Qproj = create_quadrature(ufcline, 2*order) + w1 = TimeProjector(u, order-1, Qproj) + w2 = TimeProjector(curl(u), order-1, Qproj) - bcs = [DirichletBC(Z.sub(0), Constant((0, 0, 0)), "on_boundary")] + 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 + ) - stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs, aux_indices=(1, 2)) + 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 - Q1s = [assemble(Q1)] - Q2s = [assemble(Q2)] - print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - #while float(t) < 3 * 2**(-6): - for k in range(2): +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(float(t) + float(dt)) - Q1s.append(assemble(Q1)) - Q2s.append(assemble(Q2)) - print(f"{float(t):.4e}, {assemble(Q1):.4e}, {assemble(Q2):.4e}") - - return np.array(Q1s), np.array(Q2s) + 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) -t = Constant(0) -dt = Constant(2**-10) msh.coordinates.dat.data[:, :] -= 0.5 -Re = Constant(2**8) -print("naive") -t.assign(0) -Q1s, Q2s = nse_naive(msh, order, t, dt, Re) - -print("aux") -t.assign(0) -Q1s, Q2s = nse_aux_variable(msh, order, t, dt, Re) - - - -print("project") -t.assign(0) -Q1s, Q2s = nse_project_both(msh, order, t, dt, Re) -#print(Q1s) -#print(Q2s) - +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)) From c63a366100489d55315978af95ee7095ad8e483e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 2 Jun 2025 10:53:59 +0100 Subject: [PATCH 18/24] Eliminate w1 only --- .../nse_helicity_energy.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py index 87dabe89..2baeef06 100644 --- a/demos/structure_preservation/nse_helicity_energy.py +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -148,8 +148,9 @@ def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None): ) bcs = [DirichletBC(Z.sub(i), 0, "on_boundary") for i in range(len(Z)-1)] + aux_indices = tuple(range(1, len(Z)-1)) - stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs, aux_indices=(1, 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] @@ -159,32 +160,34 @@ def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None): 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 * Q + Z = V * V * Q up = Function(Z) up.subfunctions[0].interpolate(hill_expr) - v, q = TestFunctions(Z) - u, p = split(up) + v, v2, q = TestFunctions(Z) + u, w2, 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()) - Qproj = create_quadrature(ufcline, 2*order) + # Eliminate w1 only + Qproj = create_quadrature(ufcline, 2*order-1) w1 = TimeProjector(u, order-1, Qproj) - w2 = TimeProjector(curl(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 ) - bcs = DirichletBC(Z.sub(0), 0, "on_boundary") + bcs = [DirichletBC(Z.sub(i), 0, "on_boundary") for i in range(len(Z)-1)] + aux_indices = tuple(range(1, len(Z)-1)) - stepper = GalerkinTimeStepper(F, order, t, dt, up, bcs=bcs) + 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] @@ -215,9 +218,9 @@ def run_nse(stepper, invariants): Re = Constant(2**16) solvers = { - #"naive": nse_naive, + "naive": nse_naive, "project": nse_project, - #"aux": nse_aux_variable, + "aux": nse_aux_variable, } for name, solver in solvers.items(): From 384997cbabc3296bd306a0bc01569554a6f77c5f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 5 Jun 2025 16:18:22 +0100 Subject: [PATCH 19/24] Add more axiliary variables --- .../nse_helicity_energy.py | 64 ++++++++----------- irksome/galerkin_stepper.py | 2 +- 2 files changed, 29 insertions(+), 37 deletions(-) diff --git a/demos/structure_preservation/nse_helicity_energy.py b/demos/structure_preservation/nse_helicity_energy.py index 2baeef06..843c80e1 100644 --- a/demos/structure_preservation/nse_helicity_energy.py +++ b/demos/structure_preservation/nse_helicity_energy.py @@ -1,10 +1,9 @@ from firedrake import * -from irksome import Dt, GalerkinTimeStepper, TimeStepper, GaussLegendre +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 -import math from scipy import special from firedrake.petsc import PETSc @@ -18,45 +17,34 @@ ''' Hill vortex functions ''' -# UFL-compatible Bessel function -besselJ = bessel_J - -def besselJ(alpha, x, layers=10): - return sum([ - (-1)**m / math.factorial(m) / special.gamma(m + alpha + 1) - * (x/2)**(2*m+alpha) - for m in range(layers) - ]) - - # Bessel function parameters -besselJ_root = 5.7634591968945506 -besselJ_root_threehalves = besselJ(3/2, besselJ_root) +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 * ( - besselJ(3/2, besselJ_root*rho) / rho**(3/2) - - besselJ_root_threehalves + 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 ( - besselJ_root * besselJ(5/2, besselJ_root*rho) / rho**(1/2) - + 2 * besselJ_root_threehalves - - 2 * besselJ(3/2, besselJ_root*rho) / rho**(3/2) + 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 besselJ_root * ( - besselJ(3/2, besselJ_root*rho) / rho**(3/2) - - besselJ_root_threehalves + return bessel_J_root * ( + bessel_J(3/2, bessel_J_root*rho) / rho**(3/2) + - bessel_J_root_threehalves ) * rho * sin(theta) @@ -78,7 +66,7 @@ def hill(vec, radius): zero((3,)), conditional( # If we're at the origin... le(r, 1e-13), - as_vector([0, 0, 2*((besselJ_root/2)**(3/2)/special.gamma(5/2) - besselJ_root_threehalves)]), + 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) @@ -126,12 +114,13 @@ def nse_naive(msh, order, t, dt, Re, solver_parameters=None): 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 + Z = V * V * V * Q * Q up = Function(Z) up.subfunctions[0].interpolate(hill_expr) - v, v1, v2, q = TestFunctions(Z) - u, w1, w2, p = split(up) + 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)) @@ -145,10 +134,11 @@ def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None): - 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(len(Z)-1)] - aux_indices = tuple(range(1, len(Z)-1)) + 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 @@ -160,12 +150,13 @@ def nse_aux_variable(msh, order, t, dt, Re, solver_parameters=None): 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 + Z = V * V * Q * Q up = Function(Z) up.subfunctions[0].interpolate(hill_expr) - v, v2, q = TestFunctions(Z) - u, w2, p = split(up) + 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)) @@ -182,10 +173,11 @@ def nse_project(msh, order, t, dt, Re, solver_parameters=None): - 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(len(Z)-1)] - aux_indices = tuple(range(1, len(Z)-1)) + 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 @@ -199,8 +191,8 @@ def run_nse(stepper, invariants): 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): + # while float(t) < 3 * 2**(-6): + for k in range(2): stepper.advance() t.assign(t + dt) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 0ce5256f..0e327360 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -14,7 +14,7 @@ from ufl.corealg.multifunction import MultiFunction from ufl.algorithms.map_integrands import map_integrand_dags from ufl.algorithms import expand_derivatives -from firedrake import Constant, Function, TestFunction, TensorFunctionSpace, VectorFunctionSpace, diff, dx, inner +from firedrake import Function, TestFunction, TensorFunctionSpace, VectorFunctionSpace, diff, dx, inner ufc_line = ufc_simplex(1) From 1d284665856f0f6fc8ccffd94b3883420a2309e0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 8 Nov 2025 21:33:08 +0000 Subject: [PATCH 20/24] Apply suggestions from code review --- irksome/galerkin_stepper.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index c01bc085..b9f69db5 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -185,18 +185,14 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a bcs = [] bcsnew = [] for bc in bcs: - if bc._original_arg == 0: - gbig = [bc._original_arg] * num_stages - else: - g0 = as_ufl(bc._original_arg) - u0_sub = bc2space(bc, u0) - Vg_np = np.array([replace(g0, {t: t + q * dt}) - u0_sub * phi - for q, phi in zip(qpts, trial_vals[0])]) - gbig = as_tensor(proj @ Vg_np) - + u0_sub = bc2space(bc, u0) + g0 = as_ufl(bc._original_arg) + Vg_np = np.array([replace(g0, {t: t + c * dt}) for c in qpts]) + Vg_np -= u0_sub * trial_vals[0] + g_np = proj @ Vg_np for i in range(num_stages): Vbigi = stage2spaces4bc(bc, V, Vbig, i) - bcsnew.append(bc.reconstruct(V=Vbigi, g=gbig[i])) + bcsnew.append(bc.reconstruct(V=Vbigi, g=g_np[i])) return Fnew, bcsnew From 471c55e55df66f688c34c90f736ec56f2d42221f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 8 Nov 2025 21:33:35 +0000 Subject: [PATCH 21/24] Update irksome/galerkin_stepper.py --- irksome/galerkin_stepper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index b9f69db5..47daad80 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -187,9 +187,9 @@ def getFormGalerkin(F, L_trial, L_test, Qdefault, t, dt, u0, stages, bcs=None, a for bc in bcs: u0_sub = bc2space(bc, u0) g0 = as_ufl(bc._original_arg) - Vg_np = np.array([replace(g0, {t: t + c * dt}) for c in qpts]) - Vg_np -= u0_sub * trial_vals[0] - g_np = proj @ Vg_np + Vg_np = np.array([replace(g0, {t: t + c * dt}) for c in qpts]) + Vg_np -= u0_sub * trial_vals[0] + g_np = proj @ Vg_np for i in range(num_stages): Vbigi = stage2spaces4bc(bc, V, Vbig, i) bcsnew.append(bc.reconstruct(V=Vbigi, g=g_np[i])) From 540d8ca5aefe9955525e0f46b21c85fc3cdbe0d0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 10 Nov 2025 11:11:48 +0000 Subject: [PATCH 22/24] Fix MeshSequence, fix shape --- irksome/galerkin_stepper.py | 13 ++++++++----- irksome/tools.py | 11 ++--------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index 47daad80..cde9dd7b 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -8,10 +8,11 @@ from .scheme import create_time_quadrature, ufc_line from .tools import dot, reshape, replace, vecconst -from ufl import as_tensor, Coefficient +from ufl import as_tensor, outer, Coefficient from ufl.core.operator import Operator from ufl.core.ufl_type import ufl_type from ufl.corealg.multifunction import MultiFunction +from ufl.domain import as_domain from ufl.algorithms.map_integrands import map_integrand_dags from ufl.algorithms import expand_derivatives from firedrake import Constant, Function, TestFunction, TensorFunctionSpace, VectorFunctionSpace, diff, dx, inner @@ -61,7 +62,7 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices= # internal state to be used inside projected expressions u1 = Function(u0) # symbolic Coefficient with the temporal test function - mesh = u0.function_space().mesh() + mesh = as_domain(u0.function_space().mesh()) R = VectorFunctionSpace(mesh, "Real", 0, dim=L_test.space_dimension()) phi = Coefficient(R) # apply time projectors @@ -88,7 +89,7 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices= qpts = vecconst(np.reshape(qpts, (-1,))) # set up the pieces we need to work with to do our substitutions - v_np = reshape(test, (-1, *u0.ufl_shape)) + v_np = reshape(test, (-1, *v.ufl_shape)) w_np = reshape(stages, (-1, *u0.ufl_shape)) u_np = np.concatenate((np.reshape(u0, (1, *u0.ufl_shape)), w_np)) @@ -352,7 +353,8 @@ def time_projector(self, o): Q = o.quadrature assert order+1 <= len(self.phi) f, = o.ufl_operands - R = TensorFunctionSpace(self.u0.function_space().mesh(), "DG", 0, shape=f.ufl_shape) + 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}) @@ -365,7 +367,8 @@ def time_projector(self, o): # compute modal expansion tested against c c = Coefficient(R) - test = as_tensor(np.outer(Minv[:order+1] @ self.phi, np.asarray(c, dtype=object)) / self.dt) + + 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 diff --git a/irksome/tools.py b/irksome/tools.py index 5b7a9e3d..92931d6b 100644 --- a/irksome/tools.py +++ b/irksome/tools.py @@ -3,6 +3,7 @@ import numpy from firedrake import Function, FunctionSpace, VectorSpaceBasis, MixedVectorSpaceBasis, Constant from ufl.algorithms.analysis import extract_type +from ufl.domain import as_domain from ufl import as_tensor, zero from ufl import replace as ufl_replace from pyop2.types import MixedDat @@ -10,14 +11,6 @@ from irksome.deriv import TimeDerivative -def unique_mesh(mesh): - try: - mesh, = set(mesh) - except TypeError: - pass - return mesh - - def dot(A, B): return numpy.tensordot(A, B, (-1, 0)) @@ -111,7 +104,7 @@ def is_ode(f, u): # Utility class for constants on a mesh class MeshConstant(object): def __init__(self, msh): - self.msh = unique_mesh(msh) + self.msh = as_domain(msh) self.V = FunctionSpace(self.msh, 'R', 0) def Constant(self, val=0.0): From c24081bf7d5420d715e111db4d79c4122f3852b5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 3 Mar 2026 12:33:18 +0000 Subject: [PATCH 23/24] fixes --- irksome/__init__.py | 5 ++- irksome/galerkin_stepper.py | 74 ++------------------------------- irksome/ufl/estimate_degrees.py | 5 +++ tests/test_galerkin.py | 9 ++-- 4 files changed, 16 insertions(+), 77 deletions(-) diff --git a/irksome/__init__.py b/irksome/__init__.py index a8e95ddd..84bba592 100644 --- a/irksome/__init__.py +++ b/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", ] @@ -70,7 +72,7 @@ RanaLD, IRKAuxiliaryOperatorPC, ) - from .galerkin_stepper import ContinuousPetrovGalerkinTimeStepper, TimeProjector + from .galerkin_stepper import ContinuousPetrovGalerkinTimeStepper from .discontinuous_galerkin_stepper import DiscontinuousGalerkinTimeStepper from .labeling import TimeQuadratureLabel from .stepper import TimeStepper @@ -95,7 +97,6 @@ "StageValueTimeStepper", "ContinuousPetrovGalerkinTimeStepper", "DiscontinuousGalerkinTimeStepper", - "TimeProjector", "TimeQuadratureLabel", "TimeStepper", ] diff --git a/irksome/galerkin_stepper.py b/irksome/galerkin_stepper.py index a531a529..9f9ea6d9 100644 --- a/irksome/galerkin_stepper.py +++ b/irksome/galerkin_stepper.py @@ -2,17 +2,13 @@ GaussRadau, IntegratedLegendre, Lagrange, Legendre, NodalEnrichedElement, RestrictedElement) from ufl.classes import Zero -from ufl import as_ufl, as_tensor, outer, Coefficient -from ufl.core.operator import Operator -from ufl.core.ufl_type import ufl_type -from ufl.corealg.multifunction import MultiFunction +from ufl import as_ufl, as_tensor, Coefficient from ufl.domain import as_domain -from ufl.algorithms.map_integrands import map_integrand_dags -from ufl.algorithms import expand_derivatives 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 @@ -26,7 +22,7 @@ from .stage_value import getFormStage import numpy as np -from firedrake import TestFunction, Constant, Function, TensorFunctionSpace, VectorFunctionSpace, diff, dx, inner +from firedrake import TestFunction, Constant, Function, VectorFunctionSpace from firedrake.bcs import EquationBCSplit @@ -92,7 +88,7 @@ def getTermGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, test, aux_indices= 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() @@ -457,65 +453,3 @@ def set_initial_guess(self): sbit.zero() else: sbit.assign(u0bit * dof) - - -@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,)) - - -class TimeProjectorDispatcher(MultiFunction): - def __init__(self, element, t, dt, u0, u1, stages, phi): - MultiFunction.__init__(self) - 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,)) - - def time_projector(self, o): - # 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 - - expr = MultiFunction.reuse_if_untouched - - -def expand_time_projectors(expression, element, t, dt, u0, u1, stages, phi): - rules = TimeProjectorDispatcher(element, t, dt, u0, u1, stages, phi) - return map_integrand_dags(rules, expression) 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/tests/test_galerkin.py b/tests/test_galerkin.py index b58c2c5d..6794c281 100644 --- a/tests/test_galerkin.py +++ b/tests/test_galerkin.py @@ -5,6 +5,7 @@ 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): @@ -279,7 +280,7 @@ def kepler_naive(V, order, t, dt, u0, solver_parameters): F = inner(Dt(u), test)*dx - dHdu(dot(J.T, test)) scheme = ContinuousPetrovGalerkinScheme(order, quadrature_degree="auto") stepper = TimeStepper(F, scheme, t, dt, u, solver_parameters=solver_parameters) - return stepper, [H*dx] + return stepper, [H] def kepler_aux_variable(V, order, t, dt, u0, solver_parameters): @@ -356,8 +357,6 @@ def kepler_projector(V, order, t, dt, u0, solver_parameters): dA1du = diff(A1, uv) dA2du = diff(A2, uv) - Llow = TimeQuadratureLabel(2*order-2) - Qproj = create_time_quadrature(25) w0 = TimeProjector(dHdu, order-1, Qproj) w1 = TimeProjector(dA1du, order-1, Qproj) @@ -365,9 +364,9 @@ def kepler_projector(V, order, t, dt, u0, solver_parameters): determinant_forms = [v, w0, w1, w2] tensor = as_tensor(determinant_forms) - F = Llow(inner(Dt(u), v)*dx) - (det(tensor) / (2*L*H))*dx + F = inner(Dt(u), v)*dx - (det(tensor) / (2*L*H))*dx - scheme = ContinuousPetrovGalerkinScheme(order) + scheme = ContinuousPetrovGalerkinScheme(order, quadrature_degree="auto") stepper = TimeStepper(F, scheme, t, dt, u, solver_parameters=solver_parameters) return stepper, invariants From a0f06e68ac460eeb075d11d5d3d7ed899e2180a4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 3 Mar 2026 16:49:03 +0000 Subject: [PATCH 24/24] add time_projector.py --- irksome/ufl/time_projector.py | 94 +++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 irksome/ufl/time_projector.py 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)