Solver.poisson support pre-integrated vfuncs

lapy/diffgeo.py

@@ -146,21 +146,22 @@ def compute_geodesic_f(
 
     gradf = compute_gradient(geom, vfunc)
     fem = Solver(geom, lump=True, use_cholmod=use_cholmod)
-    fem.mass = sparse.eye(fem.stiffness.shape[0], dtype=fem.stiffness.dtype)
 
+    # divf is the integrated divergence (so it is already B*div)
+    # we can solve by passing in integrate=False to avoid multiplying with B again
     if scalar_input:
         # gradf: (n_elements, 3)
         gradnorm = gradf / np.sqrt((gradf**2).sum(1))[:, np.newaxis]
         gradnorm = np.nan_to_num(gradnorm)
         divf = compute_divergence(geom, gradnorm)
-        vf = fem.poisson(divf)
+        vf = fem.poisson(divf, integrate=False)
         vf -= vf.min()
     else:
         # gradf: (n_elements, n_functions, 3) — norm along last axis
         gradnorm = gradf / np.sqrt((gradf**2).sum(-1))[:, :, np.newaxis]
         gradnorm = np.nan_to_num(gradnorm)
         divf = compute_divergence(geom, gradnorm)  # (n_vertices, n_functions)
-        vf = fem.poisson(divf)                     # (n_vertices, n_functions)
+        vf = fem.poisson(divf, integrate=False)     # (n_vertices, n_functions)
         vf -= vf.min(axis=0)
     return vf
 
@@ -198,23 +199,22 @@ def tria_compute_geodesic_f(
 
     gradf = tria_compute_gradient(tria, vfunc)
     fem = Solver(tria, lump=True, use_cholmod=use_cholmod)
-    # div is the integrated divergence (so it is already B*div);
-    # pass identity instead of B here
-    fem.mass = sparse.eye(fem.stiffness.shape[0])
 
+    # divf is the integrated divergence (so it is already B*div)
+    # we can solve by passing in integrate=False to avoid multiplying with B again
     if scalar_input:
         # gradf: (n_triangles, 3)
         gradnorm = gradf / np.sqrt((gradf**2).sum(1))[:, np.newaxis]
         gradnorm = np.nan_to_num(gradnorm)
         divf = tria_compute_divergence(tria, gradnorm)
-        vf = fem.poisson(divf)
+        vf = fem.poisson(divf, integrate=False)
         vf -= vf.min()
     else:
         # gradf: (n_triangles, n_functions, 3) — norm along last axis
         gradnorm = gradf / np.sqrt((gradf**2).sum(-1))[:, :, np.newaxis]
         gradnorm = np.nan_to_num(gradnorm)
         divf = tria_compute_divergence(tria, gradnorm)  # (n_vertices, n_functions)
-        vf = fem.poisson(divf)                          # (n_vertices, n_functions)
+        vf = fem.poisson(divf, integrate=False)          # (n_vertices, n_functions)
         vf -= vf.min(axis=0)
     return vf
 
@@ -503,20 +503,21 @@ def tria_compute_rotated_f(
     gradf = tria_compute_gradient(tria, vfunc)
     tn = tria.tria_normals()
     fem = Solver(tria, lump=True, use_cholmod=use_cholmod)
-    fem.mass = sparse.eye(fem.stiffness.shape[0], dtype=vfunc.dtype)
     dtup = (np.array([0]), np.array([0.0]))
-
+
+    # divf is the integrated divergence (so it is already B*div)
+    # we can solve by passing in integrate=False to avoid multiplying with B again
     if scalar_input:
         # gradf: (n_triangles, 3)
         gradf = np.cross(tn, gradf)
         divf = tria_compute_divergence(tria, gradf)
-        vf = fem.poisson(divf, dtup)
+        vf = fem.poisson(divf, dtup, integrate=False)
     else:
         # gradf: (n_triangles, n_functions, 3)
         # tn: (n_triangles, 3) -> broadcast via (n_triangles, 1, 3)
         gradf = np.cross(tn[:, np.newaxis, :], gradf)  # (n_triangles, n_functions, 3)
         divf = tria_compute_divergence(tria, gradf)    # (n_vertices, n_functions)
-        vf = fem.poisson(divf, dtup)                   # (n_vertices, n_functions)
+        vf = fem.poisson(divf, dtup, integrate=False)   # (n_vertices, n_functions)
     return vf
 
 

lapy/solver.py

@@ -784,13 +784,15 @@ def poisson(
         h: float | np.ndarray = 0.0,
         dtup: tuple = (),
         ntup: tuple = (),
+        integrate: bool = True,
     ) -> np.ndarray:
         """Solver for the Poisson equation with boundary conditions.
 
         This solver is based on the ``A`` and ``B`` Laplace matrices where
         ``A x = B h`` and ``A`` is a sparse symmetric positive semi-definite
         matrix of shape ``(n, n)`` and B is a sparse symmetric positive
-        definite matrix of shape ``(n, n)``.
+        definite matrix of shape ``(n, n)`` (use ``integrate=False`` to solve
+        ``A x = h`` instead).
 
         Parameters
         ----------
@@ -808,6 +810,12 @@ def poisson(
             Neumann boundary condition as a tuple containing the index and
             data arrays of same length. The default, an empty tuple,
             corresponds to Neumann on all boundaries.
+        integrate : bool, default=True
+            Whether to integrate the right hand side over the surface using
+            the mass matrix (after applying any Neumann conditions, but before
+            applying Dirichlet conditions). If True, the right hand side is
+            effectively replaced by ``B h``. If False, the right hand side is
+            used as is, which corresponds to ``A x = h``.
 
         Returns
         -------
@@ -904,8 +912,9 @@ def poisson(
                 (dim, n_rhs), dtype=dtype
             )
         # compute right hand side
-        mass = self.mass.astype(dtype, copy=False)
-        b = mass * (h - nvec)
+        b = h - nvec
+        if integrate:
+            b = self.mass.astype(dtype, copy=False) * b
         if len(didx) > 0:
             b = b - self.stiffness * dvec
         # remove Dirichlet Nodes

lapy/utils/tests/test_solver.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from scipy import sparse
 
 from ...solver import Solver
 from ...tria_mesh import TriaMesh
@@ -100,3 +101,19 @@ def test_poisson_2d_rhs_with_dirichlet(tria_mesh):
             err_msg=f"poisson 2-D Dirichlet mismatch at column {k}",
         )
 
+def test_poisson_with_integrate_false(tria_mesh):
+    """poisson with integrate=False must solve A x = h, not A x = B h."""
+    fem1 = Solver(tria_mesh, lump=True)
+    _, evec = fem1.eigs(k=5)
+    rhs = evec[:, 1:4]  # (n_vertices, 3)
+
+    res1 = fem1.poisson(rhs, integrate=False)
+
+    fem2 = Solver(tria_mesh, lump=True)
+    fem2.mass = sparse.eye(fem2.stiffness.shape[0], dtype=fem2.stiffness.dtype)
+    res2 = fem2.poisson(rhs, integrate=True)
+
+    np.testing.assert_allclose(
+        res1, res2, rtol=1e-6, atol=1e-9,
+        err_msg="poisson with integrate=False does not match poisson with identity mass matrix",
+    )