diff --git a/examples/shallow_water/moist_convective_williamson2.py b/examples/shallow_water/moist_convective_williamson2.py new file mode 100644 index 000000000..4c21ebb75 --- /dev/null +++ b/examples/shallow_water/moist_convective_williamson2.py @@ -0,0 +1,163 @@ +""" +A moist convective version of the Williamson 2 shallow water test (steady state +geostrophically-balanced flow). The saturation function depends on height, +with a constant background buoyancy/temperature field. +Vapour is initialised very close to saturation and small overshoots will +generate clouds. +""" +from gusto import * +from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, sin, cos, exp) +import sys + +# ----------------------------------------------------------------- # +# Test case parameters +# ----------------------------------------------------------------- # + +dt = 120 + +if '--running-tests' in sys.argv: + tmax = dt + dumpfreq = 1 +else: + day = 24*60*60 + tmax = 5*day + ndumps = 5 + dumpfreq = int(tmax / (ndumps*dt)) + +R = 6371220. +u_max = 20 +phi_0 = 3e4 +epsilon = 1/300 +theta_0 = epsilon*phi_0**2 +g = 9.80616 +H = phi_0/g +xi = 0 +q0 = 200 +beta1 = 110 +alpha = 16 +gamma_v = 0.98 +qprecip = 1e-4 +gamma_r = 1e-3 + +# ----------------------------------------------------------------- # +# Set up model objects +# ----------------------------------------------------------------- # + +# Domain +mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2) +degree = 1 +domain = Domain(mesh, dt, 'BDM', degree) +x = SpatialCoordinate(mesh) + +# Equations +parameters = ShallowWaterParameters(H=H, g=g) +Omega = parameters.Omega +fexpr = 2*Omega*x[2]/R + +tracers = [WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')] + +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option='vector_advection_form', + active_tracers=tracers) + +# IO +dirname = "moist_convective_williamson2" +output = OutputParameters(dirname=dirname, + dumpfreq=dumpfreq, + dumplist_latlon=['D', 'D_error'], + dump_nc=True, + dump_vtus=True) + +diagnostic_fields = [CourantNumber(), RelativeVorticity(), + PotentialVorticity(), + ShallowWaterKineticEnergy(), + ShallowWaterPotentialEnergy(parameters), + ShallowWaterPotentialEnstrophy(), + SteadyStateError('u'), SteadyStateError('D'), + SteadyStateError('water_vapour'), + SteadyStateError('cloud_water')] + +io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + +# define saturation function +def sat_func(x_in): + h = x_in.subfunctions[1] + lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) + numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma)) + denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2 + theta = numerator/denominator + return q0/(g*h) * exp(20*(theta)) + + +transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names] + +limiter = DG1Limiter(domain.spaces('DG')) + +transported_fields = [TrapeziumRule(domain, "u"), + SSPRK3(domain, "D"), + SSPRK3(domain, "water_vapour", limiter=limiter), + SSPRK3(domain, "cloud_water", limiter=limiter), + SSPRK3(domain, "rain", limiter=limiter) + ] + +linear_solver = MoistConvectiveSWSolver(eqns) + +sat_adj = SWSaturationAdjustment(eqns, sat_func, + time_varying_saturation=True, + convective_feedback=True, beta1=beta1, + gamma_v=gamma_v, time_varying_gamma_v=False, + parameters=parameters) + +inst_rain = InstantRain(eqns, qprecip, vapour_name="cloud_water", + rain_name="rain", gamma_r=gamma_r) + +physics_schemes = [(sat_adj, ForwardEuler(domain)), + (inst_rain, ForwardEuler(domain))] + +stepper = SemiImplicitQuasiNewton(eqns, io, + transport_schemes=transported_fields, + spatial_methods=transport_methods, + linear_solver=linear_solver, + physics_schemes=physics_schemes) + +# ----------------------------------------------------------------- # +# Initial conditions +# ----------------------------------------------------------------- # + +u0 = stepper.fields("u") +D0 = stepper.fields("D") +v0 = stepper.fields("water_vapour") + +lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) + +uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x) +g = parameters.g +w = Omega*R*u_max + (u_max**2)/2 +sigma = 0 + +Dexpr = H - (1/g)*(w)*((sin(phi))**2) +D_for_v = H - (1/g)*(w + sigma)*((sin(phi))**2) + +# though this set-up has no buoyancy, we use the expression for theta to set up +# the initial vapour +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 + +initial_msat = q0/(g*Dexpr) * exp(20*theta) +vexpr = (1 - xi) * initial_msat + +u0.project(uexpr) +D0.interpolate(Dexpr) +v0.interpolate(vexpr) + +# Set reference profiles +Dbar = Function(D0.function_space()).assign(H) +stepper.set_reference_profiles([('D', Dbar)]) + +# ----------------------------------------------------------------- # +# Run +# ----------------------------------------------------------------- # + +stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py new file mode 100644 index 000000000..21162a835 --- /dev/null +++ b/examples/shallow_water/moist_thermal_williamson5.py @@ -0,0 +1,146 @@ +""" +Moist flow over a mountain test case from Zerroukat and Allen (2015). Similar +to the Williamson et al test 5 but with additional thermodynamic equations. +""" +from gusto import * +from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, + as_vector, pi, sqrt, min_value, exp, cos, sin) +import sys + +# ----------------------------------------------------------------- # +# Test case parameters +# ----------------------------------------------------------------- # + +dt = 300 + +if '--running-tests' in sys.argv: + tmax = dt + dumpfreq = 1 +else: + day = 24*60*60 + tmax = 50*day + ndumps = 50 + dumpfreq = int(tmax / (ndumps*dt)) + +R = 6371220. +H = 5960. +u_max = 20. +# moist shallow water parameters +epsilon = 1/300 +SP = -40*epsilon +EQ = 30*epsilon +NP = -20*epsilon +mu1 = 0.05 +mu2 = 0.98 +q0 = 135 # chosen to give an initial max vapour of approx 0.02 +beta2 = 10 +qprecip = 1e-4 +gamma_r = 1e-3 +# topography parameters +R0 = pi/9. +R0sq = R0**2 +lamda_c = -pi/2. +phi_c = pi/6. + +# ----------------------------------------------------------------- # +# Set up model objects +# ----------------------------------------------------------------- # + +# Domain +mesh = IcosahedralSphereMesh(radius=R, + refinement_level=4, degree=1) +degree = 1 +domain = Domain(mesh, dt, "BDM", degree) +x = SpatialCoordinate(mesh) + +# Equation +parameters = ShallowWaterParameters(H=H) +Omega = parameters.Omega +fexpr = 2*Omega*x[2]/R + +# Topography +lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) +lsq = (lamda - lamda_c)**2 +thsq = (phi - phi_c)**2 +rsq = min_value(R0sq, lsq+thsq) +r = sqrt(rsq) +tpexpr = 2000 * (1 - r/R0) + +tracers = [WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')] +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=tpexpr, + thermal=True, + active_tracers=tracers) + +# I/O +dirname = "moist_thermal_williamson5" +output = OutputParameters( + dirname=dirname, + dumplist_latlon=['D'], + dumpfreq=dumpfreq, +) +diagnostic_fields = [Sum('D', 'topography'), CourantNumber()] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + +# Saturation function +def sat_func(x_in): + _, h, b = x_in.subfunctions + return (q0/(g*h + g*tpexpr)) * exp(20*(1 - b/g)) + + +# Feedback proportionality is dependent on h and b +def gamma_v(x_in): + _, h, b = x_in.subfunctions + return (1 + beta2*(20*q0/(g*h + g*tpexpr) * exp(20*(1 - b/g))))**(-1) + + +SWSaturationAdjustment(eqns, sat_func, time_varying_saturation=True, + parameters=parameters, thermal_feedback=True, + beta2=beta2, gamma_v=gamma_v, + time_varying_gamma_v=True) + +InstantRain(eqns, qprecip, vapour_name="cloud_water", rain_name="rain", + gamma_r=gamma_r) + +transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names] + +# Timestepper +stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods) + +# ----------------------------------------------------------------- # +# Initial conditions +# ----------------------------------------------------------------- # + +u0 = stepper.fields("u") +D0 = stepper.fields("D") +b0 = stepper.fields("b") +v0 = stepper.fields("water_vapour") +c0 = stepper.fields("cloud_water") +r0 = stepper.fields("rain") + +uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) + +g = parameters.g +Rsq = R**2 +Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*x[2]**2/Rsq)/g - tpexpr + +# Expression for initial buoyancy - note the bracket around 1-mu +F = (2/(pi**2))*(phi*(phi-pi/2)*SP - 2*(phi+pi/2)*(phi-pi/2)*(1-mu1)*EQ + phi*(phi+pi/2)*NP) +theta_expr = F + mu1*EQ*cos(phi)*sin(lamda) +bexpr = g * (1 - theta_expr) + +# Expression for initial vapour depends on initial saturation +initial_msat = q0/(g*D0 + g*tpexpr) * exp(20*theta_expr) +vexpr = mu2 * initial_msat + +# Initialise (cloud and rain initially zero) +u0.project(uexpr) +D0.interpolate(Dexpr) +b0.interpolate(bexpr) +v0.interpolate(vexpr) + +# ----------------------------------------------------------------- # +# Run +# ----------------------------------------------------------------- # + +stepper.run(t=0, tmax=tmax) diff --git a/gusto/diagnostics/diagnostics.py b/gusto/diagnostics/diagnostics.py index b287d6d36..d5d24466b 100644 --- a/gusto/diagnostics/diagnostics.py +++ b/gusto/diagnostics/diagnostics.py @@ -6,9 +6,10 @@ LinearVariationalProblem, LinearVariationalSolver, ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, pi, TensorFunctionSpace, SpatialCoordinate, as_vector, - Projector, Interpolator, FunctionSpace, FiniteElement, + Projector, FunctionSpace, FiniteElement, TensorProductElement) from firedrake.assign import Assigner +from firedrake.__future__ import Interpolator from ufl.domain import extract_unique_domain from abc import ABCMeta, abstractmethod, abstractproperty diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index 84b187e47..58a9902ce 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/gusto/equations/compressible_euler_equations.py @@ -40,7 +40,7 @@ def __init__(self, domain, parameters, sponge_options=None, u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, - active_tracers=None): + active_tracers=None, max_quad_deg=5): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -75,6 +75,8 @@ def __init__(self, domain, parameters, sponge_options=None, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations.. Defaults to None. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. Raises: NotImplementedError: only mixing ratio tracers are implemented. @@ -98,7 +100,8 @@ def __init__(self, domain, parameters, sponge_options=None, super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, - active_tracers=active_tracers) + active_tracers=active_tracers, + max_quad_deg=max_quad_deg) self.parameters = parameters g = parameters.g @@ -284,7 +287,8 @@ def __init__(self, domain, parameters, sponge_options=None, u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, - active_tracers=None): + active_tracers=None, + max_quad_deg=5): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -318,7 +322,9 @@ def __init__(self, domain, parameters, sponge_options=None, None. active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included - in the equations.. Defaults to None. + in the equations. Defaults to None. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. Raises: NotImplementedError: only mixing ratio tracers are implemented. @@ -330,7 +336,8 @@ def __init__(self, domain, parameters, sponge_options=None, u_transport_option=u_transport_option, diffusion_options=diffusion_options, no_normal_flow_bc_ids=no_normal_flow_bc_ids, - active_tracers=active_tracers) + active_tracers=active_tracers, + max_quad_deg=max_quad_deg) # Replace self.residual = self.residual.label_map( diff --git a/gusto/equations/prognostic_equations.py b/gusto/equations/prognostic_equations.py index b2df68e2c..3c86cd28e 100644 --- a/gusto/equations/prognostic_equations.py +++ b/gusto/equations/prognostic_equations.py @@ -25,7 +25,7 @@ class PrognosticEquation(object, metaclass=ABCMeta): """Base class for prognostic equations.""" - def __init__(self, domain, function_space, field_name): + def __init__(self, domain, function_space, field_name, max_quad_deg=5): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -33,12 +33,15 @@ def __init__(self, domain, function_space, field_name): function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. """ self.domain = domain self.function_space = function_space self.X = Function(function_space) self.field_name = field_name + self.max_quad_deg = max_quad_deg self.bcs = {} self.prescribed_fields = PrescribedFields() @@ -77,7 +80,7 @@ class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): def __init__(self, field_names, domain, space_names, linearisation_map=None, no_normal_flow_bc_ids=None, - active_tracers=None): + active_tracers=None, max_quad_deg=5): """ Args: field_names (list): a list of strings for names of the prognostic @@ -95,12 +98,15 @@ def __init__(self, field_names, domain, space_names, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations.. Defaults to None. + max_quad_deg (int): maximum quadrature degree for any form. Defaults + to 5. """ self.field_names = field_names self.space_names = space_names self.active_tracers = active_tracers self.linearisation_map = lambda t: False if linearisation_map is None else linearisation_map(t) + self.max_quad_deg = max_quad_deg # Build finite element spaces self.spaces = [domain.spaces(space_name) for space_name in diff --git a/gusto/timestepping/timestepper.py b/gusto/timestepping/timestepper.py index 3c619753b..5c906429b 100644 --- a/gusto/timestepping/timestepper.py +++ b/gusto/timestepping/timestepper.py @@ -2,7 +2,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import Function, Projector, split -from firedrake.fml import drop, Term, LabelledForm +from firedrake.fml import drop, Term, all_terms, LabelledForm from pyop2.profiling import timed_stage from gusto.equations import PrognosticEquationSet from gusto.core import TimeLevelFields, StateFields @@ -13,6 +13,7 @@ from gusto.spatial_methods.transport_methods import TransportMethod import ufl + __all__ = ["BaseTimestepper", "Timestepper", "PrescribedTransport"] @@ -110,6 +111,39 @@ def setup_equation(self, equation): for method in self.spatial_methods: method.replace_form(equation) + # -------------------------------------------------------------------- # + # Ensure the quadrature degree is not excessive for any integrals + # -------------------------------------------------------------------- # + def update_quadrature_degree(t): + # This function takes in a term and returns a new term + # with the same form and labels, the only difference being + # that any integrals in the form with no quadrature_degree + # set will have their quadrature degree set to max_quad_deg + + # This list will hold the updated integrals + new_itgs = [] + + # Loop over integrals in this term's form + for itg in t.form._integrals: + # check if the quadrature degree is not set + if 'quadrature_degree' not in itg._metadata.keys(): + # create new integral with updated quadrature degree + new_itg = itg.reconstruct( + metadata={'quadrature_degree': equation.max_quad_deg}) + new_itgs.append(new_itg) + else: + # if quadrature degree was already set, just keep + # this integral + new_itgs.append(itg) + + new_form = ufl.form.Form(new_itgs) + + return Term(new_form, t.labels) + + equation.residual = equation.residual.label_map( + all_terms, + lambda t: update_quadrature_degree(t)) + def setup_transporting_velocity(self, scheme): """ Set up the time discretisation by replacing the transporting velocity