diff --git a/examples/shallow_water/linear_thermal_galewsky_jet.py b/examples/shallow_water/linear_thermal_galewsky_jet.py new file mode 100644 index 000000000..a4fd68631 --- /dev/null +++ b/examples/shallow_water/linear_thermal_galewsky_jet.py @@ -0,0 +1,220 @@ +""" +A linearised form of the steady thermal Galewsky jet. The initial conditions are +taken from Hartney et al, 2024: ``A compatible finite element discretisation +for moist shallow water equations'' (without the perturbation). + +This uses an icosahedral mesh of the sphere. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + SpatialCoordinate, pi, assemble, dx, Constant, exp, conditional, Function, + cos +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, DefaultTransport, + DGUpwind, ForwardEuler, ShallowWaterParameters, NumericalIntegral, + LinearThermalShallowWaterEquations, GeneralIcosahedralSphereMesh, + ZonalComponent, ThermalSWSolver, lonlatr_from_xyz, xyz_vector_from_lonlatr, + RelativeVorticity, MeridionalComponent +) + +import numpy as np + +linear_thermal_galewsky_jet_defaults = { + 'ncells_per_edge': 12, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 6.*24.*60.*60., # 6 days + 'dumpfreq': 96, # once per day with default options + 'dirname': 'linear_thermal_galewsky' +} + + +def linear_thermal_galewsky_jet( + ncells_per_edge=linear_thermal_galewsky_jet_defaults['ncells_per_edge'], + dt=linear_thermal_galewsky_jet_defaults['dt'], + tmax=linear_thermal_galewsky_jet_defaults['tmax'], + dumpfreq=linear_thermal_galewsky_jet_defaults['dumpfreq'], + dirname=linear_thermal_galewsky_jet_defaults['dirname'] +): + # ----------------------------------------------------------------- # + # Parameters for test case + # ----------------------------------------------------------------- # + + R = 6371220. # planetary radius (m) + H = 10000. # reference depth (m) + u_max = 80.0 # Max amplitude of the zonal wind (m/s) + phi0 = pi/7. # latitude of the southern boundary of the jet (rad) + phi1 = pi/2. - phi0 # latitude of the northern boundary of the jet (rad) + db = 1.0 # diff in buoyancy between equator and poles (m/s^2) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ----------------------------------------------------------------- # + # Set up model objects + # ----------------------------------------------------------------- # + + # Domain + mesh = GeneralIcosahedralSphereMesh(R, ncells_per_edge, degree=2) + xyz = SpatialCoordinate(mesh) + domain = Domain(mesh, dt, 'BDM', element_order) + + # Equation + parameters = ShallowWaterParameters(H=H) + Omega = parameters.Omega + fexpr = 2*Omega*xyz[2]/R + eqns = LinearThermalShallowWaterEquations(domain, parameters, fexpr=fexpr) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_nc=True, dump_vtus=False, + dumplist=['D', 'b'] + ) + diagnostic_fields = [ + ZonalComponent('u'), MeridionalComponent('u'), RelativeVorticity() + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transport_schemes = [ForwardEuler(domain, "D")] + transport_methods = [DefaultTransport(eqns, "D"), DGUpwind(eqns, "b")] + + # Linear solver + linear_solver = ThermalSWSolver(eqns) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transport_schemes, transport_methods, + linear_solver=linear_solver + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0_field = stepper.fields("u") + D0_field = stepper.fields("D") + b0_field = stepper.fields("b") + + # Parameters + g = parameters.g + Omega = parameters.Omega + e_n = np.exp(-4./((phi1-phi0)**2)) + + _, lat, _ = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2]) + lat_VD = Function(D0_field.function_space()).interpolate(lat) + + # ------------------------------------------------------------------------ # + # Obtain u and D (by integration of analytic expression) + # ------------------------------------------------------------------------ # + + # Buoyancy + bexpr = g - db*cos(lat) + + # Wind -- UFL expression + u_zonal = conditional( + lat <= phi0, 0.0, + conditional( + lat >= phi1, 0.0, + u_max / e_n * exp(1.0 / ((lat - phi0) * (lat - phi1))) + ) + ) + uexpr = xyz_vector_from_lonlatr(u_zonal, Constant(0.0), Constant(0.0), xyz) + + # Numpy function + def u_func(y): + u_array = np.where( + y <= phi0, 0.0, + np.where( + y >= phi1, 0.0, + u_max / e_n * np.exp(1.0 / ((y - phi0) * (y - phi1))) + ) + ) + return u_array + + # Function for depth field in terms of u function + def h_func(y): + h_array = u_func(y)*R/g*( + 2*Omega*np.sin(y) + ) + + return h_array + + # Find h from numerical integral + D0_integral = Function(D0_field.function_space()) + h_integral = NumericalIntegral(-pi/2, pi/2) + h_integral.tabulate(h_func) + D0_integral.dat.data[:] = h_integral.evaluate_at(lat_VD.dat.data[:]) + Dexpr = H - D0_integral + + # Obtain fields + u0_field.project(uexpr) + D0_field.interpolate(Dexpr) + + # Adjust mean value of initial D + C = Function(D0_field.function_space()).assign(Constant(1.0)) + area = assemble(C*dx) + Dmean = assemble(D0_field*dx)/area + D0_field -= Dmean + D0_field += Constant(H) + b0_field.interpolate(bexpr) + + # Set reference profiles + Dbar = Function(D0_field.function_space()).assign(H) + bbar = Function(b0_field.function_space()).interpolate(g) + stepper.set_reference_profiles([('D', Dbar), ('b', bbar)]) + + # ----------------------------------------------------------------- # + # Run + # ----------------------------------------------------------------- # + + stepper.run(t=0, tmax=tmax) + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncells_per_edge', + help="The number of cells per edge of icosahedron", + type=int, + default=linear_thermal_galewsky_jet_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=linear_thermal_galewsky_jet_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=linear_thermal_galewsky_jet_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=linear_thermal_galewsky_jet_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=linear_thermal_galewsky_jet_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + linear_thermal_galewsky_jet(**vars(args)) diff --git a/examples/shallow_water/moist_thermal_equivb_gw.py b/examples/shallow_water/moist_thermal_equivb_gw.py new file mode 100644 index 000000000..912e9deba --- /dev/null +++ b/examples/shallow_water/moist_thermal_equivb_gw.py @@ -0,0 +1,230 @@ +""" +A gravity wave on the sphere, solved with the moist thermal shallow water +equations. The initial conditions are saturated and cloudy everywhere. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + SpatialCoordinate, pi, sqrt, min_value, cos, Constant, Function, exp, sin +) +from gusto import ( + Domain, IO, OutputParameters, DGUpwind, ShallowWaterParameters, + ThermalShallowWaterEquations, lonlatr_from_xyz, MeridionalComponent, + GeneralIcosahedralSphereMesh, SubcyclingOptions, ZonalComponent, + PartitionedCloud, RungeKuttaFormulation, SSPRK3, ThermalSWSolver, + SemiImplicitQuasiNewton, xyz_vector_from_lonlatr +) + +moist_thermal_gw_defaults = { + 'ncells_per_edge': 12, # number of cells per icosahedron edge + 'dt': 900.0, # 15 minutes + 'tmax': 5.*24.*60.*60., # 5 days + 'dumpfreq': 48, # dump twice per day + 'dirname': 'moist_thermal_equivb_gw' +} + + +def moist_thermal_gw( + ncells_per_edge=moist_thermal_gw_defaults['ncells_per_edge'], + dt=moist_thermal_gw_defaults['dt'], + tmax=moist_thermal_gw_defaults['tmax'], + dumpfreq=moist_thermal_gw_defaults['dumpfreq'], + dirname=moist_thermal_gw_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + radius = 6371220. # planetary radius (m) + mean_depth = 5960. # reference depth (m) + q0 = 0.0115 # saturation curve coefficient (kg/kg) + beta2 = 9.80616*10 # thermal feedback coefficient (m/s^2) + nu = 1.5 # dimensionless parameter in saturation curve + R0 = pi/9. # radius of perturbation (rad) + lamda_c = -pi/2. # longitudinal centre of perturbation (rad) + phi_c = pi/6. # latitudinal centre of perturbation (rad) + phi_0 = 3.0e4 # scale factor for poleward buoyancy gradient + epsilon = 1/300 # linear air expansion coeff (1/K) + u_max = 20. # max amplitude of the zonal wind (m/s) + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2) + degree = 1 + domain = Domain(mesh, dt, "BDM", degree) + xyz = SpatialCoordinate(mesh) + + # Equation parameters + parameters = ShallowWaterParameters(H=mean_depth, q0=q0, nu=nu, beta2=beta2) + Omega = parameters.Omega + fexpr = 2*Omega*xyz[2]/radius + + # Equation + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, + equivalent_buoyancy=True + ) + + # IO + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_nc=True, dump_vtus=False, + dumplist=['b_e', 'D'] + ) + diagnostic_fields = [ + ZonalComponent('u'), MeridionalComponent('u'), PartitionedCloud(eqns) + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + transport_methods = [ + DGUpwind(eqns, field_name) for field_name in eqns.field_names + ] + + linear_solver = ThermalSWSolver(eqns) + + # ------------------------------------------------------------------------ # + # Timestepper + # ------------------------------------------------------------------------ # + + subcycling_opts = SubcyclingOptions(subcycle_by_courant=0.25) + transported_fields = [ + SSPRK3(domain, "u", subcycling_options=subcycling_opts), + SSPRK3( + domain, "D", subcycling_options=subcycling_opts, + rk_formulation=RungeKuttaFormulation.linear + ), + SSPRK3(domain, "b_e", subcycling_options=subcycling_opts), + SSPRK3(domain, "q_t", subcycling_options=subcycling_opts), + ] + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver, + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + D0 = stepper.fields("D") + b0 = stepper.fields("b_e") + qt0 = stepper.fields("q_t") + + lamda, phi, _ = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2]) + + # Velocity -- a solid body rotation + uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, xyz) + + # Buoyancy -- dependent on latitude + g = parameters.g + w = Omega*radius*u_max + (u_max**2)/2 + sigma = w/10 + theta_0 = epsilon*phi_0**2 + numerator = ( + theta_0 + sigma*((cos(phi))**2) * ( + (w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma) + ) + ) + denominator = ( + phi_0**2 + (w + sigma)**2 + * (sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2 + ) + theta = numerator / denominator + b_guess = parameters.g * (1 - theta) + + # Depth -- in balance with the contribution of a perturbation + Dexpr = mean_depth - (1/g)*(w + sigma)*((sin(phi))**2) + + # Perturbation + lsq = (lamda - lamda_c)**2 + thsq = (phi - phi_c)**2 + rsq = min_value(R0**2, lsq+thsq) + r = sqrt(rsq) + pert = 2000 * (1 - r/R0) + Dexpr += pert + + # Actual initial buoyancy is specified through equivalent buoyancy + q_t = 0.03 + b_init = Function(b0.function_space()).interpolate(b_guess) + b_e_init = Function(b0.function_space()).interpolate(b_init - beta2*q_t) + q_v_init = Function(qt0.function_space()).interpolate(q_t) + + # Iterate to obtain equivalent buoyancy and saturation water vapour + n_iterations = 10 + + for _ in range(n_iterations): + q_sat_expr = q0*mean_depth/Dexpr * exp(nu*(1-b_e_init/g)) + dq_sat_dq_v_expr = nu*beta2/g*q_sat_expr + q_v_init.interpolate(q_v_init - (q_sat_expr - q_v_init)/(dq_sat_dq_v_expr - 1.0)) + b_e_init.interpolate(b_init - beta2*q_v_init) + + # Water vapour set to saturation amount + vexpr = q0*mean_depth/Dexpr * exp(nu*(1-b_e_init/g)) + + # Back out the initial buoyancy using b_e and q_v + bexpr = b_e_init + beta2*vexpr + + u0.project(uexpr) + D0.interpolate(Dexpr) + b0.interpolate(bexpr) + qt0.interpolate(Constant(0.03)) + + # Set reference profiles + Dbar = Function(D0.function_space()).assign(mean_depth) + bbar = Function(b0.function_space()).interpolate(bexpr) + stepper.set_reference_profiles([('D', Dbar), ('b_e', bbar)]) + + # ----------------------------------------------------------------- # + # Run + # ----------------------------------------------------------------- # + + stepper.run(t=0, tmax=tmax) + + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncells_per_edge', + help="The number of cells per edge of icosahedron", + type=int, + default=moist_thermal_gw_defaults['ncells_per_edge'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=moist_thermal_gw_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=moist_thermal_gw_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=moist_thermal_gw_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=moist_thermal_gw_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + moist_thermal_gw(**vars(args)) diff --git a/examples/shallow_water/moist_thermal_williamson_5.py b/examples/shallow_water/moist_thermal_williamson_5.py index 9ce7da77e..f17c278c0 100644 --- a/examples/shallow_water/moist_thermal_williamson_5.py +++ b/examples/shallow_water/moist_thermal_williamson_5.py @@ -21,7 +21,7 @@ ) from gusto import ( Domain, IO, OutputParameters, Timestepper, RK4, DGUpwind, - ShallowWaterParameters, ShallowWaterEquations, Sum, + ShallowWaterParameters, ThermalShallowWaterEquations, Sum, lonlatr_from_xyz, InstantRain, SWSaturationAdjustment, WaterVapour, CloudWater, Rain, GeneralIcosahedralSphereMesh, RelativeVorticity, ZonalComponent, MeridionalComponent @@ -99,8 +99,8 @@ def moist_thermal_williamson_5( tracers = [ WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG') ] - eqns = ShallowWaterEquations( - domain, parameters, fexpr=fexpr, bexpr=tpexpr, thermal=True, + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, topog_expr=tpexpr, active_tracers=tracers, u_transport_option=u_eqn_type ) diff --git a/examples/shallow_water/test_shallow_water_examples.py b/examples/shallow_water/test_shallow_water_examples.py index eb64e94e0..7a96c6316 100644 --- a/examples/shallow_water/test_shallow_water_examples.py +++ b/examples/shallow_water/test_shallow_water_examples.py @@ -127,3 +127,37 @@ def test_williamson_5(): @pytest.mark.parallel(nprocs=4) def test_williamson_5_parallel(): test_williamson_5() + + +def test_linear_thermal_galewsky_jet(): + from linear_thermal_galewsky_jet import linear_thermal_galewsky_jet + test_name = 'linear_thermal_galewsky_jet' + linear_thermal_galewsky_jet( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def test_linear_thermal_galewsky_jet_parallel(): + test_linear_thermal_galewsky_jet() + + +def test_moist_thermal_eqiuvb_gw(): + from moist_thermal_equivb_gw import moist_thermal_gw + test_name = 'moist_thermal_gw' + moist_thermal_gw( + ncells_per_edge=4, + dt=1800., + tmax=18000., + dumpfreq=10, + dirname=make_dirname(test_name) + ) + + +@pytest.mark.parallel(nprocs=4) +def moist_thermal_eqiuvb_gw_parallel(): + test_moist_thermal_eqiuvb_gw() diff --git a/examples/shallow_water/thermal_williamson_2.py b/examples/shallow_water/thermal_williamson_2.py index d2c4e4e29..7adc305dc 100644 --- a/examples/shallow_water/thermal_williamson_2.py +++ b/examples/shallow_water/thermal_williamson_2.py @@ -15,7 +15,7 @@ from firedrake import Function, SpatialCoordinate, sin, cos from gusto import ( Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, - TrapeziumRule, ShallowWaterParameters, ShallowWaterEquations, + TrapeziumRule, ShallowWaterParameters, ThermalShallowWaterEquations, RelativeVorticity, PotentialVorticity, SteadyStateError, ZonalComponent, MeridionalComponent, ThermalSWSolver, xyz_vector_from_lonlatr, lonlatr_from_xyz, GeneralIcosahedralSphereMesh, @@ -71,8 +71,8 @@ def thermal_williamson_2( params = ShallowWaterParameters(H=mean_depth, g=g) Omega = params.Omega fexpr = 2*Omega*z/radius - eqns = ShallowWaterEquations( - domain, params, fexpr=fexpr, u_transport_option=u_eqn_type, thermal=True + eqns = ThermalShallowWaterEquations( + domain, params, fexpr=fexpr, u_transport_option=u_eqn_type ) # IO diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index b67e5b079..efe41b183 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -73,7 +73,8 @@ def williamson_5( rsq = min_value(R0**2, (lamda - lamda_c)**2 + (phi - phi_c)**2) r = sqrt(rsq) tpexpr = mountain_height * (1 - r/R0) - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr) + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + topog_expr=tpexpr) # I/O output = OutputParameters( diff --git a/figures/shallow_water/linear_thermal_galewsky_final.png b/figures/shallow_water/linear_thermal_galewsky_final.png new file mode 100644 index 000000000..8400dc6fd Binary files /dev/null and b/figures/shallow_water/linear_thermal_galewsky_final.png differ diff --git a/figures/shallow_water/linear_thermal_galewsky_initial.png b/figures/shallow_water/linear_thermal_galewsky_initial.png new file mode 100644 index 000000000..5fdc323a0 Binary files /dev/null and b/figures/shallow_water/linear_thermal_galewsky_initial.png differ diff --git a/figures/shallow_water/moist_thermal_equivb_gw_final.png b/figures/shallow_water/moist_thermal_equivb_gw_final.png new file mode 100644 index 000000000..e5aaa2330 Binary files /dev/null and b/figures/shallow_water/moist_thermal_equivb_gw_final.png differ diff --git a/figures/shallow_water/moist_thermal_equivb_gw_initial.png b/figures/shallow_water/moist_thermal_equivb_gw_initial.png new file mode 100644 index 000000000..94468de03 Binary files /dev/null and b/figures/shallow_water/moist_thermal_equivb_gw_initial.png differ diff --git a/gusto/core/configuration.py b/gusto/core/configuration.py index 133c2c3f6..f19ef0167 100644 --- a/gusto/core/configuration.py +++ b/gusto/core/configuration.py @@ -149,6 +149,15 @@ class ShallowWaterParameters(Configuration): g = 9.80616 Omega = 7.292e-5 # rotation rate H = None # mean depth + # Factor that multiplies the vapour in the equivalent buoyancy + # formulation of the thermal shallow water equations + beta2 = None + # Scaling factor for the saturation function in the equivalent buoyancy + # formulation of the thermal shallow water equations + nu = None + # Scaling factor for the saturation function in the equivalent buoyancy + # formulation of the thermal shallow water equations + q0 = None class WrapperOptions(Configuration, metaclass=ABCMeta): diff --git a/gusto/core/fields.py b/gusto/core/fields.py index 4d204642c..5fdf011d4 100644 --- a/gusto/core/fields.py +++ b/gusto/core/fields.py @@ -28,6 +28,12 @@ def add_field(self, name, space, subfield_names=None): variables. Defaults to None. """ + if name == 'X': + raise ValueError( + 'As "X" is used for the generic field state, it is not allowed ' + + 'as the name of a field -- please choose a different name!' + ) + value = Function(space, name=name) setattr(self, name, value) self.fields.append(value) @@ -35,6 +41,14 @@ def add_field(self, name, space, subfield_names=None): if len(space) > 1: assert len(space) == len(subfield_names) for field_name, field in zip(subfield_names, value.subfunctions): + + if field_name == 'X': + raise ValueError( + 'As "X" is used for the generic field state, it is not ' + + 'allowed as the name of a field -- please choose a ' + + 'different name!' + ) + setattr(self, field_name, field) field.rename(field_name) self.fields.append(field) @@ -114,6 +128,7 @@ def __init__(self, prognostic_fields, prescribed_fields, *fields_to_dump): self.to_pick_up = [] self._field_types = [] self._field_names = [] + self.X = prognostic_fields.np1 # Add pointers to prognostic fields for field in prognostic_fields.np1.fields: diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index dce6fac69..6484b2afe 100644 --- a/gusto/diagnostics/shallow_water_diagnostics.py +++ b/gusto/diagnostics/shallow_water_diagnostics.py @@ -1,13 +1,16 @@ """Common diagnostic fields for the Shallow Water equations.""" -from firedrake import (dx, TestFunction, TrialFunction, grad, inner, curl, - LinearVariationalProblem, LinearVariationalSolver) +from firedrake import ( + dx, TestFunction, TrialFunction, grad, inner, curl, Function, Interpolator, + LinearVariationalProblem, LinearVariationalSolver, conditional +) from gusto.diagnostics.diagnostics import DiagnosticField, Energy __all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy", "ShallowWaterPotentialEnstrophy", "PotentialVorticity", - "RelativeVorticity", "AbsoluteVorticity"] + "RelativeVorticity", "AbsoluteVorticity", "PartitionedVapour", + "PartitionedCloud"] class ShallowWaterKineticEnergy(Energy): @@ -274,3 +277,104 @@ def setup(self, domain, state_fields): state_fields (:class:`StateFields`): the model's field container. """ super().setup(domain, state_fields, vorticity_type="relative") + + +class PartitionedVapour(DiagnosticField): + """ + Diagnostic for computing the vapour in the equivalent buoyancy formulation + of the moist thermal shallow water equations. + """ + name = "PartitionedVapour" + + def __init__(self, equation, name='q_t', space=None, + method='interpolate'): + """ + Args: + equation (:class:`PrognosticEquation`): the model's equation. + name (str, optional): name of the total moisture field to use to + compute the vapour from. Defaults to total moisture, q_t. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.fname = name + self.equation = equation + super().__init__(space=space, method=method, required_fields=(self.fname,)) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + q_t = state_fields(self.fname) + space = domain.spaces("DG") + self.qsat_func = Function(space) + + qsat_expr = self.equation.compute_saturation(state_fields.X( + self.equation.field_name)) + self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func) + self.expr = conditional(q_t < self.qsat_func, q_t, self.qsat_func) + + super().setup(domain, state_fields, space=space) + + def compute(self): + """Performs the computation of the diagnostic field.""" + self.qsat_interpolator.interpolate() + super().compute() + + +class PartitionedCloud(DiagnosticField): + """ + Diagnostic for computing the cloud in the equivalent buoyancy formulation + of the moist thermal shallow water equations. + """ + name = "PartitionedCloud" + + def __init__(self, equation, name='q_t', space=None, + method='interpolate'): + """ + Args: + equation (:class:`PrognosticEquation`): the model's equation. + name (str, optional): name of the total moisture field to use to + compute the vapour from. Defaults to total moisture, q_t. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.fname = name + self.equation = equation + super().__init__(space=space, method=method, required_fields=(self.fname,)) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + q_t = state_fields(self.fname) + space = domain.spaces("DG") + self.qsat_func = Function(space) + + qsat_expr = self.equation.compute_saturation(state_fields.X( + self.equation.field_name)) + self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func) + vapour = conditional(q_t < self.qsat_func, q_t, self.qsat_func) + self.expr = q_t - vapour + + super().setup(domain, state_fields, space=space) + + def compute(self): + """Performs the computation of the diagnostic field.""" + self.qsat_interpolator.interpolate() + super().compute() diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index b2df68e2c..a562795d6 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -94,7 +94,7 @@ def __init__(self, field_names, domain, space_names, None. active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included - in the equations.. Defaults to None. + in the equations. Defaults to None. """ self.field_names = field_names diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 854266952..fd61cedf7 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -1,19 +1,22 @@ """Classes for defining variants of the shallow-water equations.""" from firedrake import (inner, dx, div, FunctionSpace, FacetNormal, jump, avg, - dS, split) -from firedrake.fml import subject + dS, split, conditional, exp) +from firedrake.fml import subject, drop from gusto.core.labels import (time_derivative, transport, prognostic, linearisation, pressure_gradient, coriolis) from gusto.equations.common_forms import ( advection_form, advection_form_1d, continuity_form, continuity_form_1d, vector_invariant_form, kinetic_energy_form, advection_equation_circulation_form, diffusion_form_1d, - linear_continuity_form + linear_continuity_form, linear_advection_form ) from gusto.equations.prognostic_equations import PrognosticEquationSet + __all__ = ["ShallowWaterEquations", "LinearShallowWaterEquations", + "ThermalShallowWaterEquations", + "LinearThermalShallowWaterEquations", "ShallowWaterEquations_1d", "LinearShallowWaterEquations_1d"] @@ -21,16 +24,15 @@ class ShallowWaterEquations(PrognosticEquationSet): u""" Class for the (rotating) shallow-water equations, which evolve the velocity 'u' and the depth field 'D', via some variant of: \n - ∂u/∂t + (u.∇)u + f×u + g*∇(D+b) = 0, \n + ∂u/∂t + (u.∇)u + f×u + g*∇(D+B) = 0, \n ∂D/∂t + ∇.(D*u) = 0, \n - for Coriolis parameter 'f' and bottom surface 'b'. + for Coriolis parameter 'f' and bottom surface 'B'. """ - def __init__(self, domain, parameters, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, topog_expr=None, space_names=None, linearisation_map='default', u_transport_option='vector_invariant_form', - no_normal_flow_bc_ids=None, active_tracers=None, - thermal=False): + no_normal_flow_bc_ids=None, active_tracers=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -39,8 +41,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, the model's physical parameters. fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis parameter. Defaults to None. - bexpr (:class:`ufl.Expr`, optional): an expression for the bottom - surface of the fluid. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. space_names (dict, optional): a dictionary of strings for names of the function spaces to use for the spatial discretisation. The keys are the names of the prognostic variables. Defaults to None @@ -62,47 +64,66 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations. Defaults to None. - thermal (flag, optional): specifies whether the equations have a - thermal or buoyancy variable that feeds back on the momentum. - Defaults to False. - - Raises: - NotImplementedError: active tracers are not yet implemented. """ - self.thermal = thermal - field_names = ['u', 'D'] - - if space_names is None: - space_names = {'u': 'HDiv', 'D': 'L2'} - if active_tracers is None: active_tracers = [] - if self.thermal: - field_names.append('b') - if 'b' not in space_names.keys(): - space_names['b'] = 'L2' - if linearisation_map == 'default': # Default linearisation is time derivatives, pressure gradient and # transport term from depth equation. Don't include active tracers linearisation_map = lambda t: \ - t.get(prognostic) in ['u', 'D', 'b'] \ - and (any(t.has_label(time_derivative, pressure_gradient)) - or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport))) + t.get(prognostic) in ['u', 'D'] \ + and (any(t.has_label(time_derivative, coriolis, pressure_gradient)) + or (t.get(prognostic) in ['D'] and t.has_label(transport))) + + field_names = ['u', 'D'] + space_names = {'u': 'HDiv', 'D': 'L2'} + super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) self.parameters = parameters - g = parameters.g - H = parameters.H + self.domain = domain + self.active_tracers = active_tracers + + self._setup_residual(fexpr, topog_expr, u_transport_option) + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms( + self.residual, self.linearisation_map) + + def _setup_residual(self, fexpr, topog_expr, u_transport_option): + """ + Sets up the residual for the shallow water equations. This + is separate from the __init__ method because the thermal + shallow water equation class expands on this equation set by + adding additional fields that depend on the formulation. This + increases the size of the mixed function space and the + residual must be setup after this has happened. + + Args: + fexpr (:class:`ufl.Expr`): an expression for the Coroilis + parameter. + topog_expr (:class:`ufl.Expr`): an expression for the + bottom surface of the fluid. + u_transport_option (str): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. + """ + + g = self.parameters.g w, phi = self.tests[0:2] u, D = split(self.X)[0:2] u_trial = split(self.trials)[0] + Dbar = split(self.X_ref)[1] # -------------------------------------------------------------------- # # Time Derivative Terms @@ -114,116 +135,84 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(domain, w, u, u), 'u') + u_adv = prognostic(vector_invariant_form(self.domain, w, u, u), 'u') elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), 'u') elif u_transport_option == "circulation_form": ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') - u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form + u_adv = prognostic(advection_equation_circulation_form(self.domain, w, u, u), 'u') + ke_form else: - raise ValueError("Invalid u_transport_option: %s" % u_transport_option) + raise ValueError("Invalid u_transport_option: %s" % self.u_transport_option) # Depth transport term D_adv = prognostic(continuity_form(phi, D, u), 'D') # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): - linear_D_adv = linear_continuity_form(phi, H, u_trial) + linear_D_adv = linear_continuity_form(phi, Dbar, u_trial) # Add linearisation to D_adv D_adv = linearisation(D_adv, linear_D_adv) adv_form = subject(u_adv + D_adv, self.X) # Add transport of tracers - if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(active_tracers) - # Add transport of buoyancy, if thermal shallow water equations - if self.thermal: - gamma = self.tests[2] - b = split(self.X)[2] - b_adv = prognostic(advection_form(gamma, b, u), 'b') - adv_form += subject(b_adv, self.X) + if len(self.active_tracers) > 0: + adv_form += self.generate_tracer_transport_terms( + self.active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term # -------------------------------------------------------------------- # - # Add pressure gradient only if not doing thermal - if self.thermal: - residual = (mass_form + adv_form) - else: - pressure_gradient_form = pressure_gradient( - subject(prognostic(-g*div(w)*D*dx, 'u'), self.X)) + pressure_gradient_form = pressure_gradient( + subject(prognostic(-g*div(w)*D*dx, 'u'), self.X)) - residual = (mass_form + adv_form + pressure_gradient_form) + residual = (mass_form + adv_form + pressure_gradient_form) # -------------------------------------------------------------------- # - # Extra Terms (Coriolis, Topography and Thermal) + # Extra Terms (Coriolis, Topography) # -------------------------------------------------------------------- # # TODO: Is there a better way to store the Coriolis / topography fields? # The current approach is that these are prescribed fields, stored in # the equation, and initialised when the equation is if fexpr is not None: - V = FunctionSpace(domain.mesh, 'CG', 1) + V = FunctionSpace(self.domain.mesh, 'CG', 1) f = self.prescribed_fields('coriolis', V).interpolate(fexpr) coriolis_form = coriolis(subject( - prognostic(f*inner(domain.perp(u), w)*dx, "u"), self.X)) + prognostic(f*inner(self.domain.perp(u), w)*dx, "u"), self.X)) # Add linearisation if self.linearisation_map(coriolis_form.terms[0]): linear_coriolis = coriolis( - subject(prognostic(f*inner(domain.perp(u_trial), w)*dx, 'u'), self.X)) + subject(prognostic(f*inner(self.domain.perp(u_trial), w)*dx, 'u'), self.X)) coriolis_form = linearisation(coriolis_form, linear_coriolis) residual += coriolis_form - if bexpr is not None: - topography = self.prescribed_fields('topography', domain.spaces('DG')).interpolate(bexpr) - if self.thermal: - n = FacetNormal(domain.mesh) - topography_form = subject(prognostic - (-topography*div(b*w)*dx - + jump(b*w, n)*avg(topography)*dS, - 'u'), self.X) - else: - topography_form = subject(prognostic - (-g*div(w)*topography*dx, 'u'), - self.X) + if topog_expr is not None: + topography = self.prescribed_fields('topography', self.domain.spaces('DG')).interpolate(topog_expr) + topography_form = subject(prognostic + (-g*div(w)*topography*dx, 'u'), + self.X) residual += topography_form - # thermal source terms not involving topography. - # label these as the equivalent pressure gradient term - if self.thermal: - n = FacetNormal(domain.mesh) - source_form = pressure_gradient(subject(prognostic(-D*div(b*w)*dx - - 0.5*b*div(D*w)*dx - + jump(b*w, n)*avg(D)*dS - + 0.5*jump(D*w, n)*avg(b)*dS, - 'u'), self.X)) - residual += source_form - - # -------------------------------------------------------------------- # - # Linearise equations - # -------------------------------------------------------------------- # - # Add linearisations to equations - self.residual = self.generate_linear_terms(residual, self.linearisation_map) + self.residual = residual class LinearShallowWaterEquations(ShallowWaterEquations): u""" Class for the linear (rotating) shallow-water equations, which describe the velocity 'u' and the depth field 'D', solving some variant of: \n - ∂u/∂t + f×u + g*∇(D+b) = 0, \n + ∂u/∂t + f×u + g*∇(D+B) = 0, \n ∂D/∂t + H*∇.(u) = 0, \n - for mean depth 'H', Coriolis parameter 'f' and bottom surface 'b'. + for mean depth 'H', Coriolis parameter 'f' and bottom surface 'B'. - This is set up the from the underlying :class:`ShallowWaterEquations`, + This is set up from the underlying :class:`ShallowWaterEquations`, which is then linearised. """ - def __init__(self, domain, parameters, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, topog_expr=None, space_names=None, linearisation_map='default', u_transport_option="vector_invariant_form", - no_normal_flow_bc_ids=None, active_tracers=None, - thermal=False): + no_normal_flow_bc_ids=None, active_tracers=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -232,8 +221,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, the model's physical parameters. fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis parameter. Defaults to None. - bexpr (:class:`ufl.Expr`, optional): an expression for the bottom - surface of the fluid. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. space_names (dict, optional): a dictionary of strings for names of the function spaces to use for the spatial discretisation. The keys are the names of the prognostic variables. Defaults to None @@ -255,24 +244,345 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations. Defaults to None. - thermal (flag, optional): specifies whether the equations have a - thermal or buoyancy variable that feeds back on the momentum. - Defaults to False. """ + super().__init__(domain, parameters, + fexpr=fexpr, topog_expr=topog_expr, + space_names=space_names, + linearisation_map=linearisation_map, + u_transport_option=u_transport_option, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers) + + # Use the underlying routine to do a first linearisation of the equations + self.linearise_equation_set() + + +class ThermalShallowWaterEquations(ShallowWaterEquations): + u""" + Class for the (rotating) shallow-water equations, which evolve the velocity + 'u' and the depth field 'D' via some variant of either: \n + ∂u/∂t + (u.∇)u + f×u + b*∇(D+B) + 0.5*D*∇b= 0, \n + ∂D/∂t + ∇.(D*u) = 0, \n + ∂b/∂t + u.∇(b) = 0, \n + for Coriolis parameter 'f', bottom surface 'B' and buoyancy field b, + or, if equivalent_buoyancy=True: + ∂u/∂t + (u.∇)u + f×u + b_e*∇(D+B) + 0.5*D*∇(b_e + beta_2 q_v)= 0, \n + ∂D/∂t + ∇.(D*u) = 0, \n + ∂b_e/∂t + u.∇(b_e) = 0, \n + ∂q_t/∂t + u.∇(q_t) = 0, \n + for Coriolis parameter 'f', bottom surface 'B', equivalent buoyancy field \n + `b_e`=b-beta_2 q_v, and total moisture `q_t`=q_v+q_c, i.e. the sum of \n + water vapour and cloud water. + """ + + def __init__(self, domain, parameters, equivalent_buoyancy=False, + fexpr=None, topog_expr=None, + space_names=None, linearisation_map='default', + u_transport_option='vector_invariant_form', + no_normal_flow_bc_ids=None, active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + equivalent_buoyancy (bool, optional): switch to specify formulation + (see comments above). Defaults to False to give standard + thermal shallow water. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. + linearisation_map (func, optional): a function specifying which + terms in the equation set to linearise. If None is specified + then no terms are linearised. Defaults to the string 'default', + in which case the linearisation includes both time derivatives, + the 'D' transport term and the pressure gradient term. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + """ + + self.equivalent_buoyancy = equivalent_buoyancy + field_names = ['u', 'D'] + space_names = {'u': 'HDiv', 'D': 'L2'} + self.b_name = 'b_e' if equivalent_buoyancy else 'b' + + if equivalent_buoyancy: + for new_field in [self.b_name, 'q_t']: + field_names.append(new_field) + space_names[new_field] = 'L2' + else: + field_names.append(self.b_name) + space_names[self.b_name] = 'L2' + + if active_tracers is None: + active_tracers = [] + if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure gradient, - # Coriolis and transport term from depth equation + # Default linearisation is time derivatives, pressure + # gradient and transport terms from depth and buoyancy + # equations. Include q_t if equivalent buoyancy. Don't include + # active tracers. + linear_transported = ['D', self.b_name] + if equivalent_buoyancy: + linear_transported.append('q_t') linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis)) - or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport))) + t.get(prognostic) in field_names \ + and (any(t.has_label(time_derivative, pressure_gradient, + coriolis)) + or (t.get(prognostic) in linear_transported + and t.has_label(transport))) + + # Bypass ShallowWaterEquations.__init__ to avoid having to + # define the field_names separately + PrognosticEquationSet.__init__( + self, field_names, domain, space_names, + linearisation_map=linearisation_map, + no_normal_flow_bc_ids=no_normal_flow_bc_ids, + active_tracers=active_tracers) + + self.parameters = parameters + self.domain = domain + self.active_tracers = active_tracers + + self._setup_residual(fexpr, topog_expr, u_transport_option) + + # -------------------------------------------------------------------- # + # Linearise equations + # -------------------------------------------------------------------- # + # Add linearisations to equations + self.residual = self.generate_linear_terms( + self.residual, self.linearisation_map) + + def _setup_residual(self, fexpr, topog_expr, u_transport_option): + """ + Sets up the residual for the thermal shallow water + equations, first calling the shallow water equation + _setup_residual method to get the standard shallow water forms. + + Args: + fexpr (:class:`ufl.Expr`): an expression for the Coroilis + parameter. + topog_expr (:class:`ufl.Expr`): an expression for the + bottom surface of the fluid. + u_transport_option (str): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. + + """ + # don't pass topography to super class as we deal with those + # terms here later + super()._setup_residual(fexpr=fexpr, topog_expr=None, + u_transport_option=u_transport_option) + + w = self.tests[0] + gamma = self.tests[2] + u, D, b = split(self.X)[0:3] + Dbar, bbar = split(self.X_ref)[1:3] + u_trial, D_trial, b_trial = split(self.trials)[0:3] + n = FacetNormal(self.domain.mesh) + topog = self.prescribed_fields('topography', self.domain.spaces('DG')).interpolate(topog_expr) if topog_expr else None + self.topog = topog + if self.equivalent_buoyancy: + gamma_qt = self.tests[3] + qt = split(self.X)[3] + qtbar = split(self.X_ref)[3] + qt_trial = split(self.trials)[3] + + # -------------------------------------------------------------------- # + # Add pressure gradient-like terms to residual + # -------------------------------------------------------------------- # + # drop usual pressure gradient term + residual = self.residual.label_map( + lambda t: t.has_label(pressure_gradient), + drop) + + # add (moist) thermal source terms not involving topography - + # label these as the equivalent pressure gradient term and + # provide linearisation + if self.equivalent_buoyancy: + beta2 = self.parameters.beta2 + qsat_expr = self.compute_saturation(self.X) + qv = conditional(qt < qsat_expr, qt, qsat_expr) + qvbar = conditional(qtbar < qsat_expr, qtbar, qsat_expr) + source_form = pressure_gradient(subject(prognostic( + -D * div(b*w) * dx - 0.5 * b * div(D*w) * dx + + jump(b*w, n) * avg(D) * dS + 0.5 * jump(D*w, n) * avg(b) * dS + - beta2 * D * div(qv*w)*dx - 0.5 * beta2 * qv * div(D*w) * dx + + beta2 * jump(qv*w, n) * avg(D) * dS + + 0.5 * beta2 * jump(D*w, n) * avg(qv) * dS, + 'u'), self.X)) + linear_source_form = pressure_gradient(subject(prognostic( + -D_trial * div(bbar*w) * dx + - 0.5 * b_trial * div(Dbar*w) * dx + + jump(bbar*w, n) * avg(D_trial) * dS + + 0.5 * jump(Dbar*w, n) * avg(b_trial) * dS + - beta2 * D_trial * div(qvbar*w)*dx + - 0.5 * beta2 * qvbar * div(Dbar*w) * dx + + beta2 * jump(qvbar*w, n) * avg(D_trial) * dS + + 0.5 * beta2 * jump(Dbar*w, n) * avg(qvbar) * dS + - 0.5 * bbar * div(Dbar*w) * dx + + 0.5 * jump(Dbar*w, n) * avg(bbar) * dS + - 0.5 * bbar * div(D_trial*w) * dx + + 0.5 * jump(D_trial*w, n) * avg(bbar) * dS + - beta2 * 0.5 * qvbar * div(D_trial*w) * dx + + beta2 * 0.5 * jump(D_trial*w, n) * avg(qvbar) * dS + - beta2 * 0.5 * qt_trial * div(Dbar*w) * dx + + beta2 * 0.5 * jump(Dbar*w, n) * avg(qt_trial) * dS, + 'u'), self.X)) + else: + source_form = pressure_gradient( + subject(prognostic(-D * div(b*w) * dx + + jump(b*w, n) * avg(D) * dS + - 0.5 * b * div(D*w) * dx + + 0.5 * jump(D*w, n) * avg(b) * dS, + 'u'), self.X)) + linear_source_form = pressure_gradient( + subject(prognostic(-D_trial * div(bbar*w) * dx + + jump(bbar*w, n) * avg(D_trial) * dS + - 0.5 * b_trial * div(Dbar*w) * dx + + 0.5 * jump(Dbar*w, n) * avg(b_trial) * dS + - 0.5 * bbar * div(Dbar*w) * dx + + 0.5 * jump(Dbar*w, n) * avg(bbar) * dS + - 0.5 * bbar * div(D_trial*w) * dx + + 0.5 * jump(D_trial*w, n) * avg(bbar) * dS, + 'u'), self.X)) + source_form = linearisation(source_form, linear_source_form) + residual += source_form + + # -------------------------------------------------------------------- # + # add transport terms and their linearisations: + # -------------------------------------------------------------------- # + b_adv = prognostic(advection_form(gamma, b, u), self.b_name) + if self.linearisation_map(b_adv.terms[0]): + linear_b_adv = linear_advection_form(gamma, bbar, u_trial) + b_adv = linearisation(b_adv, linear_b_adv) + residual += subject(b_adv, self.X) + + if self.equivalent_buoyancy: + qt_adv = prognostic(advection_form(gamma_qt, qt, u), "q_t") + if self.linearisation_map(qt_adv.terms[0]): + linear_qt_adv = linear_advection_form(gamma_qt, qtbar, u_trial) + qt_adv = linearisation(qt_adv, linear_qt_adv) + residual += subject(qt_adv, self.X) + + # -------------------------------------------------------------------- # + # add topography terms: + # -------------------------------------------------------------------- # + if topog_expr is not None: + if self.equivalent_buoyancy: + topography_form = subject(prognostic( + - topog * div(b*w) * dx + + jump(b*w, n) * avg(topog) * dS + - beta2 * topog * div(qv*w) * dx + + beta2 * jump(qv*w, n) * avg(topog) * dS, + 'u'), self.X) + else: + topography_form = subject(prognostic( + - topog * div(b*w) * dx + + jump(b*w, n) * avg(topog) * dS, + 'u'), self.X) + residual += topography_form + + self.residual = residual + + def compute_saturation(self, X): + # Returns the saturation expression as a function of the + # parameters specified in self.parameters and the input + # functions X. The latter are left as inputs to the + # function so that it can also be used for initialisation + q0 = self.parameters.q0 + nu = self.parameters.nu + g = self.parameters.g + H = self.parameters.H + D, b = split(X)[1:3] + topog = self.topog + if topog is None: + sat_expr = q0*H/(D) * exp(nu*(1-b/g)) + else: + sat_expr = q0*H/(D+topog) * exp(nu*(1-b/g)) + return sat_expr + + +class LinearThermalShallowWaterEquations(ThermalShallowWaterEquations): + u""" + Class for the linear (rotating) thermal shallow-water equations, which + describe the velocity 'u' and depth field 'D', solving some variant of: \n + ∂u/∂t + f×u + bbar*∇D + 0.5*H*∇b = 0, \n + ∂D/∂t + H*∇.(u) = 0, \n + ∂b/∂t + u.∇bbar = 0, \n + ∂q_t/∂t + u.∇(q_tbar) = 0, \n + for mean depth 'H', mean buoyancy `bbar`, Coriolis parameter 'f' + + This is set up from the underlying :class:`ThermalShallowWaterEquations`, + which is then linearised. + """ + + def __init__(self, domain, parameters, equivalent_buoyancy=False, + fexpr=None, topog_expr=None, + space_names=None, linearisation_map='default', + u_transport_option="vector_invariant_form", + no_normal_flow_bc_ids=None, active_tracers=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. + equivalent_buoyancy (bool, optional): switch to specify formulation + (see comments above). Defaults to False to give standard + thermal shallow water. + fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis + parameter. Defaults to None. + topog_expr (:class:`ufl.Expr`, optional): an expression for the + bottom surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. + linearisation_map (func, optional): a function specifying which + terms in the equation set to linearise. If None is specified + then no terms are linearised. Defaults to the string 'default', + in which case the linearisation includes both time derivatives, + the 'D' transport term, pressure gradient and Coriolis terms. + u_transport_option (str, optional): specifies the transport term + used for the velocity equation. Supported options are: + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. + Defaults to 'vector_invariant_form'. + no_normal_flow_bc_ids (list, optional): a list of IDs of domain + boundaries at which no normal flow will be enforced. Defaults to + None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + """ super().__init__(domain, parameters, - fexpr=fexpr, bexpr=bexpr, space_names=space_names, + equivalent_buoyancy=equivalent_buoyancy, + fexpr=fexpr, topog_expr=topog_expr, + space_names=space_names, linearisation_map=linearisation_map, u_transport_option=u_transport_option, no_normal_flow_bc_ids=no_normal_flow_bc_ids, - active_tracers=active_tracers, thermal=thermal) + active_tracers=active_tracers) # Use the underlying routine to do a first linearisation of the equations self.linearise_equation_set() @@ -441,13 +751,6 @@ def __init__(self, domain, parameters, fexpr=None, in the equations. Defaults to None. """ - if linearisation_map == 'default': - # Default linearisation is time derivatives, pressure gradient, - # Coriolis and transport term from depth equation - linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis)) - or (t.get(prognostic) == 'D' and t.has_label(transport))) - super().__init__(domain, parameters, fexpr=fexpr, space_names=space_names, linearisation_map=linearisation_map, diff --git a/gusto/physics/shallow_water_microphysics.py b/gusto/physics/shallow_water_microphysics.py index 90cb5328d..bd3f8a02e 100644 --- a/gusto/physics/shallow_water_microphysics.py +++ b/gusto/physics/shallow_water_microphysics.py @@ -10,6 +10,7 @@ from gusto.physics.physics_parametrisation import PhysicsParametrisation from types import FunctionType + __all__ = ["InstantRain", "SWSaturationAdjustment"] diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 91361cc24..712e1c9e5 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -10,7 +10,8 @@ TestFunctions, TrialFunctions, TestFunction, TrialFunction, lhs, rhs, FacetNormal, div, dx, jump, avg, dS, dS_v, dS_h, ds_v, ds_t, ds_b, ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, cross, - BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector + BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector, + Interpolator, conditional ) from firedrake.fml import Term, drop from firedrake.petsc import flatten_parameters @@ -580,17 +581,18 @@ def solve(self, xrhs, dy): class ThermalSWSolver(TimesteppingSolver): - """ - Linear solver object for the thermal shallow water equations. + """Linear solver object for the thermal shallow water equations. - This solves a linear problem for the thermal shallow water equations with - prognostic variables u (velocity), D (depth) and b (buoyancy). It follows - the following strategy: + This solves a linear problem for the thermal shallow water + equations with prognostic variables u (velocity), D (depth) and + either b (buoyancy) or b_e (equivalent buoyancy). It follows the + following strategy: (1) Eliminate b (2) Solve the resulting system for (u, D) using a hybrid-mixed method (3) Reconstruct b - """ + + """ solver_parameters = { 'ksp_type': 'preonly', @@ -609,6 +611,8 @@ class ThermalSWSolver(TimesteppingSolver): @timed_function("Gusto:SolverSetup") def _setup_solver(self): equation = self.equations # just cutting down line length a bit + equivalent_buoyancy = equation.equivalent_buoyancy + dt = self.dt beta_u = dt*self.tau_values.get("u", self.alpha) beta_d = dt*self.tau_values.get("D", self.alpha) @@ -618,14 +622,12 @@ def _setup_solver(self): Vb = equation.domain.spaces("DG") # Check that the third field is buoyancy - if not equation.field_names[2] == 'b': - raise NotImplementedError("Field 'b' must exist to use the thermal linear solver in the SIQN scheme") + if not equation.field_names[2] == 'b' and not (equation.field_names[2] == 'b_e' and equivalent_buoyancy): + raise NotImplementedError("Field 'b' or 'b_e' must exist to use the thermal linear solver in the SIQN scheme") # Split up the rhs vector - self.xrhs = Function(self.equations.function_space) - u_in = split(self.xrhs)[0] - D_in = split(self.xrhs)[1] - b_in = split(self.xrhs)[2] + self.xrhs = Function(equation.function_space) + u_in, D_in, b_in = split(self.xrhs)[0:3] # Build the reduced function space for u, D M = MixedFunctionSpace((Vu, VD)) @@ -633,12 +635,27 @@ def _setup_solver(self): u, D = TrialFunctions(M) # Get background buoyancy and depth - Dbar = split(equation.X_ref)[1] - bbar = split(equation.X_ref)[2] + Dbar, bbar = split(equation.X_ref)[1:3] # Approximate elimination of b b = -dot(u, grad(bbar))*beta_b + b_in + if equivalent_buoyancy: + # compute q_v using q_sat to partition q_t into q_v and q_c + self.q_sat_func = Function(VD) + self.qvbar = Function(VD) + qtbar = split(equation.X_ref)[3] + + # set up interpolators that use the X_ref values for D and b_e + self.q_sat_expr_interpolator = Interpolator( + equation.compute_saturation(equation.X_ref), self.q_sat_func) + self.q_v_interpolator = Interpolator( + conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), + self.qvbar) + + # bbar was be_bar and here we correct to become bbar + bbar += equation.parameters.beta2 * self.qvbar + n = FacetNormal(equation.domain.mesh) eqn = ( @@ -694,6 +711,16 @@ def trace_nullsp(T): # Log residuals on hybridized solver self.log_ksp_residuals(self.uD_solver.snes.ksp) + @timed_function("Gusto:UpdateReferenceProfiles") + def update_reference_profiles(self): + """ + Updates the reference profiles. + """ + + if self.equations.equivalent_buoyancy: + self.q_sat_expr_interpolator.interpolate() + self.q_v_interpolator.interpolate() + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ diff --git a/integration-tests/data/sw_fplane_chkpt.h5 b/integration-tests/data/sw_fplane_chkpt.h5 index d09afda93..b752f0d70 100644 Binary files a/integration-tests/data/sw_fplane_chkpt.h5 and b/integration-tests/data/sw_fplane_chkpt.h5 differ diff --git a/integration-tests/equations/test_linear_sw_wave.py b/integration-tests/equations/test_linear_sw_wave.py index 81b6745e8..95eecdca4 100644 --- a/integration-tests/equations/test_linear_sw_wave.py +++ b/integration-tests/equations/test_linear_sw_wave.py @@ -11,7 +11,7 @@ def run_linear_sw_wave(tmpdir): - # Paramerers + # Parameters dt = 0.001 tmax = 30*dt H = 1 @@ -57,9 +57,11 @@ def run_linear_sw_wave(tmpdir): u0 = stepper.fields("u") D0 = stepper.fields("D") + Dbar = eqns.X_ref.subfunctions[1] u0.project(uexpr) D0.interpolate(Dexpr) + Dbar.interpolate(H) # --------------------------------------------------------------------- # # Run diff --git a/integration-tests/equations/test_thermal_sw.py b/integration-tests/equations/test_thermal_sw.py index 8328c308d..f201e78c6 100644 --- a/integration-tests/equations/test_thermal_sw.py +++ b/integration-tests/equations/test_thermal_sw.py @@ -40,9 +40,9 @@ def setup_sw(dirname, dt, u_transport_option): parameters = ShallowWaterParameters(H=H, g=g) Omega = parameters.Omega fexpr = 2*Omega*x[2]/R - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, - u_transport_option=u_transport_option, - thermal=True) + eqns = ThermalShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option=u_transport_option) + # I/O diagnostic_fields = [SteadyStateError('D'), SteadyStateError('u'), diff --git a/integration-tests/physics/test_sw_saturation_adjustment.py b/integration-tests/physics/test_sw_saturation_adjustment.py index 66c971519..e8c870253 100644 --- a/integration-tests/physics/test_sw_saturation_adjustment.py +++ b/integration-tests/physics/test_sw_saturation_adjustment.py @@ -49,9 +49,10 @@ def run_sw_cond_evap(dirname, process): tracers = [WaterVapour(space='DG'), CloudWater(space='DG')] - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, - u_transport_option='vector_advection_form', - thermal=True, active_tracers=tracers) + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, + u_transport_option='vector_advection_form', + active_tracers=tracers) # I/O output = OutputParameters(dirname=dirname+"/sw_cond_evap", diff --git a/integration-tests/physics/test_terminator_toy.py b/integration-tests/physics/test_terminator_toy.py index 1ae7b69df..ddd2783c4 100644 --- a/integration-tests/physics/test_terminator_toy.py +++ b/integration-tests/physics/test_terminator_toy.py @@ -32,15 +32,15 @@ def run_terminator_toy(dirname): domain = Domain(mesh, dt, 'BDM', 1) # Define the interacting species - X = ActiveTracer(name='X', space='DG', + Y = ActiveTracer(name='Y', space='DG', variable_type=TracerVariableType.mixing_ratio, transport_eqn=TransportEquationType.advective) - X2 = ActiveTracer(name='X2', space='DG', + Y2 = ActiveTracer(name='Y2', space='DG', variable_type=TracerVariableType.mixing_ratio, transport_eqn=TransportEquationType.advective) - tracers = [X, X2] + tracers = [Y, Y2] # Equation V = domain.spaces("HDiv") @@ -57,11 +57,11 @@ def run_terminator_toy(dirname): k1 = max_value(0, sin(theta)*sin(theta_c) + cos(theta)*cos(theta_c)*cos(lamda-lamda_c)) k2 = 1 - physics_schemes = [(TerminatorToy(eqn, k1=k1, k2=k2, species1_name='X', - species2_name='X2'), BackwardEuler(domain))] + physics_schemes = [(TerminatorToy(eqn, k1=k1, k2=k2, species1_name='Y', + species2_name='Y2'), BackwardEuler(domain))] transport_scheme = SSPRK3(domain) - transport_method = [DGUpwind(eqn, 'X'), DGUpwind(eqn, 'X2')] + transport_method = [DGUpwind(eqn, 'Y'), DGUpwind(eqn, 'Y2')] time_varying_velocity = True stepper = SplitPrescribedTransport( @@ -75,39 +75,39 @@ def u_t(t): stepper.setup_prescribed_expr(u_t) - X_T_0 = 4e-6 - X_0 = X_T_0 + 0*lamda - X2_0 = 0*lamda + Y_T_0 = 4e-6 + Y_0 = Y_T_0 + 0*lamda + Y2_0 = 0*lamda - stepper.fields("X").interpolate(X_0) - stepper.fields("X2").interpolate(X2_0) + stepper.fields("Y").interpolate(Y_0) + stepper.fields("Y2").interpolate(Y2_0) stepper.run(t=0, tmax=10*dt) # Compute the steady state solution to compare to steady_space = domain.spaces('DG') - X_steady = Function(steady_space) - X2_steady = Function(steady_space) + Y_steady = Function(steady_space) + Y2_steady = Function(steady_space) r = k1/(4*k2) - D_val = sqrt(r**2 + 2*X_T_0*r) + D_val = sqrt(r**2 + 2*Y_T_0*r) - X_steady.interpolate(D_val - r) - X2_steady.interpolate(0.5*(X_T_0 - D_val + r)) + Y_steady.interpolate(D_val - r) + Y2_steady.interpolate(0.5*(Y_T_0 - D_val + r)) - return stepper, X_steady, X2_steady + return stepper, Y_steady, Y2_steady def test_terminator_toy_setup(tmpdir): dirname = str(tmpdir) - stepper, X_steady, X2_steady = run_terminator_toy(dirname) - X_field = stepper.fields("X") - X2_field = stepper.fields("X2") + stepper, Y_steady, Y2_steady = run_terminator_toy(dirname) + Y_field = stepper.fields("Y") + Y2_field = stepper.fields("Y2") - print(errornorm(X_field, X_steady)/norm(X_steady)) - print(errornorm(X2_field, X2_steady)/norm(X2_steady)) + print(errornorm(Y_field, Y_steady)/norm(Y_steady)) + print(errornorm(Y2_field, Y2_steady)/norm(Y2_steady)) # Assert that the physics scheme has sufficiently moved # the species fields near their steady state solutions - assert errornorm(X_field, X_steady)/norm(X_steady) < 0.4, "The X field is not sufficiently close to the steady state profile" - assert errornorm(X2_field, X2_steady)/norm(X2_steady) < 0.4, "The X2 field is not sufficiently close to the steady state profile" + assert errornorm(Y_field, Y_steady)/norm(Y_steady) < 0.4, "The Y field is not sufficiently close to the steady state profile" + assert errornorm(Y2_field, Y2_steady)/norm(Y2_steady) < 0.4, "The Y2 field is not sufficiently close to the steady state profile" diff --git a/integration-tests/rexi/test_linear_sw.py b/integration-tests/rexi/test_linear_sw.py index c7a29aab9..77d5eafe2 100644 --- a/integration-tests/rexi/test_linear_sw.py +++ b/integration-tests/rexi/test_linear_sw.py @@ -57,9 +57,11 @@ def run_rexi_sw(tmpdir, ensemble=None): U_in = Function(eqns.function_space, name="U_in") Uexpl = Function(eqns.function_space, name="Uexpl") u, D = U_in.subfunctions + Dbar = eqns.X_ref.subfunctions[1] u.project(uexpr) D.interpolate(Dexpr) + Dbar.interpolate(H) if write_output: rexi_output.write(u, D) @@ -70,6 +72,7 @@ def run_rexi_sw(tmpdir, ensemble=None): uexpl, Dexpl = Uexpl.subfunctions u.assign(uexpl) D.assign(Dexpl) + if write_output: rexi_output.write(u, D) diff --git a/integration-tests/transport/test_zero_limiter.py b/integration-tests/transport/test_zero_limiter.py index b4f2086a3..75f0f4616 100644 --- a/integration-tests/transport/test_zero_limiter.py +++ b/integration-tests/transport/test_zero_limiter.py @@ -52,9 +52,10 @@ def setup_zero_limiter(dirname, limiter=False, rain=False): else: tracers = [WaterVapour(space='DG'), CloudWater(space='DG')] - eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, - u_transport_option='vector_advection_form', - thermal=True, active_tracers=tracers) + eqns = ThermalShallowWaterEquations( + domain, parameters, fexpr=fexpr, + u_transport_option='vector_advection_form', + active_tracers=tracers) output = OutputParameters(dirname=dirname, dumpfreq=1) diff --git a/plotting/shallow_water/plot_linear_thermal_galewsky_jet.py b/plotting/shallow_water/plot_linear_thermal_galewsky_jet.py new file mode 100644 index 000000000..de4d446a9 --- /dev/null +++ b/plotting/shallow_water/plot_linear_thermal_galewsky_jet.py @@ -0,0 +1,173 @@ +""" +Plots the linear thermal Galewsky jet test case. + """ +from os.path import abspath, dirname +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file + +results_file_name = f'{abspath(dirname(__file__))}/../../results/linear_thermal_galewsky/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/shallow_water/linear_thermal_galewsky' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D', 'RelativeVorticity', 'b'] +init_colour_schemes = ['Oranges', 'YlGnBu', 'RdBu_r', 'PuRd_r'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)', + r'$\zeta$ (s$^{-1})$', r'$b$ (m s$^{-2}$)'] +init_contours_to_remove = [None, None, 0.0, None] +init_contours = [np.linspace(0.0, 80.0, 9), + np.linspace(8900.0, 10200.0+1e-3, 12), + np.linspace(-2e-4, 2e-4, 17), + np.linspace(8.8, 9.8, 11)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'b'] +final_colour_schemes = ['RdBu_r', 'PuRd_r'] +final_field_labels = [r'$\zeta$ (s$^{-1}$)', r'$b$ (m s$^{-2}$)'] +final_contours_to_remove = [0.0, None] +final_contours = [np.linspace(-2e-4, 2e-4, 17), + np.linspace(8.8, 9.8, 11)] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +xlims = [-180, 180] +ylims = [10, 80] + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(2, 2, figsize=(16, 12), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray.flatten(), init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + lon_1d = np.linspace(-180.0, 180.0, 91) + lat_1d = np.linspace(-90.0, 90.0, 81) + lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=(12, 1), magnitude_filter=1.0, + ) + + # Labels ------------------------------------------------------------------- + if i in [0, 2]: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + if i in [2, 3]: + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(1, 2, figsize=(16, 8), sharex='all', sharey='all') +time_idx = -1 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray, final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove, final_contours)): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_moist_thermal_equivb_gravity_wave.py b/plotting/shallow_water/plot_moist_thermal_equivb_gravity_wave.py new file mode 100644 index 000000000..b771cc89b --- /dev/null +++ b/plotting/shallow_water/plot_moist_thermal_equivb_gravity_wave.py @@ -0,0 +1,187 @@ +""" +Plots the moist thermal gravity wave test case. +""" +from os.path import abspath, dirname +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file + +results_file_name = f'{abspath(dirname(__file__))}/../../results/moist_thermal_equivb_gw/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/shallow_water/moist_thermal_equivb_gw' + +beta2 = 9.80616*10 + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D', 'PartitionedCloud', 'b_e'] +init_colour_schemes = ['Oranges', 'YlGnBu', 'cividis', 'PuRd_r'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)', + r'$q_{cl}$ (kg kg$^{-1})$', r'$b_e$ (m s$^{-2}$)'] +init_contours_to_remove = [None, None, None, None] +init_contours = [np.linspace(0.0, 20.0, 9), + np.linspace(4800.0, 8000.0, 19), + np.linspace(0.01, 0.02, 11), + np.linspace(9, 10, 15)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['PartitionedCloud', 'b_e'] +final_colour_schemes = ['cividis', 'PuRd_r'] +final_field_labels = [r'$q_{cl}$ (kg kg$^{-1}$)', r'$b_e$ (m s$^{-2}$)'] +final_contours_to_remove = [None, None] +final_contours = [np.linspace(0.01, 0.02, 11), + np.linspace(9, 10, 15)] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +xlims = [-180, 180] +ylims = [-90, 90] +minmax_format = { + 'PartitionedCloud': '.1e', + 'b_e': '.2f', + 'u': '.1f', + 'D': '.0f' +} + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(2, 2, figsize=(16, 12), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray.flatten(), init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, + minmax_format=minmax_format[field_name] + ) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + lon_1d = np.linspace(-180.0, 180.0, 91) + lat_1d = np.linspace(-90.0, 90.0, 81) + lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + ) + + # Labels ------------------------------------------------------------------- + if i in [0, 2]: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + if i in [2, 3]: + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(1, 2, figsize=(16, 8), sharex='all', sharey='all') +time_idx = -1 + +for i, (ax, field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + axarray, final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove, final_contours)): + + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + + time = data_file['time'][time_idx] / (24.*60.*60.) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, + minmax_format=minmax_format[field_name] + ) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close()