Ksagiyam/add DAGTraverser #365

ksagiyam · 2025-03-26T12:38:42Z

Attempt to implement what @wence- described in this issue.

We basically would like to introduce classes (DAGTraversers) that define node processing (using singledispatchmethod) and hold caches.

We considered two DAG traversal approaches (for post-order and mostly post-order traversals).

Approach 1. Two-step approach.

Collect nodes first post-order and then process them in order.

This is the current UFL approach (cutoff_unique_post_traversal + MultiFunction wrapped in map_expr_dags).

This "bottom-up" approach is optimal, but when we process a node, we no longer know the local DAG structure around that node (e.g., we do not know its parent). So if we want to pass down some state from the parent, we need to (indirectly) specify that that parent is a "cutoff" node (so that we do not collect the child nodes in the first step), and perform a sub-DAG traversal under that parent; see RestrictionPropagator.restricted(self, o). In the current implementation, if the number of arguments to the method is two (self and o in the above case), the corresponding node type is regarded as a cutoff node type. MultiFunction constructor currently identifies cutoff node types by inspecting the signature of each method. If we wanted to do a similar thing with singledispatchmethod, we could subclass singledispatchmethod as David suggested, but we found that we would end up overwriting singledispatchmethod.register() method relying on the current implementation of singledispatchmethod, which would be a drawback.

Approach 2. Monolithic approach.

Use recursion with caching (each node is processed once and only once) described in this issue.

I observed about a few % - 20% overhead after rewriting apply_derivatives.py (please see below) presumably due to recursion, but no special consideration of cutoff is required. This approach is claimed to be more robust.

In this PR we take Approach2, and replace MultiFunction s with DAGTraversers in apply_derivatives.py. We should be able to incrementally/systematically remove all MultiFunctions in the future.

Performance checks:

holzapfel_ogden.py (holzapfel_ogden):

import time

from ufl import (
    Coefficient,
    Constant,
    FunctionSpace,
    Identity,
    Mesh,
    TestFunction,
    derivative,
    det,
    diff,
    dot,
    dx,
    exp,
    grad,
    inner,
    ln,
    tetrahedron,
    tr,
    variable,
)
from ufl.algorithms import compute_form_data
from ufl.finiteelement import FiniteElement
from ufl.pullback import identity_pullback
from ufl.sobolevspace import H1


mesh = Mesh(FiniteElement("Lagrange", tetrahedron, 1, (3,), identity_pullback, H1))
lamda = Constant(mesh)
a = Constant(mesh)
b = Constant(mesh)
a_s = Constant(mesh)
b_s = Constant(mesh)
a_f = Constant(mesh)
b_f = Constant(mesh)
a_fs = Constant(mesh)
b_fs = Constant(mesh)
e_s = Constant(mesh, shape=(3,))
e_f = Constant(mesh, shape=(3,))

def isochoric(F):
    C = F.T*F
    I_1 = tr(C)
    I4_f = dot(e_f, C*e_f)
    I4_s = dot(e_s, C*e_s)
    I8_fs = dot(e_f, C*e_s)

    def cutoff(x):
        return 1.0/(1.0 + exp(-(x - 1.0)*30.0))

    def scaled_exp(a0, a1, argument):
        return a0/(2.0*a1)*(exp(b*argument) - 1)

    def scaled_exp(a0, a1, argument):
        return a0/(2.0*a1)*(exp(b*argument) - 1)

    E_1 = scaled_exp(a, b, I_1 - 3.)
    E_f = cutoff(I4_f)*scaled_exp(a_f, b_f, (I4_f - 1.)**2)
    E_s = cutoff(I4_s)*scaled_exp(a_s, b_s, (I4_s - 1.)**2)
    E_3 = scaled_exp(a_fs, b_fs, I8_fs**2)
    E = E_1 + E_f + E_s + E_3
    return E

elem = FiniteElement("Lagrange", tetrahedron, 1, (3,), identity_pullback, H1)
V = FunctionSpace(mesh, elem)
u = Coefficient(V)
v = TestFunction(V)
I = Identity(mesh.ufl_cell().topological_dimension())
F = grad(u) + I
F = variable(F)
J = det(F)
Fbar = J**(-1.0/3.0)*F
E_volumetric = lamda*0.5*ln(J)**2
psi = isochoric(Fbar) + E_volumetric
P = diff(psi, F)
F = inner(P, grad(v))*dx
a = derivative(F, u)

ntest = 2
time_sum = 0
for _ in range(ntest):
    start = time.time()
    fd = compute_form_data(
        a,
        do_apply_function_pullbacks=True,
        do_apply_default_restrictions=True,
        do_apply_geometry_lowering=True,
        do_apply_restrictions=True,
        complex_mode=False,
    )
    end = time.time()
    time_sum += end - start
print("average time required: ", time_sum / ntest)

main:
average time required: 0.27037227153778076

this PR:
average time required: 0.27785348892211914 (+2.8%)

With #339 (<- performance regression) reverted:

main:
average time required: 0.04610729217529297

this PR:
average time required: 0.055258870124816895 (+20%)

wence_test (#69):

import time

from ufl import (
    Coefficient,
    Constant,
    FacetNormal,
    FunctionSpace,
    Identity,
    Mesh,
    avg,
    derivative,
    det,
    diff,
    dS,
    grad,
    hexahedron,
    inner,
    ln,
    outer,
    variable,
)
from ufl.algorithms import compute_form_data
from ufl.finiteelement import FiniteElement
from ufl.pullback import identity_pullback, contravariant_piola
from ufl.sobolevspace import H1, HDiv

cell = hexahedron
mesh = Mesh(FiniteElement("Q", cell, 1, (3,), identity_pullback, H1))
V = FunctionSpace(mesh, FiniteElement("NCF", cell, 2, (3,), contravariant_piola, HDiv))
u = Coefficient(V)
mu = Constant(mesh)
psi = lambda F: (mu/2) * (inner(F, F) - u.ufl_shape[0]) - mu*ln(det(F))
flux = lambda F: diff(psi(F), F)
n = FacetNormal(mesh)
uxn = outer(u('+'), n('+')) + outer(u('-'), n('-'))
eye = Identity(u.ufl_shape[0])
grad_u = grad(u)
F_ = variable(grad_u + eye)
Fm = variable(grad_u('-') + eye)
Fp = variable(grad_u('+') + eye)
avg_flux = avg(flux(F_))
U = inner(avg_flux, uxn)*dS
a = derivative(derivative(U, u), u)

ntest = 2
time_sum = 0
for _ in range(ntest):
    start = time.time()
    fd = compute_form_data(
        a,
        do_apply_function_pullbacks=True,
        do_apply_default_restrictions=True,
        do_apply_geometry_lowering=True,
        do_apply_restrictions=True,
        complex_mode=False,
    )
    end = time.time()
    time_sum += end - start
print("average time required: ", time_sum / ntest)

main:
average time required: 0.23360061645507812

this PR:
average time required: 0.2485201358795166 (+6.4%)

With #339 reverted:

main:
average time required: 0.060405850410461426

this PR:
average time required: 0.06431174278259277 (+6.5%)

Firedrake CI
firedrakeproject/firedrake#4145 (no notable performance regression).

connorjward

I think this is pretty great. I have written a lot of tree traversals and the DAGTraverser.postorder (and friends) solves an issue I've been thinking about.

It's almost a shame that this only lives in UFL.

connorjward · 2025-04-03T08:23:05Z

ufl/corealg/dag_traverser.py

+
+
+class DAGTraverser(ABC):
+    """Base class for dag traversers."""


You should document the __init__ parameters here

ufl/corealg/dag_traverser.py

connorjward · 2025-04-03T08:28:28Z

ufl/corealg/dag_traverser.py

+
+    @staticmethod
+    def postorder(method):
+        """Suppress processed operands in arguments.


Could this be clearer?

I found it hard to keep it short. Please let me know if you have a suggestion. We at least have an expanded explanation right below.

"Decorator indicating that child nodes are visited first in the traversal."?

I'm afraid I have no idea what the current docstring means.

It would also be fine to have something trivial like "Postorder decorator." provided that there is an explanation in the rest of the docstring.

Changed to "Postorder decorator".

connorjward · 2025-04-03T08:30:17Z

ufl/algorithms/apply_derivatives.py

+            Processed object.
+
+        """
+        raise AssertionError(f"UFL expression expected: got {o}")


Is this now boilerplate that needs reproducing for each subclass? Might it be better on the ABC?

Do you mean putting this base case in the base class? I think the problem is that singledispatchmethod object needs to be made in each child class. What I had to do here is probably relevant https://github.com/firedrakeproject/ufl/blob/c16f038378e17ad8bd730bf9a020028ad72a00c1/ufl/algorithms/apply_derivatives.py#L716.

connorjward · 2025-04-03T08:30:57Z

ufl/algorithms/apply_derivatives.py


-    def component_tensor(self, o, Ap, ii):
+    @process.register(ComponentTensor)
+    @DAGTraverser.postorder


This is pretty sweet

ufl/algorithms/apply_derivatives.py

ksagiyam · 2025-04-23T08:48:54Z

@jorgensd @mscroggs Could you have a look at this PR?

ufl/corealg/dag_traverser.py

jorgensd · 2025-04-23T09:12:52Z

ufl/algorithms/apply_derivatives.py

+        )
+
+    @process.register(Derivative)
+    def _(self, o: Expr) -> Expr:


Wouldn't it be better to use

@process.register def _(self, o: Derivative) -> Expr:

here and in subsequent registrations, as https://docs.python.org/3/library/functools.html#functools.singledispatch states:

To add overloaded implementations to the function, use the register() attribute of the generic function, which can be used as a decorator. For functions annotated with types, the decorator will infer the type of the first argument automatically:

I guess we would have to wait for lower version Python to be 3.11 to do this for unions (python/cpython#30017). However, personally, I think this would be a cleaner solution. What do you think?

I gave it a go with this as a separate commit, but I personally prefer:

@process.register(Grad) @process.register(CellAvg) @process.register(FacetAvg) def _(self, o) -> Expr: ...

as type names are left-aligned and easier to read for me. I also do not quite see how it plays well when decorators are nested as in:

@process.register @DAGTraverser.postorder @pending_operations_recording def external_operator(self, N: ExternalOperator, *dfs) -> Expr: ...

though functools.wraps() seems to be doing the job for us.

I revert this commit for now as Python3.10 test fails as you already pointed out. Please see the commit history for the relevant changes.

jorgensd · 2025-04-23T09:31:14Z

ufl/corealg/dag_traverser.py

+
+    """
+
+    def __init__(self, compress=True, visited_cache=None, result_cache=None) -> None:


Suggested change

def __init__(self, compress=True, visited_cache=None, result_cache=None) -> None:

def __init__(self, compress: bool=True, visited_cache=None, result_cache=None) -> None:

Currently I do not think there is any test in ufl that has a *args as a non-empty argument. i.e. is always equal to args=(). Could you add a test that has args being non-empty?

Should the type-hint for visited_cache be dict[tuple[Expr, ...], Expr]|None, and similarly for result_cache?

Added typehint.

That is right. It seems that we only need args = () in apply_derivatives.py. We will need non-empty args for instance when we want to pass down the state in apply_restrictions.py; underneath a Restricted node restricted='+' or restricted='-', otherwise restricted=None, where the method signature would look like:

@process.register(...) def _(self, o, restricted): ...

We could remove *args for this PR.

jorgensd · 2025-04-23T09:37:23Z

ufl/algorithms/apply_derivatives.py

-    def coordinate_derivative(self, o, f, w, v, cd):
+    @process.register(CoordinateDerivative)
+    @DAGTraverser.postorder_only_children([0])
+    def _(self, o: Expr, f: Expr) -> Expr:


I guess a test on coordinate-derivative should yield an *args that is non-empty.

CoordinateDerivative still only needs args=(), I think.

Yeah, I only looked at the signatures of the process method, but for the traversal of the graph, it seems like we never use args.

jorgensd · 2025-04-23T09:49:14Z

ufl/algorithms/apply_derivatives.py

@@ -582,14 +710,35 @@ def min_value(self, o, df, dg):
 class GradRuleset(GenericDerivativeRuleset):
    """Take the grad derivative."""

-    def __init__(self, geometric_dimension):
+    def __init__(self, geometric_dimension, compress=True, visited_cache=None, result_cache=None):


I would add type-hints here as well.

jorgensd · 2025-04-23T09:50:08Z

ufl/algorithms/apply_derivatives.py

@@ -1026,29 +1239,35 @@ class GateauxDerivativeRuleset(GenericDerivativeRuleset):
    D_w[v](e) = d/dtau e(w+tau v)|tau=0.
    """

-    def __init__(self, coefficients, arguments, coefficient_derivatives):
+    def __init__(


Type hints :)

jorgensd · 2025-04-23T09:58:02Z

ufl/algorithms/apply_derivatives.py

+    def _(self, o: Expr) -> Expr:
+        return self._process_coefficient(o)
+
+    def _process_coefficient(self, o: Expr) -> Expr:


Is the doc-string here strictly correct? (I know it existed prior to the PR).
As _process_coefficient supports taking a BaseFormOperator or a Cofunction it is not really a Coefficient.

I guess this goes back to the GateauxDerivativeRuleset where any ExprList can be taken as input and therefore might be out of scope of this PR.

I partially fixed this (as I added these new methods), but I think we should do more comprehensive fix in a future PR.

jorgensd · 2025-04-23T10:01:28Z

In general very positive to this change. Will run the tests later today to check for potential regressions.

jorgensd · 2025-04-23T10:47:26Z

Locally, i did not see a big regression with this PR, however, DOLFINx doesn't have the most complicated forms in the tests.
Maybe @schnellerhase , @jhale or @michalhabera has some more complicated forms in DOLFINy that we could check against.

schnellerhase · 2025-04-23T11:17:31Z

Thanks for the ping. Quick testing on https://github.com/fenics-dolfiny/dolfiny/blob/main/demo/spectral/solid_elasticity_spectral.py - modified test case attached we see the following regression:

main 
F: 0.22513866424560547s
J: 0.5185086727142334s

https://github.com/firedrakeproject/ufl.git@449b44e67efc0bda6f5d4cd99c92dbd22dafd146
F: 0.7992544174194336s
J: 4.553445982933044s

Modified example

#!/usr/bin/env python3

import time

from mpi4py import MPI
from petsc4py import PETSc

import basix
import dolfinx
import ufl
from dolfinx import default_scalar_type as scalar

import mesh_tube3d_gmshapi as mg
import numpy as np
import plot_tube3d_pyvista as pl

import dolfiny

# Basic settings
name = "solid_elasticity_spectral"
comm = MPI.COMM_WORLD

# Geometry and mesh parameters
r, t, h = 0.04, 0.01, 0.1  # [m]
nr, nt, nh = 16, 5, 8

# Create the regular mesh of a tube with given dimensions
gmsh_model, tdim = mg.mesh_tube3d_gmshapi(name, r, t, h, nr, nt, nh, do_quads=True, order=2)

# Get mesh and meshtags
mesh, mts = dolfiny.mesh.gmsh_to_dolfin(gmsh_model, tdim)

# Get merged MeshTags for each codimension
subdomains, subdomains_keys = dolfiny.mesh.merge_meshtags(mesh, mts, tdim - 0)
interfaces, interfaces_keys = dolfiny.mesh.merge_meshtags(mesh, mts, tdim - 1)

# Define shorthands for labelled tags
surface_lower = interfaces_keys["surface_lower"]
surface_upper = interfaces_keys["surface_upper"]

# Define integration measures
dx = ufl.Measure("dx", domain=mesh, subdomain_data=subdomains, metadata={"quadrature_degree": 4})
ds = ufl.Measure("ds", domain=mesh, subdomain_data=interfaces, metadata={"quadrature_degree": 4})

# Define elements
Ue = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,))

# Define function spaces
Uf = dolfinx.fem.functionspace(mesh, Ue)

# Define functions
u = dolfinx.fem.Function(Uf, name="u")

u_ = dolfinx.fem.Function(Uf, name="u_")  # boundary conditions

δu = ufl.TestFunction(Uf)

# Define state as (ordered) list of functions
m, δm = [u], [δu]

# output / visualisation
vorder = mesh.geometry.cmap.degree
uo = dolfinx.fem.Function(dolfinx.fem.functionspace(mesh, ("P", vorder, (3,))), name="u")
so = dolfinx.fem.Function(dolfinx.fem.functionspace(mesh, ("P", vorder)), name="s")  # for output

# Kinematics
F = ufl.Identity(3) + ufl.grad(u)

# Strain measure: from Cauchy strain tensor to squares of principal stretches
c, _ = dolfiny.invariants.eigenstate(F.T * F)  # spectral decomposition of C
c = ufl.as_vector(c)  # squares of principal stretches
c = ufl.variable(c)
# Variation of strain measure (squares of principal stretches)
δc = dolfiny.expression.derivative(c, m, δm)

# Elasticity parameters
E = dolfinx.fem.Constant(mesh, scalar(1.0))  # [MPa]
nu = dolfinx.fem.Constant(mesh, scalar(0.4))  # [-]
mu = E / (2 * (1 + nu))
la = E * nu / ((1 + nu) * (1 - 2 * nu))

# Define boundary stress vector (torque at upper face)
x0 = ufl.SpatialCoordinate(mesh)
n0 = ufl.FacetNormal(mesh)
λ = dolfinx.fem.Constant(mesh, scalar(0.0))
t = ufl.cross(x0 - ufl.as_vector([0.0, 0.0, h]), n0) * 4 * λ  # [N/m^2]


def strain_energy(i1, i2, i3):
    """Strain energy function
    i1, i2, i3: principal invariants of the Cauchy-Green tensor
    """
    # Determinant of configuration gradient F
    J = ufl.sqrt(i3)
    #
    # Classical St. Venant-Kirchhoff
    # Ψ = la / 8 * (i1 - 3) ** 2 + mu / 4 * ((i1 - 3) ** 2 + 4 * (i1 - 3) - 2 * (i2 - 3))
    # Modified St. Venant-Kirchhoff
    # Ψ = la / 2 * (ufl.ln(J)) ** 2 + mu / 4 * ((i1 - 3) ** 2 + 4 * (i1 - 3) - 2 * (i2 - 3))
    # Compressible neo-Hooke
    Ψ = mu / 2 * (i1 - 3 - 2 * ufl.ln(J)) + la / 2 * (J - 1) ** 2
    # Compressible Mooney-Rivlin (beta = 0)
    # Ψ = mu / 4 * (i1 - 3) + mu / 4 * (i2 - 3) - mu * ufl.ln(J) + la / 2 * (J - 1) ** 2
    #
    return Ψ


# Invariants (based on spectral decomposition of C)
i1, i2, i3 = c[0] + c[1] + c[2], c[0] * c[1] + c[1] * c[2] + c[0] * c[2], c[0] * c[1] * c[2]
# Material model (isotropic)
Ψ = strain_energy(i1, i2, i3)
# Stress measure
s = 2 * ufl.diff(Ψ, c)

# von Mises stress (output)
svm = ufl.sqrt(3 / 2 * ufl.inner(ufl.dev(ufl.diag(s)), ufl.dev(ufl.diag(s))))

# Weak form: for isotropic material, eigenprojectors of C and S are identical
form = -0.5 * ufl.inner(δc, s) * dx + ufl.inner(δu, t) * ds(surface_upper)

# Overall form (as list of forms)
forms = dolfiny.function.extract_blocks(form, δm)

form = forms[0]
n = 10
start = time.time()
for _ in range(n):
    ufl.algorithms.compute_form_data(
        form,
        do_apply_function_pullbacks=True,
        do_apply_default_restrictions=True,
        do_apply_geometry_lowering=True,
        do_apply_restrictions=True,
        complex_mode=False,
    )

print(f"F: {(time.time() - start)/n}s")


J_form = ufl.derivative(forms[0], m[0], ufl.TrialFunction(m[0].function_space))

start = time.time()
for _ in range(n):
    ufl.algorithms.compute_form_data(
        J_form,
        do_apply_function_pullbacks=True,
        do_apply_default_restrictions=True,
        do_apply_geometry_lowering=True,
        do_apply_restrictions=True,
        complex_mode=False,
    )

print(f"J: {(time.time() - start)/n}s")

jorgensd · 2025-04-23T12:37:02Z

@schnellerhase your example is quite complex.
I've here narrowed it down to a "pure" UFL example (taken bits and pieces from dolfiny)

import ufl
import time

from ufl import (
    Coefficient,
    FunctionSpace,
    Mesh,
    TestFunction,
    hexahedron,
    Constant,
)
from ufl.finiteelement import FiniteElement
from ufl.pullback import identity_pullback
from ufl.sobolevspace import H1


def invariants_principal(A):
    """Principal invariants of (real-valued) tensor A.
    https://doi.org/10.1007/978-3-7091-0174-2_3
    """
    i1 = ufl.tr(A)
    i2 = (ufl.tr(A) ** 2 - ufl.tr(A * A)) / 2
    i3 = ufl.det(A)
    return i1, i2, i3


def eigenstate3(A):
    """Eigenvalues and eigenprojectors of the 3x3 (real-valued) tensor A.
    Provides the spectral decomposition A = sum_{a=0}^{2} λ_a * E_a
    with (ordered) eigenvalues λ_a and their associated eigenprojectors E_a = n_a^R x n_a^L.

    Note: Tensor A must not have complex eigenvalues!
    """
    if ufl.shape(A) != (3, 3):
        raise RuntimeError(
            f"Tensor A of shape {ufl.shape(A)} != (3, 3) is not supported!"
        )
    #
    eps = 3.0e-16  # slightly above 2**-(53 - 1), see https://en.wikipedia.org/wiki/IEEE_754
    #
    A = ufl.variable(A)
    #
    # --- determine eigenvalues λ0, λ1, λ2

    # Invariants
    I1, _, _ = invariants_principal(A)

    # Discriminant as sum-of-products
    Δx = [
        A[0, 1] * A[1, 2] * A[2, 0] - A[0, 2] * A[1, 0] * A[2, 1],
        A[0, 1] ** 2 * A[1, 2]
        - A[0, 1] * A[0, 2] * A[1, 1]
        + A[0, 1] * A[0, 2] * A[2, 2]
        - A[0, 2] ** 2 * A[2, 1],
        A[0, 0] * A[0, 1] * A[2, 1]
        - A[0, 1] ** 2 * A[2, 0]
        - A[0, 1] * A[2, 1] * A[2, 2]
        + A[0, 2] * A[2, 1] ** 2,
        A[0, 0] * A[0, 2] * A[1, 2]
        + A[0, 1] * A[1, 2] ** 2
        - A[0, 2] ** 2 * A[1, 0]
        - A[0, 2] * A[1, 1] * A[1, 2],
        A[0, 0] * A[0, 1] * A[1, 2]
        - A[0, 1] * A[0, 2] * A[1, 0]
        - A[0, 1] * A[1, 2] * A[2, 2]
        + A[0, 2] * A[1, 2] * A[2, 1],
        A[0, 0] * A[0, 2] * A[2, 1]
        - A[0, 1] * A[0, 2] * A[2, 0]
        + A[0, 1] * A[1, 2] * A[2, 1]
        - A[0, 2] * A[1, 1] * A[2, 1],
        A[0, 1] * A[1, 0] * A[1, 2]
        - A[0, 2] * A[1, 0] * A[1, 1]
        + A[0, 2] * A[1, 0] * A[2, 2]
        - A[0, 2] * A[1, 2] * A[2, 0],
        A[0, 0] ** 2 * A[1, 2]
        - A[0, 0] * A[0, 2] * A[1, 0]
        - A[0, 0] * A[1, 1] * A[1, 2]
        - A[0, 0] * A[1, 2] * A[2, 2]
        + A[0, 1] * A[1, 0] * A[1, 2]
        + A[0, 2] * A[1, 0] * A[2, 2]
        + A[1, 1] * A[1, 2] * A[2, 2]
        - A[1, 2] ** 2 * A[2, 1],
        A[0, 0] ** 2 * A[1, 2]
        - A[0, 0] * A[0, 2] * A[1, 0]
        - A[0, 0] * A[1, 1] * A[1, 2]
        - A[0, 0] * A[1, 2] * A[2, 2]
        + A[0, 2] * A[1, 0] * A[1, 1]
        + A[0, 2] * A[1, 2] * A[2, 0]
        + A[1, 1] * A[1, 2] * A[2, 2]
        - A[1, 2] ** 2 * A[2, 1],
        A[0, 0] * A[0, 1] * A[1, 1]
        - A[0, 0] * A[0, 1] * A[2, 2]
        - A[0, 1] ** 2 * A[1, 0]
        + A[0, 1] * A[0, 2] * A[2, 0]
        - A[0, 1] * A[1, 1] * A[2, 2]
        + A[0, 1] * A[2, 2] ** 2
        + A[0, 2] * A[1, 1] * A[2, 1]
        - A[0, 2] * A[2, 1] * A[2, 2],
        A[0, 0] * A[0, 1] * A[1, 1]
        - A[0, 0] * A[0, 1] * A[2, 2]
        + A[0, 0] * A[0, 2] * A[2, 1]
        - A[0, 1] ** 2 * A[1, 0]
        - A[0, 1] * A[1, 1] * A[2, 2]
        + A[0, 1] * A[1, 2] * A[2, 1]
        + A[0, 1] * A[2, 2] ** 2
        - A[0, 2] * A[2, 1] * A[2, 2],
        A[0, 0] * A[0, 1] * A[1, 2]
        - A[0, 0] * A[0, 2] * A[1, 1]
        + A[0, 0] * A[0, 2] * A[2, 2]
        - A[0, 1] * A[1, 1] * A[1, 2]
        - A[0, 2] ** 2 * A[2, 0]
        + A[0, 2] * A[1, 1] ** 2
        - A[0, 2] * A[1, 1] * A[2, 2]
        + A[0, 2] * A[1, 2] * A[2, 1],
        A[0, 0] * A[0, 2] * A[1, 1]
        - A[0, 0] * A[0, 2] * A[2, 2]
        - A[0, 1] * A[0, 2] * A[1, 0]
        + A[0, 1] * A[1, 1] * A[1, 2]
        - A[0, 1] * A[1, 2] * A[2, 2]
        + A[0, 2] ** 2 * A[2, 0]
        - A[0, 2] * A[1, 1] ** 2
        + A[0, 2] * A[1, 1] * A[2, 2],
        A[0, 0] ** 2 * A[1, 1]
        - A[0, 0] ** 2 * A[2, 2]
        - A[0, 0] * A[0, 1] * A[1, 0]
        + A[0, 0] * A[0, 2] * A[2, 0]
        - A[0, 0] * A[1, 1] ** 2
        + A[0, 0] * A[2, 2] ** 2
        + A[0, 1] * A[1, 0] * A[1, 1]
        - A[0, 2] * A[2, 0] * A[2, 2]
        + A[1, 1] ** 2 * A[2, 2]
        - A[1, 1] * A[1, 2] * A[2, 1]
        - A[1, 1] * A[2, 2] ** 2
        + A[1, 2] * A[2, 1] * A[2, 2],
    ]
    Δy = [
        A[0, 2] * A[1, 0] * A[2, 1] - A[0, 1] * A[1, 2] * A[2, 0],
        A[1, 0] ** 2 * A[2, 1]
        - A[1, 0] * A[1, 1] * A[2, 0]
        + A[1, 0] * A[2, 0] * A[2, 2]
        - A[1, 2] * A[2, 0] ** 2,
        A[0, 0] * A[1, 0] * A[1, 2]
        - A[0, 2] * A[1, 0] ** 2
        - A[1, 0] * A[1, 2] * A[2, 2]
        + A[1, 2] ** 2 * A[2, 0],
        A[0, 0] * A[2, 0] * A[2, 1]
        - A[0, 1] * A[2, 0] ** 2
        + A[1, 0] * A[2, 1] ** 2
        - A[1, 1] * A[2, 0] * A[2, 1],
        A[0, 0] * A[1, 0] * A[2, 1]
        - A[0, 1] * A[1, 0] * A[2, 0]
        - A[1, 0] * A[2, 1] * A[2, 2]
        + A[1, 2] * A[2, 0] * A[2, 1],
        A[0, 0] * A[1, 2] * A[2, 0]
        - A[0, 2] * A[1, 0] * A[2, 0]
        + A[1, 0] * A[1, 2] * A[2, 1]
        - A[1, 1] * A[1, 2] * A[2, 0],
        A[0, 1] * A[1, 0] * A[2, 1]
        - A[0, 1] * A[1, 1] * A[2, 0]
        + A[0, 1] * A[2, 0] * A[2, 2]
        - A[0, 2] * A[2, 0] * A[2, 1],
        A[0, 0] ** 2 * A[2, 1]
        - A[0, 0] * A[0, 1] * A[2, 0]
        - A[0, 0] * A[1, 1] * A[2, 1]
        - A[0, 0] * A[2, 1] * A[2, 2]
        + A[0, 1] * A[1, 0] * A[2, 1]
        + A[0, 1] * A[2, 0] * A[2, 2]
        + A[1, 1] * A[2, 1] * A[2, 2]
        - A[1, 2] * A[2, 1] ** 2,
        A[0, 0] ** 2 * A[2, 1]
        - A[0, 0] * A[0, 1] * A[2, 0]
        - A[0, 0] * A[1, 1] * A[2, 1]
        - A[0, 0] * A[2, 1] * A[2, 2]
        + A[0, 1] * A[1, 1] * A[2, 0]
        + A[0, 2] * A[2, 0] * A[2, 1]
        + A[1, 1] * A[2, 1] * A[2, 2]
        - A[1, 2] * A[2, 1] ** 2,
        A[0, 0] * A[1, 0] * A[1, 1]
        - A[0, 0] * A[1, 0] * A[2, 2]
        - A[0, 1] * A[1, 0] ** 2
        + A[0, 2] * A[1, 0] * A[2, 0]
        - A[1, 0] * A[1, 1] * A[2, 2]
        + A[1, 0] * A[2, 2] ** 2
        + A[1, 1] * A[1, 2] * A[2, 0]
        - A[1, 2] * A[2, 0] * A[2, 2],
        A[0, 0] * A[1, 0] * A[1, 1]
        - A[0, 0] * A[1, 0] * A[2, 2]
        + A[0, 0] * A[1, 2] * A[2, 0]
        - A[0, 1] * A[1, 0] ** 2
        - A[1, 0] * A[1, 1] * A[2, 2]
        + A[1, 0] * A[1, 2] * A[2, 1]
        + A[1, 0] * A[2, 2] ** 2
        - A[1, 2] * A[2, 0] * A[2, 2],
        A[0, 0] * A[1, 0] * A[2, 1]
        - A[0, 0] * A[1, 1] * A[2, 0]
        + A[0, 0] * A[2, 0] * A[2, 2]
        - A[0, 2] * A[2, 0] ** 2
        - A[1, 0] * A[1, 1] * A[2, 1]
        + A[1, 1] ** 2 * A[2, 0]
        - A[1, 1] * A[2, 0] * A[2, 2]
        + A[1, 2] * A[2, 0] * A[2, 1],
        A[0, 0] * A[1, 1] * A[2, 0]
        - A[0, 0] * A[2, 0] * A[2, 2]
        - A[0, 1] * A[1, 0] * A[2, 0]
        + A[0, 2] * A[2, 0] ** 2
        + A[1, 0] * A[1, 1] * A[2, 1]
        - A[1, 0] * A[2, 1] * A[2, 2]
        - A[1, 1] ** 2 * A[2, 0]
        + A[1, 1] * A[2, 0] * A[2, 2],
        A[0, 0] ** 2 * A[1, 1]
        - A[0, 0] ** 2 * A[2, 2]
        - A[0, 0] * A[0, 1] * A[1, 0]
        + A[0, 0] * A[0, 2] * A[2, 0]
        - A[0, 0] * A[1, 1] ** 2
        + A[0, 0] * A[2, 2] ** 2
        + A[0, 1] * A[1, 0] * A[1, 1]
        - A[0, 2] * A[2, 0] * A[2, 2]
        + A[1, 1] ** 2 * A[2, 2]
        - A[1, 1] * A[1, 2] * A[2, 1]
        - A[1, 1] * A[2, 2] ** 2
        + A[1, 2] * A[2, 1] * A[2, 2],
    ]
    Δd = [9, 6, 6, 6, 8, 8, 8, 2, 2, 2, 2, 2, 2, 1]
    Δ = sum(Δxk * Δdk * Δyk for Δxk, Δdk, Δyk in zip(Δx, Δd, Δy))  # discriminant as sop

    # Invariant dp as sum-of-products
    Δxp = [
        A[1, 0],
        A[2, 0],
        A[2, 1],
        -A[0, 0] + A[1, 1],
        -A[0, 0] + A[2, 2],
        -A[1, 1] + A[2, 2],
    ]
    Δyp = [
        A[0, 1],
        A[0, 2],
        A[1, 2],
        -A[0, 0] + A[1, 1],
        -A[0, 0] + A[2, 2],
        -A[1, 1] + A[2, 2],
    ]
    Δdp = [6, 6, 6, 1, 1, 1]
    dp = (
        sum(Δxpk * Δdpk * Δypk for Δxpk, Δdpk, Δypk in zip(Δxp, Δdp, Δyp)) / 2
    )  # dp as sop

    # Invariant dq as sum-of-products
    Δxq = [
        A[1, 2],
        A[2, 1],
        A[0, 1] * A[1, 0],
        A[0, 2] * A[2, 0],
        A[1, 2] * A[2, 1],
        A[1, 1] + A[2, 2] - 2 * A[0, 0],
    ]
    Δyq = [
        A[0, 1] * A[2, 0],
        A[0, 2] * A[1, 0],
        A[0, 0] + A[1, 1] - 2 * A[2, 2],
        A[0, 0] + A[2, 2] - 2 * A[1, 1],
        A[1, 1] + A[2, 2] - 2 * A[0, 0],
        (A[0, 0] + A[2, 2] - 2 * A[1, 1]) * (A[0, 0] + A[1, 1] - 2 * A[2, 2]),
    ]
    Δdq = [27, 27, 9, 9, 9, -1]
    dq = sum(Δxqk * Δdqk * Δyqk for Δxqk, Δdqk, Δyqk in zip(Δxq, Δdq, Δyq))

    # Avoid dp = 0 and disc = 0, both are known with absolute error of ~eps**2
    # Required to avoid sqrt(0) derivatives and negative square roots
    dp += eps**2
    Δ += eps**2

    phi3 = ufl.atan2(ufl.sqrt(27) * ufl.sqrt(Δ), dq)

    # sorted eigenvalues: λ0 <= λ1 <= λ2
    λ = [
        (I1 + 2 * ufl.sqrt(dp) * ufl.cos((phi3 + 2 * ufl.pi * k) / 3)) / 3
        for k in range(1, 4)
    ]
    #
    # --- determine eigenprojectors E0, E1, E2
    #
    E = [ufl.diff(λk, A).T for λk in λ]

    return λ, E


def eigenstate2(A):
    """Eigenvalues and eigenprojectors of the 2x2 (real-valued) tensor A.
    Provides the spectral decomposition A = sum_{a=0}^{1} λ_a * E_a
    with (ordered) eigenvalues λ_a and their associated eigenprojectors E_a = n_a^R x n_a^L.

    Note: Tensor A must not have complex eigenvalues!
    """
    if ufl.shape(A) != (2, 2):
        raise RuntimeError(
            f"Tensor A of shape {ufl.shape(A)} != (2, 2) is not supported!"
        )
    #
    eps = 3.0e-16  # slightly above 2**-(53 - 1), see https://en.wikipedia.org/wiki/IEEE_754
    #
    A = ufl.variable(A)
    #
    # --- determine eigenvalues λ0, λ1
    #
    I1, _, _ = invariants_principal(A)
    #
    Δ = (A[0, 0] - A[1, 1]) ** 2 + 4 * A[0, 1] * A[1, 0]  # = I1**2 - 4 * I2
    # Avoid dp = 0 and disc = 0, both are known with absolute error of ~eps**2
    # Required to avoid sqrt(0) derivatives and negative square roots
    Δ += eps**2
    # sorted eigenvalues: λ0 <= λ1
    λ = (I1 - ufl.sqrt(Δ)) / 2, (I1 + ufl.sqrt(Δ)) / 2
    #
    # --- determine eigenprojectors E0, E1
    #
    E = [ufl.diff(λk, A).T for λk in λ]

    return λ, E


def eigenstate(A):
    """Eigenvalues and eigenprojectors of the (real-valued) tensor A of dimension m = 2, 3.
    Provides the spectral decomposition A = sum_{a=0}^{m} λ_a * E_a
    with (ordered) eigenvalues λ_a and their associated eigenprojectors E_a = n_a^R x n_a^L.

    Note: Tensor A must not have complex eigenvalues!
    """
    if ufl.shape(A) == (3, 3):
        return eigenstate3(A)
    elif ufl.shape(A) == (2, 2):
        return eigenstate2(A)
    else:
        raise RuntimeError(f"Tensor A of shape {ufl.shape(A)} is not supported!")


element = FiniteElement("Lagrange", hexahedron, 2, (3,), identity_pullback, H1)
domain = Mesh(FiniteElement("Lagrange", hexahedron, 2, (3,), identity_pullback, H1))


# Define integration measures
dx = ufl.Measure("dx", domain=domain, metadata={"quadrature_degree": 4})
ds = ufl.Measure("ds", domain=domain, metadata={"quadrature_degree": 4})


Uf = FunctionSpace(domain, element)
u = Coefficient(Uf)  # displacement
u_ = Coefficient(Uf)

δu = TestFunction(Uf)


# output / visualisation

# Kinematics
F = ufl.Identity(3) + ufl.grad(u)

# Strain measure: from Cauchy strain tensor to squares of principal stretches
c, _ = eigenstate(F.T * F)  # spectral decomposition of C
c = ufl.as_vector(c)  # squares of principal stretches
c = ufl.variable(c)
# Variation of strain measure (squares of principal stretches)
δc = ufl.derivative(c, u, δu)

# Elasticity parameters
E = Constant(domain)  # [MPa]
nu = Constant(domain)
mu = E / (2 * (1 + nu))
la = E * nu / ((1 + nu) * (1 - 2 * nu))

# Define boundary stress vector (torque at upper face)
x0 = ufl.SpatialCoordinate(domain)
n0 = ufl.FacetNormal(domain)
λ = Constant(domain)
h = Constant(domain)  # [m]
t = ufl.cross(x0 - ufl.as_vector([0.0, 0.0, h]), n0) * 4 * λ  # [N/m^2]


def strain_energy(i1, i2, i3):
    """Strain energy function
    i1, i2, i3: principal invariants of the Cauchy-Green tensor
    """
    # Determinant of configuration gradient F
    J = ufl.sqrt(i3)
    #
    # Classical St. Venant-Kirchhoff
    # Ψ = la / 8 * (i1 - 3) ** 2 + mu / 4 * ((i1 - 3) ** 2 + 4 * (i1 - 3) - 2 * (i2 - 3))
    # Modified St. Venant-Kirchhoff
    # Ψ = la / 2 * (ufl.ln(J)) ** 2 + mu / 4 * ((i1 - 3) ** 2 + 4 * (i1 - 3) - 2 * (i2 - 3))
    # Compressible neo-Hooke
    Ψ = mu / 2 * (i1 - 3 - 2 * ufl.ln(J)) + la / 2 * (J - 1) ** 2
    # Compressible Mooney-Rivlin (beta = 0)
    # Ψ = mu / 4 * (i1 - 3) + mu / 4 * (i2 - 3) - mu * ufl.ln(J) + la / 2 * (J - 1) ** 2
    #
    return Ψ


# Invariants (based on spectral decomposition of C)
i1, i2, i3 = (
    c[0] + c[1] + c[2],
    c[0] * c[1] + c[1] * c[2] + c[0] * c[2],
    c[0] * c[1] * c[2],
)
# Material model (isotropic)
Ψ = strain_energy(i1, i2, i3)
# Stress measure
s = 2 * ufl.diff(Ψ, c)

# von Mises stress (output)
svm = ufl.sqrt(3 / 2 * ufl.inner(ufl.dev(ufl.diag(s)), ufl.dev(ufl.diag(s))))

# Weak form: for isotropic material, eigenprojectors of C and S are identical
form = -0.5 * ufl.inner(δc, s) * dx + ufl.inner(δu, t) * ds(1)


start = time.perf_counter()
ufl.algorithms.compute_form_data(
    form,
    do_apply_function_pullbacks=True,
    do_apply_default_restrictions=True,
    do_apply_geometry_lowering=True,
    do_apply_restrictions=True,
    complex_mode=False,
)
end = time.perf_counter()
print(f"Time taken to compute form data (residual): {(end - start):.2e}")


J = ufl.derivative(form, u)

start = time.perf_counter()
ufl.algorithms.compute_form_data(
    J,
    do_apply_function_pullbacks=True,
    do_apply_default_restrictions=True,
    do_apply_geometry_lowering=True,
    do_apply_restrictions=True,
    complex_mode=False,
)
end = time.perf_counter()
print(f"Time taken to compute form data (jacobian): {(end - start):.2e}")

which results in the following timings on main:

Time taken to compute form data (residual): 1.11e-01
Time taken to compute form data (jacobian): 2.14e-01

and the following on this branch:

Time taken to compute form data (residual): 4.97e-01
Time taken to compute form data (jacobian): 3.65e+00

jorgensd · 2025-04-23T12:57:57Z

There was a small typo in the example above (didn't compile form data for jacobian).
This has now been updated, and the timings have been updated, showing the 10-fold increase in runtime for computing form data for the jacobian.

ksagiyam · 2025-04-23T23:17:06Z

Thanks! I had not updated my branch after some performance fix was merged to main. After rebasing, the numbers are like the following:

main:
Time taken to compute form data (residual): 8.51e-02
Time taken to compute form data (jacobian): 1.62e-01

this PR:
Time taken to compute form data (residual): 1.01e-01 (+19%)
Time taken to compute form data (jacobian): 2.01e-01 (+24%)

Let me also update the results for the other tests in the PR description:

holzapfel_ogden:

main:
0.04760003089904785

this PR:
0.05493581295013428 (+15%)

wence_test:

main:
0.06049168109893799

this PR:
0.06576049327850342 (+8.7%)

- replace MultiFunction with DAGTraverser in apply_derivatives.py Co-authored-by: Connor Ward <[email protected]> Co-authored-by: Jørgen Schartum Dokken <[email protected]>

This reverts commit 25e2a49.

ksagiyam force-pushed the ksagiyam/add_dag_visitor branch 4 times, most recently from bffe57a to b99a902 Compare March 28, 2025 01:32

connorjward suggested changes Apr 3, 2025

View reviewed changes

ksagiyam force-pushed the ksagiyam/add_dag_visitor branch 2 times, most recently from c16f038 to 449b44e Compare April 3, 2025 13:46

ksagiyam marked this pull request as ready for review April 3, 2025 15:23

jorgensd reviewed Apr 23, 2025

View reviewed changes

ufl/corealg/dag_traverser.py Outdated Show resolved Hide resolved

jorgensd reviewed Apr 23, 2025

View reviewed changes

ksagiyam force-pushed the ksagiyam/add_dag_visitor branch 2 times, most recently from a0442e7 to 8ecb3a3 Compare April 23, 2025 23:35

introduce DAGTraverser

a111e2e

- replace MultiFunction with DAGTraverser in apply_derivatives.py Co-authored-by: Connor Ward <[email protected]> Co-authored-by: Jørgen Schartum Dokken <[email protected]>

ksagiyam force-pushed the ksagiyam/add_dag_visitor branch from 3d63435 to a111e2e Compare April 24, 2025 13:59

ksagiyam added 2 commits April 24, 2025 16:03

singledispatch infer type

25e2a49

Revert "singledispatch infer type"

76e4f4a

This reverts commit 25e2a49.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Ksagiyam/add DAGTraverser #365

Ksagiyam/add DAGTraverser #365

ksagiyam commented Mar 26, 2025 •

edited

Loading

connorjward left a comment

connorjward Apr 3, 2025

ksagiyam Apr 3, 2025

connorjward Apr 3, 2025

ksagiyam Apr 3, 2025

connorjward Apr 3, 2025

ksagiyam Apr 3, 2025

connorjward Apr 3, 2025

ksagiyam Apr 3, 2025

connorjward Apr 3, 2025

ksagiyam commented Apr 23, 2025

jorgensd Apr 23, 2025

ksagiyam Apr 24, 2025 •

edited

Loading

jorgensd Apr 23, 2025

jorgensd Apr 23, 2025

ksagiyam Apr 24, 2025

jorgensd Apr 23, 2025

ksagiyam Apr 24, 2025

jorgensd Apr 24, 2025

jorgensd Apr 23, 2025

ksagiyam Apr 24, 2025

jorgensd Apr 23, 2025

ksagiyam Apr 24, 2025

jorgensd Apr 23, 2025 •

edited

Loading

ksagiyam Apr 24, 2025

jorgensd commented Apr 23, 2025

jorgensd commented Apr 23, 2025

schnellerhase commented Apr 23, 2025 •

edited

Loading

jorgensd commented Apr 23, 2025 •

edited

Loading

jorgensd commented Apr 23, 2025

ksagiyam commented Apr 23, 2025



		class DAGTraverser(ABC):
		"""Base class for dag traversers."""


		"""

		def __init__(self, compress=True, visited_cache=None, result_cache=None) -> None:

	def __init__(self, compress=True, visited_cache=None, result_cache=None) -> None:
	def __init__(self, compress: bool=True, visited_cache=None, result_cache=None) -> None:

Ksagiyam/add DAGTraverser #365

Are you sure you want to change the base?

Ksagiyam/add DAGTraverser #365

Conversation

ksagiyam commented Mar 26, 2025 • edited Loading

connorjward left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ksagiyam commented Apr 23, 2025

Choose a reason for hiding this comment

ksagiyam Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jorgensd Apr 23, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jorgensd commented Apr 23, 2025

jorgensd commented Apr 23, 2025

schnellerhase commented Apr 23, 2025 • edited Loading

jorgensd commented Apr 23, 2025 • edited Loading

jorgensd commented Apr 23, 2025

ksagiyam commented Apr 23, 2025

ksagiyam commented Mar 26, 2025 •

edited

Loading

ksagiyam Apr 24, 2025 •

edited

Loading

jorgensd Apr 23, 2025 •

edited

Loading

schnellerhase commented Apr 23, 2025 •

edited

Loading

jorgensd commented Apr 23, 2025 •

edited

Loading