Fix Guzman-Neilan bubbles

scripts/generate_gn.py

@@ -4,52 +4,172 @@
 import sys
 import typing
 
-import numpy as np
 import sympy
 
 import symfem
-from symfem.elements.guzman_neilan import make_piecewise_lagrange
-from symfem.functions import MatrixFunction, VectorFunction
+from symfem.functions import VectorFunction, AnyFunction
+from symfem.piecewise_functions import PiecewiseFunction
 from symfem.symbols import t, x
 
 TESTING = "test" in sys.argv
 
-line_length = 100
-if TESTING:
-    folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../_temp")
-else:
-    folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../symfem/elements")
 
+def poly(reference, k):
+    """Generate the P^perp polynomial set."""
+    if k < 2:
+        if reference.name == "triangle":
+            return [
+                VectorFunction(
+                    [x[0] ** i0 * x[1] ** i1 if d2 == d else 0 for d2 in range(reference.tdim)]
+                )
+                for d in range(reference.tdim)
+                for i0 in range(k + 1)
+                for i1 in range(k + 1 - i0)
+            ]
+        else:
+            assert reference.name == "tetrahedron"
+            return [
+                VectorFunction(
+                    [
+                        x[0] ** i0 * x[1] ** i1 * x[2] ** i2 if d2 == d else 0
+                        for d2 in range(reference.tdim)
+                    ]
+                )
+                for d in range(reference.tdim)
+                for i0 in range(k + 1)
+                for i1 in range(k + 1 - i0)
+                for i2 in range(k + 1 - i0 - i1)
+            ]
+    if k == 2:
+        assert reference.name == "tetrahedron"
 
-def find_solution(mat, aim):
-    """Solve matrix-vector problem."""
-    A_data = [[float(j) for j in i] for i in mat]
-    b_data = [float(i) for i in aim]
-    A = np.asarray(A_data, dtype=np.float64)
-    b = np.asarray(b_data)
+        poly = [
+            VectorFunction(
+                [
+                    x[0] ** i0 * x[1] ** i1 * x[2] ** i2 if d2 == d else 0
+                    for d2 in range(reference.tdim)
+                ]
+            )
+            for d in range(reference.tdim)
+            for i0 in range(k + 1)
+            for i1 in range(k + 1 - i0)
+            for i2 in range(k + 1 - i0 - i1)
+        ]
 
-    res = np.linalg.solve(A, b)
+        poly[1] -= poly[0] / 4
+        poly[2] -= poly[0] / 10
+        poly[3] -= poly[0] / 4
+        poly[4] -= poly[0] / 20
+        poly[5] -= poly[0] / 10
+        poly[6] -= poly[0] / 4
+        poly[7] -= poly[0] / 20
+        poly[8] -= poly[0] / 20
+        poly[9] -= poly[0] / 10
+        poly = poly[1:]
 
-    fractions = []
-    for i in res:
-        frac = sympy.Rational(int(round(i * 162000)), 162000)
-        assert i < 10
-        assert np.isclose(float(frac), i)
-        fractions.append(frac)
+        poly[10] -= poly[9] / 4
+        poly[11] -= poly[9] / 10
+        poly[12] -= poly[9] / 4
+        poly[13] -= poly[9] / 20
+        poly[14] -= poly[9] / 10
+        poly[15] -= poly[9] / 4
+        poly[16] -= poly[9] / 20
+        poly[17] -= poly[9] / 20
+        poly[18] -= poly[9] / 10
+        poly = poly[:9] + poly[10:]
 
-    assert MatrixFunction(mat) @ VectorFunction(fractions) == VectorFunction(aim)
+        poly[19] -= poly[18] / 4
+        poly[20] -= poly[18] / 10
+        poly[21] -= poly[18] / 4
+        poly[22] -= poly[18] / 20
+        poly[23] -= poly[18] / 10
+        poly[24] -= poly[18] / 4
+        poly[25] -= poly[18] / 20
+        poly[26] -= poly[18] / 20
+        poly[27] -= poly[18] / 10
+        poly = poly[:18] + poly[19:]
 
-    return fractions
+        poly[1] -= poly[0] * 2 / 3
+        poly[2] += poly[0] / 3
+        poly[3] -= poly[0] / 9
+        poly[4] += poly[0] * 2 / 9
+        poly[5] += poly[0] / 3
+        poly[6] -= poly[0] / 9
+        poly[7] += poly[0] / 9
+        poly[8] += poly[0] * 2 / 9
+        poly[18] -= poly[0] / 3
+        poly[19] -= poly[0] * 2 / 9
+        poly[20] -= poly[0] / 3
+        poly[21] -= poly[0] / 9
+        poly[22] -= poly[0] * 2 / 9
+        poly[23] += poly[0]
+        poly[24] += poly[0] / 9
+        poly[25] += poly[0] / 9
+        poly[26] += poly[0] * 2 / 3
+        poly = poly[1:]
 
+        poly[9] -= poly[8] * 2 / 3
+        poly[10] += poly[8] / 3
+        poly[11] -= poly[8] / 9
+        poly[12] += poly[8] * 2 / 9
+        poly[13] += poly[8] / 3
+        poly[14] -= poly[8] / 9
+        poly[15] += poly[8] / 9
+        poly[16] += poly[8] * 2 / 9
+        poly[17] -= poly[8] / 3
+        poly[18] -= poly[8] * 2 / 9
+        poly[19] += poly[8]
+        poly[20] += poly[8] / 9
+        poly[21] += poly[8] * 2 / 3
+        poly[22] -= poly[8] / 3
+        poly[23] -= poly[8] / 9
+        poly[24] += poly[8] / 9
+        poly[25] -= poly[8] * 2 / 9
+        poly = poly[:8] + poly[9:]
 
-for ref in ["triangle", "tetrahedron"]:
-    # TODO: work out why tetrahedron fails on github but passes locally
-    if TESTING and ref == "tetrahedron":
-        break
+        poly[2] -= poly[1] / 6
+        poly[3] -= poly[1] * 2 / 3
+        poly[4] += poly[1] / 2
+        poly[5] += poly[1] / 12
+        poly[6] -= poly[1] / 12
+        poly[7] += poly[1] / 3
+        poly[9] -= poly[1] / 2
+        poly[10] -= poly[1] / 12
+        poly[11] -= poly[1] / 3
+        poly[12] += poly[1]
+        poly[13] += poly[1] / 6
+        poly[14] += poly[1] / 12
+        poly[15] += poly[1] * 2 / 3
+        poly[18] -= poly[1] / 2
+        poly[19] -= poly[1] / 12
+        poly[20] -= poly[1] / 3
+        poly[21] += poly[1] / 2
+        poly[22] += poly[1] / 12
+        poly[24] += poly[1] / 3
+        poly = poly[:1] + poly[2:]
 
+        return poly
+
+    raise NotImplementedError()
+
+
+def find_solution(mat, aim):
+    """Solve mat @ x = aim."""
+    s_mat = sympy.Matrix(mat)
+    solution = (s_mat.T @ s_mat).inv() @ s_mat.T @ sympy.Matrix(aim)
+    assert s_mat @ solution == sympy.Matrix(aim)
+    return list(solution)
+
+
+if TESTING:
+    folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../_temp")
+else:
+    folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../symfem/elements")
+
+for ref in ["triangle", "tetrahedron"]:
     reference = symfem.create_reference(ref)
     br = symfem.create_element(ref, "Bernardi-Raugel", 1)
-    mid = tuple(sympy.Rational(sum(i), len(i)) for i in zip(*reference.vertices))
+    mid = reference.midpoint()
 
     sub_cells: typing.List[symfem.geometry.SetOfPoints] = []
 
@@ -62,6 +182,16 @@ def find_solution(mat, aim):
         ]
         xx, yy, zz = x
         terms = [1, xx, yy]
+
+        lamb = PiecewiseFunction(
+            {sub_cells[0]: 3 * x[1], sub_cells[1]: 3 * x[0], sub_cells[2]: 3 * (1 - x[0] - x[1])}, 2
+        )
+
+        for c, f in lamb.pieces.items():
+            assert f.subs(x, c[0]) == 0
+            assert f.subs(x, c[1]) == 0
+            assert f.subs(x, c[2]) == 1
+
     if ref == "tetrahedron":
         fs = br.get_basis_functions()[-4:]
         sub_cells = [
@@ -73,16 +203,37 @@ def find_solution(mat, aim):
         xx, yy, zz = x
         terms = [1, xx, yy, zz, xx**2, yy**2, zz**2, xx * yy, xx * zz, yy * zz]
 
-    sub_basis = make_piecewise_lagrange(sub_cells, ref, br.reference.tdim, True)
+        lamb = PiecewiseFunction(
+            {
+                sub_cells[0]: 4 * x[2],
+                sub_cells[1]: 4 * x[1],
+                sub_cells[2]: 4 * x[0],
+                sub_cells[3]: 4 * (1 - x[0] - x[1] - x[2]),
+            },
+            3,
+        )
+
+        for c, f in lamb.pieces.items():
+            assert f.subs(x, c[0]) == 0
+            assert f.subs(x, c[1]) == 0
+            assert f.subs(x, c[2]) == 0
+            assert f.subs(x, c[3]) == 1
+
+    sub_basis = [
+        p * lamb**j
+        for j in range(1, reference.tdim + 1)
+        for p in poly(reference, reference.tdim - j)
+    ]
 
     filename = os.path.join(folder, f"_guzman_neilan_{ref}.py")
-    output = '"""Values for Guzman-Neilan element."""\n' "\n" "import sympy\n" "\n" "coeffs = [\n"
+    output = '"""Values for Guzman-Neilan element."""\n\n'
+    output += "import sympy\n\nbubbles = [\n"
 
     for f in fs:
         assert isinstance(f, VectorFunction)
-        output += "    [\n"
         integrand = f.div().subs(x, t)
         fun_s = (f.div() - integrand.integral(reference) / reference.volume()).as_sympy()
+
         assert isinstance(fun_s, sympy.core.expr.Expr)
         fun = fun_s.as_coefficients_dict()
 
@@ -104,53 +255,45 @@ def find_solution(mat, aim):
                     assert term == 0 or term in terms
                 for term in terms:
                     if i < len(mat):
-                        mat[i].append(d[term] if term in d else 0)
+                        mat[i].append(d[term] if term in d else 0)  # type: ignore
                     i += 1
 
-        if ref == "triangle":
-            mat = mat[:-1]
-            aim = aim[:-1]
-            fractions = find_solution(mat, aim)
-        if ref == "tetrahedron":
-            for i in range(3):
-                row: typing.List[symfem.functions.ScalarFunction] = [sympy.Integer(0)] * 45
-                row[i] = sympy.Integer(1)
-                mat.append(row)
-            subf = f.subs(x, mid)
-            assert isinstance(subf, VectorFunction)
-            aim += [i for i in subf]
-            aim += [0, 0]
-
-            fractions = None
-            for n in range(3, 45):
-                for m in range(n + 1, 45):
-                    mat2 = [i for i in mat]
-                    for i in [n, m]:
-                        row = [sympy.Integer(0)] * 45
-                        row[i] = sympy.Integer(1)
-                        mat2.append(row)
-                    try:
-                        fractions = find_solution(mat2, aim)
-                        break
-                    except AssertionError:
-                        pass
-                else:
-                    continue
-                break
-            assert fractions is not None
-
-        line = " " * 7
-        for frac in fractions:
-            if frac.denominator == 1:
-                next = f" {frac.numerator},"
-            else:
-                next = f" sympy.Rational({frac.numerator}, {frac.denominator}),"
-            if len(line + next) > line_length:
-                output += line + "\n"
-                line = " " * 7 + next
-            else:
-                line += next
-        output += line + "\n    ],\n"
+        coeffs = find_solution(mat, aim)
+        bubble: AnyFunction = f
+        for i, j in zip(coeffs, sub_basis):
+            bubble -= i * j
+
+        assert isinstance(bubble, PiecewiseFunction)
+
+        output += "    {\n"
+        for cell, f in bubble.pieces.items():
+            output += "        "
+            output += (
+                "("
+                + ", ".join(
+                    [
+                        "("
+                        + ", ".join(
+                            [
+                                f"{c}" if isinstance(c, sympy.Integer) else f"sympy.S('{c}')"
+                                for c in p
+                            ]
+                        )
+                        + ")"
+                        for p in cell
+                    ]
+                )
+                + ")"
+            )
+            output += ": (\n"
+            output += ",\n".join(
+                [
+                    "            sympy.S('" + f"{c.as_sympy().expand()}".replace(" ", "") + "')"  # type:ignore
+                    for c in f
+                ]
+            )
+            output += "),\n"
+        output += "    },\n"
 
     output += "]\n"
 

symfem/basis_functions.py

@@ -8,7 +8,6 @@
 import sympy
 
 import symfem
-
 from symfem.functions import (
     AnyFunction,
     FunctionInput,