Multiple external operators in a single form result in the residual structure being different #46
Description
Hi, I have been trying to get external operators to work for the case where i have a external function that is a function of two variables i.e., f(a,b), where i have two coupled pdes for tracking a and b.
When i set it up so that the form has multiple external operators, the code doesnt work as expected. As a work-around, by setting it up so that each form has only one external operator, things seem to work fine.
Here is a MWE demonstrating the issue:
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import ufl
import ufl.algorithms
from basix.ufl import element, mixed_element, quadrature_element
from dolfinx import fem, mesh
from dolfinx.fem.petsc import (
create_vector,
assemble_vector,
create_matrix,
assemble_matrix,
)
from ufl import Measure, TrialFunction, TestFunctions, derivative, grad, inner, ln
from dolfinx_external_operator import (
FEMExternalOperator,
evaluate_external_operators,
evaluate_operands,
replace_external_operators,
)
# ============================================================
# USER SETTINGS
# ============================================================
Lx = 20.0
Ly = 20.0
nx = 24
ny = 24
fe_order = 2
chi12 = 12.0
chi13 = 12.0
chi23 = 12.0
N1 = 1.0
N2 = 1.0
N3 = 1.0
dt = 1.0e-3
external_quadrature_degree = 10
modes_to_test = ("raw", "reduced")
# Turn on if you also want Jacobian comparisons
compare_jacobian = False
# Print a few nonzero entries of the difference vector
print_sample_differences = False
# ============================================================
# Helper functions
# ============================================================
def kappa_coeffs(chi12, chi13, chi23, N1, N2, N3):
asym_n1n2 = N2 / N1
asym_n1n3 = N3 / N1
if asym_n1n3 > 0.1:
kappa1 = (2.0 / 3.0) * chi13
else:
kappa1 = (1.0 / 3.0) * chi13
if asym_n1n3 > 0.1 and asym_n1n2 > 0.1:
kappa2 = (2.0 / 3.0) * chi23
elif asym_n1n2 > 0.1 or asym_n1n3 > 0.1:
kappa2 = (1.0 / 3.0) * chi23
else:
kappa2 = 0.0
if asym_n1n2 > 0.1 and asym_n1n3 > 0.1:
kappa12 = (1.0 / 3.0) * chi13 + (1.0 / 3.0) * chi23 - (1.0 / 3.0) * chi12
elif asym_n1n2 > 0.1 and asym_n1n3 <= 0.1:
kappa12 = (1.0 / 6.0) * chi13 + (1.0 / 6.0) * chi23 - (1.0 / 3.0) * chi12
elif asym_n1n2 <= 0.1 and asym_n1n3 > 0.1:
kappa12 = (1.0 / 3.0) * chi13 + (1.0 / 6.0) * chi23 - (1.0 / 6.0) * chi12
else:
kappa12 = (1.0 / 6.0) * chi13 - (1.0 / 6.0) * chi12
return kappa1, kappa2, kappa12
def interpolate_state(u, u0, Lx, Ly):
"""
Build smooth admissible fields for a,b and nonzero fields for mu's.
Keeps a,b and c=1-a-b safely inside the simplex.
"""
twopi = 2.0 * np.pi
def a_fun(x):
return (
0.24
+ 0.03 * np.sin(twopi * x[0] / Lx) * np.sin(twopi * x[1] / Ly)
+ 0.01 * np.cos(2.0 * twopi * x[0] / Lx)
)
def b_fun(x):
return (
0.31
+ 0.025 * np.cos(twopi * x[0] / Lx) * np.sin(twopi * x[1] / Ly)
+ 0.01 * np.sin(2.0 * twopi * x[1] / Ly)
)
def a0_fun(x):
return (
0.24
+ 0.026 * np.sin(twopi * x[0] / Lx) * np.sin(twopi * x[1] / Ly)
+ 0.008 * np.cos(2.0 * twopi * x[0] / Lx)
)
def b0_fun(x):
return (
0.31
+ 0.022 * np.cos(twopi * x[0] / Lx) * np.sin(twopi * x[1] / Ly)
+ 0.008 * np.sin(2.0 * twopi * x[1] / Ly)
)
def mu_ab_fun(x):
return 0.07 * np.sin(twopi * x[0] / Lx) + 0.02 * np.cos(twopi * x[1] / Ly)
def mu_ac_fun(x):
return 0.05 * np.cos(twopi * x[0] / Lx) - 0.03 * np.sin(twopi * x[1] / Ly)
def mu_bc_fun(x):
return 0.04 * np.sin(twopi * x[0] / Lx) * np.cos(twopi * x[1] / Ly)
u.x.array[:] = 0.0
u0.x.array[:] = 0.0
u.sub(0).interpolate(a_fun)
u.sub(1).interpolate(b_fun)
u.sub(2).interpolate(mu_ab_fun)
u.sub(3).interpolate(mu_ac_fun)
u.sub(4).interpolate(mu_bc_fun)
u0.sub(0).interpolate(a0_fun)
u0.sub(1).interpolate(b0_fun)
u0.sub(2).interpolate(mu_ab_fun)
u0.sub(3).interpolate(mu_ac_fun)
u0.sub(4).interpolate(mu_bc_fun)
u.x.scatter_forward()
u0.x.scatter_forward()
def assemble_vector_from_form_expr(F_expr):
F_form = fem.form(F_expr)
b = create_vector(F_form)
with b.localForm() as loc:
loc.set(0.0)
assemble_vector(b, F_form)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
return b
def assemble_matrix_from_jacobian_expr(J_expr):
J_form = fem.form(J_expr)
A = create_matrix(J_form)
A.zeroEntries()
assemble_matrix(A, J_form)
A.assemble()
return A
def assemble_matrix_from_form_expr(F_expr, u, du):
J_expr = derivative(F_expr, u, du)
J_expr = ufl.algorithms.expand_derivatives(J_expr)
return assemble_matrix_from_jacobian_expr(J_expr)
def assemble_external_vector_from_form_expr(F_expr):
"""
Replace/evaluate external operators if present, then assemble the vector.
If no external operators are present, falls back to ordinary assembly.
"""
F_replaced, F_external_operators = replace_external_operators(F_expr)
if len(F_external_operators) > 0:
evaluated_operands = evaluate_operands(F_external_operators)
_ = evaluate_external_operators(F_external_operators, evaluated_operands)
return assemble_vector_from_form_expr(F_replaced)
def assemble_external_matrix_from_form_expr(F_expr, u, du):
"""
Replace/evaluate external operators if present, then assemble the Jacobian.
If no external operators are present, falls back to ordinary assembly.
"""
J_expr = derivative(F_expr, u, du)
J_expr = ufl.algorithms.expand_derivatives(J_expr)
J_replaced, J_external_operators = replace_external_operators(J_expr)
if len(J_external_operators) > 0:
evaluated_operands = evaluate_operands(J_external_operators)
_ = evaluate_external_operators(J_external_operators, evaluated_operands)
return assemble_matrix_from_jacobian_expr(J_replaced)
def vector_diff_norms(v1, v2):
diff = v1.duplicate()
v1.copy(result=diff)
diff.axpy(-1.0, v2)
return {
"ufl_l2": v1.norm(PETSc.NormType.NORM_2),
"ext_l2": v2.norm(PETSc.NormType.NORM_2),
"diff_l2": diff.norm(PETSc.NormType.NORM_2),
"ufl_inf": v1.norm(PETSc.NormType.NORM_INFINITY),
"ext_inf": v2.norm(PETSc.NormType.NORM_INFINITY),
"diff_inf": diff.norm(PETSc.NormType.NORM_INFINITY),
"diff_vec": diff,
}
def matrix_diff_norms(A1, A2):
D = A1.copy()
D.axpy(-1.0, A2, structure=PETSc.Mat.Structure.DIFFERENT_NONZERO_PATTERN)
return {
"ufl_frob": A1.norm(PETSc.NormType.NORM_FROBENIUS),
"ext_frob": A2.norm(PETSc.NormType.NORM_FROBENIUS),
"diff_frob": D.norm(PETSc.NormType.NORM_FROBENIUS),
"ufl_inf": A1.norm(PETSc.NormType.NORM_INFINITY),
"ext_inf": A2.norm(PETSc.NormType.NORM_INFINITY),
"diff_inf": D.norm(PETSc.NormType.NORM_INFINITY),
}
def print_vector_report(name, report):
print(f"\n{name}")
print("-" * len(name))
print(f" ||UFL||_2 = {report['ufl_l2']:.6e}")
print(f" ||External||_2 = {report['ext_l2']:.6e}")
print(f" ||Diff||_2 = {report['diff_l2']:.6e}")
print(f" ||UFL||_inf = {report['ufl_inf']:.6e}")
print(f" ||External||_inf = {report['ext_inf']:.6e}")
print(f" ||Diff||_inf = {report['diff_inf']:.6e}")
if print_sample_differences:
arr = report["diff_vec"].getArray(readonly=True)
nz = np.where(np.abs(arr) > 0)[0]
if nz.size > 0:
print(" first few nonzero diff entries:")
for idx in nz[:10]:
print(f" idx {idx:6d}: {arr[idx]: .6e}")
def print_matrix_report(name, report):
print(f"\n{name}")
print("-" * len(name))
print(f" ||UFL||_F = {report['ufl_frob']:.6e}")
print(f" ||External||_F = {report['ext_frob']:.6e}")
print(f" ||Diff||_F = {report['diff_frob']:.6e}")
print(f" ||UFL||_inf = {report['ufl_inf']:.6e}")
print(f" ||External||_inf = {report['ext_inf']:.6e}")
print(f" ||Diff||_inf = {report['diff_inf']:.6e}")
# ============================================================
# Build mesh and FE space
# ============================================================
domain = mesh.create_rectangle(
MPI.COMM_WORLD,
[[0.0, 0.0], [Lx, Ly]],
[nx, ny],
)
dx = Measure("dx", domain=domain)
P1 = element("Lagrange", domain.basix_cell(), fe_order)
ME = fem.functionspace(domain, mixed_element([P1, P1, P1, P1, P1]))
q1, q2, v1, v2, v3 = TestFunctions(ME)
du = TrialFunction(ME)
u = fem.Function(ME)
u0 = fem.Function(ME)
interpolate_state(u, u0, Lx, Ly)
a_sym, b_sym, mu_AB, mu_AC, mu_BC = ufl.split(u)
a0_sym, b0_sym, _, _, _ = ufl.split(u0)
kappa1, kappa2, kappa12 = kappa_coeffs(chi12, chi13, chi23, N1, N2, N3)
# ------------------------------------------------------------
# Quadrature space and matching dx_q for external operators
# ------------------------------------------------------------
Qe = quadrature_element(
domain.topology.cell_name(),
degree=external_quadrature_degree,
value_shape=(),
)
Q = fem.functionspace(domain, Qe)
dx_q = Measure(
"dx",
domain=domain,
metadata={
"quadrature_scheme": "default",
"quadrature_degree": external_quadrature_degree,
},
)
# ============================================================
# Exact analytical expressions
# ============================================================
c_sym = 1.0 - a_sym - b_sym
dfda_expr = (1.0 / N1) * ln(a_sym) + (1.0 / N1) + chi12 * b_sym + chi13 * c_sym
dfdb_expr = (1.0 / N2) * ln(b_sym) + (1.0 / N2) + chi12 * a_sym + chi23 * c_sym
dfdc_expr = (1.0 / N3) * ln(c_sym) + (1.0 / N3) + chi13 * a_sym + chi23 * b_sym
gAB_expr = dfda_expr - dfdb_expr
gAC_expr = dfda_expr - dfdc_expr
gBC_expr = dfdb_expr - dfdc_expr
# ============================================================
# External-operator analytical callbacks
# ============================================================
def make_external_operator(name, callback_factory):
op = FEMExternalOperator(a_sym, b_sym, function_space=Q)
op.external_function = callback_factory
return op
def make_exact_raw_external(which):
def value_impl(a, b):
c = 1.0 - a - b
if which == "a":
out = (1.0 / N1) * np.log(a) + (1.0 / N1) + chi12 * b + chi13 * c
elif which == "b":
out = (1.0 / N2) * np.log(b) + (1.0 / N2) + chi12 * a + chi23 * c
elif which == "c":
out = (1.0 / N3) * np.log(c) + (1.0 / N3) + chi13 * a + chi23 * b
else:
raise ValueError("which must be 'a', 'b', or 'c'.")
return out.reshape(-1)
def dx_impl(a, b):
c = 1.0 - a - b
if which == "a":
out = 1.0 / (N1 * a) - chi13
elif which == "b":
out = np.full_like(a, chi12 - chi23)
elif which == "c":
out = chi13 - 1.0 / (N3 * c)
else:
raise ValueError("which must be 'a', 'b', or 'c'.")
return out.reshape(-1)
def dy_impl(a, b):
c = 1.0 - a - b
if which == "a":
out = np.full_like(a, chi12 - chi13)
elif which == "b":
out = 1.0 / (N2 * b) - chi23
elif which == "c":
out = chi23 - 1.0 / (N3 * c)
else:
raise ValueError("which must be 'a', 'b', or 'c'.")
return out.reshape(-1)
def external(derivatives):
if derivatives == (0, 0):
return value_impl
elif derivatives == (1, 0):
return dx_impl
elif derivatives == (0, 1):
return dy_impl
else:
raise NotImplementedError(f"Unsupported derivative multi-index: {derivatives}")
return external
def make_exact_reduced_external(which):
def value_impl(a, b):
c = 1.0 - a - b
dfdphi1 = (1.0 / N1) * np.log(a) + (1.0 / N1) + chi12 * b + chi13 * c
dfdphi2 = (1.0 / N2) * np.log(b) + (1.0 / N2) + chi12 * a + chi23 * c
dfdphi3 = (1.0 / N3) * np.log(c) + (1.0 / N3) + chi13 * a + chi23 * b
if which == "AB":
out = dfdphi1 - dfdphi2
elif which == "AC":
out = dfdphi1 - dfdphi3
elif which == "BC":
out = dfdphi2 - dfdphi3
else:
raise ValueError("which must be 'AB', 'AC', or 'BC'.")
return out.reshape(-1)
def dx_impl(a, b):
c = 1.0 - a - b
if which == "AB":
out = 1.0 / (N1 * a) - chi13 - chi12 + chi23
elif which == "AC":
out = 1.0 / (N1 * a) + 1.0 / (N3 * c) - 2.0 * chi13
elif which == "BC":
out = chi12 - chi13 - chi23 + 1.0 / (N3 * c)
else:
raise ValueError("which must be 'AB', 'AC', or 'BC'.")
return out.reshape(-1)
def dy_impl(a, b):
c = 1.0 - a - b
if which == "AB":
out = chi12 - chi13 - 1.0 / (N2 * b) + chi23
elif which == "AC":
out = chi12 - chi13 - chi23 + 1.0 / (N3 * c)
elif which == "BC":
out = 1.0 / (N2 * b) + 1.0 / (N3 * c) - 2.0 * chi23
else:
raise ValueError("which must be 'AB', 'AC', or 'BC'.")
return out.reshape(-1)
def external(derivatives):
if derivatives == (0, 0):
return value_impl
elif derivatives == (1, 0):
return dx_impl
elif derivatives == (0, 1):
return dy_impl
else:
raise NotImplementedError(f"Unsupported derivative multi-index: {derivatives}")
return external
# ============================================================
# Common transport equations
# These remain on the default dx
# ============================================================
F_a = (
inner(a_sym, q1) * dx
- inner(a0_sym, q1) * dx
+ dt * a_sym * b_sym * inner(grad(mu_AB), grad(q1)) * dx
+ dt * a_sym * (1.0 - a_sym - b_sym) * inner(grad(mu_AC), grad(q1)) * dx
)
F_b = (
inner(b_sym, q2) * dx
- inner(b0_sym, q2) * dx
- dt * a_sym * b_sym * inner(grad(mu_AB), grad(q2)) * dx
+ dt * b_sym * (1.0 - a_sym - b_sym) * inner(grad(mu_BC), grad(q2)) * dx
)
# ============================================================
# Compare by mode
# ============================================================
for mode in modes_to_test:
if MPI.COMM_WORLD.rank == 0:
print("\n" + "=" * 80)
print(f"MODE = {mode.upper()}")
print("=" * 80)
if mode == "raw":
# ---------- UFL blocks ----------
F_mu_AB_ufl = (
inner(mu_AB, v1) * dx
- inner(dfda_expr, v1) * dx
+ inner(dfdb_expr, v1) * dx
- (kappa1 - kappa12) * inner(grad(a_sym), grad(v1)) * dx
+ (kappa2 - kappa12) * inner(grad(b_sym), grad(v1)) * dx
)
F_mu_AC_ufl = (
inner(mu_AC, v2) * dx
- inner(dfda_expr, v2) * dx
+ inner(dfdc_expr, v2) * dx
- kappa1 * inner(grad(a_sym), grad(v2)) * dx
- kappa12 * inner(grad(b_sym), grad(v2)) * dx
)
F_mu_BC_ufl = (
inner(mu_BC, v3) * dx
- inner(dfdb_expr, v3) * dx
+ inner(dfdc_expr, v3) * dx
- kappa12 * inner(grad(a_sym), grad(v3)) * dx
- kappa2 * inner(grad(b_sym), grad(v3)) * dx
)
# ---------- External analytical blocks ----------
dfda_ext = make_external_operator("dfda", make_exact_raw_external("a"))
dfdb_ext = make_external_operator("dfdb", make_exact_raw_external("b"))
dfdc_ext = make_external_operator("dfdc", make_exact_raw_external("c"))
F_mu_AB_ext = (
inner(mu_AB, v1) * dx
- inner(dfda_ext, v1) * dx_q
+ inner(dfdb_ext, v1) * dx_q
- (kappa1 - kappa12) * inner(grad(a_sym), grad(v1)) * dx
+ (kappa2 - kappa12) * inner(grad(b_sym), grad(v1)) * dx
)
F_mu_AC_ext = (
inner(mu_AC, v2) * dx
- inner(dfda_ext, v2) * dx_q
+ inner(dfdc_ext, v2) * dx_q
- kappa1 * inner(grad(a_sym), grad(v2)) * dx
- kappa12 * inner(grad(b_sym), grad(v2)) * dx
)
F_mu_BC_ext = (
inner(mu_BC, v3) * dx
- inner(dfdb_ext, v3) * dx_q
+ inner(dfdc_ext, v3) * dx_q
- kappa12 * inner(grad(a_sym), grad(v3)) * dx
- kappa2 * inner(grad(b_sym), grad(v3)) * dx
)
else:
# ---------- UFL blocks ----------
F_mu_AB_ufl = (
inner(mu_AB, v1) * dx
- inner(gAB_expr, v1) * dx
- (kappa1 - kappa12) * inner(grad(a_sym), grad(v1)) * dx
+ (kappa2 - kappa12) * inner(grad(b_sym), grad(v1)) * dx
)
F_mu_AC_ufl = (
inner(mu_AC, v2) * dx
- inner(gAC_expr, v2) * dx
- kappa1 * inner(grad(a_sym), grad(v2)) * dx
- kappa12 * inner(grad(b_sym), grad(v2)) * dx
)
F_mu_BC_ufl = (
inner(mu_BC, v3) * dx
- inner(gBC_expr, v3) * dx
- kappa12 * inner(grad(a_sym), grad(v3)) * dx
- kappa2 * inner(grad(b_sym), grad(v3)) * dx
)
# ---------- External analytical blocks ----------
gAB_ext = make_external_operator("gAB", make_exact_reduced_external("AB"))
gAC_ext = make_external_operator("gAC", make_exact_reduced_external("AC"))
gBC_ext = make_external_operator("gBC", make_exact_reduced_external("BC"))
F_mu_AB_ext = (
inner(mu_AB, v1) * dx
- inner(gAB_ext, v1) * dx_q
- (kappa1 - kappa12) * inner(grad(a_sym), grad(v1)) * dx
+ (kappa2 - kappa12) * inner(grad(b_sym), grad(v1)) * dx
)
F_mu_AC_ext = (
inner(mu_AC, v2) * dx
- inner(gAC_ext, v2) * dx_q
- kappa1 * inner(grad(a_sym), grad(v2)) * dx
- kappa12 * inner(grad(b_sym), grad(v2)) * dx
)
F_mu_BC_ext = (
inner(mu_BC, v3) * dx
- inner(gBC_ext, v3) * dx_q
- kappa12 * inner(grad(a_sym), grad(v3)) * dx
- kappa2 * inner(grad(b_sym), grad(v3)) * dx
)
F_total_ufl = F_a + F_b + F_mu_AB_ufl + F_mu_AC_ufl + F_mu_BC_ufl
F_total_ext = F_a + F_b + F_mu_AB_ext + F_mu_AC_ext + F_mu_BC_ext
# --------------------------------------------------------
# Assemble vectors and compare
# --------------------------------------------------------
reports = {}
reports["F_a"] = vector_diff_norms(
assemble_vector_from_form_expr(F_a),
assemble_external_vector_from_form_expr(F_a),
)
reports["F_b"] = vector_diff_norms(
assemble_vector_from_form_expr(F_b),
assemble_external_vector_from_form_expr(F_b),
)
reports["F_mu_AB"] = vector_diff_norms(
assemble_vector_from_form_expr(F_mu_AB_ufl),
assemble_external_vector_from_form_expr(F_mu_AB_ext),
)
reports["F_mu_AC"] = vector_diff_norms(
assemble_vector_from_form_expr(F_mu_AC_ufl),
assemble_external_vector_from_form_expr(F_mu_AC_ext),
)
reports["F_mu_BC"] = vector_diff_norms(
assemble_vector_from_form_expr(F_mu_BC_ufl),
assemble_external_vector_from_form_expr(F_mu_BC_ext),
)
reports["F_total"] = vector_diff_norms(
assemble_vector_from_form_expr(F_total_ufl),
assemble_external_vector_from_form_expr(F_total_ext),
)
if MPI.COMM_WORLD.rank == 0:
print(f"\nExternal quadrature degree = {external_quadrature_degree}")
print_vector_report("F_a", reports["F_a"])
print_vector_report("F_b", reports["F_b"])
print_vector_report("F_mu_AB", reports["F_mu_AB"])
print_vector_report("F_mu_AC", reports["F_mu_AC"])
print_vector_report("F_mu_BC", reports["F_mu_BC"])
print_vector_report("F_total", reports["F_total"])
# --------------------------------------------------------
# Optional Jacobian comparison
# --------------------------------------------------------
if compare_jacobian:
if MPI.COMM_WORLD.rank == 0:
print("\nComparing Jacobians...")
J_reports = {}
J_reports["J_F_mu_AB"] = matrix_diff_norms(
assemble_matrix_from_form_expr(F_mu_AB_ufl, u, du),
assemble_external_matrix_from_form_expr(F_mu_AB_ext, u, du),
)
J_reports["J_F_mu_AC"] = matrix_diff_norms(
assemble_matrix_from_form_expr(F_mu_AC_ufl, u, du),
assemble_external_matrix_from_form_expr(F_mu_AC_ext, u, du),
)
J_reports["J_F_mu_BC"] = matrix_diff_norms(
assemble_matrix_from_form_expr(F_mu_BC_ufl, u, du),
assemble_external_matrix_from_form_expr(F_mu_BC_ext, u, du),
)
J_reports["J_F_total"] = matrix_diff_norms(
assemble_matrix_from_form_expr(F_total_ufl, u, du),
assemble_external_matrix_from_form_expr(F_total_ext, u, du),
)
if MPI.COMM_WORLD.rank == 0:
print_matrix_report("J_F_mu_AB", J_reports["J_F_mu_AB"])
print_matrix_report("J_F_mu_AC", J_reports["J_F_mu_AC"])
print_matrix_report("J_F_mu_BC", J_reports["J_F_mu_BC"])
print_matrix_report("J_F_total", J_reports["J_F_total"])
if MPI.COMM_WORLD.rank == 0:
print("\nDone.")