Utility

utility

This module provides several classes and functions to perform a number of pre-, syn-, and post-processing tasks. Users incorporate utility as required in their code, depending on what they would like to achieve.

TimestepAdaptor(dt_const, u, V, target_cfl=1.0, increase_tolerance=1.5, maximum_timestep=None)

Computes timestep based on CFL condition for provided velocity field

Parameters:

Name	Type	Description	Default
dt_const		Constant whose value will be updated by the timestep adaptor	required
u		Velocity to base CFL condition on	required
V		FunctionSpace for reference velocity, usually velocity space	required
target_cfl		CFL number to target with chosen timestep	1.0
increase_tolerance		Maximum tolerance timestep is allowed to change by	1.5
maximum_timestep		Maximum allowable timestep	None

Source code in g-adopt/gadopt/utility.py

def __init__(self, dt_const, u, V, target_cfl=1.0, increase_tolerance=1.5, maximum_timestep=None):
    self.dt_const = dt_const
    self.u = u
    self.target_cfl = target_cfl
    self.increase_tolerance = increase_tolerance
    self.maximum_timestep = maximum_timestep
    self.mesh = V.mesh()

    # J^-1 u is a discontinuous expression, using op2.MAX it takes the maximum value
    # in all adjacent elements when interpolating it to a continuous function space
    # We do need to ensure we reset ref_vel to zero, as it also takes the max with any previous values
    self.ref_vel_interpolate = interpolate(abs(dot(JacobianInverse(self.mesh), self.u)), V, access=op2.MAX)

CombinedSurfaceMeasure(domain, degree)

Bases: Measure

A surface measure that combines ds_v, the integral over vertical boundary facets, and ds_t and ds_b, the integral over horizontal top and bottom facets. The vertical boundary facets are identified with the same surface ids as ds_v. The top and bottom surfaces are identified via the "top" and "bottom" ids.

Source code in g-adopt/gadopt/utility.py

def __init__(self, domain, degree):
    self.ds_v = ds_v(domain=domain, degree=degree)
    self.ds_t = ds_t(domain=domain, degree=degree)
    self.ds_b = ds_b(domain=domain, degree=degree)

__rmul__(other)

This is to handle terms to be integrated over all surfaces in the form of other*ds. Here the CombinedSurfaceMeasure ds is not called, instead we just split it up as below.

Source code in g-adopt/gadopt/utility.py

def __rmul__(self, other):
    """This is to handle terms to be integrated over all surfaces in the form of other*ds.
    Here the CombinedSurfaceMeasure ds is not called, instead we just split it up as below."""
    return other*self.ds_v + other*self.ds_t + other*self.ds_b

ExtrudedFunction(*args, mesh_3d=None, **kwargs)

Bases: Function

A 2D Function that provides a 3D view on the extruded domain.

The 3D function can be accessed as ExtrudedFunction.view_3d. The 3D function resides in V x R function space, where V is the function space of the source function. The 3D function shares the data of the 2D function.

Parameters:

Name	Type	Description	Default
mesh_3d		Extruded 3D mesh where the function will be extended to.	None

Source code in g-adopt/gadopt/utility.py

def __init__(self, *args, mesh_3d=None, **kwargs):
    # create the 2d function
    super().__init__(*args, **kwargs)
    print(*args)
    if mesh_3d is not None:
        self.view_3d = extend_function_to_3d(self, mesh_3d)

LayerAveraging(mesh, r1d=None, quad_degree=None)

A manager for computing a vertical profile of horizontal layer averages.

Parameters:

Name	Type	Description	Default
mesh		The mesh over which to compute averages	required
r1d		An array of either depth coordinates or radii, at which to compute layer averages. If not provided, and mesh is extruded, it uses the same layer heights. If mesh is not extruded, r1d is required.	None

Source code in g-adopt/gadopt/utility.py

def __init__(self, mesh, r1d=None, quad_degree=None):
    self.mesh = mesh
    XYZ = SpatialCoordinate(mesh)

    if is_cartesian(mesh):
        self.r = XYZ[len(XYZ)-1]
    else:
        self.r = sqrt(dot(XYZ, XYZ))

    self.dx = dx
    if quad_degree is not None:
        self.dx = dx(degree=quad_degree)

    if r1d is not None:
        self.r1d = r1d
    else:
        try:
            nlayers = mesh.layers
        except AttributeError:
            raise ValueError("For non-extruded mesh need to specify depths array r1d.")
        CG1 = FunctionSpace(mesh, "CG", 1)
        r_func = Function(CG1).interpolate(self.r)
        self.r1d = r_func.dat.data[:nlayers]

    self.mass = np.zeros((2, len(self.r1d)))
    self.rhs = np.zeros(len(self.r1d))
    self._assemble_mass()

get_layer_average(T)

Compute the layer averages of Function T at the predefined depths.

Returns:

Type	Description
	A numpy array containing the averages.

Source code in g-adopt/gadopt/utility.py

def get_layer_average(self, T):
    """Compute the layer averages of `Function` T at the predefined depths.

    Returns:
      A numpy array containing the averages.
    """

    self._assemble_rhs(T)
    return solveh_banded(self.mass, self.rhs, lower=True)

extrapolate_layer_average(u, avg)

Given an array of layer averages avg, extrapolate to Function u

Source code in g-adopt/gadopt/utility.py

def extrapolate_layer_average(self, u, avg):
    """Given an array of layer averages avg, extrapolate to `Function` u
    """

    r = self.r
    rc = Constant(self.r1d[0])
    rn = Constant(self.r1d[1])
    rp = Constant(0.)

    u.assign(0.0)

    phi = max_value(min_value((r - rp) / (rc - rp), (rn - r) / (rn - rc)), 0)
    val = Constant(0.)

    for a, rin in zip(avg[:-1], self.r1d[1:]):
        val.assign(a)
        rn.assign(rin)
        # reconstruct this layer according to the basis function
        u.interpolate(u + val * phi)

        rp.assign(rc)
        rc.assign(rn)

    phi = max_value(min_value(1, (r - rp) / (rn - rp)), 0)
    val.assign(avg[-1])
    u.interpolate(u + val * phi)

InteriorBC

Bases: DirichletBC

DirichletBC applied to anywhere that is not on the specified boundary

log(*args, level=None)

Log output to stdout from root processor only.

If level is provided, the message is only printed when that level is enabled by the GADOPT_LOGLEVEL environment variable.

Source code in g-adopt/gadopt/utility.py

def log(*args, level=None):
    """Log output to stdout from root processor only.

    If level is provided, the message is only printed when that level is
    enabled by the GADOPT_LOGLEVEL environment variable.
    """
    if level is not None and level < log_level:
        return
    PETSc.Sys.Print(*args)

horizontal_components(u)

Returns expression for vector u with vertical component removed, i.e. returns the vector UFL expression orthogonal to gravity.

Source code in g-adopt/gadopt/utility.py

def horizontal_components(u):
    """Returns expression for vector u with vertical component removed, i.e. returns the vector UFL expression orthogonal to gravity."""
    mesh = extract_unique_domain(u)
    return u - vertical_component(u) * upward_normal(mesh)

get_device_type(gpu_options)

Determine the available GPU device type.

Checks for a GPU device in the following order of priority: 1. The GADOPT_DEVICE_OVERRIDE environment variable. 2. The device_type key in the gpu_options dictionary. 3. Firedrake's get_device_type() function.

Returns:

Type	Description
str | None	The detected device type (e.g., CUDA, HIP). Returns None if the HOST
str | None	device is detected (i.e. no GPU) or if the device set in gpu_options or
str | None	returned by Firedrake's get_device_type() is None.

Raises:

Type	Description
ValueError	If GADOPT_DEVICE_OVERRIDE or gpu_options["device_type"]

Source code in g-adopt/gadopt/utility.py

def get_device_type(gpu_options: dict | None) -> str | None:
    """Determine the available GPU device type.

    Checks for a GPU device in the following order of priority:
    1. The `GADOPT_DEVICE_OVERRIDE` environment variable.
    2. The `device_type` key in the `gpu_options` dictionary.
    3. Firedrake's `get_device_type()` function.

    Returns:
        The detected device type (e.g., `CUDA`, `HIP`). Returns `None` if the `HOST`
        device is detected (i.e. no GPU) or if the device set in `gpu_options` or
        returned by Firedrake's `get_device_type()` is `None`.

    Raises:
        ValueError: If `GADOPT_DEVICE_OVERRIDE` or `gpu_options["device_type"]`
        specifies a device type not in `valid_devices`.
    """
    if "GADOPT_DEVICE_OVERRIDE" in os.environ:
        device_type = os.environ["GADOPT_DEVICE_OVERRIDE"]
        if device_type not in valid_devices:
            raise ValueError(
                f"Invalid device specified in GADOPT_DEVICE_OVERRIDE. Valid values are {','.join(valid_devices)}"
            )
    elif gpu_options is not None and "device_type" in gpu_options:
        device_type = gpu_options["device_type"]
        if device_type is not None and device_type not in valid_devices:
            raise ValueError(
                f"Invalid device specified in 'device_type' key. Valid values are {','.join(valid_devices)} or None"
            )
    else:
        # Set the PETSc error handler so that Firedrake can correctly intercept the
        # exception from PETSc if a GPU-enabled build calls this on a non-GPU enabled
        # system - add this to Firedrake. This function returns None for a GPU-enabled
        # build or HOST on a non-GPU enabled build when GPUs are not available.
        PETSc.Sys.pushErrorHandler("python")
        device_type = fd_get_device_type()
        PETSc.Sys.popErrorHandler()
    if device_type == "HOST":
        return None
    return device_type

depends_on(ufl_expr, terminal)

Does ufl_expr depend on terminal (Function/Constant/...)?

Source code in g-adopt/gadopt/utility.py

def depends_on(ufl_expr, terminal):
    """Does ufl_expr depend on terminal (Function/Constant/...)?"""
    return terminal in traverse_unique_terminals(ufl_expr)

tensor_jump(v, n)

Jump term for vector functions based on the tensor product

\["jump"(bb u, bb n) = (bb u^+ bb n^+) + (bb u^- bb n^-)\]

This is the discrete equivalent of grad(u) as opposed to the vectorial UFL jump operator ufl.jump which represents div(u). The equivalent of nabla_grad(u) is given by tensor_jump(n, u).

Source code in g-adopt/gadopt/utility.py

def tensor_jump(v, n):
    r"""
    Jump term for vector functions based on the tensor product

    $$"jump"(bb u, bb n) = (bb u^+ bb n^+) + (bb u^- bb n^-)$$

    This is the discrete equivalent of grad(u) as opposed to the
    vectorial UFL jump operator `ufl.jump` which represents div(u).
    The equivalent of nabla_grad(u) is given by tensor_jump(n, u).
    """
    return outer(v('+'), n('+')) + outer(v('-'), n('-'))

extend_function_to_3d(func, mesh_extruded)

Returns a 3D view of a 2D Function on the extruded domain. The 3D function resides in V x R function space, where V is the function space of the source function. The 3D function shares the data of the 2D function.

Source code in g-adopt/gadopt/utility.py

def extend_function_to_3d(func, mesh_extruded):
    """
    Returns a 3D view of a 2D `Function` on the extruded domain.
    The 3D function resides in V x R function space, where V is the function
    space of the source function. The 3D function shares the data of the 2D
    function.
    """
    fs = func.function_space()
#    assert fs.mesh().geometric_dimension == 2, 'Function must be in 2D space'
    ufl_elem = fs.ufl_element()
    family = ufl_elem.family()
    degree = ufl_elem.degree()
    name = func.name()
    if isinstance(ufl_elem, VectorElement):
        # vector function space
        fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0, dim=2, vector=True)
    else:
        fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0)
    func_extended = Function(fs_extended, name=name, val=func.dat._data)
    func_extended.source = func
    return func_extended

absv(u)

Component-wise absolute value of vector for SU stabilisation

Source code in g-adopt/gadopt/utility.py

def absv(u):
    """Component-wise absolute value of vector for SU stabilisation"""
    return as_vector([abs(ui) for ui in u])

node_coordinates(function)

Interpolates mesh coordinates at each node of the provided function.

Parameters:

Name	Type	Description	Default
function	Function	A Firedrake function	required

Returns:

Type	Description
Function	A Firedrake function for the interpolated mesh coordinates

Source code in g-adopt/gadopt/utility.py

def node_coordinates(function: Function) -> Function:
    """Interpolates mesh coordinates at each node of the provided function.

    Args:
      function: A Firedrake function

    Returns:
      A Firedrake function for the interpolated mesh coordinates
    """
    vec_space = VectorFunctionSpace(function.ufl_domain(), function.ufl_element())

    return Function(vec_space).interpolate(SpatialCoordinate(function))

interpolate_1d_profile(function, one_d_filename)

Assign a one-dimensional profile to a Function function from a file.

The function reads a one-dimensional profile (e.g., viscosity) together with the radius/height input for the profile, from a file, broadcasts the two arrays to all processes, and then interpolates this array onto the function space of function.

Parameters:

Name	Type	Description	Default
function	Function	The function onto which the 1D profile will be assigned	required
one_d_filename	str	The path to the file containing the 1D radial profile	required

Note

• This is designed to read a file with one process and distribute in parallel with MPI.
• The input file should contain an array of radius/height and an array of values, separated by a comma.

Source code in g-adopt/gadopt/utility.py

def interpolate_1d_profile(function: Function, one_d_filename: str):
    """
    Assign a one-dimensional profile to a Function `function` from a file.

    The function reads a one-dimensional profile (e.g., viscosity) together with the
    radius/height input for the profile, from a file, broadcasts the two arrays to all
    processes, and then interpolates this array onto the function space of `function`.

    Args:
        function: The function onto which the 1D profile will be assigned
        one_d_filename: The path to the file containing the 1D radial profile

    Note:
        - This is designed to read a file with one process and distribute in parallel with MPI.
        - The input file should contain an array of radius/height and an array of values, separated by a comma.
    """
    mesh = extract_unique_domain(function)

    if mesh.comm.rank == 0:
        rshl, one_d_data = np.loadtxt(one_d_filename, unpack=True, delimiter=",")
        if rshl[1] < rshl[0]:
            rshl = rshl[::-1]
            one_d_data = one_d_data[::-1]
        if not np.all(np.diff(rshl) > 0):
            raise ValueError("Data must be strictly monotonous.")
    else:
        one_d_data = None
        rshl = None

    one_d_data = mesh.comm.bcast(one_d_data, root=0)
    rshl = mesh.comm.bcast(rshl, root=0)

    X = SpatialCoordinate(mesh)

    upward_coord = vertical_component(X)

    rad = Function(function.function_space()).interpolate(upward_coord)

    averager = LayerAveraging(mesh, rshl if mesh.layers is None else None)
    interpolated_visc = np.interp(averager.get_layer_average(rad), rshl, one_d_data)
    averager.extrapolate_layer_average(function, interpolated_visc)

extruded_layer_heights(DG0_layers, radii)

Calculates layer heights for extruded mesh using rheological boundary

Parameters:

Name	Type	Description	Default
DG0_layers	int | list[int]	Number of layers per rheological layer	required
radii	list[float]	Radii of rheological boundaries to match when extruding the mesh. Assumes sequence starts with outer radii -> inner radii.	required

Returns:

Type	Description
list[float]	list of layer heights

Source code in g-adopt/gadopt/utility.py

def extruded_layer_heights(
        DG0_layers: int | list[int], radii: list[float]) -> list[float]:

    """Calculates layer heights for extruded mesh using rheological boundary

    Args:
      DG0_layers:
        Number of layers per rheological layer
      radii:
        Radii of rheological boundaries to match when extruding the mesh.
        Assumes sequence starts with outer radii -> inner radii.

    Returns:
        list of layer heights
        """

    layer_heights = []

    nz_layers = [DG0_layers] * len(radii) if isinstance(DG0_layers, int) else DG0_layers
    assert len(nz_layers) == len(radii)

    # starting at the bottom radius, work outwards
    for i in range(len(radii) - 2, -1, -1):
        dz = (radii[i] - radii[i + 1]) / nz_layers[i]
        for _ in range(nz_layers[i]):
            layer_heights.append(dz)
    return layer_heights

initialise_background_field(f, background_values, X, radii, shift=0.0)

Initialises discontinuous field with sharp jumps at rheological boundaries

Parameters:

Name	Type	Description	Default
f	Function	Function to interpolate background values into. N.b. f is modified in place.	required
background_values	list[float]	Discrete values to interpolate into f	required
X	SpatialCoordinate	Spatial coordinates associated with function f	required
radii	list[float]	Radii of rheological discontinuities. N.b. length of background_values must be one less than length of radii.	required
shift	float	Shift vertical radii by a constant, e.g. when mesh is offset from radii values.	0.0

Source code in g-adopt/gadopt/utility.py

def initialise_background_field(
        f: Function,
        background_values: list[float],
        X: ufl.geometry.SpatialCoordinate,
        radii: list[float],
        shift: float = 0.0,):
    """Initialises discontinuous field with sharp jumps at rheological boundaries

    Args:
      f:
        Function to interpolate background values into. N.b. `f` is modified in place.
      background_values:
        Discrete values to interpolate into `f`
      X:
        Spatial coordinates associated with function `f`
      radii:
        Radii of rheological discontinuities. N.b. length of `background_values` must be one less
        than length of `radii`.
      shift:
        Shift vertical radii by a constant, e.g. when mesh is offset from `radii` values.
    """

    assert len(background_values) == len(radii)-1
    for i in range(len(background_values)):
        f.interpolate(
            conditional(vertical_component(X) + shift > radii[i+1],
                        conditional(vertical_component(X) + shift <= radii[i],
                                    background_values[i], f), f))