Interface to create form_mass and form_function from a gusto Equation

case_studies/gusto/advection.py

@@ -0,0 +1,171 @@
+
+from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, as_vector,
+                       norm, Constant, conditional, sqrt, VectorFunctionSpace)
+from firedrake.petsc import PETSc
+from gusto import *
+
+from utils.serial import SerialMiniApp
+
+Print = PETSc.Sys.Print
+
+
+def asq_forms(equation):
+    from firedrake.fml import (
+        Term, all_terms, drop,
+        replace_subject, replace_test_function
+    )
+    from gusto.labels import time_derivative, prognostic
+    NullTerm = Term(None)
+
+    V = equation.function_space
+    ncpts = len(V)
+    residual = equation.residual
+
+    # split equation into mass matrix and linear operator
+    mass = residual.label_map(
+        lambda t: t.has_label(time_derivative),
+        map_if_false=drop)
+
+    function = residual.label_map(
+        lambda t: t.has_label(time_derivative),
+        map_if_true=drop)
+
+    # generate ufl for mass matrix over given trial/tests
+    def form_mass(*trials_and_tests):
+        trials = trials_and_tests[:ncpts]
+        tests = trials_and_tests[ncpts:]
+        m = mass.label_map(
+            all_terms,
+            replace_test_function(tests[0]))
+        m = m.label_map(
+            all_terms,
+            replace_subject(trials[0]))
+        return m.form
+
+    # generate ufl for linear operator over given trial/tests
+    def form_function(*trials_and_tests):
+        trials = trials_and_tests[:ncpts]
+        tests = trials_and_tests[ncpts:]
+
+        f = NullTerm
+        for i in range(ncpts):
+            fi = function.label_map(
+                lambda t: t.get(prognostic) == equation.field_name,
+                lambda t: Term(
+                    split_form(t.form)[i].form,
+                    t.labels),
+                map_if_false=drop)
+
+            fi = fi.label_map(
+                all_terms,
+                replace_test_function(tests[i]))
+
+            fi = fi.label_map(
+                all_terms,
+                replace_subject(trials[i]))
+
+            f += fi
+        f = f.label_map(lambda t: t is NullTerm, drop)
+        return f.form
+
+    return form_mass, form_function
+
+
+# Set up model objects
+# ------------------------------------------------------------------------ #
+
+# Domain
+dt = 0.02
+tmax = 1.0
+L = 10
+mesh = PeriodicIntervalMesh(20, L)
+domain = Domain(mesh, dt, "CG", 1)
+
+# Equation
+diffusion_params = DiffusionParameters(kappa=0.75, mu=5)
+V = domain.spaces("DG")
+Vu = VectorFunctionSpace(mesh, "CG", 1)
+
+equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu,
+                                      diffusion_parameters=diffusion_params)
+spatial_methods = [DGUpwind(equation, "f"),
+                   InteriorPenaltyDiffusion(equation, "f", diffusion_params)]
+
+# I/O
+output = OutputParameters(dirname=str('tmp'), dumpfreq=25)
+io = IO(domain, output)
+
+# Time stepper
+stepper = PrescribedTransport(equation, TrapeziumRule(domain), io, spatial_methods)
+
+# ------------------------------------------------------------------------ #
+# Initial conditions
+# ------------------------------------------------------------------------ #
+
+x = SpatialCoordinate(mesh)
+xc_init = 0.25*L
+xc_end = 0.75*L
+umax = 0.5*L/tmax
+# umax = 1
+
+# Get minimum distance on periodic interval to xc
+x_init = conditional(sqrt((x[0] - xc_init)**2) < 0.5*L,
+                     x[0] - xc_init, L + x[0] - xc_init)
+
+x_end = conditional(sqrt((x[0] - xc_end)**2) < 0.5*L,
+                    x[0] - xc_end, L + x[0] - xc_end)
+
+f_init = 5.0
+f_end = f_init / 2.0
+f_width_init = L / 10.0
+f_width_end = f_width_init * 2.0
+f_init_expr = f_init*exp(-(x_init / f_width_init)**2)
+f_end_expr = f_end*exp(-(x_end / f_width_end)**2)
+
+stepper.fields('f').interpolate(f_init_expr)
+stepper.fields('u').interpolate(as_vector([Constant(umax)]))
+f_end = stepper.fields('f_end', space=V)
+f_end.interpolate(f_end_expr)
+
+# ------------------------------------------------------------------------ #
+# Run Gusto stepper
+# ------------------------------------------------------------------------ #
+
+stepper.run(0, tmax=tmax)
+
+# ------------------------------------------------------------------------ #
+# Set up asQ stepper
+# ------------------------------------------------------------------------ #
+
+Print(f"{equation.function_space=}")
+Print(f"{V=}")
+params = {
+    'snes_type': 'ksponly',
+    'ksp_type': 'preonly',
+    'pc_type': 'lu',
+    'ksp_converged_reason': None,
+}
+form_mass, form_function = asq_forms(equation)
+theta = 0.5
+f0 = Function(V).interpolate(f_init_expr)
+asqstep = SerialMiniApp(dt, theta, f0, form_mass, form_function,
+                        solver_parameters=params)
+
+# ------------------------------------------------------------------------ #
+# Run Gusto stepper
+# ------------------------------------------------------------------------ #
+
+asqstep.solve(nt=stepper.step-1)
+
+# ------------------------------------------------------------------------ #
+# Check errors
+# ------------------------------------------------------------------------ #
+
+gusto_error = norm(stepper.fields('f') - f_end) / norm(f_end)
+Print(f"{gusto_error=}")
+
+asq_error = norm(asqstep.w1 - f_end) / norm(f_end)
+Print(f"{asq_error=}")
+
+gusto_asq_error = norm(stepper.fields('f') - asqstep.w1) / norm(f_end)
+Print(f"{gusto_asq_error=}")

case_studies/gusto/diffusion.py

@@ -0,0 +1,66 @@
+from firedrake import (
+    PeriodicSquareMesh, SpatialCoordinate,
+    exp, errornorm, Function)
+from gusto import *
+from firedrake.petsc import PETSc
+Print = PETSc.Sys.Print
+
+from utils.serial import SerialMiniApp
+from gusto_interface import asq_forms
+
+dt = 0.01
+L = 10.
+mesh = PeriodicSquareMesh(10, 10, L, direction='x')
+
+output = OutputParameters(dirname=str('output'), dumpfreq=25)
+domain = Domain(mesh, dt, family="CG", degree=1)
+io = IO(domain, output)
+
+tmax = 1.0
+x = SpatialCoordinate(mesh)
+f_init = exp(-((x[0]-0.5*L)**2 + (x[1]-0.5*L)**2))
+
+tol = 5.e-2
+kappa = 1.
+
+f_end_expr = (1/(1+4*tmax))*f_init**(1/(1+4*tmax))
+
+V = domain.spaces("DG")
+Print(f"{V=}")
+
+mu = 5.
+
+# gusto stepper
+
+diffusion_params = DiffusionParameters(kappa=kappa, mu=mu)
+eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params)
+diffusion_scheme = BackwardEuler(domain)
+diffusion_methods = [InteriorPenaltyDiffusion(eqn, "f", diffusion_params)]
+timestepper = Timestepper(eqn, diffusion_scheme, io, spatial_methods=diffusion_methods)
+timestepper.fields("f").interpolate(f_init)
+
+# asq stepper
+form_mass, form_function = asq_forms(eqn)
+theta = 1
+f0 = Function(V).interpolate(f_init)
+params = {
+    'snes_type': 'ksponly',
+    'ksp_type': 'preonly',
+    'pc_type': 'lu',
+    'ksp_converged_reason': None,
+}
+asqstepper = SerialMiniApp(dt, theta, f0, form_mass, form_function,
+                           solver_parameters=params)
+
+# Run gusto stepper
+timestepper.run(0., tmax)
+gusto_f_end = timestepper.fields('f')
+Print(f"{errornorm(f_end_expr, gusto_f_end)=}")
+
+# Run asQ stepper
+asqstepper.solve(nt=timestepper.step-1)
+asq_f_end = asqstepper.w1
+Print(f"{errornorm(f_end_expr, asq_f_end)=}")
+
+# error
+Print(f"{errornorm(gusto_f_end, asq_f_end)=}")

case_studies/gusto/gusto_interface.py

@@ -0,0 +1,74 @@
+try:
+    from gusto import time_derivative, prognostic, transporting_velocity
+except ModuleNotFoundError as err:
+    raise type(err)("Gusto must be installed to use the asQ-Gusto interface.") from err
+
+from firedrake.fml import (
+    Term, all_terms, drop,
+    replace_subject, replace_test_function
+)
+from firedrake.formmanipulation import split_form
+from ufl import replace
+
+
+def asq_forms(equation, transport_velocity_index=None):
+    NullTerm = Term(None)
+    ncpts = len(equation.function_space)
+
+    # split equation into mass matrix and linear operator
+    mass = equation.residual.label_map(
+        lambda t: t.has_label(time_derivative),
+        map_if_false=drop)
+
+    function = equation.residual.label_map(
+        lambda t: t.has_label(time_derivative),
+        map_if_true=drop)
+
+    # generate ufl for mass matrix over given trial/tests
+    def form_mass(*trials_and_tests):
+        trials = trials_and_tests[:ncpts] if ncpts > 1 else trials_and_tests[0]
+        tests = trials_and_tests[ncpts:] if ncpts > 1 else trials_and_tests[1]
+
+        m = mass.label_map(
+            all_terms,
+            replace_test_function(tests))
+        m = m.label_map(
+            all_terms,
+            replace_subject(trials))
+        return m.form
+
+    # generate ufl for linear operator over given trial/tests
+    def form_function(*trials_and_tests):
+        trials = trials_and_tests[:ncpts] if ncpts > 1 else trials_and_tests[0]
+        tests = trials_and_tests[ncpts:2*ncpts]
+
+        fields = equation.field_names if ncpts > 1 else [equation.field_name]
+
+        f = NullTerm
+        for i in range(ncpts):
+            fi = function.label_map(
+                lambda t: t.get(prognostic) == fields[i],
+                lambda t: Term(
+                    split_form(t.form)[i].form,
+                    t.labels),
+                map_if_false=drop)
+
+            fi = fi.label_map(
+                all_terms,
+                replace_test_function(tests[i]))
+
+            fi = fi.label_map(
+                all_terms,
+                replace_subject(trials))
+
+            if transport_velocity_index is not None:
+                transport_velocity = trials[transport_velocity_index]
+                fi = fi.label_map(
+                    lambda t: t.has_label(transporting_velocity),
+                    map_if_true=lambda t: Term(replace(t.form, {t.get(transporting_velocity): transport_velocity}), t.labels))
+
+            f += fi
+        f = f.label_map(lambda t: t is NullTerm, drop)
+        return f.form
+
+    return form_mass, form_function