From 2d7211c9b4fe18922ddbc61ab4dc464d16349d42 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 22 Jan 2023 10:44:48 +0000 Subject: [PATCH 01/28] first implementation of source/sink --- gusto/physics.py | 67 +++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 66 insertions(+), 1 deletion(-) diff --git a/gusto/physics.py b/gusto/physics.py index 11b221a05..e3d0b09ec 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -23,7 +23,7 @@ __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", - "AdvectedMoments", "InstantRain"] + "AdvectedMoments", "InstantRain", "SourceSink"] class Physics(object, metaclass=ABCMeta): @@ -52,6 +52,71 @@ def evaluate(self): pass +class SourceSink(Physics): + """ + 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. + """ + + def __init__(self, equation, variable_name, rate_expression, + method='interpolate'): + """ + Args: + equation (:class: 'PrognosticEquationSet'): the model's equation. + variable_name (str): the name of the variable + rate_expression (ufl.Expr): the saturation function, above which + excess moisture is converted. + method (str, optional): the method to use to evaluate the expression + for the source. Valid options are 'interpolate' or 'project'. + Default is 'interpolate'. + """ + + super().__init__(equation, parameters=None) + + assert method in ['interpolate', 'project'], \ + f'Method {method} for source/sink evaluation not valid' + self.method = method + + # Check the variable exists + assert variable_name in equation.field_names, \ + f'Field {variable_name} does not exist in the equation set' + + # Work out the appropriate function space + V_idx = equation.field_names.index(variable_name) + W = equation.function_space + V = W.sub(V_idx) + + # Make source/sink term + test = equation.tests[V_idx] + self.source = Function(V) + equation.residual += physics(subject(test * self.source * dx, equation.X), + self.evaluate) + + # interpolator does the conversion of vapour to rain + if self.method == 'interpolate': + self.source_interpolator = Interpolator(rate_expression, V) + else: + self.source_projector = Projector(rate_expression, V) + + 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.method == 'interpolate': + self.source.assign(self.source_interpolator.interpolate()) + else: + self.source.assign(self.source_projector.project()) + + class SaturationAdjustment(Physics): """ Represents the phase change between water vapour and cloud liquid. From fc55da7ca2b5fdbcf4a8fee5c0c44c28dc1b7c04 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 10 Aug 2023 17:28:33 +0100 Subject: [PATCH 02/28] have a universal time, threaded through the domain and start a test --- gusto/domain.py | 3 + gusto/physics.py | 46 +++++++-- gusto/timeloop.py | 2 +- integration-tests/physics/test_source_sink.py | 98 +++++++++++++++++++ 4 files changed, 138 insertions(+), 11 deletions(-) create mode 100644 integration-tests/physics/test_source_sink.py diff --git a/gusto/domain.py b/gusto/domain.py index 0528381ad..9f27bbf41 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -59,6 +59,9 @@ def __init__(self, mesh, dt, family, degree=None, else: raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + # Make a placeholder for the time + self.t = Constant(0.0) + # -------------------------------------------------------------------- # # Build compatible function spaces # -------------------------------------------------------------------- # diff --git a/gusto/physics.py b/gusto/physics.py index d59fa56cc..3ab6f2bb1 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -62,16 +62,25 @@ class SourceSink(Physics): 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, - method='interpolate'): + time_varying=False, method='interpolate'): """ Args: - equation (:class: 'PrognosticEquationSet'): the model's equation. + equation (:class:`PrognosticEquationSet`): the model's equation. variable_name (str): the name of the variable - rate_expression (ufl.Expr): the saturation function, above which - excess moisture is converted. + 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'. @@ -82,6 +91,7 @@ def __init__(self, equation, variable_name, rate_expression, assert method in ['interpolate', 'project'], \ f'Method {method} for source/sink evaluation not valid' self.method = method + self.time_varying = time_varying # Check the variable exists assert variable_name in equation.field_names, \ @@ -98,11 +108,24 @@ def __init__(self, equation, variable_name, rate_expression, equation.residual += physics(subject(test * self.source * dx, equation.X), self.evaluate) - # interpolator does the conversion of vapour to rain + # 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(rate_expression, V) + self.source_interpolator = Interpolator(expression, V) else: - self.source_projector = Projector(rate_expression, V) + 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_interpolator.interpolate() + else: + self.source_projector.project() def evaluate(self, x_in, dt): """ @@ -113,10 +136,13 @@ def evaluate(self, x_in, dt): dt: (:class:`Constant`): the timestep, which can be the time interval for the scheme. Unused. """ - if self.method == 'interpolate': - self.source.assign(self.source_interpolator.interpolate()) + if self.time_varying: + if self.method == 'interpolate': + self.source.assign(self.source_interpolator.interpolate()) + else: + self.source.assign(self.source_projector.project()) else: - self.source.assign(self.source_projector.project()) + pass class SaturationAdjustment(Physics): diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 663cc6aae..550ca72dd 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -34,7 +34,7 @@ def __init__(self, equation, io): self.equation = equation self.io = io self.dt = self.equation.domain.dt - self.t = Constant(0.0) + self.t = self.equation.domain.t self.reference_profiles_initialised = False self.setup_fields() diff --git a/integration-tests/physics/test_source_sink.py b/integration-tests/physics/test_source_sink.py new file mode 100644 index 000000000..8952711ca --- /dev/null +++ b/integration-tests/physics/test_source_sink.py @@ -0,0 +1,98 @@ +""" +This tests the source/sink +""" + +from gusto import * +from firedrake import (Constant, PeriodicSquareMesh, SpatialCoordinate, + sqrt, conditional, cos, pi, FunctionSpace) + + +def run_source_sink(dirname): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # set up mesh and domain + L = 10 + nx = 10 + mesh = PeriodicSquareMesh(nx, nx, L, quadrilateral=True) + dt = 0.1 + domain = Domain(mesh, dt, "RTCF", 1) + x, y = SpatialCoordinate(mesh) + + # Source is a Gaussian centred on a point + centre_x = L / 4.0 + + # Equation + eqns = AdvectionEquation(domain, V, "ash") + + # I/O + output = OutputParameters(dirname=dirname+"/instant_rain", + dumpfreq=1, + dumplist=['vapour', "rain"]) + diagnostic_fields = [CourantNumber()] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + transport_method = [DGUpwind(eqns, "water_vapour")] + + # Physics schemes + # define saturation function + saturation = Constant(0.5) + physics_schemes = [(SourceSink(eqns, saturation, rain_name="rain"), + ForwardEuler(domain))] + + # Time stepper + stepper = PrescribedTransport(eqns, SSPRK3(domain), io, transport_method, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + vapour0 = stepper.fields("water_vapour") + + # set up vapour + xc = L/2 + yc = L/2 + rc = L/4 + r = sqrt((x - xc)**2 + (y - yc)**2) + vapour_expr = conditional(r > rc, 0., 1 * (cos(pi * r / (rc * 2))) ** 2) + + vapour0.interpolate(vapour_expr) + + VD = FunctionSpace(mesh, "DG", 1) + initial_vapour = Function(VD).interpolate(vapour_expr) + + # TODO: This test is based on the assumption that final vapour should be + # equal to saturation, which might not be true when timestepping physics. + vapour_true = Function(VD).interpolate(saturation) + rain_true = Function(VD).interpolate(vapour0 - saturation) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=5*dt) + return stepper, saturation, initial_vapour, vapour_true, rain_true + + +def test_instant_rain_setup(tmpdir): + dirname = str(tmpdir) + stepper, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) + v = stepper.fields("water_vapour") + r = stepper.fields("rain") + + # check that the maximum of the vapour field is equal to the saturation + assert v.dat.data.max() - saturation.values() < 0.1, "The maximum of the final vapour field should be equal to saturation" + + # check that the minimum of the vapour field hasn't changed + assert v.dat.data.min() - initial_vapour.dat.data.min() < 0.1, "The minimum of the vapour field should not change" + + # check that the maximum of the excess vapour agrees with the maximum of the + # rain + VD = Function(v.function_space()) + excess_vapour = Function(VD).interpolate(initial_vapour - saturation) + assert excess_vapour.dat.data.max() - r.dat.data.max() < 0.1 + + # check that the minimum of the rain is 0 + assert r.dat.data.min() < 1e-8 From 0eacdca6829befdaf4b65195d0fdd507d0167f07 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 10 Aug 2023 19:37:27 +0100 Subject: [PATCH 03/28] get source/sink and test working --- gusto/physics.py | 24 ++-- integration-tests/physics/test_source_sink.py | 121 ++++++++++-------- 2 files changed, 81 insertions(+), 64 deletions(-) diff --git a/gusto/physics.py b/gusto/physics.py index 3ab6f2bb1..f4455fe2b 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -94,16 +94,24 @@ def __init__(self, equation, variable_name, rate_expression, self.time_varying = time_varying # Check the variable exists - assert variable_name in equation.field_names, \ - f'Field {variable_name} does not exist in the equation set' + 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 - V_idx = equation.field_names.index(variable_name) - W = equation.function_space - V = W.sub(V_idx) + 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 - test = equation.tests[V_idx] self.source = Function(V) equation.residual += physics(subject(test * self.source * dx, equation.X), self.evaluate) @@ -123,9 +131,9 @@ def __init__(self, equation, variable_name, rate_expression, # If not time-varying, evaluate for the first time here if not self.time_varying: if self.method == 'interpolate': - self.source_interpolator.interpolate() + self.source.assign(self.source_interpolator.interpolate()) else: - self.source_projector.project() + self.source.assign(self.source_projector.project()) def evaluate(self, x_in, dt): """ diff --git a/integration-tests/physics/test_source_sink.py b/integration-tests/physics/test_source_sink.py index 8952711ca..0ead67087 100644 --- a/integration-tests/physics/test_source_sink.py +++ b/integration-tests/physics/test_source_sink.py @@ -3,11 +3,11 @@ """ from gusto import * -from firedrake import (Constant, PeriodicSquareMesh, SpatialCoordinate, - sqrt, conditional, cos, pi, FunctionSpace) +from firedrake import (as_vector, PeriodicSquareMesh, SpatialCoordinate, + sqrt, sin, pi, assemble, Constant) +import pytest - -def run_source_sink(dirname): +def run_source_sink(dirname, process, time_varying): # ------------------------------------------------------------------------ # # Set up model objects @@ -18,81 +18,90 @@ def run_source_sink(dirname): 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) - # Source is a Gaussian centred on a point - centre_x = L / 4.0 - # Equation - eqns = AdvectionEquation(domain, V, "ash") + V = domain.spaces('DG') + eqn = AdvectionEquation(domain, V, "ash") # I/O - output = OutputParameters(dirname=dirname+"/instant_rain", - dumpfreq=1, - dumplist=['vapour', "rain"]) + output = OutputParameters(dirname=dirname+"/source_sink", + dumpfreq=1) diagnostic_fields = [CourantNumber()] io = IO(domain, output, diagnostic_fields=diagnostic_fields) - transport_method = [DGUpwind(eqns, "water_vapour")] + 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)) - # Physics schemes - # define saturation function - saturation = Constant(0.5) - physics_schemes = [(SourceSink(eqns, saturation, rain_name="rain"), - ForwardEuler(domain))] + if time_varying: + expression = time_varying_expression + else: + expression = basic_expression + + physics_parametrisations = [SourceSink(eqn, 'ash', expression, time_varying)] # Time stepper - stepper = PrescribedTransport(eqns, SSPRK3(domain), io, transport_method, - physics_schemes=physics_schemes) + stepper = PrescribedTransport(eqn, SSPRK3(domain), io, transport_method, + physics_parametrisations=physics_parametrisations) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # - vapour0 = stepper.fields("water_vapour") - - # set up vapour - xc = L/2 - yc = L/2 - rc = L/4 - r = sqrt((x - xc)**2 + (y - yc)**2) - vapour_expr = conditional(r > rc, 0., 1 * (cos(pi * r / (rc * 2))) ** 2) + ash0 = stepper.fields("ash") - vapour0.interpolate(vapour_expr) + 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) - VD = FunctionSpace(mesh, "DG", 1) - initial_vapour = Function(VD).interpolate(vapour_expr) - - # TODO: This test is based on the assumption that final vapour should be - # equal to saturation, which might not be true when timestepping physics. - vapour_true = Function(VD).interpolate(saturation) - rain_true = Function(VD).interpolate(vapour0 - saturation) + # Constant wind + u = stepper.fields("u") + u.project(as_vector([1.0, 0.0])) # ------------------------------------------------------------------------ # # Run # ------------------------------------------------------------------------ # - stepper.run(t=0, tmax=5*dt) - return stepper, saturation, initial_vapour, vapour_true, rain_true + stepper.run(t=0, tmax=tmax) + return stepper, initial_ash -def test_instant_rain_setup(tmpdir): +@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, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) - v = stepper.fields("water_vapour") - r = stepper.fields("rain") - - # check that the maximum of the vapour field is equal to the saturation - assert v.dat.data.max() - saturation.values() < 0.1, "The maximum of the final vapour field should be equal to saturation" - - # check that the minimum of the vapour field hasn't changed - assert v.dat.data.min() - initial_vapour.dat.data.min() < 0.1, "The minimum of the vapour field should not change" - - # check that the maximum of the excess vapour agrees with the maximum of the - # rain - VD = Function(v.function_space()) - excess_vapour = Function(VD).interpolate(initial_vapour - saturation) - assert excess_vapour.dat.data.max() - r.dat.data.max() < 0.1 - - # check that the minimum of the rain is 0 - assert r.dat.data.min() < 1e-8 + 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" From c386f0e9143355ccd5a213f7117e6761fc74af15 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 14 Aug 2023 10:09:30 +0100 Subject: [PATCH 04/28] add some physics --- gusto/configuration.py | 16 +- gusto/coordinates.py | 118 +++++++++++- gusto/domain.py | 36 +++- gusto/io.py | 15 ++ gusto/physics.py | 177 ++++++++++++++++-- gusto/timeloop.py | 23 ++- .../physics/test_saturation_adjustment.py | 8 +- .../physics/test_surface_fluxes.py | 141 ++++++++++++++ 8 files changed, 508 insertions(+), 26 deletions(-) create mode 100644 integration-tests/physics/test_surface_fluxes.py diff --git a/gusto/configuration.py b/gusto/configuration.py index 4bf3b8b2b..c54b9fd42 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" ] @@ -183,3 +183,17 @@ 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) diff --git a/gusto/coordinates.py b/gusto/coordinates.py index 82e23cc37..9916d24ed 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -4,8 +4,9 @@ """ from gusto.logging import logger -from firedrake import (SpatialCoordinate, sqrt, atan_2, asin, Function) +from firedrake import SpatialCoordinate, Function import numpy as np +import pandas as pd class Coordinates(object): @@ -147,3 +148,118 @@ 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): + """ + Reshapes a field's data into columns. + + Args: + field (:class:`Function`): the field whose data needs sorting. + + 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 + 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) + 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) + 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/domain.py b/gusto/domain.py index 9f27bbf41..1b12927b4 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -8,7 +8,8 @@ from gusto.function_spaces import Spaces, check_degree_args from gusto.perp import perp from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, - inner, grad, VectorFunctionSpace, Function) + inner, grad, VectorFunctionSpace, Function, dot, + FunctionSpace) import numpy as np @@ -122,6 +123,10 @@ 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 # -------------------------------------------------------------------- # @@ -164,3 +169,32 @@ def __init__(self, mesh, dt, family, degree=None, # Send information to other processors self.metadata = comm.bcast(self.metadata, root=0) + + + def set_height_above_surface(self): + """ + Sets a coordinate field which corresponds to height above the domain's + surface. + """ + + x = SpatialCoordinate(self.mesh) + + # Make a height field in CG1 + CG1 = self.spaces('CG1', family='CG', degree=1) + self.coords.register_space(self, 'CG1') + CG1_height = Function(CG1) + CG1_height.interpolate(dot(self.k, x)) + + # Turn height into columnwise data + columnwise_height, index_data = self.coords.get_column_data(CG1_height) + + # 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] + + height_above_surface = Function(CG1) + self.coords.set_field_from_column_data(height_above_surface, + height_above_surface_data, + index_data) + + self.height_above_surface = height_above_surface diff --git a/gusto/io.py b/gusto/io.py index 774f438a0..99594343b 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -238,6 +238,21 @@ def log_parameters(self, equation): logger.info("Physical parameters that take non-default values:") logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(equation.parameters).items())) + def log_courant(self, state_fields): + """ + Logs the maximum Courant number value. + """ + + diagnostic_names = [diagnostic.name for diagnostic in self.diagnostic_fields] + if 'CourantNumber' in diagnostic_names: + courant_idx = diagnostic_names.index('CourantNumber') + courant_diagnostic = self.diagnostic_fields[courant_idx] + courant_diagnostic.compute() + courant_field = state_fields('CourantNumber') + courant_max = self.diagnostics.max(courant_field) + + logger.info(f'Max Courant number: {courant_max}') + def setup_diagnostics(self, state_fields): """ Prepares the I/O for computing the model's global diagnostics and diff --git a/gusto/physics.py b/gusto/physics.py index f4455fe2b..b20bfebd7 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -14,8 +14,9 @@ from gusto.fml import identity, Term, subject from gusto.labels import physics, 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, + TestFunctions) from gusto import thermodynamics import ufl import math @@ -25,7 +26,7 @@ __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", "AdvectedMoments", "InstantRain", "SWSaturationAdjustment", - "SourceSink"] + "SourceSink", "SurfaceFluxes"] class Physics(object, metaclass=ABCMeta): @@ -756,10 +757,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. @@ -901,14 +902,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 @@ -1080,3 +1081,149 @@ 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(Physics): + """ + Prescribed surface temperature and moisture fluxes, to be used in aquaplanet + simulations as Sea Surface Temperature fluxes. + + 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 + # -------------------------------------------------------------------- # + assert isinstance(equation, CompressibleEulerEquations), \ + "Surface fluxes can only be used with Compressible Euler equations" + + if vapour_name is not None: + assert vapour_name in equation.field_names, \ + f"Field {vapour_name} does not exist in the equation set" + + if parameters is None: + parameters = BoundaryLayerParameters() + + super().__init__(equation, parameters=None) + + self.implicit_formulation = implicit_formulation + 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') + if vapour_name is not None: + m_v_idx = equation.field_names.index(vapour_name) + + # For implicit formulation, use internal variables. Otherwise equation's + X = self.X if implicit_formulation else equation.X + tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests + + u = X.subfunctions[u_idx] + theta_vd = X.subfunctions[T_idx] + rho = X.subfunctions[rho_idx] + test_rho = tests[rho_idx] + test_theta = tests[T_idx] + + if vapour_name is not None: + m_v = X.subfunctions[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 + # TODO: should surface be a function (like a mask?) + surface = 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 + + if implicit_formulation: + raise NotImplementedError('implicit formulation not yet implemented') + # Plan here is to first evaluate T and m_v properly + # Then back out what the new value of theta_vd will be from new T + else: + # Construct underlying expressions + kappa = equation.parameters.kappa + dT_dt = surface * 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 * C_E * u_hori_mag * (mv_sat - m_v) / z_a + source_mv_expr = test_m_v * dmv_dt * dx + equation.residual -= physics( + prognostic(subject(source_mv_expr, equation.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) + + dmv_dt * (T / (epsilon * exner) - kappa * theta_vd / (epsilon + m_v))) + + else: + dtheta_vd_dt = dT_dt * (1 / exner - kappa * theta_vd / rho) + + source_theta_expr = test_theta * dtheta_vd_dt * dx + + equation.residual -= physics( + prognostic(subject(source_theta_expr, equation.X), + 'theta'), self.evaluate) + + def evaluate(self, x_in, dt): + """ + Evaluates the source term generated by the physics. + + 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 self.implicit_formulation: + for sourcetor in self.source_interpolators: + source_interpolator.interpolate() diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 550ca72dd..1ea4771c0 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 from pyop2.profiling import timed_stage from gusto.equations import PrognosticEquationSet from gusto.fml import drop, Label, Term @@ -171,6 +171,8 @@ def run(self, t, tmax, pick_up=False): self.x.update() + self.io.log_courant(self.fields) + self.timestep() self.t.assign(self.t + self.dt) @@ -323,8 +325,8 @@ def __init__(self, equation, scheme, io, spatial_methods=None, 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" + # assert isinstance(phys_scheme, ExplicitTimeDiscretisation), \ + # "Only explicit time discretisations can be used for physics" apply_bcs = False phys_scheme.setup(equation, apply_bcs, physics) @@ -655,6 +657,21 @@ def setup_fields(self): self.fields = StateFields(self.x, self.equation.prescribed_fields, *self.io.output.dumplist) + def run(self, t, tmax, pick_up=False): + """ + Runs the model for the specified time, from t to tmax + + Args: + t (float): the start time of the run + tmax (float): the end time of the run + pick_up: (bool): specify whether to pick_up from a previous run + """ + # It's best to have evaluated the velocity before we start + if self.velocity_projection is not None: + self.velocity_projection.project() + + super().run(t, tmax, pick_up=pick_up) + def timestep(self): if self.velocity_projection is not None: self.velocity_projection.project() diff --git a/integration-tests/physics/test_saturation_adjustment.py b/integration-tests/physics/test_saturation_adjustment.py index 0b92d6f11..6218d49f7 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_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py new file mode 100644 index 000000000..bd10db04c --- /dev/null +++ b/integration-tests/physics/test_surface_fluxes.py @@ -0,0 +1,141 @@ +""" +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 os import path +from gusto import * +import gusto.thermodynamics as td +from gusto.labels import physics +from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, + SpatialCoordinate, ExtrudedMesh, Function, conditional) +from netCDF4 import Dataset +import pytest + + +def run_surface_fluxes(dirname, moist, implicit_formulation): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 2.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 + 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) + physics_schemes = [(physics_parametrisation, time_discretisation)] + + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vt = domain.spaces("theta", degree=1) + Vr = domain.spaces("DG", "DG", degree=1) + + 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(290.)) + 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.9*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])) + + # 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)), + map_if_true=identity, map_if_false=drop) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=dt) + + return eqn, stepper, T_true, mv_true + + +@pytest.mark.parametrize("moist", [True, False]) +@pytest.mark.parametrize("implicit_formulation", [True, False]) +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.01, f'Final temperature is incorrect' + + if moist: + denom = norm(mv_true) + assert norm(mv - mv_true) / denom < 0.01, f'Final water vapour is incorrect' From 504c57455f0448c52bb7136443508ec983181a60 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 22 Aug 2023 18:53:24 +0100 Subject: [PATCH 05/28] implement wind drag --- gusto/physics.py | 196 +++++++++++++++--- gusto/time_discretisation.py | 7 + .../physics/test_surface_fluxes.py | 31 +-- integration-tests/physics/test_wind_drag.py | 123 +++++++++++ 4 files changed, 319 insertions(+), 38 deletions(-) create mode 100644 integration-tests/physics/test_wind_drag.py diff --git a/gusto/physics.py b/gusto/physics.py index b20bfebd7..971a6cbb3 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -16,7 +16,7 @@ from gusto.logging import logger from firedrake import (Interpolator, conditional, Function, dx, sqrt, dot, min_value, max_value, Constant, pi, Projector, - TestFunctions) + TestFunctions, FunctionSpace, Function, split, inner) from gusto import thermodynamics import ufl import math @@ -26,7 +26,7 @@ __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", "AdvectedMoments", "InstantRain", "SWSaturationAdjustment", - "SourceSink", "SurfaceFluxes"] + "SourceSink", "SurfaceFluxes", "WindDrag"] class Physics(object, metaclass=ABCMeta): @@ -1086,7 +1086,8 @@ def evaluate(self, x_in, dt): class SurfaceFluxes(Physics): """ Prescribed surface temperature and moisture fluxes, to be used in aquaplanet - simulations as Sea Surface Temperature fluxes. + 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. @@ -1129,6 +1130,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, self.implicit_formulation = implicit_formulation self.X = Function(equation.X.function_space()) + self.dt = equation.domain.dt # -------------------------------------------------------------------- # # Extract prognostic variables @@ -1139,18 +1141,16 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, if vapour_name is not None: m_v_idx = equation.field_names.index(vapour_name) - # For implicit formulation, use internal variables. Otherwise equation's - X = self.X if implicit_formulation else equation.X + X = self.X tests = TestFunctions(X.function_space()) if implicit_formulation else equation.tests - u = X.subfunctions[u_idx] - theta_vd = X.subfunctions[T_idx] - rho = X.subfunctions[rho_idx] - test_rho = tests[rho_idx] + 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 = X.subfunctions[m_v_idx] + m_v = split(X)[m_v_idx] test_m_v = tests[m_v_idx] else: m_v = None @@ -1174,8 +1174,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, # -------------------------------------------------------------------- # z = equation.domain.height_above_surface z_a = parameters.height_surface_layer - # TODO: should surface be a function (like a mask?) - surface = conditional(z < z_a, Constant(1.0), Constant(0.0)) + 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)) @@ -1183,40 +1182,76 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, 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: - raise NotImplementedError('implicit formulation not yet implemented') - # Plan here is to first evaluate T and m_v properly - # Then back out what the new value of theta_vd will be from new T + 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 -= physics(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 -= physics(subject(prognostic(source_theta_expr, 'theta'), + X), self.evaluate) + + # General formulation ------------------------------------------------ # else: # Construct underlying expressions kappa = equation.parameters.kappa - dT_dt = surface * C_H * u_hori_mag * (T_surface_expr - T) / z_a + 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 * C_E * u_hori_mag * (mv_sat - m_v) / z_a + 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 -= physics( - prognostic(subject(source_mv_expr, equation.X), + 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) + 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) + dtheta_vd_dt = dT_dt * (1 / exner - kappa * theta_vd / (rho * T)) source_theta_expr = test_theta * dtheta_vd_dt * dx equation.residual -= physics( - prognostic(subject(source_theta_expr, equation.X), - 'theta'), self.evaluate) + subject(prognostic(source_theta_expr, 'theta'), X), + self.evaluate) def evaluate(self, x_in, dt): """ - Evaluates the source term generated by the physics. + 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. @@ -1225,5 +1260,116 @@ def evaluate(self, x_in, dt): """ if self.implicit_formulation: - for sourcetor in self.source_interpolators: + 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(Physics): + """ + 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 + # -------------------------------------------------------------------- # + assert isinstance(equation, CompressibleEulerEquations), \ + "Surface fluxes can only be used with Compressible Euler equations" + + if parameters is None: + parameters = BoundaryLayerParameters() + + super().__init__(equation, parameters=None) + + k = equation.domain.k + self.implicit_formulation = implicit_formulation + self.X = Function(equation.X.function_space()) + self.dt = equation.domain.dt + + # -------------------------------------------------------------------- # + # 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 + self.source_projector = Projector(du_expr, source_u) + + source_expr = inner(test, source_u - k*dot(source_u, k)) * dx + equation.residual -= physics(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 + + source_expr = inner(test, du_dt) * dx + + equation.residual -= physics(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. + """ + + if self.implicit_formulation: + self.X.assign(x_in) + self.dt.assign(dt) + self.source_projector.project() diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 98e8d53f5..334cff929 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -650,6 +650,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/integration-tests/physics/test_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py index bd10db04c..eeeaafe21 100644 --- a/integration-tests/physics/test_surface_fluxes.py +++ b/integration-tests/physics/test_surface_fluxes.py @@ -20,7 +20,7 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): # Set up model objects # ------------------------------------------------------------------------ # - dt = 2.0 + dt = 100.0 # declare grid shape, with length L and height H L = 500. @@ -52,9 +52,17 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): 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)), + map_if_true=identity, map_if_false=drop) + + # Time stepper scheme = ForwardEuler(domain) stepper = SplitPhysicsTimestepper(eqn, scheme, io, @@ -64,11 +72,11 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): # Initial conditions # ------------------------------------------------------------------------ # - Vt = domain.spaces("theta", degree=1) - Vr = domain.spaces("DG", "DG", degree=1) + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") surface_mask = Function(Vt) - surface_mask.interpolate(conditional(z < 100., 1.0, 0.0)) + surface_mask.interpolate(conditional(z < 100., 0.1, 0.0)) # Declare prognostic fields u0 = stepper.fields("u") @@ -77,14 +85,14 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): # Set a background state with constant pressure and temperature pressure = Function(Vr).interpolate(Constant(100000.)) - temperature = Function(Vt).interpolate(Constant(290.)) + 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.9*mv_sat) + 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: @@ -103,10 +111,6 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): # Constant horizontal wind u0.project(as_vector([5.0, 0.0])) - # 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)), - map_if_true=identity, map_if_false=drop) - # ------------------------------------------------------------------------ # # Run # ------------------------------------------------------------------------ # @@ -116,8 +120,8 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): return eqn, stepper, T_true, mv_true -@pytest.mark.parametrize("moist", [True, False]) -@pytest.mark.parametrize("implicit_formulation", [True, False]) +@pytest.mark.parametrize("moist", [False, True]) +@pytest.mark.parametrize("implicit_formulation", [False, True]) def test_surface_fluxes(tmpdir, moist, implicit_formulation): dirname = str(tmpdir) @@ -134,7 +138,8 @@ def test_surface_fluxes(tmpdir, moist, implicit_formulation): T = Function(theta_vd.function_space()) T.project(T_expr) denom = norm(T) - assert norm(T - T_true) / denom < 0.01, f'Final temperature is incorrect' + assert norm(T - T_true) / denom < 0.001, f'Final temperature is incorrect' + if moist: denom = norm(mv_true) diff --git a/integration-tests/physics/test_wind_drag.py b/integration-tests/physics/test_wind_drag.py new file mode 100644 index 000000000..202347f60 --- /dev/null +++ b/integration-tests/physics/test_wind_drag.py @@ -0,0 +1,123 @@ +""" +This tests the physics routine to apply drag to the wind. +""" + +from gusto import * +import gusto.thermodynamics as td +from gusto.labels import physics +from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, + SpatialCoordinate, ExtrudedMesh, Function, conditional) +from netCDF4 import Dataset +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)), + 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, f'Final horizontal wind is incorrect' + assert norm(u_z_final - u_z_true) < 1e-12, f'Final vertical wind is incorrect' From 2e6c9882a2e6c02032b2df497efd514e860516b3 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 28 Aug 2023 08:36:59 +0100 Subject: [PATCH 06/28] lint fixes --- gusto/coord_transforms.py | 2 +- gusto/coordinates.py | 5 ++--- gusto/domain.py | 8 +++++++- gusto/physics.py | 6 +++--- integration-tests/physics/test_source_sink.py | 1 + integration-tests/physics/test_surface_fluxes.py | 8 ++------ integration-tests/physics/test_wind_drag.py | 5 ++--- 7 files changed, 18 insertions(+), 17 deletions(-) diff --git a/gusto/coord_transforms.py b/gusto/coord_transforms.py index a8caa62f4..0299770d4 100644 --- a/gusto/coord_transforms.py +++ b/gusto/coord_transforms.py @@ -5,7 +5,7 @@ import importlib import numpy as np -from firedrake import Constant, as_vector, SpatialCoordinate +from firedrake import as_vector, SpatialCoordinate import ufl __all__ = ["xyz_from_lonlatr", "lonlatr_from_xyz", "xyz_vector_from_lonlatr", diff --git a/gusto/coordinates.py b/gusto/coordinates.py index 0c22ed23d..6e921c37f 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -156,7 +156,6 @@ def register_space(self, domain, space_name): (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): """ Reshapes a field's data into columns. @@ -210,7 +209,7 @@ def get_column_data(self, field): 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))] + & (np.isclose(data['Y'], first_Y))] else: data = data.sort_values(by=['X', 'Z']) @@ -267,6 +266,6 @@ def set_field_from_column_data(self, field, columnwise_data, index_data): _, num_levels = np.shape(columnwise_data) for lev_idx in range(num_levels): - field.dat.data[index_data[:,lev_idx]] = columnwise_data[:,lev_idx] + field.dat.data[index_data[:, lev_idx]] = columnwise_data[:, lev_idx] return field diff --git a/gusto/domain.py b/gusto/domain.py index c8ee0ba36..0d0e0f9eb 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -125,8 +125,11 @@ def __init__(self, mesh, dt, family, degree=None, if self.on_sphere: spherical_shell_mesh = mesh._base_mesh if hasattr(mesh, "_base_mesh") else mesh + xyz_shell = SpatialCoordinate(spherical_shell_mesh) + r_shell = sqrt(inner(xyz_shell, xyz_shell)) CG1 = FunctionSpace(spherical_shell_mesh, "CG", 1) - radius_field = Function(CG1).interpolate(R) + radius_field = Function(CG1) + radius_field.interpolate(r_shell) # TODO: this should use global min kernel radius = np.min(radius_field.dat.data_ro) else: @@ -158,6 +161,8 @@ def set_height_above_surface(self): surface. """ + from firedrake import dot + x = SpatialCoordinate(self.mesh) # Make a height field in CG1 @@ -180,6 +185,7 @@ def set_height_above_surface(self): self.height_above_surface = height_above_surface + def construct_domain_metadata(mesh, coords, on_sphere): """ Builds a dictionary containing metadata about the domain. diff --git a/gusto/physics.py b/gusto/physics.py index 971a6cbb3..109de0eda 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -9,6 +9,7 @@ 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 @@ -16,7 +17,7 @@ from gusto.logging import logger from firedrake import (Interpolator, conditional, Function, dx, sqrt, dot, min_value, max_value, Constant, pi, Projector, - TestFunctions, FunctionSpace, Function, split, inner) + TestFunctions, split, inner) from gusto import thermodynamics import ufl import math @@ -1245,8 +1246,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, source_theta_expr = test_theta * dtheta_vd_dt * dx equation.residual -= physics( - subject(prognostic(source_theta_expr, 'theta'), X), - self.evaluate) + subject(prognostic(source_theta_expr, 'theta'), X), self.evaluate) def evaluate(self, x_in, dt): """ diff --git a/integration-tests/physics/test_source_sink.py b/integration-tests/physics/test_source_sink.py index 0ead67087..4df01599e 100644 --- a/integration-tests/physics/test_source_sink.py +++ b/integration-tests/physics/test_source_sink.py @@ -7,6 +7,7 @@ sqrt, sin, pi, assemble, Constant) import pytest + def run_source_sink(dirname, process, time_varying): # ------------------------------------------------------------------------ # diff --git a/integration-tests/physics/test_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py index eeeaafe21..372c757fc 100644 --- a/integration-tests/physics/test_surface_fluxes.py +++ b/integration-tests/physics/test_surface_fluxes.py @@ -4,13 +4,11 @@ that afterwards the surface temperature is correct. """ -from os import path from gusto import * import gusto.thermodynamics as td from gusto.labels import physics from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, SpatialCoordinate, ExtrudedMesh, Function, conditional) -from netCDF4 import Dataset import pytest @@ -62,7 +60,6 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): eqn.residual = eqn.residual.label_map(lambda t: any(t.has_label(time_derivative, physics)), map_if_true=identity, map_if_false=drop) - # Time stepper scheme = ForwardEuler(domain) stepper = SplitPhysicsTimestepper(eqn, scheme, io, @@ -138,9 +135,8 @@ def test_surface_fluxes(tmpdir, moist, implicit_formulation): T = Function(theta_vd.function_space()) T.project(T_expr) denom = norm(T) - assert norm(T - T_true) / denom < 0.001, f'Final temperature is incorrect' - + 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, f'Final water vapour is incorrect' + 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 index 202347f60..18d857e3d 100644 --- a/integration-tests/physics/test_wind_drag.py +++ b/integration-tests/physics/test_wind_drag.py @@ -7,7 +7,6 @@ from gusto.labels import physics from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, SpatialCoordinate, ExtrudedMesh, Function, conditional) -from netCDF4 import Dataset import pytest @@ -119,5 +118,5 @@ def test_wind_drag(tmpdir, implicit_formulation): 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, f'Final horizontal wind is incorrect' - assert norm(u_z_final - u_z_true) < 1e-12, f'Final vertical wind is incorrect' + 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' From 5a87322581e2946bd371842ae33a96b6caa97b8a Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 29 Aug 2023 09:28:19 +0100 Subject: [PATCH 07/28] add brunt-vaisala diagnostic and some domain neatenings --- gusto/diagnostics.py | 46 +++++++++++++++++++++++++++++++++++++++++++- gusto/domain.py | 31 ++++++++++++++++++----------- gusto/physics.py | 2 +- 3 files changed, 66 insertions(+), 13 deletions(-) diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index 6dd7fa465..2023deb12 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -24,7 +24,8 @@ "Perturbation", "Theta_e", "InternalEnergy", "PotentialEnergy", "ThermodynamicKineticEnergy", "Dewpoint", "Temperature", "Theta_d", "RelativeHumidity", "Pressure", "Exner_Vt", "HydrostaticImbalance", "Precipitation", - "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence"] + "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence", + "BruntVaisalaFrequencySquared"] class Diagnostics(object): @@ -897,6 +898,49 @@ 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 0d0e0f9eb..214c38186 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -131,7 +131,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 @@ -170,18 +170,27 @@ def set_height_above_surface(self): 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) - - # 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] - - height_above_surface = Function(CG1) - self.coords.set_field_from_column_data(height_above_surface, - height_above_surface_data, - index_data) + try: + columnwise_height, index_data = self.coords.get_column_data(CG1_height) + + # 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) + except ModuleNotFoundError: + # If no pandas, then don't take orography into account + logger.warning( + 'Pandas not found, so using height as height above surface. ' + + 'This will be incorrect in the presence of orography') + + if self.on_sphere: + height_above_surface.assign(CG1_height - self.radius) self.height_above_surface = height_above_surface diff --git a/gusto/physics.py b/gusto/physics.py index 109de0eda..9359efb15 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -640,7 +640,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) From a447ad222c2307fc5fb20f4b205130bf631841fa Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 31 Aug 2023 09:52:15 +0100 Subject: [PATCH 08/28] introduce labels which add extra labels, which I use to give unique labels to physics schemes --- gusto/fml/form_manipulation_language.py | 1 + gusto/labels.py | 85 ++++++++- gusto/physics.py | 166 ++++++++++-------- gusto/time_discretisation.py | 7 +- gusto/timeloop.py | 19 +- .../physics/test_surface_fluxes.py | 4 +- .../physics/test_sw_saturation_adjustment.py | 2 +- integration-tests/physics/test_wind_drag.py | 6 +- 8 files changed, 198 insertions(+), 92 deletions(-) diff --git a/gusto/fml/form_manipulation_language.py b/gusto/fml/form_manipulation_language.py index 062fddb04..b9dc608f1 100644 --- a/gusto/fml/form_manipulation_language.py +++ b/gusto/fml/form_manipulation_language.py @@ -169,6 +169,7 @@ def __init__(self, *terms): self.terms = terms[0].terms else: if any([type(term) is not Term for term in list(terms)]): + import pdb; pdb.set_trace() raise TypeError('Can only pass terms or a LabelledForm to LabelledForm') self.terms = list(terms) diff --git a/gusto/labels.py b/gusto/labels.py index 74e458729..ec46009db 100644 --- a/gusto/labels.py +++ b/gusto/labels.py @@ -6,6 +6,86 @@ from gusto.fml.form_manipulation_language import Term, Label, LabelledForm from types import MethodType +dynamics_label = Label("dynamics", validator=lambda value: type(value) is str) +physics_label = Label("physics", validator=lambda value: type(value) is str) + + +class DynamicsLabel(Label): + """A label for a fluid dynamics term.""" + def __call__(self, target, value=None): + """ + Applies the label to a form or term, and additionally labels the term as + a dynamics term. + + Args: + target (:class:`ufl.Form`, :class:`Term` or :class:`LabelledForm`): + the form, term or labelled form to be labelled. + value (..., optional): the value to attach to this label. Defaults + to None. + + Returns: + :class:`Term` or :class:`LabelledForm`: a :class:`Term` is returned + if the target is a :class:`Term`, otherwise a + :class:`LabelledForm` is returned. + """ + + new_target = dynamics_label(target, self.label) + + if isinstance(new_target, LabelledForm): + # Need to be very careful in using super().__call__ method as the + # underlying __call__ method calls itself to act upon multiple terms + # in a LabelledForm. We can avoid this by special handling of the + # LabelledForm case + labelled_terms = (Label.__call__(self, t, value) for t in new_target.terms) + return LabelledForm(*labelled_terms) + else: + super().__call__(new_target, value) + + + +class PhysicsLabel(Label): + """A label for a physics parametrisation term.""" + def __init__(self, label, *, value=True, validator=lambda value: type(value) == MethodType): + """ + Args: + label (str): the name of the label. + value (..., optional): the value for the label to take. Can be any + type (subject to the validator). Defaults to True. + validator (func, optional): function to check the validity of any + value later passed to the label. Defaults to None. + """ + super().__init__(label, value=value, validator=validator) + + def __call__(self, target, value=None): + """ + Applies the label to a form or term, and additionally labels the term as + a physics term. + + Args: + target (:class:`ufl.Form`, :class:`Term` or :class:`LabelledForm`): + the form, term or labelled form to be labelled. + value (..., optional): the value to attach to this label. Defaults + to None. + + Returns: + :class:`Term` or :class:`LabelledForm`: a :class:`Term` is returned + if the target is a :class:`Term`, otherwise a + :class:`LabelledForm` is returned. + """ + + new_target = physics_label(target, self.label) + + if isinstance(new_target, LabelledForm): + # Need to be very careful in using super().__call__ method as the + # underlying __call__ method calls itself to act upon multiple terms + # in a LabelledForm. We can avoid this by special handling of the + # LabelledForm case + labelled_terms = (Label.__call__(self, t, value) for t in new_target.terms) + return LabelledForm(*labelled_terms) + else: + super().__call__(new_target, value) + + # ---------------------------------------------------------------------------- # # Common Labels # ---------------------------------------------------------------------------- # @@ -13,11 +93,10 @@ time_derivative = Label("time_derivative") transport = Label("transport", validator=lambda value: type(value) == TransportEquationType) diffusion = Label("diffusion") -physics = Label("physics", validator=lambda value: type(value) == MethodType) transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor]) prognostic = Label("prognostic", validator=lambda value: type(value) == str) -pressure_gradient = Label("pressure_gradient") -coriolis = Label("coriolis") +pressure_gradient = DynamicsLabel("pressure_gradient") +coriolis = DynamicsLabel("coriolis") linearisation = Label("linearisation", validator=lambda value: type(value) in [LabelledForm, Term]) ibp_label = Label("ibp", validator=lambda value: type(value) == IntegrateByParts) hydrostatic = Label("hydrostatic", validator=lambda value: type(value) in [LabelledForm, Term]) diff --git a/gusto/physics.py b/gusto/physics.py index 9359efb15..1a9e5cf95 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. @@ -13,7 +13,7 @@ from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations from gusto.fml import identity, Term, subject -from gusto.labels import physics, transporting_velocity, transport, prognostic +from gusto.labels import PhysicsLabel, transporting_velocity, transport, prognostic from gusto.logging import logger from firedrake import (Interpolator, conditional, Function, dx, sqrt, dot, min_value, max_value, Constant, pi, Projector, @@ -30,10 +30,10 @@ "SourceSink", "SurfaceFluxes", "WindDrag"] -class Physics(object, metaclass=ABCMeta): +class PhysicsParametrisation(object, metaclass=ABCMeta): """Base class for the parametrisation of physical processes for Gusto.""" - def __init__(self, equation, parameters=None): + def __init__(self, equation, label_name, parameters=None): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. @@ -42,6 +42,7 @@ def __init__(self, equation, parameters=None): parameters are obtained from the equation. """ + self.label = PhysicsLabel(label_name) self.equation = equation if parameters is None and hasattr(equation, 'parameters'): self.parameters = equation.parameters @@ -56,7 +57,7 @@ def evaluate(self): pass -class SourceSink(Physics): +class SourceSink(PhysicsParametrisation): """ The source or sink of some variable, described through a UFL expression. @@ -88,12 +89,14 @@ def __init__(self, equation, variable_name, rate_expression, Default is 'interpolate'. """ - super().__init__(equation, parameters=None) + label_name = f'source_sink_{variable_name}' + super().__init__(equation, label_name, parameters=None) - assert method in ['interpolate', 'project'], \ - f'Method {method} for source/sink evaluation not valid' + 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"): @@ -115,8 +118,8 @@ def __init__(self, equation, variable_name, rate_expression, # Make source/sink term self.source = Function(V) - equation.residual += physics(subject(test * self.source * dx, equation.X), - self.evaluate) + 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: @@ -155,7 +158,7 @@ def evaluate(self, x_in, dt): pass -class SaturationAdjustment(Physics): +class SaturationAdjustment(PhysicsParametrisation): """ Represents the phase change between water vapour and cloud liquid. @@ -191,7 +194,8 @@ def __init__(self, equation, vapour_name='water_vapour', CompressibleEulerEquations. """ - super().__init__(equation, parameters=parameters) + label_name = 'saturation_adjustment' + super().__init__(equation, label_name, parameters=parameters) # TODO: make a check on the variable type of the active tracers # if not a mixing ratio, we need to convert to mixing ratios @@ -199,8 +203,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 @@ -304,8 +310,8 @@ def __init__(self, equation, vapour_name='water_vapour', # Add source terms to residual for test, source in zip(tests, self.source): - equation.residual += physics(subject(test * source * dx, - equation.X), self.evaluate) + equation.residual += self.label(subject(test * source * dx, + equation.X), self.evaluate) def evaluate(self, x_in, dt): """ @@ -338,7 +344,7 @@ class AdvectedMoments(Enum): M3 = 1 # advect the mass of the distribution -class Fallout(Physics): +class Fallout(PhysicsParametrisation): """ Represents the fallout process of hydrometeors. @@ -373,8 +379,13 @@ def __init__(self, equation, rain_name, domain, transport_method, representing the number of moments of the size distribution of raindrops to be transported. Defaults to `AdvectedMoments.M3`. """ + + label_name = f'fallout_{rain_name}' + 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) @@ -410,7 +421,7 @@ def __init__(self, equation, rain_name, domain, transport_method, adv_term = transport.remove(adv_term) adv_term = prognostic(subject(adv_term, equation.X), rain_name) - equation.residual += physics(adv_term, self.evaluate) + equation.residual += self.label(adv_term, self.evaluate) # -------------------------------------------------------------------- # # Expressions for determining rainfall velocity @@ -472,7 +483,7 @@ def evaluate(self, x_in, dt): self.determine_v.project() -class Coalescence(Physics): +class Coalescence(PhysicsParametrisation): """ Represents the coalescence of cloud droplets to form rain droplets. @@ -500,9 +511,19 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', accumulation (bool, optional): whether to include the accumulation process in the parametrisation. Defaults to True. """ + + label_name = "coalescence" + if self.accretion: + label_name += "_accretion" + if self.accumulation: + label_name += "_accumulation" + 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) @@ -550,10 +571,10 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', # Add term to equation's residual test_cl = equation.tests[self.cloud_idx] test_r = equation.tests[self.rain_idx] - equation.residual += physics(subject(test_cl * self.source * dx - - test_r * self.source * dx, - equation.X), - self.evaluate) + equation.residual += self.label(subject(test_cl * self.source * dx + - test_r * self.source * dx, + equation.X), + self.evaluate) def evaluate(self, x_in, dt): """ @@ -571,7 +592,7 @@ def evaluate(self, x_in, dt): self.source.assign(self.source_interpolator.interpolate()) -class EvaporationOfRain(Physics): +class EvaporationOfRain(PhysicsParametrisation): """ Represents the evaporation of rain into water vapour. @@ -585,7 +606,7 @@ class EvaporationOfRain(Physics): """ 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. @@ -595,16 +616,14 @@ 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. """ - super().__init__(equation, parameters=parameters) + label_name = 'evaporation_of_rain' + 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 @@ -612,8 +631,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()) @@ -719,8 +740,8 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', # Add source terms to residual for test, source in zip(tests, self.source): - equation.residual += physics(subject(test * source * dx, - equation.X), self.evaluate) + equation.residual += self.label(subject(test * source * dx, + equation.X), self.evaluate) def evaluate(self, x_in, dt): """ @@ -740,7 +761,7 @@ def evaluate(self, x_in, dt): interpolator.interpolate() -class InstantRain(Physics): +class InstantRain(PhysicsParametrisation): """ The process of converting vapour above the saturation curve to rain. @@ -784,7 +805,8 @@ def __init__(self, equation, saturation_curve, parameters are obtained from the equation. """ - super().__init__(equation, parameters=None) + label_name = 'instant_rain' + super().__init__(equation, label_name, parameters=None) self.convective_feedback = convective_feedback self.time_varying_saturation = time_varying_saturation @@ -838,25 +860,24 @@ def __init__(self, equation, saturation_curve, self.saturation_curve = saturation_curve # lose proportion of vapour above the saturation curve - equation.residual += physics(subject(test_v * self.source * dx, - equation.X), - self.evaluate) + equation.residual += self.label(subject(test_v * self.source * dx, + equation.X), + self.evaluate) # if rain is not none then the excess vapour is being tracked and is # added to rain if rain_name is not None: Vr_idx = equation.field_names.index(rain_name) test_r = equation.tests[Vr_idx] - equation.residual -= physics(subject(test_r * self.source * dx, - equation.X), - self.evaluate) + equation.residual -= self.label(subject(test_r * self.source * dx, + equation.X), + self.evaluate) # if feeding back on the height adjust the height equation if convective_feedback: - equation.residual += physics(subject - (test_D * beta1 * self.source * dx, - equation.X), - self.evaluate) + equation.residual += self.label( + subject(test_D * beta1 * self.source * dx, equation.X), + self.evaluate) # interpolator does the conversion of vapour to rain self.source_interpolator = Interpolator(conditional( @@ -886,7 +907,7 @@ def evaluate(self, x_in, dt): self.source.assign(self.source_interpolator.interpolate()) -class SWSaturationAdjustment(Physics): +class SWSaturationAdjustment(PhysicsParametrisation): """ Represents the process of adjusting water vapour and cloud water according to a saturation function, via condensation and evaporation processes. @@ -950,7 +971,8 @@ def __init__(self, equation, saturation_curve, L_v=None, parameters are obtained from the equation. """ - super().__init__(equation, parameters=None) + label_name = 'saturation_adjustment' + super().__init__(equation, label_name, parameters=None) self.time_varying_saturation = time_varying_saturation self.convective_feedback = convective_feedback @@ -1052,8 +1074,8 @@ def __init__(self, equation, saturation_curve, L_v=None, # Add source terms to residual for test, source in zip(tests, self.source): - equation.residual += physics(subject(test * source * dx, - equation.X), self.evaluate) + equation.residual += self.label(subject(test * source * dx, + equation.X), self.evaluate) def evaluate(self, x_in, dt): """ @@ -1084,7 +1106,7 @@ def evaluate(self, x_in, dt): interpolator.interpolate() -class SurfaceFluxes(Physics): +class SurfaceFluxes(PhysicsParametrisation): """ Prescribed surface temperature and moisture fluxes, to be used in aquaplanet simulations as Sea Surface Temperature fluxes. This formulation is taken @@ -1117,17 +1139,18 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, # -------------------------------------------------------------------- # # Check arguments and generic initialisation # -------------------------------------------------------------------- # - assert isinstance(equation, CompressibleEulerEquations), \ - "Surface fluxes can only be used with Compressible Euler equations" + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Surface fluxes can only be used with Compressible Euler equations") if vapour_name is not None: - assert vapour_name in equation.field_names, \ - f"Field {vapour_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 parameters is None: parameters = BoundaryLayerParameters() - super().__init__(equation, parameters=None) + label_name = 'surface_flux' + super().__init__(equation, label_name, parameters=None) self.implicit_formulation = implicit_formulation self.X = Function(equation.X.function_space()) @@ -1203,8 +1226,8 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, source_mv_expr = test_m_v * source_mv * dx self.source_interpolators.append(Interpolator(dmv_expr, source_mv)) - equation.residual -= physics(subject(prognostic(source_mv_expr, vapour_name), - X), self.evaluate) + 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? @@ -1219,8 +1242,8 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, 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 -= physics(subject(prognostic(source_theta_expr, 'theta'), - X), self.evaluate) + equation.residual -= self.label(subject(prognostic(source_theta_expr, 'theta'), + X), self.evaluate) # General formulation ------------------------------------------------ # else: @@ -1232,7 +1255,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=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 -= physics( + equation.residual -= self.label( prognostic(subject(source_mv_expr, X), vapour_name), self.evaluate) @@ -1245,7 +1268,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, source_theta_expr = test_theta * dtheta_vd_dt * dx - equation.residual -= physics( + equation.residual -= self.label( subject(prognostic(source_theta_expr, 'theta'), X), self.evaluate) def evaluate(self, x_in, dt): @@ -1267,7 +1290,7 @@ def evaluate(self, x_in, dt): source_interpolator.interpolate() -class WindDrag(Physics): +class WindDrag(PhysicsParametrisation): """ A simple surface wind drag scheme. This formulation is taken from the DCMIP (2016) test case document. @@ -1293,13 +1316,14 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): # -------------------------------------------------------------------- # # Check arguments and generic initialisation # -------------------------------------------------------------------- # - assert isinstance(equation, CompressibleEulerEquations), \ - "Surface fluxes can only be used with Compressible Euler equations" + if not isinstance(equation, CompressibleEulerEquations): + raise ValueError("Surface fluxes can only be used with Compressible Euler equations") if parameters is None: parameters = BoundaryLayerParameters() - super().__init__(equation, parameters=None) + label_name = 'wind_drag' + super().__init__(equation, label_name, parameters=None) k = equation.domain.k self.implicit_formulation = implicit_formulation @@ -1346,8 +1370,8 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): self.source_projector = Projector(du_expr, source_u) source_expr = inner(test, source_u - k*dot(source_u, k)) * dx - equation.residual -= physics(subject(prognostic(source_expr, 'u'), - X), self.evaluate) + equation.residual -= self.label(subject(prognostic(source_expr, 'u'), + X), self.evaluate) # General formulation ------------------------------------------------ # else: @@ -1356,7 +1380,7 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): source_expr = inner(test, du_dt) * dx - equation.residual -= physics(subject(prognostic(source_expr, 'u'), X), self.evaluate) + equation.residual -= self.label(subject(prognostic(source_expr, 'u'), X), self.evaluate) def evaluate(self, x_in, dt): """ diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 334cff929..21ac83522 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -18,7 +18,7 @@ from gusto.fml import ( replace_subject, replace_test_function, Term, all_terms, drop ) -from gusto.labels import time_derivative, prognostic, physics +from gusto.labels import time_derivative, prognostic, physics_label from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * @@ -137,8 +137,9 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.evaluate_source = [] for t in self.residual: - if t.has_label(physics): - self.evaluate_source.append(t.get(physics)) + if t.has_label(physics_label): + physics_name = t.get(physics_label) + self.evaluate_source.append(t.labels[physics_name]) # -------------------------------------------------------------------- # # Set up Wrappers diff --git a/gusto/timeloop.py b/gusto/timeloop.py index b9137043c..912f53e3f 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -9,7 +9,7 @@ from gusto.forcing import Forcing from gusto.labels import ( transport, diffusion, time_derivative, linearisation, prognostic, - physics, transporting_velocity + physics_label, transporting_velocity ) from gusto.linear_solvers import LinearTimesteppingSolver from gusto.logging import logger @@ -249,9 +249,10 @@ def __init__(self, equation, scheme, io, spatial_methods=None, in which case the terms follow the original discretisation in the equation. physics_parametrisations: (iter, optional): an iterable of - :class:`Physics` objects that describe physical parametrisations to be included - to add to the equation. They can only be used when the time - discretisation `scheme` is explicit. Defaults to None. + :class:`PhysicsParametrisation` objects that describe physical + parametrisations to be included to add to the equation. They can + only be used when the time discretisation `scheme` is explicit. + Defaults to None. """ self.scheme = scheme if spatial_methods is not None: @@ -333,12 +334,12 @@ def __init__(self, equation, scheme, io, spatial_methods=None, else: self.physics_schemes = [] - for _, phys_scheme in self.physics_schemes: + for parametrisation, 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" apply_bcs = False - phys_scheme.setup(equation, apply_bcs, physics) + phys_scheme.setup(equation, apply_bcs, parametrisation.label) @property def transporting_velocity(self): @@ -348,7 +349,7 @@ def setup_scheme(self): self.setup_equation(self.equation) # Go through and label all non-physics terms with a "dynamics" label dynamics = Label('dynamics') - self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics)), dynamics) + self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics_label)), dynamics) apply_bcs = True self.scheme.setup(self.equation, apply_bcs, dynamics) self.setup_transporting_velocity(self.scheme) @@ -528,9 +529,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.physics_schemes: + for parametrisation, scheme in self.physics_schemes: apply_bcs = True - scheme.setup(self.equation, apply_bcs, physics) + scheme.setup(self.equation, apply_bcs, parametrisation.label) def copy_active_tracers(self, x_in, x_out): """ diff --git a/integration-tests/physics/test_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py index 372c757fc..d0506c2aa 100644 --- a/integration-tests/physics/test_surface_fluxes.py +++ b/integration-tests/physics/test_surface_fluxes.py @@ -6,7 +6,7 @@ from gusto import * import gusto.thermodynamics as td -from gusto.labels import physics +from gusto.labels import physics_label from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, SpatialCoordinate, ExtrudedMesh, Function, conditional) import pytest @@ -57,7 +57,7 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): 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)), + 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 diff --git a/integration-tests/physics/test_sw_saturation_adjustment.py b/integration-tests/physics/test_sw_saturation_adjustment.py index 43abbe269..453e77fe9 100644 --- a/integration-tests/physics/test_sw_saturation_adjustment.py +++ b/integration-tests/physics/test_sw_saturation_adjustment.py @@ -38,7 +38,7 @@ def run_sw_cond_evap(dirname, process): degree = 1 domain = Domain(mesh, dt, 'BDM', degree) x = SpatialCoordinate(mesh) - theta, lamda = latlon_coords(mesh) + lamda, theta, _ = lonlatr_from_xyz(x[0], x[1], x[2]) # saturation field (constant everywhere) sat = 100 diff --git a/integration-tests/physics/test_wind_drag.py b/integration-tests/physics/test_wind_drag.py index 18d857e3d..a00818089 100644 --- a/integration-tests/physics/test_wind_drag.py +++ b/integration-tests/physics/test_wind_drag.py @@ -4,8 +4,8 @@ from gusto import * import gusto.thermodynamics as td -from gusto.labels import physics -from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, +from gusto.labels import physics_label +from firedrake import (norm, Constant, PeriodicIntervalMesh, as_vector, dot, SpatialCoordinate, ExtrudedMesh, Function, conditional) import pytest @@ -51,7 +51,7 @@ def run_wind_drag(dirname, implicit_formulation): 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)), + 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 From 683c248cfe123b68cf92ec5746083155c228b06e Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 31 Aug 2023 10:12:36 +0100 Subject: [PATCH 09/28] fix lint and move import pandas statement --- gusto/coordinates.py | 3 ++- gusto/diagnostics.py | 5 ++++- gusto/domain.py | 5 +++-- gusto/fml/form_manipulation_language.py | 1 - gusto/labels.py | 1 - gusto/physics.py | 4 ++-- 6 files changed, 11 insertions(+), 8 deletions(-) diff --git a/gusto/coordinates.py b/gusto/coordinates.py index 6e921c37f..b0cfbae91 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -7,7 +7,6 @@ from gusto.logging import logger from firedrake import (SpatialCoordinate, Function) import numpy as np -import pandas as pd class Coordinates(object): @@ -169,6 +168,8 @@ def get_column_data(self, field): ordered column data. """ + import pandas as pd + space_name = field.function_space().name coords = self.chi_coords[space_name] diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index ca7e2c44b..824b91bbe 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -235,6 +235,7 @@ def __call__(self): class CourantNumber(DiagnosticField): """Dimensionless Courant number diagnostic field.""" name = "CourantNumber" + def __init__(self, velocity='u', component='whole', name=None, to_dump=True, space=None, method='interpolate', required_fields=()): """ @@ -337,7 +338,7 @@ def setup(self, domain, state_fields): self.cell_flux_form = 2*avg(un*test)*dS_calc + un*test*ds_calc # Final Courant number expression - self.expr = self.cell_flux *domain.dt / cell_volume + self.expr = self.cell_flux * domain.dt / cell_volume super().setup(domain, state_fields) @@ -942,6 +943,7 @@ def setup(self, 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: @@ -982,6 +984,7 @@ def setup(self, domain, state_fields): 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 214c38186..eb4ea3aaf 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -4,6 +4,7 @@ model's time interval. """ +from gusto.logging import logger from gusto.coordinates import Coordinates from gusto.function_spaces import Spaces, check_degree_args from gusto.perp import perp @@ -181,8 +182,8 @@ def set_height_above_surface(self): 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) + height_above_surface_data, + index_data) except ModuleNotFoundError: # If no pandas, then don't take orography into account logger.warning( diff --git a/gusto/fml/form_manipulation_language.py b/gusto/fml/form_manipulation_language.py index b9dc608f1..062fddb04 100644 --- a/gusto/fml/form_manipulation_language.py +++ b/gusto/fml/form_manipulation_language.py @@ -169,7 +169,6 @@ def __init__(self, *terms): self.terms = terms[0].terms else: if any([type(term) is not Term for term in list(terms)]): - import pdb; pdb.set_trace() raise TypeError('Can only pass terms or a LabelledForm to LabelledForm') self.terms = list(terms) diff --git a/gusto/labels.py b/gusto/labels.py index ec46009db..7538de5a2 100644 --- a/gusto/labels.py +++ b/gusto/labels.py @@ -42,7 +42,6 @@ def __call__(self, target, value=None): super().__call__(new_target, value) - class PhysicsLabel(Label): """A label for a physics parametrisation term.""" def __init__(self, label, *, value=True, validator=lambda value: type(value) == MethodType): diff --git a/gusto/physics.py b/gusto/physics.py index 1a9e5cf95..62887eec9 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -574,7 +574,7 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', equation.residual += self.label(subject(test_cl * self.source * dx - test_r * self.source * dx, equation.X), - self.evaluate) + self.evaluate) def evaluate(self, x_in, dt): """ @@ -877,7 +877,7 @@ def __init__(self, equation, saturation_curve, if convective_feedback: equation.residual += self.label( subject(test_D * beta1 * self.source * dx, equation.X), - self.evaluate) + self.evaluate) # interpolator does the conversion of vapour to rain self.source_interpolator = Interpolator(conditional( From dd61f3628d8a00f069021afa573d0cdee9793cac Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 1 Sep 2023 10:29:14 +0100 Subject: [PATCH 10/28] separate physics into fast/slow/final and add lots of logging statements --- gusto/diagnostics.py | 3 + gusto/linear_solvers.py | 5 + gusto/physics.py | 36 ++++++- gusto/timeloop.py | 97 ++++++++++++++----- .../equations/test_thermal_sw.py | 2 +- 5 files changed, 115 insertions(+), 28 deletions(-) diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index 824b91bbe..4373c2a7d 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", @@ -217,6 +218,8 @@ def setup(self, domain, state_fields, space=None): def compute(self): """Compute the diagnostic field from the current state.""" + logger.info(f'Computing diagnostic {self.name} with {self.method} method') + if self.method == 'interpolate': self.evaluator.interpolate() elif self.method == 'assign': diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 68ddf1477..20851106c 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -371,12 +371,15 @@ def solve(self, xrhs, dy): # TODO: can we avoid computing these each time the solver is called? with timed_region("Gusto:HybridProjectRhobar"): + logger.info('Linear solver: Rho-averaging pre-solve') self.rho_avg_solver.solve() with timed_region("Gusto:HybridProjectExnerbar"): + logger.info('Linear solver: Exner-averaging pre-solve') self.exner_avg_solver.solve() # Solve the hybridized system + logger.info('Linear solver: Compressible mixed hybrid solve') self.hybridized_solver.solve() broken_u, rho1, _ = self.urhol0.subfunctions @@ -386,6 +389,7 @@ def solve(self, xrhs, dy): u1.assign(0.0) with timed_region("Gusto:HybridProjectHDiv"): + logger.info('Linear solver: restore u continuity') self._average_kernel.apply(u1, self._weight, broken_u) # Reapply bcs to ensure they are satisfied @@ -399,6 +403,7 @@ def solve(self, xrhs, dy): # Reconstruct theta with timed_region("Gusto:ThetaRecon"): + logger.info('Linear solver: back-substitute theta') self.theta_solver.solve() # Copy into theta cpt of dy diff --git a/gusto/physics.py b/gusto/physics.py index 62887eec9..4e05f501e 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -17,7 +17,8 @@ from gusto.logging import logger from firedrake import (Interpolator, conditional, Function, dx, sqrt, dot, min_value, max_value, Constant, pi, Projector, - TestFunctions, split, inner) + TestFunctions, split, inner, TestFunction, + NonlinearVariationalProblem, NonlinearVariationalSolver) from gusto import thermodynamics import ufl import math @@ -150,6 +151,7 @@ def evaluate(self, x_in, dt): 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: @@ -321,6 +323,7 @@ def evaluate(self, x_in, dt): x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') # Update the values of internal variables self.dt.assign(dt) self.X.assign(x_in) @@ -468,7 +471,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): """ @@ -478,9 +486,10 @@ def evaluate(self, x_in, dt): x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. """ + 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): @@ -584,6 +593,8 @@ def evaluate(self, x_in, dt): x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + # Update the values of internal variables self.dt.assign(dt) self.rain.assign(x_in.subfunctions[self.rain_idx]) @@ -751,6 +762,8 @@ def evaluate(self, x_in, dt): x_in (:class:`Function`): the (mixed) field to be evolved. dt (:class:`Constant`): the time interval for the scheme. """ + logger.info(f'Evaluating physics parametrisation {self.label.label}') + # Update the values of internal variables self.dt.assign(dt) self.X.assign(x_in) @@ -897,6 +910,8 @@ def evaluate(self, x_in, dt): 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.convective_feedback: self.D.assign(x_in.subfunctions[self.VD_idx]) if self.time_varying_saturation: @@ -1089,6 +1104,7 @@ def evaluate(self, x_in, dt): 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.convective_feedback: self.D.assign(x_in.split()[self.VD_idx]) @@ -1282,6 +1298,8 @@ def evaluate(self, x_in, dt): 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) @@ -1367,7 +1385,13 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): 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 - self.source_projector = Projector(du_expr, source_u) + + # 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'), @@ -1393,7 +1417,9 @@ def evaluate(self, x_in, dt): 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.project() + self.source_projector.solve() diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 912f53e3f..0977b087a 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, 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 @@ -375,7 +375,9 @@ class SemiImplicitQuasiNewton(BaseTimestepper): def __init__(self, equation_set, io, transport_schemes, spatial_methods, auxiliary_equations_and_schemes=None, linear_solver=None, - diffusion_schemes=None, physics_schemes=None, **kwargs): + diffusion_schemes=None, physics_schemes=None, + slow_physics_schemes=None, fast_physics_schemes=None, + alpha=Constant(0.5), num_outer=4, num_inner=1): """ Args: @@ -398,26 +400,47 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, physics_schemes: (list, optional): a list of tuples of the form (:class:`Physics`, :class:`TimeDiscretisation`), pairing physics parametrisations and timestepping schemes to use for each. - Timestepping schemes for physics must be explicit. Defaults to - None. + Timestepping schemes for physics must be explicit. These schemes + are all evaluated at the end of the time step. Defaults to None. + slow_physics_schemes: (list, optional): a list of tuples of the form + (:class:`Physics`, :class:`TimeDiscretisation`), pairing physics + parametrisations and timestepping schemes to use for each. + Timestepping schemes for physics must be explicit. These schemes + are all evaluated at the end of the time step. Defaults to None. + fast_physics_schemes: (list, optional): a list of tuples of the form + (:class:`Physics`, :class:`TimeDiscretisation`), pairing physics + parametrisations and timestepping schemes to use for each. + Timestepping schemes for physics must be explicit. These schemes + are all evaluated at the end of the time step. Defaults to None. + alpha (float, optional): the :kwargs: maxk is the number of outer iterations, maxi is the number of inner iterations and alpha is the offcentering parameter """ - self.maxk = kwargs.pop("maxk", 4) - self.maxi = kwargs.pop("maxi", 1) - self.alpha = kwargs.pop("alpha", 0.5) - if kwargs: - raise ValueError("unexpected kwargs: %s" % list(kwargs.keys())) + self.num_outer = num_outer + self.num_inner = num_inner + self.alpha = alpha self.spatial_methods = spatial_methods if physics_schemes is not None: - self.physics_schemes = physics_schemes + self.final_physics_schemes = physics_schemes else: - self.physics_schemes = [] - for _, scheme in self.physics_schemes: + self.final_physics_schemes = [] + if slow_physics_schemes is not None: + self.slow_physics_schemes = slow_physics_schemes + else: + self.slow_physics_schemes = [] + if fast_physics_schemes is not None: + self.fast_physics_schemes = fast_physics_schemes + else: + self.fast_physics_schemes = [] + self.all_physics_schemes = (self.slow_physics_schemes + + self.fast_physics_schemes + + self.final_physics_schemes) + + for _, 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" @@ -479,6 +502,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.field_name = equation_set.field_name W = equation_set.function_space self.xrhs = Function(W) + self.xrhs_phys = Function(W) self.dy = Function(W) if linear_solver is None: self.linear_solver = LinearTimesteppingSolver(equation_set, self.alpha) @@ -507,7 +531,7 @@ def transporting_velocity(self): def setup_fields(self): """Sets up time levels n, star, p and np1""" self.x = TimeLevelFields(self.equation, 1) - self.x.add_fields(self.equation, levels=("star", "p")) + self.x.add_fields(self.equation, levels=("star", "p", "after_slow", "after_fast")) for aux_eqn, _ in self.auxiliary_equations_and_schemes: self.x.add_fields(aux_eqn) # Prescribed fields for auxiliary eqns should come from prognostics of @@ -529,7 +553,7 @@ def setup_scheme(self): apply_bcs = True for _, scheme in self.diffusion_schemes: scheme.setup(self.equation, apply_bcs, diffusion) - for parametrisation, scheme in self.physics_schemes: + for parametrisation, scheme in self.all_physics_schemes: apply_bcs = True scheme.setup(self.equation, apply_bcs, parametrisation.label) @@ -553,40 +577,68 @@ def timestep(self): xnp1 = self.x.np1 xstar = self.x.star xp = self.x.p + x_after_slow = self.x.after_slow + x_after_fast = self.x.after_fast + xrhs = self.xrhs + xrhs_phys = self.xrhs_phys dy = self.dy + + x_after_slow(self.field_name).assign(xn(self.field_name)) + if len(self.slow_physics_schemes) > 0: + with timed_stage("Slow physics"): + logger.info('SIQN: Slow physics') + for _, scheme in self.slow_physics_schemes: + scheme.apply(x_after_slow(scheme.field_name), x_after_slow(scheme.field_name)) + with timed_stage("Apply forcing terms"): - self.forcing.apply(xn, xn, xstar(self.field_name), "explicit") + 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") xp(self.field_name).assign(xstar(self.field_name)) - for k in range(self.maxk): + for outer in range(self.num_outer): with timed_stage("Transport"): self.io.log_courant(self.fields, 'transporting_velocity', - message=f'transporting velocity, outer iteration {k}') + message=f'transporting velocity, outer iteration {outer}') for name, scheme in self.active_transport: + logger.info(f'SIQN: Transport {outer}: {name}') # transports a field from xstar and puts result in xp scheme.apply(xp(name), xstar(name)) + 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(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)) + xrhs.assign(0.) # xrhs is the residual which goes in the linear solve + xrhs_phys.assign(x_after_fast(self.field_name) - xp(self.field_name)) - for i in range(self.maxi): + for inner in range(self.num_inner): 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) + xrhs += xrhs_phys with timed_stage("Implicit solve"): + logger.info(f'SIQN: Mixed solve {(outer, inner)}') self.linear_solver.solve(xrhs, dy) # solves linear system and places result in dy xnp1X = xnp1(self.field_name) 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() @@ -598,9 +650,10 @@ def timestep(self): for name, scheme in self.diffusion_schemes: scheme.apply(xnp1(name), xnp1(name)) - with timed_stage("Physics"): - for _, scheme in self.physics_schemes: - scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name)) + if len(self.final_physics_schemes) > 0: + with timed_stage("Final Physics"): + for _, scheme in self.final_physics_schemes: + scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name)) def run(self, t, tmax, pick_up=False): """ diff --git a/integration-tests/equations/test_thermal_sw.py b/integration-tests/equations/test_thermal_sw.py index 6d7e1ef08..bbb1b29d9 100644 --- a/integration-tests/equations/test_thermal_sw.py +++ b/integration-tests/equations/test_thermal_sw.py @@ -81,7 +81,7 @@ def set_up_initial_conditions(domain, equation, stepper): bexpr = g * (1 - theta) - u0.project(uexpr) + u0.project(as_vector(uexpr)) D0.interpolate(Dexpr) b0.interpolate(bexpr) From 22f6cdec63157296f9fec3a08154e304d0642193 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 1 Sep 2023 13:23:54 +0000 Subject: [PATCH 11/28] fix bug in coalescence --- gusto/physics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gusto/physics.py b/gusto/physics.py index 4e05f501e..cd77a6c55 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -522,9 +522,9 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', """ label_name = "coalescence" - if self.accretion: + if accretion: label_name += "_accretion" - if self.accumulation: + if accumulation: label_name += "_accumulation" super().__init__(equation, label_name, parameters=None) From d85ac4c7ace03d7d55966bc8e182e07056b27537 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 2 Sep 2023 18:44:10 +0100 Subject: [PATCH 12/28] tiny fix to prevent all ranks trying to make directory --- gusto/logging.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gusto/logging.py b/gusto/logging.py index dc116275b..c6eef1e5e 100644 --- a/gusto/logging.py +++ b/gusto/logging.py @@ -186,7 +186,6 @@ def update_logfile_location(new_path, comm): fh.close() logger.removeHandler(fh) - os.makedirs(new_path, exist_ok=True) if parallel_log in ["FILE", "BOTH"]: # If all ranks are logging wait here in case a directory is being created comm.Barrier() From 9224e512aa88f2251d4ac8dec9a625b40086511c Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 3 Sep 2023 17:01:59 +0100 Subject: [PATCH 13/28] fix checkpointing test, avoid multiple physics evaluations and more logging statements --- gusto/io.py | 14 ++++++++---- gusto/time_discretisation.py | 5 ++++- gusto/timeloop.py | 10 ++++++--- integration-tests/model/test_checkpointing.py | 22 +++++++++---------- 4 files changed, 32 insertions(+), 19 deletions(-) diff --git a/gusto/io.py b/gusto/io.py index 69ab3475d..a08d38738 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -484,7 +484,7 @@ def setup_dump(self, state_fields, t, pick_up=False): # dump initial fields if not pick_up: - self.dump(state_fields, t) + self.dump(state_fields, t, step=1) def pick_up_from_checkpoint(self, state_fields): """ @@ -555,8 +555,9 @@ def pick_up_from_checkpoint(self, state_fields): except AttributeError: initial_steps = None - # Finally pick up time + # Finally pick up time and step number t = chk.read_attribute("/", "time") + step = chk.read_attribute("/", "step") else: with CheckpointFile(chkfile, 'r') as chk: @@ -586,6 +587,7 @@ def pick_up_from_checkpoint(self, state_fields): # Finally pick up time t = chk.get_attr("/", "time") + step = chk.get_attr("/", "step") # If we have picked up from a non-standard file, reset this name # so that we will checkpoint using normal file name from now on @@ -593,9 +595,9 @@ def pick_up_from_checkpoint(self, state_fields): else: raise ValueError("Must set checkpoint True if picking up") - return t, reference_profiles, initial_steps + return t, reference_profiles, step, initial_steps - def dump(self, state_fields, t, initial_steps=None): + def dump(self, state_fields, t, step, initial_steps=None): """ Dumps all of the required model output. @@ -606,6 +608,7 @@ def dump(self, state_fields, t, initial_steps=None): Args: state_fields (:class:`StateFields`): the model's field container. t (float): the simulation's current time. + step (int): the number of time steps. initial_steps (int, optional): the number of initial time steps completed by a multi-level time scheme. Defaults to None. """ @@ -630,6 +633,7 @@ def dump(self, state_fields, t, initial_steps=None): for field_name in self.to_pick_up: self.chkpt.store(state_fields(field_name), name=field_name) self.chkpt.write_attribute("/", "time", t) + self.chkpt.write_attribute("/", "step", step) if initial_steps is not None: self.chkpt.write_attribute("/", "initial_steps", initial_steps) else: @@ -638,8 +642,10 @@ def dump(self, state_fields, t, initial_steps=None): for field_name in self.to_pick_up: chk.save_function(state_fields(field_name), name=field_name) chk.set_attr("/", "time", t) + chk.set_attr("/", "step", step) if initial_steps is not None: chk.set_attr("/", "initial_steps", initial_steps) + logger.info(f'Just checkpointed, count {self.chkptcount}') if (next(self.dumpcount) % output.dumpfreq) == 0: if output.dump_nc: diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index eb88fa279..e8007c765 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -135,10 +135,13 @@ def setup(self, equation, apply_bcs=True, *active_labels): map_if_false=drop) self.evaluate_source = [] + self.physics_names = [] for t in self.residual: if t.has_label(physics_label): physics_name = t.get(physics_label) - self.evaluate_source.append(t.labels[physics_name]) + if t.labels[physics_name] not in self.physics_names: + self.evaluate_source.append(t.labels[physics_name]) + self.physics_names.append(t.labels[physics_name]) # -------------------------------------------------------------------- # # Set up Wrappers diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 0977b087a..1c70b34c8 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -162,9 +162,11 @@ def run(self, t, tmax, pick_up=False): if pick_up: # Pick up fields, and return other info to be picked up - t, reference_profiles, initial_timesteps = self.io.pick_up_from_checkpoint(self.fields) + t, reference_profiles, self.step, initial_timesteps = self.io.pick_up_from_checkpoint(self.fields) self.set_reference_profiles(reference_profiles) self.set_initial_timesteps(initial_timesteps) + else: + self.step = 1 # Set up dump, which may also include an initial dump with timed_stage("Dump output"): @@ -174,7 +176,8 @@ def run(self, t, tmax, pick_up=False): # Time loop while float(self.t) < tmax - 0.5*float(self.dt): - logger.info(f'at start of timestep, t={float(self.t)}, dt={float(self.dt)}') + logger.info('*'*80) + logger.info(f'at start of timestep {self.step}, t={float(self.t)}, dt={float(self.dt)}') self.x.update() @@ -186,9 +189,10 @@ def run(self, t, tmax, pick_up=False): self.timestep() self.t.assign(self.t + self.dt) + self.step += 1 with timed_stage("Dump output"): - self.io.dump(self.fields, float(self.t), self.get_initial_timesteps()) + self.io.dump(self.fields, float(self.t), self.step, self.get_initial_timesteps()) if self.io.output.checkpoint and self.io.output.checkpoint_method == 'dumbcheckpoint': self.io.chkpt.close() diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index f8a39b69c..d53a17a0f 100644 --- a/integration-tests/model/test_checkpointing.py +++ b/integration-tests/model/test_checkpointing.py @@ -100,7 +100,7 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): # Set up mesh nlayers = 5 # horizontal layers - columns = 15 # number of columns + columns = 5 # number of columns L = 3.e5 m = PeriodicIntervalMesh(columns, L) dt = 0.2 @@ -175,16 +175,6 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): assert error < 5e-16, \ f'Checkpointed and picked up field {field_name} is not equal' - # ------------------------------------------------------------------------ # - # Pick up from checkpoint and run *same* timestepper for 2 more time steps - # ------------------------------------------------------------------------ # - - # Wipe fields from second time stepper - if checkpoint_method == 'dumbcheckpoint': - # Get an error when picking up fields with the same stepper with new method - initialise_fields(eqns_2, stepper_2) - stepper_2.run(t=2*dt, tmax=4*dt, pick_up=True) - # ------------------------------------------------------------------------ # # Run *new* timestepper for 2 time steps # ------------------------------------------------------------------------ # @@ -201,6 +191,16 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): stepper_3, _ = set_up_model_objects(mesh, dt, output_3, stepper_type) stepper_3.run(t=2*dt, tmax=4*dt, pick_up=True) + # ------------------------------------------------------------------------ # + # Pick up from checkpoint and run *same* timestepper for 2 more time steps + # ------------------------------------------------------------------------ # + # Done after stepper_3, as we don't want to checkpoint again + # Wipe fields from second time stepper + if checkpoint_method == 'dumbcheckpoint': + # Get an error when picking up fields with the same stepper with new method + initialise_fields(eqns_2, stepper_2) + stepper_2.run(t=2*dt, tmax=4*dt, pick_up=True) + # ------------------------------------------------------------------------ # # Compare fields against saved values for run without checkpointing # ------------------------------------------------------------------------ # From f62ad5f2f89528dfaa1c5b752fb705aba28f831d Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 5 Sep 2023 09:16:38 +0100 Subject: [PATCH 14/28] prevent extra evaluation of physics --- gusto/time_discretisation.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index e8007c765..68e42f962 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -334,8 +334,6 @@ def apply(self, x_out, x_in): """ self.x0.assign(x_in) for i in range(self.ncycles): - for evaluate in self.evaluate_source: - evaluate(x_in, self.dt) self.apply_cycle(self.x1, self.x0) self.x0.assign(self.x1) x_out.assign(self.x1) From fa437d4bca66cef1e2f5dfb81955d4e8a3dbd152 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 6 Sep 2023 09:48:01 +0100 Subject: [PATCH 15/28] add static adjustment routine and allow time discretisation to work with no RHS --- gusto/io.py | 1 - gusto/physics.py | 110 ++++++++++++++++-- gusto/time_discretisation.py | 14 ++- gusto/timeloop.py | 16 ++- .../physics/test_static_adjustment.py | 109 +++++++++++++++++ .../physics/test_surface_fluxes.py | 4 +- 6 files changed, 238 insertions(+), 16 deletions(-) create mode 100644 integration-tests/physics/test_static_adjustment.py diff --git a/gusto/io.py b/gusto/io.py index a08d38738..4ccb31144 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -645,7 +645,6 @@ def dump(self, state_fields, t, step, initial_steps=None): chk.set_attr("/", "step", step) if initial_steps is not None: chk.set_attr("/", "initial_steps", initial_steps) - logger.info(f'Just checkpointed, count {self.chkptcount}') if (next(self.dumpcount) % output.dumpfreq) == 0: if output.dump_nc: diff --git a/gusto/physics.py b/gusto/physics.py index cd77a6c55..0f0e4f3b0 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -28,7 +28,7 @@ __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", "AdvectedMoments", "InstantRain", "SWSaturationAdjustment", - "SourceSink", "SurfaceFluxes", "WindDrag"] + "SourceSink", "SurfaceFluxes", "WindDrag", "StaticAdjustment"] class PhysicsParametrisation(object, metaclass=ABCMeta): @@ -259,7 +259,7 @@ def __init__(self, equation, vapour_name='water_vapour', liquid_water += self.X.subfunctions[liq_idx] # define some parameters as attributes - self.dt = Constant(0.0) + self.dt = Constant(1.0) R_d = parameters.R_d cp = parameters.cp cv = parameters.cv @@ -546,7 +546,7 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', self.source = Function(Vt) # define some parameters as attributes - self.dt = Constant(0.0) + self.dt = Constant(1.0) # TODO: should these parameters be hard-coded or configurable? k_1 = Constant(0.001) # accretion rate in 1/s k_2 = Constant(2.2) # accumulation rate in 1/s @@ -692,7 +692,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', liquid_water += self.X.subfunctions[liq_idx] # define some parameters as attributes - self.dt = Constant(0.0) + self.dt = Constant(1.0) R_d = parameters.R_d cp = parameters.cp cv = parameters.cv @@ -1170,7 +1170,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, self.implicit_formulation = implicit_formulation self.X = Function(equation.X.function_space()) - self.dt = equation.domain.dt + self.dt = Constant(1.0) # -------------------------------------------------------------------- # # Extract prognostic variables @@ -1335,7 +1335,7 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): # Check arguments and generic initialisation # -------------------------------------------------------------------- # if not isinstance(equation, CompressibleEulerEquations): - raise ValueError("Surface fluxes can only be used with Compressible Euler equations") + raise ValueError("Wind drag can only be used with Compressible Euler equations") if parameters is None: parameters = BoundaryLayerParameters() @@ -1346,7 +1346,7 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): k = equation.domain.k self.implicit_formulation = implicit_formulation self.X = Function(equation.X.function_space()) - self.dt = equation.domain.dt + self.dt = Constant(1.0) # -------------------------------------------------------------------- # # Extract prognostic variables @@ -1423,3 +1423,99 @@ def evaluate(self, x_in, dt): 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". + """ + + 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(1.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 + # -------------------------------------------------------------------- # + + self.get_column_data = partial(equation.domain.coords.get_column_data, field=self.theta_to_sort) + self.set_column_data = equation.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() diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 68e42f962..a6ff88194 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -7,7 +7,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import (Function, TestFunction, NonlinearVariationalProblem, - NonlinearVariationalSolver, DirichletBC) + NonlinearVariationalSolver, DirichletBC, Constant) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property @@ -419,6 +419,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): diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 1c70b34c8..9727f8810 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -140,6 +140,14 @@ def setup_transporting_velocity(self, scheme): scheme.residual = transporting_velocity.update_value(scheme.residual, uadv) + def log_timestep(self): + """ + Logs the start of a time step. + """ + logger.info('') + logger.info('*'*80) + logger.info(f'at start of timestep {self.step}, t={float(self.t)}, dt={float(self.dt)}') + def run(self, t, tmax, pick_up=False): """ Runs the model for the specified time, from t to tmax @@ -176,8 +184,7 @@ def run(self, t, tmax, pick_up=False): # Time loop while float(self.t) < tmax - 0.5*float(self.dt): - logger.info('*'*80) - logger.info(f'at start of timestep {self.step}, t={float(self.t)}, dt={float(self.dt)}') + self.log_timestep() self.x.update() @@ -588,7 +595,6 @@ def timestep(self): xrhs_phys = self.xrhs_phys dy = self.dy - x_after_slow(self.field_name).assign(xn(self.field_name)) if len(self.slow_physics_schemes) > 0: with timed_stage("Slow physics"): @@ -598,8 +604,8 @@ 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 + # 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") xp(self.field_name).assign(xstar(self.field_name)) diff --git a/integration-tests/physics/test_static_adjustment.py b/integration-tests/physics/test_static_adjustment.py new file mode 100644 index 000000000..404d96047 --- /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) + + # 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_surface_fluxes.py b/integration-tests/physics/test_surface_fluxes.py index d0506c2aa..09d424af4 100644 --- a/integration-tests/physics/test_surface_fluxes.py +++ b/integration-tests/physics/test_surface_fluxes.py @@ -23,8 +23,8 @@ def run_surface_fluxes(dirname, moist, implicit_formulation): # declare grid shape, with length L and height H L = 500. H = 500. - nlayers = int(H / 5.) - ncolumns = int(L / 5.) + nlayers = 5 + ncolumns = 5 # make mesh and domain m = PeriodicIntervalMesh(ncolumns, L) From 1f5cd227af72556c6891084f8dba3e263a18694b Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 6 Sep 2023 22:24:25 +0100 Subject: [PATCH 16/28] add a scheme to provide a spin-up period in which vertical winds are suppressed --- gusto/physics.py | 79 ++++++++++++++++++++++++++++++++++++++++++++++- gusto/timeloop.py | 18 ++++++++--- 2 files changed, 91 insertions(+), 6 deletions(-) diff --git a/gusto/physics.py b/gusto/physics.py index 0f0e4f3b0..a9d630fc1 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -28,7 +28,8 @@ __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", "AdvectedMoments", "InstantRain", "SWSaturationAdjustment", - "SourceSink", "SurfaceFluxes", "WindDrag", "StaticAdjustment"] + "SourceSink", "SurfaceFluxes", "WindDrag", "StaticAdjustment", + "SuppressVerticalWind"] class PhysicsParametrisation(object, metaclass=ABCMeta): @@ -1519,3 +1520,79 @@ def evaluate(self, x_in, dt): 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. + """ + + # -------------------------------------------------------------------- # + # 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(1.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 diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 9727f8810..c49cd99a9 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -423,11 +423,19 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, parametrisations and timestepping schemes to use for each. Timestepping schemes for physics must be explicit. These schemes are all evaluated at the end of the time step. Defaults to None. - alpha (float, optional): the - - :kwargs: maxk is the number of outer iterations, maxi is the number - of inner iterations and alpha is the offcentering parameter - """ + alpha (`ufl.Constant`, optional): the semi-implicit off-centering + parameter. A value of 1 corresponds to fully implicit, while 0 + corresponds to fully explicit. Defaults to Constant(0.5). + num_outer (int, optional): number of outer iterations in the semi- + implicit algorithm. The outer loop includes transport and any + fast physics schemes. Defaults to 4. Note that default used by + the Met Office's ENDGame and GungHo models is 2. + num_inner (int, optional): number of inner iterations in the semi- + implicit algorithm. The inner loop includes the evaluation of + implicit forcing (pressure gradient and Coriolis) terms, and the + linear solve. Defaults to 1. Note that default used by the Met + Office's ENDGame and GungHo models is 2. + """ self.num_outer = num_outer self.num_inner = num_inner From 98ab175af543640eabe8d08d4b3856c70006360d Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 7 Sep 2023 21:05:11 +0100 Subject: [PATCH 17/28] implement adaptive sub-cycling for transport schemes --- gusto/io.py | 12 ++++- gusto/time_discretisation.py | 54 +++++++++++++------ gusto/timeloop.py | 3 ++ .../transport/test_subcycling.py | 14 ++--- 4 files changed, 60 insertions(+), 23 deletions(-) diff --git a/gusto/io.py b/gusto/io.py index 4ccb31144..a7ef44296 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -7,7 +7,7 @@ import time from gusto.diagnostics import Diagnostics, CourantNumber from gusto.meshes import get_flat_latlon_mesh -from firedrake import (Function, functionspaceimpl, File, +from firedrake import (Function, functionspaceimpl, Constant, File, DumbCheckpoint, FILE_CREATE, FILE_READ, CheckpointFile) import numpy as np from gusto.logging import logger, update_logfile_location @@ -226,6 +226,9 @@ def __init__(self, domain, output, diagnostics=None, diagnostic_fields=None): self.dumpfile = None self.to_pick_up = None + if output.log_courant: + self.courant_max = Constant(0.0) + def log_parameters(self, equation): """ Logs an equation's physical parameters that take non-default values. @@ -302,6 +305,13 @@ def log_courant(self, state_fields, name='u', component="whole", message=None): else: logger.info(f'Max Courant {message}: {courant_max:.2e}') + if component == 'whole': + # TODO: this will update the Courant number more than we need to + # and possibly with the wrong Courant number + # we could make self.courant_max a dict with keys depending on + # the field to take the Courant number of + self.courant_max.assign(courant_max) + def setup_diagnostics(self, state_fields): """ Prepares the I/O for computing the model's global diagnostics and diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index a6ff88194..be0dc9d87 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -72,8 +72,10 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.equation = None self.dt = domain.dt + self.original_dt = Constant(self.dt) self.options = options self.limiter = limiter + self.courant_max = None if options is not None: self.wrapper_name = options.name @@ -241,7 +243,7 @@ def apply(self, x_out, x_in): class ExplicitTimeDiscretisation(TimeDiscretisation): """Base class for explicit time discretisations.""" - def __init__(self, domain, field_name=None, subcycles=None, + def __init__(self, domain, field_name=None, subcycles=None, subcycle_by=None, solver_parameters=None, limiter=None, options=None): """ Args: @@ -249,8 +251,15 @@ def __init__(self, domain, field_name=None, subcycles=None, mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. - subcycles (int, optional): the number of sub-steps to perform. - Defaults to None. + subcycles (int, optional): the fixed number of sub-steps to perform. + Defaults to None. This option cannot be specified with the + `subcycle_by` argument. + subcycle_by (float, optional): specifying this option will make the + scheme performing adaptive sub-cycling based on the Courant + number. The specified argument is the maximum Courant for one + sub-cycle. Defaults to None, in which case adaptive sub-cycling + is not used. This option cannot be specified with the + `subcycles` argument. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -264,7 +273,11 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=solver_parameters, limiter=limiter, options=options) + if subcycles is not None and subcycle_by is not None: + raise ValueError('Cannot specify both subcycle and subcycle_by ' + + 'arguments to a time discretisation') self.subcycles = subcycles + self.subcycle_by = subcycle_by def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -332,6 +345,12 @@ def apply(self, x_out, x_in): x_out (:class:`Function`): the output field to be computed. x_in (:class:`Function`): the input field. """ + import math + # If doing adaptive subcycles, update dt and ncycles here + if self.subcycle_by is not None: + self.ncycles = math.ceil(float(self.courant_max)/self.subcycle_by) + self.dt = self.original_dt/self.ncycles + self.x0.assign(x_in) for i in range(self.ncycles): self.apply_cycle(self.x1, self.x0) @@ -376,10 +395,11 @@ class ExplicitMultistage(ExplicitTimeDiscretisation): """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, - limiter=None, options=None, butcher_matrix=None): + def __init__(self, domain, field_name=None, subcycles=None, subcycle_by=None, + solver_parameters=None, limiter=None, options=None, + butcher_matrix=None): super().__init__(domain, field_name=field_name, subcycles=subcycles, - solver_parameters=solver_parameters, + subcycle_by=subcycle_by, solver_parameters=solver_parameters, limiter=limiter, options=options) if butcher_matrix is not None: self.butcher_matrix = butcher_matrix @@ -479,10 +499,10 @@ class ForwardEuler(ExplicitMultistage): k0 = F[y^n] y^(n+1) = y^n + dt*k0 """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, - limiter=None, options=None, butcher_matrix=None): + def __init__(self, domain, field_name=None, subcycles=None, subcycle_by=None, + solver_parameters=None, limiter=None, options=None, butcher_matrix=None): super().__init__(domain, field_name=field_name, subcycles=subcycles, - solver_parameters=solver_parameters, + subcycle_by=subcycle_by, solver_parameters=solver_parameters, limiter=limiter, options=options, butcher_matrix=butcher_matrix) self.butcher_matrix = np.array([1.]).reshape(1, 1) self.nbutcher = int(np.shape(self.butcher_matrix)[0]) @@ -498,10 +518,11 @@ class SSPRK3(ExplicitMultistage): k2 = F[y^n + (1/4)*dt*(k0+k1)] y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2) """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, + def __init__(self, domain, field_name=None, subcycles=None, + subcycle_by=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None): super().__init__(domain, field_name=field_name, subcycles=subcycles, - solver_parameters=solver_parameters, + subcycle_by=subcycle_by, solver_parameters=solver_parameters, limiter=limiter, options=options, butcher_matrix=butcher_matrix) self.butcher_matrix = np.array([[1., 0., 0.], [1./4., 1./4., 0.], [1./6., 1./6., 2./3.]]) self.nbutcher = int(np.shape(self.butcher_matrix)[0]) @@ -522,10 +543,11 @@ class RK4(ExplicitMultistage): where superscripts indicate the time-level. """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, + def __init__(self, domain, field_name=None, subcycles=None, + subcycle_by=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None): super().__init__(domain, field_name=field_name, subcycles=subcycles, - solver_parameters=solver_parameters, + subcycle_by=subcycle_by, solver_parameters=solver_parameters, limiter=limiter, options=options, butcher_matrix=butcher_matrix) self.butcher_matrix = np.array([[0.5, 0., 0., 0.], [0., 0.5, 0., 0.], [0., 0., 1., 0.], [1./6., 1./3., 1./3., 1./6.]]) self.nbutcher = int(np.shape(self.butcher_matrix)[0]) @@ -544,9 +566,11 @@ class Heun(ExplicitMultistage): where superscripts indicate the time-level and subscripts indicate the stage number. """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, + def __init__(self, domain, field_name=None, subcycles=None, + subcycle_by=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None): - super().__init__(domain, field_name, + super().__init__(domain, field_name, subcycles=subcycles, + subcycle_by=subcycle_by, solver_parameters=solver_parameters, limiter=limiter, options=options) self.butcher_matrix = np.array([[1., 0.], [0.5, 0.5]]) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index c49cd99a9..2153f6207 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -297,6 +297,7 @@ def setup_scheme(self): self.setup_equation(self.equation) self.scheme.setup(self.equation) self.setup_transporting_velocity(self.scheme) + self.scheme.courant_max = self.io.courant_max def timestep(self): """ @@ -364,6 +365,7 @@ def setup_scheme(self): apply_bcs = True self.scheme.setup(self.equation, apply_bcs, dynamics) self.setup_transporting_velocity(self.scheme) + self.scheme.courant_max = self.io.courant_max def timestep(self): @@ -568,6 +570,7 @@ def setup_scheme(self): for _, scheme in self.active_transport: scheme.setup(self.equation, apply_bcs, transport) self.setup_transporting_velocity(scheme) + scheme.courant_max = self.io.courant_max apply_bcs = True for _, scheme in self.diffusion_schemes: diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 5894050b5..e15afb9bf 100644 --- a/integration-tests/transport/test_subcycling.py +++ b/integration-tests/transport/test_subcycling.py @@ -13,18 +13,18 @@ def run(timestepper, tmax, f_end): return norm(timestepper.fields("f") - f_end) / norm(f_end) -@pytest.mark.parametrize("equation_form", ["advective", "continuity"]) -def test_subcyling(tmpdir, equation_form, tracer_setup): +@pytest.mark.parametrize("subcycling", ["fixed", "adaptive"]) +def test_subcyling(tmpdir, subcycling, tracer_setup): geometry = "slice" setup = tracer_setup(tmpdir, geometry) domain = setup.domain V = domain.spaces("DG") - if equation_form == "advective": - eqn = AdvectionEquation(domain, V, "f") - else: - eqn = ContinuityEquation(domain, V, "f") + eqn = AdvectionEquation(domain, V, "f") - transport_scheme = SSPRK3(domain, subcycles=2) + if subcycling == "fixed": + transport_scheme = SSPRK3(domain, subcycles=2) + elif subcycling == "adaptive": + transport_scheme = SSPRK3(domain, subcycle_by=0.25) transport_method = DGUpwind(eqn, "f") timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) From 1daffc05ef32c4bc822861599065195d3e0bc886 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 8 Sep 2023 15:23:37 +0100 Subject: [PATCH 18/28] add test for suppressing vertical wind --- .../physics/test_suppress_vertical_wind.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 integration-tests/physics/test_suppress_vertical_wind.py 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..684aa2e51 --- /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 \ No newline at end of file From a536491d9756a9ed9884c2cfea0aa5bafa7195b4 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 9 Sep 2023 15:45:51 +0100 Subject: [PATCH 19/28] add boundary layer mixing parametrisation (not working for wind) --- gusto/common_forms.py | 4 +- gusto/configuration.py | 1 + gusto/coordinates.py | 6 +- gusto/domain.py | 2 +- gusto/physics.py | 159 +++++++++++++++++- gusto/time_discretisation.py | 6 - .../physics/test_boundary_layer_mixing.py | 144 ++++++++++++++++ .../physics/test_static_adjustment.py | 2 +- 8 files changed, 306 insertions(+), 18 deletions(-) create mode 100644 integration-tests/physics/test_boundary_layer_mixing.py 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 532258c4a..443db6297 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -198,3 +198,4 @@ class BoundaryLayerParameters(Configuration): 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 \ No newline at end of file diff --git a/gusto/coordinates.py b/gusto/coordinates.py index b0cfbae91..3e053639f 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -155,12 +155,14 @@ def register_space(self, domain, space_name): (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): + 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 @@ -171,6 +173,8 @@ def get_column_data(self, field): import pandas as pd 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) diff --git a/gusto/domain.py b/gusto/domain.py index eb4ea3aaf..54a8a27fa 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -175,7 +175,7 @@ def set_height_above_surface(self): # Turn height into columnwise data try: - columnwise_height, index_data = self.coords.get_column_data(CG1_height) + 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) diff --git a/gusto/physics.py b/gusto/physics.py index a9d630fc1..c23b5f090 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -16,8 +16,9 @@ from gusto.labels import PhysicsLabel, transporting_velocity, transport, prognostic from gusto.logging import logger from firedrake import (Interpolator, conditional, Function, dx, sqrt, dot, - min_value, max_value, Constant, pi, Projector, - TestFunctions, split, inner, TestFunction, + 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 @@ -29,7 +30,7 @@ __all__ = ["SaturationAdjustment", "Fallout", "Coalescence", "EvaporationOfRain", "AdvectedMoments", "InstantRain", "SWSaturationAdjustment", "SourceSink", "SurfaceFluxes", "WindDrag", "StaticAdjustment", - "SuppressVerticalWind"] + "SuppressVerticalWind", "BoundaryLayerMixing"] class PhysicsParametrisation(object, metaclass=ABCMeta): @@ -1283,7 +1284,8 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, else: dtheta_vd_dt = dT_dt * (1 / exner - kappa * theta_vd / (rho * T)) - source_theta_expr = test_theta * dtheta_vd_dt * dx + 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) @@ -1403,7 +1405,8 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): # Construct underlying expressions du_dt = -surface_expr * C_D * u_hori_mag * u_hori / z_a - source_expr = inner(test, du_dt) * dx + 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) @@ -1484,8 +1487,11 @@ def __init__(self, equation, theta_variable='theta_vd'): # Set up routines to reshape data # -------------------------------------------------------------------- # - self.get_column_data = partial(equation.domain.coords.get_column_data, field=self.theta_to_sort) - self.set_column_data = equation.domain.coords.set_field_from_column_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 @@ -1596,3 +1602,142 @@ def evaluate(self, x_in, 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) + d_dz = lambda q: outer(domain.k, dot(domain.k, grad(q))) + 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 + inner(d_dz(test/rho_averaged), d_dz(rho_averaged*K*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 be0dc9d87..e48ba3d23 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -601,13 +601,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', 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..dafea47ac --- /dev/null +++ b/integration-tests/physics/test_boundary_layer_mixing.py @@ -0,0 +1,144 @@ +""" +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_DG1 = VectorFunctionSpace(mesh, 'DG', 1) + recovery_opts = RecoveryOptions(embedding_space=Vec_DG1, + 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), +# ('u', True, True)]) +@pytest.mark.parametrize("field_name, recovered, semi_implicit", + [('u', True, True)]) +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('u') + 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 + import pdb; pdb.set_trace() + assert field_data[0,0] < 0.995*initial_data[0,0] diff --git a/integration-tests/physics/test_static_adjustment.py b/integration-tests/physics/test_static_adjustment.py index 404d96047..18ee13cea 100644 --- a/integration-tests/physics/test_static_adjustment.py +++ b/integration-tests/physics/test_static_adjustment.py @@ -101,7 +101,7 @@ def test_static_adjustment(tmpdir, theta_variable): else: theta = theta_vd - column_data, _ = domain.coords.get_column_data(theta) + 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:])) From 345c1a365340347557cbe89934d8a593006510f5 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 9 Sep 2023 16:30:49 +0100 Subject: [PATCH 20/28] add a boundary layer mixing scheme --- gusto/forcing.py | 4 ++-- gusto/timeloop.py | 4 ++-- .../physics/test_boundary_layer_mixing.py | 10 ++++------ 3 files changed, 8 insertions(+), 10 deletions(-) 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/timeloop.py b/gusto/timeloop.py index 2153f6207..e7208a0a1 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -463,8 +463,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, for _, 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" + # assert isinstance(scheme, ExplicitTimeDiscretisation), \ + # "Only explicit time discretisations can be used for physics" self.active_transport = [] for scheme in transport_schemes: diff --git a/integration-tests/physics/test_boundary_layer_mixing.py b/integration-tests/physics/test_boundary_layer_mixing.py index dafea47ac..4f9a6afa7 100644 --- a/integration-tests/physics/test_boundary_layer_mixing.py +++ b/integration-tests/physics/test_boundary_layer_mixing.py @@ -113,13 +113,11 @@ def run_boundary_layer_mixing(dirname, field_name, recovered, semi_implicit): return domain, stepper, initial_field -# @pytest.mark.parametrize("field_name, recovered, semi_implicit", -# [('theta', False, False), -# ('theta', False, True), -# ('u', False, False), -# ('u', True, True)]) @pytest.mark.parametrize("field_name, recovered, semi_implicit", - [('u', True, True)]) + [('theta', False, False), + ('theta', False, True), + ('u', False, False), + ('u', True, True)]) def test_boundary_layer_mixing(tmpdir, field_name, recovered, semi_implicit): dirname = str(tmpdir) From 726aaa73748973e68b987d753161c4bffeda87ba Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 17 Sep 2023 11:20:09 +0100 Subject: [PATCH 21/28] fix boundary layer mixing for wind --- gusto/physics.py | 6 ++++-- .../physics/test_boundary_layer_mixing.py | 12 ++++++------ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/gusto/physics.py b/gusto/physics.py index c23b5f090..6779bb096 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -1693,7 +1693,9 @@ def __init__(self, equation, field_name, parameters=None): dx_reduced = dx(degree=4) dS_v_reduced = dS_v(degree=4) - d_dz = lambda q: outer(domain.k, dot(domain.k, grad(q))) + # 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) @@ -1716,7 +1718,7 @@ def __init__(self, equation, field_name, parameters=None): # Interior penalty discretisation of vertical diffusion source_expr = ( # Volume term - inner(d_dz(test/rho_averaged), d_dz(rho_averaged*K*field))*dx_reduced + 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 diff --git a/integration-tests/physics/test_boundary_layer_mixing.py b/integration-tests/physics/test_boundary_layer_mixing.py index 4f9a6afa7..b7f736636 100644 --- a/integration-tests/physics/test_boundary_layer_mixing.py +++ b/integration-tests/physics/test_boundary_layer_mixing.py @@ -48,8 +48,8 @@ def run_boundary_layer_mixing(dirname, field_name, recovered, semi_implicit): if recovered: # Only implemented for u Vec_CG1 = VectorFunctionSpace(mesh, 'CG', 1) - Vec_DG1 = VectorFunctionSpace(mesh, 'DG', 1) - recovery_opts = RecoveryOptions(embedding_space=Vec_DG1, + 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) @@ -117,7 +117,8 @@ def run_boundary_layer_mixing(dirname, field_name, recovered, semi_implicit): [('theta', False, False), ('theta', False, True), ('u', False, False), - ('u', True, True)]) + ('u', True, True) + ]) def test_boundary_layer_mixing(tmpdir, field_name, recovered, semi_implicit): dirname = str(tmpdir) @@ -126,7 +127,7 @@ def test_boundary_layer_mixing(tmpdir, field_name, recovered, semi_implicit): if field_name == 'u': # Need to project horizontal wind into W3 - wind_2d = stepper.fields('u') + 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 @@ -138,5 +139,4 @@ def test_boundary_layer_mixing(tmpdir, field_name, recovered, semi_implicit): initial_data, _ = domain.coords.get_column_data(initial_field, domain) # Check first column - import pdb; pdb.set_trace() - assert field_data[0,0] < 0.995*initial_data[0,0] + assert field_data[0,0] < 0.999*initial_data[0,0] From 9f111545d7b6fdaf29b86ec75d0446d84db8d0ce Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 17 Sep 2023 20:23:29 +0100 Subject: [PATCH 22/28] allow test to xfail --- integration-tests/physics/test_boundary_layer_mixing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/integration-tests/physics/test_boundary_layer_mixing.py b/integration-tests/physics/test_boundary_layer_mixing.py index b7f736636..a9446f150 100644 --- a/integration-tests/physics/test_boundary_layer_mixing.py +++ b/integration-tests/physics/test_boundary_layer_mixing.py @@ -117,7 +117,8 @@ def run_boundary_layer_mixing(dirname, field_name, recovered, semi_implicit): [('theta', False, False), ('theta', False, True), ('u', False, False), - ('u', True, True) + 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): From eb7366199b599d0a95bd8a7387fa08b5fbc01a02 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 8 Oct 2023 18:56:04 +0100 Subject: [PATCH 23/28] add pandas requirement --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 7d438f194..a8302376a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ netCDF4 +pandas From 9c7f7ce74f59ad46bf4deaac012ae958c428c89f Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 8 Oct 2023 20:06:11 +0100 Subject: [PATCH 24/28] lint fixes --- gusto/configuration.py | 2 +- integration-tests/physics/test_boundary_layer_mixing.py | 2 +- integration-tests/physics/test_suppress_vertical_wind.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/gusto/configuration.py b/gusto/configuration.py index 901b6d6d3..3694ae9ea 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -199,4 +199,4 @@ class BoundaryLayerParameters(Configuration): 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 \ No newline at end of file + mu = 100. # Interior penalty coefficient for vertical diffusion diff --git a/integration-tests/physics/test_boundary_layer_mixing.py b/integration-tests/physics/test_boundary_layer_mixing.py index a9446f150..1b0499e70 100644 --- a/integration-tests/physics/test_boundary_layer_mixing.py +++ b/integration-tests/physics/test_boundary_layer_mixing.py @@ -140,4 +140,4 @@ def test_boundary_layer_mixing(tmpdir, field_name, recovered, semi_implicit): initial_data, _ = domain.coords.get_column_data(initial_field, domain) # Check first column - assert field_data[0,0] < 0.999*initial_data[0,0] + assert field_data[0, 0] < 0.999*initial_data[0, 0] diff --git a/integration-tests/physics/test_suppress_vertical_wind.py b/integration-tests/physics/test_suppress_vertical_wind.py index 684aa2e51..250916227 100644 --- a/integration-tests/physics/test_suppress_vertical_wind.py +++ b/integration-tests/physics/test_suppress_vertical_wind.py @@ -92,4 +92,4 @@ def test_suppress_vertical_wind(tmpdir): vertical_wind.interpolate(dot(u, domain.k)) tol = 1e-12 - assert norm(vertical_wind) < tol \ No newline at end of file + assert norm(vertical_wind) < tol From 0232df200fb5ecf1c52337c8bf9fb0b11c08089b Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 8 Oct 2023 20:11:14 +0100 Subject: [PATCH 25/28] update interface to CG1 space in domain given function space upgrades --- gusto/domain.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gusto/domain.py b/gusto/domain.py index 25d45cc11..4300920b3 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -168,7 +168,8 @@ def set_height_above_surface(self): x = SpatialCoordinate(self.mesh) # Make a height field in CG1 - CG1 = self.spaces('CG1', family='CG', degree=1) + 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)) From d158b4584cf8a917eed4a8915700a1ff763bdb06 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 9 Oct 2023 08:20:16 +0100 Subject: [PATCH 26/28] fix failing tests (improve coord rounding and remove undefined as_vector) --- examples/compressible/skamarock_klemp_hydrostatic.py | 2 +- gusto/coordinates.py | 11 +++++++---- integration-tests/equations/test_thermal_sw.py | 2 +- 3 files changed, 9 insertions(+), 6 deletions(-) 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/coordinates.py b/gusto/coordinates.py index 33e9ae24a..f8cf94f2d 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -189,13 +189,16 @@ def get_column_data(self, field, domain): # 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) - digits = int(np.floor(-np.log10(data_range / num_points)) + 3) - coords_X = coords_X.round(digits) + if data_range > 1e-16: + 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) - digits = int(np.floor(-np.log10(data_range / num_points)) + 3) - coords_Y = coords_Y.round(digits) + if data_range > 1e-16: + # 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 diff --git a/integration-tests/equations/test_thermal_sw.py b/integration-tests/equations/test_thermal_sw.py index 266fc8201..8328c308d 100644 --- a/integration-tests/equations/test_thermal_sw.py +++ b/integration-tests/equations/test_thermal_sw.py @@ -80,7 +80,7 @@ def set_up_initial_conditions(domain, equation, stepper): bexpr = g * (1 - theta) - u0.project(as_vector(uexpr)) + u0.project(uexpr) D0.interpolate(Dexpr) b0.interpolate(bexpr) From afa859c527f5747f5b9135ae0b501ced7af636b6 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 22 Oct 2023 20:13:04 +0100 Subject: [PATCH 27/28] Restore check for explicit physics schemes and a few other tidyings --- gusto/coordinates.py | 7 +++---- gusto/diagnostics.py | 2 +- gusto/domain.py | 28 +++++++++------------------- gusto/logging.py | 1 + gusto/physics.py | 31 ++++++++++++++++++++----------- gusto/timeloop.py | 16 ++++++++++------ 6 files changed, 44 insertions(+), 41 deletions(-) diff --git a/gusto/coordinates.py b/gusto/coordinates.py index f8cf94f2d..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): @@ -170,8 +171,6 @@ def get_column_data(self, field, domain): ordered column data. """ - import pandas as pd - space_name = field.function_space().name if space_name not in self.chi_coords.keys(): self.register_space(domain, space_name) @@ -189,13 +188,13 @@ def get_column_data(self, field, domain): # 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 > 1e-16: + 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 > 1e-16: + 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) diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index cfcde1a7d..024a42fa9 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -222,7 +222,7 @@ def setup(self, domain, state_fields, space=None): def compute(self): """Compute the diagnostic field from the current state.""" - logger.info(f'Computing diagnostic {self.name} with {self.method} method') + logger.debug(f'Computing diagnostic {self.name} with {self.method} method') if self.method == 'interpolate': self.evaluator.interpolate() diff --git a/gusto/domain.py b/gusto/domain.py index 4300920b3..8f22083c6 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -4,7 +4,6 @@ model's time interval. """ -from gusto.logging import logger from gusto.coordinates import Coordinates from gusto.function_spaces import Spaces, check_degree_args from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, @@ -176,24 +175,15 @@ def set_height_above_surface(self): height_above_surface = Function(CG1) # Turn height into columnwise data - try: - 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) - except ModuleNotFoundError: - # If no pandas, then don't take orography into account - logger.warning( - 'Pandas not found, so using height as height above surface. ' - + 'This will be incorrect in the presence of orography') - - if self.on_sphere: - height_above_surface.assign(CG1_height - self.radius) + 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 diff --git a/gusto/logging.py b/gusto/logging.py index c6eef1e5e..dc116275b 100644 --- a/gusto/logging.py +++ b/gusto/logging.py @@ -186,6 +186,7 @@ def update_logfile_location(new_path, comm): fh.close() logger.removeHandler(fh) + os.makedirs(new_path, exist_ok=True) if parallel_log in ["FILE", "BOTH"]: # If all ranks are logging wait here in case a directory is being created comm.Barrier() diff --git a/gusto/physics.py b/gusto/physics.py index 658163005..76c8ee4c3 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -200,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 @@ -262,7 +263,7 @@ def __init__(self, equation, vapour_name='water_vapour', liquid_water += self.X.subfunctions[liq_idx] # define some parameters as attributes - self.dt = Constant(1.0) + self.dt = Constant(0.0) R_d = parameters.R_d cp = parameters.cp cv = parameters.cv @@ -439,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') @@ -524,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" @@ -549,7 +552,7 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', self.source = Function(Vt) # define some parameters as attributes - self.dt = Constant(1.0) + self.dt = Constant(0.0) # TODO: should these parameters be hard-coded or configurable? k_1 = Constant(0.001) # accretion rate in 1/s k_2 = Constant(2.2) # accumulation rate in 1/s @@ -635,6 +638,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', CompressibleEulerEquations. """ + self.explicit_only = True label_name = 'evaporation_of_rain' super().__init__(equation, label_name, parameters=None) @@ -694,7 +698,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', liquid_water += self.X.subfunctions[liq_idx] # define some parameters as attributes - self.dt = Constant(1.0) + self.dt = Constant(0.0) R_d = parameters.R_d cp = parameters.cp cv = parameters.cv @@ -819,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) @@ -986,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) @@ -1144,8 +1150,8 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, 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. + 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. @@ -1169,7 +1175,7 @@ def __init__(self, equation, T_surface_expr, vapour_name=None, self.implicit_formulation = implicit_formulation self.X = Function(equation.X.function_space()) - self.dt = Constant(1.0) + self.dt = Constant(0.0) # -------------------------------------------------------------------- # # Extract prognostic variables @@ -1324,8 +1330,8 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): 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. + 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. @@ -1346,7 +1352,7 @@ def __init__(self, equation, implicit_formulation=False, parameters=None): k = equation.domain.k self.implicit_formulation = implicit_formulation self.X = Function(equation.X.function_space()) - self.dt = Constant(1.0) + self.dt = Constant(0.0) # -------------------------------------------------------------------- # # Extract prognostic variables @@ -1441,6 +1447,7 @@ def __init__(self, equation, theta_variable='theta_vd'): to "theta_vd". """ + self.explicit_only = True from functools import partial # -------------------------------------------------------------------- # @@ -1457,7 +1464,7 @@ def __init__(self, equation, theta_variable='theta_vd'): super().__init__(equation, label_name, parameters=equation.parameters) self.X = Function(equation.X.function_space()) - self.dt = Constant(1.0) + self.dt = Constant(0.0) # -------------------------------------------------------------------- # # Extract prognostic variables @@ -1540,6 +1547,8 @@ def __init__(self, equation, spin_up_period): period. """ + self.explicit_only = True + # -------------------------------------------------------------------- # # Check arguments and generic initialisation # -------------------------------------------------------------------- # @@ -1554,7 +1563,7 @@ def __init__(self, equation, spin_up_period): super().__init__(equation, label_name, parameters=equation.parameters) self.X = Function(equation.X.function_space()) - self.dt = Constant(1.0) + self.dt = Constant(0.0) self.t = domain.t self.spin_up_period = Constant(spin_up_period) self.strength = Constant(1.0) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 7252d563c..c020c8bb8 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -347,9 +347,11 @@ def __init__(self, equation, scheme, io, spatial_methods=None, self.physics_schemes = [] for parametrisation, 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" + # 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, parametrisation.label) @@ -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: From 51f4f9142ea72a790b497652a732e130ef2f7d16 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 24 Oct 2023 14:08:06 +0100 Subject: [PATCH 28/28] remove chevrons from merge --- gusto/timeloop.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index c020c8bb8..a7fc661b3 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -750,10 +750,6 @@ def setup_fields(self): def run(self, t, tmax, pick_up=False): """ Runs the model for the specified time, from t to tmax -<<<<<<< HEAD - -======= ->>>>>>> origin/main Args: t (float): the start time of the run tmax (float): the end time of the run