add multiple seeds sets in heat diffusion

examples/Test_TriaMesh_Geodesics.ipynb

@@ -256,9 +256,11 @@
     }
    ],
    "source": [
+    "import numpy as np\n",
+    "\n",
     "from lapy import heat\n",
     "\n",
-    "bvert = T.boundary_loops()\n",
+    "bvert = np.concatenate(T.boundary_loops())\n",
     "u = heat.diffusion(T, bvert, m=1)"
    ]
   },
@@ -604,4 +606,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

lapy/heat.py

@@ -4,9 +4,10 @@
 mesh geometries (tet or tria mesh) for heat diffusion.
 """
 
+from __future__ import annotations
+
 import importlib
 import logging
-from typing import Optional, Union
 
 import numpy as np
 
@@ -16,7 +17,7 @@
 
 
 def diagonal(
-    t: Union[float, np.ndarray],
+    t: float | np.ndarray,
     x: np.ndarray,
     evecs: np.ndarray,
     evals: np.ndarray,
@@ -60,7 +61,7 @@ def diagonal(
 
 
 def kernel(
-    t: Union[float, np.ndarray],
+    t: float | np.ndarray,
     vfix: int,
     evecs: np.ndarray,
     evals: np.ndarray,
@@ -112,9 +113,9 @@ def kernel(
 
 def diffusion(
     geometry: object,
-    vids: Union[int, np.ndarray],
+    vids: int | np.ndarray | list[int] | list[list[int]] | list[np.ndarray],
     m: float = 1.0,
-    aniso: Optional[int] = None,
+    aniso: int | None = None,
     use_cholmod: bool = False,
 ) -> np.ndarray:
     """Compute the heat diffusion from initial vertices in vids.
@@ -126,8 +127,32 @@ def diffusion(
     ----------
     geometry : TriaMesh or TetMesh
         Geometric object on which to run diffusion.
-    vids : int or np.ndarray
+    vids : int or array-like
         Vertex index or indices where initial heat is applied.
+
+        The nesting level determines whether one or multiple seed sets are
+        computed:
+
+        - **Single seed set** — an ``int``, a 1-D ``np.ndarray``, or a plain
+          ``list[int]``: heat is seeded at those vertices and a 1-D result
+          array of shape ``(n_vertices,)`` is returned.
+        - **Multiple seed sets** — a ``list[list[int]]`` or a
+          ``list[np.ndarray]``: each inner sequence defines one independent
+          seed set.  The heat matrix is factorised only once and reused for
+          all sets.  A 2-D array of shape ``(n_vertices, n_cases)`` is
+          returned, where column ``k`` corresponds to ``vids[k]``.
+
+        .. note::
+            The direct output of :func:`~lapy.TriaMesh.boundary_loops` is a
+            ``list[list[int]]`` and therefore treated as **multiple seed
+            sets** (one diffusion per boundary loop).  To compute the
+            distance from *all* boundary vertices at once, concatenate the
+            loops first::
+
+                loops = tria.boundary_loops()
+                all_bnd = np.concatenate(loops)
+                vfunc = diffusion(tria, all_bnd)
+
     m : float, default=1.0
         Factor to compute time of heat evolution.
     aniso : int, default=None
@@ -139,7 +164,8 @@ def diffusion(
     Returns
     -------
     np.ndarray
-        Heat diffusion values at vertices, shape (n_vertices,).
+        Heat diffusion values at vertices.  Shape ``(n_vertices,)`` for a
+        single seed set, or ``(n_vertices, n_cases)`` for multiple seed sets.
 
     Raises
     ------
@@ -156,18 +182,38 @@ def diffusion(
     from . import Solver
 
     nv = len(geometry.v)
-    vids = np.asarray(vids, dtype=int)
-    if np.any(vids < 0) or np.any(vids >= nv):
-        raise ValueError("vids contains out-of-range vertex indices")
+
+    # ------------------------------------------------------------------
+    # Nesting level determines single vs. multiple seed sets:
+    #   Multiple: list[list[int]] or list[np.ndarray]  → scalar_input=False
+    #   Single  : int, np.ndarray, or list[int]        → scalar_input=True
+    # ------------------------------------------------------------------
+    if isinstance(vids, list) and len(vids) > 0 and isinstance(
+        vids[0], (list, np.ndarray)
+    ):
+        scalar_input = False
+        vids_list = [np.asarray(v, dtype=int).ravel() for v in vids]
+    else:
+        scalar_input = True
+        vids_list = [np.asarray(vids, dtype=int).ravel()]
+
+    for v in vids_list:
+        if np.any(v < 0) or np.any(v >= nv):
+            raise ValueError("vids contains out-of-range vertex indices")
+
     fem = Solver(geometry, lump=True, aniso=aniso)
     # time of heat evolution:
     t = m * geometry.avg_edge_length() ** 2
-    # backward Euler matrix:
+    # backward Euler matrix (same for all cases):
     hmat = fem.mass + t * fem.stiffness
-    # set initial heat
-    b0 = np.zeros((nv,))
-    b0[vids] = 1.0
-    # solve H x = b0
+
+    # Build rhs matrix — each column is one initial-heat vector
+    n_cases = len(vids_list)
+    b0 = np.zeros((nv, n_cases))
+    for k, v in enumerate(vids_list):
+        b0[v, k] = 1.0
+
+    # Solve — both splu.solve and cholmod accept 2-D rhs natively
     logger.debug("Matrix Format: %s", hmat.getformat())
     if use_cholmod:
         logger.info("Solver: Cholesky decomposition from scikit-sparse cholmod")
@@ -178,5 +224,9 @@ def diffusion(
 
         logger.info("Solver: LU decomposition via splu")
         lu = splu(hmat)
-        vfunc = lu.solve(np.float32(b0))
+        vfunc = lu.solve(b0.astype(np.float64))
+
+    # Return 1-D array for the single-case (backward compatible) path
+    if scalar_input:
+        return vfunc[:, 0]
     return vfunc

lapy/utils/tests/test_TriaMesh_Geodesics.py

@@ -119,7 +119,7 @@ def test_heat_diffusion_shape(load_square_mesh):
         None
     """
     T = load_square_mesh
-    bvert = T.boundary_loops()
+    bvert = np.concatenate(T.boundary_loops())
     u = diffusion(T, bvert, m=1)
     expected_shape = (len(T.v),)
     assert u.shape == expected_shape
@@ -130,7 +130,7 @@ def test_Geodesics_format(loaded_data, load_square_mesh):
     Test geodesics format and accuracy.
     """
     T = load_square_mesh
-    bvert = T.boundary_loops()
+    bvert = np.concatenate(T.boundary_loops())
     u = diffusion(T, bvert, m=1)
     # compute gradient of heat diffusion
     tfunc = compute_gradient(T, u)

lapy/utils/tests/test_heat.py

@@ -0,0 +1,72 @@
+"""Tests for lapy.heat — backward compatibility and multi-case diffusion."""
+
+import numpy as np
+import pytest
+
+from ...heat import diffusion
+from ...tria_mesh import TriaMesh
+
+
+@pytest.fixture
+def tria_mesh():
+    return TriaMesh.read_off("data/square-mesh.off")
+
+
+# ---------------------------------------------------------------------------
+# Backward compatibility — single vids set must still return 1-D
+# ---------------------------------------------------------------------------
+
+
+def test_diffusion_single_returns_1d(tria_mesh):
+    """int, list[int], or 1-D array vids must return a 1-D (n_vertices,) array."""
+    T = tria_mesh
+    bvert = np.concatenate(T.boundary_loops())
+
+    assert diffusion(T, bvert, m=1).shape == (len(T.v),)
+    assert diffusion(T, bvert.tolist(), m=1).shape == (len(T.v),)
+    assert diffusion(T, 0, m=1).shape == (len(T.v),)
+
+
+# ---------------------------------------------------------------------------
+# Multi-case — list of arrays/lists returns 2-D, each column matches single call
+# ---------------------------------------------------------------------------
+
+
+def test_diffusion_multi_matches_single(tria_mesh):
+    """Multi-case diffusion columns must match independent single calls."""
+    T = tria_mesh
+    seeds = [np.concatenate(T.boundary_loops()), np.array([0]), np.array([1, 2])]
+
+    u_batch = diffusion(T, seeds, m=1)
+
+    assert u_batch.shape == (len(T.v), 3)
+    for k, s in enumerate(seeds):
+        np.testing.assert_allclose(
+            u_batch[:, k], diffusion(T, s, m=1), rtol=1e-6, atol=1e-9
+        )
+
+
+def test_diffusion_boundary_loops_multi(tria_mesh):
+    """boundary_loops() output (list[list[int]]) triggers multi-case path."""
+    T = tria_mesh
+    loops = T.boundary_loops()
+
+    u_batch = diffusion(T, loops, m=1)
+
+    assert u_batch.ndim == 2
+    assert u_batch.shape == (len(T.v), len(loops))
+
+
+# ---------------------------------------------------------------------------
+# Error handling
+# ---------------------------------------------------------------------------
+
+
+def test_diffusion_out_of_range_raises(tria_mesh):
+    """Out-of-range vertex indices must raise ValueError."""
+    T = tria_mesh
+    with pytest.raises(ValueError, match="out-of-range"):
+        diffusion(T, np.array([len(T.v)]), m=1)
+
+    with pytest.raises(ValueError, match="out-of-range"):
+        diffusion(T, [np.array([0]), np.array([len(T.v)])], m=1)