Skip to content

MMR (mean mixing ratio) limiter update #754

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tommbendall merged 4 commits into from
Aug 21, 2026
Merged

MMR (mean mixing ratio) limiter update #754

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8e4017f
update the mmr implementation
ta440 Jul 14, 2026
04917bd
mmr updates and lint
ta440 Jul 15, 2026
0e68c62
Update gusto/spatial_methods/augmentation.py
ta440 Aug 13, 2026
b7dc49d
Merge branch 'main' into mmr_update
tommbendall Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions gusto/core/kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
"""

from firedrake import dx
from firedrake.parloops import par_loop, READ, WRITE, INC, MIN, MAX, op2
from firedrake.parloops import par_loop, READ, WRITE, RW, MIN, MAX, op2
import numpy as np


Expand Down Expand Up @@ -158,7 +158,7 @@ def apply(self, lamda, mX_field, mean_field):
lives in the continuous target space.
"""
par_loop(self._kernel, dx,
{"lamda": (lamda, INC),
{"lamda": (lamda, RW),
"mX_field": (mX_field, READ),
"mean_field": (mean_field, READ)})

Expand Down
239 changes: 50 additions & 189 deletions gusto/spatial_methods/augmentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,10 @@
transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div,
Constant, sqrt, cross, curl, FunctionSpace, assemble, DirichletBC
)
from firedrake.fml import (
subject, all_terms, replace_subject, replace_test_function,
drop, Term, LabelledForm
)
from firedrake.fml import subject
from gusto import (
time_derivative, transport, transporting_velocity, TransportEquationType,
logger, prognostic, mass_weighted
logger
)
from gusto.spatial_methods.limiters import MeanLimiter
from gusto.core.conservative_projection import ConservativeProjector
Expand Down Expand Up @@ -257,9 +254,10 @@ def update(self, x_in_mixed):

class MeanMixingRatio(Augmentation):
"""
This augments a transport problem involving a mixing ratio to
include a mean mixing ratio field. This enables posivity to be
ensured during conservative transport.
This augments a transport problem involving a k=1 mixing ratio, by adding
a mean mixing ratio field. This enables posivity to be
Comment thread
ta440 marked this conversation as resolved.
Outdated
ensured after each conservative transport step by blending the k=1 and
mean fields.

Args:
domain (:class:`Domain`): The domain object.
Expand All @@ -278,33 +276,26 @@ def __init__(

# Store information about original equation set
self.field_names = []
self.orig_spaces = []
for i in np.arange(len(eqns.field_names)):
self.field_names.append(eqns.field_names[i])
self.orig_spaces.append(eqns.spaces[i])

self.eqn_orig = eqns
self.domain = domain
exist_spaces = eqns.spaces
self.idx_orig = len(exist_spaces)
self.idx_orig = len(self.orig_spaces)

DG0 = FunctionSpace(domain.mesh, "DG", 0)
DG1 = FunctionSpace(domain.mesh, "DG", 1)

# Set up fields and names for each mixing ratio
self.mean_names = []
self.mean_idxs = []
self.mX_idxs = []
mX_spaces = []
mean_spaces = []
self.rho_idxs = []

for i in range(self.mX_num):
mX_name = mX_names[i]
self.mean_names.append('mean_'+mX_name)
self.field_names.append(self.mean_names[-1])
mean_spaces.append(DG0)
exist_spaces.append(DG0)

self.mean_idxs.append(self.idx_orig + i)

# Extract the mixing ratio in question:
mX_idx = eqns.field_names.index(mX_name)
Expand All @@ -313,13 +304,14 @@ def __init__(

# Determine if this is a conservatively transported tracer.
# If so, extract the corresponding density name, if not
# set this to None.
# raise an error as this limiter shouldn't be used.
for tracer in eqns.active_tracers:
if tracer.name == mX_name:
if tracer.density_name is not None:
self.rho_idxs.append(eqns.field_names.index(tracer.density_name))
else:
self.rho_idxs.append('None')
raise ValueError("Mean mixing ratio can't be used for non-conservatively "
+ f"transported mixing ratio {mX_name}")

# Define a limiter using the mean mixing ratios
self.limiters = MeanLimiter(mX_spaces)
Expand All @@ -328,162 +320,36 @@ def __init__(
self.DG1_field = Function(DG1)
self.rho_field = Function(DG1)
self.DG0_field = Function(DG0)
self.compute_mean_mX = ConservativeProjector(self.rho_field, self.rho_field, self.DG1_field, self.DG0_field)

# Create the new mixed function space
self.fs = MixedFunctionSpace(exist_spaces)
# Construct a conservative, consistent, projection operator for computing
# the mean mixing ratios.
self.compute_mean_mX = ConservativeProjector(self.rho_field, self.rho_field, self.DG1_field, self.DG0_field, subtract_mean=True)

self.X = Function(self.fs)
self.tests = TestFunctions(self.fs)
# Set up the function space and residual as
# per the original equation set
self.residual = eqns.residual
self.X = eqns.X
self.tests = eqns.tests
self.fs = MixedFunctionSpace(self.orig_spaces)
self.x_in = Function(self.fs)
self.x_out = Function(self.fs)

self.bcs = None

# Make mean fields
mean_fs = MixedFunctionSpace(mean_spaces)
self.mean_fields = Function(mean_fs)

def setup_residual(self, equation):
"""
Create a new residual for the augmented equation set,
using the larger mixed function space that includes the mean
mixing ratios, and new residual terms for the mean mixing ratio
equations
Copy the residual to the augmentation

Args:
equation (:class:`PrognosticEquationSet`): The overarching equation set.
Note, this does not include the mean mixing ratios.
"""

# Copy the existing residual
new_residual = equation.residual

# Replace test and trial functions of original residual with
# those from the new (larger) mixed function space.
# The indices of the original fields
# are the same in the new mixed space.
for idx in range(self.idx_orig):
new_residual = new_residual.label_map(
all_terms,
replace_subject(self.X, old_idx=idx, new_idx=idx)
)
new_residual = new_residual.label_map(
all_terms,
replace_test_function(self.tests, old_idx=idx, new_idx=idx)
)

# Loop over each mean mixing ratio,
# copy the residual terms relating to the original mixing ratio,
# update the test and trial functions, then add to the new residual.
for i in range(self.mX_num):
mean_residual = new_residual.label_map(
lambda t: t.get(prognostic) == self.mX_names[i],
map_if_false=drop
)

# Replace all instances of original mixing ratios with
# the mean versions
for j in range(self.mX_num):
mean_residual = mean_residual.label_map(
all_terms,
replace_subject(self.X, old_idx=self.mX_idxs[j], new_idx=self.mean_idxs[j])
)

mean_residual = mean_residual.label_map(
all_terms,
replace_test_function(self.tests, old_idx=self.mX_idxs[i], new_idx=self.mean_idxs[i])
)

# Update the name to be that of the mean mixing ratio
mean_residual = mean_residual.label_map(
all_terms,
lambda t: prognostic.update_value(t, self.mean_names[i])
)

# Append to the new residual
new_residual += mean_residual

self.residual = subject(new_residual, self.X)

# For any mass_weighted (conservative transport) terms,
# replace the subject and test functions for
# the mass-weighted form, and update the label to
# point to the new form.
for term in self.residual:
if term.has_label(mass_weighted):
field = term.get(prognostic)
mass_term = term.get(mass_weighted)

# Extract the previous labels for the
# mass-weighted term
if term.has_label(transport):
old_mass_weighted_labels = mass_term.labels

# Transport terms are Terms not LabelledForms,
# so this change this to use the label_map
mass_term = LabelledForm(mass_term)
else:
old_mass_weighted_labels = mass_term.terms[0].labels

if field in self.mX_names:
list_idx = self.mX_names.index(field)
mX_idx = self.mX_idxs[list_idx]
rho_idx = self.rho_idxs[list_idx]

# Replace mixing ratio
mass_term_new = mass_term.label_map(
all_terms,
replace_subject(self.X, old_idx=mX_idx, new_idx=mX_idx)
)

# Replace original density
mass_term_new = mass_term_new.label_map(
all_terms,
replace_subject(self.X, old_idx=rho_idx, new_idx=rho_idx)
)

# Replace test function
mass_term_new = mass_term_new.label_map(
all_terms,
replace_test_function(self.tests, old_idx=mX_idx, new_idx=mX_idx)
)

elif field in self.mean_names:
list_idx = self.mean_names.index(field)
mX_idx = self.mX_idxs[list_idx]
mean_idx = self.mean_idxs[list_idx]
rho_idx = self.rho_idxs[list_idx]

# Replace mixing ratio
mass_term_new = mass_term.label_map(
all_terms,
replace_subject(self.X, old_idx=mX_idx, new_idx=mean_idx)
)

# Replace density
mass_term_new = mass_term_new.label_map(
all_terms,
replace_subject(self.X, old_idx=rho_idx, new_idx=rho_idx)
)

# Replace test function
mass_term_new = mass_term_new.label_map(
all_terms,
replace_test_function(self.tests, old_idx=mX_idx, new_idx=mean_idx)
)

# Create a new mass-weighted term, which has the correct labels
mass_term_new = Term(mass_term_new.form, old_mass_weighted_labels)
mass_term_new = subject(mass_term_new, self.X)

# Make a new term, that links to the new mass-weighted term
new_term = Term(term.form, term.labels)
new_term = mass_weighted.update_value(new_term, mass_term_new)

# Put this new term back in the residual:
self.residual = self.residual.label_map(
lambda t: t == term,
map_if_true=lambda t: new_term
)

self.residual = subject(self.residual, self.X)
self.residual = equation.residual

def pre_apply(self, x_in):
"""
Expand All @@ -492,9 +358,7 @@ def pre_apply(self, x_in):
Args:
x_in (:class:`Function`): The input fields
"""

for idx in range(self.idx_orig):
self.x_in.subfunctions[idx].assign(x_in.subfunctions[idx])
self.x_in.assign(x_in)

def post_apply(self, x_out):
"""
Expand All @@ -503,42 +367,24 @@ def post_apply(self, x_out):
Args:
x_out (:class:`Function`): The output fields
"""

for idx in range(self.idx_orig):
x_out.subfunctions[idx].assign(self.x_out.subfunctions[idx])
x_out.assign(self.x_out)

def update(self, x_in_mixed):
"""
Compute the mean mixing ratio field by conservative projection,
Compute the mean mixing ratio fields by conservative projection,
where both the target and source density are in the higher-order
space.

Args:
x_in_mixed (:class:`Function`): The mixed function, containing
mean fields to update.
"""

# Update the density field for the mean mixing ratios
self.rho_field.assign(x_in_mixed.subfunctions[0])

# Update the mean mixing ratios:
for i in range(self.mX_num):
# Extract the reference density:
self.rho_field.assign(x_in_mixed.subfunctions[self.rho_idxs[i]])

# Compute the mean mixing ratio with conservative projection
self.DG1_field.assign(x_in_mixed.subfunctions[self.mX_idxs[i]])
self.compute_mean_mX.project()

# Clip any minuscule negative values in the mean mixing ratio
self.limiters._clip_means_kernel.apply(self.DG0_field, self.DG0_field)

x_in_mixed.subfunctions[self.mean_idxs[i]].assign(self.DG0_field)
pass

def limit(self, x_in_mixed):
"""
Limit the mixing ratios using a blended limiter with
the mean mixing ratios
Limit k=1 mixing ratios using a limiter that blends the
k=1 field and its mean field.

Args:
x_in_mixed (:class:`Function`): The mixed function, containing
Expand All @@ -549,12 +395,27 @@ def limit(self, x_in_mixed):
mX_pre = []
means = []

# Compute the new mean mixing ratio
for i in range(self.mX_num):
self.rho_field.assign(x_in_mixed.subfunctions[self.rho_idxs[i]])

# Compute the mean mixing ratio with conservative projection
self.DG1_field.assign(x_in_mixed.subfunctions[self.mX_idxs[i]])
self.compute_mean_mX.project()

# Clip any minuscule negative values in the mean mixing ratio
# that arise from numerical error
self.limiters._clip_means_kernel.apply(self.DG0_field, self.DG0_field)

self.mean_fields.subfunctions[i].assign(self.DG0_field)

mX_pre.append(x_in_mixed.subfunctions[self.mX_idxs[i]])
means.append(x_in_mixed.subfunctions[self.mean_idxs[i]])
means.append(self.mean_fields.subfunctions[i])

# Limit all mixing ratios at the same time to have
# the same blending weights.
self.limiters.apply(mX_pre, means)

# Update the mixing ratios with the limited version
for i in range(self.mX_num):
self.limiters._clip_DG1_field.apply(mX_pre[i], mX_pre[i])
x_in_mixed.subfunctions[self.mX_idxs[i]].assign(mX_pre[i])
8 changes: 3 additions & 5 deletions gusto/spatial_methods/limiters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,6 @@ def __init__(self, spaces):

DG1_equispaced = FunctionSpace(mesh, DG1_element)
DG0 = FunctionSpace(mesh, 'DG', 0)
DG1 = FunctionSpace(mesh, 'DG', 1)

self.lamda = Function(DG0)
self.mX_field = Function(DG1_equispaced)
Expand All @@ -312,10 +311,9 @@ def __init__(self, spaces):

self._lamda_kernel = MeanMixingRatioWeights(DG1_equispaced)

# Also construct kernels to clip any very small negatives
# that arise from numerical error. These are used in the
# mean mixing ratio augmentation limit routine.
self._clip_DG1_field = ClipZero(DG1)
# Also construct a kernels to clip any very small negatives
# that arise from numerical error when computing the
# mean mixing ratio.
self._clip_means_kernel = ClipZero(DG0)

def apply(self, mX_fields, mean_fields):
Expand Down
Loading
Loading