Question about scaling of bubble enriched basis

Edit so that bottom line is clear: How does ufl/basix scale bubble functions in enriched elements?

I am wondering about how ufl and basix go about the creation of bubble enriched elements, specifically how the bubble basis is scaled to preserve the partition of unity. Consider a P2+B3 element. The test below passes for the standard P2 basis, but when I add and scale the bubble function, I no longer get what dolfin outputs. I assume my scaling is incorrect, but https://defelement.com doesn’t have this element and I can’t seem to find the relevant source code in either ufl or basix to compare. Thank you!

import numpy as np
import dolfin as d

mesh = d.UnitTriangleMesh.create()
P2 = d.FiniteElement("P", d.triangle, degree=2)
bubble = d.FiniteElement("B", d.triangle, degree=3)
VE = d.NodalEnrichedElement(P2, bubble)
V = d.FunctionSpace(mesh, P2)
W = d.FunctionSpace(mesh, VE)

def P2_basis(pt):
    # barycentric coords
    L2 = pt[:,0]
    L3 = pt[:,1]
    L1 = (1 - L2 - L3)
    
    basis = np.array([2*L1**2 - L1,
                      2*L2**2 - L2, 
                      2*L3**2 - L3,
                      4*L2*L3,
                      4*L1*L3,
                      4*L1*L2])
    return basis

def P2_B3_basis(pt):
    L2 = pt[:,0]
    L3 = pt[:,1]
    L1 = (1 - L2 - L3) 
    
    basis = P2_basis(pt)
    basis.resize(7,1) 
    B3 = 27*L1*L2*L3 # bubble
    basis[-1] = B3
    scale = (np.array([-1/6., -1/6., -1/6., -1/6., -1/6., -1/6., 0.])*B3).reshape(7,1)
    scaled_basis = basis + scale
    return scaled_basis

for i in range(10):
    x = np.random.random(1)/3
    y = np.random.random(1)/3
    pt = np.array([x,y]).T
    my_P2 = P2_basis(pt)
    my_P2B3 = P2_B3_basis(pt)
    dolfin_vals_P2 = []
    dolfin_vals_P2B3 = []
    cell = d.Cell(mesh, 0)
    for j in range(6):
        dolfin_vals_P2.append(V.element().evaluate_basis(j, pt, cell.get_vertex_coordinates(), cell.orientation())[0])
    np.testing.assert_almost_equal(np.array([dolfin_vals_P2]).T, my_P2)  # OK
    for j in range(7):
        dolfin_vals_P2B3.append(W.element().evaluate_basis(j, pt, cell.get_vertex_coordinates(), cell.orientation())[0])
    np.testing.assert_almost_equal(np.array([dolfin_vals_P2B3]).T, my_P2B3)  # NOT OK

This element is on DefElement under “bubble enriched Lagrange”: DefElement edit: sorry that’s P2 + B4. I’ve opened an issue for adding P2 + B3: Add P2 + B3 element · Issue #273 · mscroggs/defelement.com · GitHub

Currently, Basix doesn’t ensure that enriched elements satisfy partition of unity. Is this a necessary requirement in your application?

There’s an open issue in Basix about improving the enriched element interface (Handle EnrichedElements better · Issue #633 · FEniCS/basix · GitHub), so I’ll add a note about this there. This issue already suggests adding an equivalent of NodalEnrichedElement so that might already be what you’re after

@mscroggs Hi Matthew, thanks for looking into this. To give more context, I actually need the P2+B3 basis functions in explicit form as in my code above, and I do need them to match the output of legacy dolfin NodalEnrichedElement, which does seem satisfy to the partition of unity as far as I can tell. (This is all part of a larger JAX.numpy program that generates inputs to my PDE that I solve with fenics.) So I guess my question is really just how NodalEnrichedElement constructs that basis and what those polynomials are