Expanding on vectors and simplifying/speeding up forward for directsumtps (#34)

* Expanding on vectors and simplifying/speeding up forward for directsumtps

* Remove unused fixtures from test_to_from_orthogonal

* Increase tolerance for 3D to_from_orthogonal test

* Caching inverse of stencil matrix when computed the first time

Author: mikaem

examples/poisson1D.py

@@ -12,8 +12,6 @@
 from jaxfun.galerkin.Chebyshev import Chebyshev as space
 from jaxfun.galerkin.functionspace import FunctionSpace
 from jaxfun.galerkin.inner import inner
-
-# from jaxfun.Jacobi import Jacobi as space
 from jaxfun.operators import Div, Grad
 from jaxfun.utils.common import lambdify, n, ulp
 
@@ -28,7 +26,8 @@
 ue = 1 - x**2  # * sp.exp(sp.cos(2 * sp.pi * x))
 
 # A = inner(v*sp.Derivative(u, x, 2), sparse=True)
-# A = inner(-Dot(Grad(v), Grad(u)), sparse=True)
+# A = inner(-Dot(sp.sqrt(1-x**2)*Grad(v/sp.sqrt(1-x**2)), Grad(u)), sparse=True) # Cheb
+# A = inner(-Dot(Grad(v), Grad(u)), sparse=True) # Legendre
 # A = inner(v*Div(Grad(u)), sparse=True)
 # b = inner(v*sp.Derivative(ue, x, 2))
 A, b = inner(

src/jaxfun/galerkin/Chebyshev.py

@@ -270,7 +270,7 @@ def derivative_coeffs(self, c: Array, k: int = 0) -> Array:
         N: int = c.shape[0] - 1
         x0: Array = jnp.array(0.0, dtype=float)
         if N == 0:
-            return x0
+            return jnp.array([x0])
         x1: Array = c[-1] * N * 2
         if N == 1:
             return jnp.array([x1, x0])
@@ -282,7 +282,7 @@ def inner_loop(
             x2 = 2 * (n + 1) * c[n + 1] + x0
             return (x1, x2), x2
 
-        _, xs = jax.lax.scan(inner_loop, (x0, x1), jnp.arange(N - 2, -1, -1))
+        xs = jax.lax.scan(inner_loop, (x0, x1), jnp.arange(N - 2, -1, -1))[1]
         return jnp.concatenate(
             (jnp.array([xs[-1] / 2]), xs[-2::-1], jnp.array([x1, x0]))
         )

src/jaxfun/galerkin/Legendre.py

@@ -168,7 +168,7 @@ def derivative_coeffs(self, c: Array, k: int = 0) -> Array:
         N: int = c.shape[0] - 1
         x0: Array = jnp.array(0.0)
         if N == 0:
-            return x0
+            return jnp.array([x0])
         x1: Array = c[-1] * (2 * N - 1)
         if N == 1:
             return jnp.array([x1, x0])

src/jaxfun/galerkin/arguments.py

@@ -786,48 +786,11 @@ def evaluate_mesh(self, x: Array | list[Array]) -> Array:
         return self.functionspace.evaluate_mesh(x, self.array, True)
 
 
-def evaluate_jaxfunction_expr(
-    a: Basic, xj: Array | tuple[Array, ...], jaxf: JAXFunction | None = None
-) -> Array:
-    """Evaluate a symbolic JAXFunction expression on a mesh.
-
-    Input coordinates ``xj`` are always given in the physical (true) domain.
-    """
-
-    if jaxf is None:
-        from .forms import get_jaxfunctions
-
-        jaxf: set[JAXFunction] = get_jaxfunctions(a)
-        assert len(jaxf) == 1, "Single JAXFunction not found in expression."
-        jaxf: JAXFunction = jaxf.pop()
-
-    V = jaxf.functionspace
-
-    if isinstance(a, sp.Pow):
-        wa = a.args[0]
-        variables = getattr(wa, "variables", ())
-        var = tuple(int(variables.count(s)) for s in V.system.base_scalars())
-        var = var[0] if V.dims == 1 else var
-        h = V.evaluate_derivative(xj, jaxf.array, k=var)
-        return h ** int(a.exp)
-
-    if isinstance(a, sp.Derivative):
-        variables = getattr(a, "variables", ())
-        var = tuple(int(variables.count(s)) for s in V.system.base_scalars())
-        var = var[0] if V.dims == 1 else var
-        return V.evaluate_derivative(xj, jaxf.array, k=var)
-
-    if not isinstance(V, OrthogonalSpace | DirectSum):
-        return V.evaluate_mesh(xj, jaxf.array, True)
-
-    assert isinstance(xj, Array)
-    return V.evaluate(xj, jaxf.array)
-
-
 def evaluate_jaxfunction_expr_quad(
     a: Basic, jaxf: JAXFunction | None = None, N: int | tuple[int, ...] | None = None
 ) -> Array:
     """Evaluate a symbolic JAXFunction expression on the quadrature mesh."""
+    from jaxfun.integrators.nonlinear import compile_nonlinear_evaluator
 
     if jaxf is None:
         from .forms import get_jaxfunctions
@@ -837,19 +800,5 @@ def evaluate_jaxfunction_expr_quad(
         jaxf: JAXFunction = jaxf.pop()
 
     V = jaxf.functionspace
-
-    if isinstance(a, sp.Pow):
-        wa = a.args[0]
-        variables = getattr(wa, "variables", ())
-        var = tuple(int(variables.count(s)) for s in V.system.base_scalars())
-        var = var[0] if V.dims == 1 else var
-        h = V.backward_primitive(jaxf.array, k=var, N=N)
-        return h ** int(a.exp)
-
-    if isinstance(a, sp.Derivative):
-        variables = getattr(a, "variables", ())
-        var = tuple(int(variables.count(s)) for s in V.system.base_scalars())
-        var = var[0] if V.dims == 1 else var
-        return V.backward_primitive(jaxf.array, k=var, N=N)
-
-    return V.backward(jaxf.array, N=N)
+    fun = compile_nonlinear_evaluator(cast(sp.Expr, a), V, cast(AppliedUndef, jaxf))
+    return fun(jaxf.array, N)

src/jaxfun/galerkin/composite.py

@@ -252,6 +252,11 @@ def to_orthogonal(self, a: Array) -> Array:
         """Map composite coefficients -> underlying orthogonal coefficients."""
         return a @ self.S
 
+    @jax.jit(static_argnums=0)
+    def from_orthogonal(self, a: Array) -> Array:
+        """Map underlying orthogonal coefficients -> composite coefficients."""
+        return a @ self.get_inverse_stencil()
+
     @jax.jit(static_argnums=0)
     def apply_stencil_galerkin(self, b: Array) -> Array:
         """Apply stencil on both sides (Galerkin mass-like transform)."""
@@ -439,9 +444,11 @@ def __init__(self, a: Composite, b: BCGeneric) -> None:
         self.map_reference_domain = a.map_reference_domain
         self.map_true_domain = a.map_true_domain
         self.quad_points_and_weights = a.quad_points_and_weights
+        self.num_quad_points = a.num_quad_points
         self.get_orthogonal = a.get_orthogonal
         self.dims = a.dims
         self.rank = a.rank
+        self.domain = a.domain
 
     @overload
     def __getitem__(self, i: Literal[0]) -> Composite: ...
@@ -486,8 +493,16 @@ def to_orthogonal(self, c: Array) -> Array:
         """Map direct-sum coefficients -> underlying orthogonal coefficients."""
         c_a = self[0].to_orthogonal(c)
         c_b = self[1].to_orthogonal(self[1].bnd_vals())
-        Nd = c_b.shape[0]
-        return jnp.concatenate((c_a[:Nd] + c_b, c_a[Nd:]))
+        c_b = jnp.pad(c_b, (0, c_a.shape[0] - c_b.shape[0]))
+        return c_a + c_b
+
+    @jax.jit(static_argnums=0)
+    def from_orthogonal(self, a: Array) -> Array:
+        """Map underlying orthogonal coefficients -> direct-sum (inhomogeneous) coefficients."""  # noqa: E501
+        c_b = self[1].to_orthogonal(self[1].bnd_vals())
+        c_b = jnp.pad(c_b, (0, a.shape[0] - c_b.shape[0]))
+        c_a = a - c_b
+        return self[0].from_orthogonal(c_a)
 
     @jax.jit(static_argnums=0)
     def evaluate(self, x: Array, c: Array) -> Array:
@@ -512,15 +527,10 @@ def backward_primitive(
         """Return backward transform for k-th derivative."""
         return self.orthogonal.backward_primitive(self.to_orthogonal(c), k, kind, N)
 
+    @jax.jit(static_argnums=0)
     def forward(self, uj: Array) -> Array:
-        from .arguments import TestFunction, TrialFunction
-        from .inner import inner
-
-        u = TrialFunction(self)
-        v = TestFunction(self)
-        M, b = inner(v * u)
-        b += self[0].scalar_product(uj)
-        return jnp.linalg.solve(M, b)
+        a = self[0].orthogonal.forward(uj)
+        return self.from_orthogonal(a)
 
     @jax.jit(static_argnums=0)
     def scalar_product(self, uj: Array) -> Array:

src/jaxfun/galerkin/forms.py

@@ -150,7 +150,19 @@ def check_if_nonlinear_in_jaxfunction(a: sp.Basic) -> bool:
         return False
     assert len(jaxfunctions) <= 1, "Multiple JAXFunctions found"
     ad = a.doit(linear=True)  # assume linear
-    jf = ad.atoms(Jaxc).pop()
+    if ad.is_Vector:
+        for comp in ad.components.values():
+            jf = comp.atoms(Jaxc)
+            if len(jf) > 1:
+                return True
+            jf = jf.pop()
+            if sp.diff(comp, jf, 2) != 0:
+                return True
+        return False
+    jf = ad.atoms(Jaxc)
+    if len(jf) > 1:
+        return True
+    jf = jf.pop()
     return sp.diff(ad, jf, 2) != 0
 
 
@@ -263,7 +275,7 @@ def _split(ms: sp.Expr) -> InnerResultDict:
                 if arg.atoms(Jaxc):
                     jaxc.append(arg)
                     continue
-                if len(arg.free_symbols) == 1:
+                if len(arg.free_symbols) <= 1:
                     rest.append(arg)
                 else:
                     multivar.append(arg)
@@ -276,7 +288,7 @@ def _split(ms: sp.Expr) -> InnerResultDict:
                 if len(multivar) > 0:
                     d["multivar"] = sp.Mul(*multivar)
                 if len(jfun) > 0:
-                    d["jaxfunction"] = jfun[0]
+                    d["jaxfunction"] = sp.Mul(*jfun)
                 if len(jaxc) > 0:
                     d["coeff"] = jaxc[0]
         if d is None:

src/jaxfun/galerkin/inner.py

@@ -6,6 +6,7 @@
 from jax import Array
 from jax.experimental.sparse import BCOO
 
+from jaxfun.galerkin import JAXFunction
 from jaxfun.typing import TrialSpaceType
 from jaxfun.utils.common import lambdify, matmat, tosparse
 
@@ -28,6 +29,7 @@
 )
 from .orthogonal import OrthogonalSpace
 from .tensorproductspace import (
+    BlockTPMatrix,
     DirectSumTPS,
     TensorMatrix,
     TensorProductSpace,
@@ -160,6 +162,7 @@ def inner(
 
             if isinstance(z, tuple):  # multivar
                 mats.append((z, global_indices))
+                sc = 1
                 continue
 
             if z.size == 0:
@@ -202,7 +205,6 @@ def inner(
                 scales.append(
                     evaluate_jaxfunction_expr_quad(a0["jaxfunction"], N=num_quad_points)
                 )
-
             Am = assemble_multivar(mats_, scales, test_space)
             if has_bcs:
                 sign = 1 if all_linear else -1
@@ -224,7 +226,6 @@ def inner(
                     bresults.append(vectorize_bresult(res, test_space, gi[0][0]))
 
                 if "bilinear" in coeffs:
-                    assert coeffs["bilinear"] == 1
                     assert isinstance(trial_space, TensorProductSpace)
                     aresults.append(
                         TensorMatrix(
@@ -312,6 +313,7 @@ def inner(
             sc = sc * (-1)
 
         bs = []
+
         for key, bi in b0.items():
             if key in ("coeff", "multivar", "jaxfunction"):
                 continue
@@ -346,16 +348,17 @@ def inner(
             if isinstance(bs[0], tuple):
                 assert isinstance(num_quad_points, tuple)
                 # multivar or JAXFunction
+                uj = jnp.array(1.0)
                 if "multivar" in b0:
                     s = test_space.system.base_scalars()
-                    uj = lambdify(s, b0["multivar"], modules="jax")(
+                    uj *= lambdify(s, b0["multivar"], modules="jax")(
                         *test_space.mesh(N=num_quad_points)
                     )
-                elif "jaxfunction" in b0:
-                    uj = evaluate_jaxfunction_expr_quad(
+                if "jaxfunction" in b0:
+                    uj *= evaluate_jaxfunction_expr_quad(
                         b0["jaxfunction"], N=num_quad_points
                     )
-                else:
+                if "jaxfunction" not in b0 and "multivar" not in b0:
                     raise ValueError("Expected multivar or jaxfunction key in b0")
                 res = bs[0][0].T @ uj @ bs[1][0]
                 bresults.append(vectorize_bresult(res, test_space, global_index))
@@ -683,7 +686,7 @@ def assemble_multivar(
         test_space: Tensor product space (for mesh / variable order).
 
     Returns:
-        Dense matrix of shape (i*j, k*l) assembled from factors.
+        Dense matrix of shape (i, k, j, l) assembled from factors.
     """
     P0, P1 = mats[0]
     P2, P3 = mats[1]
@@ -730,18 +733,44 @@ def project(ue: sp.Expr, V: TrialSpaceType) -> Array:
     Returns:
         Coefficient array shaped to V.num_dofs.
     """
+    from scipy import sparse as scipy_sparse
+
+    from jaxfun.operators import Dot
+
     if V.dims == 1:
         assert isinstance(V, OrthogonalSpace | Composite | DirectSum)
         return project1D(ue, V)
 
-    if V.is_orthogonal:
+    if len(get_jaxfunctions(ue)) == 0:
         assert not isinstance(V, OrthogonalSpace | Composite | DirectSum)
-        uj = lambdify(V.system.base_scalars(), ue, modules="jax")(*V.mesh())
-        uj = jnp.broadcast_to(uj, V.num_dofs)
+        if V.rank == 0:
+            uj = lambdify(V.system.base_scalars(), ue, modules="jax")(*V.mesh())
+            uj = jnp.broadcast_to(uj, V.num_quad_points)
+        elif V.rank == 1:
+            assert isinstance(V, VectorTensorProductSpace)
+            s = V.system.base_scalars()
+            bv = V.system.base_vectors()
+            uj = (lambdify(s, Dot(ue, n).doit())(*V.mesh()) for n in bv)
+            uj = jnp.stack(
+                [jnp.broadcast_to(ui, V.tensorspaces[0].num_quad_points) for ui in uj],
+                axis=0,
+            )
         return V.forward(uj)
 
     u = TrialFunction(V)
     v = TestFunction(V)
-    M, b = inner(v * (u - ue))
-    uh = jnp.linalg.solve(M[0].mat, b.flatten()).reshape(V.num_dofs)
+    if V.rank == 0:
+        M, b = inner(v * (u - ue))
+        uh = jnp.linalg.solve(M[0].mat, b.flatten()).reshape(V.num_dofs)
+
+    elif V.rank == 1:
+        assert isinstance(ue, sp.Mul | sp.Add | JAXFunction), (
+            "Projection requires unevaluated expressions"
+        )  # noqa: E501
+        assert isinstance(V, VectorTensorProductSpace)
+        M, b = inner(Dot(v, (u - ue)))
+        A = BlockTPMatrix(M, V, V)
+        C = A.block_array()
+        uh = jnp.array(scipy_sparse.linalg.spsolve(C, b.ravel()).reshape(b.shape))
+
     return uh

src/jaxfun/galerkin/orthogonal.py

@@ -76,6 +76,7 @@ def __init__(
         self.S = sparse.BCOO(
             (jnp.ones(N), jnp.vstack((jnp.arange(N),) * 2).T), shape=(N, N)
         )
+        self.S_inv: Array | None = None
         super().__init__(system, name, fun_str)
 
     @abstractmethod
@@ -424,3 +425,10 @@ def get_padded(self, N: int) -> Self:
     def get_orthogonal(self) -> Self:
         """Return self (orthogonal space is self; overridden in Composite)."""
         return self
+
+    @jax.jit(static_argnums=0)
+    def get_inverse_stencil(self) -> Array:
+        """Return inverse of stencil matrix S."""
+        if self.S_inv is None:
+            self.S_inv = jnp.linalg.pinv(self.S.todense())
+        return self.S_inv