From 573ad527e91b2f67e470f11451748c76d0e74183 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Wed, 20 Dec 2023 02:59:59 +0000 Subject: [PATCH 01/25] started MixedOptions and an integration test for this --- gusto/time_discretisation.py | 3 + gusto/wrappers.py | 44 ++- .../transport/test_mixed_fs_options.py | 313 ++++++++++++++++++ 3 files changed, 359 insertions(+), 1 deletion(-) create mode 100644 integration-tests/transport/test_mixed_fs_options.py diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 79c8af94e..869b61931 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -87,6 +87,9 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.courant_max = None if options is not None: + if options.wrapper_type = Mixed: + pass + else: self.wrapper_name = options.name if self.wrapper_name == "embedded_dg": self.wrapper = EmbeddedDGWrapper(self, options) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 77d46979a..1fd695685 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -15,7 +15,7 @@ from gusto.labels import transporting_velocity import ufl -__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper"] +__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedOptions"] class Wrapper(object, metaclass=ABCMeta): @@ -33,6 +33,7 @@ def __init__(self, time_discretisation, wrapper_options): self.time_discretisation = time_discretisation self.options = wrapper_options self.solver_parameters = None + self.field_name = None @abstractmethod def setup(self): @@ -361,3 +362,44 @@ def label_terms(self, residual): new_residual = transporting_velocity.update_value(new_residual, self.transporting_velocity) return new_residual + + +class MixedOptions(object): + """ + An object to hold a dictionary with different options for different + tracers. This means that different tracers can be solved simultaneously + using a CoupledTransportEquation, whilst being in different spaces + and needing different implementation options. + """ + + def __init__(self, equation, suboptions): + """ + Args: + equation (:class: `PrognosticEquationSet`): the prognostic equation(s) + sublimiters (dict): A dictionary holding limiters defined for individual prognostic variables + Raises: + ValueError: If an option is defined for a field that is not in the prognostic variable set + """ + + self.suboptions = suboptions + self.wrapper_type = Mixed + + for field, suboption in suboptions: + # Check that the field is in the prognostic variable set: + if field not in equation.field_names: + raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") + else: + # check that a valid wrapper has been given + + + self.suboptions[field].idx = equation.field_names.index(field) + + + def apply(self, fields): + """ + Apply the individual limiters to specific prognostic variables + """ + + for _, suboption in self.suboptions.items(): + field = fields.subfunctions[suboption.idx] + sublimiter.apply(field) diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py new file mode 100644 index 000000000..325ef645c --- /dev/null +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -0,0 +1,313 @@ +""" +This tests limiter options for different transport schemes. +A sharp bubble of warm air is generated in a vertical slice and then transported +by a prescribed transport scheme. If the limiter is working, the transport +should have produced no new maxima or minima. +""" + +from gusto import * +from firedrake import (as_vector, PeriodicIntervalMesh, pi, SpatialCoordinate, + ExtrudedMesh, FunctionSpace, Function, norm, + conditional, sqrt) +from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter +import numpy as np +import pytest + + +def setup_limiters(dirname, space_A, space_B): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + Ld = 1. + tmax = 0.2 + dt = tmax / 40 + rotations = 0.25 + + # ------------------------------------------------------------------------ # + # Build model objects + # ------------------------------------------------------------------------ # + + # Domain + m = PeriodicIntervalMesh(20, Ld) + mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20)) + degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] or space_B in ['DG0', 'Vtheta_degree_0'] else 1 + + domain = Domain(mesh, dt, family="CG", degree=degree) + + # Transporting velocity space + V = domain.spaces('HDiv') + + # Tracer A spaces + if space_A == 'DG0': + VA = domain.spaces('DG') + VCG1_A = FunctionSpace(mesh, 'CG', 1) + VDG1_A = domain.spaces('DG1_equispaced') + space_A_string = 'DG' + elif space_A == 'DG1': + VA = domain.spaces('DG') + space_A_string = 'DG' + elif space_A == 'DG1_equispaced': + VA = domain.spaces('DG1_equispaced') + space_A_string = 'DG1_equispaced' + elif space_A == 'Vtheta_degree_0': + VA = domain.spaces('theta') + VCG1_A = FunctionSpace(mesh, 'CG', 1) + VDG1_A = domain.spaces('DG1_equispaced') + space_A_string = 'theta' + elif space_A == 'Vtheta_degree_1': + VA = domain.spaces('theta') + space_A_string = 'theta' + else: + raise NotImplementedError + + # Tracer B spaces + if space_B == 'DG0': + VB = domain.spaces('DG') + VCG1_B = FunctionSpace(mesh, 'CG', 1) + VDG1_B = domain.spaces('DG1_equispaced') + space_B_string = 'DG' + elif space_B == 'DG1': + VB = domain.spaces('DG') + space_B_string = 'DG' + elif space_B == 'DG1_equispaced': + VB = domain.spaces('DG1_equispaced') + space_B_string = 'DG1_equispaced' + elif space_B == 'Vtheta_degree_0': + VB = domain.spaces('theta') + VCG1_B = FunctionSpace(mesh, 'CG', 1) + VDG1_B = domain.spaces('DG1_equispaced') + space_B_string = 'theta' + elif space_B == 'Vtheta_degree_1': + VB = domain.spaces('theta') + space_B_string = 'theta' + else: + raise NotImplementedError + + Vpsi = domain.spaces('H1') + + tracerA = ActiveTracer(name='tracerA', space=space_A_string, + variable_type=TracerVariableType.mixing_ratio, + transport_eqn=TransportEquationType.advective) + + tracerB = ActiveTracer(name='tracerB', space=space_B_string, + variable_type=TracerVariableType.mixing_ratio, + transport_eqn=TransportEquationType.advective) + + + tracers = [tracerA, tracerB] + + eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V) + output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1, + dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB']) + + io = IO(domain, output) + + # ------------------------------------------------------------------------ # + # Set up transport scheme with options + # ------------------------------------------------------------------------ # + + #suboptions = {} + sublimiters = {} + + # Options and limiters for tracer_A + + if space_A in ['DG0', 'Vtheta_degree_0']: + recover_opts = RecoveryOptions(embedding_space=VDG1_A, + recovered_space=VCG1_A, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor) + + sublimiters.update({'tracerA': VertexBasedLimiter(VDG1_A)}) + + elif space_A == 'DG1': + sublimiters.update({'tracerA': DG1Limiter(VA)}) + + elif space_A == 'DG1_equispaced': + sublimiters.update({'tracerA': VertexBasedLimiter(VA)}) + + elif space_A == 'Vtheta_degree_1': + opts = EmbeddedDGOptions() + sublimiters.update({'tracerA': ThetaLimiter(VA)}) + + else: + raise NotImplementedError + + # Options and limiters for tracer_B + + if space_B in ['DG0', 'Vtheta_degree_0']: + recover_opts = RecoveryOptions(embedding_space=VDG1_B, + recovered_space=VCG1_B, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor) + + sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)}) + + elif space_B == 'DG1': + sublimiters.update({'tracerB': DG1Limiter(VB)}) + + elif space_B == 'DG1_equispaced': + sublimiters.update({'tracerB': VertexBasedLimiter(VB)}) + + elif space_B == 'Vtheta_degree_1': + opts = EmbeddedDGOptions() + sublimiters.update({'tracerB': ThetaLimiter(VB)}) + + else: + raise NotImplementedError + + #MixedOptions = MixedFSOptions(eqn, suboptions) + MixedLimiter = MixedFSLimiter(eqn, sublimiters) + + # Give the scheme for the coupled transport + transport_schemes = SSPRK3(domain, limiter=MixedLimiter) + + # DG Upwind transport for both tracers: + transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] + + # Build time stepper + stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method) + + # ------------------------------------------------------------------------ # + # Initial condition + # ------------------------------------------------------------------------ # + + tracerA_0 = stepper.fields('tracerA', space=VA) + tracerB_0 = stepper.fields('tracerB', space=VB) + true_fieldA = stepper.fields('true_tracerA', space=VA) + true_fieldB = stepper.fields('true_tracerB', space=VB) + + x, z = SpatialCoordinate(mesh) + + tracer_min = 12.6 + dtracer = 3.2 + + # First time do initial conditions, second time do final conditions + for i in range(2): + + if i == 0: + x1_lower = 2 * Ld / 5 + x1_upper = 3 * Ld / 5 + z1_lower = 6 * Ld / 10 + z1_upper = 8 * Ld / 10 + x2_lower = 6 * Ld / 10 + x2_upper = 8 * Ld / 10 + z2_lower = 2 * Ld / 5 + z2_upper = 3 * Ld / 5 + elif i == 1: + # Rotated anti-clockwise by 90 degrees (x -> z, z -> -x) + x1_lower = 2 * Ld / 10 + x1_upper = 4 * Ld / 10 + z1_lower = 2 * Ld / 5 + z1_upper = 3 * Ld / 5 + x2_lower = 2 * Ld / 5 + x2_upper = 3 * Ld / 5 + z2_lower = 6 * Ld / 10 + z2_upper = 8 * Ld / 10 + else: + raise ValueError + + expr_1 = conditional(x > x1_lower, + conditional(x < x1_upper, + conditional(z > z1_lower, + conditional(z < z1_upper, dtracer, 0.0), + 0.0), + 0.0), + 0.0) + + expr_2 = conditional(x > x2_lower, + conditional(x < x2_upper, + conditional(z > z2_lower, + conditional(z < z2_upper, dtracer, 0.0), + 0.0), + 0.0), + 0.0) + + if i == 0: + tracerA_0.interpolate(Constant(tracer_min) + expr_1 + expr_2) + tracerB_0.interpolate(Constant(tracer_min) + expr_1 + expr_2) + elif i == 1: + true_fieldA.interpolate(Constant(tracer_min) + expr_1 + expr_2) + true_fieldB.interpolate(Constant(tracer_min) + expr_1 + expr_2) + else: + raise ValueError + + # ------------------------------------------------------------------------ # + # Velocity profile + # ------------------------------------------------------------------------ # + + psi = Function(Vpsi) + u = stepper.fields('u') + + # set up solid body rotation for transport + # we do this slightly complicated stream function to make the velocity 0 at edges + # thus we avoid odd effects at boundaries + xc = Ld / 2 + zc = Ld / 2 + r = sqrt((x - xc) ** 2 + (z - zc) ** 2) + omega = rotations * 2 * pi / tmax + r_out = 9 * Ld / 20 + r_in = 2 * Ld / 5 + A = omega * r_in / (2 * (r_in - r_out)) + B = - omega * r_in * r_out / (r_in - r_out) + C = omega * r_in ** 2 * r_out / (r_in - r_out) / 2 + psi_expr = conditional(r < r_in, + omega * r ** 2 / 2, + conditional(r < r_out, + A * r ** 2 + B * r + C, + A * r_out ** 2 + B * r_out + C)) + psi.interpolate(psi_expr) + + gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)]) + u.project(gradperp(psi)) + + return stepper, tmax, true_fieldA, true_fieldB + + +@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', + 'DG1', 'DG1_equispaced']) +#Remove Dg1-dg1 and other easy ones after debugging +@pytest.mark.parametrize('space_B', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', + 'DG1', 'DG1_equispaced']) + + +def test_limiters(tmpdir, space_A, space_B): + + # Setup and run + dirname = str(tmpdir) + + stepper, tmax, true_fieldA, true_fieldB = setup_limiters(dirname, space_A, space_B) + + stepper.run(t=0, tmax=tmax) + + tol = 1e-9 + + final_fieldA = stepper.fields('tracerA') + final_fieldB = stepper.fields('tracerB') + + # Check tracer is roughly in the correct place + assert norm(true_fieldA - final_fieldA) / norm(true_fieldA) < 0.05, \ + 'Something is wrong with the DG space tracer using a mixed limiter' + + # Check tracer is roughly in the correct place + assert norm(true_fieldB - final_fieldB) / norm(true_fieldB) < 0.05, \ + 'Something is wrong with the DG1 equispaced tracer using a mixed limiter' + + # Check for no new overshoots in A + assert np.max(final_fieldA.dat.data) <= np.max(true_fieldA.dat.data) + tol, \ + 'Application of the DG space limiter in the mixed limiter has not prevented overshoots' + + # Check for no new undershoots in A + assert np.min(final_fieldA.dat.data) >= np.min(true_fieldA.dat.data) - tol, \ + 'Application of the DG space limiter in the mixed limiter has not prevented undershoots' + + # Check for no new overshoots in B + assert np.max(final_fieldB.dat.data) <= np.max(true_fieldB.dat.data) + tol, \ + 'Application of the DG1 equispaced limiter in the mixed limiter has not prevented overshoots' + + # Check for no new undershoots in B + assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \ + 'Application of the DG1 equispaced limiter in the mixed limiter has not prevented undershoots' + + \ No newline at end of file From e04c73ac8ee9dcd5ee6bde0e1fe88f0636924b69 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Thu, 21 Dec 2023 20:18:24 +0000 Subject: [PATCH 02/25] MixedOptions dictionary added --- gusto/time_discretisation.py | 24 +++++------ gusto/wrappers.py | 43 +++++++++++++------ .../transport/test_mixed_fs_options.py | 18 ++++---- 3 files changed, 52 insertions(+), 33 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 869b61931..ab3432b39 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -87,19 +87,19 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.courant_max = None if options is not None: - if options.wrapper_type = Mixed: - pass - else: - self.wrapper_name = options.name - if self.wrapper_name == "embedded_dg": - self.wrapper = EmbeddedDGWrapper(self, options) - elif self.wrapper_name == "recovered": - self.wrapper = RecoveryWrapper(self, options) - elif self.wrapper_name == "supg": - self.wrapper = SUPGWrapper(self, options) + if options.wrapper_type == 'mixed': + self.wrapper = options else: - raise RuntimeError( - f'Time discretisation: wrapper {self.wrapper_name} not implemented') + self.wrapper_name = options.name + if self.wrapper_name == "embedded_dg": + self.wrapper = EmbeddedDGWrapper(self, options) + elif self.wrapper_name == "recovered": + self.wrapper = RecoveryWrapper(self, options) + elif self.wrapper_name == "supg": + self.wrapper = SUPGWrapper(self, options) + else: + raise RuntimeError( + f'Time discretisation: wrapper {self.wrapper_name} not implemented') else: self.wrapper = None self.wrapper_name = None diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 1fd695685..71d97493e 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -381,19 +381,38 @@ def __init__(self, equation, suboptions): ValueError: If an option is defined for a field that is not in the prognostic variable set """ - self.suboptions = suboptions - self.wrapper_type = Mixed - - for field, suboption in suboptions: - # Check that the field is in the prognostic variable set: - if field not in equation.field_names: - raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") - else: - # check that a valid wrapper has been given - - - self.suboptions[field].idx = equation.field_names.index(field) + self.suboptions = suboptions + self.wrapper_type = 'mixed' + self.wrappers = [] + + for field, suboption in suboptions: + # Check that the field is in the prognostic variable set: + if field not in equation.field_names: + raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") + else: + # check that a valid wrapper has been given + wrapper_name = suboption.name + + if wrapper_name == "embedded_dg": + self.wrappers.append(EmbeddedDGWrapper(self, options)) + self.wrapper_fields.append(field) + elif wrapper_name == "recovered": + self.wrappers.append(RecoveryWrapper(self, options)) + self.wrapper_fields.append(field) + elif wrapper_name == "supg": + self.wrappers.append(SUPGWrapper(self, options)) + self.wrapper_fields.append(field) + else: + raise RuntimeError( + f'Time discretisation: suboption wrapper {wrapper_name} not implemented') + + #Initialise the wrapper and associate with a field: + + self.suboptions[field].idx = equation.field_names.index(field) + def setup(self): + # This is done in the suboption wrappers themselves + pass def apply(self, fields): """ diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 325ef645c..94ec51bb6 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -108,16 +108,16 @@ def setup_limiters(dirname, space_A, space_B): # Set up transport scheme with options # ------------------------------------------------------------------------ # - #suboptions = {} + suboptions = {} sublimiters = {} # Options and limiters for tracer_A if space_A in ['DG0', 'Vtheta_degree_0']: - recover_opts = RecoveryOptions(embedding_space=VDG1_A, - recovered_space=VCG1_A, - project_low_method='recover', - boundary_method=BoundaryMethod.taylor) + suboptions.update({'tracerA': RecoveryOptions(embedding_space=VDG1_A, + recovered_space=VCG1_A, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor)}) sublimiters.update({'tracerA': VertexBasedLimiter(VDG1_A)}) @@ -157,11 +157,11 @@ def setup_limiters(dirname, space_A, space_B): else: raise NotImplementedError - #MixedOptions = MixedFSOptions(eqn, suboptions) + opts = MixedOptions(eqn, suboptions) MixedLimiter = MixedFSLimiter(eqn, sublimiters) # Give the scheme for the coupled transport - transport_schemes = SSPRK3(domain, limiter=MixedLimiter) + transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter) # DG Upwind transport for both tracers: transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] @@ -268,8 +268,8 @@ def setup_limiters(dirname, space_A, space_B): @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', 'DG1', 'DG1_equispaced']) #Remove Dg1-dg1 and other easy ones after debugging -@pytest.mark.parametrize('space_B', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', - 'DG1', 'DG1_equispaced']) +@pytest.mark.parametrize('space_B', ['Vtheta_degree_0'])#, 'Vtheta_degree_1', 'DG0', + #'DG1', 'DG1_equispaced']) def test_limiters(tmpdir, space_A, space_B): From 4a13df3a326ce8404ec18ef655a523e289885280 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Fri, 22 Dec 2023 02:43:21 +0000 Subject: [PATCH 03/25] improvements to MixedOptions --- gusto/time_discretisation.py | 56 ++++++--- gusto/wrappers.py | 47 +++++--- .../transport/test_mixed_fs_options.py | 106 ++++++++++-------- 3 files changed, 131 insertions(+), 78 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index ab3432b39..593bbe830 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -37,6 +37,20 @@ def wrapper_apply(original_apply): """Decorator to add steps for using a wrapper around the apply method.""" def get_apply(self, x_out, x_in): + #if type(self.wrapper) == MixedOptions: + + # Subwrappers for different tracers have been defined + # Use the apply defined in the MixedOptions object + + # print('aha') + # self.wrapper.pre_apply(x_in) + + # for _, subwrapper in self.wrapper.suboptions.items(): + # original_apply(self, self.subwrapper.x_out, self.subwrapper.x_in) + + # self.wrapper.post_apply(x_out) + + if self.wrapper is not None: def new_apply(self, x_out, x_in): @@ -87,8 +101,11 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.courant_max = None if options is not None: - if options.wrapper_type == 'mixed': + if type(options) == MixedOptions: + print('jahjah') self.wrapper = options + # Or do I need to initialise everything here + # like is done for a single wrapper? else: self.wrapper_name = options.name if self.wrapper_name == "embedded_dg": @@ -162,21 +179,28 @@ def setup(self, equation, apply_bcs=True, *active_labels): # -------------------------------------------------------------------- # if self.wrapper is not None: - self.wrapper.setup() - self.fs = self.wrapper.function_space - if self.solver_parameters is None: - self.solver_parameters = self.wrapper.solver_parameters - new_test = TestFunction(self.wrapper.test_space) - # SUPG has a special wrapper - if self.wrapper_name == "supg": - new_test = self.wrapper.test - - # Replace the original test function with the one from the wrapper - self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test)) - - self.residual = self.wrapper.label_terms(self.residual) + if type(self.wrapper) == MixedOptions: + # Subwrappers are defined. + # Set these up with ? + for _, subwrapper in self.wrapper.suboptions.items(): + print(subwrapper) + pass + else: + self.wrapper.setup() + self.fs = self.wrapper.function_space + if self.solver_parameters is None: + self.solver_parameters = self.wrapper.solver_parameters + new_test = TestFunction(self.wrapper.test_space) + # SUPG has a special wrapper + if self.wrapper_name == "supg": + new_test = self.wrapper.test + + # Replace the original test function with the one from the wrapper + self.residual = self.residual.label_map( + all_terms, + map_if_true=replace_test_function(new_test)) + + self.residual = self.wrapper.label_terms(self.residual) # -------------------------------------------------------------------- # # Make boundary conditions diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 71d97493e..b14a23b6b 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -380,28 +380,33 @@ def __init__(self, equation, suboptions): Raises: ValueError: If an option is defined for a field that is not in the prognostic variable set """ - + print(equation.function_space) + self.x_in = Function(equation.function_space) + self.x_out = Function(equation.function_space) + + print('Initialising MixedOptions') + self.suboptions = suboptions - self.wrapper_type = 'mixed' - self.wrappers = [] - for field, suboption in suboptions: + print(suboptions.items()) + + self.wrapper = {} + + for field, suboption in suboptions.items(): # Check that the field is in the prognostic variable set: if field not in equation.field_names: raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") else: # check that a valid wrapper has been given wrapper_name = suboption.name + print(wrapper_name) if wrapper_name == "embedded_dg": - self.wrappers.append(EmbeddedDGWrapper(self, options)) - self.wrapper_fields.append(field) + self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)}) elif wrapper_name == "recovered": - self.wrappers.append(RecoveryWrapper(self, options)) - self.wrapper_fields.append(field) + self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption) elif wrapper_name == "supg": - self.wrappers.append(SUPGWrapper(self, options)) - self.wrapper_fields.append(field) + self.suboptions[field].subwrapper = SUPGWrapper(self, suboption) else: raise RuntimeError( f'Time discretisation: suboption wrapper {wrapper_name} not implemented') @@ -409,16 +414,30 @@ def __init__(self, equation, suboptions): #Initialise the wrapper and associate with a field: self.suboptions[field].idx = equation.field_names.index(field) + #self.suboptions[field].fs = equation.field_names.function_space(field) + # self.subwrapper.x_in = Function(self.suboptions[field].fs) def setup(self): # This is done in the suboption wrappers themselves pass - def apply(self, fields): + def pre_apply(self, x_in): """ - Apply the individual limiters to specific prognostic variables + Perform the pre-applications from all subwrappers + """ + + for _, subwrapper in self.suboptions.items(): + print(subwrapper) + field = x_in.subfunctions[subwrapper.idx] + subwrapper.pre_apply(field) + + def post_apply(self, x_in): + """ + Perform the post-applications from all subwrappers """ for _, suboption in self.suboptions.items(): - field = fields.subfunctions[suboption.idx] - sublimiter.apply(field) + print(suboption) + field = x_in.subfunctions[suboption.idx] + suboption.subwrapper.post_apply(field) + diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 94ec51bb6..a82136923 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -32,7 +32,7 @@ def setup_limiters(dirname, space_A, space_B): # Domain m = PeriodicIntervalMesh(20, Ld) mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20)) - degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] or space_B in ['DG0', 'Vtheta_degree_0'] else 1 + degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] else 1 domain = Domain(mesh, dt, family="CG", degree=degree) @@ -63,25 +63,28 @@ def setup_limiters(dirname, space_A, space_B): raise NotImplementedError # Tracer B spaces - if space_B == 'DG0': - VB = domain.spaces('DG') - VCG1_B = FunctionSpace(mesh, 'CG', 1) - VDG1_B = domain.spaces('DG1_equispaced') - space_B_string = 'DG' - elif space_B == 'DG1': - VB = domain.spaces('DG') - space_B_string = 'DG' - elif space_B == 'DG1_equispaced': - VB = domain.spaces('DG1_equispaced') - space_B_string = 'DG1_equispaced' - elif space_B == 'Vtheta_degree_0': - VB = domain.spaces('theta') - VCG1_B = FunctionSpace(mesh, 'CG', 1) - VDG1_B = domain.spaces('DG1_equispaced') - space_B_string = 'theta' - elif space_B == 'Vtheta_degree_1': - VB = domain.spaces('theta') - space_B_string = 'theta' + if space_B == 'DG': + if degree == 0: + VB = domain.spaces('DG') + VCG1_B = FunctionSpace(mesh, 'CG', 1) + VDG1_B = domain.spaces('DG1_equispaced') + space_B_string = 'DG' + elif degree == 1: + VB = domain.spaces('DG') + space_B_string = 'DG' + else: + raise NotImplementedError + elif space_B == 'Vtheta': + if degree == 0: + VB = domain.spaces('theta') + VCG1_B = FunctionSpace(mesh, 'CG', 1) + VDG1_B = domain.spaces('DG1_equispaced') + space_B_string = 'theta' + elif degree == 1: + VB = domain.spaces('theta') + space_B_string = 'theta' + else: + raise NotImplementedError else: raise NotImplementedError @@ -128,32 +131,39 @@ def setup_limiters(dirname, space_A, space_B): sublimiters.update({'tracerA': VertexBasedLimiter(VA)}) elif space_A == 'Vtheta_degree_1': - opts = EmbeddedDGOptions() + suboptions.update({'tracerA': EmbeddedDGOptions()}) sublimiters.update({'tracerA': ThetaLimiter(VA)}) else: raise NotImplementedError # Options and limiters for tracer_B - - if space_B in ['DG0', 'Vtheta_degree_0']: - recover_opts = RecoveryOptions(embedding_space=VDG1_B, - recovered_space=VCG1_B, - project_low_method='recover', - boundary_method=BoundaryMethod.taylor) + + if space_B == 'DG': + if degree == 0: + suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B, + recovered_space=VCG1_B, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor)}) - sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)}) - - elif space_B == 'DG1': - sublimiters.update({'tracerB': DG1Limiter(VB)}) - - elif space_B == 'DG1_equispaced': - sublimiters.update({'tracerB': VertexBasedLimiter(VB)}) - - elif space_B == 'Vtheta_degree_1': - opts = EmbeddedDGOptions() - sublimiters.update({'tracerB': ThetaLimiter(VB)}) - + sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)}) + elif degree == 1: + sublimiters.update({'tracerB': DG1Limiter(VB)}) + else: + raise NotImplementedError + elif space_B == 'Vtheta': + if degree == 0: + suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B, + recovered_space=VCG1_B, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor)}) + + sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)}) + elif degree == 1: + opts = EmbeddedDGOptions() + sublimiters.update({'tracerB': ThetaLimiter(VB)}) + else: + raise NotImplementedError else: raise NotImplementedError @@ -162,6 +172,7 @@ def setup_limiters(dirname, space_A, space_B): # Give the scheme for the coupled transport transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter) + #transport_schemes = SSPRK3(domain, limiter=MixedLimiter) # DG Upwind transport for both tracers: transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] @@ -267,9 +278,8 @@ def setup_limiters(dirname, space_A, space_B): @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', 'DG1', 'DG1_equispaced']) -#Remove Dg1-dg1 and other easy ones after debugging -@pytest.mark.parametrize('space_B', ['Vtheta_degree_0'])#, 'Vtheta_degree_1', 'DG0', - #'DG1', 'DG1_equispaced']) +# It only makes sense to use the same degree for tracer B +@pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) def test_limiters(tmpdir, space_A, space_B): @@ -288,26 +298,26 @@ def test_limiters(tmpdir, space_A, space_B): # Check tracer is roughly in the correct place assert norm(true_fieldA - final_fieldA) / norm(true_fieldA) < 0.05, \ - 'Something is wrong with the DG space tracer using a mixed limiter' + f"Something is wrong with the {space_A} space tracer using a mixed limiter" # Check tracer is roughly in the correct place assert norm(true_fieldB - final_fieldB) / norm(true_fieldB) < 0.05, \ - 'Something is wrong with the DG1 equispaced tracer using a mixed limiter' + f"Something is wrong with the {space_B} space tracer using a mixed limiter" # Check for no new overshoots in A assert np.max(final_fieldA.dat.data) <= np.max(true_fieldA.dat.data) + tol, \ - 'Application of the DG space limiter in the mixed limiter has not prevented overshoots' + f"Application of the {space_A} space limiter in the mixed limiter has not prevented overshoots" # Check for no new undershoots in A assert np.min(final_fieldA.dat.data) >= np.min(true_fieldA.dat.data) - tol, \ - 'Application of the DG space limiter in the mixed limiter has not prevented undershoots' + f"Application of the {space_B} space limiter in the mixed limiter has not prevented undershoots" # Check for no new overshoots in B assert np.max(final_fieldB.dat.data) <= np.max(true_fieldB.dat.data) + tol, \ - 'Application of the DG1 equispaced limiter in the mixed limiter has not prevented overshoots' + f"Application of the {space_A} limiter in the mixed limiter has not prevented overshoots" # Check for no new undershoots in B assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \ - 'Application of the DG1 equispaced limiter in the mixed limiter has not prevented undershoots' + f"Application of the {space_B} equispaced limiter in the mixed limiter has not prevented undershoots" \ No newline at end of file From 616505790ecbe48cdbe4ebfa7478fca8598d21f0 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Tue, 9 Jan 2024 04:01:49 +0000 Subject: [PATCH 04/25] more mixed options changes --- gusto/time_discretisation.py | 44 +++++++++++- gusto/wrappers.py | 69 ++++++++++++------- .../transport/test_mixed_fs_options.py | 7 +- 3 files changed, 91 insertions(+), 29 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 593bbe830..c537b317f 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -56,7 +56,10 @@ def get_apply(self, x_out, x_in): def new_apply(self, x_out, x_in): self.wrapper.pre_apply(x_in) - original_apply(self, self.wrapper.x_out, self.wrapper.x_in) + if type(self.wrapper) == MixedOptions: + original_apply(self, x_out, x_in) + else: + original_apply(self, self.wrapper.x_out, self.wrapper.x_in) self.wrapper.post_apply(x_out) return new_apply(self, x_out, x_in) @@ -103,9 +106,39 @@ def __init__(self, domain, field_name=None, solver_parameters=None, if options is not None: if type(options) == MixedOptions: print('jahjah') + print(options) self.wrapper = options + #self.subwrappers = {} # Or do I need to initialise everything here # like is done for a single wrapper? + for field, suboption in self.wrapper.suboptions.items(): + print(field) + print(suboption) + + #Replace options with a wrapper? + + #if suboption.name == 'embedded_dg': + # self.suboptions[field].wrapper = EmbeddedDGWrapper(self, suboption) + #elif suboption.name == "recovered": + # self.suboptions[field].wrapper = RecoveryWrapper(self, suboption) + #elif suboption.name == "supg": + # self.suboptions[field].wrapper = SUPGWrapper(self, suboption) + #else: + # raise RuntimeError( + # f'Time discretisation: suboption wrapper {wrapper_name} not implemented') + if suboption.name == 'embedded_dg': + self.wrapper.subwrappers.update({field:EmbeddedDGWrapper(self, suboption)}) + elif suboption.name == "recovered": + self.wrapper.subwrappers.update({field:RecoveryWrapper(self, suboption)}) + elif suboption.name == "supg": + self.wrapper.subwrappers.update({field:SUPGWrapper(self, suboption)}) + else: + raise RuntimeError( + f'Time discretisation: suboption wrapper {wrapper_name} not implemented') + print(self.wrapper) + print(self.wrapper.suboptions) + print(self.wrapper.subwrappers) + #self.wrapper_name = 'mixed' else: self.wrapper_name = options.name if self.wrapper_name == "embedded_dg": @@ -180,10 +213,17 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.wrapper is not None: if type(self.wrapper) == MixedOptions: + #if self.wrapper_name == 'mixed': # Subwrappers are defined. # Set these up with ? - for _, subwrapper in self.wrapper.suboptions.items(): + for field, subwrapper in self.wrapper.subwrappers.items(): + #Set up field idxs here. + print(field) print(subwrapper) + self.wrapper.subwrappers[field].idx = equation.field_names.index(field) + self.wrapper.subwrappers[field].mixed_options = True + self.wrapper.subwrappers[field].setup() + self.fs = self.wrapper.subwrappers[field].function_space pass else: self.wrapper.setup() diff --git a/gusto/wrappers.py b/gusto/wrappers.py index b14a23b6b..52da12e4f 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -33,7 +33,9 @@ def __init__(self, time_discretisation, wrapper_options): self.time_discretisation = time_discretisation self.options = wrapper_options self.solver_parameters = None - self.field_name = None + #self.field_name = None + self.idx = None + self.mixed_options = False @abstractmethod def setup(self): @@ -162,8 +164,10 @@ class RecoveryWrapper(Wrapper): def setup(self): """Sets up function spaces and fields needed for this wrapper.""" + print(self.options) + assert isinstance(self.options, RecoveryOptions), \ - 'Embedded DG wrapper can only be used with Recovery Options' + 'Recovery wrapper can only be used with Recovery Options' original_space = self.time_discretisation.fs domain = self.time_discretisation.domain @@ -177,6 +181,7 @@ def setup(self): V_elt = BrokenElement(original_space.ufl_element()) self.function_space = FunctionSpace(domain.mesh, V_elt) else: + print('using embedded space') self.function_space = self.options.embedding_space self.test_space = self.function_space @@ -185,11 +190,23 @@ def setup(self): # Internal variables to be used # -------------------------------------------------------------------- # - self.x_in_tmp = Function(self.time_discretisation.fs) + if self.mixed_options == True: + self.x_in_tmp = Function(self.function_space) + else: + self.x_in_tmp = Function(self.time_discretisation.fs) + + print(self.function_space) + print(self.time_discretisation.fs) + self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) + + #self.x_out = Function(equation.function_space) + #self.x_out = Function(self.time_discretisation.fs) + if self.time_discretisation.idx is None: - self.x_projected = Function(equation.function_space) + #self.x_projected = Function(equation.function_space) + self.x_projected = Function(self.function_space) else: self.x_projected = Function(equation.spaces[self.time_discretisation.idx]) @@ -381,18 +398,20 @@ def __init__(self, equation, suboptions): ValueError: If an option is defined for a field that is not in the prognostic variable set """ print(equation.function_space) - self.x_in = Function(equation.function_space) - self.x_out = Function(equation.function_space) + #self.x_in = Function(equation.function_space) + #self.x_out = Function(equation.function_space) print('Initialising MixedOptions') self.suboptions = suboptions - print(suboptions.items()) + #print(suboptions.items()) - self.wrapper = {} + self.subwrappers = {} for field, suboption in suboptions.items(): + print(field) + print(suboption) # Check that the field is in the prognostic variable set: if field not in equation.field_names: raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") @@ -401,19 +420,19 @@ def __init__(self, equation, suboptions): wrapper_name = suboption.name print(wrapper_name) - if wrapper_name == "embedded_dg": - self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)}) - elif wrapper_name == "recovered": - self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption) - elif wrapper_name == "supg": - self.suboptions[field].subwrapper = SUPGWrapper(self, suboption) - else: - raise RuntimeError( - f'Time discretisation: suboption wrapper {wrapper_name} not implemented') + #if wrapper_name == "embedded_dg": + # self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)}) + #elif wrapper_name == "recovered": + # self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption) + #elif wrapper_name == "supg": + # self.suboptions[field].subwrapper = SUPGWrapper(self, suboption) + #else: + # raise RuntimeError( + # f'Time discretisation: suboption wrapper {wrapper_name} not implemented') #Initialise the wrapper and associate with a field: - - self.suboptions[field].idx = equation.field_names.index(field) + #self.suboptions[field].wrapper_name = suboption.name + #self.suboptions[field].idx = equation.field_names.index(field) #self.suboptions[field].fs = equation.field_names.function_space(field) # self.subwrapper.x_in = Function(self.suboptions[field].fs) @@ -426,8 +445,9 @@ def pre_apply(self, x_in): Perform the pre-applications from all subwrappers """ - for _, subwrapper in self.suboptions.items(): + for _, subwrapper in self.subwrappers.items(): print(subwrapper) + print('pre') field = x_in.subfunctions[subwrapper.idx] subwrapper.pre_apply(field) @@ -436,8 +456,9 @@ def post_apply(self, x_in): Perform the post-applications from all subwrappers """ - for _, suboption in self.suboptions.items(): - print(suboption) - field = x_in.subfunctions[suboption.idx] - suboption.subwrapper.post_apply(field) + for _, subwrapper in self.subwrappers.items(): + print(subwrapper) + print('post') + field = x_in.subfunctions[subwrapper.idx] + subwrapper.post_apply(field) diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index a82136923..9701e3dac 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -160,7 +160,7 @@ def setup_limiters(dirname, space_A, space_B): sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)}) elif degree == 1: - opts = EmbeddedDGOptions() + suboptions.update({'tracerB': EmbeddedDGOptions()}) sublimiters.update({'tracerB': ThetaLimiter(VB)}) else: raise NotImplementedError @@ -276,8 +276,9 @@ def setup_limiters(dirname, space_A, space_B): return stepper, tmax, true_fieldA, true_fieldB -@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', - 'DG1', 'DG1_equispaced']) +#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', +# 'DG1', 'DG1_equispaced']) +@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced']) # It only makes sense to use the same degree for tracer B @pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) From 775d3b2101c73f4b3a3de040de73a587d5010b10 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Wed, 10 Jan 2024 01:39:48 +0000 Subject: [PATCH 05/25] more changes to mixed options --- gusto/equations.py | 3 + gusto/limiters.py | 2 + gusto/time_discretisation.py | 33 ++++++++-- gusto/wrappers.py | 66 ++++++++++++++----- .../transport/test_mixed_fs_options.py | 9 +-- 5 files changed, 87 insertions(+), 26 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index bd61e795b..c30124c73 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -426,6 +426,8 @@ def add_tracers_to_prognostics(self, domain, active_tracers): else: self.space_names[tracer.name] = tracer.space self.spaces.append(domain.spaces(tracer.space)) + #print('self.spaces is', self.spaces) + print(len(self.spaces)) else: raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}') @@ -542,6 +544,7 @@ def __init__(self, domain, active_tracers, Vu=None): # Build finite element spaces self.spaces = [] + # Add active tracers to the list of prognostics if active_tracers is None: active_tracers = [] diff --git a/gusto/limiters.py b/gusto/limiters.py index a39c6688d..16a2d3b40 100644 --- a/gusto/limiters.py +++ b/gusto/limiters.py @@ -203,6 +203,8 @@ def __init__(self, equation, sublimiters): raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") else: self.sublimiters[field].idx = equation.field_names.index(field) + + # Give the sublimiter a subspace? def apply(self, fields): """ diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index c537b317f..ad45fd84c 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -57,7 +57,8 @@ def new_apply(self, x_out, x_in): self.wrapper.pre_apply(x_in) if type(self.wrapper) == MixedOptions: - original_apply(self, x_out, x_in) + #original_apply(self, x_out, x_in) + original_apply(self, self.wrapper.x_out, self.wrapper.x_in) else: original_apply(self, self.wrapper.x_out, self.wrapper.x_in) self.wrapper.post_apply(x_out) @@ -105,8 +106,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, if options is not None: if type(options) == MixedOptions: - print('jahjah') - print(options) + #print('jahjah') + #print(options) self.wrapper = options #self.subwrappers = {} # Or do I need to initialise everything here @@ -216,14 +217,36 @@ def setup(self, equation, apply_bcs=True, *active_labels): #if self.wrapper_name == 'mixed': # Subwrappers are defined. # Set these up with ? + + # Give more than one fs? + fields = [] + + print(self.wrapper.wrapper_spaces) + for field, subwrapper in self.wrapper.subwrappers.items(): - #Set up field idxs here. + # Set up field idxs here. print(field) print(subwrapper) self.wrapper.subwrappers[field].idx = equation.field_names.index(field) self.wrapper.subwrappers[field].mixed_options = True + + # Store the original space of the tracer + self.wrapper.subwrappers[field].tracer_fs = self.equation.spaces[equation.field_names.index(field)] self.wrapper.subwrappers[field].setup() - self.fs = self.wrapper.subwrappers[field].function_space + + # Update the function space to that needed by the wrapper + self.wrapper.wrapper_spaces.update({field: self.wrapper.subwrappers[field].function_space}) + fields.append(self.wrapper.subwrappers[field].function_space) + # Append the new function space from the wrapper + #print('self.tracer_fs is', self.wrapper.subwrappers[field].tracer_fs) + print(self.wrapper.wrapper_spaces) + #Setup the mixed wrapper + #yo_spaces = [] + #for _, space_name in self.wrapper.wrapper_spaces.items(): + # yo_spaces.append(self.domain.spaces(space_name)) + #print(yo_spaces) + self.wrapper.wrapper_spaces = fields + self.wrapper.setup() pass else: self.wrapper.setup() diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 52da12e4f..d4fa6d4dd 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -7,7 +7,7 @@ from abc import ABCMeta, abstractmethod from firedrake import ( FunctionSpace, Function, BrokenElement, Projector, Interpolator, - VectorElement, Constant, as_ufl, dot, grad, TestFunction + VectorElement, Constant, as_ufl, dot, grad, TestFunction, MixedFunctionSpace ) from firedrake.fml import Term from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions @@ -36,6 +36,7 @@ def __init__(self, time_discretisation, wrapper_options): #self.field_name = None self.idx = None self.mixed_options = False + self.tracer_fs = None @abstractmethod def setup(self): @@ -84,10 +85,16 @@ def setup(self): assert isinstance(self.options, EmbeddedDGOptions), \ 'Embedded DG wrapper can only be used with Embedded DG Options' - - original_space = self.time_discretisation.fs + domain = self.time_discretisation.domain equation = self.time_discretisation.equation + + if self.mixed_options == True: + #print('self.tracer_fs is', self.tracer_fs) + #print('self.time_discretisation.fs is', self.time_discretisation.fs) + original_space = self.tracer_fs + else: + original_space = self.time_discretisation.fs # -------------------------------------------------------------------- # # Set up spaces to be used with wrapper @@ -107,7 +114,9 @@ def setup(self): self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - if self.time_discretisation.idx is None: + if self.mixed_options == True: + self.x_projected = Function(self.tracer_fs) + elif self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) else: self.x_projected = Function(equation.spaces[self.time_discretisation.idx]) @@ -169,9 +178,13 @@ def setup(self): assert isinstance(self.options, RecoveryOptions), \ 'Recovery wrapper can only be used with Recovery Options' - original_space = self.time_discretisation.fs domain = self.time_discretisation.domain equation = self.time_discretisation.equation + + if self.mixed_options == True: + original_space = self.tracer_fs + else: + original_space = self.time_discretisation.fs # -------------------------------------------------------------------- # # Set up spaces to be used with wrapper @@ -191,12 +204,13 @@ def setup(self): # -------------------------------------------------------------------- # if self.mixed_options == True: - self.x_in_tmp = Function(self.function_space) + self.x_in_tmp = Function(self.tracer_fs) + #self.x_in_tmp = Function(equation.spaces[self.idx]) else: self.x_in_tmp = Function(self.time_discretisation.fs) - print(self.function_space) - print(self.time_discretisation.fs) + #print(self.function_space) + #print(self.time_discretisation.fs) self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) @@ -204,9 +218,10 @@ def setup(self): #self.x_out = Function(equation.function_space) #self.x_out = Function(self.time_discretisation.fs) - if self.time_discretisation.idx is None: - #self.x_projected = Function(equation.function_space) - self.x_projected = Function(self.function_space) + if self.mixed_options == True: + self.x_projected = Function(self.tracer_fs) + elif self.time_discretisation.idx is None: + self.x_projected = Function(equation.function_space) else: self.x_projected = Function(equation.spaces[self.time_discretisation.idx]) @@ -397,7 +412,14 @@ def __init__(self, equation, suboptions): Raises: ValueError: If an option is defined for a field that is not in the prognostic variable set """ - print(equation.function_space) + #print(equation.function_space) + #print(equation.active_tracers) + #print(equation.active_tracers[0]) + #print(equation.active_tracers[1]) + print(equation.space_names) + + self.wrapper_spaces = equation.space_names + #self.x_in = Function(equation.function_space) #self.x_out = Function(equation.function_space) @@ -412,13 +434,17 @@ def __init__(self, equation, suboptions): for field, suboption in suboptions.items(): print(field) print(suboption) - # Check that the field is in the prognostic variable set: + + # Check that the field is in the prognostic variable set: if field not in equation.field_names: raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") else: # check that a valid wrapper has been given wrapper_name = suboption.name - print(wrapper_name) + #print(wrapper_name) + + # Extract the space? + #if wrapper_name == "embedded_dg": # self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)}) @@ -438,7 +464,13 @@ def __init__(self, equation, suboptions): def setup(self): # This is done in the suboption wrappers themselves - pass + # Or, determine the new mixed function space + self.function_space = MixedFunctionSpace(self.wrapper_spaces) + self.x_in = Function(self.function_space) + self.x_out = Function(self.function_space) + + + #pass def pre_apply(self, x_in): """ @@ -450,6 +482,7 @@ def pre_apply(self, x_in): print('pre') field = x_in.subfunctions[subwrapper.idx] subwrapper.pre_apply(field) + print('success') def post_apply(self, x_in): """ @@ -460,5 +493,4 @@ def post_apply(self, x_in): print(subwrapper) print('post') field = x_in.subfunctions[subwrapper.idx] - subwrapper.post_apply(field) - + subwrapper.post_apply(field) \ No newline at end of file diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 9701e3dac..699618a64 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -171,6 +171,7 @@ def setup_limiters(dirname, space_A, space_B): MixedLimiter = MixedFSLimiter(eqn, sublimiters) # Give the scheme for the coupled transport + #transport_schemes = SSPRK3(domain, options=opts) transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter) #transport_schemes = SSPRK3(domain, limiter=MixedLimiter) @@ -276,14 +277,14 @@ def setup_limiters(dirname, space_A, space_B): return stepper, tmax, true_fieldA, true_fieldB -#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', -# 'DG1', 'DG1_equispaced']) -@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced']) +@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', + 'DG1', 'DG1_equispaced']) +#@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced']) # It only makes sense to use the same degree for tracer B @pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) -def test_limiters(tmpdir, space_A, space_B): +def test_mixed_fs_options(tmpdir, space_A, space_B): # Setup and run dirname = str(tmpdir) From 793768ef7919b8832dfb4cc5cd495982c3ae83ab Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Wed, 10 Jan 2024 22:58:18 +0000 Subject: [PATCH 06/25] working on new test functions with mixed options --- gusto/time_discretisation.py | 97 ++++++++++++++++++------------------ gusto/wrappers.py | 42 +++++++++++----- 2 files changed, 78 insertions(+), 61 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index ad45fd84c..90f73a502 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -37,30 +37,21 @@ def wrapper_apply(original_apply): """Decorator to add steps for using a wrapper around the apply method.""" def get_apply(self, x_out, x_in): - #if type(self.wrapper) == MixedOptions: - - # Subwrappers for different tracers have been defined - # Use the apply defined in the MixedOptions object - - # print('aha') - # self.wrapper.pre_apply(x_in) - - # for _, subwrapper in self.wrapper.suboptions.items(): - # original_apply(self, self.subwrapper.x_out, self.subwrapper.x_in) - - # self.wrapper.post_apply(x_out) - - if self.wrapper is not None: def new_apply(self, x_out, x_in): self.wrapper.pre_apply(x_in) - if type(self.wrapper) == MixedOptions: + original_apply(self, self.wrapper.x_out, self.wrapper.x_in) + + #if type(self.wrapper) == MixedOptions: + #print('x_out', x_out) + #print('self.wrapper.x_out', self.wrapper.x_out) #original_apply(self, x_out, x_in) - original_apply(self, self.wrapper.x_out, self.wrapper.x_in) - else: - original_apply(self, self.wrapper.x_out, self.wrapper.x_in) + # original_apply(self, self.wrapper.x_out, self.wrapper.x_in) + #else: + # original_apply(self, self.wrapper.x_out, self.wrapper.x_in) + self.wrapper.post_apply(x_out) return new_apply(self, x_out, x_in) @@ -116,17 +107,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None, print(field) print(suboption) - #Replace options with a wrapper? - - #if suboption.name == 'embedded_dg': - # self.suboptions[field].wrapper = EmbeddedDGWrapper(self, suboption) - #elif suboption.name == "recovered": - # self.suboptions[field].wrapper = RecoveryWrapper(self, suboption) - #elif suboption.name == "supg": - # self.suboptions[field].wrapper = SUPGWrapper(self, suboption) - #else: - # raise RuntimeError( - # f'Time discretisation: suboption wrapper {wrapper_name} not implemented') if suboption.name == 'embedded_dg': self.wrapper.subwrappers.update({field:EmbeddedDGWrapper(self, suboption)}) elif suboption.name == "recovered": @@ -136,9 +116,9 @@ def __init__(self, domain, field_name=None, solver_parameters=None, else: raise RuntimeError( f'Time discretisation: suboption wrapper {wrapper_name} not implemented') - print(self.wrapper) - print(self.wrapper.suboptions) - print(self.wrapper.subwrappers) + #print(self.wrapper) + #print(self.wrapper.suboptions) + #print(self.wrapper.subwrappers) #self.wrapper_name = 'mixed' else: self.wrapper_name = options.name @@ -219,35 +199,54 @@ def setup(self, equation, apply_bcs=True, *active_labels): # Set these up with ? # Give more than one fs? - fields = [] + #fields = [] print(self.wrapper.wrapper_spaces) for field, subwrapper in self.wrapper.subwrappers.items(): - # Set up field idxs here. - print(field) - print(subwrapper) - self.wrapper.subwrappers[field].idx = equation.field_names.index(field) + field_idx = equation.field_names.index(field) + self.wrapper.subwrappers[field].idx = field_idx self.wrapper.subwrappers[field].mixed_options = True # Store the original space of the tracer self.wrapper.subwrappers[field].tracer_fs = self.equation.spaces[equation.field_names.index(field)] + self.wrapper.subwrappers[field].setup() # Update the function space to that needed by the wrapper - self.wrapper.wrapper_spaces.update({field: self.wrapper.subwrappers[field].function_space}) - fields.append(self.wrapper.subwrappers[field].function_space) - # Append the new function space from the wrapper - #print('self.tracer_fs is', self.wrapper.subwrappers[field].tracer_fs) - print(self.wrapper.wrapper_spaces) - #Setup the mixed wrapper - #yo_spaces = [] - #for _, space_name in self.wrapper.wrapper_spaces.items(): - # yo_spaces.append(self.domain.spaces(space_name)) - #print(yo_spaces) - self.wrapper.wrapper_spaces = fields + self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space + + # Replace with new test function + # STILL TO SORT OUT. + #new_test = TestFunction(self.wrapper.subwrappers[field].test_space) + + #self.residual = self.residual.label_map( + # all_terms, + # map_if_true=replace_test_function(new_test, old_idx=field_idx)) + + self.wrapper.setup() - pass + + # Check if mixed function spcae has changed: + if self.wrapper.function_space == equation.function_space: + print('same') + else: + print('different') + + self.fs = self.wrapper.function_space + + # Or replace test functions here?? + new_test = TestFunction(self.wrapper.test_space) + + self.residual = self.residual.label_map( + all_terms, + map_if_true=replace_test_function(new_test)) + + # Only call this if with SUPG, so should put this into the + # previous section + self.residual = self.wrapper.label_terms(self.residual) + + else: self.wrapper.setup() self.fs = self.wrapper.function_space diff --git a/gusto/wrappers.py b/gusto/wrappers.py index d4fa6d4dd..0e6c655d5 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -33,7 +33,6 @@ def __init__(self, time_discretisation, wrapper_options): self.time_discretisation = time_discretisation self.options = wrapper_options self.solver_parameters = None - #self.field_name = None self.idx = None self.mixed_options = False self.tracer_fs = None @@ -93,6 +92,7 @@ def setup(self): #print('self.tracer_fs is', self.tracer_fs) #print('self.time_discretisation.fs is', self.time_discretisation.fs) original_space = self.tracer_fs + #original_space = self.time_discretisation.fs else: original_space = self.time_discretisation.fs @@ -114,6 +114,7 @@ def setup(self): self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) + if self.mixed_options == True: self.x_projected = Function(self.tracer_fs) elif self.time_discretisation.idx is None: @@ -182,7 +183,8 @@ def setup(self): equation = self.time_discretisation.equation if self.mixed_options == True: - original_space = self.tracer_fs + #original_space = self.tracer_fs + original_space = self.time_discretisation.fs else: original_space = self.time_discretisation.fs @@ -416,9 +418,12 @@ def __init__(self, equation, suboptions): #print(equation.active_tracers) #print(equation.active_tracers[0]) #print(equation.active_tracers[1]) - print(equation.space_names) + #print(equation.space_names) + #print(equation.spaces) - self.wrapper_spaces = equation.space_names + #self.wrapper_spaces = equation.space_names + self.wrapper_spaces = equation.spaces + print(len(equation.spaces)) #self.x_in = Function(equation.function_space) #self.x_out = Function(equation.function_space) @@ -432,8 +437,8 @@ def __init__(self, equation, suboptions): self.subwrappers = {} for field, suboption in suboptions.items(): - print(field) - print(suboption) + #print(field) + #print(suboption) # Check that the field is in the prognostic variable set: if field not in equation.field_names: @@ -465,32 +470,45 @@ def __init__(self, equation, suboptions): def setup(self): # This is done in the suboption wrappers themselves # Or, determine the new mixed function space + + #Loop over all active variables? self.function_space = MixedFunctionSpace(self.wrapper_spaces) self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - - + self.test_space = self.function_space #pass def pre_apply(self, x_in): """ Perform the pre-applications from all subwrappers """ + #self.x_in.assign(x_in) + self.x_in = x_in for _, subwrapper in self.subwrappers.items(): print(subwrapper) print('pre') field = x_in.subfunctions[subwrapper.idx] subwrapper.pre_apply(field) - print('success') + #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in + + x_in_sub = self.x_in.subfunctions[subwrapper.idx] + #x_in_sub.assign(subwrapper.x_in) + x_in_sub = subwrapper.x_in + #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in - def post_apply(self, x_in): + + def post_apply(self, x_out): """ Perform the post-applications from all subwrappers """ + x_out.assign(self.x_out) for _, subwrapper in self.subwrappers.items(): print(subwrapper) print('post') - field = x_in.subfunctions[subwrapper.idx] - subwrapper.post_apply(field) \ No newline at end of file + field = x_out.subfunctions[subwrapper.idx] + subwrapper.post_apply(field) + + x_out_sub = x_out.subfunctions[subwrapper.idx] + x_out_sub.assign(subwrapper.x_out) \ No newline at end of file From 5f05367a13b13fd92fda235e3c92295140447881 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Thu, 11 Jan 2024 04:09:03 +0000 Subject: [PATCH 07/25] more changes to test functions for multiple wrappers --- gusto/time_discretisation.py | 16 ++++++++-------- gusto/wrappers.py | 4 ++-- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 90f73a502..e77dff92e 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -216,14 +216,14 @@ def setup(self, equation, apply_bcs=True, *active_labels): # Update the function space to that needed by the wrapper self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space - # Replace with new test function - # STILL TO SORT OUT. - #new_test = TestFunction(self.wrapper.subwrappers[field].test_space) - - #self.residual = self.residual.label_map( - # all_terms, - # map_if_true=replace_test_function(new_test, old_idx=field_idx)) + # Replace the test function space + # This currently won't work: supg test + # is a function, not a space + if self.wrapper.suboptions[field].name == "supg": + self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test + else: + self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test_space self.wrapper.setup() @@ -244,7 +244,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): # Only call this if with SUPG, so should put this into the # previous section - self.residual = self.wrapper.label_terms(self.residual) + # self.residual = self.wrapper.label_terms(self.residual) else: diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 0e6c655d5..052ddbc2e 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -423,6 +423,7 @@ def __init__(self, equation, suboptions): #self.wrapper_spaces = equation.space_names self.wrapper_spaces = equation.spaces + self.test_spaces = equation.spaces print(len(equation.spaces)) #self.x_in = Function(equation.function_space) @@ -475,8 +476,7 @@ def setup(self): self.function_space = MixedFunctionSpace(self.wrapper_spaces) self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - self.test_space = self.function_space - #pass + self.test_space = MixedFunctionSpace(self.test_spaces) def pre_apply(self, x_in): """ From 5c13db13af2e641d5cb74390f0a020ef170cf8b4 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Tue, 16 Jan 2024 22:26:08 +0000 Subject: [PATCH 08/25] more changes to mixed options with replacing test functions --- gusto/time_discretisation.py | 93 +++++++--------- gusto/wrappers.py | 105 ++++++++++-------- .../transport/test_mixed_fs_options.py | 10 +- 3 files changed, 102 insertions(+), 106 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index e77dff92e..8410a0677 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -21,7 +21,7 @@ from gusto.configuration import EmbeddedDGOptions, RecoveryOptions from gusto.labels import (time_derivative, prognostic, physics_label, - implicit, explicit) + implicit, explicit, transport) from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * @@ -43,15 +43,6 @@ def new_apply(self, x_out, x_in): self.wrapper.pre_apply(x_in) original_apply(self, self.wrapper.x_out, self.wrapper.x_in) - - #if type(self.wrapper) == MixedOptions: - #print('x_out', x_out) - #print('self.wrapper.x_out', self.wrapper.x_out) - #original_apply(self, x_out, x_in) - # original_apply(self, self.wrapper.x_out, self.wrapper.x_in) - #else: - # original_apply(self, self.wrapper.x_out, self.wrapper.x_in) - self.wrapper.post_apply(x_out) return new_apply(self, x_out, x_in) @@ -97,12 +88,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, if options is not None: if type(options) == MixedOptions: - #print('jahjah') - #print(options) self.wrapper = options - #self.subwrappers = {} - # Or do I need to initialise everything here - # like is done for a single wrapper? + for field, suboption in self.wrapper.suboptions.items(): print(field) print(suboption) @@ -116,10 +103,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, else: raise RuntimeError( f'Time discretisation: suboption wrapper {wrapper_name} not implemented') - #print(self.wrapper) - #print(self.wrapper.suboptions) - #print(self.wrapper.subwrappers) - #self.wrapper_name = 'mixed' + else: self.wrapper_name = options.name if self.wrapper_name == "embedded_dg": @@ -194,58 +178,61 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.wrapper is not None: if type(self.wrapper) == MixedOptions: - #if self.wrapper_name == 'mixed': - # Subwrappers are defined. - # Set these up with ? - - # Give more than one fs? - #fields = [] - - print(self.wrapper.wrapper_spaces) + + #print(self.wrapper.wrapper_spaces) for field, subwrapper in self.wrapper.subwrappers.items(): field_idx = equation.field_names.index(field) - self.wrapper.subwrappers[field].idx = field_idx - self.wrapper.subwrappers[field].mixed_options = True + + subwrapper.idx = field_idx + subwrapper.mixed_options = True # Store the original space of the tracer - self.wrapper.subwrappers[field].tracer_fs = self.equation.spaces[equation.field_names.index(field)] + subwrapper.tracer_fs = self.equation.spaces[field_idx] - self.wrapper.subwrappers[field].setup() + subwrapper.setup() # Update the function space to that needed by the wrapper - self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space + self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space + # Store test space? + #self.wrapper.test_spaces[field_idx] = subwrapper.function_space # Replace the test function space - # This currently won't work: supg test - # is a function, not a space if self.wrapper.suboptions[field].name == "supg": - self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test + new_test = subwrapper.test else: - self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test_space + new_test = TestFunction(subwrapper.test_space) + + self.residual = self.residual.label_map( + lambda t: t.has_label(transport) and t.get(prognostic) == field, + map_if_true=replace_test_function(new_test, old_idx=field_idx)) + + self.residual = subwrapper.label_terms(self.residual) + + # Currently can only use one set of solver parameters ... + if self.solver_parameters is None: + self.solver_parameters = subwrapper.solver_parameters self.wrapper.setup() - - # Check if mixed function spcae has changed: - if self.wrapper.function_space == equation.function_space: - print('same') - else: - print('different') self.fs = self.wrapper.function_space - # Or replace test functions here?? - new_test = TestFunction(self.wrapper.test_space) - - self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test)) - - # Only call this if with SUPG, so should put this into the - # previous section - # self.residual = self.wrapper.label_terms(self.residual) - + #new_test_mixed = TestFunction(self.fs) + + #for field, subwrapper in self.wrapper.subwrappers.items(): + # field_idx = equation.field_names.index(field) + + # if self.wrapper.suboptions[field].name == "supg": + # new_test_mixed[field_idx] = subwrapper.test + + # self.residual = self.residual.label_map( + # lambda t: t.has_label(transport) and t.get(prognostic) == field, + # map_if_true=replace_test_function(new_test_mixed[field_idx], old_idx=field_idx)) + + # self.residual = subwrapper.label_terms(self.residual) + + else: self.wrapper.setup() diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 052ddbc2e..4508ffde3 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -414,39 +414,21 @@ def __init__(self, equation, suboptions): Raises: ValueError: If an option is defined for a field that is not in the prognostic variable set """ - #print(equation.function_space) - #print(equation.active_tracers) - #print(equation.active_tracers[0]) - #print(equation.active_tracers[1]) - #print(equation.space_names) - #print(equation.spaces) - - #self.wrapper_spaces = equation.space_names self.wrapper_spaces = equation.spaces self.test_spaces = equation.spaces - print(len(equation.spaces)) - - #self.x_in = Function(equation.function_space) - #self.x_out = Function(equation.function_space) - - print('Initialising MixedOptions') - + self.field_names = equation.field_names self.suboptions = suboptions - - #print(suboptions.items()) - self.subwrappers = {} for field, suboption in suboptions.items(): - #print(field) - #print(suboption) # Check that the field is in the prognostic variable set: if field not in equation.field_names: raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") - else: + #else: # check that a valid wrapper has been given - wrapper_name = suboption.name + + # wrapper_name = suboption.name #print(wrapper_name) # Extract the space? @@ -469,10 +451,8 @@ def __init__(self, equation, suboptions): # self.subwrapper.x_in = Function(self.suboptions[field].fs) def setup(self): - # This is done in the suboption wrappers themselves - # Or, determine the new mixed function space - #Loop over all active variables? + # Compute the new mixed function space self.function_space = MixedFunctionSpace(self.wrapper_spaces) self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) @@ -480,35 +460,64 @@ def setup(self): def pre_apply(self, x_in): """ - Perform the pre-applications from all subwrappers + Perform the pre-applications for all fields + with an associated subwrapper. """ - #self.x_in.assign(x_in) - self.x_in = x_in - - for _, subwrapper in self.subwrappers.items(): - print(subwrapper) - print('pre') - field = x_in.subfunctions[subwrapper.idx] - subwrapper.pre_apply(field) - #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in - x_in_sub = self.x_in.subfunctions[subwrapper.idx] - #x_in_sub.assign(subwrapper.x_in) - x_in_sub = subwrapper.x_in - #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in + for field_name in self.field_names: + #print(field_name) + #print(self.subwrappers) + + field_idx = self.field_names.index(field_name) + #print(field_idx) + + field = x_in.subfunctions[field_idx] + x_in_sub = self.x_in.subfunctions[field_idx] + + if field_name in self.subwrappers: + subwrapper = self.subwrappers[field_name] + print(subwrapper) + print('pre') + #field = x_in.subfunctions[subwrapper.idx] + subwrapper.pre_apply(field) + #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in + + #x_in_sub = self.x_in.subfunctions[subwrapper.idx] + x_in_sub.assign(subwrapper.x_in) + #x_in_sub = subwrapper.x_in + #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in + else: + x_in_sub.assign(field) def post_apply(self, x_out): """ - Perform the post-applications from all subwrappers + Perform the post-applications for all fields + with an associated subwrapper. """ - x_out.assign(self.x_out) + #x_out.assign(self.x_out) - for _, subwrapper in self.subwrappers.items(): - print(subwrapper) - print('post') - field = x_out.subfunctions[subwrapper.idx] - subwrapper.post_apply(field) + #for _, subwrapper in self.subwrappers.items(): + # print(subwrapper) + # print('post') + # field = x_out.subfunctions[subwrapper.idx] + # subwrapper.post_apply(field) - x_out_sub = x_out.subfunctions[subwrapper.idx] - x_out_sub.assign(subwrapper.x_out) \ No newline at end of file + # x_out_sub = x_out.subfunctions[subwrapper.idx] + # x_out_sub.assign(subwrapper.x_out) + + for field_name in self.field_names: + + field_idx = self.field_names.index(field_name) + + field = self.x_out.subfunctions[field_idx] + x_out_sub = x_out.subfunctions[field_idx] + + if field_name in self.subwrappers: + subwrapper = self.subwrappers[field_name] + print(subwrapper) + print('post') + subwrapper.post_apply(field) + x_out_sub.assign(subwrapper.x_out) + else: + x_out_sub.assign(field) \ No newline at end of file diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 699618a64..5e6c4739b 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -64,25 +64,23 @@ def setup_limiters(dirname, space_A, space_B): # Tracer B spaces if space_B == 'DG': + space_B_string = 'DG' if degree == 0: VB = domain.spaces('DG') VCG1_B = FunctionSpace(mesh, 'CG', 1) VDG1_B = domain.spaces('DG1_equispaced') - space_B_string = 'DG' elif degree == 1: VB = domain.spaces('DG') - space_B_string = 'DG' else: raise NotImplementedError elif space_B == 'Vtheta': + space_B_string = 'theta' if degree == 0: VB = domain.spaces('theta') VCG1_B = FunctionSpace(mesh, 'CG', 1) VDG1_B = domain.spaces('DG1_equispaced') - space_B_string = 'theta' elif degree == 1: VB = domain.spaces('theta') - space_B_string = 'theta' else: raise NotImplementedError else: @@ -178,6 +176,9 @@ def setup_limiters(dirname, space_A, space_B): # DG Upwind transport for both tracers: transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] + # Need to give SUPG options to the above, if using supg ... + # Need to test SUPG here! + # Build time stepper stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method) @@ -279,7 +280,6 @@ def setup_limiters(dirname, space_A, space_B): @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', 'DG1', 'DG1_equispaced']) -#@pytest.mark.parametrize('space_A', ['DG0'])#, 'DG1_equispaced']) # It only makes sense to use the same degree for tracer B @pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) From 806a49ef5701c33660d6916ac48622b6f7194291 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Wed, 17 Jan 2024 21:38:21 +0000 Subject: [PATCH 09/25] mixed options works for embedded dg and recovery --- gusto/equations.py | 2 - gusto/limiters.py | 2 - gusto/time_discretisation.py | 50 +++++++----- gusto/wrappers.py | 81 +++++-------------- .../transport/test_mixed_fs_options.py | 21 ++++- 5 files changed, 70 insertions(+), 86 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index c30124c73..1f3b90afd 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -426,8 +426,6 @@ def add_tracers_to_prognostics(self, domain, active_tracers): else: self.space_names[tracer.name] = tracer.space self.spaces.append(domain.spaces(tracer.space)) - #print('self.spaces is', self.spaces) - print(len(self.spaces)) else: raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}') diff --git a/gusto/limiters.py b/gusto/limiters.py index 16a2d3b40..a39c6688d 100644 --- a/gusto/limiters.py +++ b/gusto/limiters.py @@ -203,8 +203,6 @@ def __init__(self, equation, sublimiters): raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") else: self.sublimiters[field].idx = equation.field_names.index(field) - - # Give the sublimiter a subspace? def apply(self, fields): """ diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 8410a0677..d883ae216 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -10,7 +10,7 @@ import numpy as np from firedrake import ( - Function, TestFunction, NonlinearVariationalProblem, + Function, TestFunction, TestFunctions, NonlinearVariationalProblem, NonlinearVariationalSolver, DirichletBC, split, Constant ) from firedrake.fml import ( @@ -199,39 +199,53 @@ def setup(self, equation, apply_bcs=True, *active_labels): #self.wrapper.test_spaces[field_idx] = subwrapper.function_space # Replace the test function space - if self.wrapper.suboptions[field].name == "supg": - new_test = subwrapper.test - else: - new_test = TestFunction(subwrapper.test_space) + #if self.wrapper.suboptions[field].name == "supg": + # new_test = subwrapper.test + #else: + # new_test = TestFunction(subwrapper.test_space) - self.residual = self.residual.label_map( - lambda t: t.has_label(transport) and t.get(prognostic) == field, - map_if_true=replace_test_function(new_test, old_idx=field_idx)) + #self.residual = self.residual.label_map( + # lambda t: t.has_label(transport) and t.get(prognostic) == field, + # map_if_true=replace_test_function(new_test, old_idx=field_idx)) - self.residual = subwrapper.label_terms(self.residual) - - # Currently can only use one set of solver parameters ... - if self.solver_parameters is None: - self.solver_parameters = subwrapper.solver_parameters + #self.residual = subwrapper.label_terms(self.residual) self.wrapper.setup() self.fs = self.wrapper.function_space - #new_test_mixed = TestFunction(self.fs) + new_test_mixed = TestFunctions(self.fs) - #for field, subwrapper in self.wrapper.subwrappers.items(): - # field_idx = equation.field_names.index(field) + # If use supg, then change the required test function - # if self.wrapper.suboptions[field].name == "supg": - # new_test_mixed[field_idx] = subwrapper.test + # Replace one-by-one in this case? + + #for field, subwrapper in self.wrapper.subwrappers.items(): + # if self.wrapper.suboptions[field].name == "supg": + # field_idx = equation.field_names.index(field) + # test_list = list(new_test_mixed) + # test_list[field_idx] = subwrapper.test + # new_test_mixed = tuple(test_list) + # Change test functions with the new space: + #for field, subwrapper in self.wrapper.subwrappers.items(): # self.residual = self.residual.label_map( # lambda t: t.has_label(transport) and t.get(prognostic) == field, # map_if_true=replace_test_function(new_test_mixed[field_idx], old_idx=field_idx)) # self.residual = subwrapper.label_terms(self.residual) + # Or, all-at-once: + #self.residual = self.residual.label_map( + # lambda t: t.has_label(transport), + # map_if_true=replace_test_function(new_test_mixed)) + + self.residual = self.residual.label_map( + all_terms, + map_if_true=replace_test_function(new_test_mixed)) + + + else: diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 4508ffde3..956f409d2 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -89,10 +89,7 @@ def setup(self): equation = self.time_discretisation.equation if self.mixed_options == True: - #print('self.tracer_fs is', self.tracer_fs) - #print('self.time_discretisation.fs is', self.time_discretisation.fs) original_space = self.tracer_fs - #original_space = self.time_discretisation.fs else: original_space = self.time_discretisation.fs @@ -207,19 +204,12 @@ def setup(self): if self.mixed_options == True: self.x_in_tmp = Function(self.tracer_fs) - #self.x_in_tmp = Function(equation.spaces[self.idx]) else: self.x_in_tmp = Function(self.time_discretisation.fs) - #print(self.function_space) - #print(self.time_discretisation.fs) - self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - #self.x_out = Function(equation.function_space) - #self.x_out = Function(self.time_discretisation.fs) - if self.mixed_options == True: self.x_projected = Function(self.tracer_fs) elif self.time_discretisation.idx is None: @@ -306,9 +296,23 @@ def setup(self): 'SUPG wrapper can only be used with SUPG Options' domain = self.time_discretisation.domain + + #if self.mixed_options == True: + # self.function_space = self.tracer_fs + #else: + # self.function_space = self.time_discretisation.fs + self.function_space = self.time_discretisation.fs - self.test_space = self.function_space - self.x_out = Function(self.function_space) + + if self.mixed_options == True: + self.test_space = self.tracer_fs + self.x_out = Function(self.tracer_fs) + else: + self.test_space = self.function_space + self.x_out = Function(self.function_space) + + #self.test_space = self.function_space + #self.x_out = Function(self.function_space) # -------------------------------------------------------------------- # # Work out SUPG parameter @@ -415,40 +419,15 @@ def __init__(self, equation, suboptions): ValueError: If an option is defined for a field that is not in the prognostic variable set """ self.wrapper_spaces = equation.spaces - self.test_spaces = equation.spaces + #self.test_spaces = equation.spaces self.field_names = equation.field_names self.suboptions = suboptions self.subwrappers = {} for field, suboption in suboptions.items(): - # Check that the field is in the prognostic variable set: if field not in equation.field_names: raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") - #else: - # check that a valid wrapper has been given - - # wrapper_name = suboption.name - #print(wrapper_name) - - # Extract the space? - - - #if wrapper_name == "embedded_dg": - # self.wrapper.update({field:EmbeddedDGWrapper(self, suboption)}) - #elif wrapper_name == "recovered": - # self.suboptions[field].subwrapper = RecoveryWrapper(self, suboption) - #elif wrapper_name == "supg": - # self.suboptions[field].subwrapper = SUPGWrapper(self, suboption) - #else: - # raise RuntimeError( - # f'Time discretisation: suboption wrapper {wrapper_name} not implemented') - - #Initialise the wrapper and associate with a field: - #self.suboptions[field].wrapper_name = suboption.name - #self.suboptions[field].idx = equation.field_names.index(field) - #self.suboptions[field].fs = equation.field_names.function_space(field) - # self.subwrapper.x_in = Function(self.suboptions[field].fs) def setup(self): @@ -456,7 +435,7 @@ def setup(self): self.function_space = MixedFunctionSpace(self.wrapper_spaces) self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - self.test_space = MixedFunctionSpace(self.test_spaces) + #self.test_space = MixedFunctionSpace(self.test_spaces) def pre_apply(self, x_in): """ @@ -476,16 +455,8 @@ def pre_apply(self, x_in): if field_name in self.subwrappers: subwrapper = self.subwrappers[field_name] - print(subwrapper) - print('pre') - #field = x_in.subfunctions[subwrapper.idx] subwrapper.pre_apply(field) - #x_in.subfunctions[subwrapper.idx] = subwrapper.x_in - - #x_in_sub = self.x_in.subfunctions[subwrapper.idx] x_in_sub.assign(subwrapper.x_in) - #x_in_sub = subwrapper.x_in - #self.x_in.subfunctions[subwrapper.idx] = subwrapper.x_in else: x_in_sub.assign(field) @@ -495,16 +466,6 @@ def post_apply(self, x_out): Perform the post-applications for all fields with an associated subwrapper. """ - #x_out.assign(self.x_out) - - #for _, subwrapper in self.subwrappers.items(): - # print(subwrapper) - # print('post') - # field = x_out.subfunctions[subwrapper.idx] - # subwrapper.post_apply(field) - - # x_out_sub = x_out.subfunctions[subwrapper.idx] - # x_out_sub.assign(subwrapper.x_out) for field_name in self.field_names: @@ -515,9 +476,7 @@ def post_apply(self, x_out): if field_name in self.subwrappers: subwrapper = self.subwrappers[field_name] - print(subwrapper) - print('post') - subwrapper.post_apply(field) - x_out_sub.assign(subwrapper.x_out) + subwrapper.x_out.assign(field) + subwrapper.post_apply(x_out_sub) else: x_out_sub.assign(field) \ No newline at end of file diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 5e6c4739b..301c260e5 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -59,6 +59,12 @@ def setup_limiters(dirname, space_A, space_B): elif space_A == 'Vtheta_degree_1': VA = domain.spaces('theta') space_A_string = 'theta' + elif space_A == 'HDiv': + VA = domain.spaces('HDiv') + space_A_string = 'HDiv' + #elif space_A == 'CG1' + #VA = FunctionSpace(mesh, 'CG', 1) + #space_A_string = 'CG' else: raise NotImplementedError @@ -131,7 +137,9 @@ def setup_limiters(dirname, space_A, space_B): elif space_A == 'Vtheta_degree_1': suboptions.update({'tracerA': EmbeddedDGOptions()}) sublimiters.update({'tracerA': ThetaLimiter(VA)}) - + elif space_A == 'HDiv': + ibp_A = IntegrateByParts.TWICE + suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)}) else: raise NotImplementedError @@ -174,7 +182,13 @@ def setup_limiters(dirname, space_A, space_B): #transport_schemes = SSPRK3(domain, limiter=MixedLimiter) # DG Upwind transport for both tracers: - transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] + if space_A == 'HDiv': + transport_method_A = DGUpwind(eqn, 'tracerA', ibp=ibp_A) + else: + transport_method_A = DGUpwind(eqn, 'tracerA') + transport_method_B = DGUpwind(eqn, 'tracerB') + + transport_method = [transport_method_A, transport_method_B] # Need to give SUPG options to the above, if using supg ... # Need to test SUPG here! @@ -278,8 +292,9 @@ def setup_limiters(dirname, space_A, space_B): return stepper, tmax, true_fieldA, true_fieldB +#@pytest.mark.parametrize('space_A', ['HDiv'])#, @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', - 'DG1', 'DG1_equispaced']) + 'DG1', 'DG1_equispaced', 'HDiv']) # It only makes sense to use the same degree for tracer B @pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) From 2be135486d0ef1d29e0648d25658bcff3967027b Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Mon, 22 Jan 2024 02:09:55 +0000 Subject: [PATCH 10/25] moved DG1-DG1 equispaced test from test_limiters to test_mixed_fs_options --- gusto/equations.py | 1 - gusto/time_discretisation.py | 50 +++---- gusto/wrappers.py | 11 +- integration-tests/transport/test_limiters.py | 124 ++++-------------- .../transport/test_mixed_fs_options.py | 18 ++- 5 files changed, 67 insertions(+), 137 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index 1f3b90afd..bd61e795b 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -542,7 +542,6 @@ def __init__(self, domain, active_tracers, Vu=None): # Build finite element spaces self.spaces = [] - # Add active tracers to the list of prognostics if active_tracers is None: active_tracers = [] diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index d883ae216..0ce9dc68e 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -21,7 +21,7 @@ from gusto.configuration import EmbeddedDGOptions, RecoveryOptions from gusto.labels import (time_derivative, prognostic, physics_label, - implicit, explicit, transport) + implicit, explicit) from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * @@ -89,21 +89,21 @@ def __init__(self, domain, field_name=None, solver_parameters=None, if options is not None: if type(options) == MixedOptions: self.wrapper = options - + for field, suboption in self.wrapper.suboptions.items(): print(field) print(suboption) - + if suboption.name == 'embedded_dg': - self.wrapper.subwrappers.update({field:EmbeddedDGWrapper(self, suboption)}) + self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)}) elif suboption.name == "recovered": - self.wrapper.subwrappers.update({field:RecoveryWrapper(self, suboption)}) + self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)}) elif suboption.name == "supg": - self.wrapper.subwrappers.update({field:SUPGWrapper(self, suboption)}) + self.wrapper.subwrappers.update({field: SUPGWrapper(self, suboption)}) else: raise RuntimeError( - f'Time discretisation: suboption wrapper {wrapper_name} not implemented') - + f'Time discretisation: suboption wrapper {wrapper_name} not implemented') + else: self.wrapper_name = options.name if self.wrapper_name == "embedded_dg": @@ -178,23 +178,21 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.wrapper is not None: if type(self.wrapper) == MixedOptions: - - #print(self.wrapper.wrapper_spaces) - + for field, subwrapper in self.wrapper.subwrappers.items(): field_idx = equation.field_names.index(field) - + subwrapper.idx = field_idx subwrapper.mixed_options = True - + # Store the original space of the tracer subwrapper.tracer_fs = self.equation.spaces[field_idx] - + subwrapper.setup() - + # Update the function space to that needed by the wrapper self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space - + # Store test space? #self.wrapper.test_spaces[field_idx] = subwrapper.function_space @@ -219,13 +217,21 @@ def setup(self, equation, apply_bcs=True, *active_labels): # If use supg, then change the required test function # Replace one-by-one in this case? + print(type(new_test_mixed)) - #for field, subwrapper in self.wrapper.subwrappers.items(): - # if self.wrapper.suboptions[field].name == "supg": - # field_idx = equation.field_names.index(field) - # test_list = list(new_test_mixed) - # test_list[field_idx] = subwrapper.test - # new_test_mixed = tuple(test_list) + for field, subwrapper in self.wrapper.subwrappers.items(): + if self.wrapper.suboptions[field].name == "supg": + field_idx = equation.field_names.index(field) + test_list = list(new_test_mixed) + test_list[field_idx] = subwrapper.test + new_test_mixed2 = tuple(test_list) + + print(type(new_test_mixed2)) + + if new_test_mixed == new_test_mixed2: + print('same') + else: + print('different') # Change test functions with the new space: #for field, subwrapper in self.wrapper.subwrappers.items(): diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 956f409d2..900224733 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -180,8 +180,7 @@ def setup(self): equation = self.time_discretisation.equation if self.mixed_options == True: - #original_space = self.tracer_fs - original_space = self.time_discretisation.fs + original_space = self.tracer_fs else: original_space = self.time_discretisation.fs @@ -302,7 +301,9 @@ def setup(self): #else: # self.function_space = self.time_discretisation.fs + # But this will create some problems when defining the new mixed space... self.function_space = self.time_discretisation.fs + #self.function_space = self.tracer_fs if self.mixed_options == True: self.test_space = self.tracer_fs @@ -444,12 +445,8 @@ def pre_apply(self, x_in): """ for field_name in self.field_names: - #print(field_name) - #print(self.subwrappers) - + field_idx = self.field_names.index(field_name) - #print(field_idx) - field = x_in.subfunctions[field_idx] x_in_sub = self.x_in.subfunctions[field_idx] diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index 9853b11c3..94d83ae92 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -50,32 +50,17 @@ def setup_limiters(dirname, space): VDG1 = domain.spaces('DG1_equispaced') elif space == 'Vtheta_degree_1': V = domain.spaces('theta') - elif space == 'mixed_FS': - V = domain.spaces('HDiv') - VA = domain.spaces('DG') - VB = domain.spaces('DG1_equispaced') else: raise NotImplementedError Vpsi = domain.spaces('H1') # Equation - if space == 'mixed_FS': - tracerA = ActiveTracer(name='tracerA', space='DG', - variable_type=TracerVariableType.mixing_ratio, - transport_eqn=TransportEquationType.advective) - tracerB = ActiveTracer(name='tracerB', space='DG1_equispaced', - variable_type=TracerVariableType.mixing_ratio, - transport_eqn=TransportEquationType.advective) - tracers = [tracerA, tracerB] - eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V) - output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1, - dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB']) - else: - eqn = AdvectionEquation(domain, V, 'tracer') - output = OutputParameters(dirname=dirname+'/limiters', - dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) + eqn = AdvectionEquation(domain, V, 'tracer') + # I/O + output = OutputParameters(dirname=dirname+'/limiters', + dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) io = IO(domain, output) # ------------------------------------------------------------------------ # @@ -99,19 +84,10 @@ def setup_limiters(dirname, space): elif space == 'Vtheta_degree_1': opts = EmbeddedDGOptions() transport_schemes = SSPRK3(domain, options=opts, limiter=ThetaLimiter(V)) - - elif space == 'mixed_FS': - sublimiters = {'tracerA': DG1Limiter(domain.spaces('DG')), - 'tracerB': VertexBasedLimiter(domain.spaces('DG1_equispaced'))} - MixedLimiter = MixedFSLimiter(eqn, sublimiters) - transport_schemes = SSPRK3(domain, limiter=MixedLimiter) else: raise NotImplementedError - if space == 'mixed_FS': - transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] - else: - transport_method = DGUpwind(eqn, "tracer") + transport_method = DGUpwind(eqn, "tracer") # Build time stepper stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method) @@ -120,14 +96,8 @@ def setup_limiters(dirname, space): # Initial condition # ------------------------------------------------------------------------ # - if space == 'mixed_FS': - tracerA_0 = stepper.fields('tracerA', space=VA) - tracerB_0 = stepper.fields('tracerB', space=VB) - true_fieldA = stepper.fields('true_tracerA', space=VA) - true_fieldB = stepper.fields('true_tracerB', space=VB) - else: - tracer0 = stepper.fields('tracer', V) - true_field = stepper.fields('true_tracer', space=V) + tracer0 = stepper.fields('tracer', V) + true_field = stepper.fields('true_tracer', space=V) x, z = SpatialCoordinate(mesh) @@ -176,17 +146,9 @@ def setup_limiters(dirname, space): 0.0) if i == 0: - if space == 'mixed_FS': - tracerA_0.interpolate(Constant(tracer_min) + expr_1 + expr_2) - tracerB_0.interpolate(Constant(tracer_min) + expr_1 + expr_2) - else: - tracer0.interpolate(Constant(tracer_min) + expr_1 + expr_2) + tracer0.interpolate(Constant(tracer_min) + expr_1 + expr_2) elif i == 1: - if space == 'mixed_FS': - true_fieldA.interpolate(Constant(tracer_min) + expr_1 + expr_2) - true_fieldB.interpolate(Constant(tracer_min) + expr_1 + expr_2) - else: - true_field.interpolate(Constant(tracer_min) + expr_1 + expr_2) + true_field.interpolate(Constant(tracer_min) + expr_1 + expr_2) else: raise ValueError @@ -219,67 +181,29 @@ def setup_limiters(dirname, space): gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)]) u.project(gradperp(psi)) - if space == 'mixed_FS': - return stepper, tmax, true_fieldA, true_fieldB - else: - return stepper, tmax, true_field + return stepper, tmax, true_field -@pytest.mark.parametrize('space', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', - 'DG1', 'DG1_equispaced', 'mixed_FS']) +@pytest.mark.parametrize('space', ['Vtheta_degree_0', 'Vtheta_degree_1', + 'DG0', 'DG1', 'DG1_equispaced']) def test_limiters(tmpdir, space): # Setup and run dirname = str(tmpdir) - - if space == 'mixed_FS': - stepper, tmax, true_fieldA, true_fieldB = setup_limiters(dirname, space) - else: - stepper, tmax, true_field = setup_limiters(dirname, space) - + stepper, tmax, true_field = setup_limiters(dirname, space) stepper.run(t=0, tmax=tmax) + final_field = stepper.fields('tracer') - tol = 1e-9 - - if space == 'mixed_FS': - final_fieldA = stepper.fields('tracerA') - final_fieldB = stepper.fields('tracerB') + # Check tracer is roughly in the correct place + assert norm(true_field - final_field) / norm(true_field) < 0.05, \ + 'Something appears to have gone wrong with transport of tracer using a limiter' - # Check tracer is roughly in the correct place - assert norm(true_fieldA - final_fieldA) / norm(true_fieldA) < 0.05, \ - 'Something is wrong with the DG space tracer using a mixed limiter' - - # Check tracer is roughly in the correct place - assert norm(true_fieldB - final_fieldB) / norm(true_fieldB) < 0.05, \ - 'Something is wrong with the DG1 equispaced tracer using a mixed limiter' - - # Check for no new overshoots in A - assert np.max(final_fieldA.dat.data) <= np.max(true_fieldA.dat.data) + tol, \ - 'Application of the DG space limiter in the mixed limiter has not prevented overshoots' - - # Check for no new undershoots in A - assert np.min(final_fieldA.dat.data) >= np.min(true_fieldA.dat.data) - tol, \ - 'Application of the DG space limiter in the mixed limiter has not prevented undershoots' - - # Check for no new overshoots in B - assert np.max(final_fieldB.dat.data) <= np.max(true_fieldB.dat.data) + tol, \ - 'Application of the DG1 equispaced limiter in the mixed limiter has not prevented overshoots' - - # Check for no new undershoots in B - assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \ - 'Application of the DG1 equispaced limiter in the mixed limiter has not prevented undershoots' - - else: - final_field = stepper.fields('tracer') - - # Check tracer is roughly in the correct place - assert norm(true_field - final_field) / norm(true_field) < 0.05, \ - 'Something appears to have gone wrong with transport of tracer using a limiter' + tol = 1e-9 - # Check for no new overshoots - assert np.max(final_field.dat.data) <= np.max(true_field.dat.data) + tol, \ - 'Application of limiter has not prevented overshoots' + # Check for no new overshoots + assert np.max(final_field.dat.data) <= np.max(true_field.dat.data) + tol, \ + 'Application of limiter has not prevented overshoots' - # Check for no new undershoots - assert np.min(final_field.dat.data) >= np.min(true_field.dat.data) - tol, \ - 'Application of limiter has not prevented undershoots' + # Check for no new undershoots + assert np.min(final_field.dat.data) >= np.min(true_field.dat.data) - tol, \ + 'Application of limiter has not prevented undershoots' \ No newline at end of file diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 301c260e5..69c1c55ac 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -62,9 +62,9 @@ def setup_limiters(dirname, space_A, space_B): elif space_A == 'HDiv': VA = domain.spaces('HDiv') space_A_string = 'HDiv' - #elif space_A == 'CG1' - #VA = FunctionSpace(mesh, 'CG', 1) - #space_A_string = 'CG' + elif space_A == 'CG1': + VA = FunctionSpace(mesh, 'CG', 1) + space_A_string = 'CG' else: raise NotImplementedError @@ -140,6 +140,9 @@ def setup_limiters(dirname, space_A, space_B): elif space_A == 'HDiv': ibp_A = IntegrateByParts.TWICE suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)}) + elif space_A == 'CG1': + ibp_A = IntegrateByParts.NEVER + suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)}) else: raise NotImplementedError @@ -182,7 +185,8 @@ def setup_limiters(dirname, space_A, space_B): #transport_schemes = SSPRK3(domain, limiter=MixedLimiter) # DG Upwind transport for both tracers: - if space_A == 'HDiv': + if space_A == 'HDiv' or space_A == 'CG1': + #Pass SUPG options to the transport method transport_method_A = DGUpwind(eqn, 'tracerA', ibp=ibp_A) else: transport_method_A = DGUpwind(eqn, 'tracerA') @@ -292,9 +296,9 @@ def setup_limiters(dirname, space_A, space_B): return stepper, tmax, true_fieldA, true_fieldB -#@pytest.mark.parametrize('space_A', ['HDiv'])#, -@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', - 'DG1', 'DG1_equispaced', 'HDiv']) +@pytest.mark.parametrize('space_A', ['HDiv'])#, 'CG1'])#, +#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', +# 'DG1', 'DG1_equispaced', 'HDiv']) # It only makes sense to use the same degree for tracer B @pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) From 6d2fc9b5eb7ce57892e858bddfae3f943ff2e13e Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Mon, 22 Jan 2024 23:37:56 +0000 Subject: [PATCH 11/25] finalising mixed options test script --- gusto/time_discretisation.py | 64 ++----------------- gusto/wrappers.py | 54 ++++++---------- .../transport/test_mixed_fs_options.py | 41 +++--------- 3 files changed, 34 insertions(+), 125 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 0ce9dc68e..2b8f7bb57 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -99,7 +99,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, elif suboption.name == "recovered": self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)}) elif suboption.name == "supg": - self.wrapper.subwrappers.update({field: SUPGWrapper(self, suboption)}) + raise RuntimeError( + f'Time discretisation: suboption SUPG is currently not implemented within MixedOptions') else: raise RuntimeError( f'Time discretisation: suboption wrapper {wrapper_name} not implemented') @@ -180,11 +181,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): if type(self.wrapper) == MixedOptions: for field, subwrapper in self.wrapper.subwrappers.items(): - field_idx = equation.field_names.index(field) - subwrapper.idx = field_idx subwrapper.mixed_options = True + field_idx = equation.field_names.index(field) + # Store the original space of the tracer subwrapper.tracer_fs = self.equation.spaces[field_idx] @@ -193,67 +194,16 @@ def setup(self, equation, apply_bcs=True, *active_labels): # Update the function space to that needed by the wrapper self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space - # Store test space? - #self.wrapper.test_spaces[field_idx] = subwrapper.function_space - - # Replace the test function space - #if self.wrapper.suboptions[field].name == "supg": - # new_test = subwrapper.test - #else: - # new_test = TestFunction(subwrapper.test_space) - - #self.residual = self.residual.label_map( - # lambda t: t.has_label(transport) and t.get(prognostic) == field, - # map_if_true=replace_test_function(new_test, old_idx=field_idx)) - - #self.residual = subwrapper.label_terms(self.residual) - self.wrapper.setup() - self.fs = self.wrapper.function_space - new_test_mixed = TestFunctions(self.fs) - # If use supg, then change the required test function - - # Replace one-by-one in this case? - print(type(new_test_mixed)) - - for field, subwrapper in self.wrapper.subwrappers.items(): - if self.wrapper.suboptions[field].name == "supg": - field_idx = equation.field_names.index(field) - test_list = list(new_test_mixed) - test_list[field_idx] = subwrapper.test - new_test_mixed2 = tuple(test_list) - - print(type(new_test_mixed2)) - - if new_test_mixed == new_test_mixed2: - print('same') - else: - print('different') - - # Change test functions with the new space: - #for field, subwrapper in self.wrapper.subwrappers.items(): - # self.residual = self.residual.label_map( - # lambda t: t.has_label(transport) and t.get(prognostic) == field, - # map_if_true=replace_test_function(new_test_mixed[field_idx], old_idx=field_idx)) - - # self.residual = subwrapper.label_terms(self.residual) - - # Or, all-at-once: - #self.residual = self.residual.label_map( - # lambda t: t.has_label(transport), - # map_if_true=replace_test_function(new_test_mixed)) - + # Replace the original test function with one from the new + # function space defined by the subwrappers self.residual = self.residual.label_map( all_terms, map_if_true=replace_test_function(new_test_mixed)) - - - - - + else: self.wrapper.setup() self.fs = self.wrapper.function_space diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 900224733..880630cd4 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -295,25 +295,9 @@ def setup(self): 'SUPG wrapper can only be used with SUPG Options' domain = self.time_discretisation.domain - - #if self.mixed_options == True: - # self.function_space = self.tracer_fs - #else: - # self.function_space = self.time_discretisation.fs - - # But this will create some problems when defining the new mixed space... self.function_space = self.time_discretisation.fs - #self.function_space = self.tracer_fs - - if self.mixed_options == True: - self.test_space = self.tracer_fs - self.x_out = Function(self.tracer_fs) - else: - self.test_space = self.function_space - self.x_out = Function(self.function_space) - - #self.test_space = self.function_space - #self.x_out = Function(self.function_space) + self.test_space = self.function_space + self.x_out = Function(self.function_space) # -------------------------------------------------------------------- # # Work out SUPG parameter @@ -332,8 +316,11 @@ def setup(self): # so that we don't apply supg in that direction if is_cg(self.function_space): vals = default_vals + print('is cg') else: + print('is not cg') space = self.function_space.ufl_element().sobolev_space + print(space.name) if space.name in ["HDiv", "DirectionalH"]: vals = [default_vals[i] if space[i].name == "H1" else 0. for i in range(dim)] @@ -415,12 +402,11 @@ def __init__(self, equation, suboptions): """ Args: equation (:class: `PrognosticEquationSet`): the prognostic equation(s) - sublimiters (dict): A dictionary holding limiters defined for individual prognostic variables + suboptions (dict): A dictionary holding options defined for individual prognostic variables Raises: ValueError: If an option is defined for a field that is not in the prognostic variable set """ self.wrapper_spaces = equation.spaces - #self.test_spaces = equation.spaces self.field_names = equation.field_names self.suboptions = suboptions self.subwrappers = {} @@ -431,49 +417,45 @@ def __init__(self, equation, suboptions): raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") def setup(self): - - # Compute the new mixed function space + """ Compute the new mixed function space from the subwrappers """ + self.function_space = MixedFunctionSpace(self.wrapper_spaces) self.x_in = Function(self.function_space) - self.x_out = Function(self.function_space) - #self.test_space = MixedFunctionSpace(self.test_spaces) - + self.x_out = Function(self.function_space) + def pre_apply(self, x_in): """ Perform the pre-applications for all fields with an associated subwrapper. """ - + for field_name in self.field_names: - field_idx = self.field_names.index(field_name) field = x_in.subfunctions[field_idx] x_in_sub = self.x_in.subfunctions[field_idx] - + if field_name in self.subwrappers: subwrapper = self.subwrappers[field_name] subwrapper.pre_apply(field) x_in_sub.assign(subwrapper.x_in) else: - x_in_sub.assign(field) - - + x_in_sub.assign(field) + + def post_apply(self, x_out): """ Perform the post-applications for all fields with an associated subwrapper. """ - + for field_name in self.field_names: - field_idx = self.field_names.index(field_name) - field = self.x_out.subfunctions[field_idx] x_out_sub = x_out.subfunctions[field_idx] - + if field_name in self.subwrappers: subwrapper = self.subwrappers[field_name] subwrapper.x_out.assign(field) subwrapper.post_apply(x_out_sub) else: - x_out_sub.assign(field) \ No newline at end of file + x_out_sub.assign(field) diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 69c1c55ac..d00195db2 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -32,7 +32,7 @@ def setup_limiters(dirname, space_A, space_B): # Domain m = PeriodicIntervalMesh(20, Ld) mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20)) - degree = 0 if space_A in ['DG0', 'Vtheta_degree_0'] else 1 + degree = 0 if space_A in ['DG0', 'Vtheta_degree_0', 'Hdiv'] else 1 domain = Domain(mesh, dt, family="CG", degree=degree) @@ -59,12 +59,6 @@ def setup_limiters(dirname, space_A, space_B): elif space_A == 'Vtheta_degree_1': VA = domain.spaces('theta') space_A_string = 'theta' - elif space_A == 'HDiv': - VA = domain.spaces('HDiv') - space_A_string = 'HDiv' - elif space_A == 'CG1': - VA = FunctionSpace(mesh, 'CG', 1) - space_A_string = 'CG' else: raise NotImplementedError @@ -137,12 +131,6 @@ def setup_limiters(dirname, space_A, space_B): elif space_A == 'Vtheta_degree_1': suboptions.update({'tracerA': EmbeddedDGOptions()}) sublimiters.update({'tracerA': ThetaLimiter(VA)}) - elif space_A == 'HDiv': - ibp_A = IntegrateByParts.TWICE - suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)}) - elif space_A == 'CG1': - ibp_A = IntegrateByParts.NEVER - suboptions.update({'tracerA': SUPGOptions(ibp=ibp_A)}) else: raise NotImplementedError @@ -176,27 +164,16 @@ def setup_limiters(dirname, space_A, space_B): else: raise NotImplementedError + # Create the mixed options and mixed limiter objects opts = MixedOptions(eqn, suboptions) MixedLimiter = MixedFSLimiter(eqn, sublimiters) # Give the scheme for the coupled transport - #transport_schemes = SSPRK3(domain, options=opts) transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter) - #transport_schemes = SSPRK3(domain, limiter=MixedLimiter) - - # DG Upwind transport for both tracers: - if space_A == 'HDiv' or space_A == 'CG1': - #Pass SUPG options to the transport method - transport_method_A = DGUpwind(eqn, 'tracerA', ibp=ibp_A) - else: - transport_method_A = DGUpwind(eqn, 'tracerA') - transport_method_B = DGUpwind(eqn, 'tracerB') - - transport_method = [transport_method_A, transport_method_B] - - # Need to give SUPG options to the above, if using supg ... - # Need to test SUPG here! - + + # DG Upwind transport for both tracers: + transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] + # Build time stepper stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method) @@ -296,9 +273,9 @@ def setup_limiters(dirname, space_A, space_B): return stepper, tmax, true_fieldA, true_fieldB -@pytest.mark.parametrize('space_A', ['HDiv'])#, 'CG1'])#, -#@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', -# 'DG1', 'DG1_equispaced', 'HDiv']) + +@pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', + 'DG1', 'DG1_equispaced']) # It only makes sense to use the same degree for tracer B @pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) From 25ed24320a64bf100814a9531189d010cf398616 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Mon, 22 Jan 2024 23:44:30 +0000 Subject: [PATCH 12/25] lint --- gusto/time_discretisation.py | 13 +++++-------- gusto/wrappers.py | 10 +++++----- .../transport/test_mixed_fs_options.py | 2 -- 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 2b8f7bb57..10ddc8526 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -91,16 +91,13 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.wrapper = options for field, suboption in self.wrapper.suboptions.items(): - print(field) - print(suboption) - if suboption.name == 'embedded_dg': self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)}) elif suboption.name == "recovered": self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)}) elif suboption.name == "supg": raise RuntimeError( - f'Time discretisation: suboption SUPG is currently not implemented within MixedOptions') + 'Time discretisation: suboption SUPG is currently not implemented within MixedOptions') else: raise RuntimeError( f'Time discretisation: suboption wrapper {wrapper_name} not implemented') @@ -201,8 +198,8 @@ def setup(self, equation, apply_bcs=True, *active_labels): # Replace the original test function with one from the new # function space defined by the subwrappers self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test_mixed)) + all_terms, + map_if_true=replace_test_function(new_test_mixed)) else: self.wrapper.setup() @@ -213,12 +210,12 @@ def setup(self, equation, apply_bcs=True, *active_labels): # SUPG has a special wrapper if self.wrapper_name == "supg": new_test = self.wrapper.test - + # Replace the original test function with the one from the wrapper self.residual = self.residual.label_map( all_terms, map_if_true=replace_test_function(new_test)) - + self.residual = self.wrapper.label_terms(self.residual) # -------------------------------------------------------------------- # diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 880630cd4..dc2923e9d 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -84,10 +84,10 @@ def setup(self): assert isinstance(self.options, EmbeddedDGOptions), \ 'Embedded DG wrapper can only be used with Embedded DG Options' - + domain = self.time_discretisation.domain equation = self.time_discretisation.equation - + if self.mixed_options == True: original_space = self.tracer_fs else: @@ -111,8 +111,8 @@ def setup(self): self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - - if self.mixed_options == True: + + if self.mixed_options: self.x_projected = Function(self.tracer_fs) elif self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) @@ -178,7 +178,7 @@ def setup(self): domain = self.time_discretisation.domain equation = self.time_discretisation.equation - + if self.mixed_options == True: original_space = self.tracer_fs else: diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index d00195db2..ac27f4bd7 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -317,5 +317,3 @@ def test_mixed_fs_options(tmpdir, space_A, space_B): # Check for no new undershoots in B assert np.min(final_fieldB.dat.data) >= np.min(true_fieldB.dat.data) - tol, \ f"Application of the {space_B} equispaced limiter in the mixed limiter has not prevented undershoots" - - \ No newline at end of file From f410d2cc01efa49be4f08ccd90d3d14352d351d4 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Mon, 22 Jan 2024 23:46:44 +0000 Subject: [PATCH 13/25] lint --- gusto/wrappers.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index dc2923e9d..4d42d1dac 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -88,7 +88,7 @@ def setup(self): domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if self.mixed_options == True: + if self.mixed_options: original_space = self.tracer_fs else: original_space = self.time_discretisation.fs @@ -179,7 +179,7 @@ def setup(self): domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if self.mixed_options == True: + if self.mixed_options: original_space = self.tracer_fs else: original_space = self.time_discretisation.fs @@ -201,15 +201,15 @@ def setup(self): # Internal variables to be used # -------------------------------------------------------------------- # - if self.mixed_options == True: + if self.mixed_options: self.x_in_tmp = Function(self.tracer_fs) else: self.x_in_tmp = Function(self.time_discretisation.fs) - + self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - - if self.mixed_options == True: + + if self.mixed_options: self.x_projected = Function(self.tracer_fs) elif self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) @@ -388,16 +388,16 @@ def label_terms(self, residual): new_residual = transporting_velocity.update_value(new_residual, self.transporting_velocity) return new_residual - + class MixedOptions(object): """ An object to hold a dictionary with different options for different tracers. This means that different tracers can be solved simultaneously using a CoupledTransportEquation, whilst being in different spaces - and needing different implementation options. + and needing different implementation options. """ - + def __init__(self, equation, suboptions): """ Args: @@ -410,12 +410,12 @@ def __init__(self, equation, suboptions): self.field_names = equation.field_names self.suboptions = suboptions self.subwrappers = {} - + for field, suboption in suboptions.items(): # Check that the field is in the prognostic variable set: if field not in equation.field_names: raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") - + def setup(self): """ Compute the new mixed function space from the subwrappers """ From d166a224bdfa40d8ffd988823fa9b80565601c2c Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Mon, 22 Jan 2024 23:48:35 +0000 Subject: [PATCH 14/25] lint --- gusto/wrappers.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 4d42d1dac..7cbcc5af7 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -412,9 +412,9 @@ def __init__(self, equation, suboptions): self.subwrappers = {} for field, suboption in suboptions.items(): - # Check that the field is in the prognostic variable set: - if field not in equation.field_names: - raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") + # Check that the field is in the prognostic variable set: + if field not in equation.field_names: + raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") def setup(self): """ Compute the new mixed function space from the subwrappers """ @@ -441,7 +441,6 @@ def pre_apply(self, x_in): else: x_in_sub.assign(field) - def post_apply(self, x_out): """ Perform the post-applications for all fields From aa3d7c07b25ec7022ac5c6d78b8eea52f25e4822 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Mon, 22 Jan 2024 23:53:44 +0000 Subject: [PATCH 15/25] lint --- integration-tests/transport/test_limiters.py | 2 +- .../transport/test_mixed_fs_options.py | 53 ++++++++----------- 2 files changed, 22 insertions(+), 33 deletions(-) diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index 94d83ae92..8255ad2cd 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -206,4 +206,4 @@ def test_limiters(tmpdir, space): # Check for no new undershoots assert np.min(final_field.dat.data) >= np.min(true_field.dat.data) - tol, \ - 'Application of limiter has not prevented undershoots' \ No newline at end of file + 'Application of limiter has not prevented undershoots' diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index ac27f4bd7..96b28c43b 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -35,10 +35,10 @@ def setup_limiters(dirname, space_A, space_B): degree = 0 if space_A in ['DG0', 'Vtheta_degree_0', 'Hdiv'] else 1 domain = Domain(mesh, dt, family="CG", degree=degree) - + # Transporting velocity space V = domain.spaces('HDiv') - + # Tracer A spaces if space_A == 'DG0': VA = domain.spaces('DG') @@ -61,7 +61,7 @@ def setup_limiters(dirname, space_A, space_B): space_A_string = 'theta' else: raise NotImplementedError - + # Tracer B spaces if space_B == 'DG': space_B_string = 'DG' @@ -75,7 +75,7 @@ def setup_limiters(dirname, space_A, space_B): raise NotImplementedError elif space_B == 'Vtheta': space_B_string = 'theta' - if degree == 0: + if degree == 0: VB = domain.spaces('theta') VCG1_B = FunctionSpace(mesh, 'CG', 1) VDG1_B = domain.spaces('DG1_equispaced') @@ -85,23 +85,22 @@ def setup_limiters(dirname, space_A, space_B): raise NotImplementedError else: raise NotImplementedError - - Vpsi = domain.spaces('H1') - + + Vpsi = domain.spaces('H1') + tracerA = ActiveTracer(name='tracerA', space=space_A_string, - variable_type=TracerVariableType.mixing_ratio, - transport_eqn=TransportEquationType.advective) - + variable_type=TracerVariableType.mixing_ratio, + transport_eqn=TransportEquationType.advective) + tracerB = ActiveTracer(name='tracerB', space=space_B_string, - variable_type=TracerVariableType.mixing_ratio, - transport_eqn=TransportEquationType.advective) - - + variable_type=TracerVariableType.mixing_ratio, + transport_eqn=TransportEquationType.advective) + tracers = [tracerA, tracerB] eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V) output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1, - dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB']) + dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB']) io = IO(domain, output) @@ -113,21 +112,17 @@ def setup_limiters(dirname, space_A, space_B): sublimiters = {} # Options and limiters for tracer_A - if space_A in ['DG0', 'Vtheta_degree_0']: suboptions.update({'tracerA': RecoveryOptions(embedding_space=VDG1_A, - recovered_space=VCG1_A, - project_low_method='recover', - boundary_method=BoundaryMethod.taylor)}) - + recovered_space=VCG1_A, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor)}) + sublimiters.update({'tracerA': VertexBasedLimiter(VDG1_A)}) - elif space_A == 'DG1': sublimiters.update({'tracerA': DG1Limiter(VA)}) - elif space_A == 'DG1_equispaced': sublimiters.update({'tracerA': VertexBasedLimiter(VA)}) - elif space_A == 'Vtheta_degree_1': suboptions.update({'tracerA': EmbeddedDGOptions()}) sublimiters.update({'tracerA': ThetaLimiter(VA)}) @@ -135,14 +130,12 @@ def setup_limiters(dirname, space_A, space_B): raise NotImplementedError # Options and limiters for tracer_B - if space_B == 'DG': if degree == 0: suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B, - recovered_space=VCG1_B, - project_low_method='recover', - boundary_method=BoundaryMethod.taylor)}) - + recovered_space=VCG1_B, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor)}) sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)}) elif degree == 1: sublimiters.update({'tracerB': DG1Limiter(VB)}) @@ -273,13 +266,9 @@ def setup_limiters(dirname, space_A, space_B): return stepper, tmax, true_fieldA, true_fieldB - @pytest.mark.parametrize('space_A', ['Vtheta_degree_0', 'Vtheta_degree_1', 'DG0', 'DG1', 'DG1_equispaced']) -# It only makes sense to use the same degree for tracer B @pytest.mark.parametrize('space_B', ['Vtheta', 'DG']) - - def test_mixed_fs_options(tmpdir, space_A, space_B): # Setup and run From f1e958a8101e5b8fa2805d55d6df0dc98b7f84bf Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Mon, 22 Jan 2024 23:55:18 +0000 Subject: [PATCH 16/25] lint --- .../transport/test_mixed_fs_options.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 96b28c43b..1c4f545d2 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -101,7 +101,7 @@ def setup_limiters(dirname, space_A, space_B): eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=V) output = OutputParameters(dirname=dirname+'/limiters', dumpfreq=1, dumplist=['u', 'tracerA', 'tracerB', 'true_tracerA', 'true_tracerB']) - + io = IO(domain, output) # ------------------------------------------------------------------------ # @@ -128,7 +128,7 @@ def setup_limiters(dirname, space_A, space_B): sublimiters.update({'tracerA': ThetaLimiter(VA)}) else: raise NotImplementedError - + # Options and limiters for tracer_B if space_B == 'DG': if degree == 0: @@ -142,12 +142,11 @@ def setup_limiters(dirname, space_A, space_B): else: raise NotImplementedError elif space_B == 'Vtheta': - if degree == 0: + if degree == 0: suboptions.update({'tracerB': RecoveryOptions(embedding_space=VDG1_B, - recovered_space=VCG1_B, - project_low_method='recover', - boundary_method=BoundaryMethod.taylor)}) - + recovered_space=VCG1_B, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor)}) sublimiters.update({'tracerB': VertexBasedLimiter(VDG1_B)}) elif degree == 1: suboptions.update({'tracerB': EmbeddedDGOptions()}) @@ -164,7 +163,7 @@ def setup_limiters(dirname, space_A, space_B): # Give the scheme for the coupled transport transport_schemes = SSPRK3(domain, options=opts, limiter=MixedLimiter) - # DG Upwind transport for both tracers: + # DG Upwind transport for both tracers: transport_method = [DGUpwind(eqn, 'tracerA'), DGUpwind(eqn, 'tracerB')] # Build time stepper @@ -178,7 +177,7 @@ def setup_limiters(dirname, space_A, space_B): tracerB_0 = stepper.fields('tracerB', space=VB) true_fieldA = stepper.fields('true_tracerA', space=VA) true_fieldB = stepper.fields('true_tracerB', space=VB) - + x, z = SpatialCoordinate(mesh) tracer_min = 12.6 From 51f6f7b511c7851c2e7f187d3aba78c81545e511 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Tue, 23 Jan 2024 01:29:29 +0000 Subject: [PATCH 17/25] remove dubugging statements --- gusto/wrappers.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 7cbcc5af7..110b774df 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -192,7 +192,6 @@ def setup(self): V_elt = BrokenElement(original_space.ufl_element()) self.function_space = FunctionSpace(domain.mesh, V_elt) else: - print('using embedded space') self.function_space = self.options.embedding_space self.test_space = self.function_space @@ -316,11 +315,8 @@ def setup(self): # so that we don't apply supg in that direction if is_cg(self.function_space): vals = default_vals - print('is cg') else: - print('is not cg') space = self.function_space.ufl_element().sobolev_space - print(space.name) if space.name in ["HDiv", "DirectionalH"]: vals = [default_vals[i] if space[i].name == "H1" else 0. for i in range(dim)] From 55e0856b7027924cec6a1a58f1a946a91cac1c86 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Tue, 20 Feb 2024 02:40:47 +0000 Subject: [PATCH 18/25] separate MixedFSOptions and MixedFSWrapper to align with exising code --- gusto/configuration.py | 13 +++++- gusto/time_discretisation.py | 36 ++++++++------- gusto/wrappers.py | 44 +++++++------------ .../transport/test_mixed_fs_options.py | 3 +- 4 files changed, 48 insertions(+), 48 deletions(-) diff --git a/gusto/configuration.py b/gusto/configuration.py index 3694ae9ea..21e16513f 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -7,7 +7,7 @@ __all__ = [ "IntegrateByParts", "TransportEquationType", "OutputParameters", "CompressibleParameters", "ShallowWaterParameters", - "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", + "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", "MixedFSOptions", "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters" ] @@ -172,6 +172,17 @@ class SUPGOptions(WrapperOptions): ibp = IntegrateByParts.TWICE +class MixedFSOptions(WrapperOptions): + """Specifies options for a mixed finite element formulation + where different suboptions are applied to different + prognostic variables.""" + + name = "mixed_options" + #tracer_fs = None + suboptions = {} + #subwrappers = {} + + class SpongeLayerParameters(Configuration): """Specifies parameters describing a 'sponge' (damping) layer.""" diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 10ddc8526..2d010ffcb 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -87,10 +87,11 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.courant_max = None if options is not None: - if type(options) == MixedOptions: - self.wrapper = options + self.wrapper_name = options.name + if self.wrapper_name == "mixed_options": + self.wrapper = MixedFSWrapper() - for field, suboption in self.wrapper.suboptions.items(): + for field, suboption in options.suboptions.items(): if suboption.name == 'embedded_dg': self.wrapper.subwrappers.update({field: EmbeddedDGWrapper(self, suboption)}) elif suboption.name == "recovered": @@ -101,18 +102,15 @@ def __init__(self, domain, field_name=None, solver_parameters=None, else: raise RuntimeError( f'Time discretisation: suboption wrapper {wrapper_name} not implemented') - + elif self.wrapper_name == "embedded_dg": + self.wrapper = EmbeddedDGWrapper(self, options) + elif self.wrapper_name == "recovered": + self.wrapper = RecoveryWrapper(self, options) + elif self.wrapper_name == "supg": + self.wrapper = SUPGWrapper(self, options) else: - self.wrapper_name = options.name - if self.wrapper_name == "embedded_dg": - self.wrapper = EmbeddedDGWrapper(self, options) - elif self.wrapper_name == "recovered": - self.wrapper = RecoveryWrapper(self, options) - elif self.wrapper_name == "supg": - self.wrapper = SUPGWrapper(self, options) - else: - raise RuntimeError( - f'Time discretisation: wrapper {self.wrapper_name} not implemented') + raise RuntimeError( + f'Time discretisation: wrapper {self.wrapper_name} not implemented') else: self.wrapper = None self.wrapper_name = None @@ -175,11 +173,15 @@ def setup(self, equation, apply_bcs=True, *active_labels): # -------------------------------------------------------------------- # if self.wrapper is not None: - if type(self.wrapper) == MixedOptions: + if self.wrapper_name == "mixed_options": + + self.wrapper.wrapper_spaces = equation.spaces + self.wrapper.field_names = equation.field_names for field, subwrapper in self.wrapper.subwrappers.items(): - - subwrapper.mixed_options = True + + if field not in equation.field_names: + raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set") field_idx = equation.field_names.index(field) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 110b774df..0ac2728c3 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -15,7 +15,7 @@ from gusto.labels import transporting_velocity import ufl -__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedOptions"] +__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper", "MixedFSWrapper"] class Wrapper(object, metaclass=ABCMeta): @@ -33,8 +33,6 @@ def __init__(self, time_discretisation, wrapper_options): self.time_discretisation = time_discretisation self.options = wrapper_options self.solver_parameters = None - self.idx = None - self.mixed_options = False self.tracer_fs = None @abstractmethod @@ -88,7 +86,7 @@ def setup(self): domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if self.mixed_options: + if self.tracer_fs is not None: original_space = self.tracer_fs else: original_space = self.time_discretisation.fs @@ -112,7 +110,7 @@ def setup(self): self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - if self.mixed_options: + if self.tracer_fs is not None: self.x_projected = Function(self.tracer_fs) elif self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) @@ -179,7 +177,7 @@ def setup(self): domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if self.mixed_options: + if self.tracer_fs is not None: original_space = self.tracer_fs else: original_space = self.time_discretisation.fs @@ -200,7 +198,7 @@ def setup(self): # Internal variables to be used # -------------------------------------------------------------------- # - if self.mixed_options: + if self.tracer_fs is not None: self.x_in_tmp = Function(self.tracer_fs) else: self.x_in_tmp = Function(self.time_discretisation.fs) @@ -208,7 +206,7 @@ def setup(self): self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - if self.mixed_options: + if self.tracer_fs is not None: self.x_projected = Function(self.tracer_fs) elif self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) @@ -386,31 +384,19 @@ def label_terms(self, residual): return new_residual -class MixedOptions(object): +class MixedFSWrapper(object): """ - An object to hold a dictionary with different options for different - tracers. This means that different tracers can be solved simultaneously - using a CoupledTransportEquation, whilst being in different spaces - and needing different implementation options. + An object to hold a subwrapper dictionary with different wrappers for + different tracers. This means that different tracers can be solved + simultaneously using a CoupledTransportEquation, whilst being in + different spaces and needing different implementation options. """ - def __init__(self, equation, suboptions): - """ - Args: - equation (:class: `PrognosticEquationSet`): the prognostic equation(s) - suboptions (dict): A dictionary holding options defined for individual prognostic variables - Raises: - ValueError: If an option is defined for a field that is not in the prognostic variable set - """ - self.wrapper_spaces = equation.spaces - self.field_names = equation.field_names - self.suboptions = suboptions - self.subwrappers = {} + def __init__(self): - for field, suboption in suboptions.items(): - # Check that the field is in the prognostic variable set: - if field not in equation.field_names: - raise ValueError(f"The limiter defined for {field} is for a field that does not exist in the equation set") + self.wrapper_spaces = None + self.field_names = None + self.subwrappers = {} def setup(self): """ Compute the new mixed function space from the subwrappers """ diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 1c4f545d2..097751212 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -157,7 +157,8 @@ def setup_limiters(dirname, space_A, space_B): raise NotImplementedError # Create the mixed options and mixed limiter objects - opts = MixedOptions(eqn, suboptions) + #opts = MixedFSOptions(eqn, suboptions) + opts = MixedFSOptions(suboptions=suboptions) MixedLimiter = MixedFSLimiter(eqn, sublimiters) # Give the scheme for the coupled transport From 164041cbd71d1c4811e27b7aff61d417a8c62095 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Tue, 20 Feb 2024 02:43:54 +0000 Subject: [PATCH 19/25] lint --- gusto/configuration.py | 2 -- integration-tests/transport/test_mixed_fs_options.py | 1 - 2 files changed, 3 deletions(-) diff --git a/gusto/configuration.py b/gusto/configuration.py index 21e16513f..5c909aa59 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -178,9 +178,7 @@ class MixedFSOptions(WrapperOptions): prognostic variables.""" name = "mixed_options" - #tracer_fs = None suboptions = {} - #subwrappers = {} class SpongeLayerParameters(Configuration): diff --git a/integration-tests/transport/test_mixed_fs_options.py b/integration-tests/transport/test_mixed_fs_options.py index 097751212..fc33dae6d 100644 --- a/integration-tests/transport/test_mixed_fs_options.py +++ b/integration-tests/transport/test_mixed_fs_options.py @@ -157,7 +157,6 @@ def setup_limiters(dirname, space_A, space_B): raise NotImplementedError # Create the mixed options and mixed limiter objects - #opts = MixedFSOptions(eqn, suboptions) opts = MixedFSOptions(suboptions=suboptions) MixedLimiter = MixedFSLimiter(eqn, sublimiters) From edcc929e9518f996e0f68ce0c09f7440b34723ef Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Tue, 20 Feb 2024 02:45:43 +0000 Subject: [PATCH 20/25] lint --- gusto/time_discretisation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 2d010ffcb..49cc46e1d 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -174,7 +174,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.wrapper is not None: if self.wrapper_name == "mixed_options": - + self.wrapper.wrapper_spaces = equation.spaces self.wrapper.field_names = equation.field_names From b317f39024bda5137f3aaa459df1f1728fa3ae1c Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Tue, 20 Feb 2024 02:47:24 +0000 Subject: [PATCH 21/25] lint --- gusto/time_discretisation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 49cc46e1d..5d9297c75 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -179,7 +179,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.wrapper.field_names = equation.field_names for field, subwrapper in self.wrapper.subwrappers.items(): - + if field not in equation.field_names: raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set") From 07003818f7c00bfb7f82652541243a4024b7b3dc Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Wed, 28 Feb 2024 03:25:54 +0000 Subject: [PATCH 22/25] set up original spaces for MixedFSWrapper within the subwrappers --- gusto/time_discretisation.py | 6 +----- gusto/wrappers.py | 27 ++++++++++++--------------- 2 files changed, 13 insertions(+), 20 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 5d9297c75..2f6cc009e 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -184,11 +184,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set") field_idx = equation.field_names.index(field) - - # Store the original space of the tracer - subwrapper.tracer_fs = self.equation.spaces[field_idx] - - subwrapper.setup() + subwrapper.setup(self.equation.spaces[field_idx]) # Update the function space to that needed by the wrapper self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 0ac2728c3..670fcacd1 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -33,7 +33,6 @@ def __init__(self, time_discretisation, wrapper_options): self.time_discretisation = time_discretisation self.options = wrapper_options self.solver_parameters = None - self.tracer_fs = None @abstractmethod def setup(self): @@ -77,7 +76,7 @@ class EmbeddedDGWrapper(Wrapper): the original space. """ - def setup(self): + def setup(self, previous_space=None): """Sets up function spaces and fields needed for this wrapper.""" assert isinstance(self.options, EmbeddedDGOptions), \ @@ -86,8 +85,8 @@ def setup(self): domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if self.tracer_fs is not None: - original_space = self.tracer_fs + if previous_space is not None: + original_space = previous_space else: original_space = self.time_discretisation.fs @@ -110,8 +109,8 @@ def setup(self): self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - if self.tracer_fs is not None: - self.x_projected = Function(self.tracer_fs) + if previous_space is not None: + self.x_projected = Function(previous_space) elif self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) else: @@ -166,19 +165,17 @@ class RecoveryWrapper(Wrapper): field is then returned to the original space. """ - def setup(self): + def setup(self, previous_space=None): """Sets up function spaces and fields needed for this wrapper.""" - print(self.options) - assert isinstance(self.options, RecoveryOptions), \ 'Recovery wrapper can only be used with Recovery Options' domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if self.tracer_fs is not None: - original_space = self.tracer_fs + if previous_space is not None: + original_space = previous_space else: original_space = self.time_discretisation.fs @@ -198,16 +195,16 @@ def setup(self): # Internal variables to be used # -------------------------------------------------------------------- # - if self.tracer_fs is not None: - self.x_in_tmp = Function(self.tracer_fs) + if previous_space is not None: + self.x_in_tmp = Function(previous_space) else: self.x_in_tmp = Function(self.time_discretisation.fs) self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - if self.tracer_fs is not None: - self.x_projected = Function(self.tracer_fs) + if previous_space is not None: + self.x_projected = Function(previous_space) elif self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) else: From 2f517e1df8f06b93c6b7f6811728e85d3cc94dac Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Wed, 28 Feb 2024 19:30:26 +0000 Subject: [PATCH 23/25] make original_space a base wrapper property --- gusto/wrappers.py | 61 +++++++++++++++++++++++------------------------ 1 file changed, 30 insertions(+), 31 deletions(-) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 670fcacd1..c4867e37b 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -33,14 +33,27 @@ def __init__(self, time_discretisation, wrapper_options): self.time_discretisation = time_discretisation self.options = wrapper_options self.solver_parameters = None + self.original_space = None @abstractmethod - def setup(self): + def setup(self, original_space=None): """ - Performs standard set up routines, and is to be called by the setup - method of the underlying time discretisation. + Store the original function space of the prognostic variable. + + Within each child wrapper, setup performs standard set up routines, + and is to be called by the setup method of the underlying + time discretisation. + + Args: + original_space (:class:`FunctionSpace`): the space that the + prognostic variable is defined on. This is a subset space of + a mixed function space when using a MixedFSWrapper. Defaults + to None, when not using a MixedFSWrapper. """ - pass + if original_space is not None: + self.original_space = original_space + else: + self.original_space = self.time_discretisation.fs @abstractmethod def pre_apply(self): @@ -76,26 +89,23 @@ class EmbeddedDGWrapper(Wrapper): the original space. """ - def setup(self, previous_space=None): + def setup(self, original_space=None): """Sets up function spaces and fields needed for this wrapper.""" assert isinstance(self.options, EmbeddedDGOptions), \ 'Embedded DG wrapper can only be used with Embedded DG Options' + super().setup(original_space) + domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if previous_space is not None: - original_space = previous_space - else: - original_space = self.time_discretisation.fs - # -------------------------------------------------------------------- # # Set up spaces to be used with wrapper # -------------------------------------------------------------------- # if self.options.embedding_space is None: - V_elt = BrokenElement(original_space.ufl_element()) + V_elt = BrokenElement(self.original_space.ufl_element()) self.function_space = FunctionSpace(domain.mesh, V_elt) else: self.function_space = self.options.embedding_space @@ -109,10 +119,8 @@ def setup(self, previous_space=None): self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - if previous_space is not None: - self.x_projected = Function(previous_space) - elif self.time_discretisation.idx is None: - self.x_projected = Function(equation.function_space) + if self.time_discretisation.idx is None: + self.x_projected = Function(self.original_space) else: self.x_projected = Function(equation.spaces[self.time_discretisation.idx]) @@ -165,26 +173,23 @@ class RecoveryWrapper(Wrapper): field is then returned to the original space. """ - def setup(self, previous_space=None): + def setup(self, original_space=None): """Sets up function spaces and fields needed for this wrapper.""" assert isinstance(self.options, RecoveryOptions), \ 'Recovery wrapper can only be used with Recovery Options' + super().setup(original_space) + domain = self.time_discretisation.domain equation = self.time_discretisation.equation - if previous_space is not None: - original_space = previous_space - else: - original_space = self.time_discretisation.fs - # -------------------------------------------------------------------- # # Set up spaces to be used with wrapper # -------------------------------------------------------------------- # if self.options.embedding_space is None: - V_elt = BrokenElement(original_space.ufl_element()) + V_elt = BrokenElement(self.original_space.ufl_element()) self.function_space = FunctionSpace(domain.mesh, V_elt) else: self.function_space = self.options.embedding_space @@ -195,18 +200,12 @@ def setup(self, previous_space=None): # Internal variables to be used # -------------------------------------------------------------------- # - if previous_space is not None: - self.x_in_tmp = Function(previous_space) - else: - self.x_in_tmp = Function(self.time_discretisation.fs) - + self.x_in_tmp = Function(self.original_space) self.x_in = Function(self.function_space) self.x_out = Function(self.function_space) - if previous_space is not None: - self.x_projected = Function(previous_space) - elif self.time_discretisation.idx is None: - self.x_projected = Function(equation.function_space) + if self.time_discretisation.idx is None: + self.x_projected = Function(self.original_space) else: self.x_projected = Function(equation.spaces[self.time_discretisation.idx]) From 5c2980dbf9b7339c9e65651a5faa70af1a566892 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Thu, 29 Feb 2024 19:09:37 +0000 Subject: [PATCH 24/25] make original_space a required argument for wrapper.setup() --- gusto/time_discretisation.py | 7 +++++-- gusto/wrappers.py | 14 +++++--------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 2f6cc009e..9a30e337d 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -184,7 +184,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set") field_idx = equation.field_names.index(field) - subwrapper.setup(self.equation.spaces[field_idx]) + subwrapper.setup(equation.spaces[field_idx]) # Update the function space to that needed by the wrapper self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space @@ -200,7 +200,10 @@ def setup(self, equation, apply_bcs=True, *active_labels): map_if_true=replace_test_function(new_test_mixed)) else: - self.wrapper.setup() + if self.wrapper_name == "supg": + self.wrapper.setup() + else: + self.wrapper.setup(self.fs) self.fs = self.wrapper.function_space if self.solver_parameters is None: self.solver_parameters = self.wrapper.solver_parameters diff --git a/gusto/wrappers.py b/gusto/wrappers.py index c4867e37b..f3fd54c2f 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -36,7 +36,7 @@ def __init__(self, time_discretisation, wrapper_options): self.original_space = None @abstractmethod - def setup(self, original_space=None): + def setup(self, original_space): """ Store the original function space of the prognostic variable. @@ -47,13 +47,9 @@ def setup(self, original_space=None): Args: original_space (:class:`FunctionSpace`): the space that the prognostic variable is defined on. This is a subset space of - a mixed function space when using a MixedFSWrapper. Defaults - to None, when not using a MixedFSWrapper. + a mixed function space when using a MixedFSWrapper. """ - if original_space is not None: - self.original_space = original_space - else: - self.original_space = self.time_discretisation.fs + self.original_space = original_space @abstractmethod def pre_apply(self): @@ -89,7 +85,7 @@ class EmbeddedDGWrapper(Wrapper): the original space. """ - def setup(self, original_space=None): + def setup(self, original_space): """Sets up function spaces and fields needed for this wrapper.""" assert isinstance(self.options, EmbeddedDGOptions), \ @@ -173,7 +169,7 @@ class RecoveryWrapper(Wrapper): field is then returned to the original space. """ - def setup(self, original_space=None): + def setup(self, original_space): """Sets up function spaces and fields needed for this wrapper.""" assert isinstance(self.options, RecoveryOptions), \ From 222c396edeeac69e5f5465c4a2ca4d3904e45b56 Mon Sep 17 00:00:00 2001 From: Tim Andrews Date: Thu, 29 Feb 2024 19:11:25 +0000 Subject: [PATCH 25/25] lint --- gusto/wrappers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index f3fd54c2f..3ca9c4347 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -47,7 +47,7 @@ def setup(self, original_space): Args: original_space (:class:`FunctionSpace`): the space that the prognostic variable is defined on. This is a subset space of - a mixed function space when using a MixedFSWrapper. + a mixed function space when using a MixedFSWrapper. """ self.original_space = original_space