Using numexpr and/or einsumt inside form definitions

[bhaveshshrimali]

Hi @adtzlr,

Is it permissible to use numexpr and/or einsumt inside form definitions, e.g. something like

import numpy as np
import felupe as fem
from felupe.math import det, inv, transpose, identity
import pypardiso
from einsumt import einsumt
from time import time
from meshio import xdmf
import matplotlib.pyplot as plt
from numexpr import evaluate as nev


t1 = time()


def F1(F, p, mu, lmbda):
    J = det(F)
    Finv = inv(F)
    t1 = p * J * transpose(Finv)
    t2 = nev("mu * F")
    return nev("t1 + t2")


def F2(F, p, mu, lmbda):
    J = det(F)
    Js = nev(
        "0.5 * (lmbda + p + 2.0 * (lmbda * mu + 0.25 * (lmbda + p) ** 2))**(0.5) / lmbda"
    )
    dJsdp = nev(
        "(0.25 * lmbda + 0.25 * p + 0.5 * (lmbda * mu + 0.25 * (lmbda + p) ** 2))**(0.5) / (lmbda * (lmbda * mu + 0.25 * (lmbda + p) ** 2)**(0.5))"
    )
    ans = nev("J - (Js + (p + mu / Js - lmbda * (Js - 1)) * dJsdp)")
    return ans


def A11(F, p, mu, lmbda):
    eye = identity(F)
    J = det(F)
    Finv = inv(F)
    einsum_lkji = einsumt("lkab,jiab->ijklab", Finv, Finv)
    einsum_jkli = einsumt("jkab,liab->ijklab", Finv, Finv)
    einsum_ikjl = einsumt("ikab,jlab->ijklab", eye, eye)
    L_nev = nev("p * J * einsum_lkji - p * J * einsum_jkli + mu * einsum_ikjl")
    # L = (
    #     p * J * einsumt("lkab,jiab->ijklab", Finv, Finv)
    #     - p * J * einsumt("jkab,liab->ijklab", Finv, Finv)
    #     + mu * einsumt("ikab,jlab->ijklab", eye, eye)
    # )
    return L_nev


def A12(F, p, mu, lmbda):
    J = det(F)
    Finv = inv(F)
    FinvT = transpose(Finv)
    ans = nev("J * FinvT")
    return ans


def A22(F, p, mu, lmbda):
    Js = nev(
        """(
        0.5
        * (lmbda + p + 2.0 * (lmbda * mu + 0.25 * (lmbda + p) ** 2) ** (0.5))
        / lmbda
    )"""
    )
    dJsdp = nev(
        """(
        0.25 * lmbda + 0.25 * p + 0.5 * (lmbda * mu + 0.25 * (lmbda + p) ** 2) ** (0.5)
    ) / (lmbda * (lmbda * mu + 0.25 * (lmbda + p) ** 2) ** (0.5))"""
    )
    d2Jdp2 = nev("""0.25 * mu / (lmbda * mu + 0.25 * (lmbda + p) ** 2) ** (3.0 / 2)""")
    L = nev(
        """(
        -2.0 * dJsdp
        - p * d2Jdp2
        + mu / Js**2 * dJsdp**2
        - mu / Js * d2Jdp2
        + lmbda * (Js - 1.0) * d2Jdp2
        + lmbda * dJsdp**2
    )"""
    )
    return L


def stress(x, mu, lmbda):
    F, p, statevars = x[0], x[1], x[-1]
    return [F1(F, p, mu, lmbda), F2(F, p, mu, lmbda), statevars]


def elasticity(x, mu, lmbda):
    F, p = x[0], x[1]
    return [A11(F, p, mu, lmbda), A12(F, p, mu, lmbda), A22(F, p, mu, lmbda)]


# mesh = fem.Cube(n=5).triangulate().add_midpoints_edges()

meshio_mesh = xdmf.read("trialMesh.xdmf")
mesh = fem.Mesh(
    points=meshio_mesh.points,
    cells=meshio_mesh.get_cells_type("tetra"),
    cell_type=meshio_mesh.cells[0].type,
)

region = fem.RegionTetra(mesh=mesh)
# print(dummy_region.dV)
print(mesh.ncells)
mesh = mesh.flip(np.any(region.dV < 0, axis=0))
region.reload(mesh)
assert np.all(region.dV > 0)
region_u = fem.RegionQuadraticTetra(mesh=mesh.add_midpoints_edges())
region_p = fem.RegionTetra(
    mesh=fem.mesh.dual(mesh, disconnect=False),
    quadrature=region_u.quadrature,
    grad=False,
)

displacement = fem.Field(region_u, dim=3)
pressure = fem.Field(region_p)
field = fem.FieldContainer([displacement, pressure])

# mark the boundary facets and bcs
left_x_zero = fem.Boundary(field[0], fx=0.0, skip=[0, 1, 1], name="leftZero")
bottom_y_zero = fem.Boundary(field[0], fy=0.0, skip=[1, 0, 1], name="bottomZero")
back_z_zero = fem.Boundary(field[0], fz=0.0, skip=[1, 1, 0], name="backZero")
right_x_pull = fem.Boundary(field[0], fx=1.0, skip=[0, 1, 1], name="rightPull")

boundaries, loadcase = fem.dof.uniaxial(field, clamped=False)

boundaries_man = {
    bd.name: bd for bd in [left_x_zero, bottom_y_zero, back_z_zero, right_x_pull]
}

umat = fem.Material(stress, elasticity, mu=1, lmbda=1e4)
solid = fem.SolidBody(umat, field)

stiffness = solid._matrix()
force = solid._vector()

move = fem.math.linsteps([0, 2], num=51)
step = fem.Step(
    items=[solid], ramp={boundaries_man["rightPull"]: move}, boundaries=boundaries_man
)

job = fem.CharacteristicCurve(steps=[step], boundary=boundaries_man["rightPull"])
job.evaluate(
    tol=1e-6,
    solver=pypardiso.spsolve,
    filename="feluperesults.xdmf",
    verbose=True,
    kwargs={"parallel": True},
)
print(time() - t1)
fig, ax = job.plot(
    xlabel="Displacement $u$ in mm $\longrightarrow$",
    ylabel="Normal Force $F$ in N $\longrightarrow$",
    color="blue",
    marker="o",
)
ax.grid(linestyle="--")
fig.savefig("felupeEx36ref.png")

line = ax.lines[0]

xdata, ydata = line.get_xdata(), line.get_ydata()

np.savetxt("dataFE.txt", np.vstack((xdata, ydata)))


ax2 = field.imshow("Principal Values of Logarithmic Strain")
plt.show()

While this is clumsy and coding per-se, it does speed up things even more dramatically.... I wanted to first check with you if this is permissible before stepping it through the debugger..

An initial thought is that it may not be (just comparing the force deformation curves) from Ex.36 reimplementation, but the error isn't so bad, especially at the end. Not sure if it is too simple a constitutive model, but thoughts?

[adtzlr]

Hi @bhaveshshrimali!

Of course, feel free to use all packages you like inside your functions as long as you return numpy arrays.

There is a typo in Js, it should be:

def F2(F, p, mu, lmbda):
    J = det(F)
    Js = nev(
        "0.5 * (lmbda + p + 2.0 * (lmbda * mu + 0.25 * (lmbda + p) ** 2)**(0.5)) / lmbda"
    )
    dJsdp = nev(
        "(0.25 * lmbda + 0.25 * p + 0.5 * (lmbda * mu + 0.25 * (lmbda + p) ** 2))**(0.5) / (lmbda * (lmbda * mu + 0.25 * (lmbda + p) ** 2)**(0.5))"
    )
    ans = nev("J - (Js + (p + mu / Js - lmbda * (Js - 1)) * dJsdp)")
    return ans

[bhaveshshrimali]

Thanks for spotting that. That was a bit sloppy on my part. Indeed it seems to be working good..

[adtzlr]

Hi @bhaveshshrimali - JAX with JIT+vmap could also be a nice alternative - see #869 for an example.

[bhaveshshrimali]

That's great @adtzlr I will try it out later in the week. Any performance/feature benefits compared to e.g. using matADi/tensortrax?

[adtzlr]

Hi @bhaveshshrimali, actually - I haven't tried it yet beside the simple Neo-Hookean material. This is also just a very basic approach. If it is useful, I'm thinking of something like felupe.Hyperelastic(..., backend="jax"). For now, it is available in the main-branch as felupe.constitution.jax.Hyperelastic. I still have to figure out where it will be stored in the next release(s).

I was not able to successfully test it with JIT. I could activate it but it was slower than without JIT. Done!

The main advantages of JAX are

1. its popularity,
2. a JIT compiler, so no further math-function-modifications (to support batches of trailing axes) are necessary,
3. its numpy- & scipy-compatible API makes it really easy to use.

A benefit in performance is to be expected.

A quick test on the (quite costly) MORPH material formulation gives a 6x speedup compared to tensortrax. No matter how complicated the material model is, everything is nearly as fast as the Neo Hookean material. It's fantastic!