Custom Degree 2 vector Lagrange on a triangle

Hi, everyone.

I have tried to create a custom Degree 2 vector Lagrange on a triangle and use it for a 2D linear elastic problem.

code

import numpy as np
import basix
from basix import CellType, MapType, PolynomialType, PolysetType, SobolevSpace
from mpi4py import MPI
from basix.ufl import element, mixed_element
from ufl import (Identity, TestFunction, TrialFunction, 
                 dx, inner, grad, sym, tr)
from dolfinx import fem, mesh, io, plot
from dolfinx.fem.petsc import LinearProblem
from petsc4py import PETSc

import dolfinx
import ufl
from dolfinx.fem.petsc import apply_lifting, assemble_matrix, assemble_vector, set_bc


# Set coefficients for the basis functions
wcoeffs = np.zeros((12, 12))
# Set coefficients for the basis functions
wcoeffs[0, 0] = 1   # (1,0)
wcoeffs[1, 6] = 1   # (0,1)
wcoeffs[2, 3] = 1   # (x,0)
wcoeffs[3, 9] = 1   # (0,x)
wcoeffs[4, 4] = 1   # (x^2,0)
wcoeffs[5, 10] = 1   # (0,x^2)
wcoeffs[6, 1] = 1   # (y,0)
wcoeffs[7, 7] = 1   # (0,y)
wcoeffs[8, 5] = 1   # (xy,0)
wcoeffs[9, 11] = 1   # (0,xy)
wcoeffs[10, 2] = 1 # (y^2,0)
wcoeffs[11, 8] = 1 # (0,y^2)

# Interpolation points
# DOFs are point evaluations at vertices: (0,0), (1,0), (0,1)
# and edge midpoints: (0.5,0), (0.5,0.5), (0,0.5) for x- and y-components.
x = [[], [], [], []]
# Vertices
x[0].append(np.array([[0.0, 0.0]]))  # Vertex 0
x[0].append(np.array([[1.0, 0.0]]))  # Vertex 1
x[0].append(np.array([[0.0, 1.0]]))  # Vertex 2
# Edge midpoints
x[1].append(np.array([[0.5, 0.0]]))  # Edge 0: midpoint between (0,0) and (1,0)
x[1].append(np.array([[0.5, 0.5]]))  # Edge 1: midpoint between (1,0) and (0,1)
x[1].append(np.array([[0.0, 0.5]]))  # Edge 2: midpoint between (0,1) and (0,0)
# Empty for cell interior
x[2].append(np.zeros((0, 2)))

# Interpolation matrices
# Each point (vertex or edge midpoint) has two DOFs: x-component (v . (1,0)) and y-component (v . (0,1)).
# Shape per point: (2, 2, 1, 1) for 2 DOFs, value size 2, 1 point, 0 derivatives.
M = [[], [], [], []]
for _ in range(3):  # One matrix per vertex
    mat = np.zeros((2, 2, 1, 1))
    mat[0, :, 0, 0] = [1.0, 0.0]  # x-component: v . (1,0)
    mat[1, :, 0, 0] = [0.0, 1.0]  # y-component: v . (0,1)
    M[0].append(mat)
for _ in range(3):  # One matrix per edge midpoint
    mat = np.zeros((2, 2, 1, 1))
    mat[0, :, 0, 0] = [1.0, 0.0]  # x-component: v . (1,0)
    mat[1, :, 0, 0] = [0.0, 1.0]  # y-component: v . (0,1)
    M[1].append(mat)
M[2].append(np.zeros((0, 2, 0, 1)))  # Empty for cell interior

# Create the custom element
basix_element = basix.create_custom_element(
    CellType.triangle,           # Cell type
    [2],                        # Value shape (vector in 2D)
    wcoeffs,                    # Polynomial coefficients
    x,                          # Interpolation points
    M,                          # Interpolation matrices
    0,                          # Number of derivatives (point evaluations, so 0)
    MapType.identity,           # Identity map for Lagrange elements
    SobolevSpace.H1,            # H^1-conforming
    False,                      # Not discontinuous
    2,                          # Embedded subdegree (contains all degree 2 polynomials)
    2,                          # Embedded superdegree (highest degree is 2)
    PolysetType.standard,       # Standard polynomial set
)



ufl_el = basix.ufl.wrap_element(basix_element)


# Create mesh
L = 25.0
H = 1.0
Nx = 40
Ny = 2
msh = mesh.create_rectangle(
    MPI.COMM_WORLD, 
    [[0.0, 0.0], [L, H]], 
    [Nx, Ny], 
    mesh.CellType.triangle
)


#ufl_element = element(basix_element)

V = fem.functionspace(msh, ufl_el)

# Define material properties
E = 1e5
nu = 0.3
model = "plane_stress"

mu = E / 2 / (1 + nu)
lmbda = E * nu /( (1 + nu) *(1 - 2 * nu))
if model == "plane_stress":
    lmbda = 2 * mu * lmbda / (lmbda + 2 * mu)


# Define strain and stress functions
def eps(v):
    
    return sym(grad(v))

def sigma(v):
    return lmbda * tr(eps(v)) * Identity(2) + 2.0 * mu * eps(v)

# Define body force
rho_g = 1e-3
f = fem.Constant(msh, (0.0, -rho_g))

# Define variational problem
u = TrialFunction(V)
v = TestFunction(V)
a = inner(sigma(u), eps(v)) * dx
L = inner(f, v) * dx

# Define boundary condition
def left(x):
    return np.isclose(x[0], 0.0)

fdim = msh.topology.dim - 1
left_facets = mesh.locate_entities_boundary(msh, fdim, left)
left_dofs = fem.locate_dofs_topological(V, fdim, left_facets)
zero = fem.Function(V)
zero.x.array[:] = 0.0
bc = fem.dirichletbc(zero, left_dofs)

# Solve problem
problem = LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
uh = problem.solve()
uh.name = "Displacement"

def project(e, target_func, bcs=[]):
    """Project UFL expression.

    Note
    ----
    This method solves a linear system (using KSP defaults).

    """


    # Ensure we have a mesh and attach to measure
    V = target_func.function_space
    dx = ufl.dx(V.mesh)

    # Define variational problem for projection
    w = ufl.TestFunction(V)
    v = ufl.TrialFunction(V)
    a = dolfinx.fem.form(ufl.inner(v, w) * dx)
    L = dolfinx.fem.form(ufl.inner(e, w) * dx)

    # Assemble linear system
    A = assemble_matrix(a, bcs)
    A.assemble()
    b = assemble_vector(L)
    apply_lifting(b, [a], [bcs])
    b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b, bcs)

    # Solve linear system
    solver = PETSc.KSP().create(A.getComm())
    solver.setType("bcgs")
    solver.getPC().setType("bjacobi")
    solver.rtol = 1.0e-05
    solver.setOperators(A)
    solver.solve(b, target_func.x.petsc_vec)
    assert solver.reason > 0
    target_func.x.scatter_forward()
    

    # Destroy PETSc linear algebra objects and solver
    solver.destroy()
    A.destroy()
    b.destroy()

V_vector = fem.functionspace(msh, ("DG", 1, (2,)))
displacement = fem.Function(V_vector)
project(uh, displacement)
print("Solution array contains NaN:", np.any(np.isnan(uh.x.array)))
print("Solution array contains inf:", np.any(np.isinf(uh.x.array)))
print("Solution array max/min:", np.max(np.abs(uh.x.array)), np.min(np.abs(uh.x.array)))
# from dolfinx.io import VTXWriter
# outputfoldername = "results_2D_custom_triangle_p1/"
# u_bp = VTXWriter(MPI.COMM_WORLD, f"{outputfoldername}displacement.bp", [displacement], engine="BP4")
# u_bp.write(0.0)
# u_bp.close()
print("Maximum displacement:", np.max(np.linalg.norm(uh.x.array.reshape(-1, 2), axis=1)))

output:

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
File ~/FEniCSx/unbenannt2.py:209
    207 V_vector = fem.functionspace(msh, ("DG", 1, (2,)))
    208 displacement = fem.Function(V_vector)
--> 209 project(uh, displacement)
    210 print("Solution array contains NaN:", np.any(np.isnan(uh.x.array)))
    211 print("Solution array contains inf:", np.any(np.isinf(uh.x.array)))

File ~/FEniCSx/unbenannt2.py:198, in project(e, target_func, bcs)
    196 solver.setOperators(A)
    197 solver.solve(b, target_func.x.petsc_vec)
--> 198 assert solver.reason > 0
    199 target_func.x.scatter_forward()
    202 # Destroy PETSc linear algebra objects and solver

AssertionError: 

When i tried to comment the projection and run the code to check the displacement value, i got the displacement value as inf.

The code ran fine when tried with custom degree 1 vector lagrange.

any suggestions would be appreciated.

Thanks!

I would take a step back starting with a custom degree 1 vector Lagrange on a triangle, and verify its correctness through interpolation:

import numpy as np
import basix
from basix import CellType, MapType, PolysetType, SobolevSpace
from mpi4py import MPI
from dolfinx import fem, mesh
import dolfinx


# Set coefficients for the basis functions
wcoeffs = np.zeros((6, 6))
# Set coefficients for the basis functions
wcoeffs[0, 0] = 1   # (1,0)
wcoeffs[1, 3] = 1   # (0,1)
wcoeffs[2, 2] = 1   # (x,0)
wcoeffs[3, 5] = 1   # (0,x)
wcoeffs[4, 1] = 1   # (y,0)
wcoeffs[5, 4] = 1   # (0,y)

# Interpolation points
# DOFs are point evaluations at vertices: (0,0), (1,0), (0,1)
x = [[], [], [], []]
# Vertices
x[0].append(np.array([[0.0, 0.0]]))  # Vertex 0
x[0].append(np.array([[1.0, 0.0]]))  # Vertex 1
x[0].append(np.array([[0.0, 1.0]]))  # Vertex 2
for _ in range(3):
    x[1].append(np.zeros((0, 2)))
x[2].append(np.zeros((0, 2)))

# Interpolation matrices
# Each point (vertex or edge midpoint) has two DOFs: x-component (v . (1,0)) and y-component (v . (0,1)).
# Shape per point: (2, 2, 1, 1) for 2 DOFs, value size 2, 1 point, 0 derivatives.
M = [[], [], [], []]
for _ in range(3):  # One matrix per vertex
    mat = np.zeros((2, 2, 1, 1))
    mat[0, :, 0, 0] = [1.0, 0.0]  # x-component: v . (1,0)
    mat[1, :, 0, 0] = [0.0, 1.0]  # y-component: v . (0,1)
    M[0].append(mat)

for _ in range(3):
    M[1].append(np.zeros((0, 2, 0, 1)))
M[2].append(np.zeros((0, 2, 0, 1)))  # Empty for cell interior

# Create the custom element
basix_element = basix.create_custom_element(
    CellType.triangle,           # Cell type
    [2],                        # Value shape (vector in 2D)
    wcoeffs,                    # Polynomial coefficients
    x,                          # Interpolation points
    M,                          # Interpolation matrices
    0,                          # Number of derivatives (point evaluations, so 0)
    MapType.identity,           # Identity map for Lagrange elements
    SobolevSpace.H1,            # H^1-conforming
    False,                      # Not discontinuous
    1,                          # Embedded subdegree (contains all degree 2 polynomials)
    1,                          # Embedded superdegree (highest degree is 2)
    PolysetType.standard,       # Standard polynomial set
)





def solve_problem(mesh, el):
    # Create mesh
    L = 25.0
    H = 1.0
    Nx = 40
    Ny = 2
    msh = mesh.create_rectangle(
        MPI.COMM_WORLD, 
        [[0.0, 0.0], [L, H]], 
        [Nx, Ny], 
        mesh.CellType.triangle
    )

    V = fem.functionspace(msh, el)

    def f(x):
        return -x[1], x[0]

    u = dolfinx.fem.Function(V, name=f"u+{el!r}")
    u.interpolate(lambda x: f(x))

    print(u.name)
    return u


el = basix.ufl.element("Lagrange", CellType.triangle, 1, shape=(2,))
ref_sol = solve_problem(mesh, el)


ufl_el = basix.ufl.wrap_element(basix_element)
custom_sol = solve_problem(mesh, ufl_el)
np.testing.assert_allclose(ref_sol.x.array, custom_sol.x.array)

Then, carefully extending this to degree two yields the following:

import numpy as np
import basix
from basix import CellType, MapType, PolysetType, SobolevSpace
from mpi4py import MPI
from dolfinx import fem, mesh
import dolfinx


# Set coefficients for the basis functions
wcoeffs = np.zeros((12, 12))
# Set coefficients for the basis functions
# Basix polyset ordering: 1, y, x, y^2, xy, x^2
# https://docs.fenicsproject.org/basix/main/polyset-order.html
wcoeffs[0, 0] = 1   # (1,0)
wcoeffs[1, 6] = 1   # (0,1)
wcoeffs[2, 2] = 1   # (x,0)
wcoeffs[3, 8] = 1   # (0,x)
wcoeffs[4, 1] = 1   # (y,0)
wcoeffs[5, 7] = 1   # (0,y)
wcoeffs[6, 5] = 1   # (x^2,0)
wcoeffs[7, 11] = 1 # (0,x^2)
wcoeffs[8, 4] = 1   # (xy,0)
wcoeffs[9, 10] = 1   # (0,xy)
wcoeffs[10, 3] = 1  # (y^2,0)
wcoeffs[11, 9] = 1  # (0,y^2)

# Interpolation points
# DOFs are point evaluations at vertices: (0,0), (1,0), (0,1)
x = [[], [], [], []]
# Vertices
x[0].append(np.array([[0.0, 0.0]]))  # Vertex 0
x[0].append(np.array([[1.0, 0.0]]))  # Vertex 1
x[0].append(np.array([[0.0, 1.0]]))  # Vertex 2
x[1].append(np.array([[0.5, 0.5]]))  # Edge 0: midpoint between (0,0) and (1,0)
x[1].append(np.array([[0.0, 0.5]]))  # Edge 1: midpoint between (1,0) and (0,1)
x[1].append(np.array([[0.5, 0.0]]))  # Edge 2: midpoint between (0,1) and (0,0)

x[2].append(np.zeros((0, 2)))

# Interpolation matrices
# Each point (vertex or edge midpoint) has two DOFs: x-component (v . (1,0)) and y-component (v . (0,1)).
# Shape per point: (2, 2, 1, 1) for 2 DOFs, value size 2, 1 point, 0 derivatives.
M = [[], [], [], []]
for _ in range(3):  # One matrix per vertex
    mat = np.zeros((2, 2, 1, 1))
    mat[0, :, 0, 0] = [1.0, 0.0]  # x-component: v . (1,0)
    mat[1, :, 0, 0] = [0.0, 1.0]  # y-component: v . (0,1)
    M[0].append(mat)

for _ in range(3):  # One matrix per edge midpoint
    mat = np.zeros((2, 2, 1, 1))
    mat[0, :, 0, 0] = [1.0, 0.0]  # x-component: v . (1,0)
    mat[1, :, 0, 0] = [0.0, 1.0]  # y-component: v . (0,1)
    M[1].append(mat)
M[2].append(np.zeros((0, 2, 0, 1)))  # Empty for cell interior

# Create the custom element
basix_element = basix.create_custom_element(
    CellType.triangle,           # Cell type
    [2],                        # Value shape (vector in 2D)
    wcoeffs,                    # Polynomial coefficients
    x,                          # Interpolation points
    M,                          # Interpolation matrices
    0,                          # Number of derivatives (point evaluations, so 0)
    MapType.identity,           # Identity map for Lagrange elements
    SobolevSpace.H1,            # H^1-conforming
    False,                      # Not discontinuous
    2,                          # Embedded subdegree (contains all degree 2 polynomials)
    2,                          # Embedded superdegree (highest degree is 2)
    PolysetType.standard,       # Standard polynomial set
)




def solve_problem(mesh, el):
    # Create mesh
    L = 25.0
    H = 1.0
    Nx = 40
    Ny = 2
    msh = mesh.create_rectangle(
        MPI.COMM_WORLD, 
        [[0.0, 0.0], [L, H]], 
        [Nx, Ny], 
        mesh.CellType.triangle
    )

    V = fem.functionspace(msh, el)

    def f(x):
        return -x[1]**2, x[0]*x[1]

    u = dolfinx.fem.Function(V, name=f"u+{el!r}")
    u.interpolate(lambda x: f(x))
    return u


el = basix.ufl.element("Lagrange", CellType.triangle, 2, shape=(2,))
ref_sol = solve_problem(mesh, el)


ufl_el = basix.ufl.wrap_element(basix_element)

custom_sol = solve_problem(mesh, ufl_el)
np.testing.assert_allclose(ref_sol.x.array, custom_sol.x.array)

Note a few things:

  1. Order of polysets in my code differs from your code.
  2. The order of the edges differ, see basix documentation for this:

Thanks for your reply and clarification.

I am bit confused in the ordering of polynomials sets. In your code you have the order of (1, x, y, x², xy, y²), but in the DefElement, the span of “V” is defined as of the order (1, x, x², y, xy, y²).

[quote=“dokken, post:2, topic:18227”]

# https://docs.fenicsproject.org/basix/main/polyset-order.html
wcoeffs[0, 0] = 1   # (1,0)
wcoeffs[1, 6] = 1   # (0,1)
wcoeffs[2, 2] = 1   # (x,0)
wcoeffs[3, 8] = 1   # (0,x)
wcoeffs[4, 1] = 1   # (y,0)
wcoeffs[5, 7] = 1   # (0,y)
wcoeffs[6, 5] = 1   # (x^2,0)
wcoeffs[7, 11] = 1 # (0,x^2)
wcoeffs[8, 4] = 1   # (xy,0)
wcoeffs[9, 10] = 1   # (0,xy)
wcoeffs[10, 3] = 1  # (y^2,0)
wcoeffs[11, 9] = 1  # (0,y^2)

also in the interpolation points, your ordering of edges is different as compared with the figure shown here.

can you give me a clarification on the ordering.

Thanks!

You need to look at the order of polysets in basix when you are implementing an element in basix, as pointed out in:

On the second comment

The only difference is that I did not move the comments around.
As you can see from the actual numbers, Edge 0 should be the midpoint between (1,0), (0,1), i.e.

x[1].append(np.array([[0.5, 0.5]]))  # Edge 0: midpoint between (1,0) and (0,1)
x[1].append(np.array([[0.0, 0.5]]))  # Edge 1: midpoint between (0,0) and (0,1)
x[1].append(np.array([[0.5, 0.0]]))  # Edge 2: midpoint between (0,0) and (1,0)