Example usage of firedrake.fml.form_manipulation_language

[jmbigman]

I am interested in using form labeling to easily add/remove terms from a PDE.

I searched the Firedrake site but found only the API reference and no examples.

Are there examples on the site I'm missing? Alternatively, does anyone have an example script they would be willing to share?

Thanks for the help

[jshipton]

Hi @jbigmanUT ! Apologies, creating an example for this is on my to do list - and has been for quite a while to be honest! Here's an example I've put together for a paper I'm writing on how we use this in Gusto. The example sets up an advection diffusion equation and uses a splitting scheme for time stepping where the advection (or transport) is done explicitly and the diffusion is done implicitly. I'm posting as-is in the interest of speed, but I'd be happy to send you also the text from the (not yet finished!) paper that describes what is actually happening.

from firedrake import (PeriodicIntervalMesh, SpatialCoordinate, Constant,
                       FunctionSpace, Function, TestFunction, TrialFunction,
                       LinearVariationalProblem, LinearVariationalSolver,
                       exp, dx)
from firedrake.fml import all_terms, drop, keep, replace_subject, subject
from firedrake.pyplot import plot
from gusto.core import time_derivative, transport, diffusion
import matplotlib.pyplot as plt

L = 1
nx = 20
mesh = PeriodicIntervalMesh(nx, L)
x = SpatialCoordinate(mesh)[0]

V = FunctionSpace(mesh, "Lagrange", 2)
d = 0.05
qn = Function(V).interpolate(exp(-((x-0.5*L)/d)**2))
plot(qn)
plt.show()
c = Constant(0.5)
kappa = Constant(0.01)

# Equation
q = Function(V)
v = TestFunction(V)
residual = (
    time_derivative(q * v * dx) +
    transport(c * q.dx(0) * v * dx) +
    diffusion(kappa * q.dx(0) * v.dx(0) * dx)
    )
residual = subject(residual, q)

# Set up solvers
q_trial = TrialFunction(V)
qstar = Function(V)
qnp1 = Function(V)
dt = 0.01

# Set up solver for explicit transport step
lhs_explicit = residual.label_map(lambda t: t.has_label(time_derivative),
                                  map_if_true=keep,
                                  map_if_false=drop)
lhs_explicit = lhs_explicit.label_map(all_terms,
                                      map_if_true=replace_subject(q_trial))
rhs_explicit = residual.label_map(lambda t: t.has_label(diffusion),
                                  map_if_true=drop)
rhs_explicit = rhs_explicit.label_map(lambda t: t.has_label(time_derivative),
                                      map_if_false=lambda t: -dt*t)
rhs_explicit = rhs_explicit.label_map(all_terms,
                                      map_if_true=replace_subject(qn))
prob_explicit = LinearVariationalProblem(lhs_explicit.form,
                                         rhs_explicit.form,
                                         qstar)
solver_explicit = LinearVariationalSolver(prob_explicit)

# Set up solver for implicit transport step
lhs_implicit = residual.label_map(lambda t: t.has_label(transport),
                                  map_if_true=drop)
lhs_implicit = lhs_implicit.label_map(lambda t: t.has_label(time_derivative),
                                      map_if_false=lambda t: dt*t)
lhs_implicit = lhs_implicit.label_map(all_terms,
                                      map_if_true=replace_subject(q_trial))
rhs_implicit = residual.label_map(lambda t: t.has_label(time_derivative),
                                      map_if_false=drop)
rhs_implicit = rhs_implicit.label_map(all_terms,
                                      map_if_true=replace_subject(qstar))
prob_implicit = LinearVariationalProblem(lhs_implicit.form,
                                         rhs_implicit.form,
                                         qnp1)
solver_implicit = LinearVariationalSolver(prob_implicit)

# Timestepping loop
tmax = 0.5
t = 0
while t < tmax:
    t += dt
    solver_explicit.solve()
    solver_implicit.solve()
    qn.assign(qnp1)
    plot(qn)
    plt.show()

[jmbigman]

Thank you! I'm also working on operator splits and have an advection-diffusion test code, so this was very helpful.

Do you know if there's a way to add multiple labels to a term in one command? I'm adding multiple labels so I can switch on/off different phenomena and mark them as explicit/implicit.

[tommbendall]

There's actually also an test in Firedrake that is a basic IMEX style example: https://github.com/firedrakeproject/firedrake/blob/master/tests/firedrake/regression/test_fml.py

And yes you can definitely add multiple labels in one line! Or you can apply a new label to a term that already has a label. For instance in that Firedrake test, we have:

transport_form = transport(subject(H*inner(div(u0), phi)*dx, X))

which applies two labels in one go, then later on we add the explicit label to that term.

And for more complicated examples, Gusto has lots of different labels: https://github.com/firedrakeproject/gusto/blob/main/gusto/core/labels.py

[rckirby]

FYI, importing this mechanism into Irksome is a near-term todo on an undergrad research project.On Feb 21, 2025, at 4:29 PM, Joshua Bigman ***@***.***> wrote: Thank you! I'm also working on operator splits and have an advection-diffusion test code, so this was very helpful. Do you know if there's a way to add multiple labels to a term in one command? I'm adding multiple labels so I can switch on/off different phenomena and mark them as explicit/implicit. —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>