vectorize 2-D vfuncs in diffgeo and solver

lapy/solver.py

@@ -729,8 +729,11 @@ def poisson(
         Parameters
         ----------
         h : float or np.ndarray, default=0.0
-            Right hand side, can be constant or array with vertex values.
-            The default ``0.0`` corresponds to Laplace equation ``A x = 0``.
+            Right hand side, can be a scalar, a 1-D array of shape
+            ``(n_vertices,)``, or a 2-D array of shape
+            ``(n_vertices, n_functions)`` to solve for multiple functions
+            simultaneously.  The default ``0.0`` corresponds to the Laplace
+            equation ``A x = 0``.
         dtup : tuple, default=()
             Dirichlet boundary condition as a tuple containing the index and
             data arrays of same length. The default, an empty tuple,
@@ -743,7 +746,9 @@ def poisson(
         Returns
         -------
         np.ndarray
-            Array with vertex values of the solution, shape (n_vertices,).
+            Array with vertex values of the solution, shape ``(n_vertices,)``
+            for a scalar or 1-D input, or ``(n_vertices, n_functions)`` for a
+            2-D input.
 
         Raises
         ------
@@ -768,17 +773,30 @@ def poisson(
                 "Error: Square input matrices should have same number of rows and "
                 "columns."
             )
-        # create vector h
+        # create vector/matrix h
         if np.isscalar(h):
             h = np.full((dim, 1), h, dtype=dtype)
         else:
             h = np.asarray(h, dtype=dtype)
-            if h.size != dim:
+            if h.ndim == 1:
+                if h.size != dim:
+                    raise ValueError(
+                        "h should be either scalar or column vector with row num of A"
+                    )
+                h = h[:, np.newaxis]
+            elif h.ndim == 2:
+                if h.shape[0] != dim:
+                    raise ValueError(
+                        "h should be either scalar or array with first dim matching A"
+                    )
+            else:
                 raise ValueError(
-                    "h should be either scalar or column vector with row num of A"
+                    "h should be either scalar or 1-D/2-D array"
                 )
-        if h.ndim == 1:
-            h = h[:, np.newaxis]
+        # number of right-hand sides
+        n_rhs = h.shape[1]
+        scalar_rhs = n_rhs == 1  # remember whether caller passed 1-D / scalar
+
         # create vector d
         didx = []
         dvec = []
@@ -794,8 +812,12 @@ def poisson(
                 raise ValueError(
                     "dtup should contain index and data arrays (same lengths > 0)"
                 )
+            # broadcast Dirichlet values across all rhs columns
             dvec = sparse.csc_matrix(
-                (ddat, (didx, np.zeros(len(didx), dtype=np.uint32))), (dim, 1), dtype=dtype
+                (np.tile(ddat, n_rhs),
+                 (np.tile(didx, n_rhs),
+                  np.repeat(np.arange(n_rhs), len(didx)))),
+                (dim, n_rhs), dtype=dtype
             )
 
         # create vector n
@@ -810,7 +832,10 @@ def poisson(
                     "dtup should contain index and data arrays (same lengths > 0)"
                 )
             nvec = sparse.csc_matrix(
-                (ndat, (nidx, np.zeros(len(nidx), dtype=np.uint32))), (dim, 1), dtype=dtype
+                (np.tile(ndat, n_rhs),
+                 (np.tile(nidx, n_rhs),
+                  np.repeat(np.arange(n_rhs), len(nidx)))),
+                (dim, n_rhs), dtype=dtype
             )
         # compute right hand side
         mass = self.mass.astype(dtype, copy=False)
@@ -836,7 +861,7 @@ def poisson(
                 raise ValueError("A matrix needs to be sparse CSC or CSR")
         else:
             a = self.stiffness
-        # solve A x = b
+        # solve A x = b  (both splu.solve and cholmod accept 2-D b natively)
         logger.debug("Matrix format now: %s", a.getformat())
         if self.use_cholmod:
             logger.info("Solver: Cholesky decomposition from scikit-sparse cholmod ...")
@@ -848,11 +873,15 @@ def poisson(
             logger.info("Solver: spsolve (LU decomposition) ...")
             lu = splu(a)
             x = lu.solve(b.astype(dtype))
-        x = np.squeeze(np.array(x))
-        # pad Dirichlet nodes
+        # pad Dirichlet nodes back in
         if len(didx) > 0:
-            xfull = np.zeros(dim)
-            xfull[mask] = x
-            xfull[didx] = ddat
-            return xfull
+            xfull = np.zeros((dim, n_rhs), dtype=dtype)
+            xfull[mask, :] = x
+            xfull[didx, :] = dvec[didx, :].toarray()
+            x = xfull
+        # squeeze back to 1-D if the caller passed a scalar or 1-D rhs
+        if scalar_rhs:
+            x = np.squeeze(np.array(x))
+        else:
+            x = np.array(x)
         return x

lapy/tria_mesh.py

@@ -498,8 +498,7 @@ def tria_normals(self) -> np.ndarray:
         n = np.cross(v1mv0, v2mv0)
         ln = np.sqrt(np.sum(n * n, axis=1))
         ln[ln < sys.float_info.epsilon] = 1  # avoid division by zero
-        n = n / ln.reshape(-1, 1)
-        return n
+        return n / ln.reshape(-1, 1)
 
     def vertex_normals(self) -> np.ndarray:
         """Compute vertex normals.
@@ -537,16 +536,13 @@ def vertex_normals(self) -> np.ndarray:
         cr0 = np.cross(v1mv0, -v0mv2)
         cr1 = np.cross(v2mv1, -v1mv0)
         cr2 = np.cross(v0mv2, -v2mv1)
-        # Add normals at each vertex (there can be duplicate indices in t at vertex i)
         n = np.zeros(self.v.shape)
         np.add.at(n, self.t[:, 0], cr0)
         np.add.at(n, self.t[:, 1], cr1)
         np.add.at(n, self.t[:, 2], cr2)
-        # Normalize normals
         ln = np.sqrt(np.sum(n * n, axis=1))
         ln[ln < sys.float_info.epsilon] = 1  # avoid division by zero
-        n = n / ln.reshape(-1, 1)
-        return n
+        return n / ln.reshape(-1, 1)
 
     def has_free_vertices(self) -> bool:
         """Check if the vertex list has more vertices than what is used in triangles.
@@ -987,12 +983,8 @@ def curvature_tria(
         # print(np.max(np.abs(np.sum(tumin * tumax, axis=1))))
         # print(np.sum(tumin * tumax, axis=1))
 
-        # project onto triangle plane:
-        e0 = self.v[self.t[:, 1], :] - self.v[self.t[:, 0], :]
-        e1 = self.v[self.t[:, 2], :] - self.v[self.t[:, 0], :]
-        tn = np.cross(e0, e1)
-        tnl = np.sqrt(np.sum(tn * tn, axis=1)).reshape(-1, 1)
-        tn = tn / np.maximum(tnl, 1e-8)
+        # project onto triangle plane using tria normals (already unit length)
+        tn = self.tria_normals()
         # project tumin back to tria plane and normalize
         tumin2 = tumin - tn * (np.sum(tn * tumin, axis=1)).reshape(-1, 1)
         tuminl = np.sqrt(np.sum(tumin2 * tumin2, axis=1)).reshape(-1, 1)

lapy/utils/tests/test_TetMesh_Geodesics.py

@@ -74,13 +74,6 @@ def test_evals_evec_dimension(load_tet_mesh, loaded_data):
 def test_gradients_normalization_and_divergence(load_tet_mesh, loaded_data):
     """
     Test the computation of gradients, normalization, and divergence for a TetMesh.
-
-    Parameters:
-        load_tet_mesh (fixture): Fixture to load a TetMesh for testing.
-        loaded_data (dict): Dictionary containing loaded test data.
-
-    Raises:
-        AssertionError: If any test condition is not met.
     """
     T = load_tet_mesh
     tria = T.boundary_tria()
@@ -90,20 +83,12 @@ def test_gradients_normalization_and_divergence(load_tet_mesh, loaded_data):
     # Compute gradients
     tfunc = compute_gradient(T, u)
 
-    # Define the expected shape of tfunc (gradient)
-    expected_tfunc_shape = (48000, 3)
-
-    # Assert that the shape of tfunc matches the expected shape
-    assert tfunc.shape == expected_tfunc_shape
+    assert tfunc.shape == (48000, 3)
 
     # Flip and normalize
     X = -tfunc / np.sqrt((tfunc**2).sum(1))[:, np.newaxis]
 
-    # Define the expected shape of X (normalized gradient)
-    expected_X_shape = (48000, 3)
-
-    # Assert that the shape of X matches the expected shape
-    assert X.shape == expected_X_shape
+    assert X.shape == (48000, 3)
 
     # Load the expected maximum and minimum values for each column of X
     expected_max_col_values = loaded_data["expected_outcomes"][
@@ -113,60 +98,41 @@ def test_gradients_normalization_and_divergence(load_tet_mesh, loaded_data):
         "test_TetMesh_Geodesics"
     ]["expected_min_col_values"]
 
-    # Assert maximum and minimum values of each column of X match the expected values
     for col in range(X.shape[1]):
         assert np.allclose(np.max(X[:, col]), expected_max_col_values[col], atol=1e-6)
         assert np.allclose(np.min(X[:, col]), expected_min_col_values[col], atol=1e-6)
 
     # Compute divergence
     divx = compute_divergence(T, X)
-
-    # Define the expected shape of divx (divergence)
-    expected_divx_shape = (9261,)
-
-    # Assert that the shape of divx matches the expected shape
-    assert divx.shape == expected_divx_shape
+    assert divx.shape == (9261,)
 
 
 def test_tetMesh_Geodesics_format(load_tet_mesh, loaded_data):
     """
-    Test if matrix format, solver settings, max distance,
-    and computed values match the expected outcomes.
-
-    Parameters:
-    - loaded_data (dict): Dictionary containing loaded test data.
-
-    Raises:
-    - AssertionError: If any test condition is not met.
+    Test matrix format, solver settings, max distance, and computed values.
     """
-
     T = load_tet_mesh
     tria = T.boundary_tria()
     bvert = np.unique(tria.t)
     u = diffusion(T, bvert, m=1)
 
-    # get gradients
     tfunc = compute_gradient(T, u)
-    # flip and normalize
     X = -tfunc / np.sqrt((tfunc**2).sum(1))[:, np.newaxis]
     X = np.nan_to_num(X)
-    # compute divergence
     divx = compute_divergence(T, X)
 
-    # compute distance
     useCholmod = True
     try:
         from sksparse.cholmod import cholesky
     except ImportError:
         useCholmod = False
 
     fem = Solver(T, lump=True)
-    A, B = fem.stiffness, fem.mass  # computed above when creating Solver
+    A, B = fem.stiffness, fem.mass
 
     H = A
     b0 = -divx
 
-    # solve H x = b0
     if useCholmod:
         print("Solver: cholesky decomp - performance optimal ...")
         chol = cholesky(H)
@@ -178,12 +144,9 @@ def test_tetMesh_Geodesics_format(load_tet_mesh, loaded_data):
 
     x = x - np.min(x)
 
-    # get heat diffusion
-
     v1func = T.v[:, 0] * T.v[:, 0] + T.v[:, 1] * T.v[:, 1] + T.v[:, 2] * T.v[:, 2]
     grad = compute_gradient(T, v1func)
     glength = np.sqrt(np.sum(grad * grad, axis=1))
-    A, B = fem.stiffness, fem.mass
     Bi = B.copy()
     Bi.data **= -1
     divx2 = Bi * divx

lapy/utils/tests/test_TriaMesh_Geodesics.py

@@ -190,3 +190,4 @@ def test_Geodesics_format(loaded_data, load_square_mesh):
     ]
     computed_max_abs_diff = max(abs(gf - x))
     assert np.allclose(computed_max_abs_diff, expected_max_abs_diff)
+