diff --git a/examples/compressible/dcmip_3_1_meanflow_quads.py b/examples/compressible/dcmip_3_1_meanflow_quads.py index 13cc2f532..93bb12d8c 100644 --- a/examples/compressible/dcmip_3_1_meanflow_quads.py +++ b/examples/compressible/dcmip_3_1_meanflow_quads.py @@ -87,8 +87,8 @@ # Transport schemes transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "rho", subcycles=2), - SSPRK3(domain, "theta", options=SUPGOptions(), subcycles=2)] + SSPRK3(domain, "rho", fixed_subcycles=2), + SSPRK3(domain, "theta", options=SUPGOptions(), fixed_subcycles=2)] transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta"]] # Linear solver diff --git a/examples/compressible/mountain_hydrostatic.py b/examples/compressible/mountain_hydrostatic.py index b90630884..98d221db6 100644 --- a/examples/compressible/mountain_hydrostatic.py +++ b/examples/compressible/mountain_hydrostatic.py @@ -66,7 +66,7 @@ dumpfreq=dumpfreq, dumplist=['u'], ) -diagnostic_fields = [CourantNumber(), VelocityZ(), HydrostaticImbalance(eqns), +diagnostic_fields = [CourantNumber(), ZComponent('u'), HydrostaticImbalance(eqns), Perturbation('theta'), Perturbation('rho')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/compressible/skamarock_klemp_nonhydrostatic.py b/examples/compressible/skamarock_klemp_nonhydrostatic.py index c6c32af7a..8e9273b2e 100644 --- a/examples/compressible/skamarock_klemp_nonhydrostatic.py +++ b/examples/compressible/skamarock_klemp_nonhydrostatic.py @@ -9,7 +9,7 @@ from gusto import * import itertools from firedrake import (as_vector, SpatialCoordinate, PeriodicIntervalMesh, - ExtrudedMesh, exp, sin, Function, pi) + ExtrudedMesh, exp, sin, Function, pi, COMM_WORLD) import numpy as np import sys @@ -51,13 +51,28 @@ points_z = [H/2.] points = np.array([p for p in itertools.product(points_x, points_z)]) dirname = 'skamarock_klemp_nonlinear' -output = OutputParameters( - dirname=dirname, - dumpfreq=dumpfreq, - pddumpfreq=dumpfreq, - dumplist=['u'], - point_data=[('theta_perturbation', points)], -) + +# Dumping point data using legacy PointDataOutput is not supported in parallel +if COMM_WORLD.size == 1: + output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + pddumpfreq=dumpfreq, + dumplist=['u'], + point_data=[('theta_perturbation', points)], + ) +else: + logger.warning( + 'Dumping point data using legacy PointDataOutput is not' + ' supported in parallel\nDisabling PointDataOutput' + ) + output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + pddumpfreq=dumpfreq, + dumplist=['u'], + ) + diagnostic_fields = [CourantNumber(), Gradient("u"), Perturbation('theta'), Gradient("theta_perturbation"), Perturbation('rho'), RichardsonNumber("theta", parameters.g/Tsurf), Gradient("theta")] diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py index acb73e7a2..09381cc7d 100644 --- a/examples/shallow_water/moist_thermal_williamson5.py +++ b/examples/shallow_water/moist_thermal_williamson5.py @@ -60,7 +60,7 @@ fexpr = 2*Omega*x[2]/R # Topography -phi, lamda = latlon_coords(mesh) +lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) lsq = (lamda - lamda_c)**2 thsq = (phi - phi_c)**2 rsq = min_value(R0sq, lsq+thsq) diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py index 28f3a1d8f..e9fcd94ac 100644 --- a/examples/shallow_water/thermal_williamson2.py +++ b/examples/shallow_water/thermal_williamson2.py @@ -1,5 +1,5 @@ from gusto import * -from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, sin, cos) +from firedrake import IcosahedralSphereMesh, SpatialCoordinate, sin, cos import sys # ----------------------------------------------------------------- # @@ -71,9 +71,9 @@ D0 = stepper.fields("D") b0 = stepper.fields("b") -phi, lamda = latlon_coords(mesh) +lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) -uexpr = sphere_to_cartesian(mesh, u_max*cos(phi), 0) +uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x) g = params.g w = Omega*R*u_max + (u_max**2)/2 sigma = w/10 diff --git a/examples/shallow_water/williamson_2.py b/examples/shallow_water/williamson_2.py index 47563a632..8efa7fc27 100644 --- a/examples/shallow_water/williamson_2.py +++ b/examples/shallow_water/williamson_2.py @@ -7,7 +7,7 @@ """ from gusto import * -from firedrake import IcosahedralSphereMesh, SpatialCoordinate, as_vector, pi +from firedrake import IcosahedralSphereMesh, SpatialCoordinate, sin, cos, pi import sys # ---------------------------------------------------------------------------- # @@ -28,6 +28,7 @@ # setup shallow water parameters R = 6371220. H = 5960. +rotated_pole = (0.0, pi/3) # setup input that doesn't change with ref level or dt parameters = ShallowWaterParameters(H=H) @@ -42,11 +43,13 @@ mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2) x = SpatialCoordinate(mesh) - domain = Domain(mesh, dt, 'BDM', 1) + domain = Domain(mesh, dt, 'BDM', 1, rotated_pole=rotated_pole) # Equation Omega = parameters.Omega - fexpr = 2*Omega*x[2]/R + _, lat, _ = rotated_lonlatr_coords(x, rotated_pole) + e_lon, _, _ = rotated_lonlatr_vectors(x, rotated_pole) + fexpr = 2*Omega*sin(lat) eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr) # I/O @@ -63,12 +66,14 @@ ShallowWaterKineticEnergy(), ShallowWaterPotentialEnergy(parameters), ShallowWaterPotentialEnstrophy(), - SteadyStateError('u'), SteadyStateError('D')] + SteadyStateError('u'), SteadyStateError('D'), + MeridionalComponent('u', rotated_pole), + ZonalComponent('u', rotated_pole)] io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes transported_fields = [TrapeziumRule(domain, "u"), - SSPRK3(domain, "D", subcycles=2)] + SSPRK3(domain, "D", fixed_subcycles=2)] transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] # Time stepper @@ -82,9 +87,9 @@ D0 = stepper.fields("D") x = SpatialCoordinate(mesh) u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) - uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) + uexpr = u_max*cos(lat)*e_lon g = parameters.g - Dexpr = H - ((R * Omega * u_max + u_max*u_max/2.0)*(x[2]*x[2]/(R*R)))/g + Dexpr = H - (R * Omega * u_max + u_max*u_max/2.0)*(sin(lat))**2/g u0.project(uexpr) D0.interpolate(Dexpr) diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index b71fddab1..366a5a903 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -47,7 +47,7 @@ # Equation Omega = parameters.Omega fexpr = 2*Omega*x[2]/R - theta, lamda = latlon_coords(mesh) + lamda, theta, _ = lonlatr_from_xyz(x[0], x[1], x[2]) R0 = pi/9. R0sq = R0**2 lamda_c = -pi/2. diff --git a/gusto/__init__.py b/gusto/__init__.py index 3d0b4f81d..d7dc21faa 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -6,6 +6,7 @@ from gusto.active_tracers import * # noqa from gusto.common_forms import * # noqa from gusto.configuration import * # noqa +from gusto.coord_transforms import * # noqa from gusto.domain import * # noqa from gusto.diagnostics import * # noqa from gusto.diffusion_methods import * # noqa diff --git a/gusto/configuration.py b/gusto/configuration.py index 0ab9cdb37..039b32740 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -88,7 +88,8 @@ class OutputParameters(Configuration): dumplist = None dumplist_latlon = [] dump_diagnostics = True - checkpoint = True + diagfreq = 1 + checkpoint = False checkpoint_method = 'checkpointfile' checkpoint_pickup_filename = None chkptfreq = 1 diff --git a/gusto/coord_transforms.py b/gusto/coord_transforms.py new file mode 100644 index 000000000..f40456819 --- /dev/null +++ b/gusto/coord_transforms.py @@ -0,0 +1,533 @@ +""" +Stores some common routines to transform coordinates between spherical and +Cartesian systems. +""" + +import importlib +import numpy as np +from firedrake import SpatialCoordinate +import ufl + +__all__ = ["xyz_from_lonlatr", "lonlatr_from_xyz", "xyz_vector_from_lonlatr", + "lonlatr_components_from_xyz", "rodrigues_rotation", "pole_rotation", + "rotated_lonlatr_vectors", "rotated_lonlatr_coords", + "periodic_distance", "great_arc_angle"] + + +def firedrake_or_numpy(variable): + """ + A function internal to this module, used to determine whether to import + other routines from firedrake or numpy. + + Args: + variable (:class:`np.ndarray` or :class:`ufl.Expr`): a variable to be + used in the coordinate transform routines. + + Returns: + tuple (:class:`module`, str): either the firedrake or numpy module, with + its name. + """ + + if isinstance(variable, ufl.core.expr.Expr): + module_name = 'firedrake' + else: + module_name = 'numpy' + + module = importlib.import_module(module_name) + + return module, module_name + + +def magnitude(u): + """ + A function internal to this module for returning the pointwise magnitude of + a vector-valued field. + + Args: + u (:class:`np.ndarray` or :class:`ufl.Expr`): the vector-valued field to + take the magnitude of. + + Returns: + :class:`np.ndarray` or :class:`ufl.Expr`: |u|, the pointwise magntiude + of the vector field. + """ + + # Determine whether to use firedrake or numpy functions + module = firedrake_or_numpy(u) + sqrt = module.sqrt + dot = module.dot + + return sqrt(dot(u, u)) + + +def xyz_from_lonlatr(lon, lat, r, angle_units='rad'): + """ + Returns the geocentric Cartesian coordinates x, y, z from spherical lon, lat + and r coordinates. + + Args: + lon (:class:`np.ndarray` or :class:`ufl.Expr`): longitude coordinate. + lat (:class:`np.ndarray` or :class:`ufl.Expr`): latitude coordinate. + r (:class:`np.ndarray` or :class:`ufl.Expr`): radial coordinate. + angle_units (str, optional): the units used for the angle. Valid + options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. + + Returns: + tuple of :class`np.ndarray` or tuple of :class:`ufl.Expr`: the tuple + of (x, y, z) coordinates in the appropriate form given the + provided arguments. + """ + + if angle_units not in ['rad', 'deg']: + raise ValueError(f'angle_units arg {angle_units} not valid') + + # Import routines + module, _ = firedrake_or_numpy(lon) + cos = module.cos + sin = module.sin + pi = module.pi + + if angle_units == 'deg': + unit_factor = pi/180.0 + if angle_units == 'rad': + unit_factor = 1.0 + + lon = lon*unit_factor + lat = lat*unit_factor + + x = r * cos(lon) * cos(lat) + y = r * sin(lon) * cos(lat) + z = r * sin(lat) + + return x, y, z + + +def lonlatr_from_xyz(x, y, z, angle_units='rad'): + """ + Returns the spherical lon, lat and r coordinates from the global geocentric + Cartesian x, y, z coordinates. + + Args: + x (:class:`np.ndarray` or :class:`ufl.Expr`): x-coordinate. + y (:class:`np.ndarray` or :class:`ufl.Expr`): y-coordinate. + z (:class:`np.ndarray` or :class:`ufl.Expr`): z-coordinate. + angle_units (str, optional): the units to use for the angle. Valid + options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. + + Returns: + tuple of :class`np.ndarray` or tuple of :class:`ufl.Expr`: the tuple + of (lon, lat, r) coordinates in the appropriate form given the + provided arguments. + """ + + if angle_units not in ['rad', 'deg']: + raise ValueError(f'angle_units arg {angle_units} not valid') + + # Determine whether to use firedrake or numpy functions + module, _ = firedrake_or_numpy(x) + atan2 = module.atan2 if hasattr(module, "atan2") else module.arctan2 + sqrt = module.sqrt + pi = module.pi + + if angle_units == 'deg': + unit_factor = 180./pi + if angle_units == 'rad': + unit_factor = 1.0 + + lon = atan2(y, x) + r = sqrt(x**2 + y**2 + z**2) + l = sqrt(x**2 + y**2) + lat = atan2(z, l) + + return lon*unit_factor, lat*unit_factor, r + + +def xyz_vector_from_lonlatr(lon_component, lat_component, r_component, + position_vector, position_units="xyz"): + """ + Returns the Cartesian geocentric x, y and z components of a vector from a + vector whose components are in lon, lat and r spherical coordinates. If + dealing with Firedrake, a vector expression is returned. + + Args: + lon_component (:class:`np.ndarray` or :class:`ufl.Expr`): the zonal + component of the input vector. + lat_component (:class:`np.ndarray` or :class:`ufl.Expr`): the meridional + component of the input vector. + r_component (:class:`np.ndarray` or :class:`ufl.Expr`): the radial + component of the input vector. + position_vector (:class:`np.ndarray` or :class:`ufl.Expr`): the position + vector, either as (x, y, z) or (lon, lat, radius) coordinates, + subject to the `position_units` argument. Should match the shape of + the input vector. + position_units (str, optional): in which units the provided position + vector is. Valid options are ["xyz", "lonlatr_rad", "lonlatr_deg"]. + Defaults to "xyz". + + Returns: + :class:`np.ndarray` or :class:`ufl.as_vector`: (x, y, z) components of the + input vector. + """ + + # Check position units argument is valid + if position_units not in ["xyz", "lonlatr_rad", "lonlatr_deg"]: + raise ValueError('xyz_vector_from_lonlatr: the `position_units` arg ' + + 'must be one of "xyz", "lonlatr_rad", "lonlatr_deg "' + + f'but {position_units} was provided') + + # Determine whether to use firedrake or numpy functions + module, module_name = firedrake_or_numpy(position_vector[0]) + pi = module.pi + cos = module.cos + sin = module.sin + + # Convert position to lonlatr_rad + if position_units == 'xyz': + lon, lat, _ = lonlatr_from_xyz(position_vector[0], position_vector[1], + position_vector[2]) + elif position_units == 'lonlatr_rad': + lon, lat, _ = position_vector + elif position_units == 'lonlatr_deg': + lon, lat = position_vector[0]*pi/180., position_vector[1]*pi/180. + + # f is our vector, e_i is the ith unit basis vector + # f = f_r * e_r + f_lon * e_lon + f_lat * e_lat + # We want f = f_x * e_x + f_y * e_y + f_z * e_z + + # f_x = dot(f, e_x) + # e_x = cos(lon)*cos(lat) * e_r - sin(lon) * e_lon - cos(lon)*sin(lat) * e_lat + x_component = (cos(lon)*cos(lat) * r_component + - sin(lon) * lon_component + - cos(lon)*sin(lat) * lat_component) + + # f_y = dot(f, e_y) + # e_y = sin(lon)*cos(lat) * e_r + cos(lon) * e_lon - sin(lon)*sin(lat) * e_lat + y_component = (sin(lon)*cos(lat) * r_component + + cos(lon) * lon_component + - sin(lon)*sin(lat) * lat_component) + + # f_z = dot(f, e_z) + # e_z = sin(lat) * e_r + cos(lat) * e_lat + z_component = (sin(lat) * r_component + + cos(lat) * lat_component) + + if module_name == 'firedrake': + return module.as_vector((x_component, y_component, z_component)) + else: + return (x_component, y_component, z_component) + + +def lonlatr_components_from_xyz(xyz_vector, position_vector, position_units='xyz'): + """ + Returns the spherical (zonal, meridional, radial) components of a vector- + valued field from a vector which is expressed in geocentric Cartesian + (x, y, z) components. + + Args: + xyz_vector (:class:`np.ndarray` or :class:`ufl.Expr`): the input vector + in geocentric Cartesian (x, y, z) components. + position_vector (:class:`np.ndarray` or :class:`ufl.Expr`): the position + vector, either as (x, y, z) or (lon, lat, radius) coordinates, + subject to the `position_units` argument. Should match the shape of + the input vector. + position_units (str, optional): in which units the provided position + vector is. Valid options are ["xyz", "lonlatr_rad", "lonlatr_deg"]. + Defaults to "xyz". + + Returns: + :class:`np.ndarray` or :class:`ufl.Expr`: (zonal, meridional, radial) + components of the input vector. + """ + + # Check position units argument is valid + if position_units not in ["xyz", "lonlatr_rad", "lonlatr_deg"]: + raise ValueError('xyz_vector_from_lonlatr: the `position_units` arg ' + + 'must be one of "xyz", "lonlatr_rad", "lonlatr_deg "' + + f'but {position_units} was provided') + + # Determine whether to use firedrake or numpy functions + module, _ = firedrake_or_numpy(position_vector[0]) + pi = module.pi + cos = module.cos + sin = module.sin + + # Convert position to lonlatr_rad + if position_units == 'xyz': + lon, lat, _ = lonlatr_from_xyz(position_vector[0], position_vector[1], + position_vector[2]) + elif position_units == 'lonlatr_rad': + lon, lat, _ = position_vector + elif position_units == 'lonlatr_deg': + lon, lat = position_vector[0]*pi/180., position_vector[1]*pi/180. + + # f is our vector, e_i is the ith unit basis vector + # f = f_x * e_x + f_y * e_y + f_z * e_z + # We want f = f_r * e_r + f_lon * e_lon + f_lat * e_lat + + # f_lon = dot(f, e_lon) + # e_lon = -y/l * e_x + x/l * e_y + zonal_component = (-sin(lon) * xyz_vector[0] + + cos(lon) * xyz_vector[1]) + + # f_lat = dot(f, e_lat) + # e_lat = -x*z/(r*l) * e_x - y*z/(r*l) * e_y + l/r * e_z + meridional_component = (- cos(lon) * sin(lat) * xyz_vector[0] + - sin(lon) * sin(lat) * xyz_vector[1] + + cos(lat) * xyz_vector[2]) + + # f_r = dot(f, e_r) + # e_r = x/r * e_x + y/r * e_y + z/r * e_z + radial_component = (cos(lon) * cos(lat) * xyz_vector[0] + + sin(lon) * cos(lat) * xyz_vector[1] + + sin(lat) * xyz_vector[2]) + + return (zonal_component, meridional_component, radial_component) + + +def rodrigues_rotation(old_vector, rot_axis, rot_angle): + u""" + Performs the rotation of a vector v about some axis k by some angle ϕ, as + given by Rodrigues' rotation formula: \n + + v_rot = v * cos(ϕ) + (k cross v) sin(ϕ) + k (k . v)*(1 - cos(ϕ)) \n + + Returns a new vector. All components must be (x,y,z) components. + + Args: + old_vector (:class:`np.ndarray` or :class:`ufl.Expr`): the original + vector or vector-valued field to be rotated, to be expressed in + geocentric Cartesian (x,y,z) components in the original coordinate + system. + rot_axis (tuple or :class:`ufl.as_vector`): the vector representing the + axis to rotate around, expressed in geocentric Cartesian (x,y,z) + components (in the frame before the rotation). + rot_angle (float): the angle to rotate by. + + Returns: + :class:`np.ndarray` or :class:`ufl.Expr`: the rotated vector or + vector-valued field. + """ + + # Determine whether to use firedrake or numpy functions + module, module_name = firedrake_or_numpy(old_vector) + cos = module.cos + sin = module.sin + cross = module.cross + + # Numpy vector may need reshaping + if module_name == 'numpy' and np.shape(rot_axis) != np.shape(old_vector): + # Construct shape for tiling vector + tile_shape = [dim for dim in np.shape(old_vector)[:-1]]+[1] + # Tile rot_axis vector to create an ndarray + rot_axis = np.tile(rot_axis, tile_shape) + + # Replace dot product routine with something that does elementwise dot + def dot(a, b): + dotted_vectors = np.einsum('ij,ij->i', a, b) + # Add new axis to allow further multiplication by a vector + return dotted_vectors[:, np.newaxis] + else: + dot = module.dot + + new_vector = (old_vector * cos(rot_angle) + + cross(rot_axis, old_vector) * sin(rot_angle) + + rot_axis * dot(rot_axis, old_vector) * (1 - cos(rot_angle))) + + return new_vector + + +def pole_rotation(new_pole): + """ + Computes the rotation axis and angle associated with rotating the pole from + lon = 0 and lat = pi / 2 to a new longitude and latitude. Returns the + rotation axis and angle for use in the Rodrigues rotation. + + Args: + new_pole (tuple): a tuple of floats (lon, lat) of the new pole. The + longitude and latitude must be expressed in radians. + + Returns: + tuple: (rot_axis, rot_angle). This describes the rotation axis (a tuple + or :class:`as_vector` of (x, y, z) components of the rotation axis, + and a float describing the rotation angle. + """ + + import numpy as np + + # We assume that the old pole is old_lon_p = 0 and old_lat_p = pi / 2. + old_lat_p = np.pi / 2 + + # Then moving the pole to new_lon_p, new_lat_p is akin to rotating the pole + # about lon_rot = new_lon + pi / 2, lat_rot = 0 + new_lon_p, new_lat_p = new_pole + lon_rot = new_lon_p + np.pi / 2 + lat_rot = 0.0 + + # The rotation angle is only in the latitudinal direction + rot_angle = old_lat_p - new_lat_p + + # Turn rotation axis into a vector + # it points in the radial direction and has a length of one + rot_axis = xyz_vector_from_lonlatr(0, 0, 1, (lon_rot, lat_rot, 1), + position_units='lonlatr_rad') + + return rot_axis, rot_angle + + +def rotated_lonlatr_vectors(xyz, new_pole): + """ + Returns the (X,Y,Z) components of rotated (lon,lat,r) unit basis vectors, + given a rotation axis and the (X,Y,Z) coordinates and the old (lon,lat,r) + unit basis vectors. Only implemented for Firedrake. + + Args: + xyz (:class:`SpatialCoordinate`): Original geocentric Cartesian + coordinates. + new_pole (tuple): a tuple of floats (lon, lat) of the new pole, in the + original coordinate system. The longitude and latitude must be + expressed in radians. + + Returns: + tuple of :class:`ufl.Expr`: the rotated basis vectors (e_lon, e_lat, e_r). + """ + + from firedrake import sqrt, dot, as_vector, Constant + + rot_axis, rot_angle = pole_rotation(new_pole) + rot_axis = as_vector(rot_axis) + new_xyz = rodrigues_rotation(xyz, rot_axis, rot_angle) + + # Compute e_lon, e_lat vectors in terms of new (x,y,z) components + e_lon_new_xyz = xyz_vector_from_lonlatr(Constant(1.0), Constant(0.0), Constant(0.0), new_xyz) + e_lat_new_xyz = xyz_vector_from_lonlatr(Constant(0.0), Constant(1.0), Constant(0.0), new_xyz) + + # Rotate back to original (x,y,z) components + new_e_lon = rodrigues_rotation(e_lon_new_xyz, rot_axis, -rot_angle) + new_e_lat = rodrigues_rotation(e_lat_new_xyz, rot_axis, -rot_angle) + + # e_r isn't rotated + new_e_r = xyz_vector_from_lonlatr(Constant(0.0), Constant(0.0), Constant(1.0), xyz) + + # Normalise + new_e_lon /= sqrt(dot(new_e_lon, new_e_lon)) + new_e_lat /= sqrt(dot(new_e_lat, new_e_lat)) + new_e_r /= sqrt(dot(new_e_r, new_e_r)) + + return (new_e_lon, new_e_lat, new_e_r) + + +def rotated_lonlatr_coords(xyz, new_pole): + """ + Returns the rotated (lon,lat,r) coordinates, given a rotation axis and the + (X,Y,Z) coordinates. + + Args: + xyz (tuple of :class:`np.ndarray` or :class:`SpatialCoordinate`): + Original geocentric Cartesian coordinates. + new_pole (tuple): a tuple of floats (lon, lat) of the new pole, in the + original coordinate system. The longitude and latitude must be + expressed in radians. + + tuple of :class:`np.ndarray` or :class:`ufl.Expr`: the rotated + (lon,lat,r) coordinates. + """ + + rot_axis, rot_angle = pole_rotation(new_pole) + + # If numpy, shape (x,y,z) array as a vector + module, module_name = firedrake_or_numpy(xyz[0]) + if module_name == 'numpy': + old_xyz_vector = np.transpose(xyz) + else: + assert isinstance(xyz, SpatialCoordinate), 'Rotated lonlatr ' \ + + 'coordinates require xyz to be a SpatialCoordinate object' + old_xyz_vector = xyz + rot_axis = module.as_vector(rot_axis) + + # Do rotations to get new coordinates + new_xyz_vector = rodrigues_rotation(old_xyz_vector, rot_axis, rot_angle) + + if module_name == 'numpy': + new_xyz = np.transpose(new_xyz_vector) + else: + new_xyz = new_xyz_vector + + new_lonlatr = lonlatr_from_xyz(new_xyz[0], new_xyz[1], new_xyz[2]) + + return new_lonlatr + + +def periodic_distance(x1, x2, max_x, min_x=0.0): + """ + Finds the shortest distance between two points x1 and x2, on a periodic + interval of length Lx. + + Args: + x1 (:class:`np.ndarray` or :class:`ufl.Expr`): first set of position + values. + x2 (:class:`np.ndarray` or :class:`ufl.Expr`): second set of position + values. + max_x (:class:`Constant` or float): maximum coordinate on the domain. + min_x (:class:`Constant` or float, optional): minimum coordinate on the + domain. Defaults to None. + + Returns: + :class:`np.ndarray` or :class:`ufl.Expr`: the shortest distance between + the two points. + """ + + module, _ = firedrake_or_numpy(x1) + + # Use firedrake.conditional or numpy.where + conditional = module.conditional if hasattr(module, "conditional") else module.where + + Lx = max_x - min_x + longest_dist = Lx / 2 + trial_dist = x1 - x2 + dist = conditional(trial_dist > longest_dist, trial_dist - Lx, + conditional(trial_dist < - longest_dist, trial_dist + Lx, + trial_dist)) + + return dist + + +def great_arc_angle(lon1, lat1, lon2, lat2, units='rad'): + """ + Finds the arc angle along a great circle between two points on the sphere. + + Args: + lon1 (:class:`np.ndarray` or :class:`ufl.Expr`): first longitude value + or set of longitude values. + lat1 (:class:`np.ndarray` or :class:`ufl.Expr`): first latitude value or + set of latitude values. + lon2 (:class:`np.ndarray` or :class:`ufl.Expr`): second longitude value + or set of longitude values. + lat2 (:class:`np.ndarray` or :class:`ufl.Expr`): second latitude value + or set of latitude values. + units (str, optional): which units the angles are expressed in. Should + be "deg" or "rad". Defaults to "rad". + + Returns: + :class:`np.ndarray` or :class:`ufl.Expr`: the great-circle arc angle + values between the two points. + """ + + # Determine whether to use firedrake or numpy functions + module, _ = firedrake_or_numpy(lon1) + cos = module.cos + sin = module.sin + acos = module.acos if hasattr(module, "acos") else module.arccos + pi = module.pi + + if units == 'deg': + lon1 *= pi / 180.0 + lat1 *= pi / 180.0 + lon2 *= pi / 180.0 + lat2 *= pi / 180.0 + + arc_length = acos(cos(lon1 - lon2)*cos(lat1)*cos(lat2) + sin(lat1)*sin(lat2)) + + if units == 'deg': + arc_length *= 180.0 / pi + + return arc_length diff --git a/gusto/coordinates.py b/gusto/coordinates.py index 7d9abc407..65d115f1c 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -3,8 +3,9 @@ Coordinate fields are stored in specified VectorFunctionSpaces. """ +from gusto.coord_transforms import lonlatr_from_xyz, rotated_lonlatr_coords from gusto.logging import logger -from firedrake import (SpatialCoordinate, sqrt, atan2, asin, Function) +from firedrake import SpatialCoordinate, Function import numpy as np @@ -12,34 +13,38 @@ class Coordinates(object): """ An object for holding and setting up coordinate fields. """ - def __init__(self, mesh): + def __init__(self, mesh, on_sphere=False, rotated_pole=None, radius=None): """ Args: mesh (:class:`Mesh`): the model's domain object. + on_sphere (bool, optional): whether the domain is on the surface of + a sphere. If False, the domain is assumed to be Cartesian. + Defaults to False. + rotated_pole (tuple, optional): a tuple of floats (lon, lat) of the + location to use as the north pole in a spherical coordinate + system. These are expressed in the original coordinate system. + The longitude and latitude must be expressed in radians. + Defaults to None. This is unused for non-spherical domains. + radius (float, optional): the radius of a spherical domain. Defaults + to None. This is unused for non-spherical domains. """ self.mesh = mesh - # TODO: is this the best way of determining whether we are on a sphere? - if hasattr(mesh, "_base_mesh") and hasattr(mesh._base_mesh, 'geometric_dimension'): - on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) - else: - on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) - # -------------------------------------------------------------------- # # Set up spatial coordinate # -------------------------------------------------------------------- # if on_sphere: xyz = SpatialCoordinate(mesh) - r = sqrt(xyz[0]**2 + xyz[1]**2 + xyz[2]**2) - lon = atan2(xyz[1], xyz[0]) - lat = asin(xyz[2]/r) + if rotated_pole is not None: + lon, lat, r = rotated_lonlatr_coords(xyz, rotated_pole) + else: + lon, lat, r = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2]) if mesh.extruded: - # TODO: would we prefer to store height instead of radius? - self.coords = (lon, lat, r) - self.coords_name = ['lon', 'lat', 'r'] + self.coords = (lon, lat, r-radius) + self.coords_name = ['lon', 'lat', 'h'] else: self.coords = (lon, lat) self.coords_name = ['lon', 'lat'] diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index e1d3eff58..d6e1be4cc 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -3,20 +3,21 @@ from firedrake import op2, assemble, dot, dx, Function, sqrt, \ TestFunction, TrialFunction, Constant, grad, inner, curl, \ LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \ - ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, \ + ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, pi, \ TensorFunctionSpace, SpatialCoordinate, as_vector, \ - Projector, Interpolator + Projector, Interpolator, FunctionSpace from firedrake.assign import Assigner from abc import ABCMeta, abstractmethod, abstractproperty import gusto.thermodynamics as tde +from gusto.coord_transforms import rotated_lonlatr_vectors from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations from gusto.active_tracers import TracerVariableType, Phases import numpy as np -__all__ = ["Diagnostics", "CourantNumber", "VelocityX", "VelocityZ", "VelocityY", "Gradient", - "SphericalComponent", "MeridionalComponent", "ZonalComponent", "RadialComponent", +__all__ = ["Diagnostics", "CourantNumber", "Gradient", "XComponent", "YComponent", + "ZComponent", "MeridionalComponent", "ZonalComponent", "RadialComponent", "RichardsonNumber", "Energy", "KineticEnergy", "ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy", "ShallowWaterPotentialEnstrophy", "CompressibleKineticEnergy", "Exner", "Sum", "Difference", "SteadyStateError", @@ -187,7 +188,10 @@ def setup(self, domain, state_fields, space=None): if not self._initialised: if self.space is None: if space is None: - space = domain.spaces("DG0", "DG", 0) + if not hasattr(domain.spaces, "DG0"): + space = domain.spaces.create_space("DG0", "DG", 0) + else: + space = domain.spaces("DG0") self.space = space else: space = self.space @@ -195,7 +199,8 @@ def setup(self, domain, state_fields, space=None): # Add space to domain assert space.name is not None, \ f'Diagnostics {self.name} is using a function space which does not have a name' - domain.spaces(space.name, V=space) + if not hasattr(domain.spaces, space.name): + domain.spaces.add_space(space.name, space) self.field = state_fields(self.name, space=space, dump=self.to_dump, pick_up=False) @@ -300,7 +305,7 @@ def setup(self, domain, state_fields): state_fields (:class:`StateFields`): the model's field container. """ - V = domain.spaces("DG0", "DG", 0) + V = FunctionSpace(domain.mesh, "DG", 0) test = TestFunction(V) cell_volume = Function(V) self.cell_flux = Function(V) @@ -348,58 +353,6 @@ def compute(self): super().compute() -# TODO: unify all component diagnostics -class VelocityX(DiagnosticField): - """The geocentric Cartesian X component of the velocity field.""" - name = "VelocityX" - - 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. - """ - u = state_fields("u") - self.expr = u[0] - super().setup(domain, state_fields) - - -class VelocityZ(DiagnosticField): - """The geocentric Cartesian Z component of the velocity field.""" - name = "VelocityZ" - - 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. - """ - u = state_fields("u") - self.expr = u[domain.mesh.geometric_dimension() - 1] - super().setup(domain, state_fields) - - -class VelocityY(DiagnosticField): - """The geocentric Cartesian Y component of the velocity field.""" - name = "VelocityY" - - 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. - """ - u = state_fields("u") - self.expr = u[1] - super().setup(domain, state_fields) - - class Gradient(DiagnosticField): """Diagnostic for computing the gradient of fields.""" def __init__(self, name, space=None, method='solve'): @@ -495,8 +448,8 @@ def setup(self, domain, state_fields): super().setup(domain, state_fields, space=space) -class SphericalComponent(DiagnosticField): - """Base diagnostic for computing spherical-polar components of fields.""" +class VectorComponent(DiagnosticField): + """Base diagnostic for orthogonal components of vector-valued fields.""" def __init__(self, name, space=None, method='interpolate'): """ Args: @@ -511,31 +464,104 @@ def __init__(self, name, space=None, method='interpolate'): self.fname = name super().__init__(space=space, method=method, required_fields=(name,)) - # TODO: these routines must be moved to somewhere more available generally - # (e.g. initialisation tools?) - def _spherical_polar_unit_vectors(self, domain): + def setup(self, domain, state_fields, unit_vector): """ - Generate ufl expressions for the spherical polar unit vectors. + Sets up the :class:`Function` for the diagnostic field. Args: - domain (:class:`Domain`): the model's domain, containing its mesh. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + unit_vector (:class:`ufl.Expr`): the unit vector to extract the + component for. This assumes an orthogonal coordinate system. + """ + f = state_fields(self.fname) + self.expr = dot(f, unit_vector) + super().setup(domain, state_fields) - Returns: - tuple of (:class:`ufl.Expr`): the zonal, meridional and radial unit - vectors. + +class XComponent(VectorComponent): + """The geocentric Cartesian x-component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_x" + + 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. + """ + dim = domain.mesh.topological_dimension() + e_x = as_vector([Constant(1.0)]+[Constant(0.0)]*(dim-1)) + super().setup(domain, state_fields, e_x) + + +class YComponent(VectorComponent): + """The geocentric Cartesian y-component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_y" + + 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. """ - x, y, z = SpatialCoordinate(domain.mesh) - x_hat = Constant(as_vector([1.0, 0.0, 0.0])) - y_hat = Constant(as_vector([0.0, 1.0, 0.0])) - z_hat = Constant(as_vector([0.0, 0.0, 1.0])) - R = sqrt(x**2 + y**2) # distance from z axis - r = sqrt(x**2 + y**2 + z**2) # distance from origin + assert domain.metadata['domain_type'] not in ['interval', 'vertical_slice'], \ + f'Y-component diagnostic cannot be used with domain {domain.metadata["domain_type"]}' + dim = domain.mesh.topological_dimension() + e_y = as_vector([Constant(0.0), Constant(1.0)]+[Constant(0.0)]*(dim-2)) + super().setup(domain, state_fields, e_y) - lambda_hat = (x * y_hat - y * x_hat) / R - phi_hat = (-x*z/R * x_hat - y*z/R * y_hat + R * z_hat) / r - r_hat = (x * x_hat + y * y_hat + z * z_hat) / r - return lambda_hat, phi_hat, r_hat +class ZComponent(VectorComponent): + """The geocentric Cartesian z-component of a vector-valued field.""" + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_z" + + 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. + """ + assert domain.metadata['domain_type'] not in ['interval', 'plane'], \ + f'Z-component diagnostic cannot be used with domain {domain.metadata["domain_type"]}' + dim = domain.mesh.topological_dimension() + e_x = as_vector([Constant(0.0)]*(dim-1)+[Constant(1.0)]) + super().setup(domain, state_fields, e_x) + + +class SphericalComponent(VectorComponent): + """Base diagnostic for computing spherical-polar components of fields.""" + def __init__(self, name, rotated_pole=None, space=None, method='interpolate'): + """ + Args: + name (str): name of the field to compute the component of. + rotated_pole (tuple of floats, optional): a tuple of floats + (lon, lat) of the new pole, in the original coordinate system. + The longitude and latitude must be expressed in radians. + Defaults to None, corresponding to a pole of (0, pi/2). + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.rotated_pole = (0.0, pi/2) if rotated_pole is None else rotated_pole + super().__init__(name=name, space=space, method=method) def _check_args(self, domain, field): """ @@ -572,9 +598,9 @@ def setup(self, domain, state_fields): """ f = state_fields(self.fname) self._check_args(domain, f) - _, phi_hat, _ = self._spherical_polar_unit_vectors(domain) - self.expr = dot(f, phi_hat) - super().setup(domain, state_fields) + xyz = SpatialCoordinate(domain.mesh) + _, e_lat, _ = rotated_lonlatr_vectors(xyz, self.rotated_pole) + super().setup(domain, state_fields, e_lat) class ZonalComponent(SphericalComponent): @@ -594,9 +620,9 @@ def setup(self, domain, state_fields): """ f = state_fields(self.fname) self._check_args(domain, f) - lambda_hat, _, _ = self._spherical_polar_unit_vectors(domain) - self.expr = dot(f, lambda_hat) - super().setup(domain, state_fields) + xyz = SpatialCoordinate(domain.mesh) + e_lon, _, _ = rotated_lonlatr_vectors(xyz, self.rotated_pole) + super().setup(domain, state_fields, e_lon) class RadialComponent(SphericalComponent): @@ -616,9 +642,9 @@ def setup(self, domain, state_fields): """ f = state_fields(self.fname) self._check_args(domain, f) - _, _, r_hat = self._spherical_polar_unit_vectors(domain) - self.expr = dot(f, r_hat) - super().setup(domain, state_fields) + xyz = SpatialCoordinate(domain.mesh) + _, _, e_r = rotated_lonlatr_vectors(xyz, self.rotated_pole) + super().setup(domain, state_fields, e_r) class RichardsonNumber(DiagnosticField): @@ -1399,7 +1425,11 @@ def compute(self): class Precipitation(DiagnosticField): - """The total precipitation falling through the domain's bottom surface.""" + """ + The total precipitation falling through the domain's bottom surface. + + This is normalised by unit area, giving a result in kg / m^2. + """ name = "Precipitation" def __init__(self): @@ -1415,33 +1445,43 @@ def setup(self, domain, state_fields): domain (:class:`Domain`): the model's domain object. state_fields (:class:`StateFields`): the model's field container. """ - space = domain.spaces("DG0", "DG", 0) - assert space.extruded, 'Cannot compute precipitation on a non-extruded mesh' + if not hasattr(domain.spaces, "DG0"): + DG0 = domain.spaces.create_space("DG0", "DG", 0) + else: + DG0 = domain.spaces("DG0") + assert DG0.extruded, 'Cannot compute precipitation on a non-extruded mesh' + self.space = DG0 + + # Gather fields rain = state_fields('rain') rho = state_fields('rho') v = state_fields('rainfall_velocity') # Set up problem - self.phi = TestFunction(space) - flux = TrialFunction(space) - n = FacetNormal(domain.mesh) - un = 0.5 * (dot(v, n) + abs(dot(v, n))) - self.flux = Function(space) + self.phi = TestFunction(DG0) + flux = TrialFunction(DG0) + self.flux = Function(DG0) # Flux to solve for + area = Function(DG0) # Need to compute normalisation (area) - a = self.phi * flux * dx - L = self.phi * rain * un * rho * (ds_b + ds_t + ds_v) + eqn_lhs = self.phi * flux * dx + area_rhs = self.phi * ds_b + eqn_rhs = domain.dt * self.phi * (rain * dot(- v, domain.k) * rho / area) * ds_b + + # Compute area normalisation + area_prob = LinearVariationalProblem(eqn_lhs, area_rhs, area) + area_solver = LinearVariationalSolver(area_prob) + area_solver.solve() # setup solver - problem = LinearVariationalProblem(a, L, self.flux) - self.solver = LinearVariationalSolver(problem) - self.space = space - self.field = state_fields(self.name, space=space, dump=True, pick_up=False) - # TODO: might we want to pick up this field? Otherwise initialise to zero + rain_prob = LinearVariationalProblem(eqn_lhs, eqn_rhs, self.flux) + self.solver = LinearVariationalSolver(rain_prob) + self.field = state_fields(self.name, space=DG0, dump=True, pick_up=True) + # Initialise field to zero, if picking up this will be overridden self.field.assign(0.0) def compute(self): - """Compute the diagnostic field from the current state.""" + """Increment the precipitation diagnostic.""" self.solver.solve() - self.field.assign(self.field + assemble(self.flux * self.phi * dx)) + self.field.assign(self.field + self.flux) class Vorticity(DiagnosticField): diff --git a/gusto/domain.py b/gusto/domain.py index 090a5c01e..42ade2b53 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -7,7 +7,8 @@ from gusto.coordinates import Coordinates from gusto.function_spaces import Spaces, check_degree_args from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, - inner, grad, VectorFunctionSpace, Function, perp) + inner, grad, VectorFunctionSpace, Function, FunctionSpace, + perp) import numpy as np @@ -24,7 +25,8 @@ class Domain(object): the same then this can be specified through the "degree" argument. """ def __init__(self, mesh, dt, family, degree=None, - horizontal_degree=None, vertical_degree=None): + horizontal_degree=None, vertical_degree=None, + rotated_pole=None): """ Args: mesh (:class:`Mesh`): the model's mesh. @@ -40,6 +42,11 @@ def __init__(self, mesh, dt, family, degree=None, horizontal part of the DG space. Defaults to None. vertical_degree (int, optional): the element degree used for the vertical part of the DG space. Defaults to None. + rotated_pole (tuple, optional): a tuple of floats (lon, lat) of the + location to use as the north pole in a spherical coordinate + system. These are expressed in the original coordinate system. + The longitude and latitude must be expressed in radians. + Defaults to None. This is unused for non-spherical domains. Raises: ValueError: if incompatible degrees are specified (e.g. specifying @@ -58,6 +65,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 # -------------------------------------------------------------------- # @@ -73,6 +83,7 @@ def __init__(self, mesh, dt, family, degree=None, self.spaces = Spaces(mesh) self.spaces.build_compatible_spaces(self.family, self.horizontal_degree, self.vertical_degree) + self.spaces.build_dg1_equispaced() # -------------------------------------------------------------------- # # Determine some useful aspects of domain @@ -93,7 +104,7 @@ def __init__(self, mesh, dt, family, degree=None, R = sqrt(inner(x, x)) self.k = grad(R) if dim == 2: - if hasattr(mesh, "_bash_mesh"): + if hasattr(mesh, "_base_mesh"): sphere_degree = mesh._base_mesh.coordinates.function_space().ufl_element().degree() else: if not hasattr(mesh, "_cell_orientations"): @@ -109,11 +120,28 @@ def __init__(self, mesh, dt, family, degree=None, if dim == 2: self.perp = perp + # -------------------------------------------------------------------- # + # Construct information relating to height/radius + # -------------------------------------------------------------------- # + + 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) + radius_field.interpolate(r_shell) + # TODO: this should use global min kernel + radius = np.min(radius_field.dat.data_ro) + else: + radius = None + # -------------------------------------------------------------------- # # Set up coordinates # -------------------------------------------------------------------- # - self.coords = Coordinates(mesh) + self.coords = Coordinates(mesh, on_sphere=self.on_sphere, + rotated_pole=rotated_pole, radius=radius) # Set up DG1 equispaced space, used for making metadata _ = self.spaces('DG1_equispaced') self.coords.register_space(self, 'DG1_equispaced') @@ -122,41 +150,58 @@ def __init__(self, mesh, dt, family, degree=None, # Construct metadata about domain # -------------------------------------------------------------------- # - self.metadata = {} - self.metadata['extruded'] = mesh.extruded - - if self.on_sphere and hasattr(mesh, "_base_mesh"): - self.metadata['domain_type'] = 'extruded_spherical_shell' - elif self.on_sphere: - self.metadata['domain_type'] = 'spherical_shell' - elif mesh.geometric_dimension() == 1 and mesh.topological_dimension() == 1: - self.metadata['domain_type'] = 'interval' - elif mesh.geometric_dimension() == 2 and mesh.topological_dimension() == 2 and mesh.extruded: - self.metadata['domain_type'] = 'vertical_slice' - elif mesh.geometric_dimension() == 2 and mesh.topological_dimension() == 2: - self.metadata['domain_type'] = 'plane' - elif mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 3 and mesh.extruded: - self.metadata['domain_type'] = 'extruded_plane' - else: - raise ValueError('Unable to determine domain type') - - comm = self.mesh.comm - my_rank = comm.Get_rank() + self.metadata = construct_domain_metadata(mesh, self.coords, self.on_sphere) - # Properties of domain will be determined from full coords, so need - # doing on the first processor then broadcasting to others - if my_rank == 0: - chi = self.coords.global_chi_coords['DG1_equispaced'] - if not self.on_sphere: - self.metadata['domain_extent_x'] = np.max(chi[0, :]) - np.min(chi[0, :]) - if self.metadata['domain_type'] in ['plane', 'extruded_plane']: - self.metadata['domain_extent_y'] = np.max(chi[1, :]) - np.min(chi[1, :]) - if mesh.extruded: - self.metadata['domain_extent_z'] = np.max(chi[-1, :]) - np.min(chi[-1, :]) +def construct_domain_metadata(mesh, coords, on_sphere): + """ + Builds a dictionary containing metadata about the domain. - else: - self.metadata = {} + Args: + mesh (:class:`Mesh`): the model's mesh. + coords (:class:`Coordinates`): the model's coordinate object. + on_sphere (bool): whether the domain is on the sphere or not. - # Send information to other processors - self.metadata = comm.bcast(self.metadata, root=0) + Returns: + dict: a dictionary of metadata relating to the domain. + """ + metadata = {} + metadata['extruded'] = mesh.extruded + + if on_sphere and hasattr(mesh, "_base_mesh"): + metadata['domain_type'] = 'extruded_spherical_shell' + elif on_sphere: + metadata['domain_type'] = 'spherical_shell' + elif mesh.geometric_dimension() == 1 and mesh.topological_dimension() == 1: + metadata['domain_type'] = 'interval' + elif mesh.geometric_dimension() == 2 and mesh.topological_dimension() == 2 and mesh.extruded: + metadata['domain_type'] = 'vertical_slice' + elif mesh.geometric_dimension() == 2 and mesh.topological_dimension() == 2: + metadata['domain_type'] = 'plane' + elif mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 3 and mesh.extruded: + metadata['domain_type'] = 'extruded_plane' + else: + raise ValueError('Unable to determine domain type') + + comm = mesh.comm + my_rank = comm.Get_rank() + + # Properties of domain will be determined from full coords, so need + # doing on the first processor then broadcasting to others + + if my_rank == 0: + chi = coords.global_chi_coords['DG1_equispaced'] + if not on_sphere: + metadata['domain_extent_x'] = np.max(chi[0, :]) - np.min(chi[0, :]) + if metadata['domain_type'] in ['plane', 'extruded_plane']: + metadata['domain_extent_y'] = np.max(chi[1, :]) - np.min(chi[1, :]) + if mesh.extruded: + metadata['domain_extent_z'] = np.max(chi[-1, :]) - np.min(chi[-1, :]) + + else: + metadata = {} + + # Send information to other processors + metadata = comm.bcast(metadata, root=0) + + return metadata diff --git a/gusto/equations.py b/gusto/equations.py index dd4a9ceb4..1c50bd4e6 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -84,9 +84,8 @@ def __init__(self, domain, function_space, field_name, Vu=None): super().__init__(domain, function_space, field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu, overwrite_space=True) - else: - V = domain.spaces("HDiv") + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") u = self.prescribed_fields("u", V) test = self.test @@ -115,9 +114,8 @@ def __init__(self, domain, function_space, field_name, Vu=None): super().__init__(domain, function_space, field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu, overwrite_space=True) - else: - V = domain.spaces("HDiv") + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") u = self.prescribed_fields("u", V) test = self.test @@ -174,9 +172,8 @@ def __init__(self, domain, function_space, field_name, Vu=None, super().__init__(domain, function_space, field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu, overwrite_space=True) - else: - V = domain.spaces("HDiv") + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") u = self.prescribed_fields("u", V) test = self.test @@ -552,9 +549,8 @@ def __init__(self, domain, active_tracers, Vu=None): PrognosticEquation.__init__(self, domain, W, full_field_name) if Vu is not None: - V = domain.spaces("HDiv", V=Vu, overwrite_space=True) - else: - V = domain.spaces("HDiv") + domain.spaces.add_space("HDiv", Vu, overwrite_space=True) + V = domain.spaces("HDiv") _ = self.prescribed_fields("u", V) self.tests = TestFunctions(W) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 8f23f2a0d..185273420 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -19,7 +19,9 @@ 'RTCF': 'RTCE', 'RTCE': 'RTCE', 'CG': 'DG', - 'BDFM': None} + 'BDFM': None, + 'BDMCF': 'BDMCE', + 'BDMCE': 'BDMCE'} # HCurl spaces are keys, HDiv spaces are values # Can't just reverse the other dictionary as values are not necessarily unique @@ -32,7 +34,9 @@ 'RTCE': 'RTCF', 'RTCF': 'RTCF', 'CG': 'CG', - 'BDFM': 'BDFM'} + 'BDFM': 'BDFM', + 'BDMCF': 'BDMCF', + 'BDMCE': 'BDMCF'} # Degree to use for H1 space for a particular family @@ -49,7 +53,7 @@ def h1_degree(family, l2_degree): """ if family in ['CG', 'RT', 'RTE', 'RTF', 'RTCE', 'RTCF']: return l2_degree + 1 - elif family in ['BDM', 'BDME', 'BDMF']: + elif family in ['BDM', 'BDME', 'BDMF', 'BDMCE', 'BDMCF']: return l2_degree + 2 elif family == 'BDFM': return l2_degree + 1 @@ -66,38 +70,53 @@ def __init__(self, mesh): """ self.mesh = mesh self.extruded_mesh = hasattr(mesh, "_base_mesh") - self._initialised_base_spaces = False + self.de_rham_complex = {} - def __call__(self, name, family=None, degree=None, - horizontal_degree=None, vertical_degree=None, - V=None, overwrite_space=False): + def __call__(self, name): """ - Returns a space, and also creates it if it is not created yet. + Returns a space from the space container. - If a space needs creating, it may be that more arguments (such as the - family and degree) need to be provided. Alternatively a space can be - passed in to be stored in the space container. + Args: + name (str): the name of the space. + + Returns: + :class:`FunctionSpace`: the desired function space. + """ - For extruded meshes, it is possible to seperately specify the horizontal - and vertical degrees of the elements. Alternatively, if these degrees - should be the same then this can be specified through the "degree" - argument. + if hasattr(self, name): + return getattr(self, name) + + else: + raise ValueError(f'The space container has no space {name}') + + def add_space(self, name, space, overwrite_space=False): + """ + Adds a function space to the container. Args: name (str): the name of the space. - family (str, optional): name of the finite element family to be - created. Defaults to None. - degree (int, optional): the element degree used for the space. - Defaults to None, in which case the horizontal degree must be - provided. - horizontal_degree (int, optional): the horizontal degree of the - finite element space to be created. Defaults to None. - vertical_degree (int, optional): the vertical degree of the - finite element space to be created. Defaults to None. - V (:class:`FunctionSpace`, optional): an existing space, to be - stored in the creator object. If this is provided, it will be - added to the creator and no other action will be taken. This - space will be returned. Defaults to None. + space (:class:`FunctionSpace`): the function space to be added to + the Space container.. + overwrite_space (bool, optional): Logical to allow space existing in + container to be overwritten by an incoming space. Defaults to + False. + """ + + if hasattr(self, name) and not overwrite_space: + raise RuntimeError(f'Space {name} already exists. If you really ' + + 'to create it then set `overwrite_space` as ' + + 'to be True') + + setattr(self, name, space) + + def create_space(self, name, family, degree, overwrite_space=False): + """ + Creates a space and adds it to the container.. + + Args: + name (str): the name to give to the space. + family (str): name of the finite element family to be created. + degree (int): the element degree used for the space. overwrite_space (bool, optional): Logical to allow space existing in container to be overwritten by an incoming space. Defaults to False. @@ -106,57 +125,160 @@ def __call__(self, name, family=None, degree=None, :class:`FunctionSpace`: the desired function space. """ + if hasattr(self, name) and not overwrite_space: + raise RuntimeError(f'Space {name} already exists. If you really ' + + 'to create it then set `overwrite_space` as ' + + 'to be True') + + space = FunctionSpace(self.mesh, family, degree, name=name) + setattr(self, name, space) + return space + + def build_compatible_spaces(self, family, horizontal_degree, + vertical_degree=None, complex_name=None): + """ + Builds the compatible spaces associated with some de Rham complex, and + sets the spaces as attributes to the underlying :class:`Spaces` object. + + Args: + family (str): the family of the compatible spaces. This is either + the horizontal element of the HDiv or HCurl spaces. + horizontal_degree (int): the horizontal degree of the L2 space. + vertical_degree (int, optional): the vertical degree of the L2 + space. Defaults to None. + complex_name (str, optional): optional name to pass to the + :class:`DeRhamComplex` object to be created. Defaults to None. + """ + + if vertical_degree is not None: + degree_key = (horizontal_degree, vertical_degree) + else: + degree_key = horizontal_degree + + # Determine suffix to spaces + if complex_name is None and len(self.de_rham_complex.keys()) > 0: + if vertical_degree is not None and horizontal_degree != vertical_degree: + complex_name = f'{horizontal_degree}_{vertical_degree}' + else: + complex_name = f'{horizontal_degree}' + elif complex_name is None: + complex_name = '' + + if degree_key in self.de_rham_complex.keys(): + raise RuntimeError(f'de Rham complex for degree {degree_key} has ' + + 'already been created. Cannot create again.') + + # Create the compatible space objects + de_rham_complex = DeRhamComplex(self.mesh, family, horizontal_degree, + vertical_degree=vertical_degree, + complex_name=complex_name) + + # Set the spaces as attributes of the space container + self.de_rham_complex[degree_key] = de_rham_complex + setattr(self, "H1"+complex_name, de_rham_complex.H1) + setattr(self, "HCurl"+complex_name, de_rham_complex.HCurl) + setattr(self, "HDiv"+complex_name, de_rham_complex.HDiv) + setattr(self, "L2"+complex_name, de_rham_complex.L2) + # Register L2 space as DG also + setattr(self, "DG"+complex_name, de_rham_complex.L2) + if hasattr(de_rham_complex, "theta"+complex_name): + setattr(self, "theta"+complex_name, de_rham_complex.theta) + + def build_dg1_equispaced(self): + """ + Builds the equispaced variant of the DG1 function space, which is used in + recovered finite element schemes. + + Returns: + (:class:`FunctionSpace`): the equispaced DG1 function space. + """ + + if self.extruded_mesh: + cell = self.mesh._base_mesh.ufl_cell().cellname() + hori_elt = FiniteElement('DG', cell, 1, variant='equispaced') + vert_elt = FiniteElement('DG', interval, 1, variant='equispaced') + V_elt = TensorProductElement(hori_elt, vert_elt) + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement('DG', cell, 1, variant='equispaced') + + space = FunctionSpace(self.mesh, V_elt, name='DG1_equispaced') + setattr(self, 'DG1_equispaced', space) + return space + + +class DeRhamComplex(object): + """Constructs and stores the function spaces forming a de Rham complex.""" + def __init__(self, mesh, family, horizontal_degree, vertical_degree=None, + complex_name=None): + """ + Args: + mesh (:class:`Mesh`): _description_ + family (str): name of the finite element family to be + created. Defaults to None. + horizontal_degree (int): the horizontal degree of the finite element + space to be created. + vertical_degree (int, optional): the vertical degree of the + finite element space to be created. Defaults to None. + complex_name (str, optional): suffix to give to the spaces created + in this complex. Defaults to None. + """ + implemented_families = ["DG", "CG", "RT", "RTF", "RTE", "RTCF", "RTCE", - "BDM", "BDMF", "BDME", "BDFM"] + "BDM", "BDMF", "BDME", "BDFM", "BDMCE", "BDMCF"] if family not in [None]+implemented_families: raise NotImplementedError(f'family {family} either not recognised ' + 'or implemented in Gusto') - if hasattr(self, name) and (V is None or not overwrite_space): - # We have requested a space that should already have been created - if V is not None: - assert getattr(self, name) == V, \ - f'There is a conflict between the space {name} already ' + \ - 'existing in the space container, and the space being passed to it' - return getattr(self, name) + self.mesh = mesh + self.extruded_mesh = hasattr(mesh, '_base_mesh') + self.family = family + self.complex_name = complex_name + self.build_base_spaces(family, horizontal_degree, vertical_degree) + self.build_compatible_spaces() + def build_base_spaces(self, family, horizontal_degree, vertical_degree=None): + """ + Builds the :class:`FiniteElement` objects for the base mesh. + + Args: + family (str): the family of the horizontal part of either the HDiv + or HCurl space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the L2 space. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the L2 space. Defaults to None. + """ + + if self.extruded_mesh: + cell = self.mesh._base_mesh.ufl_cell().cellname() else: - # Space does not exist in creator - if V is not None: - # The space itself has been provided (to add it to the creator) - value = V + cell = self.mesh.ufl_cell().cellname() - elif name == "DG1_equispaced": - # Special case as no degree arguments need providing - value = self.build_l2_space(1, 1, variant='equispaced', name='DG1_equispaced') + hdiv_family = hcurl_hdiv_dict[family] + hcurl_family = hdiv_hcurl_dict[family] - else: - check_degree_args('Spaces', self.mesh, degree, horizontal_degree, vertical_degree) - - # Convert to horizontal and vertical degrees - horizontal_degree = degree if horizontal_degree is None else horizontal_degree - vertical_degree = degree if vertical_degree is None else vertical_degree - - # Loop through name and family combinations - if name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF", "RTF", "BDMF"]: - hdiv_family = hcurl_hdiv_dict[family] - value = self.build_hdiv_space(hdiv_family, horizontal_degree, vertical_degree) - elif name == "HCurl" and family in ["BDM", "RT", "CG", "RTCE", "RTE", "BDME"]: - hcurl_family = hdiv_hcurl_dict[family] - value = self.build_hcurl_space(hcurl_family, horizontal_degree, vertical_degree) - elif name == "theta": - value = self.build_theta_space(horizontal_degree, vertical_degree) - elif family == "DG": - value = self.build_l2_space(horizontal_degree, vertical_degree, name=name) - elif family == "CG": - value = self.build_h1_space(horizontal_degree, vertical_degree, name=name) - else: - raise ValueError(f'There is no space corresponding to {name}') - setattr(self, name, value) - return value + # Horizontal base spaces + self.base_elt_hori_hdiv = FiniteElement(hdiv_family, cell, horizontal_degree+1) + if hcurl_family is not None: + self.base_elt_hori_hcurl = FiniteElement(hcurl_family, cell, horizontal_degree+1) + self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) - def build_compatible_spaces(self, family, horizontal_degree, - vertical_degree=None): + # Special handling of CG horizontal element based on family + if hdiv_family == 'BDMCF': + self.base_elt_hori_cg = FiniteElement("S", cell, h1_degree(family, horizontal_degree)) + else: + self.base_elt_hori_cg = FiniteElement("CG", cell, h1_degree(family, horizontal_degree)) + if hdiv_family == "BDFM": + # Need to enhance this element with bubble element + self.base_elt_hori_cg += FiniteElement("Bubble", cell, horizontal_degree+2) + + # Vertical base spaces + if vertical_degree is not None: + self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) + self.base_elt_vert_dg = FiniteElement("DG", interval, vertical_degree) + + def build_compatible_spaces(self): """ Builds the sequence of compatible finite element spaces for the mesh. @@ -177,26 +299,19 @@ def build_compatible_spaces(self, family, horizontal_degree, Returns: tuple: the created compatible :class:`FunctionSpace` objects. """ - if self._initialised_base_spaces: - pass - elif self.extruded_mesh and not self._initialised_base_spaces: - # Base spaces need building, while horizontal and vertical degrees + if self.extruded_mesh: + # Horizontal and vertical degrees # need specifying separately. Vtheta needs returning. - self.build_base_spaces(family, horizontal_degree, vertical_degree) - Vcg = self.build_h1_space(h1_degree(family, horizontal_degree), - h1_degree(family, vertical_degree), name='H1') + Vcg = self.build_h1_space() setattr(self, "H1", Vcg) - hcurl_family = hdiv_hcurl_dict[family] - Vcurl = self.build_hcurl_space(hcurl_family, horizontal_degree, vertical_degree) + Vcurl = self.build_hcurl_space() setattr(self, "HCurl", Vcurl) - hdiv_family = hcurl_hdiv_dict[family] - Vu = self.build_hdiv_space(hdiv_family, horizontal_degree, vertical_degree) + Vu = self.build_hdiv_space() setattr(self, "HDiv", Vu) - Vdg = self.build_l2_space(horizontal_degree, vertical_degree, name='L2') + Vdg = self.build_l2_space() setattr(self, "L2", Vdg) - setattr(self, "DG", Vdg) # Register this as "L2" and "DG" - Vth = self.build_theta_space(horizontal_degree, vertical_degree) + Vth = self.build_theta_space() setattr(self, "theta", Vth) return Vcg, Vcurl, Vu, Vdg, Vth @@ -205,70 +320,30 @@ def build_compatible_spaces(self, family, horizontal_degree, # 2D: two de Rham complexes (hcurl or hdiv) with 3 spaces # 3D: one de Rham complexes with 4 spaces # either way, build all spaces - Vcg = self.build_h1_space(h1_degree(family, horizontal_degree), name='H1') + Vcg = self.build_h1_space() setattr(self, "H1", Vcg) - hcurl_family = hdiv_hcurl_dict[family] - Vcurl = self.build_hcurl_space(hcurl_family, horizontal_degree+1) + Vcurl = self.build_hcurl_space() setattr(self, "HCurl", Vcurl) - hdiv_family = hcurl_hdiv_dict[family] - Vu = self.build_hdiv_space(family, horizontal_degree+1) + Vu = self.build_hdiv_space() setattr(self, "HDiv", Vu) - Vdg = self.build_l2_space(horizontal_degree, vertical_degree, name='L2') + Vdg = self.build_l2_space() setattr(self, "L2", Vdg) - setattr(self, "DG", Vdg) # Register this as "L2" and "DG" return Vcg, Vcurl, Vu, Vdg else: # 1D domain, de Rham complex has 2 spaces # CG, hdiv and hcurl spaces should be the same - Vcg = self.build_h1_space(horizontal_degree+1, name='H1') + Vcg = self.build_h1_space() setattr(self, "H1", Vcg) setattr(self, "HCurl", None) setattr(self, "HDiv", Vcg) - Vdg = self.build_l2_space(horizontal_degree, name='L2') + Vdg = self.build_l2_space() setattr(self, "L2", Vdg) - setattr(self, "DG", Vdg) # Register this as "L2" and "DG" return Vcg, Vdg - def build_base_spaces(self, family, horizontal_degree, vertical_degree): - """ - Builds the :class:`FiniteElement` objects for the base mesh. - - Args: - family (str): the family of the horizontal part of either the HDiv - or HCurl space. - horizontal_degree (int): the polynomial degree of the horizontal - part of the L2 space. - vertical_degree (int): the polynomial degree of the vertical part of - the L2 space. - """ - - if family == 'BDFM': - # Need a special implementation of base spaces here as it does not - # fit the same pattern as other spaces - self.build_bdfm_base_spaces(horizontal_degree, vertical_degree) - return - - cell = self.mesh._base_mesh.ufl_cell().cellname() - - hdiv_family = hcurl_hdiv_dict[family] - hcurl_family = hdiv_hcurl_dict[family] - - # horizontal base spaces - self.base_elt_hori_hdiv = FiniteElement(hdiv_family, cell, horizontal_degree+1) - self.base_elt_hori_hcurl = FiniteElement(hcurl_family, cell, horizontal_degree+1) - self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) - self.base_elt_hori_cg = FiniteElement("CG", cell, h1_degree(family, horizontal_degree)) - - # vertical base spaces - self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) - self.base_elt_vert_dg = FiniteElement("DG", interval, vertical_degree) - - self._initialised_base_spaces = True - - def build_hcurl_space(self, family, horizontal_degree, vertical_degree=None): + def build_hcurl_space(self): """ Builds and returns the HCurl :class:`FunctionSpace`. @@ -283,98 +358,55 @@ def build_hcurl_space(self, family, horizontal_degree, vertical_degree=None): Returns: :class:`FunctionSpace`: the HCurl space. """ - if family is None: + if hdiv_hcurl_dict[self.family] is None: logger.warning('There is no HCurl space for this family. Not creating one') return None if self.extruded_mesh: - if not self._initialised_base_spaces: - if vertical_degree is None: - raise ValueError('vertical_degree must be specified to create HCurl space on an extruded mesh') - self.build_base_spaces(family, horizontal_degree, vertical_degree) - Vh_elt = HCurl(TensorProductElement(self.base_elt_hori_hcurl, self.base_elt_vert_cg)) - Vv_elt = HCurl(TensorProductElement(self.base_elt_hori_cg, self.base_elt_vert_dg)) + Vh_elt = HCurl(TensorProductElement(self.base_elt_hori_hcurl, + self.base_elt_vert_cg)) + Vv_elt = HCurl(TensorProductElement(self.base_elt_hori_cg, + self.base_elt_vert_dg)) V_elt = Vh_elt + Vv_elt else: - cell = self.mesh.ufl_cell().cellname() - hcurl_family = hdiv_hcurl_dict[family] - V_elt = FiniteElement(hcurl_family, cell, horizontal_degree) + V_elt = self.base_elt_hori_hcurl - return FunctionSpace(self.mesh, V_elt, name='HCurl') + return FunctionSpace(self.mesh, V_elt, name='HCurl'+self.complex_name) - def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): + def build_hdiv_space(self): """ Builds and returns the HDiv :class:`FunctionSpace`. - Args: - family (str): the family of the horizontal part of the HDiv space. - horizontal_degree (int): the polynomial degree of the horizontal - part of the L2 space from the de Rham complex. - vertical_degree (int, optional): the polynomial degree of the - vertical part of the L2 space from the de Rham complex. - Defaults to None. Must be specified if the mesh is extruded. - Returns: :class:`FunctionSpace`: the HDiv space. """ if self.extruded_mesh: - if not self._initialised_base_spaces: - if vertical_degree is None: - raise ValueError('vertical_degree must be specified to create HDiv space on an extruded mesh') - self.build_base_spaces(family, horizontal_degree, vertical_degree) - Vh_elt = HDiv(TensorProductElement(self.base_elt_hori_hdiv, self.base_elt_vert_dg)) - Vt_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_cg) + Vh_elt = HDiv(TensorProductElement(self.base_elt_hori_hdiv, + self.base_elt_vert_dg)) + Vt_elt = TensorProductElement(self.base_elt_hori_dg, + self.base_elt_vert_cg) Vv_elt = HDiv(Vt_elt) V_elt = Vh_elt + Vv_elt else: - cell = self.mesh.ufl_cell().cellname() - hdiv_family = hcurl_hdiv_dict[family] - V_elt = FiniteElement(hdiv_family, cell, horizontal_degree) - return FunctionSpace(self.mesh, V_elt, name='HDiv') + V_elt = self.base_elt_hori_hdiv + return FunctionSpace(self.mesh, V_elt, name='HDiv'+self.complex_name) - def build_l2_space(self, horizontal_degree, vertical_degree=None, variant=None, name='L2'): + def build_l2_space(self): """ Builds and returns the discontinuous L2 :class:`FunctionSpace`. - Args: - horizontal_degree (int): the polynomial degree of the horizontal - part of the L2 space. - vertical_degree (int, optional): the polynomial degree of the - vertical part of the L2 space. Defaults to None. Must be - specified if the mesh is extruded. - variant (str, optional): the variant of the underlying - :class:`FiniteElement` to use. Defaults to None, which will call - the default variant. - name (str, optional): name to assign to the function space. Default - is "L2". - Returns: :class:`FunctionSpace`: the L2 space. """ - assert not hasattr(self, name), f'There already exists a function space with name {name}' if self.extruded_mesh: - if vertical_degree is None: - raise ValueError('vertical_degree must be specified to create L2 space on an extruded mesh') - if (not self._initialised_base_spaces - or self.base_elt_vert_dg.degree() != vertical_degree - or self.base_elt_vert_dg.variant() != variant - or self.base_elt_hori_dg.degree() != horizontal_degree - or self.base_elt_hori_dg.degree() != variant): - cell = self.mesh._base_mesh.ufl_cell().cellname() - base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree, variant=variant) - base_elt_vert_dg = FiniteElement("DG", interval, vertical_degree, variant=variant) - else: - base_elt_hori_dg = self.base_elt_hori_dg - base_elt_vert_dg = self.base_elt_vert_dg - V_elt = TensorProductElement(base_elt_hori_dg, base_elt_vert_dg) + V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_dg) else: - cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement("DG", cell, horizontal_degree, variant=variant) + V_elt = self.base_elt_hori_dg - return FunctionSpace(self.mesh, V_elt, name=name) + return FunctionSpace(self.mesh, V_elt, name='L2'+self.complex_name) - def build_theta_space(self, horizontal_degree, vertical_degree): + def build_theta_space(self): """ Builds and returns the 'theta' space. @@ -382,12 +414,6 @@ def build_theta_space(self, horizontal_degree, vertical_degree): of the velocity. The space will be discontinuous in the horizontal but continuous in the vertical. - Args: - horizontal_degree (int): the polynomial degree of the horizontal - part of the L2 space from the de Rham complex. - vertical_degree (int): the polynomial degree of the vertical part of - the L2 space from the de Rham complex. - Raises: AssertionError: the mesh is not extruded. @@ -395,79 +421,30 @@ def build_theta_space(self, horizontal_degree, vertical_degree): :class:`FunctionSpace`: the 'theta' space. """ assert self.extruded_mesh, 'Cannot create theta space if mesh is not extruded' - if not self._initialised_base_spaces: - cell = self.mesh._base_mesh.ufl_cell().cellname() - self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) - self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) + V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_cg) - return FunctionSpace(self.mesh, V_elt, name='theta') - def build_h1_space(self, horizontal_degree, vertical_degree=None, name='H1'): + return FunctionSpace(self.mesh, V_elt, name='theta'+self.complex_name) + + def build_h1_space(self): """ Builds the continuous scalar space at the top of the de Rham complex. Args: - horizontal_degree (int): the polynomial degree of the horizontal - part of the H1 space. - vertical_degree (int, optional): the polynomial degree of the - vertical part of the H1 space. Defaults to None. Must be - specified if the mesh is extruded. name (str, optional): name to assign to the function space. Default is "H1". Returns: :class:`FunctionSpace`: the continuous space. """ - assert not hasattr(self, name), f'There already exists a function space with name {name}' if self.extruded_mesh: - if vertical_degree is None: - raise ValueError('vertical_degree must be specified to create H1 space on an extruded mesh') - if (not self._initialised_base_spaces - or self.base_elt_vert_cg.degree() != vertical_degree - or self.base_elt_hori_cg.degree() != horizontal_degree): - cell = self.mesh._base_mesh.ufl_cell().cellname() - base_elt_hori_cg = FiniteElement("CG", cell, horizontal_degree) - base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree) - else: - base_elt_hori_cg = self.base_elt_hori_cg - base_elt_vert_cg = self.base_elt_vert_cg - V_elt = TensorProductElement(base_elt_hori_cg, base_elt_vert_cg) - else: - cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement("CG", cell, horizontal_degree) - - return FunctionSpace(self.mesh, V_elt, name=name) - - def build_bdfm_base_spaces(self, horizontal_degree, vertical_degree): - """ - Builds the :class:`FiniteElement` objects for the base mesh when using - the . - - Args: - horizontal_degree (int): the polynomial degree of the horizontal - part of the L2 space. - vertical_degree (int): the polynomial degree of the vertical part of - the L2 space. - """ - - cell = self.mesh._base_mesh.ufl_cell().cellname() + V_elt = TensorProductElement(self.base_elt_hori_cg, self.base_elt_vert_cg) - hdiv_family = 'BDFM' - - # horizontal base spaces - self.base_elt_hori_hdiv = FiniteElement(hdiv_family, cell, horizontal_degree+1) - self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) - - # Add bubble space - self.base_elt_hori_cg = FiniteElement("CG", cell, horizontal_degree+1) - self.base_elt_hori_cg += FiniteElement("Bubble", cell, horizontal_degree+2) - - # vertical base spaces - self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) - self.base_elt_vert_dg = FiniteElement("DG", interval, vertical_degree) + else: + V_elt = self.base_elt_hori_cg - self._initialised_base_spaces = True + return FunctionSpace(self.mesh, V_elt, name='H1'+self.complex_name) def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): diff --git a/gusto/initialisation_tools.py b/gusto/initialisation_tools.py index b4a0da7e2..1a5c5db93 100644 --- a/gusto/initialisation_tools.py +++ b/gusto/initialisation_tools.py @@ -1,65 +1,22 @@ """Tools for computing initial conditions, such as hydrostatic balance.""" from firedrake import MixedFunctionSpace, TrialFunctions, TestFunctions, \ - TestFunction, TrialFunction, SpatialCoordinate, \ + TestFunction, TrialFunction, \ FacetNormal, inner, div, dx, ds_b, ds_t, DirichletBC, \ Function, Constant, \ LinearVariationalProblem, LinearVariationalSolver, \ NonlinearVariationalProblem, NonlinearVariationalSolver, split, solve, \ - sin, cos, sqrt, asin, atan2, as_vector, min_value, max_value, FunctionSpace, \ - errornorm, zero + FunctionSpace, errornorm, zero from gusto import thermodynamics from gusto.logging import logger from gusto.recovery import Recoverer, BoundaryMethod -__all__ = ["latlon_coords", "sphere_to_cartesian", "incompressible_hydrostatic_balance", +__all__ = ["incompressible_hydrostatic_balance", "compressible_hydrostatic_balance", "remove_initial_w", "saturated_hydrostatic_balance", "unsaturated_hydrostatic_balance"] -# TODO: maybe coordinate transforms could go elsewhere -def latlon_coords(mesh): - """ - Gets expressions for the latitude and longitude fields. - - Args: - mesh (:class:`Mesh`): the model's mesh. - - Returns: - tuple of :class:`ufl.Expr`: expressions for the latitude and longitude - fields, in radians. - """ - x0, y0, z0 = SpatialCoordinate(mesh) - unsafe = z0/sqrt(x0*x0 + y0*y0 + z0*z0) - safe = min_value(max_value(unsafe, -1.0), 1.0) # avoid silly roundoff errors - theta = asin(safe) # latitude - lamda = atan2(y0, x0) # longitude - return theta, lamda - - -def sphere_to_cartesian(mesh, u_zonal, u_merid): - """ - Convert the horizontal spherical-polar components of a vector into - geocentric Cartesian components. - - Args: - mesh (:class:`Mesh`): _description_ - u_zonal (:class:`ufl.Expr`): the zonal component of the vector. - u_merid (:class:`ufl.Expr`): the meridional component of the vector. - - Returns: - _type_: _description_ - """ - theta, lamda = latlon_coords(mesh) - - cartesian_u_expr = -u_zonal*sin(lamda) - u_merid*sin(theta)*cos(lamda) - cartesian_v_expr = u_zonal*cos(lamda) - u_merid*sin(theta)*sin(lamda) - cartesian_w_expr = u_merid*cos(theta) - - return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr)) - - def incompressible_hydrostatic_balance(equation, b0, p0, top=False, params=None): """ Gives a pressure field in hydrostatic-balance for the Incompressible eqns. diff --git a/gusto/io.py b/gusto/io.py index cd08bdad2..e1b41e23b 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -7,14 +7,19 @@ 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, File, Constant, DumbCheckpoint, FILE_CREATE, FILE_READ, CheckpointFile) +from pyop2.mpi import MPI import numpy as np from gusto.logging import logger, update_logfile_location __all__ = ["pick_up_mesh", "IO"] +class GustoIOError(IOError): + pass + + def pick_up_mesh(output, mesh_name): """ Picks up a checkpointed mesh. This must be the first step of any model being @@ -72,6 +77,8 @@ def __init__(self, filename, field_points, description, self.field_points = field_points self.tolerance = tolerance self.comm = comm + if self.comm.size > 1: + raise GustoIOError("PointDataOutput does not work in parallel") if not create: return if self.comm.rank == 0: @@ -226,6 +233,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 +312,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 @@ -352,9 +369,12 @@ def setup_dump(self, state_fields, t, pick_up=False): state from a checkpointing file. Defaults to False. Raises: - IOError: if the results directory already exists, and the model is + GustoIOError: if the results directory already exists, and the model is not picking up or running in test mode. """ + # Use 0 for okay, 1 for internal exception 2 for external exception + raise_parallel_exception = 0 + error = None if any([self.output.dump_vtus, self.output.dump_nc, self.output.dumplist_latlon, self.output.dump_diagnostics, @@ -363,14 +383,30 @@ def setup_dump(self, state_fields, t, pick_up=False): self.dumpdir = path.join("results", self.output.dirname) running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname + # Raising exceptions needs to be done in parallel if self.mesh.comm.Get_rank() == 0: # Create results directory if it doesn't already exist if not path.exists(self.dumpdir): - makedirs(self.dumpdir) + try: + makedirs(self.dumpdir) + except OSError as e: + error = e + raise_parallel_exception = 2 elif not (running_tests or pick_up): # Throw an error if directory already exists, unless we # are picking up or running tests - raise IOError(f'results directory {self.dumpdir} already exists') + raise_parallel_exception = 1 + + # Gather errors from each rank and raise appropriate error everywhere + # This allreduce also ensures that all ranks are in sync wrt the results dir + raise_exception = self.mesh.comm.allreduce(raise_parallel_exception, op=MPI.MAX) + if raise_exception == 1: + raise GustoIOError(f'results directory {self.dumpdir} already exists') + elif raise_exception == 2: + if error: + raise error + else: + raise OSError('Check error message on rank 0') update_logfile_location(self.dumpdir, self.mesh.comm) @@ -380,6 +416,9 @@ def setup_dump(self, state_fields, t, pick_up=False): # make dump counter self.dumpcount = itertools.count() + # if picking-up, don't do initial dump + if pick_up: + next(self.dumpcount) if self.output.dump_vtus: # setup pvd output file @@ -400,6 +439,11 @@ def setup_dump(self, state_fields, t, pick_up=False): nc_field_file = Dataset(self.nc_filename, 'r') self.field_t_idx = len(nc_field_file['time'][:]) nc_field_file.close() + else: + self.field_t_idx = None + # Send information to other processors + self.field_t_idx = self.mesh.comm.bcast(self.field_t_idx, root=0) + else: # File needs creating self.create_nc_dump(self.nc_filename, space_names) @@ -434,6 +478,11 @@ def setup_dump(self, state_fields, t, pick_up=False): self.mesh.comm, create=to_create) + # if picking-up, don't do initial dump + self.diagcount = itertools.count() + if pick_up: + next(self.diagcount) + if len(self.output.point_data) > 0: # set up point data output pointdata_filename = self.dumpdir+"/point_data.nc" @@ -448,6 +497,9 @@ def setup_dump(self, state_fields, t, pick_up=False): # make point data dump counter self.pddumpcount = itertools.count() + # if picking-up, don't do initial dump + if pick_up: + next(self.pddumpcount) # set frequency of point data output - defaults to # dumpfreq if not set by user @@ -472,10 +524,13 @@ def setup_dump(self, state_fields, t, pick_up=False): # make a checkpoint counter self.chkptcount = itertools.count() + # if picking-up, don't do initial dump + if pick_up: + next(self.chkptcount) # 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): """ @@ -546,8 +601,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: @@ -577,6 +633,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 @@ -584,9 +641,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. @@ -597,6 +654,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. """ @@ -607,7 +665,7 @@ def dump(self, state_fields, t, initial_steps=None): for field in self.diagnostic_fields: field.compute() - if output.dump_diagnostics: + if output.dump_diagnostics and (next(self.diagcount) % output.diagfreq) == 0: # Output diagnostic data self.diagnostic_output.dump(state_fields, t) @@ -621,6 +679,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: @@ -629,6 +688,7 @@ 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) diff --git a/gusto/labels.py b/gusto/labels.py index 74e458729..7538de5a2 100644 --- a/gusto/labels.py +++ b/gusto/labels.py @@ -6,6 +6,85 @@ 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 +92,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/linear_solvers.py b/gusto/linear_solvers.py index 68ddf1477..f991706a6 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('Compressible linear solver: rho average solve') self.rho_avg_solver.solve() with timed_region("Gusto:HybridProjectExnerbar"): + logger.info('Compressible linear solver: Exner average solve') self.exner_avg_solver.solve() # Solve the hybridized system + logger.info('Compressible linear solver: hybridized 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('Compressible linear solver: restore 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('Compressible linear solver: theta solve') self.theta_solver.solve() # Copy into theta cpt of dy @@ -529,6 +534,7 @@ def solve(self, xrhs, dy): self.xrhs.assign(xrhs) with timed_region("Gusto:VelocityPressureSolve"): + logger.info('Incompressible linear solver: mixed solve') self.up_solver.solve() u1, p1 = self.up.subfunctions @@ -537,6 +543,7 @@ def solve(self, xrhs, dy): p.assign(p1) with timed_region("Gusto:BuoyancyRecon"): + logger.info('Incompressible linear solver: buoyancy reconstruction') self.b_solver.solve() b.assign(self.b) diff --git a/gusto/physics.py b/gusto/physics.py index 80ce276f5..1d62be0d8 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -12,7 +12,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, min_value, max_value, Constant, pi, Projector) @@ -27,18 +27,20 @@ "AdvectedMoments", "InstantRain", "SWSaturationAdjustment"] -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. + label_name (str): name of physics scheme, to be passed to its label. parameters (:class:`Configuration`, optional): parameters containing the values of gas constants. Defaults to None, in which case the 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 @@ -53,7 +55,7 @@ def evaluate(self): pass -class SaturationAdjustment(Physics): +class SaturationAdjustment(PhysicsParametrisation): """ Represents the phase change between water vapour and cloud liquid. @@ -89,7 +91,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 @@ -202,8 +205,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): """ @@ -213,6 +216,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) @@ -236,7 +240,7 @@ class AdvectedMoments(Enum): M3 = 1 # advect the mass of the distribution -class Fallout(Physics): +class Fallout(PhysicsParametrisation): """ Represents the fallout process of hydrometeors. @@ -271,6 +275,10 @@ 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" @@ -308,7 +316,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 @@ -365,12 +373,13 @@ 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() -class Coalescence(Physics): +class Coalescence(PhysicsParametrisation): """ Represents the coalescence of cloud droplets to form rain droplets. @@ -398,6 +407,14 @@ 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 accretion: + label_name += "_accretion" + if 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" @@ -448,10 +465,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): """ @@ -461,6 +478,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.rain.assign(x_in.subfunctions[self.rain_idx]) @@ -469,7 +487,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. @@ -502,7 +520,8 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', CompressibleEulerEquations. """ - super().__init__(equation, parameters=parameters) + label_name = 'evaporation_of_rain' + 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 @@ -617,8 +636,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): """ @@ -628,6 +647,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) @@ -638,7 +658,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. @@ -682,7 +702,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=parameters) self.convective_feedback = convective_feedback self.time_varying_saturation = time_varying_saturation @@ -736,25 +757,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( @@ -774,6 +794,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.subfunctions[self.VD_idx]) if self.time_varying_saturation: @@ -784,7 +805,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. @@ -848,7 +869,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=parameters) self.time_varying_saturation = time_varying_saturation self.convective_feedback = convective_feedback @@ -950,8 +972,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): """ @@ -965,7 +987,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]) if self.thermal_feedback: diff --git a/gusto/preconditioners.py b/gusto/preconditioners.py index 6148da642..92a9d7b8c 100644 --- a/gusto/preconditioners.py +++ b/gusto/preconditioners.py @@ -45,7 +45,7 @@ def initialize(self, pc): Args: pc (:class:`PETSc.PC`): preconditioner object to initialize. """ - from firedrake import (FunctionSpace, Function, Constant, + from firedrake import (FunctionSpace, Function, Constant, Cofunction, FiniteElement, TensorProductElement, TrialFunction, TrialFunctions, TestFunction, DirichletBC, interval, MixedElement, BrokenElement) @@ -104,7 +104,7 @@ def initialize(self, pc): V_d = FunctionSpace(mesh, broken_elements) # Set up relevant functions - self.broken_solution = Function(V_d) + self.broken_solution = Function(V_d.dual()) self.broken_residual = Function(V_d) self.trace_solution = Function(Vv_tr) self.unbroken_solution = Function(V) @@ -175,7 +175,7 @@ def initialize(self, pc): K = Tensor(Kform) # Assemble the Schur complement operator and right-hand side - self.schur_rhs = Function(Vv_tr) + self.schur_rhs = Cofunction(Vv_tr.dual()) self._assemble_Srhs = OneFormAssembler( K * Atilde.inv * AssembledVector(self.broken_residual), tensor=self.schur_rhs, diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 1e1ff41b1..c6d67e0f6 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, split) + NonlinearVariationalSolver, DirichletBC, split, Constant) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property @@ -15,9 +15,10 @@ 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 * +import math import numpy as np @@ -71,9 +72,13 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.field_name = field_name self.equation = None - self.dt = domain.dt + self.dt = Constant(0.0) + self.dt.assign(domain.dt) + self.original_dt = Constant(0.0) + self.original_dt.assign(self.dt) self.options = options self.limiter = limiter + self.courant_max = None if options is not None: self.wrapper_name = options.name @@ -135,9 +140,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): - self.evaluate_source.append(t.get(physics)) + if t.has_label(physics_label): + physics_name = t.get(physics_label) + 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 @@ -375,16 +384,24 @@ 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, - solver_parameters=None, limiter=None, options=None): + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, solver_parameters=None, limiter=None, + options=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the 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. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform 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 + `fixed_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 @@ -398,7 +415,11 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=solver_parameters, limiter=limiter, options=options) - self.subcycles = subcycles + if fixed_subcycles is not None and subcycle_by_courant is not None: + raise ValueError('Cannot specify both subcycle and subcycle_by ' + + 'arguments to a time discretisation') + self.fixed_subcycles = fixed_subcycles + self.subcycle_by_courant = subcycle_by_courant def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -413,11 +434,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): """ super().setup(equation, apply_bcs, *active_labels) - # if user has specified a number of subcycles, then save this + # if user has specified a number of fixed subcycles, then save this # and rescale dt accordingly; else perform just one cycle using dt - if self.subcycles is not None: - self.dt = self.dt/self.subcycles - self.ncycles = self.subcycles + if self.fixed_subcycles is not None: + self.dt.assign(self.dt/self.fixed_subcycles) + self.ncycles = self.fixed_subcycles else: self.dt = self.dt self.ncycles = 1 @@ -441,8 +462,9 @@ def solver(self): problem = NonlinearVariationalProblem(self.lhs - self.rhs, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ # If snes_type not specified by user, set this to ksp only to avoid outer Newton iteration - return NonlinearVariationalSolver(problem, solver_parameters={'snes_type': 'ksponly'} - | self.solver_parameters, options_prefix=solver_name) + self.solver_parameters.setdefault('snes_type', 'ksponly') + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) @abstractmethod def apply_cycle(self, x_out, x_in): @@ -464,10 +486,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. """ + # If doing adaptive subcycles, update dt and ncycles here + if self.subcycle_by_courant is not None: + self.ncycles = math.ceil(float(self.courant_max)/self.subcycle_by_courant) + self.dt.assign(self.original_dt/self.ncycles) + 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) @@ -514,8 +539,9 @@ class ExplicitMultistage(ExplicitTimeDiscretisation): # [ b_0 b_1 . b_s] # --------------------------------------------------------------------------- - def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None, - solver_parameters=None, limiter=None, options=None): + def __init__(self, domain, butcher_matrix, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, solver_parameters=None, + limiter=None, options=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -524,8 +550,15 @@ def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None, a butcher tableau defining a given Runge Kutta time discretisation. 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. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform 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 + `fixed_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 @@ -535,7 +568,9 @@ def __init__(self, domain, butcher_matrix, field_name=None, subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - super().__init__(domain, field_name=field_name, subcycles=subcycles, + super().__init__(domain, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, solver_parameters=solver_parameters, limiter=limiter, options=options) self.butcher_matrix = butcher_matrix @@ -623,7 +658,8 @@ class ForwardEuler(ExplicitMultistage): k0 = F[y^n] \n y^(n+1) = y^n + dt*k0 \n """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, solver_parameters=None, limiter=None, options=None): """ Args: @@ -631,8 +667,15 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No 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. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform 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 + `fixed_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 @@ -643,8 +686,9 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No recovery method. Defaults to None. """ butcher_matrix = np.array([1.]).reshape(1, 1) - super().__init__(domain, butcher_matrix=butcher_matrix, - field_name=field_name, subcycles=subcycles, + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -659,7 +703,8 @@ class SSPRK3(ExplicitMultistage): k2 = F[y^n + (1/4)*dt*(k0+k1)] \n y^(n+1) = y^n + (1/6)*dt*(k0 + k1 + 4*k2) \n """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, solver_parameters=None, limiter=None, options=None): """ Args: @@ -667,8 +712,15 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No 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. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform 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 + `fixed_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 @@ -679,8 +731,10 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No recovery method. Defaults to None. """ butcher_matrix = np.array([[1., 0., 0.], [1./4., 1./4., 0.], [1./6., 1./6., 2./3.]]) - super().__init__(domain, butcher_matrix=butcher_matrix, field_name=field_name, - subcycles=subcycles, solver_parameters=solver_parameters, + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -699,7 +753,8 @@ class RK4(ExplicitMultistage): where superscripts indicate the time-level. \n """ - def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, + def __init__(self, domain, field_name=None, fixed_subcycles=None, + subcycle_by_courant=None, solver_parameters=None, limiter=None, options=None): """ Args: @@ -707,8 +762,15 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No 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. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform 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 + `fixed_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 @@ -719,8 +781,10 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No recovery method. Defaults to None. """ 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.]]) - super().__init__(domain, butcher_matrix=butcher_matrix, field_name=field_name, - subcycles=subcycles, solver_parameters=solver_parameters, + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -737,7 +801,8 @@ 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, fixed_subcycles=None, + subcycle_by_courant=None, solver_parameters=None, limiter=None, options=None): """ Args: @@ -745,8 +810,15 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No 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. + fixed_subcycles (int, optional): the fixed number of sub-steps to + perform. This option cannot be specified with the + `subcycle_by_courant` argument. Defaults to None. + subcycle_by_courant (float, optional): specifying this option will + make the scheme perform 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 + `fixed_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 @@ -757,8 +829,9 @@ def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=No recovery method. Defaults to None. """ butcher_matrix = np.array([[1., 0.], [0.5, 0.5]]) - super().__init__(domain, butcher_matrix=butcher_matrix, - field_name=field_name, subcycles=subcycles, + super().__init__(domain, butcher_matrix, field_name=field_name, + fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -778,7 +851,7 @@ def __init__(self, domain, field_name=None, solver_parameters=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. + fixed_subcycles (int, optional): the number of sub-steps to perform. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 1b5024f28..48073a030 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 @@ -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() @@ -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('='*40) + 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 @@ -162,9 +170,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 +184,7 @@ 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)}') + self.log_timestep() self.x.update() @@ -186,9 +196,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() @@ -249,9 +260,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: @@ -285,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): """ @@ -318,10 +331,10 @@ def __init__(self, equation, scheme, io, spatial_methods=None, in which case the terms follow the original discretisation in the equation. 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. + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`), + pairing physics parametrisations and timestepping schemes to use + for each. Timestepping schemes for physics must be explicit. + Defaults to None. """ # As we handle physics differently to the Timestepper, these are not @@ -338,7 +351,7 @@ def __init__(self, equation, scheme, io, spatial_methods=None, 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, physics_label) @property def transporting_velocity(self): @@ -348,10 +361,11 @@ 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) + self.scheme.courant_max = self.io.courant_max def timestep(self): @@ -374,7 +388,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: @@ -392,31 +408,59 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, linear_solver (:class:`TimesteppingSolver`, optional): the object to use for the linear solve. Defaults to None. diffusion_schemes (iter, optional): iterable of pairs of the form - ``(field_name, scheme)`` indicating the fields to diffuse, and the + ``(field_name, scheme)`` indicating the fields to diffuse, and the :class:`~.TimeDiscretisation` to use. Defaults to None. 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 + (:class:`PhysicsParametrisation`, :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. + slow_physics_schemes: (list, optional): a list of tuples of the form + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`). + These schemes are all evaluated at the start of the time step. + Defaults to None. + fast_physics_schemes: (list, optional): a list of tuples of the form + (:class:`PhysicsParametrisation`, :class:`TimeDiscretisation`). + These schemes are evaluated within the outer loop. Defaults to None. + 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. + """ - :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" @@ -478,6 +522,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) @@ -506,7 +551,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 @@ -524,13 +569,14 @@ 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: scheme.setup(self.equation, apply_bcs, diffusion) - for _, scheme in self.physics_schemes: + for _, scheme in self.all_physics_schemes: apply_bcs = True - scheme.setup(self.equation, apply_bcs, physics) + scheme.setup(self.equation, apply_bcs, physics_label) def copy_active_tracers(self, x_in, x_out): """ @@ -552,33 +598,57 @@ 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') + # 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('SIQN: Fast physics') + 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)}') 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) @@ -597,9 +667,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): """ @@ -639,9 +710,10 @@ def __init__(self, equation, scheme, io, transport_method, Timestepping schemes for physics must be explicit. Defaults to None. 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. """ if isinstance(transport_method, TransportMethod): @@ -669,6 +741,20 @@ 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/data/dry_compressible_chkpt.h5 b/integration-tests/data/dry_compressible_chkpt.h5 index 9f9245928..da6537054 100644 Binary files a/integration-tests/data/dry_compressible_chkpt.h5 and b/integration-tests/data/dry_compressible_chkpt.h5 differ diff --git a/integration-tests/data/incompressible_chkpt.h5 b/integration-tests/data/incompressible_chkpt.h5 index 637a3bf84..5406c1282 100644 Binary files a/integration-tests/data/incompressible_chkpt.h5 and b/integration-tests/data/incompressible_chkpt.h5 differ diff --git a/integration-tests/data/linear_sw_wave_chkpt.h5 b/integration-tests/data/linear_sw_wave_chkpt.h5 index 6e7331825..17bd256c8 100644 Binary files a/integration-tests/data/linear_sw_wave_chkpt.h5 and b/integration-tests/data/linear_sw_wave_chkpt.h5 differ diff --git a/integration-tests/data/moist_compressible_chkpt.h5 b/integration-tests/data/moist_compressible_chkpt.h5 index 1e7ad4bb6..69e5649dc 100644 Binary files a/integration-tests/data/moist_compressible_chkpt.h5 and b/integration-tests/data/moist_compressible_chkpt.h5 differ diff --git a/integration-tests/data/sw_fplane_chkpt.h5 b/integration-tests/data/sw_fplane_chkpt.h5 index ff40a7d4e..5ace1eca0 100644 Binary files a/integration-tests/data/sw_fplane_chkpt.h5 and b/integration-tests/data/sw_fplane_chkpt.h5 differ diff --git a/integration-tests/diffusion/test_diffusion.py b/integration-tests/diffusion/test_diffusion.py index e946052cc..fd570107d 100644 --- a/integration-tests/diffusion/test_diffusion.py +++ b/integration-tests/diffusion/test_diffusion.py @@ -26,9 +26,9 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): f_end_expr = (1/(1+4*tmax))*f_init**(1/(1+4*tmax)) if DG: - V = domain.spaces("DG", "DG", degree=1) + V = domain.spaces("DG") else: - V = domain.spaces("theta", degree=1) + V = domain.spaces("theta") mu = 5. @@ -60,7 +60,7 @@ def test_vector_diffusion(tmpdir, DG, tracer_setup): if DG: V = VectorFunctionSpace(domain.mesh, "DG", 1) else: - V = domain.spaces("HDiv", "CG", 1) + V = domain.spaces("HDiv") f_init = as_vector([f_init, 0.]) f_end_expr = as_vector([f_end_expr, 0.]) diff --git a/integration-tests/equations/test_dry_compressible.py b/integration-tests/equations/test_dry_compressible.py index 2e293e474..d9eb21713 100644 --- a/integration-tests/equations/test_dry_compressible.py +++ b/integration-tests/equations/test_dry_compressible.py @@ -35,7 +35,7 @@ def run_dry_compressible(tmpdir): # I/O output = OutputParameters(dirname=tmpdir+"/dry_compressible", - dumpfreq=2, chkptfreq=2) + dumpfreq=2, chkptfreq=2, checkpoint=True) io = IO(domain, output) # Transport schemes @@ -93,7 +93,8 @@ def run_dry_compressible(tmpdir): checkpoint_name = 'dry_compressible_chkpt.h5' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_output = OutputParameters(dirname=tmpdir+"/dry_compressible", - checkpoint_pickup_filename=new_path) + checkpoint_pickup_filename=new_path, + checkpoint=True) check_mesh = pick_up_mesh(check_output, mesh_name) check_domain = Domain(check_mesh, dt, "CG", 1) check_eqn = CompressibleEulerEquations(check_domain, parameters) diff --git a/integration-tests/equations/test_incompressible.py b/integration-tests/equations/test_incompressible.py index a3ac3d8d0..b70b88598 100644 --- a/integration-tests/equations/test_incompressible.py +++ b/integration-tests/equations/test_incompressible.py @@ -34,7 +34,7 @@ def run_incompressible(tmpdir): # I/O output = OutputParameters(dirname=tmpdir+"/incompressible", - dumpfreq=2, chkptfreq=2) + dumpfreq=2, chkptfreq=2, checkpoint=True) io = IO(domain, output) # Transport Schemes @@ -87,7 +87,8 @@ def run_incompressible(tmpdir): checkpoint_name = 'incompressible_chkpt.h5' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_output = OutputParameters(dirname=tmpdir+"/incompressible", - checkpoint_pickup_filename=new_path) + checkpoint_pickup_filename=new_path, + checkpoint=True) check_mesh = pick_up_mesh(check_output, mesh_name) check_domain = Domain(check_mesh, dt, "CG", 1) check_eqn = IncompressibleBoussinesqEquations(check_domain, parameters) diff --git a/integration-tests/equations/test_linear_sw_wave.py b/integration-tests/equations/test_linear_sw_wave.py index 47b03e347..81b6745e8 100644 --- a/integration-tests/equations/test_linear_sw_wave.py +++ b/integration-tests/equations/test_linear_sw_wave.py @@ -36,6 +36,7 @@ def run_linear_sw_wave(tmpdir): output = OutputParameters( dirname=str(tmpdir)+"/linear_sw_wave", dumpfreq=1, + checkpoint=True ) io = IO(domain, output) transport_methods = [DefaultTransport(eqns, "D")] @@ -70,7 +71,8 @@ def run_linear_sw_wave(tmpdir): checkpoint_name = 'linear_sw_wave_chkpt.h5' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_output = OutputParameters(dirname=tmpdir+"/linear_sw_wave", - checkpoint_pickup_filename=new_path) + checkpoint_pickup_filename=new_path, + checkpoint=True) check_mesh = pick_up_mesh(check_output, mesh_name) check_domain = Domain(check_mesh, dt, 'BDM', 1) check_eqn = ShallowWaterEquations(check_domain, parameters, fexpr=fexpr) diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index e65a27157..04e931777 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -36,7 +36,7 @@ def run_moist_compressible(tmpdir): # I/O output = OutputParameters(dirname=tmpdir+"/moist_compressible", - dumpfreq=2, chkptfreq=2) + dumpfreq=2, checkpoint=True, chkptfreq=2) io = IO(domain, output) # Transport schemes @@ -99,7 +99,8 @@ def run_moist_compressible(tmpdir): checkpoint_name = 'moist_compressible_chkpt.h5' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_output = OutputParameters(dirname=tmpdir+"/moist_compressible", - checkpoint_pickup_filename=new_path) + checkpoint_pickup_filename=new_path, + checkpoint=True) check_mesh = pick_up_mesh(check_output, mesh_name) check_domain = Domain(check_mesh, dt, "CG", 1) check_eqn = CompressibleEulerEquations(check_domain, parameters, active_tracers=tracers) diff --git a/integration-tests/equations/test_sw_fplane.py b/integration-tests/equations/test_sw_fplane.py index ccb78ac75..7fd472ca3 100644 --- a/integration-tests/equations/test_sw_fplane.py +++ b/integration-tests/equations/test_sw_fplane.py @@ -32,6 +32,7 @@ def run_sw_fplane(tmpdir): output = OutputParameters( dirname=str(tmpdir)+"/sw_fplane", dumpfreq=1, + checkpoint=True ) io = IO(domain, output, diagnostic_fields=[CourantNumber()]) @@ -105,7 +106,8 @@ def run_sw_fplane(tmpdir): checkpoint_name = 'sw_fplane_chkpt.h5' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_output = OutputParameters(dirname=tmpdir+"/sw_fplane", - checkpoint_pickup_filename=new_path) + checkpoint_pickup_filename=new_path, + checkpoint=True) check_mesh = pick_up_mesh(check_output, mesh_name) check_domain = Domain(check_mesh, dt, 'RTCF', 1) check_eqn = ShallowWaterEquations(check_domain, parameters, fexpr=fexpr) diff --git a/integration-tests/equations/test_thermal_sw.py b/integration-tests/equations/test_thermal_sw.py index c48f97939..8328c308d 100644 --- a/integration-tests/equations/test_thermal_sw.py +++ b/integration-tests/equations/test_thermal_sw.py @@ -8,8 +8,7 @@ from os import path from gusto import * -from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, - pi, sin, cos) +from firedrake import IcosahedralSphereMesh, SpatialCoordinate, pi, sin, cos from netCDF4 import Dataset R = 6371220. @@ -65,9 +64,10 @@ def set_up_initial_conditions(domain, equation, stepper): g = equation.parameters.g Omega = equation.parameters.Omega - phi, lamda = latlon_coords(domain.mesh) + x = SpatialCoordinate(domain.mesh) + _, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2]) - uexpr = sphere_to_cartesian(domain.mesh, u_max*cos(phi), 0) + uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x) w = Omega*R*u_max + (u_max**2)/2 sigma = 0 Dexpr = H - (1/g)*(w + sigma)*((sin(phi))**2) diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index f8a39b69c..863369634 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 @@ -116,12 +116,14 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): output_1 = OutputParameters( dirname=dirname_1, dumpfreq=1, + checkpoint=True, checkpoint_method=checkpoint_method, chkptfreq=4, ) output_2 = OutputParameters( dirname=dirname_2, dumpfreq=1, + checkpoint=True, checkpoint_method=checkpoint_method, chkptfreq=2, ) @@ -154,6 +156,7 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): dirname=dirname_3, dumpfreq=1, chkptfreq=2, + checkpoint=True, checkpoint_method=checkpoint_method, checkpoint_pickup_filename=chkpt_2_path, ) @@ -175,16 +178,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 # ------------------------------------------------------------------------ # @@ -193,6 +186,7 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): dirname=dirname_3, dumpfreq=1, chkptfreq=2, + checkpoint=True, checkpoint_method=checkpoint_method, checkpoint_pickup_filename=chkpt_2_path ) @@ -201,6 +195,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 # ------------------------------------------------------------------------ # diff --git a/integration-tests/model/test_nc_outputting.py b/integration-tests/model/test_nc_outputting.py index ef704cacf..9daf01479 100644 --- a/integration-tests/model/test_nc_outputting.py +++ b/integration-tests/model/test_nc_outputting.py @@ -8,8 +8,8 @@ VectorFunctionSpace, ExtrudedMesh, SpatialCoordinate, as_vector, exp) from gusto import (Domain, IO, PrescribedTransport, AdvectionEquation, - ForwardEuler, OutputParameters, VelocityX, VelocityY, - VelocityZ, MeridionalComponent, ZonalComponent, + ForwardEuler, OutputParameters, XComponent, YComponent, + ZComponent, MeridionalComponent, ZonalComponent, RadialComponent, DGUpwind) from netCDF4 import Dataset import pytest @@ -83,13 +83,13 @@ def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details): # Make velocity components for this geometry if geometry == "interval": - diagnostic_fields = [VelocityX()] + diagnostic_fields = [XComponent('u')] elif geometry == "vertical_slice": - diagnostic_fields = [VelocityX(), VelocityZ()] + diagnostic_fields = [XComponent('u'), ZComponent('u')] elif geometry == "plane": - diagnostic_fields = [VelocityX(), VelocityY()] + diagnostic_fields = [XComponent('u'), YComponent('u')] elif geometry == "extruded_plane": - diagnostic_fields = [VelocityX(), VelocityY(), VelocityZ()] + diagnostic_fields = [XComponent('u'), YComponent('u'), ZComponent('u')] elif geometry == "spherical_shell": diagnostic_fields = [ZonalComponent('u'), MeridionalComponent('u')] elif geometry == "extruded_spherical_shell": diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index ebe9993b1..6f114fe90 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -43,7 +43,8 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): transport_scheme = TR_BDF2(domain, gamma=0.5) elif scheme == "Leapfrog": # Leapfrog unstable with DG - Vf = domain.spaces("CG", "CG", 1) + domain.spaces.create_space("CG1", "CG", 1) + Vf = domain.spaces("CG1") eqn = AdvectionEquation(domain, Vf, "f") transport_scheme = Leapfrog(domain) elif scheme == "AdamsBashforth": diff --git a/integration-tests/physics/test_saturation_adjustment.py b/integration-tests/physics/test_saturation_adjustment.py index 0b92d6f11..922e7e40e 100644 --- a/integration-tests/physics/test_saturation_adjustment.py +++ b/integration-tests/physics/test_saturation_adjustment.py @@ -60,8 +60,8 @@ def run_cond_evap(dirname, process): # Initial conditions # ------------------------------------------------------------------------ # - Vt = domain.spaces("theta", degree=1) - Vr = domain.spaces("DG", "DG", degree=1) + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") # Declare prognostic fields rho0 = stepper.fields("rho") 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/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index f68d45a2d..8255ad2cd 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -53,7 +53,7 @@ def setup_limiters(dirname, space): else: raise NotImplementedError - Vpsi = domain.spaces('CG', 'CG', degree+1) + Vpsi = domain.spaces('H1') # Equation eqn = AdvectionEquation(domain, V, 'tracer') diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 5894050b5..2247c59b7 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, fixed_subcycles=2) + elif subcycling == "adaptive": + transport_scheme = SSPRK3(domain, subcycle_by_courant=0.25) transport_method = DGUpwind(eqn, "f") timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) diff --git a/unit-tests/test_distances.py b/unit-tests/test_distances.py new file mode 100644 index 000000000..e7895314a --- /dev/null +++ b/unit-tests/test_distances.py @@ -0,0 +1,91 @@ +""" +Test the formulae for finding shortest distances between two points. +""" +import importlib +import numpy as np +import firedrake as fd +from gusto.coord_transforms import * +import pytest + +tol = 1e-12 + + +# Structure of values for testing Firedrake and numpy routines are different +def setup_values(values_list, module_name, mesh=None): + + if module_name == "firedrake": + # Put data into fields + DG0 = fd.FunctionSpace(mesh, 'DG', 0) + new_values = fd.Function(DG0) + new_values.dat.data[:] = values_list[:] + else: + new_values = np.array(values_list) + return new_values + + +# Checks values and prints an error message if they aren't correct +def check_values(new_values, answers, module_name, routine): + + if module_name == 'firedrake': + # Unpack Firedrake field to array of data + V = answers.function_space() + field = fd.Function(V) + field.interpolate(new_values) + new_values = field.dat.data[:] + answers = answers.dat.data[:] + + assert np.allclose(new_values, answers, atol=tol), \ + f'Incorrect answer for {module_name} module and {routine} ' + + +@pytest.mark.parametrize("module_name", ["numpy", "firedrake"]) +def test_great_arc_angle(module_name): + + module = importlib.import_module(module_name) + pi = module.pi + sqrt = module.sqrt + acos = module.acos if hasattr(module, "acos") else module.arccos + + lon2 = pi / 4.0 + lat2 = pi / 4.0 + raw_lon1 = [0.0, 0.0, pi/4, pi/4, pi/4] + raw_lat1 = [0.0, pi/4, 0.0, pi/2, -pi/2] + raw_answers = [pi/3, acos(0.5*(1+1/sqrt(2))), pi/4, pi/4, 3*pi/4] + + # Test for firedrake routines requires a mesh + if module_name == "firedrake": + mesh = fd.UnitIntervalMesh(len(raw_answers)) + else: + mesh = None + + # Put the coordinates in numpy arrays or firedrake functions + lon1 = setup_values(raw_lon1, module_name, mesh=mesh) + lat1 = setup_values(raw_lat1, module_name, mesh=mesh) + answers = setup_values(raw_answers, module_name, mesh=mesh) + + distances = great_arc_angle(lon1, lat1, lon2, lat2) + + check_values(distances, answers, module_name, 'great_arc') + + +@pytest.mark.parametrize("module_name", ["numpy", "firedrake"]) +def test_periodic_distance(module_name): + + max_x = 10.0 + x2 = 6.0 + raw_x1 = [0.0, 1.0, 2.0, 4.0, 5.0, 6.0, 8.0, 10.0] + raw_answers = [4.0, -5.0, -4.0, -2.0, -1.0, 0.0, 2.0, 4.0] + + # Test for firedrake routines requires a mesh + if module_name == "firedrake": + mesh = fd.UnitIntervalMesh(len(raw_answers)) + else: + mesh = None + + # Put the coordinates in numpy arrays or firedrake functions + x1 = setup_values(raw_x1, module_name, mesh=mesh) + answers = setup_values(raw_answers, module_name, mesh=mesh) + + distances = periodic_distance(x1, x2, max_x) + + check_values(distances, answers, module_name, "periodic_distance") diff --git a/unit-tests/test_function_spaces.py b/unit-tests/test_function_spaces.py index e72789f28..b39aa8cc5 100644 --- a/unit-tests/test_function_spaces.py +++ b/unit-tests/test_function_spaces.py @@ -12,9 +12,11 @@ domain_family_dict = {'interval': ['CG'], 'vertical_slice': ['CG'], 'plane': ['BDM', 'BDMF', 'BDME', 'BDFM', - 'RT', 'RTF', 'RTE', 'RTCF', 'RTCE'], + 'RT', 'RTF', 'RTE', 'RTCF', 'RTCE', + 'BDMCE', 'BDMCF'], 'extruded_plane': ['BDM', 'BDMF', 'BDME', 'BDFM', - 'RT', 'RTF', 'RTE', 'RTCF', 'RTCE']} + 'RT', 'RTF', 'RTE', 'RTCF', 'RTCE', + 'BDMCE', 'BDMCF']} reduced_domain_family_dict = {'interval': ['CG'], 'vertical_slice': ['CG'], @@ -36,7 +38,7 @@ def set_up_mesh(domain, family): if family in ['BDM', 'BDMF', 'BDME', 'BDFM', 'RT', 'RTE', 'RTF']: quadrilateral = False - elif family in ['RTCF', 'RTCE']: + elif family in ['RTCF', 'RTCE', 'BDMCE', 'BDMCF']: quadrilateral = True if domain == 'interval': @@ -73,11 +75,15 @@ def test_de_rham_spaces(domain, family): spaces.build_compatible_spaces(family, degree) # Work out correct CG degree - if family in ['BDM', 'BDME', 'BDMF']: + if family in ['BDM', 'BDME', 'BDMF', 'BDMCE', 'BDMCF']: if domain in ['vertical_slice', 'extruded_plane']: - cg_degree = (degree[0]+2, degree[1]+2) + cg_degree = (degree[0]+2, degree[1]+1) else: cg_degree = degree + 2 + elif family == 'BDFM' and domain in ['vertical_slice', 'extruded_plane']: + cg_degree = (degree[0] + 2, degree[1]+1) + elif family == 'BDFM': + cg_degree = degree + 2 elif domain in ['vertical_slice', 'extruded_plane']: cg_degree = (degree[0]+1, degree[1]+1) else: @@ -103,6 +109,7 @@ def test_dg_equispaced(domain, family): mesh = set_up_mesh(domain, family) spaces = Spaces(mesh) + spaces.build_dg1_equispaced() DG1 = spaces('DG1_equispaced') elt = DG1.ufl_element() @@ -121,7 +128,7 @@ def test_dg0(domain, family): mesh = set_up_mesh(domain, family) spaces = Spaces(mesh) - DG0 = spaces('DG0', 'DG', degree=0) + DG0 = spaces.create_space('DG0', 'DG', degree=0) elt = DG0.ufl_element() assert elt.degree() in [0, (0, 0)], '"DG0" space does not seem to be' \ + f'degree 0. Found degree {elt.degree()}' @@ -135,14 +142,11 @@ def test_cg(domain, family): mesh = set_up_mesh(domain, family) spaces = Spaces(mesh) + degree = 3 - if domain in ['vertical_slice', 'extruded_plane']: - degree = (1, 2) - CG = spaces('CG', 'CG', horizontal_degree=degree[0], vertical_degree=degree[1]) - else: - degree = 3 - CG = spaces('CG', 'CG', degree=degree) + CG = spaces.create_space('CG', 'CG', degree=degree) elt = CG.ufl_element() - assert elt.degree() == degree, '"CG" space does not seem to be degree ' \ - + f'{degree}. Found degree {elt.degree()}' + assert elt.degree() == degree or elt.degree() == (degree, degree), \ + (f'"CG" space does not seem to be degree {degree}. ' + + f'Found degree {elt.degree()}') diff --git a/unit-tests/test_h1_spaces.py b/unit-tests/test_h1_spaces.py new file mode 100644 index 000000000..8b616eae5 --- /dev/null +++ b/unit-tests/test_h1_spaces.py @@ -0,0 +1,89 @@ +""" +Check that the H1 spaces are correct by: +- creating HDiv spaces from a stream function in H1, and checking that their + divergence in L2 is zero +- checking that this velocity gives a sensible answer +- projecting the velocity back to obtain the vorticity field, and checking + that this is correct +""" + +from firedrake import (PeriodicRectangleMesh, SpatialCoordinate, norm, + Function, sin, cos, pi, as_vector, grad, div, + TestFunction, TrialFunction, dx, inner, errornorm, + LinearVariationalProblem, LinearVariationalSolver) +from gusto import Domain +import pytest + +hdiv_families = ["RTF", "BDMF", "BDFM", "RTCF", "BDMCF"] +degrees = [1, 2] + + +def combos(families, degs): + # Form all combinations of families/degrees + # This is done because for BDFM family, only degree 1 is possible + all_combos = [] + for family in families: + for deg in degs: + if not (family == 'BDFM' and deg != 1): + all_combos.append((family, deg)) + return all_combos + + +@pytest.mark.parametrize("hdiv_family, degree", combos(hdiv_families, degrees)) +def test_h1_spaces(hdiv_family, degree): + # Set up mesh and domain + dt = 2.0 + Lx = 10 + Ly = 10 + nx = int(20 / degree) + ny = int(20 / degree) + quadrilateral = True if hdiv_family in ['RTCF', 'BDMCF'] else False + mesh = PeriodicRectangleMesh(nx, ny, Lx, Ly, quadrilateral=quadrilateral) + domain = Domain(mesh, dt, hdiv_family, degree) + x, y = SpatialCoordinate(mesh) + + # Declare spaces and functions + H1_space = domain.spaces('H1') + HDiv_space = domain.spaces('HDiv') + L2_space = domain.spaces('L2') + streamfunc = Function(H1_space) + velocity = Function(HDiv_space) + velocity_true = Function(HDiv_space) + vorticity = Function(H1_space) + vorticity_true = Function(H1_space) + divergence = Function(L2_space) + test = TestFunction(H1_space) + trial = TrialFunction(H1_space) + + # Expressions + streamfunc_expr = sin(2*pi*x/Lx)*cos(4*pi*y/Ly) + velocity_expr = as_vector([4*pi/Ly*sin(4*pi*y/Ly)*sin(2*pi*x/Lx), + 2*pi/Lx*cos(2*pi*x/Lx)*cos(4*pi*y/Ly)]) + vorticity_expr = - ((2*pi/Lx)**2 + (4*pi/Ly)**2)*streamfunc_expr + + # Evaluate velocity from stream function + gradperp = lambda q: domain.perp(grad(q)) + streamfunc.project(streamfunc_expr) + velocity.project(gradperp(streamfunc)) + velocity_true.project(velocity_expr) + divergence.project(div(velocity)) + velocity_norm = errornorm(velocity, velocity_true) / norm(velocity_true) + divergence_norm = norm(divergence) / (Lx*Ly) + + # Evaluate vorticity from velocity + velocity.project(velocity_expr) + vorticity_true.project(vorticity_expr) + eqn_lhs = trial * test * dx + eqn_rhs = -inner(gradperp(test), velocity) * dx + problem = LinearVariationalProblem(eqn_lhs, eqn_rhs, vorticity) + solver = LinearVariationalSolver(problem) + solver.solve() + vorticity_norm = errornorm(vorticity, vorticity_true) / norm(vorticity_true) + + # Check values + assert divergence_norm < 1e-8, \ + f'Divergence for family {hdiv_family} degree {degree} is too large' + assert velocity_norm < 0.015, \ + f'Error in velocity for family {hdiv_family} degree {degree} is too large' + assert vorticity_norm < 2e-6, \ + f'Error in vorticity for family {hdiv_family} degree {degree} is too large' diff --git a/unit-tests/test_rotated_coords.py b/unit-tests/test_rotated_coords.py new file mode 100644 index 000000000..aa01b7f47 --- /dev/null +++ b/unit-tests/test_rotated_coords.py @@ -0,0 +1,60 @@ +""" +Test the formulae for rotating spherical coordinates. +""" +import numpy as np +from gusto.coord_transforms import * + +tol = 1e-12 + + +def test_rotated_lonlatr_coords_numpy(): + + pi = np.pi + + rotated_pole = (pi/2, 0.0) + + llr_coords = np.array(([pi, pi/2, 0.0], + [0.0, 0.0, 0.0], + [0.5, 10.0, 0.0])) + + xyz_coords = np.array(([-0.5, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 10.0, 0.0])) + + new_llr = rotated_lonlatr_coords(xyz_coords, rotated_pole) + + assert np.allclose(new_llr, llr_coords), \ + 'Incorrect answer for numpy and rotated_lonlatr' + + +def test_rotated_lonlatr_coords_firedrake(): + + from firedrake import (CubedSphereMesh, pi, SpatialCoordinate, Function, + FunctionSpace) + + rotated_pole = (pi/2, 0.0) + + radius = 10.0 + mesh = CubedSphereMesh(radius) + + xyz = SpatialCoordinate(mesh) + rotated_llr = rotated_lonlatr_coords(xyz, rotated_pole) + + # Check against new (X,Y,Z) + new_xyz = xyz_from_lonlatr(rotated_llr[0], rotated_llr[1], rotated_llr[2]) + + V = FunctionSpace(mesh, "CG", 1) + + x = Function(V).interpolate(xyz[0]) + y = Function(V).interpolate(xyz[1]) + z = Function(V).interpolate(xyz[2]) + new_x = Function(V).interpolate(new_xyz[0]) + new_y = Function(V).interpolate(new_xyz[1]) + new_z = Function(V).interpolate(new_xyz[2]) + + assert np.allclose(x.dat.data, new_x.dat.data), \ + 'Incorrect answer for firedrake rotated x coordinates' + assert np.allclose(y.dat.data, -new_z.dat.data), \ + 'Incorrect answer for firedrake rotated y coordinates' + assert np.allclose(z.dat.data, new_y.dat.data), \ + 'Incorrect answer for firedrake rotated z coordinates' diff --git a/unit-tests/test_rotated_vectors.py b/unit-tests/test_rotated_vectors.py new file mode 100644 index 000000000..e326f0938 --- /dev/null +++ b/unit-tests/test_rotated_vectors.py @@ -0,0 +1,57 @@ +""" +Test the formulae for rotating spherical vectors. +""" +import numpy as np +from gusto.coord_transforms import * + +tol = 1e-12 + + +def test_rotated_lonlatr_vectors_firedrake(): + + from firedrake import (CubedSphereMesh, pi, SpatialCoordinate, Function, + VectorFunctionSpace, as_vector, grad, sqrt, dot, + atan) + + new_pole = (pi/4, pi/4) + + radius = 10.0 + mesh = CubedSphereMesh(radius) + + xyz = SpatialCoordinate(mesh) + + rot_axis, rot_angle = pole_rotation(new_pole) + new_xyz = rodrigues_rotation(xyz, as_vector(rot_axis), rot_angle) + + # TODO: this should be + # new_lonlatr = lonlatr_from_xyz(new_xyz[0], new_xyz[1], new_xyz[2]) + # but when atan_2 became atan2 in UFL we lost the derivative of atan2 + # therefore define this here with atan + lon = atan(new_xyz[1]/new_xyz[0]) + l = sqrt(new_xyz[0]**2 + new_xyz[1]**2) + lat = atan(new_xyz[2]/l) + + # Do an alternative calculation based on gradients of new coordinates + answer_e_lon = grad(lon) + answer_e_lat = grad(lat) + answer_e_r = grad(sqrt(dot(xyz, xyz))) + + # Normalise + answer_e_lon /= sqrt(dot(answer_e_lon, answer_e_lon)) + answer_e_lat /= sqrt(dot(answer_e_lat, answer_e_lat)) + answer_e_r /= sqrt(dot(answer_e_r, answer_e_r)) + + new_e_lon, new_e_lat, new_e_r = rotated_lonlatr_vectors(xyz, new_pole) + + # Check answers + V = VectorFunctionSpace(mesh, "CG", 1) + + for new_vector, answer_vector, component in zip([new_e_lon, new_e_lat, new_e_r], + [answer_e_lon, answer_e_lat, answer_e_r], + ['lon', 'lat', 'r']): + + new_field = Function(V).interpolate(new_vector) + answer_field = Function(V).interpolate(as_vector(answer_vector)) + + assert np.allclose(new_field.dat.data, answer_field.dat.data), \ + f'Incorrect answer for firedrake rotated {component} vector' diff --git a/unit-tests/test_spherical_coord_transforms.py b/unit-tests/test_spherical_coord_transforms.py new file mode 100644 index 000000000..486c950e2 --- /dev/null +++ b/unit-tests/test_spherical_coord_transforms.py @@ -0,0 +1,90 @@ +""" +Test the formulae for transforming between spherical and Cartesian coordinates. +""" +import importlib +import numpy as np +import firedrake as fd +from gusto.coord_transforms import * +import pytest + +tol = 1e-12 + + +# Structure of coordinates for testing Firedrake and numpy routines are different +def setup_coordinates(coords_list, module_name, mesh=None): + # Coords should be a list of lists + all_coords = np.array(coords_list) + _, num_coords = np.shape(all_coords) + + if module_name == "firedrake": + # Put data into fields + DG0 = fd.FunctionSpace(mesh, 'DG', 0) + coord_fields = [] + for i in range(num_coords): + coord_field = fd.Function(DG0) + coord_field.dat.data[:] = all_coords[:, i] + coord_fields.append(coord_field) + return coord_fields + else: + # Transform to list of arrays + new_coords_list = [] + for i in range(num_coords): + new_coords_list.append(all_coords[:, i]) + return new_coords_list + + +# Checks coordinate values and prints an error message if they aren't correct +def check_coords(new_values, answers, module_name, routine): + + if module_name == 'firedrake': + V = answers[0].function_space() + fields = [fd.Function(V) for _ in new_values] + for field, new_value in zip(fields, new_values): + field.interpolate(new_value) + new_values = [field.dat.data[:] for field in fields] + answers = [answer.dat.data[:] for answer in answers] + + for i, (new_value, answer) in enumerate(zip(new_values, answers)): + assert np.allclose(new_value, answer, atol=tol), \ + f'Incorrect answer for {module_name} module and {routine} ' \ + + f'routine, coord {i}' + + +@pytest.mark.parametrize("module_name", ["numpy", "firedrake"]) +def test_xyz_and_lonlatr(module_name): + + module = importlib.import_module(module_name) + pi = module.pi + + # Use the following sets of coordinates: + # (r, lon, lat) <--> (x, y, z) + # (2,0,pi/2) <--> (0, 0, 2) + # (0.5,pi,0) <--> (-0.5, 0, 0) + # (10,-pi/2,0) <--> (0,-10, 0) + # (0,0,0) <--> (0, 0, 0) + + raw_llr_coords = [[0.0, pi/2, 2.0], + [pi, 0.0, 0.5], + [-pi/2, 0.0, 10], + [0.0, 0.0, 0.0]] + + raw_xyz_coords = [[0.0, 0.0, 2.0], + [-0.5, 0.0, 0.0], + [0.0, -10.0, 0.0], + [0.0, 0.0, 0.0]] + + # Test for firedrake routines requires a mesh + if module_name == "firedrake": + mesh = fd.UnitIntervalMesh(len(raw_llr_coords)) + else: + mesh = None + + # Put the coordinates in numpy arrays or firedrake functions + llr_coords = setup_coordinates(raw_llr_coords, module_name, mesh=mesh) + xyz_coords = setup_coordinates(raw_xyz_coords, module_name, mesh=mesh) + + new_llr = lonlatr_from_xyz(xyz_coords[0], xyz_coords[1], xyz_coords[2]) + new_xyz = xyz_from_lonlatr(llr_coords[0], llr_coords[1], llr_coords[2]) + + check_coords(new_llr, llr_coords, module_name, 'lonlatr_from_xyz') + check_coords(new_xyz, xyz_coords, module_name, 'xyz_from_lonlatr') diff --git a/unit-tests/test_spherical_rotations.py b/unit-tests/test_spherical_rotations.py new file mode 100644 index 000000000..07063d457 --- /dev/null +++ b/unit-tests/test_spherical_rotations.py @@ -0,0 +1,112 @@ +""" +Test the formulae for rotating spherical vectors. +""" +import importlib +import numpy as np +import firedrake as fd +from gusto.coord_transforms import * +import pytest + +tol = 1e-12 + + +# Structure of values for testing Firedrake and numpy routines are different +def setup_values(values, config_name, len_array, mesh=None): + if config_name == "numpy_vector": + # Transform to array + vector = np.array(values) + return vector + + elif config_name == "numpy_field": + # Transform to list of arrays + vector_field = np.zeros((len_array, 3)) + for i in range(len_array): + vector_field[i, :] = values[:] + return vector_field + + else: + # Firedrake + # Put data into a single vector-valued field + Vec_DG0 = fd.VectorFunctionSpace(mesh, 'DG', 0, dim=3) + vector_field = fd.Function(Vec_DG0) + for i in range(len_array): + vector_field.dat.data[i, :] = values[:] + return vector_field + + +def test_pole_rotation(): + + pi = np.pi + + new_poles = [(0.0, pi/2), (0.0, 0.0), (pi/2, 0.0)] + answer_axes = [(0.0, 1.0, 0.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)] + answer_angles = [0.0, pi/2, pi/2] + + for i, (new_pole, answer_axis, answer_angle) in \ + enumerate(zip(new_poles, answer_axes, answer_angles)): + rot_axis, rot_angle = pole_rotation(new_pole) + + assert abs(rot_angle - answer_angle) < tol, f'pole_rotation {i}: ' \ + + f'rotation angle not correct. Got {rot_angle} but expected {answer_angle}' + + assert np.allclose(np.array(rot_axis), np.array(answer_axis), atol=tol), \ + f'pole_rotation {i}: rotation axis not correct' + + +# Checks values and prints an error message if they aren't correct +def check_values(new_value, answer, config_name, routine): + + if config_name == 'firedrake': + # Need to interpolate expression into vector field + new_field = fd.Function(answer.function_space()) + new_field.interpolate(new_value) + assert np.allclose(new_field.dat.data, answer.dat.data), \ + f'Incorrect answer for {config_name} and {routine}' + else: + assert np.allclose(new_value, answer), \ + f'Incorrect answer for {config_name} and {routine}' + + +@pytest.mark.parametrize("config_name", ["numpy_field", "numpy_vector", "firedrake"]) +def test_rodrigues_rotation(config_name): + + module_name = "firedrake" if config_name == "firedrake" else "numpy" + module = importlib.import_module(module_name) + pi = module.pi + sqrt = module.sqrt + len_array = 5 + + rot_axes = [(1.0, 0.0, 0.0), (0.0, 0.0, 1.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)] + rot_angles = [0.0, pi/4, pi/2, pi/2] + + raw_orig_vectors = [(19.4, -14.6, 12.1), + (2.0, 0.0, -7.5), + (-101.4, 0.6, 7.0), + (12.3, -4.5, 0.09)] + + raw_answer_vectors = [(19.4, -14.6, 12.1), + (sqrt(2.0), sqrt(2.0), -7.5), + (7.0, 0.6, 101.4), + (12.3, 0.09, 4.5)] + + # Test for firedrake routines requires a mesh + if module_name == "firedrake": + mesh = fd.UnitIntervalMesh(len_array) + else: + mesh = None + + for j, (raw_orig_vector, raw_answer_vector, rot_axis, rot_angle) in \ + enumerate(zip(raw_orig_vectors, raw_answer_vectors, rot_axes, rot_angles)): + + if module_name == 'firedrake': + rot_axis = fd.as_vector(rot_axis) + else: + rot_axis = np.array(rot_axis) + + # Put the coordinates in numpy arrays or firedrake functions + orig_vector = setup_values(raw_orig_vector, config_name, len_array, mesh=mesh) + answer_vector = setup_values(raw_answer_vector, config_name, len_array, mesh=mesh) + + new_vector = rodrigues_rotation(orig_vector, rot_axis, rot_angle) + + check_values(new_vector, answer_vector, config_name, f'rodrigues_rotation {j}') diff --git a/unit-tests/test_spherical_vector_transforms.py b/unit-tests/test_spherical_vector_transforms.py new file mode 100644 index 000000000..b309db1fa --- /dev/null +++ b/unit-tests/test_spherical_vector_transforms.py @@ -0,0 +1,130 @@ +""" +Test the formulae for transforming vectors between spherical and Cartesian +components. This is tested for numpy arrays, tuples of Firedrake scalars and +Firedrake vector function spaces. +""" +import importlib +import numpy as np +import firedrake as fd +from gusto.coord_transforms import * +import pytest + +tol = 1e-12 + + +# Structure of values for testing Firedrake and numpy routines are different +def setup_values(values_list, config_name, mesh=None): + # Values should be a list of lists + all_values = np.array(values_list) + _, num_coords = np.shape(all_values) + + if config_name == "firedrake_scalars": + # Put data into a series of scalar fields + DG0 = fd.FunctionSpace(mesh, 'DG', 0) + scalar_fields = [] + for i in range(num_coords): + scalar_field = fd.Function(DG0) + scalar_field.dat.data[:] = all_values[:, i] + scalar_fields.append(scalar_field) + return scalar_fields + elif config_name == "firedrake_vectors": + # Put data into a single vector-valued field + Vec_DG0 = fd.VectorFunctionSpace(mesh, 'DG', 0, dim=3) + vector_field = fd.Function(Vec_DG0) + vector_field.dat.data[:, :] = all_values[:, :] + return vector_field + + else: + # Transform to list of arrays + new_values_list = [] + for i in range(num_coords): + new_values_list.append(all_values[:, i]) + return new_values_list + + +# Checks values and prints an error message if they aren't correct +def check_values(new_values, answers, config_name, routine): + + if config_name == 'firedrake_scalars': + V = answers[0].function_space() + fields = [fd.Function(V) for _ in new_values] + for field, new_value in zip(fields, new_values): + field.interpolate(new_value) + new_values = [field.dat.data[:] for field in fields] + answers = [answer.dat.data[:] for answer in answers] + + elif config_name == 'firedrake_vectors': + # Interpolate answer into vector function space + vector_field = fd.Function(answers.function_space()) + vector_field.interpolate(new_values) + + answer_shape = np.shape(answers.dat.data) + new_values = [vector_field.dat.data[:, i] for i in range(answer_shape[1])] + answers = [answers.dat.data[:, i] for i in range(answer_shape[1])] + + for i, (new_value, answer) in enumerate(zip(new_values, answers)): + assert np.all(np.isclose(new_value, answer, atol=tol)), \ + f'Incorrect answer for {config_name} module and {routine} ' \ + + f'routine, coord {i}' + + +@pytest.mark.parametrize("config_name", ["numpy", "firedrake_scalars", "firedrake_vectors"]) +@pytest.mark.parametrize("position_units", ["xyz", "lonlatr_rad"]) +def test_xyz_and_lonlatr_vectors(config_name, position_units): + + module_name = "numpy" if config_name == "numpy" else "firedrake" + coord_config = "numpy" if config_name == "numpy" else "firedrake_scalars" + module = importlib.import_module(module_name) + pi = module.pi + + # Consider the following vectors: + # (r,lon,lat) components <--> (x,y,z) components at (x,y,z) or (r,lon,lat) + # (10,-6,0.5) <--> (10,-6,0.5) at (5,0,0) or (5,0,0) + # (0.7,3,1.2) <--> (3,-0.7,1.2) at (0,-0.5,0) or (0.5,-pi/2,0) + # (2,0,5) <--> (5,0,-2) at (0,0,-15) or (15,0,-pi/2) + + raw_llr_coords = [[0.0, 0.0, 5.0], + [-pi/2, 0.0, 0.5], + [0.0, -pi/2, 15.0], + [0.0, 0.0, 0.0]] + + raw_xyz_coords = [[5.0, 0.0, 0.0], + [0.0, -0.5, 0.0], + [0.0, 0.0, -15.0], + [0.0, 0.0, 0.0]] + + raw_llr_vectors = [[-6.0, 0.5, 10.0], + [3.0, 1.2, 0.7], + [0.0, 5.0, 2.0], + [0.0, 0.0, 0.0]] + + raw_xyz_vectors = [[10.0, -6.0, 0.5], + [3.0, -0.7, 1.2], + [5.0, 0.0, -2.0], + [0.0, 0.0, 0.0]] + + # Test for firedrake routines requires a mesh + if module_name == "firedrake": + mesh = fd.UnitIntervalMesh(len(raw_llr_coords)) + else: + mesh = None + + # Put the coordinates in numpy arrays or firedrake functions + # Coordinates are Firedrake scalars when testing Firedrake vectors + llr_coords = setup_values(raw_llr_coords, coord_config, mesh=mesh) + xyz_coords = setup_values(raw_xyz_coords, coord_config, mesh=mesh) + llr_vectors = setup_values(raw_llr_vectors, config_name, mesh=mesh) + xyz_vectors = setup_values(raw_xyz_vectors, config_name, mesh=mesh) + + position_vector = xyz_coords if position_units == 'xyz' else llr_coords + + new_llr_vectors = lonlatr_components_from_xyz(xyz_vectors, position_vector, position_units) + new_xyz_vectors = xyz_vector_from_lonlatr(llr_vectors[0], llr_vectors[1], llr_vectors[2], + position_vector, position_units) + + if config_name != 'numpy': + llr_vectors = fd.as_vector(llr_vectors) + new_llr_vectors = fd.as_vector(new_llr_vectors) + + check_values(new_llr_vectors, llr_vectors, config_name, 'lonlatr_components_from_xyz') + check_values(new_xyz_vectors, xyz_vectors, config_name, 'xyz_vector_from_lonlatr')