diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index 7369f5ee2..b2df68e2c 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -305,12 +305,13 @@ def add_tracers_to_prognostics(self, domain, active_tracers): name of the active tracer. """ - # Check if there are any conservatively transported tracers. - # If so, ensure that the reference density is indexed before this tracer. + # If there are any conservatively transported tracers, ensure + # that the reference density, if it is also an active tracer, + # is indexed earlier. for i in range(len(active_tracers) - 1): tracer = active_tracers[i] if tracer.transport_eqn == TransportEquationType.tracer_conservative: - ref_density = next(x for x in active_tracers if x.name == tracer.density_name) + ref_density = next((x for x in active_tracers if x.name == tracer.density_name), tracer) j = active_tracers.index(ref_density) if j > i: # Swap the indices of the tracer and the reference density diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index df108a615..4f72af317 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -132,26 +132,58 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.residual = equation.residual if self.field_name is not None and hasattr(equation, "field_names"): - self.idx = equation.field_names.index(self.field_name) - self.fs = equation.spaces[self.idx] - self.residual = self.residual.label_map( - lambda t: t.get(prognostic) == self.field_name, - lambda t: Term( - split_form(t.form)[self.idx].form, - t.labels), - drop) + if isinstance(self.field_name, list): + # Multiple fields are being solved for simultaneously. + # This enables conservative transport to be implemented with SIQN. + # Use the full mixed space for self.fs, with the + # field_name, residual, and BCs being set up later. + self.fs = equation.function_space + self.idx = None + else: + self.idx = equation.field_names.index(self.field_name) + self.fs = equation.spaces[self.idx] + self.residual = self.residual.label_map( + lambda t: t.get(prognostic) == self.field_name, + lambda t: Term( + split_form(t.form)[self.idx].form, + t.labels), + drop) else: self.field_name = equation.field_name self.fs = equation.function_space self.idx = None - bcs = equation.bcs[self.field_name] - if len(active_labels) > 0: - self.residual = self.residual.label_map( - lambda t: any(t.has_label(time_derivative, *active_labels)), - map_if_false=drop) + if isinstance(self.field_name, list): + # Multiple fields are being solved for simultaneously. + # Keep all time derivative terms: + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + + # Only keep active labels for prognostics in the list + # of simultaneously transported variables: + for subname in self.field_name: + field_residual = self.residual.label_map( + lambda t: t.get(prognostic) == subname, + map_if_false=drop) + + residual += field_residual.label_map( + lambda t: t.has_label(*active_labels), + map_if_false=drop) + + self.residual = residual + else: + self.residual = self.residual.label_map( + lambda t: any(t.has_label(time_derivative, *active_labels)), + map_if_false=drop) + + # Set the field name if using simultaneous transport. + if isinstance(self.field_name, list): + self.field_name = equation.field_name + + bcs = equation.bcs[self.field_name] self.evaluate_source = [] self.physics_names = [] @@ -175,7 +207,10 @@ def setup(self, equation, apply_bcs=True, *active_labels): # timestepper should be used instead. if len(field_terms.label_map(lambda t: t.has_label(mass_weighted), map_if_false=drop)) > 0: if len(field_terms.label_map(lambda t: not t.has_label(mass_weighted), map_if_false=drop)) > 0: - raise ValueError(f"Mass-weighted and non-mass-weighted terms are present in a timestepping equation for {field}. As these terms cannot be solved for simultaneously, a split timestepping method should be used instead.") + raise ValueError('Mass-weighted and non-mass-weighted terms are present in a ' + + f'timestepping equation for {field}. As these terms cannot ' + + 'be solved for simultaneously, a split timestepping method ' + + 'should be used instead.') else: # Replace the terms with a mass_weighted label with the # mass_weighted form. It is important that the labels from @@ -199,10 +234,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): 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") + 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(equation.spaces[field_idx], wrapper_bcs) + subwrapper.setup(equation.spaces[field_idx], equation.bcs[field]) # Update the function space to that needed by the wrapper self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space @@ -244,9 +280,19 @@ def setup(self, equation, apply_bcs=True, *active_labels): if not apply_bcs: self.bcs = None elif self.wrapper is not None: - # Transfer boundary conditions onto test function space - self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) - for bc in bcs] + if self.wrapper_name == 'mixed_options': + # Define new Dirichlet BCs on the wrapper-modified + # mixed function space. + self.bcs = [] + for idx, field_name in enumerate(self.equation.field_names): + for bc in equation.bcs[field_name]: + self.bcs.append(DirichletBC(self.fs.sub(idx), + bc.function_arg, + bc.sub_domain)) + else: + # Transfer boundary conditions onto test function space + self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) + for bc in bcs] else: self.bcs = bcs diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py index 76c517a7f..7c4de236b 100644 --- a/gusto/timestepping/semi_implicit_quasi_newton.py +++ b/gusto/timestepping/semi_implicit_quasi_newton.py @@ -115,6 +115,9 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.reference_update_freq = reference_update_freq self.to_update_ref_profile = False + # Flag for if we have simultaneous transport + self.simult = False + # default is to not offcentre transporting velocity but if it # is offcentred then use the same value as alpha self.alpha_u = Constant(alpha) if off_centred_u else Constant(0.5) @@ -148,15 +151,30 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.transported_fields = [] for scheme in transport_schemes: assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" - assert scheme.field_name in equation_set.field_names - self.active_transport.append((scheme.field_name, scheme)) - self.transported_fields.append(scheme.field_name) - # Check that there is a corresponding transport method - method_found = False - for method in spatial_methods: - if scheme.field_name == method.variable and method.term_label == transport: - method_found = True - assert method_found, f'No transport method found for variable {scheme.field_name}' + if isinstance(scheme.field_name, list): + # This means that multiple fields are being transported simultaneously + self.simult = True + for subfield in scheme.field_name: + assert subfield in equation_set.field_names + + # Check that there is a corresponding transport method for + # each field in the list + method_found = False + for method in spatial_methods: + if subfield == method.variable and method.term_label == transport: + method_found = True + assert method_found, f'No transport method found for variable {scheme.field_name}' + self.active_transport.append((scheme.field_name, scheme)) + else: + assert scheme.field_name in equation_set.field_names + + # Check that there is a corresponding transport method + method_found = False + for method in spatial_methods: + if scheme.field_name == method.variable and method.term_label == transport: + method_found = True + self.active_transport.append((scheme.field_name, scheme)) + assert method_found, f'No transport method found for variable {scheme.field_name}' self.diffusion_schemes = [] if diffusion_schemes is not None: @@ -240,7 +258,11 @@ def transporting_velocity(self): def setup_fields(self): """Sets up time levels n, star, p and np1""" self.x = TimeLevelFields(self.equation, 1) - self.x.add_fields(self.equation, levels=("star", "p", "after_slow", "after_fast")) + if self.simult is True: + # If there is any simultaneous transport, add an extra 'simult' field: + self.x.add_fields(self.equation, levels=("star", "p", "simult", "after_slow", "after_fast")) + else: + self.x.add_fields(self.equation, levels=("star", "p", "after_slow", "after_fast")) for aux_eqn, _ in self.auxiliary_equations_and_schemes: self.x.add_fields(aux_eqn) # Prescribed fields for auxiliary eqns should come from prognostics of @@ -282,32 +304,44 @@ def copy_active_tracers(self, x_in, x_out): for name in self.tracers_to_copy: x_out(name).assign(x_in(name)) - def transport_field(self, name, scheme, xstar, xp): + def transport_fields(self, outer, xstar, xp): """ - Performs the transport of a field in xstar, placing the result in xp. + Transports all fields in xstar with a transport scheme + and places the result in xp. Args: - name (str): the name of the field to be transported. - scheme (:class:`TimeDiscretisation`): the time discretisation used - for the transport. + outer (int): the outer loop iteration number xstar (:class:`Fields`): the collection of state fields to be transported. xp (:class:`Fields`): the collection of state fields resulting from the transport. """ - - if name == self.predictor: - # Pre-multiply this variable by (1 - dt*beta*div(u)) - V = xstar(name).function_space() - field_out = Function(V) - self.predictor_interpolator.interpolate() - scheme.apply(field_out, self.predictor_field_in) - - # xp is xstar plus the increment from the transported predictor - xp(name).assign(xstar(name) + field_out - self.predictor_field_in) - else: - # Standard transport - scheme.apply(xp(name), xstar(name)) + for name, scheme in self.active_transport: + if isinstance(name, list): + # Transport multiple fields from xstar simultaneously. + # We transport the mixed function space from xstar to xsimult, then + # extract the updated fields and pass them to xp; this avoids overwriting + # any previously transported fields. + logger.info(f'Semi-implicit Quasi Newton: Transport {outer}: ' + + f'Simultaneous transport of {name}') + scheme.apply(self.x.simult(self.field_name), xstar(self.field_name)) + for field_name in name: + xp(field_name).assign(self.x.simult(field_name)) + else: + logger.info(f'Semi-implicit Quasi Newton: Transport {outer}: {name}') + # transports a single field from xstar and puts the result in xp + if name == self.predictor: + # Pre-multiply this variable by (1 - dt*beta*div(u)) + V = xstar(name).function_space() + field_out = Function(V) + self.predictor_interpolator.interpolate() + scheme.apply(field_out, self.predictor_field_in) + + # xp is xstar plus the increment from the transported predictor + xp(name).assign(xstar(name) + field_out - self.predictor_field_in) + else: + # Standard transport + scheme.apply(xp(name), xstar(name)) def update_reference_profiles(self): """ @@ -367,10 +401,7 @@ def timestep(self): with timed_stage("Transport"): self.io.log_courant(self.fields, 'transporting_velocity', message=f'transporting velocity, outer iteration {outer}') - for name, scheme in self.active_transport: - logger.info(f'Semi-implicit Quasi Newton: Transport {outer}: {name}') - # transports a field from xstar and puts result in xp - self.transport_field(name, scheme, xstar, xp) + self.transport_fields(outer, xstar, xp) # Fast physics ----------------------------------------------------- x_after_fast(self.field_name).assign(xp(self.field_name)) diff --git a/integration-tests/data/simult_SIQN_order0_chkpt.h5 b/integration-tests/data/simult_SIQN_order0_chkpt.h5 new file mode 100644 index 000000000..0d1367606 Binary files /dev/null and b/integration-tests/data/simult_SIQN_order0_chkpt.h5 differ diff --git a/integration-tests/data/simult_SIQN_order1_chkpt.h5 b/integration-tests/data/simult_SIQN_order1_chkpt.h5 new file mode 100644 index 000000000..0309870ea Binary files /dev/null and b/integration-tests/data/simult_SIQN_order1_chkpt.h5 differ diff --git a/integration-tests/model/test_simultaneous_SIQN.py b/integration-tests/model/test_simultaneous_SIQN.py new file mode 100644 index 000000000..ea3a12327 --- /dev/null +++ b/integration-tests/model/test_simultaneous_SIQN.py @@ -0,0 +1,263 @@ +""" +This tests the use of simultaneous transport with the SIQN +timestepping method. A few timesteps are taken with the Bryan-Fritsch +bubble test case, which solves the Compressible Euler Equations. +The two tracers of water vapour and cloud water are being tranpsported +conservatively, which means they need to be transported simultaneously +with the density. +Degree 0 and 1 configurations are tested to ensure that the simultaneous +transport is working with the different wrappers. +""" + +from os.path import join, abspath, dirname +from firedrake import ( + PeriodicIntervalMesh, ExtrudedMesh, SpatialCoordinate, conditional, cos, pi, + sqrt, NonlinearVariationalProblem, NonlinearVariationalSolver, TestFunction, + dx, TrialFunction, Function, as_vector, LinearVariationalProblem, + LinearVariationalSolver, Constant, BrokenElement +) +from gusto import * +import pytest + + +def run_simult_SIQN(tmpdir, order): + + if order == 0: + ncolumns = 20 + nlayers = 20 + u_eqn_type = "vector_advection_form" + else: + ncolumns = 10 + nlayers = 10 + u_eqn_type = "vector_invariant_form" + + dt = 2.0 + tmax = 10.0 + + domain_width = 10000. # domain width, in m + domain_height = 10000. # domain height, in m + zc = 2000. # vertical centre of bubble, in m + rc = 2000. # radius of bubble, in m + Tdash = 2.0 # strength of temperature perturbation, in K + Tsurf = 320.0 # background theta_e value, in K + total_water = 0.02 # total moisture mixing ratio, in kg/kg + + # Domain + mesh_name = 'bryan_fritsch_mesh' + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh( + base_mesh, layers=nlayers, layer_height=domain_height/nlayers, name=mesh_name + ) + domain = Domain(mesh, dt, 'CG', order) + + # Set up the tracers and their transport schemes + V_rho = domain.spaces('DG') + V_theta = domain.spaces('theta') + + tracers = [WaterVapour(space='theta', + transport_eqn=TransportEquationType.tracer_conservative, + density_name='rho'), + CloudWater(space='theta', + transport_eqn=TransportEquationType.tracer_conservative, + density_name='rho')] + + # Equation + params = CompressibleParameters() + eqns = CompressibleEulerEquations( + domain, params, active_tracers=tracers, u_transport_option=u_eqn_type + ) + + # I/O + output_dirname = tmpdir+"/simult_SIQN_order"+str(order) + + output = OutputParameters( + dirname=output_dirname, dumpfreq=5, chkptfreq=5, checkpoint=True + ) + io = IO(domain, output) + + # Set up transport schemes + if order == 0: + VDG1 = domain.spaces("DG1_equispaced") + VCG1 = FunctionSpace(mesh, "CG", 1) + Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) + Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) + + u_opts = RecoveryOptions(embedding_space=Vu_DG1, + recovered_space=Vu_CG1, + boundary_method=BoundaryMethod.taylor) + theta_opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1) + + suboptions = {'rho': RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + boundary_method=BoundaryMethod.taylor), + 'water_vapour': ConservativeRecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + rho_name="rho", + orig_rho_space=V_rho), + 'cloud_water': ConservativeRecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + rho_name="rho", + orig_rho_space=V_rho)} + else: + theta_opts = EmbeddedDGOptions() + Vt_brok = FunctionSpace(mesh, BrokenElement(V_theta.ufl_element())) + suboptions = {'rho': EmbeddedDGOptions(embedding_space=Vt_brok), + 'water_vapour': ConservativeEmbeddedDGOptions(embedding_space=Vt_brok, + rho_name="rho", + orig_rho_space=V_rho), + 'cloud_water': ConservativeEmbeddedDGOptions(embedding_space=Vt_brok, + rho_name="rho", + orig_rho_space=V_rho)} + + transported_fields = [SSPRK3(domain, "theta", options=theta_opts)] + + mixed_opts = MixedFSOptions(suboptions=suboptions) + transported_fields.append(SSPRK3(domain, ["rho", "water_vapour", "cloud_water"], options=mixed_opts, rk_formulation=RungeKuttaFormulation.predictor)) + + if order == 0: + transported_fields.append(SSPRK3(domain, 'u', options=u_opts)) + else: + transported_fields.append(TrapeziumRule(domain, 'u')) + + transport_methods = [ + DGUpwind(eqns, field) for field in + ["u", "rho", "theta", "water_vapour", "cloud_water"] + ] + + # Linear solver + linear_solver = CompressibleSolver(eqns) + + # Physics schemes (condensation/evaporation) + physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver, physics_schemes=physics_schemes + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + water_v0 = stepper.fields("water_vapour") + water_c0 = stepper.fields("cloud_water") + + # spaces + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + x, z = SpatialCoordinate(mesh) + quadrature_degree = (4, 4) + dxp = dx(degree=(quadrature_degree)) + + # Define constant theta_e and water_t + theta_e = Function(Vt).assign(Tsurf) + water_t = Function(Vt).assign(total_water) + + # Calculate hydrostatic fields + saturated_hydrostatic_balance(eqns, stepper.fields, theta_e, water_t) + + # make mean fields + theta_b = Function(Vt).assign(theta0) + rho_b = Function(Vr).assign(rho0) + water_vb = Function(Vt).assign(water_v0) + water_cb = Function(Vt).assign(water_t - water_vb) + + # define perturbation + xc = domain_width / 2 + r = sqrt((x - xc) ** 2 + (z - zc) ** 2) + theta_pert = Function(Vt).interpolate( + conditional( + r > rc, + 0.0, + Tdash * (cos(pi * r / (2.0 * rc))) ** 2 + ) + ) + + # define initial theta + theta0.interpolate(theta_b * (theta_pert / 300.0 + 1.0)) + + # find perturbed rho + gamma = TestFunction(Vr) + rho_trial = TrialFunction(Vr) + a = gamma * rho_trial * dxp + L = gamma * (rho_b * theta_b / theta0) * dxp + rho_problem = LinearVariationalProblem(a, L, rho0) + rho_solver = LinearVariationalSolver(rho_problem) + rho_solver.solve() + + # find perturbed water_v + w_v = Function(Vt) + phi = TestFunction(Vt) + rho_averaged = Function(Vt) + rho_recoverer = Recoverer(rho0, rho_averaged) + rho_recoverer.project() + + exner = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) + p = thermodynamics.p(eqns.parameters, exner) + T = thermodynamics.T(eqns.parameters, theta0, exner, r_v=w_v) + w_sat = thermodynamics.r_sat(eqns.parameters, T, p) + + w_functional = (phi * w_v * dxp - phi * w_sat * dxp) + w_problem = NonlinearVariationalProblem(w_functional, w_v) + w_solver = NonlinearVariationalSolver(w_problem) + w_solver.solve() + + water_v0.assign(w_v) + water_c0.assign(water_t - water_v0) + + # wind initially zero + u0.project(as_vector( + [Constant(0.0, domain=mesh), Constant(0.0, domain=mesh)] + )) + + stepper.set_reference_profiles( + [ + ('rho', rho_b), + ('theta', theta_b), + ('water_vapour', water_vb), + ('cloud_water', water_cb) + ] + ) + + # --------------------------------------------------------------------- # + # Run + # --------------------------------------------------------------------- # + + stepper.run(t=0, tmax=tmax) + + # State for checking checkpoints + checkpoint_name = 'simult_SIQN_order'+str(order)+'_chkpt.h5' + new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') + check_output = OutputParameters(dirname=output_dirname, + checkpoint_pickup_filename=new_path, + checkpoint=True) + check_mesh = pick_up_mesh(check_output, mesh_name) + check_domain = Domain(check_mesh, dt, "CG", order) + check_eqn = CompressibleEulerEquations(check_domain, params, active_tracers=tracers, u_transport_option=u_eqn_type) + check_io = IO(check_domain, check_output) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, [], []) + check_stepper.io.pick_up_from_checkpoint(check_stepper.fields) + + return stepper, check_stepper + + +@pytest.mark.parametrize("order", [0, 1]) +def test_simult_SIQN(tmpdir, order): + + dirname = str(tmpdir) + stepper, check_stepper = run_simult_SIQN(dirname, order) + + for variable in ['u', 'rho', 'theta', 'water_vapour', 'cloud_water']: + new_variable = stepper.fields(variable) + check_variable = check_stepper.fields(variable) + diff_array = new_variable.dat.data - check_variable.dat.data + error = np.linalg.norm(diff_array) / np.linalg.norm(check_variable.dat.data) + + # Slack values chosen to be robust to different platforms + assert error < 1e-10, f'Values for {variable} in the ' + \ + f'order {order} elements test do not match KGO values'