Mixed formulation for Incompressible Hyperelasticity

[bhaveshshrimali]

Following #439 and inputs from ex18, I started working yesterday on a mixed formulation for incompressible hyperelasticity using the Taylor-Hood elements. The idea is to work on imposing Dirichlet BC's later today and then test a simple uniaxial tension later in the week. This should be fairly self contained. This is what I have so far. Please let me know whenever you get time to take a look. No rush from my end

from numpy import einsum, linalg as nla, zeros, zeros_like, concatenate, abs as npabs, split as npsplit
from scipy.sparse import bmat
from sympy import Symbol, simplify, lambdify, diff, sqrt
from skfem.helpers import grad, transpose
from typing import List, Optional
from numba import jit, prange
from skfem import *

@jit(cache=True, nopython=True, nogil=True, parallel=True)
def vdet(A):
    detA = zeros_like(A[0, 0])
    detA = A[0, 0] * (A[1, 1] * A[2, 2] -
                    A[1, 2] * A[2, 1]) -\
        A[0, 1] * (A[1, 0] * A[2, 2] -
                    A[1, 2] * A[2, 0]) +\
        A[0, 2] * (A[1, 0] * A[2, 1] -
                    A[1, 1] * A[2, 0])
    return detA

@jit(cache=True, nopython=True, nogil=True, parallel=True)
def vinv(A):
    invA = zeros_like(A)
    detA = vdet(A)
    invA[0, 0] = (-A[1, 2] * A[2, 1] +
                A[1, 1] * A[2, 2]) / detA
    invA[1, 0] = (A[1, 2] * A[2, 0] -
                A[1, 0] * A[2, 2]) / detA
    invA[2, 0] = (-A[1, 1] * A[2, 0] +
                A[1, 0] * A[2, 1]) / detA
    invA[0, 1] = (A[0, 2] * A[2, 1] -
                A[0, 1] * A[2, 2]) / detA
    invA[1, 1] = (-A[0, 2] * A[2, 0] +
                A[0, 0] * A[2, 2]) / detA
    invA[2, 1] = (A[0, 1] * A[2, 0] -
                A[0, 0] * A[2, 1]) / detA
    invA[0, 2] = (-A[0, 2] * A[1, 1] +
                A[0, 1] * A[1, 2]) / detA
    invA[1, 2] = (A[0, 2] * A[1, 0] -
                A[0, 0] * A[1, 2]) / detA
    invA[2, 2] = (-A[0, 1] * A[1, 0] +
                A[0, 0] * A[1, 1]) / detA
    return invA, detA


def F1(w):
    u = w["w"][0]
    p = w["w"][1]
    # print(w["w"][0])
    F = zeros_like(grad(u))
    for i in prange(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return p * J * transpose(Finv) + mu * F

def F2(w):
    u = w["w"][0]
    p = w["w"][1]
    F = zeros_like(grad(u))
    for i in prange(3):
        F[i,i] += 1.
    F += grad(u)
    J = vdet(F)
    Js = Jstar(p)
    return J-(Js + (p + mu/Js - lmbda*(Js - 1))*dJst_dp(p))

def A11(w):
    u = w["w"][0]
    p = w["w"][1]
    F = zeros_like(grad(u))
    eye = zeros_like(grad(u))
    for i in prange(3):
        F[i,i] += 1.
        eye[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    L = p*J*einsum("lk...,ji...->ijkl...",Finv, Finv) - p * J * einsum("jk...,li...->ijkl...", Finv, Finv) + mu * einsum("ik...,jl...->ijkl...", eye, eye)
    return L

def A12(w):
    u = w["w"][0]
    p = w["w"][1]
    F = zeros_like(grad(u))
    for i in prange(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A21(w):
    u = w["w"][0]
    p = w["w"][1]
    F = zeros_like(grad(u))
    for i in prange(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A22(w):
    u = w["w"][0]
    p = w["w"][1]
    # F = zeros_like(grad(u))
    # for i in prange(3):
    #     F[i,i] += 1.
    # F += grad(u)
    # Finv, J = vinv(F)
    Js = Jstar(p)
    dJdp = dJst_dp(p)
    d2Jdp2 = d2Jst_dp2(p)

    L = -2.*dJdp - p * d2Jdp2 + mu/Js**2 * dJdp**2 - mu/Js*d2Jdp2 + lmbda * (Js - 1.) * d2Jdp2 + lmbda * dJdp**2
    return L
    

mesh = MeshTet()
uelem = ElementVectorH1(ElementTetP2())
pelem = ElementTetP1()
elems = {
    "u": uelem,
    "p": pelem
}
basis = {
    field: InteriorBasis(mesh, e, intorder=3)
    for field, e in elems.items()
}

du = zeros(basis["u"].N)
dp = zeros(basis["p"].N)
stretch_ = 1.5

dofs = {
    "left":basis["u"].get_dofs(lambda x: x[0] < 1.e-6),
    "bottom":basis["u"].get_dofs(lambda x: x[1] < 1.e-6),
    "back":basis["u"].get_dofs(lambda x: x[2] < 1.e-6),
    "front":basis["u"].get_dofs(lambda x: x[2] > 1. - 1.e-6)
}
# print(dofs)
du[dofs["left"].nodal["u^1"]] = 0.
du[dofs["bottom"].nodal["u^2"]] = 0.
du[dofs["back"].nodal["u^3"]] = 0.
du[dofs["front"].nodal["u^3"]] = stretch_

I = basis["u"].complement_dofs(dofs)
# material parameters
mu, lmbda = 1., 1.e3

p = Symbol("p", real=True)
Jst = (lmbda + p + sqrt((lmbda+p)**2 + 4*lmbda*mu ))/(2*lmbda)
dJst_dp = lambdify(p, simplify(Jst.diff(p)), "numpy")
d2Jst_dp2 = lambdify(p, simplify(Jst.diff(p,2)), "numpy")
Jstar = lambdify(p, simplify(Jst), "numpy")
w = (
    basis["u"].interpolate(du),
    basis["p"].interpolate(dp)
)

@LinearForm
def a1(v, w):
    return einsum("ij...,ij...",F1(w), grad(v))

@LinearForm
def a2(v, w):
    return F2(w) * v

@BilinearForm
def b11(u, v, w):
    return einsum("ijkl...,ij...,kl...",A11(w), grad(u), grad(v))

@BilinearForm
def b12(u, v, w):
    return einsum("ij...,ij...",A12(w), grad(v)) * u

@BilinearForm
def b21(u, v, w):
    return einsum("ij...,ij...",A21(w), grad(v)) * u

@BilinearForm
def b22(u, v, w):
    return A22(w) * u * v

for itr in range(100):
    w = (
        basis["u"].interpolate(du),
        basis["p"].interpolate(dp)
    )

    K11 = asm(b11, basis["u"], basis["u"], w=w)
    K12 = asm(b12, basis["p"], basis["u"], w=w)
    K21 = asm(b21, basis["p"], basis["u"], w=w)
    K22 = asm(b22, basis["p"], basis["p"], w=w)
    f = concatenate((
        asm(a1, basis["u"], w=w),
        asm(a2, basis["p"], w=w)
    ))
    K = bmat(
        [[K11, K12],
         [K21.T, K22]], "csr"
    )
    uvp = solve(*condense(K, -f, I=I), use_umfpack=True)
    delu, delp = npsplit(uvp, [K11.shape[0]])
    du += delu
    dp += delp
    normu = nla.norm(delu)
    normp = nla.norm(delp)
    print(f"{itr+1}, norm_du: {normu}, norm_dp: {normp}")
    if normu < 1.e-6 and normp < 1.e-6:
        break

import meshio
from vedo import show as vShow, Mesh as vMesh
meshio.xdmf.write(
    "mixedFEM.xdmf",
    meshio.Mesh(mesh.p.T, {"tetra":mesh.t.T},point_data={"u":du[basis["u"].nodal_dofs].T})
)

Once I work through the details, and it happens to be of interest to you I'd be glad to contribute it as an (advanced) example.

[bhaveshshrimali]

I took a brief look at ex27. That helped me clean up the code a little bit. I had a question regarding writing the form jacobian correctly.

Writing the jacobian for a two-field mixed problem that depends on both the fields

Does the following look correct syntax-wise? That is, does w["disp"] and w["press"] to access the DiscreteField disp and press respectively make sense?

mesh = MeshTet()
mesh.refine(2)
uelem = ElementVectorH1(ElementTetP2())
pelem = ElementTetP1()
elems = {
    "u": uelem,
    "p": pelem
}
basis = {
    field: InteriorBasis(mesh, e, intorder=2)
    for field, e in elems.items()
}


def A11(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    eye = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
        eye[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    L = p*J*einsum("lk...,ji...->ijkl...",Finv, Finv) - p * J * einsum("jk...,li...->ijkl...", Finv, Finv) + mu * einsum("ik...,jl...->ijkl...", eye, eye)
    return L


@BilinearForm
def b11(u, v, w):
    return einsum("ijkl...,ij...,kl...",A11(w), grad(u), grad(v))

du = zeros(basis["u"].N)
dp = zeros(basis["p"].N)

uv = basis["u"].interpolate(du)
pv = basis["p"].interpolate(dp)

K11 = asm(b11, basis["u"], basis["u"], disp=uv, press = pv)

I can post a a MWE in sometime if the above snippet isn't sufficient. Sorry for posting too many questions at once. Please take your time, no hurries at my end.

[gdmcbain]

I think that looks O. K.

[bhaveshshrimali]

Thanks

[bhaveshshrimali]

Thanks so much for your help ! I'm attaching a code just in case if you are free and want to take a look. I think I should be able to debug it over the weekend and have a working neo-hookean incompressible hyperelasticity code. I don't think it should take me longer than that (unless I made a glaring mistake somewhere.. TBD)

The displacement made sense for a small test mesh but upon refinement resulted in incorrect values (Newton still converges). The exact value of the transverse displacement at the lateral face should be -0.29289321881345254.

Edit: I realized that it would be better to maybe keep this as-is since there is some tedious algebra in deriving the weak forms. I can post a documented code with the mathematical details (similar to one of the examples) or as a jupyter notebook later on via binder.

Small test mesh:

Newton:

1, norm_du: 5.663582879103, norm_dp: 2.5607547877285115
2, norm_du: 2.173881337517441, norm_dp: 3.342211577316037
3, norm_du: 0.39512278071861107, norm_dp: 1.2539219799665047
4, norm_du: 0.028389652739808858, norm_dp: 0.1021083592514163
5, norm_du: 0.00038749151295292125, norm_dp: 0.0005095287203490282
6, norm_du: 4.0873362256525696e-08, norm_dp: 1.4204691012559674e-07

After one refinement:

Newton:

1, norm_du: 15.410765449205675, norm_dp: 4.998414451293624
2, norm_du: 4.500520843456108, norm_dp: 8.936979546463057
3, norm_du: 0.7992009431110418, norm_dp: 4.213765171323738
4, norm_du: 0.07089776688817662, norm_dp: 0.564193714836833
5, norm_du: 0.004787999473443613, norm_dp: 0.010087501562038267
6, norm_du: 5.075643963084136e-05, norm_dp: 0.00018744046276803278
7, norm_du: 6.298006279559812e-09, norm_dp: 4.9291176582605015e-08

Complete code (for reference)

from numpy import einsum, linalg as nla, zeros, zeros_like, concatenate, split as npsplit, hstack
from scipy.sparse import bmat
from sympy import Symbol, simplify, lambdify, diff, sqrt
from skfem.helpers import grad, transpose
from typing import List, Optional
from numba import jit
from skfem import *

# material parameters
mu, lmbda = 1., 1.e3

p = Symbol("p", real=True)
Jst = (lmbda + p + sqrt((lmbda+p)**2 + 4*lmbda*mu ))/(2*lmbda)
dJst_dp = lambdify(p, simplify(Jst.diff(p)), "numpy")
d2Jst_dp2 = lambdify(p, simplify(Jst.diff(p,2)), "numpy")
Jstar = lambdify(p, simplify(Jst), "numpy")


@jit(cache=True, nopython=True)
def vdet(A):
    detA = zeros_like(A[0, 0])
    detA = A[0, 0] * (A[1, 1] * A[2, 2] -
                    A[1, 2] * A[2, 1]) -\
        A[0, 1] * (A[1, 0] * A[2, 2] -
                    A[1, 2] * A[2, 0]) +\
        A[0, 2] * (A[1, 0] * A[2, 1] -
                    A[1, 1] * A[2, 0])
    return detA

@jit(cache=True, nopython=True)
def vinv(A):
    invA = zeros_like(A)
    detA = vdet(A)
    invA[0, 0] = (-A[1, 2] * A[2, 1] +
                A[1, 1] * A[2, 2]) / detA
    invA[1, 0] = (A[1, 2] * A[2, 0] -
                A[1, 0] * A[2, 2]) / detA
    invA[2, 0] = (-A[1, 1] * A[2, 0] +
                A[1, 0] * A[2, 1]) / detA
    invA[0, 1] = (A[0, 2] * A[2, 1] -
                A[0, 1] * A[2, 2]) / detA
    invA[1, 1] = (-A[0, 2] * A[2, 0] +
                A[0, 0] * A[2, 2]) / detA
    invA[2, 1] = (A[0, 1] * A[2, 0] -
                A[0, 0] * A[2, 1]) / detA
    invA[0, 2] = (-A[0, 2] * A[1, 1] +
                A[0, 1] * A[1, 2]) / detA
    invA[1, 2] = (A[0, 2] * A[1, 0] -
                A[0, 0] * A[1, 2]) / detA
    invA[2, 2] = (-A[0, 1] * A[1, 0] +
                A[0, 0] * A[1, 1]) / detA
    return invA, detA


def F1(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return p * J * transpose(Finv) + mu * F

def F2(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    J = vdet(F)
    Js = Jstar(p)
    return J-(Js + (p + mu/Js - lmbda*(Js - 1))*dJst_dp(p))

def A11(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    eye = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
        eye[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    L = p*J*einsum("lk...,ji...->ijkl...",Finv, Finv) - p * J * einsum("jk...,li...->ijkl...", Finv, Finv) + mu * einsum("ik...,jl...->ijkl...", eye, eye)
    return L

def A12(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A21(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A22(w):
    u = w["disp"]
    p = w["press"]
    Js = Jstar(p)
    dJdp = dJst_dp(p)
    d2Jdp2 = d2Jst_dp2(p)
    L = -2.*dJdp - p * d2Jdp2 + mu/Js**2 * dJdp**2 - mu/Js*d2Jdp2 + lmbda * (Js - 1.) * d2Jdp2 + lmbda * dJdp**2
    return L
    

mesh = MeshTet()
mesh.refine()
uelem = ElementVectorH1(ElementTetP2())
pelem = ElementTetP1()
elems = {
    "u": uelem,
    "p": pelem
}
basis = {
    field: InteriorBasis(mesh, e, intorder=4)
    for field, e in elems.items()
}

du = zeros(basis["u"].N)
dp = zeros(basis["p"].N)
stretch_ = 1.

dofsold = [
    basis["u"].find_dofs({"left":mesh.facets_satisfying(lambda x: x[0] == 0.)}, skip=["u^2", "u^3"]),
    basis["u"].find_dofs({"bottom":mesh.facets_satisfying(lambda x: x[1] == 0.)}, skip=["u^1", "u^3"]),
    basis["u"].find_dofs({"back":mesh.facets_satisfying(lambda x: x[2] == 0.)}, skip=["u^1", "u^2"]),
    basis["u"].find_dofs({"front":mesh.facets_satisfying(lambda x: x[2] == 1.)}, skip=["u^1", "u^2"])
]

dofs = {}
for dof in dofsold:
    dofs.update(dof)

du[dofs["left"].all()] = 0.
du[dofs["bottom"].all()] = 0.
du[dofs["back"].all()] = 0.
du[dofs["front"].all()] = stretch_

I = hstack((
    basis["u"].complement_dofs(dofs),
    basis["u"].N + basis["p"].find_dofs()["all"].all()
))

@LinearForm
def a1(v, w):
    return einsum("ij...,ij...", F1(w), grad(v))

@LinearForm
def a2(v, w):
    return F2(w) * v

@BilinearForm
def b11(u, v, w):
    return einsum("ijkl...,ij...,kl...", A11(w), grad(u), grad(v))

@BilinearForm
def b12(u, v, w):
    return einsum("ij...,ij...", A12(w), grad(v)) * u

@BilinearForm
def b21(u, v, w):
    return einsum("ij...,ij...", A21(w), grad(v)) * u

@BilinearForm
def b22(u, v, w):
    return A22(w) * u * v


for itr in range(100):
    uv = basis["u"].interpolate(du)
    pv = basis["p"].interpolate(dp) 

    K11 = asm(b11, basis["u"], basis["u"], disp=uv, press = pv)
    K12 = asm(b12, basis["p"], basis["u"], disp=uv, press = pv)
    K21 = asm(b21, basis["p"], basis["u"], disp=uv, press = pv)
    K22 = asm(b22, basis["p"], basis["p"], disp=uv, press = pv)

    f = concatenate((
        asm(a1, basis["u"], disp=uv, press = pv),
        asm(a2, basis["p"], disp=uv, press = pv)
    ))
    K = bmat(
        [[K11, K12],
         [K21.T, K22]], "csr"
    )
    uvp = solve(*condense(K, -f, I=I), use_umfpack=True)
    delu, delp = npsplit(uvp, [du.shape[0]])
    du += delu
    dp += delp
    normu = nla.norm(delu)
    normp = nla.norm(delp)
    print(f"{itr+1}, norm_du: {normu}, norm_dp: {normp}")
    if normu < 1.e-6 and normp < 1.e-6:
        break

import meshio
meshio.xdmf.write(
    "mixedFEM.xdmf",
    meshio.Mesh(mesh.p.T, {"tetra":mesh.t.T},point_data={"u":du[basis["u"].nodal_dofs].T})
)
print(f"transverse disp: {1./(1+stretch_)**(0.5)-1.}")

[bhaveshshrimali]

I made another attempt at debugging the above error, but in doing so I guess I had a question regarding the global ordering of dofs. I think this should depend on how I setup the block matrix A in Ax = b, but still wanted to make sure if there is a problem in how I am interpreting the dof ordering.

from numpy import zeros, abs as npabs, hstack
from skfem import *
mesh = MeshTet()
uelem = ElementVectorH1(ElementTetP2())
pelem = ElementTetP1()
elems = {
    "u": uelem,
    "p": pelem
}
basis = {
    field: InteriorBasis(mesh, e, intorder=2)
    for field, e in elems.items()
}

du = zeros(basis["u"].N)
dp = zeros(basis["p"].N)

dofsold = [
    basis["u"].find_dofs({"left":mesh.facets_satisfying(lambda x: x[0] < 1.e-6)}, skip=["u^2", "u^3"]),
    basis["u"].find_dofs({"bottom":mesh.facets_satisfying(lambda x: x[1] < 1.e-6)}, skip=["u^1", "u^3"]),
    basis["u"].find_dofs({"back":mesh.facets_satisfying(lambda x: x[2] < 1.e-6)}, skip=["u^1", "u^2"]),
    basis["u"].find_dofs({"front":mesh.facets_satisfying(lambda x: npabs(x[2]-1.) < 1.e-6 )}, skip=["u^1", "u^2"])
]

dofs = {}
for dof in dofsold:
    dofs.update(dof)

I = hstack((
    basis["u"].complement_dofs(dofs),
    basis["u"].N + basis["p"].find_dofs()["all"].all()
))

Does I correctly capture all the free dofs in the system (all free displacement followed by all pressure dofs)? Or should I change anything here ?

[kinnala]

I cannot see any obvious mistake in the definition of I.

[kinnala]

Did you try plotting p? Usually in Stokes that reveals instabilities, if any. Sometimes even might pinpoint a source of an issue, e.g., oscillating near a corner with wrong BC.

[bhaveshshrimali]

Did you try plotting p? Usually in Stokes that reveals instabilities, if any.

Thanks a lot. Will do it now. I checked the variational formulation and it seems to be fine.

[kinnala]

I'll check it again with fresh eyes tomorrow. One question: do K12 and K21 have the same values?

[bhaveshshrimali]

I'll check it again with fresh eyes tomorrow.

Sure! No hurries at my end. I would also spend sometime today on debugging everything.

One question: do K12 and K21 have the same values?

Yes they have the same values. I will write down the math in a notebook and post it here (Do you recommend anything else?). That way it would be easy to follow for anyone who takes a look at this in the future (including myself:)).

[bhaveshshrimali]

FEniCS vs scikit-FEM

I plotted the pressure obtained from solving the problem in FEniCS (right) and in scikit-FEM (left) for the same mesh

scikit-FEM fine (left) vs coarse (right) mesh

However when I changed the mesh to a coarser one, the pressure values match exactly to what they are in FEniCS. Printing the pressure dofs gives

[-0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5
 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5 -0.5]

and the plot:

I will debug it with a fresh look tomorrow...

[bhaveshshrimali]

This is the last updated code in case if you want to take a look.

r"""
The variational problem is: 
    
          \min_u \max_p : p * (J - J*) + mu /2 * (I_1 - 3) - mu * ln(J*) + lmbda/2 * (J* - 1)^2

with 
u_1(0,y,z) = 0
u_2(x,0,z) = 0
u_3(x,y,0) = 0
u_3(x,y,1) = 1.
"""

from numpy import einsum, linalg as nla, zeros, zeros_like, concatenate, split as npsplit, hstack, abs as npabs
from scipy.sparse import bmat
from sympy import Symbol, simplify, lambdify, diff, sqrt
from skfem.helpers import grad, transpose
from typing import List, Optional
from numba import jit, prange
from skfem import *

# material parameters
mu, lmbda = 1., 1.e3

p = Symbol("p", real=True)
Jst = (lmbda + p + sqrt((lmbda + p)**2 + 4.*lmbda*mu ))/(2.*lmbda)
dJst_dp = lambdify(p, simplify(Jst.diff(p)), "numpy")
d2Jst_dp2 = lambdify(p, simplify(Jst.diff(p,2)), "numpy")
Jstar = lambdify(p, simplify(Jst), "numpy")

@jit(nopython=True, cache=True, fastmath=True, nogil=True, parallel=True)
def vdet(F):
    J = zeros_like(F[0,0])
    for a in prange(J.shape[0]):
        for b in prange(J.shape[1]):
            J[a,b] += F[0, 0, a, b] * (F[1, 1, a, b] * F[2, 2, a, b] -
                                        F[1, 2, a, b] * F[2, 1, a, b]) -\
                    F[0, 1, a, b] * (F[1, 0, a, b] * F[2, 2, a, b] -
                                        F[1, 2, a, b] * F[2, 0, a, b]) +\
                    F[0, 2, a, b] * (F[1, 0,a ,b] * F[2, 1, a, b] -
                                        F[1, 1, a, b] * F[2, 0, a, b])
    return J


@jit(nopython=True, cache=True, fastmath=True, nogil=True, parallel=True)
def vinv(F):
    J = vdet(F)
    Finv = zeros_like(F)
    for a in prange(J.shape[0]):
        for b in prange(J.shape[1]):
            Finv[0, 0, a, b] += (-F[1, 2, a, b] * F[2, 1, a, b] +
                                F[1, 1, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[1, 0, a, b] += (F[1, 2, a, b] * F[2, 0, a, b] -
                                F[1, 0, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[2, 0, a, b] += (-F[1, 1, a, b] * F[2, 0, a, b] +
                                F[1, 0, a, b] * F[2, 1, a, b]) / J[a, b]
            Finv[0, 1, a, b] += (F[0, 2, a, b] * F[2, 1, a, b] -
                                F[0, 1, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[1, 1, a, b] += (-F[0, 2, a, b] * F[2, 0, a, b] +
                                F[0, 0, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[2, 1, a, b] += (F[0, 1, a, b] * F[2, 0, a, b] -
                                F[0, 0, a, b] * F[2, 1, a, b]) / J[a, b]
            Finv[0, 2, a, b] += (-F[0, 2, a, b] * F[1, 1, a, b] +
                                F[0, 1, a, b] * F[1, 2, a, b]) / J[a, b]
            Finv[1, 2, a, b] += (F[0, 2, a, b] * F[1, 0, a, b] -
                                F[0, 0, a, b] * F[1, 2, a, b]) / J[a, b]
            Finv[2, 2, a, b] += (-F[0, 1, a, b] * F[1, 0, a, b] +
                                F[0, 0, a, b] * F[1, 1, a, b]) / J[a, b]
    return Finv, J

def F1(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return p * J * transpose(Finv) + mu * F

def F2(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    J = vdet(F)
    Js = Jstar(p)
    return J - (Js + (p + mu/Js - lmbda*(Js - 1))*dJst_dp(p))

def A11(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    eye = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
        eye[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    L = p * J * einsum("lk...,ji...->ijkl...", Finv, Finv) - p * J * einsum("jk...,li...->ijkl...", Finv, Finv) + mu * einsum("ik...,jl...->ijkl...", eye, eye)
    return L

def A12(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A21(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A22(w):
    u = w["disp"]
    p = w["press"]
    Js = Jstar(p)
    dJdp = dJst_dp(p)
    d2Jdp2 = d2Jst_dp2(p)
    L = -2. * dJdp - p * d2Jdp2 + mu/Js**2 * dJdp**2 - mu/Js * d2Jdp2 + lmbda * (Js - 1.) * d2Jdp2 + lmbda * dJdp**2
    return L
    

mesh = MeshTet()
mesh.refine()
uelem = ElementVectorH1(ElementTetP2())
pelem = ElementTetP1()
elems = {
    "u": uelem,
    "p": pelem
}
basis = {
    field: InteriorBasis(mesh, e, intorder=2)
    for field, e in elems.items()
}

du = zeros(basis["u"].N)
dp = zeros(basis["p"].N)
stretch_ = 1.

dofsold = [
    basis["u"].find_dofs({"left":mesh.facets_satisfying(lambda x: x[0] < 1.e-6)}, skip=["u^2", "u^3"]),
    basis["u"].find_dofs({"bottom":mesh.facets_satisfying(lambda x: x[1] < 1.e-6)}, skip=["u^1", "u^3"]),
    basis["u"].find_dofs({"back":mesh.facets_satisfying(lambda x: x[2] < 1.e-6)}, skip=["u^1", "u^2"]),
    basis["u"].find_dofs({"front":mesh.facets_satisfying(lambda x: npabs(x[2]-1.) < 1.e-6 )}, skip=["u^1", "u^2"])
]

dofs = {}
for dof in dofsold:
    dofs.update(dof)

du[dofs["left"].all()] = 0.
du[dofs["bottom"].all()] = 0.
du[dofs["back"].all()] = 0.
du[dofs["front"].all()] = stretch_

I = hstack((
    basis["u"].complement_dofs(dofs),
    basis["u"].N + basis["p"].find_dofs()["all"].all()
))

@LinearForm
def a1(v, w):
    return einsum("ij...,ij...", F1(w), grad(v))

@LinearForm
def a2(v, w):
    return F2(w) * v

@BilinearForm
def b11(u, v, w):
    return einsum("ijkl...,ij...,kl...", A11(w), grad(u), grad(v))

@BilinearForm
def b12(u, v, w):
    return einsum("ij...,ij...", A12(w), grad(v)) * u

@BilinearForm
def b21(u, v, w):
    return einsum("ij...,ij...", A21(w), grad(v)) * u

@BilinearForm
def b22(u, v, w):
    return A22(w) * u * v


for itr in range(100):
    uv = basis["u"].interpolate(du)
    pv = basis["p"].interpolate(dp) 

    K11 = asm(b11, basis["u"], basis["u"], disp=uv, press = pv)
    K12 = asm(b12, basis["p"], basis["u"], disp=uv, press = pv)
    # K21 = asm(b21, basis["p"], basis["u"], disp=uv, press = pv)
    K22 = asm(b22, basis["p"], basis["p"], disp=uv, press = pv)
    # print(K11.shape, K12.shape, K21.shape, K22.shape)
    f = concatenate((
        asm(a1, basis["u"], disp=uv, press = pv),
        asm(a2, basis["p"], disp=uv, press = pv)
    ))
    K = bmat(
        [[K11, K12],
         [K12.T, K22]], "csr"
    )
    uvp = solve(*condense(K, -f, I=I), use_umfpack=True)
    delu, delp = npsplit(uvp, [du.shape[0]])
    du += delu
    dp += delp
    normu = nla.norm(delu)
    normp = nla.norm(delp)
    print(f"{itr+1}, norm_du: {normu}, norm_dp: {normp}")
    if normu < 1.e-8 and normp < 1.e-8:
        break

import meshio
with meshio.xdmf.TimeSeriesWriter("mixedFEMSingle.xdmf") as writer:
    writer.write_points_cells(points=mesh.p.T, cells={"tetra":mesh.t.T})
    for t in range(2):
        writer.write_data(t, point_data={"u":t*du[basis["u"].nodal_dofs].T,
                                         "p":t*dp[basis["p"].nodal_dofs].T})
print(dp)
print(f"trans stretch: {1./(1+stretch_)**(0.5)-1.}")

[kinnala]

I made a simple test but found nothing yet: it's probably not related to the edge DOF's of ElementTetP2 because the same behaviour happens with ElementTetMini which has no edge DOF's.

[kinnala]

So maybe it was something related to the DOF's after all? I changed

    uvp = solve(*condense(K, -f, I=I), use_umfpack=True)

to

    uvp = solve(*condense(K, -f, D=dofs), use_umfpack=True)

and now I get

Edit: Forgot to mention that the above uses ElementTetMini but I now verified that I get the same correct result with ElementTetP2.

[bhaveshshrimali]

    uvp = solve(*condense(K, -f, D=dofs), use_umfpack=True)

Thanks a lot @kinnala! That was a bit sloppy on my side. I still will figure out what was wrong with assigning free dofs (I).

[bhaveshshrimali]

So I was using

basis["u"].N + basis["p"].find_dofs()["all"].all() 

but forgot to check if it indeed returns all the pressure dofs which is what I need in this case. Went back and read the documentation once again

scikit-fem/skfem/assembly/basis/basis.py

Lines 109 to 113 in 5d07cad

	Parameters
	----------
	facets
	A dictionary of facets. If ``None``, use ``self.mesh.boundaries``
	if set or otherwise use ``{'all': self.mesh.boundary_facets()}``.

and also

np.setdiff1d(np.arange(basis["p"].N), basis["p"].find_dofs()["all"].all())

to check that it wasn't selecting all the pressure dofs in I.

TLDR;

I = np.hstack((
    basis["u"].complement_dofs(dofs),
    basis["u"].N + arange(basis["p"].N)
))

works as expected. Thanks so much for your help again. Finding an error at this place wasn't very high in my list of possibilities.

I will clean up the complete code and post it here for reference. The timings are OK for this one (~20s) for this example, but I would still go ahead and write a numba version for the jacobian and residual as it could be useful.

Next target implementing #437 for this example instead of taylor-hood.

[kinnala]

but forgot to check if it indeed returns all the pressure dofs which is what I need in this case.

I see, right now find_dofs (etc.) are targeted more towards finding boundary DOF's than interior DOF's although the key all may suggest something else. I still haven't quite figured out how to design a proper interface for finding and handling DOF's.

[bhaveshshrimali]

I think basis["u"].N + arange(basis["p"].N) is a more natural way that someone would think of for this case (which itself may be pretty niche). It was just when I was writing things the first time, and having used find_dofs right before that to assign Dirichlet BC's, I tried to use find_dofs to select free dofs as well (without properly reading the docstring).

Edit: maybe all_boundary could be a more intuitive key (although the docstring is clear enough on it's own)

[bhaveshshrimali]

So here are the cleaned up versions of the code with and without numba

Without Numba (in calculation of the jacobian blocks explicitly)

from numpy import einsum, linalg as nla, zeros, zeros_like, concatenate, split as npsplit, hstack, abs as npabs, arange
from scipy.sparse import bmat
from sympy import symbols, simplify, lambdify, diff, sqrt, log
from skfem.helpers import grad, transpose
from numba import jit
from skfem import *

mu, lmbda = 1., 1.e3

p, I1, J = symbols("p, I1, J", real=True)

Jst = (lmbda + p + sqrt((lmbda + p)**2 + 4.*lmbda*mu ))/(2.*lmbda)
psiHat = p * (J - Jst) + mu/2. * (I1 - 3.) - mu * log(Jst) + lmbda / 2. * (Jst - 1.)**2
dpsi_dp = lambdify((p, J), simplify(psiHat.diff(p)), "numpy")
d2psi_dp2 = lambdify(p, simplify(psiHat.diff(p, 2)), "numpy")
dJst_dp = lambdify(p, simplify(Jst.diff(p)), "numpy")
d2Jst_dp2 = lambdify(p, simplify(Jst.diff(p,2)), "numpy")
Jstar = lambdify(p, simplify(Jst), "numpy")


@jit(cache=True, nopython=True)
def vdet(A):
    detA = zeros_like(A[0, 0])
    detA = A[0, 0] * (A[1, 1] * A[2, 2] -
                    A[1, 2] * A[2, 1]) -\
        A[0, 1] * (A[1, 0] * A[2, 2] -
                    A[1, 2] * A[2, 0]) +\
        A[0, 2] * (A[1, 0] * A[2, 1] -
                    A[1, 1] * A[2, 0])
    return detA

@jit(cache=True, nopython=True)
def vinv(A):
    invA = zeros_like(A)
    detA = vdet(A)
    invA[0, 0] = (-A[1, 2] * A[2, 1] +
                A[1, 1] * A[2, 2]) / detA
    invA[1, 0] = (A[1, 2] * A[2, 0] -
                A[1, 0] * A[2, 2]) / detA
    invA[2, 0] = (-A[1, 1] * A[2, 0] +
                A[1, 0] * A[2, 1]) / detA
    invA[0, 1] = (A[0, 2] * A[2, 1] -
                A[0, 1] * A[2, 2]) / detA
    invA[1, 1] = (-A[0, 2] * A[2, 0] +
                A[0, 0] * A[2, 2]) / detA
    invA[2, 1] = (A[0, 1] * A[2, 0] -
                A[0, 0] * A[2, 1]) / detA
    invA[0, 2] = (-A[0, 2] * A[1, 1] +
                A[0, 1] * A[1, 2]) / detA
    invA[1, 2] = (A[0, 2] * A[1, 0] -
                A[0, 0] * A[1, 2]) / detA
    invA[2, 2] = (-A[0, 1] * A[1, 0] +
                A[0, 0] * A[1, 1]) / detA
    return invA, detA

def F1(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return p * J * transpose(Finv) + mu * F

def F2(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    J = vdet(F)
    Js = Jstar(p)
    return J - (Js + (p + mu/Js - lmbda*(Js - 1))*dJst_dp(p))

def A11(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    eye = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
        eye[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    L = p * J * einsum("lk...,ji...->ijkl...", Finv, Finv) - p * J * einsum("jk...,li...->ijkl...", Finv, Finv) + mu * einsum("ik...,jl...->ijkl...", eye, eye)
    return L

def A12(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A21(w):
    u = w["disp"]
    p = w["press"]
    F = zeros_like(grad(u))
    for i in range(3):
        F[i,i] += 1.
    F += grad(u)
    Finv, J = vinv(F)
    return J * transpose(Finv)

def A22(w):
    u = w["disp"]
    p = w["press"]
    
    Js = Jstar(p)
    dJdp = dJst_dp(p)
    d2Jdp2 = d2Jst_dp2(p)
    L = -2. * dJdp - p * d2Jdp2 + mu/Js**2 * dJdp**2 - mu/Js * d2Jdp2 + lmbda * (Js - 1.) * d2Jdp2 + lmbda * dJdp**2
    return L
    

mesh = MeshTet()
mesh.refine(3)
uelem = ElementVectorH1(ElementTetP2())
pelem = ElementTetP1()
elems = {
    "u": uelem,
    "p": pelem
}
basis = {
    field: InteriorBasis(mesh, e, intorder=2)
    for field, e in elems.items()
}

du = zeros(basis["u"].N)
dp = zeros(basis["p"].N)
stretch_ = 1.

dofsold = [
    basis["u"].find_dofs({"left":mesh.facets_satisfying(lambda x: x[0] < 1.e-6)}, skip=["u^2", "u^3"]),
    basis["u"].find_dofs({"bottom":mesh.facets_satisfying(lambda x: x[1] < 1.e-6)}, skip=["u^1", "u^3"]),
    basis["u"].find_dofs({"back":mesh.facets_satisfying(lambda x: x[2] < 1.e-6)}, skip=["u^1", "u^2"]),
    basis["u"].find_dofs({"front":mesh.facets_satisfying(lambda x: npabs(x[2]-1.) < 1.e-6 )}, skip=["u^1", "u^2"])
]

dofs = {}
for dof in dofsold:
    dofs.update(dof)

du[dofs["left"].all()] = 0.
du[dofs["bottom"].all()] = 0.
du[dofs["back"].all()] = 0.
du[dofs["front"].all()] = stretch_

I = hstack((
    basis["u"].complement_dofs(dofs),
    basis["u"].N + arange(basis["p"].N)
))

@LinearForm
def a1(v, w):
    return einsum("ij...,ij...", F1(w), grad(v))

@LinearForm
def a2(v, w):
    return F2(w) * v

@BilinearForm
def b11(u, v, w):
    return einsum("ijkl...,ij...,kl...", A11(w), grad(u), grad(v))

@BilinearForm
def b12(u, v, w):
    return einsum("ij...,ij...", A12(w), grad(v)) * u

@BilinearForm
def b21(u, v, w):
    return einsum("ij...,ij...", A21(w), grad(v)) * u

@BilinearForm
def b22(u, v, w):
    return A22(w) * u * v

for itr in range(12):
    uv = basis["u"].interpolate(du)
    pv = basis["p"].interpolate(dp) 

    K11 = asm(b11, basis["u"], basis["u"], disp=uv, press = pv)
    K12 = asm(b12, basis["p"], basis["u"], disp=uv, press = pv)
    K22 = asm(b22, basis["p"], basis["p"], disp=uv, press = pv)
    f = concatenate((
        asm(a1, basis["u"], disp=uv, press = pv),
        asm(a2, basis["p"], disp=uv, press = pv)
    ))
    K = bmat(
        [[K11, K12],
         [K12.T, K22]], "csr"
    )
    uvp = solve(*condense(K, -f, I=I), use_umfpack=True)
    delu, delp = npsplit(uvp, [du.shape[0]])
    du += delu
    dp += delp
    normu = nla.norm(delu)
    normp = nla.norm(delp)
    # print(f"{itr+1}, norm_du: {normu}, norm_dp: {normp}")
    if normu < 1.e-8 and normp < 1.e-8:
        break

import meshio, os
numbaResults = os.path.join(os.getcwd(), "normalResults")
os.makedirs(numbaResults, exist_ok=True)
filw = os.path.join(numbaResults, "mixedFEMSingle.xdmf")
with meshio.xdmf.TimeSeriesWriter(filw) as writer:
    writer.write_points_cells(points=mesh.p.T, cells={"tetra":mesh.t.T})
    for t in range(2):
        writer.write_data(t, point_data={"u":t*du[basis["u"].nodal_dofs].T,
                                         "p":t*dp[basis["p"].nodal_dofs].T})

With Numba

from numpy import einsum, linalg as nla, zeros, zeros_like, sqrt, eye, concatenate, split as npsplit, hstack, abs as npabs, arange
from scipy.sparse import bmat
from skfem.helpers import grad
from numba import jit, prange
from skfem import *

# material parameters
mu, lmbda = 1., 1.e3

@jit(nopython=True, cache=True)
def Jstar(p):
    return (lmbda + p + sqrt((lmbda + p)**2. + 4.*lmbda*mu ))/(2.*lmbda)

@jit(nopython=True, cache=True)
def d2Jst_dp2(p):
    return 2.*mu/(4.*lmbda*mu + (lmbda + p)**2.)**(3./2.)

@jit(nopython=True, cache=True)
def dJst_dp(p):
    return (lmbda + p + sqrt(4.*lmbda*mu + (lmbda + p)**2.))/(2.*lmbda*sqrt(4.*lmbda*mu + (lmbda + p)**2.))

@jit(nopython=True, cache=True, fastmath=True, nogil=True, parallel=True)
def vdet(F):
    J = zeros_like(F[0,0])
    for a in prange(J.shape[0]):
        for b in prange(J.shape[1]):
            J[a,b] += F[0, 0, a, b] * (F[1, 1, a, b] * F[2, 2, a, b] -
                                        F[1, 2, a, b] * F[2, 1, a, b]) -\
                    F[0, 1, a, b] * (F[1, 0, a, b] * F[2, 2, a, b] -
                                        F[1, 2, a, b] * F[2, 0, a, b]) +\
                    F[0, 2, a, b] * (F[1, 0,a ,b] * F[2, 1, a, b] -
                                        F[1, 1, a, b] * F[2, 0, a, b])
    return J


@jit(nopython=True, cache=True, fastmath=True, nogil=True, parallel=True)
def vinv(F):
    J = vdet(F)
    Finv = zeros_like(F)
    for a in prange(J.shape[0]):
        for b in prange(J.shape[1]):
            Finv[0, 0, a, b] += (-F[1, 2, a, b] * F[2, 1, a, b] +
                                F[1, 1, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[1, 0, a, b] += (F[1, 2, a, b] * F[2, 0, a, b] -
                                F[1, 0, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[2, 0, a, b] += (-F[1, 1, a, b] * F[2, 0, a, b] +
                                F[1, 0, a, b] * F[2, 1, a, b]) / J[a, b]
            Finv[0, 1, a, b] += (F[0, 2, a, b] * F[2, 1, a, b] -
                                F[0, 1, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[1, 1, a, b] += (-F[0, 2, a, b] * F[2, 0, a, b] +
                                F[0, 0, a, b] * F[2, 2, a, b]) / J[a, b]
            Finv[2, 1, a, b] += (F[0, 1, a, b] * F[2, 0, a, b] -
                                F[0, 0, a, b] * F[2, 1, a, b]) / J[a, b]
            Finv[0, 2, a, b] += (-F[0, 2, a, b] * F[1, 1, a, b] +
                                F[0, 1, a, b] * F[1, 2, a, b]) / J[a, b]
            Finv[1, 2, a, b] += (F[0, 2, a, b] * F[1, 0, a, b] -
                                F[0, 0, a, b] * F[1, 2, a, b]) / J[a, b]
            Finv[2, 2, a, b] += (-F[0, 1, a, b] * F[1, 0, a, b] +
                                F[0, 0, a, b] * F[1, 1, a, b]) / J[a, b]
    return Finv, J

@jit(nopython=True, cache=True, fastmath=True, nogil=True, parallel=True)
def F1(dv, du, p):
    
    out = zeros_like(du[0, 0])
    F = zeros_like(du)
    for a in prange(du.shape[2]):
        for b in prange(du.shape[3]):
            F[0, 0, a, b] += 1.
            F[1, 1, a, b] += 1.
            F[2, 2, a, b] += 1.
    
    F += du
    Finv, J = vinv(F)
    for i in range(du.shape[0]):
        for j in range(du.shape[1]):
            for a in prange(J.shape[0]):
                for b in prange(J.shape[1]):
                    out[a,b] += (mu * F[i, j, a, b] + p[a,b] * J[a,b] * Finv[j, i, a, b]) * dv[i, j, a, b]
    return out

@jit(nopython=True, cache=True, fastmath=True, nogil=True, parallel=True)
def F2(dv, du, p):
    out = zeros_like(du[0,0])
    Js = zeros_like(du[0,0])
    dJdp = zeros_like(du[0,0])
    F = zeros_like(du)
    
    for a in prange(du.shape[2]):
        for b in prange(du.shape[3]):
            F[0, 0, a, b] += 1.
            F[1, 1, a, b] += 1.
            F[2, 2, a, b] += 1.
    
    F += du

    J = vdet(F)
    for a in prange(p.shape[0]):
        for b in prange(p.shape[1]):
            Js[a,b] = (lmbda + p[a,b] + sqrt((lmbda + p[a,b])**2. + 4.*lmbda*mu ))/(2.*lmbda)
            dJdp[a, b] = (lmbda + p[a,b] + sqrt(4.*lmbda*mu + (lmbda + p[a,b])**2.))/(2.*lmbda*sqrt(4.*lmbda*mu + (lmbda + p[a,b])**2.))
            out[a, b] += (J[a,b] - (Js[a,b] + (p[a,b] + mu/Js[a,b] - lmbda*(Js[a,b] - 1.)) * dJdp[a,b])) * dv[a,b]

    return out

@jit(nopython=True, cache=True, nogil=True, parallel=True)
def A11(du, dv, dw, p):
    out = zeros_like(p)
    F = zeros_like(dw)
    for a in prange(dw.shape[2]):
        for b in prange(dw.shape[3]):
            F[0, 0, a, b] += 1.
            F[1, 1, a, b] += 1.
            F[2, 2, a, b] += 1.
    
    F += dw
    Finv, J = vinv(F)

    kron = eye(dw.shape[0])
    for i in range(du.shape[0]):
        for j in range(du.shape[1]):
            for k in range(du.shape[0]):
                for l in range(du.shape[1]):
                    for a in prange(J.shape[0]):
                        for b in prange(J.shape[1]):
                            out[a, b] += (
                                p[a,b] * J[a,b] * Finv[l,k,a,b] * Finv[j, i, a, b] -\
                                    p[a,b] * J[a,b] * Finv[j,k,a,b] * Finv[l, i, a, b] +\
                                        mu * kron[i,k] * kron[j,l]
                            ) * du[i,j,a,b] * dv[k,l,a,b]

    return out

@jit(nopython=True, cache=True, nogil=True, parallel=True)
def A12(du, dv, dw):
    out = zeros_like(dw[0,0])
    F = zeros_like(dw)
    for a in prange(dw.shape[2]):
        for b in prange(dw.shape[3]):
            F[0, 0, a, b] += 1.
            F[1, 1, a, b] += 1.
            F[2, 2, a, b] += 1.
    
    F += dw
    Finv, J = vinv(F)
    for i in range(dw.shape[0]):
        for j in range(dw.shape[1]):
            for a in prange(J.shape[0]):
                for b in prange(J.shape[1]):
                    out[a, b] += (
                        J[a,b] * Finv[j, i, a, b]
                    ) * dv[i, j, a, b] * du[a, b]
 
    return out

@jit(nopython=True, cache=True, nogil=True, parallel=True)
def A22(du, dv, p):
    Js = Jstar(p)
    dJdp = dJst_dp(p)
    d2Jdp2 = d2Jst_dp2(p)
    
    out = zeros_like(p)
    for a in prange(p.shape[0]):
        for b in prange(p.shape[1]):
            out[a,b] += (-2. * dJdp[a, b] - p[a,b] * d2Jdp2[a,b] + mu/Js[a,b]**2 * dJdp[a,b]**2 - mu/Js[a,b] * d2Jdp2[a,b] + lmbda * (Js[a,b] - 1.) * d2Jdp2[a,b] + lmbda * dJdp[a,b]**2) * du[a,b] * dv[a,b]
    return out
    

mesh = MeshTet()
mesh.refine(3)
uelem = ElementVectorH1(ElementTetP2())
pelem = ElementTetP1()
elems = {
    "u": uelem,
    "p": pelem
}
basis = {
    field: InteriorBasis(mesh, e, intorder=2)
    for field, e in elems.items()
}

du = zeros(basis["u"].N)
dp = zeros(basis["p"].N)
stretch_ = 1.

dofsold = [
    basis["u"].find_dofs({"left":mesh.facets_satisfying(lambda x: x[0] < 1.e-6)}, skip=["u^2", "u^3"]),
    basis["u"].find_dofs({"bottom":mesh.facets_satisfying(lambda x: x[1] < 1.e-6)}, skip=["u^1", "u^3"]),
    basis["u"].find_dofs({"back":mesh.facets_satisfying(lambda x: x[2] < 1.e-6)}, skip=["u^1", "u^2"]),
    basis["u"].find_dofs({"front":mesh.facets_satisfying(lambda x: npabs(x[2]-1.) < 1.e-6 )}, skip=["u^1", "u^2"])
]

dofs = {}
for dof in dofsold:
    dofs.update(dof)

du[dofs["left"].all()] = 0.
du[dofs["bottom"].all()] = 0.
du[dofs["back"].all()] = 0.
du[dofs["front"].all()] = stretch_

I = hstack((
    basis["u"].complement_dofs(dofs),
    basis["u"].N + arange(basis["p"].N)
))

@LinearForm
def a1(v, w):
    return F1(*(grad(v), grad(w["disp"]), w["press"].value))

@LinearForm
def a2(v, w):
    return F2(*(v.value, grad(w["disp"]), w["press"].value))

@BilinearForm
def b11(u, v, w):
    return A11(*(grad(u), grad(v), grad(w["disp"]), w["press"].value))

@BilinearForm
def b12(u, v, w):
    return A12(*(u.value, grad(v), grad(w["disp"])))

@BilinearForm
def b22(u, v, w):
    return A22(u.value, v.value, w["press"].value)

for itr in range(12):
    uv = basis["u"].interpolate(du)
    pv = basis["p"].interpolate(dp) 

    K11 = asm(b11, basis["u"], basis["u"], disp=uv, press = pv)
    K12 = asm(b12, basis["p"], basis["u"], disp=uv, press = pv)
    K22 = asm(b22, basis["p"], basis["p"], disp=uv, press = pv)
    f = concatenate((
        asm(a1, basis["u"], disp=uv, press = pv),
        asm(a2, basis["p"], disp=uv, press = pv)
    ))
    K = bmat(
        [[K11, K12],
         [K12.T, K22]], "csr"
    )
    uvp = solve(*condense(K, -f, I=I), use_umfpack=True)
    delu, delp = npsplit(uvp, [du.shape[0]])
    du += delu
    dp += delp
    normu = nla.norm(delu)
    normp = nla.norm(delp)
    if normu < 1.e-8 and normp < 1.e-8:
        break


import meshio, os
numbaResults = os.path.join(os.getcwd(), "numbaResults")
os.makedirs(numbaResults, exist_ok=True)
filw = os.path.join(numbaResults, "mixedFEMSingle.xdmf")
with meshio.xdmf.TimeSeriesWriter(filw) as writer:
    writer.write_points_cells(points=mesh.p.T, cells={"tetra":mesh.t.T})
    for t in range(2):
        writer.write_data(t, point_data={"u":t*du[basis["u"].nodal_dofs].T,
                                         "p":t*dp[basis["p"].nodal_dofs].T})

[bhaveshshrimali]

A slightly odd thing that I noticed when comparing the timings of the two codes is that unlike #439 the numba version takes significantly longer than the normal version. I will investigate this later this week

Edit2 : prange adds quite a bit of overhead in this case. Removing it makes the results look more plausible

Numba (updated)

         12964617 function calls (11992312 primitive calls) in 65.578 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     8100   36.979    0.005   36.979    0.005 hyperElasticityMixedNumba.py:106(A11)
        9    8.125    0.903    8.125    0.903 {built-in method scikits.umfpack.__umfpack.umfpack_dl_numeric}
    12379    4.113    0.000    4.115    0.000 ffi.py:111(__call__)
     1080    1.351    0.001    1.351    0.001 hyperElasticityMixedNumba.py:134(A12)
    10238    0.960    0.000    0.960    0.000 {method 'reduce' of 'numpy.ufunc' objects}
624981/103873    0.691    0.000    0.965    0.000 ir.py:313(_rec_list_vars)

Normal

         12747915 function calls (12049534 primitive calls) in 103.773 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    35215   35.523    0.001   35.523    0.001 {built-in method numpy.core._multiarray_umath.c_einsum}
     8100   30.023    0.004   72.992    0.009 hyperElasticityMixed.py:77(A11)
     9450   10.274    0.001   10.274    0.001 hyperElasticityMixed.py:32(vinv)
        9    8.233    0.915    8.233    0.915 {built-in method scikits.umfpack.__umfpack.umfpack_dl_numeric}
     3794    2.247    0.001    2.248    0.001 ffi.py:111(__call__)
105356/69656    1.380    0.000   38.159    0.001 {built-in method numpy.core._multiarray_umath.implement_array_function}

Edit: As expected though, assembling the block A11 takes the longest. The system is ~13k dofs

[kinnala]

Do you wish to contribute the non-numba version as an example? There is some work involved such as adding it to the test suite, writing a short description to the docs, taking a nice picture. I'd also like to have those detA and invA routines added to skfem.helpers as generic versions. If you are up for the job, feel free to open a PR and we will support you. I can also do it when I have time if you feel it's too much right now.

[bhaveshshrimali]

Do you wish to contribute the non-numba version as an example? There is some work involved such as adding it to the test suite, writing a short description to the docs, taking a nice picture. I'd also like to have those detA and invA routines added to skfem.helpers as generic versions. If you are up for the job, feel free to open a PR and we will support you. I can also do it when I have time if you feel it's too much right now.

Thanks @kinnala

I'll be happy to contribute this as an example.

I can take a first stab at putting it together this weekend and we can take it from there.

[bhaveshshrimali]

I started working on writing the example. Since scikit-fem does not have sympy as a dependency, do you want me to remove the portions that depend on it? I can write down the expressions explicitly as done for the Numba version above.

I will create a PR sometime tomorrow or monday once I finish documenting it.

[kinnala]

I started working on writing the example. Since scikit-fem does not have sympy as a dependency, do you want me to remove the portions that depend on it? I can write down the expressions explicitly as done for the Numba version above.

If it's easy to do, then that would be nice.

[bhaveshshrimali]

This is a trial figure for the example ? I tried to conform the color map to what is used in the majority of the examples. I can change it based on what you feel is more appropriate and then add it to #431 once everything is done

Let me know.

I have added the inv and det functions to skfem.helpers as well as added a test to compare the dofs to -0.5. The analytical pressure should be given by -mu/l where l is the applied stretch (which in this case is equal to 2). Do you feel anything else could be a better test (like computing the deformed volume ?)

[bhaveshshrimali]

I would close this and continue with some experiments on #437