Skip to content

Commit

Permalink
Merge branch 'main' into 1dsw
Browse files Browse the repository at this point in the history
  • Loading branch information
@tommbendall
tommbendall authored Mar 28, 2024
2 parents bae2aa1 + 6133055 commit 3de144d
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 11 deletions.
30 changes: 24 additions & 6 deletions gusto/active_tracers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,8 @@ class WaterVapour(ActiveTracer):
"""An object encoding the details of water vapour as a tracer."""
def __init__(self, name='water_vapour', space='theta',
variable_type=TracerVariableType.mixing_ratio,
transport_eqn=TransportEquationType.advective):
transport_eqn=TransportEquationType.advective,
density_name=None):
"""
Args:
name (str, optional): the variable's name. Defaults to
Expand All @@ -113,16 +114,22 @@ def __init__(self, name='water_vapour', space='theta',
transport_eqn (:class:`TransportEquationType`, optional): enumerator
indicating the type of transport equation to be solved (e.g.
advective). Defaults to `TransportEquationType.advective`.
density_name (str): the name of the associated density for a mixing
ratio when using the tracer_conservative transport. Defaults to None,
as this argument is not needed for other transport types.
"""
super().__init__(f'{name}', space, variable_type,
transport_eqn=transport_eqn, phase=Phases.gas, chemical='H2O')
transport_eqn=transport_eqn,
density_name=density_name,
phase=Phases.gas, chemical='H2O')


class CloudWater(ActiveTracer):
"""An object encoding the details of cloud water as a tracer."""
def __init__(self, name='cloud_water', space='theta',
variable_type=TracerVariableType.mixing_ratio,
transport_eqn=TransportEquationType.advective):
transport_eqn=TransportEquationType.advective,
density_name=None):
"""
Args:
name (str, optional): the variable name. Default is 'cloud_water'.
Expand All @@ -134,16 +141,22 @@ def __init__(self, name='cloud_water', space='theta',
transport_eqn (:class:`TransportEquationType`, optional): enumerator
indicating the type of transport equation to be solved (e.g.
advective). Defaults to `TransportEquationType.advective`.
density_name (str): the name of the associated density for a mixing
ratio when using the tracer_conservative transport. Defaults to None,
as this argument is not needed for other transport types.
"""
super().__init__(f'{name}', space, variable_type,
transport_eqn=transport_eqn, phase=Phases.liquid, chemical='H2O')
transport_eqn=transport_eqn,
density_name=density_name,
phase=Phases.liquid, chemical='H2O')


class Rain(ActiveTracer):
"""An object encoding the details of rain as a tracer."""
def __init__(self, name='rain', space='theta',
variable_type=TracerVariableType.mixing_ratio,
transport_eqn=TransportEquationType.advective):
transport_eqn=TransportEquationType.advective,
density_name=None):
"""
Args:
name (str, optional): the name for the variable. Defaults to 'rain'.
Expand All @@ -155,6 +168,11 @@ def __init__(self, name='rain', space='theta',
transport_eqn (:class:`TransportEquationType`, optional): enumerator
indicating the type of transport equation to be solved (e.g.
advective). Defaults to `TransportEquationType.advective`.
density_name (str): the name of the associated density for a mixing
ratio when using the tracer_conservative transport. Defaults to None,
as this argument is not needed for other transport types.
"""
super().__init__(f'{name}', space, variable_type,
transport_eqn=transport_eqn, phase=Phases.liquid, chemical='H2O')
transport_eqn=transport_eqn,
density_name=density_name,
phase=Phases.liquid, chemical='H2O')
55 changes: 50 additions & 5 deletions gusto/diagnostics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \
ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, pi, \
TensorFunctionSpace, SpatialCoordinate, as_vector, \
Projector, Interpolator, FunctionSpace
Projector, Interpolator, FunctionSpace, FiniteElement, \
TensorProductElement
from firedrake.assign import Assigner

from abc import ABCMeta, abstractmethod, abstractproperty
Expand Down Expand Up @@ -1691,7 +1692,12 @@ class TracerDensity(DiagnosticField):
"""Diagnostic for computing the density of a tracer. This is
computed as the product of a mixing ratio and dry density"""

name = "TracerDensity"
@property
def name(self):
"""Gives the name of this diagnostic field. This records
the mixing ratio and density names, in case multiple tracer
densities are used."""
return "TracerDensity_"+self.mixing_ratio_name+'_'+self.density_name

def __init__(self, mixing_ratio_name, density_name, space=None, method='interpolate'):
"""
Expand All @@ -1700,12 +1706,15 @@ def __init__(self, mixing_ratio_name, density_name, space=None, method='interpol
density_name (str): the name of the tracer density variable
space (:class:`FunctionSpace`, optional): the function space to
evaluate the diagnostic field in. Defaults to None, in which
case a default space will be chosen for this diagnostic.
case a new space will be constructed for this diagnostic. This
space will have enough a high enough degree to accurately compute
the product of the mixing ratio and density.
method (str, optional): a string specifying the method of evaluation
for this diagnostic. Valid options are 'interpolate', 'project' and
'assign'. Defaults to 'interpolate'.
"""
super().__init__(method=method, required_fields=(mixing_ratio_name, density_name))
super().__init__(space=space, method=method, required_fields=(mixing_ratio_name, density_name))

self.mixing_ratio_name = mixing_ratio_name
self.density_name = density_name

Expand All @@ -1720,4 +1729,40 @@ def setup(self, domain, state_fields):
m_X = state_fields(self.mixing_ratio_name)
rho_d = state_fields(self.density_name)
self.expr = m_X*rho_d
super().setup(domain, state_fields)

if self.space is None:
# Construct a space for the diagnostic that has enough
# degrees to accurately capture the tracer density. This
# will be the sum of the degrees of the individual mixing ratio
# and density function spaces.
m_X_space = m_X.function_space()
rho_d_space = rho_d.function_space()

if domain.spaces.extruded_mesh:
# Extract the base horizontal and vertical elements
# for the mixing ratio and density.
m_X_horiz = m_X_space.ufl_element().sub_elements[0]
m_X_vert = m_X_space.ufl_element().sub_elements[1]
rho_d_horiz = rho_d_space.ufl_element().sub_elements[0]
rho_d_vert = rho_d_space.ufl_element().sub_elements[1]

horiz_degree = m_X_horiz.degree() + rho_d_horiz.degree()
vert_degree = m_X_vert.degree() + rho_d_vert.degree()

cell = domain.mesh._base_mesh.ufl_cell().cellname()
horiz_elt = FiniteElement('DG', cell, horiz_degree)
vert_elt = FiniteElement('DG', cell, vert_degree)
elt = TensorProductElement(horiz_elt, vert_elt)
else:
m_X_degree = m_X_space.ufl_element().degree()
rho_d_degree = rho_d_space.ufl_element().degree()
degree = m_X_degree + rho_d_degree

cell = domain.mesh.ufl_cell().cellname()
elt = FiniteElement('DG', cell, degree)

tracer_density_space = FunctionSpace(domain.mesh, elt, name='tracer_density_space')
super().setup(domain, state_fields, space=tracer_density_space)

else:
super().setup(domain, state_fields)

0 comments on commit 3de144d

Please sign in to comment.