From 3d10abad3ac2107e45a27e095b3b20a29a2c8a08 Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 26 Jul 2024 14:57:31 +0100 Subject: [PATCH 1/4] set quadrature degree for forms in equation residual to ensure maximum is not exceeded --- .../equations/compressible_euler_equations.py | 17 ++++++--- gusto/equations/prognostic_equations.py | 10 ++++-- gusto/timestepping/timestepper.py | 36 ++++++++++++++++++- 3 files changed, 55 insertions(+), 8 deletions(-) diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index b6b89cd8..acf5aa37 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/gusto/equations/compressible_euler_equations.py @@ -40,7 +40,7 @@ def __init__(self, domain, parameters, Omega=None, sponge_options=None, u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, - active_tracers=None): + active_tracers=None, max_quad_deg=5): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -77,6 +77,8 @@ def __init__(self, domain, parameters, Omega=None, sponge_options=None, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations.. Defaults to None. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. Raises: NotImplementedError: only mixing ratio tracers are implemented. @@ -100,7 +102,8 @@ def __init__(self, domain, parameters, Omega=None, sponge_options=None, super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, - active_tracers=active_tracers) + active_tracers=active_tracers, + max_quad_deg=max_quad_deg) self.parameters = parameters g = parameters.g @@ -284,7 +287,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, - active_tracers=None): + active_tracers=None, + max_quad_deg=5): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -319,7 +323,9 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, None. active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included - in the equations.. Defaults to None. + in the equations. Defaults to None. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. Raises: NotImplementedError: only mixing ratio tracers are implemented. @@ -331,7 +337,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, u_transport_option=u_transport_option, diffusion_options=diffusion_options, no_normal_flow_bc_ids=no_normal_flow_bc_ids, - active_tracers=active_tracers) + active_tracers=active_tracers, + max_quad_deg=max_quad_deg) self.residual = self.residual.label_map( lambda t: t.has_label(time_derivative), diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index bf8dd784..326966d0 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -24,7 +24,7 @@ class PrognosticEquation(object, metaclass=ABCMeta): """Base class for prognostic equations.""" - def __init__(self, domain, function_space, field_name): + def __init__(self, domain, function_space, field_name, max_quad_deg=5): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -32,12 +32,15 @@ def __init__(self, domain, function_space, field_name): function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. """ self.domain = domain self.function_space = function_space self.X = Function(function_space) self.field_name = field_name + self.max_quad_deg = max_quad_deg self.bcs = {} self.prescribed_fields = PrescribedFields() @@ -76,7 +79,7 @@ class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): def __init__(self, field_names, domain, space_names, linearisation_map=None, no_normal_flow_bc_ids=None, - active_tracers=None): + active_tracers=None, max_quad_deg=5): """ Args: field_names (list): a list of strings for names of the prognostic @@ -94,12 +97,15 @@ def __init__(self, field_names, domain, space_names, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations.. Defaults to None. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. """ self.field_names = field_names self.space_names = space_names self.active_tracers = active_tracers self.linearisation_map = lambda t: False if linearisation_map is None else linearisation_map(t) + self.max_quad_deg = max_quad_deg # Build finite element spaces self.spaces = [domain.spaces(space_name) for space_name in diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py index e8e1a1d3..3158977a 100644 --- a/gusto/timestepping/timestepper.py +++ b/gusto/timestepping/timestepper.py @@ -2,7 +2,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import Function, Projector, split -from firedrake.fml import drop, Term +from firedrake.fml import drop, Term, all_terms from pyop2.profiling import timed_stage from gusto.equations import PrognosticEquationSet from gusto.core import TimeLevelFields, StateFields @@ -12,6 +12,7 @@ from gusto.spatial_methods.transport_methods import TransportMethod import ufl + __all__ = ["BaseTimestepper", "Timestepper", "PrescribedTransport"] @@ -108,6 +109,39 @@ def setup_equation(self, equation): for method in self.spatial_methods: method.replace_form(equation) + # -------------------------------------------------------------------- # + # Ensure the quadrature degree is not excessive for any integrals + # -------------------------------------------------------------------- # + def update_quadrature_degree(t): + # This function takes in a term and returns a new term + # with the same form and labels, the only difference being + # that any integrals in the form with no quadrature_degree + # set will have their quadrature degree set to max_quad_deg + + # This list will hold the updated integrals + new_itgs = [] + + # Loop over integrals in this term's form + for itg in t.form._integrals: + # check if the quadrature degree is not set + if 'quadrature_degree' not in itg._metadata.keys(): + # create new integral with updated quadrature degree + new_itg = itg.reconstruct( + metadata={'quadrature_degree': equation.max_quad_deg}) + new_itgs.append(new_itg) + else: + # if quadrature degree was already set, just keep + # this integral + new_itgs.append(itg) + + new_form = ufl.form.Form(new_itgs) + + return Term(new_form, t.labels) + + equation.residual = equation.residual.label_map( + all_terms, + lambda t: update_quadrature_degree(t)) + def setup_transporting_velocity(self, scheme): """ Set up the time discretisation by replacing the transporting velocity From 3bb7df58752b995df1952de6450acd5c80859e20 Mon Sep 17 00:00:00 2001 From: jshipton Date: Sun, 28 Jul 2024 12:03:49 +0100 Subject: [PATCH 2/4] get rid of .split --- examples/shallow_water/moist_convective_williamson2.py | 2 +- examples/shallow_water/moist_thermal_williamson5.py | 6 ++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/examples/shallow_water/moist_convective_williamson2.py b/examples/shallow_water/moist_convective_williamson2.py index 22ca6b78..4c21ebb7 100644 --- a/examples/shallow_water/moist_convective_williamson2.py +++ b/examples/shallow_water/moist_convective_williamson2.py @@ -82,7 +82,7 @@ # define saturation function def sat_func(x_in): - h = x_in.split()[1] + h = x_in.subfunctions[1] lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma)) denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2 diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py index f022e105..21162a83 100644 --- a/examples/shallow_water/moist_thermal_williamson5.py +++ b/examples/shallow_water/moist_thermal_williamson5.py @@ -84,15 +84,13 @@ # Saturation function def sat_func(x_in): - h = x_in.split()[1] - b = x_in.split()[2] + _, h, b = x_in.subfunctions return (q0/(g*h + g*tpexpr)) * exp(20*(1 - b/g)) # Feedback proportionality is dependent on h and b def gamma_v(x_in): - h = x_in.split()[1] - b = x_in.split()[2] + _, h, b = x_in.subfunctions return (1 + beta2*(20*q0/(g*h + g*tpexpr) * exp(20*(1 - b/g))))**(-1) From ca0bcd9402fb7183e7761d8ec917f1937d5587de Mon Sep 17 00:00:00 2001 From: jshipton Date: Sun, 28 Jul 2024 16:31:05 +0100 Subject: [PATCH 3/4] shut up some interpolation warnings --- gusto/diagnostics/diagnostics.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/gusto/diagnostics/diagnostics.py b/gusto/diagnostics/diagnostics.py index b287d6d3..00c8ba08 100644 --- a/gusto/diagnostics/diagnostics.py +++ b/gusto/diagnostics/diagnostics.py @@ -9,6 +9,7 @@ Projector, Interpolator, FunctionSpace, FiniteElement, TensorProductElement) from firedrake.assign import Assigner +from firedrake.__future__ import interpolate from ufl.domain import extract_unique_domain from abc import ABCMeta, abstractmethod, abstractproperty @@ -191,9 +192,7 @@ def setup(self, domain, state_fields, space=None): f"The expression for diagnostic {self.name} has not been specified" # Solve method must be declared in diagnostic's own setup routine - if self.method == 'interpolate': - self.evaluator = Interpolator(self.expr, self.field) - elif self.method == 'project': + if self.method == 'project': self.evaluator = Projector(self.expr, self.field) elif self.method == 'assign': self.evaluator = Assigner(self.field, self.expr) @@ -206,7 +205,7 @@ def compute(self): logger.debug(f'Computing diagnostic {self.name} with {self.method} method') if self.method == 'interpolate': - self.evaluator.interpolate() + assemble(interpolate(self.expr, self.field)) elif self.method == 'assign': self.evaluator.assign() elif self.method == 'project': From 586be15aeba747392f0eab03d5e973f853fc733c Mon Sep 17 00:00:00 2001 From: jshipton Date: Sun, 28 Jul 2024 17:02:44 +0100 Subject: [PATCH 4/4] actually, I think this is how to remove the interpolation warnings --- gusto/diagnostics/diagnostics.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/gusto/diagnostics/diagnostics.py b/gusto/diagnostics/diagnostics.py index 00c8ba08..d5d24466 100644 --- a/gusto/diagnostics/diagnostics.py +++ b/gusto/diagnostics/diagnostics.py @@ -6,10 +6,10 @@ LinearVariationalProblem, LinearVariationalSolver, ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, pi, TensorFunctionSpace, SpatialCoordinate, as_vector, - Projector, Interpolator, FunctionSpace, FiniteElement, + Projector, FunctionSpace, FiniteElement, TensorProductElement) from firedrake.assign import Assigner -from firedrake.__future__ import interpolate +from firedrake.__future__ import Interpolator from ufl.domain import extract_unique_domain from abc import ABCMeta, abstractmethod, abstractproperty @@ -192,7 +192,9 @@ def setup(self, domain, state_fields, space=None): f"The expression for diagnostic {self.name} has not been specified" # Solve method must be declared in diagnostic's own setup routine - if self.method == 'project': + if self.method == 'interpolate': + self.evaluator = Interpolator(self.expr, self.field) + elif self.method == 'project': self.evaluator = Projector(self.expr, self.field) elif self.method == 'assign': self.evaluator = Assigner(self.field, self.expr) @@ -205,7 +207,7 @@ def compute(self): logger.debug(f'Computing diagnostic {self.name} with {self.method} method') if self.method == 'interpolate': - assemble(interpolate(self.expr, self.field)) + self.evaluator.interpolate() elif self.method == 'assign': self.evaluator.assign() elif self.method == 'project':