diff --git a/examples/compressible/skamarock_klemp_hydrostatic.py b/examples/compressible/skamarock_klemp_hydrostatic.py index 90f2bd700..36b500540 100644 --- a/examples/compressible/skamarock_klemp_hydrostatic.py +++ b/examples/compressible/skamarock_klemp_hydrostatic.py @@ -17,7 +17,7 @@ dt = 25. if '--running-tests' in sys.argv: nlayers = 5 # horizontal layers - columns = 50 # number of columns + columns = 10 # number of columns tmax = dt dumpfreq = 1 else: diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 0a3790147..bb9de1d94 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -183,7 +183,7 @@ def advection_equation_circulation_form(domain, test, q, ubar): def diffusion_form(test, q, kappa): u""" - The diffusion form, ∇.(κ∇q) for diffusivity κ and variable q. + The diffusion form, -∇.(κ∇q) for diffusivity κ and variable q. Args: test (:class:`TestFunction`): the test function. @@ -191,6 +191,6 @@ def diffusion_form(test, q, kappa): kappa: (:class:`ufl.Expr`): the diffusivity value. """ - form = inner(test, div(kappa*grad(q)))*dx + form = -inner(test, div(kappa*grad(q)))*dx return diffusion(form) diff --git a/gusto/configuration.py b/gusto/configuration.py index 039b32740..3694ae9ea 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -8,7 +8,7 @@ "IntegrateByParts", "TransportEquationType", "OutputParameters", "CompressibleParameters", "ShallowWaterParameters", "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", - "SpongeLayerParameters", "DiffusionParameters" + "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters" ] @@ -185,3 +185,18 @@ class DiffusionParameters(Configuration): kappa = None mu = None + + +class BoundaryLayerParameters(Configuration): + """ + Parameters for the idealised wind drag, surface flux and boundary layer + mixing schemes. + """ + + coeff_drag_0 = 7e-4 # Zeroth drag coefficient (dimensionless) + coeff_drag_1 = 6.5e-5 # First drag coefficient (s/m) + coeff_drag_2 = 2e-3 # Second drag coefficient (dimensionless) + coeff_heat = 1.1e-3 # Dimensionless surface sensible heat coefficient + coeff_evap = 1.1e-3 # Dimensionless surface evaporation coefficient + height_surface_layer = 75. # Height (m) of surface level (usually lowest level) + mu = 100. # Interior penalty coefficient for vertical diffusion diff --git a/gusto/coordinates.py b/gusto/coordinates.py index 65d115f1c..571effbe8 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -7,6 +7,7 @@ from gusto.logging import logger from firedrake import SpatialCoordinate, Function import numpy as np +import pandas as pd class Coordinates(object): @@ -154,3 +155,124 @@ def register_space(self, domain, space_name): new_coords = comm.recv(source=procid, tag=my_tag) (low_lim, up_lim) = self.parallel_array_lims[space_name][procid][:] self.global_chi_coords[space_name][i, low_lim:up_lim+1] = new_coords + + def get_column_data(self, field, domain): + """ + Reshapes a field's data into columns. + + Args: + field (:class:`Function`): the field whose data needs sorting. + domain (:class:`Domain`): the domain used to register coordinates + if this hasn't already been done. + + Returns: + tuple of :class:`numpy.ndarray`: a 2D array of data, arranged in + columns, and the data pairing the indices of the data with the + ordered column data. + """ + + space_name = field.function_space().name + if space_name not in self.chi_coords.keys(): + self.register_space(domain, space_name) + coords = self.chi_coords[space_name] + + data_is_3d = (len(coords) == 3) + coords_X = coords[0].dat.data + coords_Y = coords[1].dat.data if data_is_3d else None + coords_Z = coords[-1].dat.data + + # ------------------------------------------------------------------------ # + # Round data to ensure sorting in dataframe is OK + # ------------------------------------------------------------------------ # + + # Work out digits to round to, based on number of points and range of coords + num_points = np.size(coords_X) + data_range = np.max(coords_X) - np.min(coords_X) + if data_range > np.finfo(type(data_range)).tiny: + digits = int(np.floor(-np.log10(data_range / num_points)) + 3) + coords_X = coords_X.round(digits) + + if data_is_3d: + data_range = np.max(coords_Y) - np.min(coords_Y) + if data_range > np.finfo(type(data_range)).tiny: + # Only round if there is already some range + digits = int(np.floor(-np.log10(data_range / num_points)) + 3) + coords_Y = coords_Y.round(digits) + + # -------------------------------------------------------------------- # + # Make data frame + # -------------------------------------------------------------------- # + + data_dict = {'field': field.dat.data, 'X': coords_X, 'Z': coords_Z, + 'index': range(len(field.dat.data))} + if coords_Y is not None: + data_dict['Y'] = coords_Y + + # Put everything into a pandas dataframe + data = pd.DataFrame(data_dict) + + # Sort array by X and Y coordinates + if data_is_3d: + data = data.sort_values(by=['X', 'Y', 'Z']) + first_X, first_Y = data['X'].values[0], data['Y'].values[0] + first_point = data[(np.isclose(data['X'], first_X)) + & (np.isclose(data['Y'], first_Y))] + + else: + data = data.sort_values(by=['X', 'Z']) + first_X = data['X'].values[0] + first_point = data[np.isclose(data['X'], first_X)] + + # Number of levels should correspond to the number of points with the first + # coordinate values + num_levels = len(first_point) + assert len(data) % num_levels == 0, 'Unable to nicely divide data into levels' + + # -------------------------------------------------------------------- # + # Create new arrays to store structured data + # -------------------------------------------------------------------- # + + num_hori_points = int(len(data) / num_levels) + field_data = np.zeros((num_hori_points, num_levels)) + coords_X = np.zeros((num_hori_points, num_levels)) + if data_is_3d: + coords_Y = np.zeros((num_hori_points, num_levels)) + coords_Z = np.zeros((num_hori_points, num_levels)) + index_data = np.zeros((num_hori_points, num_levels), dtype=int) + + # -------------------------------------------------------------------- # + # Fill arrays, on the basis of the dataframe already being sorted + # -------------------------------------------------------------------- # + + for lev_idx in range(num_levels): + data_slice = slice(lev_idx, num_hori_points*num_levels+lev_idx, num_levels) + field_data[:, lev_idx] = data['field'].values[data_slice] + coords_X[:, lev_idx] = data['X'].values[data_slice] + if data_is_3d: + coords_Y[:, lev_idx] = data['Y'].values[data_slice] + coords_Z[:, lev_idx] = data['Z'].values[data_slice] + index_data[:, lev_idx] = data['index'].values[data_slice] + + return field_data, index_data + + def set_field_from_column_data(self, field, columnwise_data, index_data): + """ + Fills in field data from some columnwise data. + + Args: + field (:class:`Function`): the field whose data shall be filled. + columnwise_data (:class:`numpy.ndarray`): the field data arranged + into columns, to be written into the field. + index_data (:class:`numpy.ndarray`): the indices of the original + field data, arranged like the columnwise data. + + Returns: + :class:`Function`: the updated field. + """ + + _, num_levels = np.shape(columnwise_data) + + for lev_idx in range(num_levels): + field.dat.data[index_data[:, lev_idx]] = columnwise_data[:, lev_idx] + + return field diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index d6e1be4cc..024a42fa9 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -14,6 +14,7 @@ from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations from gusto.active_tracers import TracerVariableType, Phases +from gusto.logging import logger import numpy as np __all__ = ["Diagnostics", "CourantNumber", "Gradient", "XComponent", "YComponent", @@ -25,7 +26,7 @@ "ThermodynamicKineticEnergy", "Dewpoint", "Temperature", "Theta_d", "RelativeHumidity", "Pressure", "Exner_Vt", "HydrostaticImbalance", "Precipitation", "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence", - "TracerDensity"] + "BruntVaisalaFrequencySquared", "TracerDensity"] class Diagnostics(object): @@ -221,6 +222,8 @@ def setup(self, domain, state_fields, space=None): def compute(self): """Compute the diagnostic field from the current state.""" + logger.debug(f'Computing diagnostic {self.name} with {self.method} method') + if self.method == 'interpolate': self.evaluator.interpolate() elif self.method == 'assign': @@ -944,6 +947,51 @@ def setup(self, domain, state_fields): super().setup(domain, state_fields) +class BruntVaisalaFrequencySquared(DiagnosticField): + """The diagnostic for the Brunt-Väisälä frequency.""" + name = "Brunt-Vaisala_squared" + + def __init__(self, equations, space=None, method='interpolate'): + """ + Args: + equations (:class:`PrognosticEquationSet`): the equation set being + solved by the model. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + 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.parameters = equations.parameters + # Work out required fields + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['theta'] + if equations.active_tracers is not None and len(equations.active_tracers) > 1: + # TODO: I think theta here should be theta_e, which would be + # easiest if this is a ThermodynamicDiagnostic. But in the dry + # case, our numerical theta_e does not reduce to the numerical + # dry theta + raise NotImplementedError( + 'Brunt-Vaisala diagnostic not implemented for moist equations') + else: + raise NotImplementedError( + f'Brunt-Vaisala diagnostic not implemented for {type(equations)}') + super().__init__(space=space, method=method, required_fields=tuple(required_fields)) + + 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. + """ + theta = state_fields('theta') + self.expr = self.parameters.g/theta * dot(domain.k, grad(theta)) + super().setup(domain, state_fields) + + class Sum(DiagnosticField): """Base diagnostic for computing the sum of two fields.""" def __init__(self, field_name1, field_name2): diff --git a/gusto/domain.py b/gusto/domain.py index 42ade2b53..8f22083c6 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -132,7 +132,7 @@ def __init__(self, mesh, dt, family, degree=None, radius_field = Function(CG1) radius_field.interpolate(r_shell) # TODO: this should use global min kernel - radius = np.min(radius_field.dat.data_ro) + radius = Constant(np.min(radius_field.dat.data_ro)) else: radius = None @@ -146,12 +146,47 @@ def __init__(self, mesh, dt, family, degree=None, _ = self.spaces('DG1_equispaced') self.coords.register_space(self, 'DG1_equispaced') + # Set height above surface (requires coordinates) + if hasattr(mesh, "_base_mesh"): + self.set_height_above_surface() + # -------------------------------------------------------------------- # # Construct metadata about domain # -------------------------------------------------------------------- # self.metadata = construct_domain_metadata(mesh, self.coords, self.on_sphere) + def set_height_above_surface(self): + """ + Sets a coordinate field which corresponds to height above the domain's + surface. + """ + + from firedrake import dot + + x = SpatialCoordinate(self.mesh) + + # Make a height field in CG1 + CG1 = FunctionSpace(self.mesh, "CG", 1, name='CG1') + self.spaces.add_space('CG1', CG1) + self.coords.register_space(self, 'CG1') + CG1_height = Function(CG1) + CG1_height.interpolate(dot(self.k, x)) + height_above_surface = Function(CG1) + + # Turn height into columnwise data + columnwise_height, index_data = self.coords.get_column_data(CG1_height, self) + + # Find minimum height in each column + surface_height_1d = np.min(columnwise_height, axis=1) + height_above_surface_data = columnwise_height - surface_height_1d[:, None] + + self.coords.set_field_from_column_data(height_above_surface, + height_above_surface_data, + index_data) + + self.height_above_surface = height_above_surface + def construct_domain_metadata(mesh, coords, on_sphere): """ diff --git a/gusto/forcing.py b/gusto/forcing.py index 662463515..71ab8bfd0 100644 --- a/gusto/forcing.py +++ b/gusto/forcing.py @@ -6,7 +6,7 @@ ) from gusto.fml import drop, replace_subject, name from gusto.labels import ( - transport, diffusion, time_derivative, hydrostatic + transport, diffusion, time_derivative, hydrostatic, physics_label ) from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual @@ -47,7 +47,7 @@ def __init__(self, equation, alpha): # drop terms relating to transport and diffusion residual = equation.residual.label_map( - lambda t: any(t.has_label(transport, diffusion, return_tuple=True)), drop) + lambda t: any(t.has_label(transport, diffusion, physics_label, return_tuple=True)), drop) # the lhs of both of the explicit and implicit solvers is just # the time derivative form diff --git a/gusto/physics.py b/gusto/physics.py index 1d62be0d8..76c8ee4c3 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -1,7 +1,7 @@ """ Objects to perform parametrisations of physical processes, or "physics". -"Physics" schemes are routines to compute updates to prognostic fields that +"PhysicsParametrisation" schemes are routines to compute updates to prognostic fields that represent the action of non-fluid processes, or those fluid processes that are unresolved. This module contains a set of these processes in the form of classes with "evaluate" methods. @@ -9,13 +9,17 @@ from abc import ABCMeta, abstractmethod from gusto.active_tracers import Phases, TracerVariableType +from gusto.configuration import BoundaryLayerParameters from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations from gusto.fml import identity, Term, subject from gusto.labels import PhysicsLabel, transporting_velocity, transport, prognostic from gusto.logging import logger -from firedrake import (Interpolator, conditional, Function, dx, - min_value, max_value, Constant, pi, Projector) +from firedrake import (Interpolator, conditional, Function, dx, sqrt, dot, + min_value, max_value, Constant, pi, Projector, grad, + TestFunctions, split, inner, TestFunction, exp, avg, + outer, FacetNormal, SpatialCoordinate, dS_v, + NonlinearVariationalProblem, NonlinearVariationalSolver) from gusto import thermodynamics import ufl import math @@ -24,7 +28,9 @@ __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", - "AdvectedMoments", "InstantRain", "SWSaturationAdjustment"] + "AdvectedMoments", "InstantRain", "SWSaturationAdjustment", + "SourceSink", "SurfaceFluxes", "WindDrag", "StaticAdjustment", + "SuppressVerticalWind", "BoundaryLayerMixing"] class PhysicsParametrisation(object, metaclass=ABCMeta): @@ -55,6 +61,108 @@ def evaluate(self): pass +class SourceSink(PhysicsParametrisation): + """ + The source or sink of some variable, described through a UFL expression. + + A scheme representing the general source or sink of a variable, described + through a UFL expression. The expression should be for the rate of change + of the variable. It is intended that the source/sink is independent of the + prognostic variables. + + The expression can also be a time-varying expression. In which case a + function should be provided, taking a :class:`Constant` as an argument (to + represent the time.) + """ + + def __init__(self, equation, variable_name, rate_expression, + time_varying=False, method='interpolate'): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + variable_name (str): the name of the variable + rate_expression (:class:`ufl.Expr` or func): an expression giving + the rate of change of the variable. If a time-varying expression + is needed, this should be a function taking a single argument + representing the time. Then the `time_varying` argument must + be set to True. + time_varying (bool, optional): whether the source/sink expression + varies with time. Defaults to False. + method (str, optional): the method to use to evaluate the expression + for the source. Valid options are 'interpolate' or 'project'. + Default is 'interpolate'. + """ + + label_name = f'source_sink_{variable_name}' + super().__init__(equation, label_name, parameters=None) + + if method not in ['interpolate', 'project']: + raise ValueError(f'Method {method} for source/sink evaluation not valid') + self.method = method + self.time_varying = time_varying + self.variable_name = variable_name + + # Check the variable exists + if hasattr(equation, "field_names"): + assert variable_name in equation.field_names, \ + f'Field {variable_name} does not exist in the equation set' + else: + assert variable_name == equation.field_name, \ + f'Field {variable_name} does not exist in the equation' + + # Work out the appropriate function space + if hasattr(equation, "field_names"): + V_idx = equation.field_names.index(variable_name) + W = equation.function_space + V = W.sub(V_idx) + test = equation.tests[V_idx] + else: + V = equation.function_space + test = equation.test + + # Make source/sink term + self.source = Function(V) + equation.residual += self.label(subject(test * self.source * dx, equation.X), + self.evaluate) + + # Handle whether the expression is time-varying or not + if self.time_varying: + expression = rate_expression(equation.domain.t) + else: + expression = rate_expression + + # Handle method of evaluating source/sink + if self.method == 'interpolate': + self.source_interpolator = Interpolator(expression, V) + else: + self.source_projector = Projector(expression, V) + + # If not time-varying, evaluate for the first time here + if not self.time_varying: + if self.method == 'interpolate': + self.source.assign(self.source_interpolator.interpolate()) + else: + self.source.assign(self.source_projector.project()) + + def evaluate(self, x_in, dt): + """ + Evalutes the source term generated by the physics. + + Args: + x_in: (:class:`Function`): the (mixed) field to be evolved. Unused. + dt: (:class:`Constant`): the timestep, which can be the time + interval for the scheme. Unused. + """ + if self.time_varying: + logger.info(f'Evaluating physics parametrisation {self.label.label}') + if self.method == 'interpolate': + self.source.assign(self.source_interpolator.interpolate()) + else: + self.source.assign(self.source_projector.project()) + else: + pass + + class SaturationAdjustment(PhysicsParametrisation): """ Represents the phase change between water vapour and cloud liquid. @@ -92,6 +200,7 @@ def __init__(self, equation, vapour_name='water_vapour', """ label_name = 'saturation_adjustment' + self.explicit_only = True super().__init__(equation, label_name, parameters=parameters) # TODO: make a check on the variable type of the active tracers @@ -100,8 +209,10 @@ def __init__(self, equation, vapour_name='water_vapour', # active tracer metadata corresponding to variable names # Check that fields exist - assert vapour_name in equation.field_names, f"Field {vapour_name} does not exist in the equation set" - assert cloud_name in equation.field_names, f"Field {cloud_name} does not exist in the equation set" + if vapour_name not in equation.field_names: + raise ValueError(f"Field {vapour_name} does not exist in the equation set") + if cloud_name not in equation.field_names: + raise ValueError(f"Field {cloud_name} does not exist in the equation set") # Make prognostic for physics scheme parameters = self.parameters @@ -280,7 +391,8 @@ def __init__(self, equation, rain_name, domain, transport_method, super().__init__(equation, label_name, parameters=None) # Check that fields exist - assert rain_name in equation.field_names, f"Field {rain_name} does not exist in the equation set" + if rain_name not in equation.field_names: + raise ValueError(f"Field {rain_name} does not exist in the equation set") # Check if variable is a mixing ratio rain_tracer = equation.get_active_tracer(rain_name) @@ -328,6 +440,7 @@ def __init__(self, equation, rain_name, domain, transport_method, terminal_velocity = Constant(5) # in m/s v.project(-terminal_velocity*domain.k) elif moments == AdvectedMoments.M3: + self.explicit_only = True # this advects the third moment M3 of the raindrop # distribution, which corresponds to the mean mass rho_idx = equation.field_names.index('rho') @@ -363,7 +476,12 @@ def __init__(self, equation, rain_name, domain, transport_method, + 'AdvectedMoments.M0 and AdvectedMoments.M3') if moments != AdvectedMoments.M0: - self.determine_v = Projector(-v_expression*domain.k, v) + # TODO: introduce reduced projector + test = TestFunction(Vu) + dx_reduced = dx(degree=4) + proj_eqn = inner(test, v + v_expression*domain.k)*dx_reduced + proj_prob = NonlinearVariationalProblem(proj_eqn, v) + self.determine_v = NonlinearVariationalSolver(proj_prob) def evaluate(self, x_in, dt): """ @@ -376,7 +494,7 @@ def evaluate(self, x_in, dt): logger.info(f'Evaluating physics parametrisation {self.label.label}') self.X.assign(x_in) if self.moments != AdvectedMoments.M0: - self.determine_v.project() + self.determine_v.solve() class Coalescence(PhysicsParametrisation): @@ -408,6 +526,7 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', process in the parametrisation. Defaults to True. """ + self.explicit_only = True label_name = "coalescence" if accretion: label_name += "_accretion" @@ -416,8 +535,10 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', super().__init__(equation, label_name, parameters=None) # Check that fields exist - assert cloud_name in equation.field_names, f"Field {cloud_name} does not exist in the equation set" - assert rain_name in equation.field_names, f"Field {rain_name} does not exist in the equation set" + if cloud_name not in equation.field_names: + raise ValueError(f"Field {cloud_name} does not exist in the equation set") + if rain_name not in equation.field_names: + raise ValueError(f"Field {rain_name} does not exist in the equation set") self.cloud_idx = equation.field_names.index(cloud_name) self.rain_idx = equation.field_names.index(rain_name) @@ -501,7 +622,7 @@ class EvaporationOfRain(PhysicsParametrisation): """ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', - latent_heat=True, parameters=None): + latent_heat=True): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. @@ -511,17 +632,15 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', Defaults to 'water_vapour'. latent_heat (bool, optional): whether to have latent heat exchange feeding back from the phase change. Defaults to True. - parameters (:class:`Configuration`, optional): parameters containing - the values of gas constants. Defaults to None, in which case the - parameters are obtained from the equation. Raises: NotImplementedError: currently this is only implemented for the CompressibleEulerEquations. """ + self.explicit_only = True label_name = 'evaporation_of_rain' - super().__init__(equation, label_name, parameters=parameters) + super().__init__(equation, label_name, parameters=None) # TODO: make a check on the variable type of the active tracers # if not a mixing ratio, we need to convert to mixing ratios @@ -529,8 +648,10 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', # active tracer metadata corresponding to variable names # Check that fields exist - assert vapour_name in equation.field_names, f"Field {vapour_name} does not exist in the equation set" - assert rain_name in equation.field_names, f"Field {rain_name} does not exist in the equation set" + if vapour_name not in equation.field_names: + raise ValueError(f"Field {vapour_name} does not exist in the equation set") + if rain_name not in equation.field_names: + raise ValueError(f"Field {rain_name} does not exist in the equation set") # Make prognostic for physics scheme self.X = Function(equation.X.function_space()) @@ -557,7 +678,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', V_idxs.append(theta_idx) # need to evaluate rho at theta-points, and do this via recovery - boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 1 else None + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None rho_averaged = Function(V) self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) @@ -676,10 +797,10 @@ def __init__(self, equation, saturation_curve, parameters=None): """ Args: - equation (:class: 'PrognosticEquationSet'): the model's equation. - saturation_curve (ufl expression or :class: `function`): the - curve above which excess moisture is converted to rain. Is - either prescribed or dependent on a prognostic field. + equation (:class:`PrognosticEquationSet`): the model's equation. + saturation_curve (:class:`ufl.Expr` or func): the curve above which + excess moisture is converted to rain. Is either prescribed or + dependent on a prognostic field. time_varying_saturation (bool, optional): set this to True if the saturation curve is changing in time. Defaults to False. vapour_name (str, optional): name of the water vapour variable. @@ -702,6 +823,7 @@ def __init__(self, equation, saturation_curve, parameters are obtained from the equation. """ + self.explicit_only = True label_name = 'instant_rain' super().__init__(equation, label_name, parameters=parameters) @@ -822,14 +944,14 @@ def __init__(self, equation, saturation_curve, L_v=None, parameters=None): """ Args: - equation (:class: 'PrognosticEquationSet'): the model's equation - saturation_curve (ufl expression or :class: `function`): the - curve which dictates when phase changes occur. In a saturated - atmosphere vapour above the saturation curve becomes cloud, - and if the atmosphere is sub-saturated and there is cloud - present cloud will become vapour until the saturation curve is - reached. The saturation curve is either prescribed or - dependent on a prognostic field. + equation (:class:`PrognosticEquationSet`): the model's equation + saturation_curve (:class:`ufl.Expr` or func): the curve which + dictates when phase changes occur. In a saturated atmosphere + vapour above the saturation curve becomes cloud, and if the + atmosphere is sub-saturated and there is cloud present cloud + will become vapour until the saturation curve is reached. The + saturation curve is either prescribed or dependent on a + prognostic field. time_varying_saturation (bool, optional): set this to True if the saturation curve is changing in time. Defaults to False. L_v (float, optional): The air expansion factor multiplied by the @@ -869,6 +991,7 @@ def __init__(self, equation, saturation_curve, L_v=None, parameters are obtained from the equation. """ + self.explicit_only = True label_name = 'saturation_adjustment' super().__init__(equation, label_name, parameters=parameters) @@ -1002,3 +1125,627 @@ def evaluate(self, x_in, dt): self.gamma_v.interpolate(self.gamma_v_computation(x_in)) for interpolator in self.source_interpolators: interpolator.interpolate() + + +class SurfaceFluxes(PhysicsParametrisation): + """ + Prescribed surface temperature and moisture fluxes, to be used in aquaplanet + simulations as Sea Surface Temperature fluxes. This formulation is taken + from the DCMIP (2016) test case document. + + These can be used either with an in-built implicit formulation, or with a + generic time discretisation. + + Written to assume that dry density is unchanged by the parametrisation. + """ + + def __init__(self, equation, T_surface_expr, vapour_name=None, + implicit_formulation=False, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + T_surface_expr (:class:`ufl.Expr`): the surface temperature. + vapour_name (str, optional): name of the water vapour variable. + Defaults to None, in which case no moisture fluxes are applied. + implicit_formulation (bool, optional): whether the scheme is already + put into a Backwards Euler formulation (which allows this scheme + to actually be used with a Forwards Euler or other explicit time + discretisation). Otherwise, this is formulated more generally + and can be used with any time stepper. Defaults to False. + parameters (:class:`BoundaryLayerParameters`): configuration object + giving the parameters for boundary and surface level schemes. + Defaults to None, in which case default values are used. + """ + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Surface fluxes can only be used with Compressible Euler equations") + + if vapour_name is not None: + if vapour_name not in equation.field_names: + raise ValueError(f"Field {vapour_name} does not exist in the equation set") + + if parameters is None: + parameters = BoundaryLayerParameters() + + label_name = 'surface_flux' + super().__init__(equation, label_name, parameters=None) + + self.implicit_formulation = implicit_formulation + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + u_idx = equation.field_names.index('u') + T_idx = equation.field_names.index('theta') + rho_idx = equation.field_names.index('rho') + if vapour_name is not None: + m_v_idx = equation.field_names.index(vapour_name) + + X = self.X + tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests + + u = split(X)[u_idx] + rho = split(X)[rho_idx] + theta_vd = split(X)[T_idx] + test_theta = tests[T_idx] + + if vapour_name is not None: + m_v = split(X)[m_v_idx] + test_m_v = tests[m_v_idx] + else: + m_v = None + + if implicit_formulation: + # Need to evaluate rho at theta-points + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None + rho_averaged = Function(equation.function_space.sub(T_idx)) + self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) + exner = thermodynamics.exner_pressure(equation.parameters, rho_averaged, theta_vd) + else: + # Exner is more general expression + exner = thermodynamics.exner_pressure(equation.parameters, rho, theta_vd) + + # Alternative variables + T = thermodynamics.T(equation.parameters, theta_vd, exner, r_v=m_v) + p = thermodynamics.p(equation.parameters, exner) + + # -------------------------------------------------------------------- # + # Expressions for surface fluxes + # -------------------------------------------------------------------- # + z = equation.domain.height_above_surface + z_a = parameters.height_surface_layer + surface_expr = conditional(z < z_a, Constant(1.0), Constant(0.0)) + + u_hori = u - equation.domain.k*dot(u, equation.domain.k) + u_hori_mag = sqrt(dot(u_hori, u_hori)) + + C_H = parameters.coeff_heat + C_E = parameters.coeff_evap + + # Implicit formulation ----------------------------------------------- # + # For use with ForwardEuler only, as implicit solution is hand-written + if implicit_formulation: + self.source_interpolators = [] + + # First specify T_np1 expression + Vtheta = equation.spaces[T_idx] + T_np1_expr = ((T + C_H*u_hori_mag*T_surface_expr*self.dt/z_a) + / (1 + C_H*u_hori_mag*self.dt/z_a)) + + # If moist formulation, determine next vapour value + if vapour_name is not None: + source_mv = Function(Vtheta) + mv_sat = thermodynamics.r_sat(equation.parameters, T, p) + mv_np1_expr = ((m_v + C_E*u_hori_mag*mv_sat*self.dt/z_a) + / (1 + C_E*u_hori_mag*self.dt/z_a)) + dmv_expr = surface_expr * (mv_np1_expr - m_v) / self.dt + source_mv_expr = test_m_v * source_mv * dx + + self.source_interpolators.append(Interpolator(dmv_expr, source_mv)) + equation.residual -= self.label(subject(prognostic(source_mv_expr, vapour_name), + X), self.evaluate) + + # Moisture needs including in theta_vd expression + # NB: still using old pressure here, which implies constant p? + epsilon = equation.parameters.R_d / equation.parameters.R_v + theta_np1_expr = (thermodynamics.theta(equation.parameters, T_np1_expr, p) + * (1 + mv_np1_expr / epsilon)) + + else: + theta_np1_expr = thermodynamics.theta(equation.parameters, T_np1_expr, p) + + source_theta_vd = Function(Vtheta) + dtheta_vd_expr = surface_expr * (theta_np1_expr - theta_vd) / self.dt + source_theta_expr = test_theta * source_theta_vd * dx + self.source_interpolators.append(Interpolator(dtheta_vd_expr, source_theta_vd)) + equation.residual -= self.label(subject(prognostic(source_theta_expr, 'theta'), + X), self.evaluate) + + # General formulation ------------------------------------------------ # + else: + # Construct underlying expressions + kappa = equation.parameters.kappa + dT_dt = surface_expr * C_H * u_hori_mag * (T_surface_expr - T) / z_a + + if vapour_name is not None: + mv_sat = thermodynamics.r_sat(equation.parameters, T, p) + dmv_dt = surface_expr * C_E * u_hori_mag * (mv_sat - m_v) / z_a + source_mv_expr = test_m_v * dmv_dt * dx + equation.residual -= self.label( + prognostic(subject(source_mv_expr, X), + vapour_name), self.evaluate) + + # Theta expression depends on dmv_dt + epsilon = equation.parameters.R_d / equation.parameters.R_v + dtheta_vd_dt = (dT_dt * ((1 + m_v / epsilon) / exner - kappa * theta_vd / (rho * T)) + + dmv_dt * (T / (epsilon * exner) - kappa * theta_vd / (epsilon + m_v))) + else: + dtheta_vd_dt = dT_dt * (1 / exner - kappa * theta_vd / (rho * T)) + + dx_reduced = dx(degree=4) + source_theta_expr = test_theta * dtheta_vd_dt * dx_reduced + + equation.residual -= self.label( + subject(prognostic(source_theta_expr, 'theta'), X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. This does nothing if + the implicit formulation is not used. + + Args: + x_in: (:class: 'Function'): the (mixed) field to be evolved. + dt: (:class: 'Constant'): the timestep, which can be the time + interval for the scheme. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + if self.implicit_formulation: + self.X.assign(x_in) + self.dt.assign(dt) + self.rho_recoverer.project() + for source_interpolator in self.source_interpolators: + source_interpolator.interpolate() + + +class WindDrag(PhysicsParametrisation): + """ + A simple surface wind drag scheme. This formulation is taken from the DCMIP + (2016) test case document. + + These can be used either with an in-built implicit formulation, or with a + generic time discretisation. + """ + + def __init__(self, equation, implicit_formulation=False, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + implicit_formulation (bool, optional): whether the scheme is already + put into a Backwards Euler formulation (which allows this scheme + to actually be used with a Forwards Euler or other explicit time + discretisation). Otherwise, this is formulated more generally + and can be used with any time stepper. Defaults to False. + parameters (:class:`BoundaryLayerParameters`): configuration object + giving the parameters for boundary and surface level schemes. + Defaults to None, in which case default values are used. + """ + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Wind drag can only be used with Compressible Euler equations") + + if parameters is None: + parameters = BoundaryLayerParameters() + + label_name = 'wind_drag' + super().__init__(equation, label_name, parameters=None) + + k = equation.domain.k + self.implicit_formulation = implicit_formulation + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + u_idx = equation.field_names.index('u') + + X = self.X + tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests + + test = tests[u_idx] + + u = split(X)[u_idx] + u_hori = u - k*dot(u, k) + u_hori_mag = sqrt(dot(u_hori, u_hori)) + + # -------------------------------------------------------------------- # + # Expressions for wind drag + # -------------------------------------------------------------------- # + z = equation.domain.height_above_surface + z_a = parameters.height_surface_layer + surface_expr = conditional(z < z_a, Constant(1.0), Constant(0.0)) + + C_D0 = parameters.coeff_drag_0 + C_D1 = parameters.coeff_drag_1 + C_D2 = parameters.coeff_drag_2 + + C_D = conditional(u_hori_mag < 20.0, C_D0 + C_D1*u_hori_mag, C_D2) + + # Implicit formulation ----------------------------------------------- # + # For use with ForwardEuler only, as implicit solution is hand-written + if implicit_formulation: + + # First specify T_np1 expression + Vu = equation.spaces[u_idx] + source_u = Function(Vu) + u_np1_expr = u_hori / (1 + C_D*u_hori_mag*self.dt/z_a) + + du_expr = surface_expr * (u_np1_expr - u_hori) / self.dt + + # TODO: introduce reduced projector + test_Vu = TestFunction(Vu) + dx_reduced = dx(degree=4) + proj_eqn = inner(test_Vu, source_u - du_expr)*dx_reduced + proj_prob = NonlinearVariationalProblem(proj_eqn, source_u) + self.source_projector = NonlinearVariationalSolver(proj_prob) + + source_expr = inner(test, source_u - k*dot(source_u, k)) * dx + equation.residual -= self.label(subject(prognostic(source_expr, 'u'), + X), self.evaluate) + + # General formulation ------------------------------------------------ # + else: + # Construct underlying expressions + du_dt = -surface_expr * C_D * u_hori_mag * u_hori / z_a + + dx_reduced = dx(degree=4) + source_expr = inner(test, du_dt) * dx_reduced + + equation.residual -= self.label(subject(prognostic(source_expr, 'u'), X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. This does nothing if + the implicit formulation is not used. + + Args: + x_in: (:class: 'Function'): the (mixed) field to be evolved. + dt: (:class: 'Constant'): the timestep, which can be the time + interval for the scheme. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + if self.implicit_formulation: + self.X.assign(x_in) + self.dt.assign(dt) + self.source_projector.solve() + + +class StaticAdjustment(PhysicsParametrisation): + """ + A scheme to provide static adjustment, by sorting the potential temperature + values in a column so that they are increasing with height. + """ + + def __init__(self, equation, theta_variable='theta_vd'): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + theta_variable (str, optional): which theta variable to sort the + profile for. Valid options are "theta" or "theta_vd". Defaults + to "theta_vd". + """ + + self.explicit_only = True + from functools import partial + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Static adjustment can only be used with Compressible Euler equations") + + if theta_variable not in ['theta', 'theta_vd']: + raise ValueError('Static adjustment: theta variable ' + + f'{theta_variable} not valid') + + label_name = 'static_adjustment' + super().__init__(equation, label_name, parameters=equation.parameters) + + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + + theta_idx = equation.field_names.index('theta') + Vt = equation.spaces[theta_idx] + self.theta_to_sort = Function(Vt) + sorted_theta = Function(Vt) + theta = self.X.subfunctions[theta_idx] + + if theta_variable == 'theta' and 'water_vapour' in equation.field_names: + Rv = equation.parameters.R_v + Rd = equation.parameters.R_d + mv_idx = equation.field_names.index('water_vapour') + mv = self.X.subfunctions[mv_idx] + self.get_theta_variable = Interpolator(theta / (1 + mv*Rv/Rd), self.theta_to_sort) + self.set_theta_variable = Interpolator(self.theta_to_sort * (1 + mv*Rv/Rd), sorted_theta) + else: + self.get_theta_variable = Interpolator(theta, self.theta_to_sort) + self.set_theta_variable = Interpolator(self.theta_to_sort, sorted_theta) + + # -------------------------------------------------------------------- # + # Set up routines to reshape data + # -------------------------------------------------------------------- # + + domain = equation.domain + self.get_column_data = partial(domain.coords.get_column_data, + field=self.theta_to_sort, + domain=domain) + self.set_column_data = domain.coords.set_field_from_column_data + + # -------------------------------------------------------------------- # + # Set source term + # -------------------------------------------------------------------- # + + tests = TestFunctions(self.X.function_space()) + test = tests[theta_idx] + + source_expr = inner(test, sorted_theta - theta) / self.dt * dx + + equation.residual -= self.label(subject(prognostic(source_expr, 'theta'), equation.X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. This does nothing if + the implicit formulation is not used. + + Args: + x_in: (:class: 'Function'): the (mixed) field to be evolved. + dt: (:class: 'Constant'): the timestep, which can be the time + interval for the scheme. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + self.X.assign(x_in) + self.dt.assign(dt) + + self.get_theta_variable.interpolate() + theta_column_data, index_data = self.get_column_data() + for col in range(theta_column_data.shape[0]): + theta_column_data[col].sort() + self.set_column_data(self.theta_to_sort, theta_column_data, index_data) + self.set_theta_variable.interpolate() + + +class SuppressVerticalWind(PhysicsParametrisation): + """ + Suppresses the model's vertical wind, reducing it to zero. This is used for + instance in a model's spin up period. + """ + + def __init__(self, equation, spin_up_period): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + spin_up_period (`ufl.Constant`): this parametrisation is applied + while the time is less than this -- corresponding to a spin up + period. + """ + + self.explicit_only = True + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + + domain = equation.domain + + if not domain.mesh.extruded: + raise RuntimeError("Suppress vertical wind can only be used with " + + "extruded meshes") + + label_name = 'suppress_vertical_wind' + super().__init__(equation, label_name, parameters=equation.parameters) + + self.X = Function(equation.X.function_space()) + self.dt = Constant(0.0) + self.t = domain.t + self.spin_up_period = Constant(spin_up_period) + self.strength = Constant(1.0) + self.spin_up_done = False + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + + u_idx = equation.field_names.index('u') + wind = self.X.subfunctions[u_idx] + + # -------------------------------------------------------------------- # + # Set source term + # -------------------------------------------------------------------- # + + tests = TestFunctions(self.X.function_space()) + test = tests[u_idx] + + # The sink should be just the value of the current vertical wind + source_expr = -self.strength * inner(test, domain.k*dot(domain.k, wind)) / self.dt * dx + + equation.residual -= self.label(subject(prognostic(source_expr, 'u'), equation.X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. This does nothing if + the implicit formulation is not used. + + Args: + x_in: (:class: 'Function'): the (mixed) field to be evolved. + dt: (:class: 'Constant'): the timestep, which can be the time + interval for the scheme. + """ + + if float(self.t) < float(self.spin_up_period): + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + self.X.assign(x_in) + self.dt.assign(dt) + + elif not self.spin_up_done: + self.strength.assign(0.0) + self.spin_up_done = True + + +class BoundaryLayerMixing(PhysicsParametrisation): + """ + A simple boundary layer mixing scheme. This acts like a vertical diffusion, + using an interior penalty method. + """ + + def __init__(self, equation, field_name, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + field_name (str): name of the field to apply the diffusion to. + ibp (:class:`IntegrateByParts`, optional): how many times to + integrate the term by parts. Defaults to IntegrateByParts.ONCE. + parameters (:class:`BoundaryLayerParameters`): configuration object + giving the parameters for boundary and surface level schemes. + Defaults to None, in which case default values are used. + """ + + # -------------------------------------------------------------------- # + # Check arguments and generic initialisation + # -------------------------------------------------------------------- # + + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Boundary layer mixing can only be used with Compressible Euler equations") + + if field_name not in equation.field_names: + raise ValueError(f'field {field_name} not found in equation') + + if parameters is None: + parameters = BoundaryLayerParameters() + + label_name = f'boundary_layer_mixing_{field_name}' + super().__init__(equation, label_name, parameters=None) + + self.X = Function(equation.X.function_space()) + + # -------------------------------------------------------------------- # + # Extract prognostic variables + # -------------------------------------------------------------------- # + + u_idx = equation.field_names.index('u') + T_idx = equation.field_names.index('theta') + rho_idx = equation.field_names.index('rho') + + u = split(self.X)[u_idx] + rho = split(self.X)[rho_idx] + theta_vd = split(self.X)[T_idx] + + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None + rho_averaged = Function(equation.function_space.sub(T_idx)) + self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) + exner = thermodynamics.exner_pressure(equation.parameters, rho_averaged, theta_vd) + + # Alternative variables + p = thermodynamics.p(equation.parameters, exner) + p_top = Constant(85000.) + p_strato = Constant(10000.) + + # -------------------------------------------------------------------- # + # Expressions for diffusivity coefficients + # -------------------------------------------------------------------- # + z_a = parameters.height_surface_layer + + domain = equation.domain + u_hori = u - domain.k*dot(u, domain.k) + u_hori_mag = sqrt(dot(u_hori, u_hori)) + + if field_name == 'u': + C_D0 = parameters.coeff_drag_0 + C_D1 = parameters.coeff_drag_1 + C_D2 = parameters.coeff_drag_2 + + C_D = conditional(u_hori_mag < 20.0, C_D0 + C_D1*u_hori_mag, C_D2) + K = conditional(p > p_top, + C_D*u_hori_mag*z_a, + C_D*u_hori_mag*z_a*exp(-((p_top - p)/p_strato)**2)) + + else: + C_E = parameters.coeff_evap + K = conditional(p > p_top, + C_E*u_hori_mag*z_a, + C_E*u_hori_mag*z_a*exp(-((p_top - p)/p_strato)**2)) + + # -------------------------------------------------------------------- # + # Make source expression + # -------------------------------------------------------------------- # + + dx_reduced = dx(degree=4) + dS_v_reduced = dS_v(degree=4) + # Need to be careful with order of operations here, to correctly index + # when the field is vector-valued + d_dz = lambda q: outer(domain.k, dot(grad(q), domain.k)) + n = FacetNormal(domain.mesh) + # Work out vertical height + xyz = SpatialCoordinate(domain.mesh) + Vr = domain.spaces('L2') + dz = Function(Vr) + dz.interpolate(dot(domain.k, d_dz(dot(domain.k, xyz)))) + mu = parameters.mu + + X = self.X + tests = equation.tests + + idx = equation.field_names.index(field_name) + test = tests[idx] + field = X.subfunctions[idx] + + if field_name == 'u': + # Horizontal diffusion only + field = field - domain.k*dot(field, domain.k) + + # Interior penalty discretisation of vertical diffusion + source_expr = ( + # Volume term + rho_averaged*K*inner(d_dz(test/rho_averaged), d_dz(field))*dx_reduced + # Interior penalty surface term + - 2*inner(avg(outer(K*field, n)), avg(d_dz(test)))*dS_v_reduced + - 2*inner(avg(outer(test, n)), avg(d_dz(K*field)))*dS_v_reduced + + 4*mu*avg(dz)*inner(avg(outer(K*field, n)), avg(outer(test, n)))*dS_v_reduced + ) + + equation.residual += self.label( + subject(prognostic(source_expr, field_name), X), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. This only recovers + the density field. + + Args: + x_in: (:class: 'Function'): the (mixed) field to be evolved. + dt: (:class: 'Constant'): the timestep, which can be the time + interval for the scheme. + """ + + logger.info(f'Evaluating physics parametrisation {self.label.label}') + + self.X.assign(x_in) + self.rho_recoverer.project() diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 50efad05d..ec6428228 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -444,6 +444,18 @@ def rhs(self): map_if_true=drop, map_if_false=lambda t: -1*t) + # If there are no active labels, we may have no terms at this point + # So that we can still do xnp1 = xn, put in a zero term here + if len(r.terms) == 0: + logger.warning('No terms detected for RHS of explicit problem. ' + + 'Adding a zero term to avoid failure.') + null_term = Constant(0.0)*self.residual.label_map( + lambda t: t.has_label(time_derivative), + # Drop label from this + map_if_true=lambda t: time_derivative.remove(t), + map_if_false=drop) + r += null_term + return r.form def solve_stage(self, x0, stage): @@ -604,13 +616,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, options (:class:`AdvectionOptions`, optional): an object containing options to either be passed to the spatial discretisation, or to control the "wrapper" methods. Defaults to None. - - Raises: - NotImplementedError: if options is an instance of - EmbeddedDGOptions or RecoveryOptions """ - if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): - raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") if not solver_parameters: # default solver parameters solver_parameters = {'ksp_type': 'gmres', @@ -649,6 +655,13 @@ def apply(self, x_out, x_in): x_out (:class:`Function`): the output field to be computed. x_in (:class:`Function`): the input field. """ + for evaluate in self.evaluate_source: + evaluate(x_in, self.dt) + + if len(self.evaluate_source) > 0: + # If we have physics, use x_in as first guess + self.x_out.assign(x_in) + self.x1.assign(x_in) self.solver.solve() x_out.assign(self.x_out) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 48073a030..a7fc661b3 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,7 +1,7 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector, Constant, split +from firedrake import Function, Projector, split, Constant from pyop2.profiling import timed_stage from gusto.equations import PrognosticEquationSet from gusto.fml import drop, Label, Term @@ -346,12 +346,14 @@ def __init__(self, equation, scheme, io, spatial_methods=None, else: self.physics_schemes = [] - for _, phys_scheme in self.physics_schemes: - # check that the supplied schemes for physics are explicit - assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \ - "Only explicit time discretisations can be used for physics" + for parametrisation, phys_scheme in self.physics_schemes: + # check that the supplied schemes for physics are valid + if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: + assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \ + ("Only explicit time discretisations can be used with " + + f"physics scheme {parametrisation.label.label}") apply_bcs = False - phys_scheme.setup(equation, apply_bcs, physics_label) + phys_scheme.setup(equation, apply_bcs, parametrisation.label) @property def transporting_velocity(self): @@ -460,10 +462,12 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, + self.fast_physics_schemes + self.final_physics_schemes) - for _, scheme in self.all_physics_schemes: + for parametrisation, scheme in self.all_physics_schemes: assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" - assert isinstance(scheme, ExplicitTimeDiscretisation), \ - "Only explicit time discretisations can be used for physics" + if hasattr(parametrisation, "explicit_only") and parametrisation.explicit_only: + assert isinstance(scheme, ExplicitTimeDiscretisation), \ + ("Only explicit time discretisations can be used with " + + f"physics scheme {parametrisation.label.label}") self.active_transport = [] for scheme in transport_schemes: @@ -574,9 +578,9 @@ def setup_scheme(self): apply_bcs = True for _, scheme in self.diffusion_schemes: scheme.setup(self.equation, apply_bcs, diffusion) - for _, scheme in self.all_physics_schemes: + for parametrisation, scheme in self.all_physics_schemes: apply_bcs = True - scheme.setup(self.equation, apply_bcs, physics_label) + scheme.setup(self.equation, apply_bcs, parametrisation.label) def copy_active_tracers(self, x_in, x_out): """ @@ -613,6 +617,7 @@ def timestep(self): with timed_stage("Apply forcing terms"): logger.info('SIQN: Explicit forcing') + # TODO: check if forcing is applied to x_after_slow or xn # Put explicit forcing into xstar self.forcing.apply(x_after_slow, x_after_slow, xstar(self.field_name), "explicit") @@ -631,7 +636,7 @@ def timestep(self): x_after_fast(self.field_name).assign(xp(self.field_name)) if len(self.fast_physics_schemes) > 0: with timed_stage("Fast physics"): - logger.info('SIQN: Fast physics') + logger.info(f'SIQN: Fast physics {outer}') for _, scheme in self.fast_physics_schemes: scheme.apply(x_after_fast(scheme.field_name), x_after_fast(scheme.field_name)) @@ -642,6 +647,7 @@ def timestep(self): with timed_stage("Apply forcing terms"): logger.info(f'SIQN: Implicit forcing {(outer, inner)}') + # TODO: why don't we update xnp1 with a better guess here? self.forcing.apply(xp, xnp1, xrhs, "implicit") xrhs -= xnp1(self.field_name) @@ -655,7 +661,7 @@ def timestep(self): xnp1X += dy # Update xnp1 values for active tracers not included in the linear solve - self.copy_active_tracers(xp, xnp1) + self.copy_active_tracers(x_after_fast, xnp1) self._apply_bcs() diff --git a/integration-tests/physics/test_boundary_layer_mixing.py b/integration-tests/physics/test_boundary_layer_mixing.py new file mode 100644 index 000000000..1b0499e70 --- /dev/null +++ b/integration-tests/physics/test_boundary_layer_mixing.py @@ -0,0 +1,143 @@ +""" +This tests the physics routine to mix fields in the boundary layer. +""" + +from gusto import * +from gusto.labels import physics_label +from firedrake import (VectorFunctionSpace, PeriodicIntervalMesh, as_vector, + exp, SpatialCoordinate, ExtrudedMesh, Function) +import pytest + + +def run_boundary_layer_mixing(dirname, field_name, recovered, semi_implicit): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + element_degree = 1 if field_name == 'u' and not recovered else 0 + dt = 100.0 + + # declare grid shape, with length L and height H + L = 500. + H = 500. + nlayers = 5 + ncolumns = 3 + + # make mesh and domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", element_degree) + + _, z = SpatialCoordinate(mesh) + + # Set up equation + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters) + + # I/O + output = OutputParameters(dirname=dirname+"/boundary_layer_mixing", + dumpfreq=1, + dumplist=[field_name]) + io = IO(domain, output) + + # Physics scheme + surf_params = BoundaryLayerParameters() + physics_parametrisation = BoundaryLayerMixing(eqn, field_name, surf_params) + + if recovered: + # Only implemented for u + Vec_CG1 = VectorFunctionSpace(mesh, 'CG', 1) + Vec_CG1 = VectorFunctionSpace(mesh, 'DG', 1) + recovery_opts = RecoveryOptions(embedding_space=Vec_CG1, + recovered_space=Vec_CG1, + boundary_method=BoundaryMethod.taylor) + implicit_discretisation = BackwardEuler(domain, field_name=field_name, options=recovery_opts) + else: + implicit_discretisation = BackwardEuler(domain) + + if semi_implicit: + if recovered: + explicit_discretisation = ForwardEuler(domain, field_name=field_name, options=recovery_opts) + else: + explicit_discretisation = ForwardEuler(domain) + + # Use half of the time discretisation for each + explicit_discretisation.dt.assign(domain.dt/2) + implicit_discretisation.dt.assign(domain.dt/2) + physics_schemes = [(physics_parametrisation, explicit_discretisation), + (physics_parametrisation, implicit_discretisation)] + else: + physics_schemes = [(physics_parametrisation, implicit_discretisation)] + + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vt = domain.spaces("theta") + Vu = domain.spaces("HDiv") + + # Declare prognostic fields + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Set prognostic variables + theta0.interpolate(300.*exp(-z/(2*H))) + rho0.interpolate(1.1*exp(-z/(5*H))) + + u0.project(as_vector([5.0*exp(-z/(0.5*H)), 0.0])) + + if field_name == 'theta': + initial_field = Function(Vt) + initial_field.assign(theta0) + elif field_name == 'u': + initial_field = Function(Vu) + initial_field.assign(u0) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return domain, stepper, initial_field + + +@pytest.mark.parametrize("field_name, recovered, semi_implicit", + [('theta', False, False), + ('theta', False, True), + ('u', False, False), + pytest.param('u', True, True, + marks=pytest.mark.xfail(reason='recovered physics not implemented')) + ]) +def test_boundary_layer_mixing(tmpdir, field_name, recovered, semi_implicit): + + dirname = str(tmpdir) + domain, stepper, initial_field = \ + run_boundary_layer_mixing(dirname, field_name, recovered, semi_implicit) + + if field_name == 'u': + # Need to project horizontal wind into W3 + wind_2d = stepper.fields(field_name) + field = Function(domain.spaces('L2')).project(wind_2d[0]) + initial_1d = Function(domain.spaces('L2')).project(initial_field[0]) + # Relabel initial field + initial_field = initial_1d + else: + field = stepper.fields(field_name) + + field_data, _ = domain.coords.get_column_data(field, domain) + initial_data, _ = domain.coords.get_column_data(initial_field, domain) + + # Check first column + assert field_data[0, 0] < 0.999*initial_data[0, 0] diff --git a/integration-tests/physics/test_saturation_adjustment.py b/integration-tests/physics/test_saturation_adjustment.py index 922e7e40e..767e0c787 100644 --- a/integration-tests/physics/test_saturation_adjustment.py +++ b/integration-tests/physics/test_saturation_adjustment.py @@ -1,7 +1,7 @@ """ -# This tests the Condensation routine. It creates a bubble of water vapour that -# is advected by a prescribed velocity. The test passes if the integral -# of the water mixing ratio is conserved. +This tests the SaturationAdjustment routine. It creates a bubble of water vapour +that is advected by a prescribed velocity. The test passes if the integral +of the water mixing ratio is conserved. """ from os import path @@ -35,8 +35,6 @@ def run_cond_evap(dirname, process): x, z = SpatialCoordinate(mesh) - # spaces - # Set up equation tracers = [WaterVapour(), CloudWater()] parameters = CompressibleParameters() diff --git a/integration-tests/physics/test_source_sink.py b/integration-tests/physics/test_source_sink.py new file mode 100644 index 000000000..4df01599e --- /dev/null +++ b/integration-tests/physics/test_source_sink.py @@ -0,0 +1,108 @@ +""" +This tests the source/sink +""" + +from gusto import * +from firedrake import (as_vector, PeriodicSquareMesh, SpatialCoordinate, + sqrt, sin, pi, assemble, Constant) +import pytest + + +def run_source_sink(dirname, process, time_varying): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # set up mesh and domain + L = 10 + nx = 10 + mesh = PeriodicSquareMesh(nx, nx, L, quadrilateral=True) + dt = 0.1 + tmax = 5*dt + domain = Domain(mesh, dt, "RTCF", 1) + x, y = SpatialCoordinate(mesh) + + # Equation + V = domain.spaces('DG') + eqn = AdvectionEquation(domain, V, "ash") + + # I/O + output = OutputParameters(dirname=dirname+"/source_sink", + dumpfreq=1) + diagnostic_fields = [CourantNumber()] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + transport_method = [DGUpwind(eqn, "ash")] + + # Physics scheme --------------------------------------------------------- # + # Source is a Lorentzian centred on a point + centre_x = L / 4.0 + centre_y = 3*L / 4.0 + width = L / 8.0 + dist_x = periodic_distance(x, centre_x, L) + dist_y = periodic_distance(y, centre_y, L) + dist = sqrt(dist_x**2 + dist_y**2) + # Lorentzian function + basic_expression = width / (dist**2 + width**2) + + if process == 'source': + basic_expression = -basic_expression + + def time_varying_expression(t): + return 2*basic_expression*sin(pi*t/(2.0*tmax)) + + if time_varying: + expression = time_varying_expression + else: + expression = basic_expression + + physics_parametrisations = [SourceSink(eqn, 'ash', expression, time_varying)] + + # Time stepper + stepper = PrescribedTransport(eqn, SSPRK3(domain), io, transport_method, + physics_parametrisations=physics_parametrisations) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + ash0 = stepper.fields("ash") + + if process == "source": + # Start with no ash + background_ash = Constant(0.0) + elif process == "sink": + # Take ash away + background_ash = Constant(100.0) + ash0.interpolate(background_ash) + initial_ash = Function(V).assign(ash0) + + # Constant wind + u = stepper.fields("u") + u.project(as_vector([1.0, 0.0])) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=tmax) + return stepper, initial_ash + + +@pytest.mark.parametrize("process", ["source", "sink"]) +@pytest.mark.parametrize("time_varying", [False, True]) +def test_source_sink(tmpdir, process, time_varying): + dirname = str(tmpdir) + stepper, initial_ash = run_source_sink(dirname, process, time_varying) + final_ash = stepper.fields("ash") + + initial_total_ash = assemble(initial_ash*dx) + final_total_ash = assemble(final_ash*dx) + + tol = 1.0 + if process == "source": + assert final_total_ash > initial_total_ash + tol, \ + "Source process does not appear to have created tracer" + else: + assert final_total_ash < initial_total_ash - tol, \ + "Sink process does not appear to have removed tracer" diff --git a/integration-tests/physics/test_static_adjustment.py b/integration-tests/physics/test_static_adjustment.py new file mode 100644 index 000000000..18ee13cea --- /dev/null +++ b/integration-tests/physics/test_static_adjustment.py @@ -0,0 +1,109 @@ +""" +This tests the physics routine to apply static adjustment. A column initially +has a decreasing theta profile (which would be unstable). The static adjustment +should then sort this to make it increasing with height. +""" + +from gusto import * +from gusto.labels import physics_label +from firedrake import (Constant, PeriodicIntervalMesh, + SpatialCoordinate, ExtrudedMesh, Function) +import pytest + + +def run_static_adjustment(dirname, theta_variable): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 100.0 + + # declare grid shape, with length L and height H + L = 500. + H = 500. + nlayers = 5 + ncolumns = 5 + + # make mesh and domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 0) + domain.coords.register_space(domain, 'theta') + + _, z = SpatialCoordinate(mesh) + + # Set up equation + tracers = [WaterVapour()] + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + + # I/O + output = OutputParameters(dirname=dirname+"/static_adjustment", + dumpfreq=1, + dumplist=['theta']) + io = IO(domain, output, diagnostic_fields=[Perturbation('theta')]) + + # Physics scheme + physics_parametrisation = StaticAdjustment(eqn, theta_variable) + + time_discretisation = ForwardEuler(domain) + + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] + + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + # Declare prognostic fields + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Set prognostic variables -- decreasing theta profile + water_v0 = stepper.fields("water_vapour") + water_v0.interpolate(Constant(0.01) + 0.001*z/H) + theta0.interpolate(Constant(300) - 20*z/H) + rho0.interpolate(Constant(1.0)) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return domain, eqn, stepper + + +@pytest.mark.parametrize("theta_variable", ["theta", "theta_vd"]) +def test_static_adjustment(tmpdir, theta_variable): + + dirname = str(tmpdir) + domain, eqn, stepper = run_static_adjustment(dirname, theta_variable) + + # Back out temperature from prognostic fields + theta_vd = stepper.fields('theta') + if theta_variable == 'theta': + Rv = eqn.parameters.R_v + Rd = eqn.parameters.R_d + mv = stepper.fields('water_vapour') + theta = Function(theta_vd.function_space()) + theta.interpolate(theta_vd / (1 + Rv*mv/Rd)) + else: + theta = theta_vd + + column_data, _ = domain.coords.get_column_data(theta, domain) + + # Check first column + is_increasing = all(i < j for i, j in zip(column_data[0, :], column_data[0, 1:])) + assert is_increasing, \ + 'static adjustment has not worked: data in column is not increasing' diff --git a/integration-tests/physics/test_suppress_vertical_wind.py b/integration-tests/physics/test_suppress_vertical_wind.py new file mode 100644 index 000000000..250916227 --- /dev/null +++ b/integration-tests/physics/test_suppress_vertical_wind.py @@ -0,0 +1,95 @@ +""" +This tests the physics routine to apply suppress vertical wind in a model's spin +up period. +""" + +from gusto import * +from gusto.labels import physics_label +from firedrake import (Constant, PeriodicIntervalMesh, as_vector, sin, norm, + SpatialCoordinate, ExtrudedMesh, Function, dot) + + +def run_suppress_vertical_wind(dirname): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 100.0 + spin_up_period = 5*dt + + # declare grid shape, with length L and height H + L = 500. + H = 500. + nlayers = 5 + ncolumns = 5 + + # make mesh and domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 0) + + _, z = SpatialCoordinate(mesh) + + # Set up equation + tracers = [WaterVapour()] + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + + # I/O + output = OutputParameters(dirname=dirname+"/static_adjustment", + dumpfreq=1, + dumplist=['theta']) + io = IO(domain, output, diagnostic_fields=[Perturbation('theta')]) + + # Physics scheme + physics_parametrisation = SuppressVerticalWind(eqn, spin_up_period) + + time_discretisation = ForwardEuler(domain) + + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] + + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + # Declare prognostic fields + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Set prognostic variables -- there is initially some vertical wind + theta0.interpolate(Constant(300)) + rho0.interpolate(Constant(1.0)) + u0.project(as_vector([Constant(0.0), sin(pi*z/H)])) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return domain, stepper + + +def test_suppress_vertical_wind(tmpdir): + + dirname = str(tmpdir) + domain, stepper = run_suppress_vertical_wind(dirname) + + u = stepper.fields('u') + vertical_wind = Function(domain.spaces('theta')) + vertical_wind.interpolate(dot(u, domain.k)) + + tol = 1e-12 + assert norm(vertical_wind) < tol diff --git a/integration-tests/physics/test_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py new file mode 100644 index 000000000..09d424af4 --- /dev/null +++ b/integration-tests/physics/test_surface_fluxes.py @@ -0,0 +1,142 @@ +""" +This tests the physics routine to provide surface fluxes. The initial fields are +set to correspond to a different temperature at the surface -- we then check +that afterwards the surface temperature is correct. +""" + +from gusto import * +import gusto.thermodynamics as td +from gusto.labels import physics_label +from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, + SpatialCoordinate, ExtrudedMesh, Function, conditional) +import pytest + + +def run_surface_fluxes(dirname, moist, implicit_formulation): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 100.0 + + # declare grid shape, with length L and height H + L = 500. + H = 500. + nlayers = 5 + ncolumns = 5 + + # make mesh and domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 0) + + _, z = SpatialCoordinate(mesh) + + # Set up equation + tracers = [WaterVapour()] if moist else None + vapour_name = 'water_vapour' if moist else None + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + + # I/O + output = OutputParameters(dirname=dirname+"/surface_fluxes", + dumpfreq=1, + dumplist=['u']) + io = IO(domain, output) + + # Physics scheme + surf_params = BoundaryLayerParameters() + T_surf = Constant(300.0) + physics_parametrisation = SurfaceFluxes(eqn, T_surf, vapour_name, + implicit_formulation, surf_params) + + time_discretisation = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) + + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] + + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + surface_mask = Function(Vt) + surface_mask.interpolate(conditional(z < 100., 0.1, 0.0)) + + # Declare prognostic fields + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Set a background state with constant pressure and temperature + pressure = Function(Vr).interpolate(Constant(100000.)) + temperature = Function(Vt).interpolate(Constant(295.)) + theta_d = td.theta(parameters, temperature, pressure) + mv_sat = td.r_sat(parameters, temperature, pressure) + + # Set prognostic variables + if moist: + water_v0 = stepper.fields("water_vapour") + water_v0.interpolate(0.95*mv_sat) + theta0.project(theta_d*(1 + water_v0 * parameters.R_v / parameters.R_d)) + rho0.interpolate(pressure / (temperature*parameters.R_d * (1 + water_v0 * parameters.R_v / parameters.R_d))) + else: + theta0.project(theta_d) + rho0.interpolate(pressure / (temperature*parameters.R_d)) + + T_true = Function(Vt) + T_true.interpolate(surface_mask*T_surf + (1-surface_mask)*temperature) + + if moist: + mv_true = Function(Vt) + mv_true.interpolate(surface_mask*mv_sat + (1-surface_mask)*water_v0) + else: + mv_true = None + + # Constant horizontal wind + u0.project(as_vector([5.0, 0.0])) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return eqn, stepper, T_true, mv_true + + +@pytest.mark.parametrize("moist", [False, True]) +@pytest.mark.parametrize("implicit_formulation", [False, True]) +def test_surface_fluxes(tmpdir, moist, implicit_formulation): + + dirname = str(tmpdir) + eqn, stepper, T_true, mv_true = run_surface_fluxes(dirname, moist, implicit_formulation) + + # Back out temperature from prognostic fields + theta_vd = stepper.fields('theta') + rho = stepper.fields('rho') + exner = td.exner_pressure(eqn.parameters, rho, theta_vd) + mv = stepper.fields('water_vapour') if moist else None + T_expr = td.T(eqn.parameters, theta_vd, exner, r_v=mv) + + # Project T_expr + T = Function(theta_vd.function_space()) + T.project(T_expr) + denom = norm(T) + assert norm(T - T_true) / denom < 0.001, 'Final temperature is incorrect' + + if moist: + denom = norm(mv_true) + assert norm(mv - mv_true) / denom < 0.01, 'Final water vapour is incorrect' diff --git a/integration-tests/physics/test_wind_drag.py b/integration-tests/physics/test_wind_drag.py new file mode 100644 index 000000000..a00818089 --- /dev/null +++ b/integration-tests/physics/test_wind_drag.py @@ -0,0 +1,122 @@ +""" +This tests the physics routine to apply drag to the wind. +""" + +from gusto import * +import gusto.thermodynamics as td +from gusto.labels import physics_label +from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, dot, + SpatialCoordinate, ExtrudedMesh, Function, conditional) +import pytest + + +def run_wind_drag(dirname, implicit_formulation): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 100.0 + + # declare grid shape, with length L and height H + L = 500. + H = 500. + nlayers = int(H / 5.) + ncolumns = int(L / 5.) + + # make mesh and domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 0) + + _, z = SpatialCoordinate(mesh) + + # Set up equation + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters) + + # I/O + output = OutputParameters(dirname=dirname+"/surface_fluxes", + dumpfreq=1, + dumplist=['u']) + io = IO(domain, output) + + # Physics scheme + surf_params = BoundaryLayerParameters() + physics_parametrisation = WindDrag(eqn, implicit_formulation, surf_params) + + time_discretisation = ForwardEuler(domain) if implicit_formulation else BackwardEuler(domain) + + # time_discretisation = ForwardEuler(domain) + physics_schemes = [(physics_parametrisation, time_discretisation)] + + # Only want time derivatives and physics terms in equation, so drop the rest + eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics_label)), + map_if_true=identity, map_if_false=drop) + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vu = domain.spaces("HDiv") + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + surface_mask = Function(Vt) + surface_mask.interpolate(conditional(z < 100., 1.0, 0.0)) + + # Declare prognostic fields + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # Set a background state with constant pressure and temperature + pressure = Function(Vr).interpolate(Constant(100000.)) + temperature = Function(Vt).interpolate(Constant(295.)) + theta_d = td.theta(parameters, temperature, pressure) + + theta0.project(theta_d) + rho0.interpolate(pressure / (temperature*parameters.R_d)) + + # Constant horizontal wind + u0.project(as_vector([15.0, 0.0])) + + # Answer: slower winds than initially + u_true = Function(Vu) + u_true.project(surface_mask*as_vector([14.53, 0.0]) + (1-surface_mask)*u0) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return mesh, stepper, u_true + + +@pytest.mark.parametrize("implicit_formulation", [False, True]) +def test_wind_drag(tmpdir, implicit_formulation): + + dirname = str(tmpdir) + mesh, stepper, u_true = run_wind_drag(dirname, implicit_formulation) + + u_final = stepper.fields('u') + + # Project into CG1 to get sensible values + e_x = as_vector([1.0, 0.0]) + e_z = as_vector([0.0, 1.0]) + + DG0 = FunctionSpace(mesh, "DG", 0) + u_x_final = Function(DG0).project(dot(u_final, e_x)) + u_x_true = Function(DG0).project(dot(u_true, e_x)) + u_z_final = Function(DG0).project(dot(u_final, e_z)) + u_z_true = Function(DG0).project(dot(u_true, e_z)) + + denom = norm(u_x_true) + assert norm(u_x_final - u_x_true) / denom < 0.01, 'Final horizontal wind is incorrect' + assert norm(u_z_final - u_z_true) < 1e-12, 'Final vertical wind is incorrect' diff --git a/requirements.txt b/requirements.txt index 7d438f194..a8302376a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ netCDF4 +pandas