Thank you very much dokken!
A naive question in __slots__ does the order of the arguments matter? Because in your answer, the arguments are in the opposite order to the __init__ or ufl_operands.
In addition, there were several issues with the Operator DZ, when trying to compile dolfinx.fem.form(a). I solved that by including a handle for ffcx and a lambda function.
I leave here the complete minimal example
import dolfinx
from mpi4py import MPI
import basix.ufl
import ufl
import dolfinx.fem
import numpy as np
from ufl.constantvalue import as_ufl, Zero
from ufl.core.operator import Operator
from ufl.core.ufl_type import ufl_type
from ufl.algorithms.apply_derivatives import GenericDerivativeRuleset
from ufl.corealg.dag_traverser import DAGTraverser
from ffcx.ir.analysis import reconstruct
from ufl.constantvalue import (
ComplexValue,
ConstantValue,
FloatValue,
IntValue,
RealValue,
Zero,
as_ufl,
is_true_ufl_scalar,
)
from ffcx.codegeneration import lnodes
@ufl_type(num_ops=1, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
class DZ(Operator):
"""Custom operator."""
__slots__ = ("_f", "_k")
def __init__(self, f, k):
"""Create a custom operator."""
self._f = as_ufl(f)
self._k = as_ufl(k)
@property
def ufl_operands(self):
return (self._f, self._k)
@property
def _hash(self)->int:
return hash(("dz", self._f,self._k))
def __hash__(self)-> int:
return self._hash()
def dz(f, k):
"""Create a custom operator."""
k = as_ufl(k)
f = as_ufl(f)
return DZ(f, k)
@DAGTraverser.postorder
def custom_derivative(self, o, df, dk):
f, k, = o.ufl_operands
return k*df
def handle_dz(o, ops):
"""Handle a binary operator."""
if len(ops) != 2:
raise RuntimeError("Expecting two operands.")
return [o._ufl_expr_reconstruct_(a, b) for a, b in zip(ops[0], ops[1])]
# Need to register the derivative, the reconstruction handle for ffcx and the lambda function (lnodes)
lnodes._ufl_call_lookup[DZ] = lambda x, a, b: 0.0 * a
reconstruct._reconstruct_call_lookup[DZ] = handle_dz
GenericDerivativeRuleset.process.register(DZ, custom_derivative)
domain = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
s_el = basix.ufl.element("Lagrange", basix.CellType.interval, 2, shape=(4,))
V = dolfinx.fem.functionspace(domain, s_el)
u = dolfinx.fem.Function(V)
u.interpolate(lambda x: [x[0], x[0], x[1], x[1]])
q = dz(u[0], 2.0)
a = q*ufl.dx(domain=domain)
F = dolfinx.fem.form(a)
f = dolfinx.fem.assemble_scalar(F)
du = ufl.TestFunction(V)
J = ufl.derivative(a, u, du)
J_compiled = dolfinx.fem.form(J)
vec = dolfinx.fem.assemble_vector(J_compiled)
ref = dolfinx.fem.assemble_vector(dolfinx.fem.form(2*du[0]*ufl.dx(domain=domain)))
np.testing.assert_allclose(vec.array, ref.array)
np.testing.assert_allclose(f, 0.0)
Thanks!
The motivation for the definition of a new ufl operator it was to account for the Fourier derivative of the linearization of nonlinear equations that are homogeneous in the z direction. For instance, if we consider a 2D field f_0(x,y) but we want to analyse a linearization of this field
f(x,y,z) \mapsto f_0(x,y) + \epsilon f'(x,y) \exp (i \beta z) for \beta \in \mathbb{R} and \epsilon \ll 1,
it induces a slow derivative \partial_z f \mapsto i \beta f'. However, the evaluation of \partial_z f_0 \mapsto 0.