sublevel-flood 0.1__tar.gz

sublevel_flood-0.1/.gitignore

@@ -0,0 +1,5 @@
+.gitignore
+.idea/
+dist/
+__pycache__/
+check.py

sublevel_flood-0.1/PKG-INFO

@@ -0,0 +1,68 @@
+Metadata-Version: 2.5
+Name: sublevel-flood
+Version: 0.1
+Summary: A package that computes the sublevel Flood bifiltration, a scalable approach for using 2-parameter persistent homology with large-scale, low-dimensional point sets.
+Project-URL: Homepage, https://github.com/MClemot/SLFlood
+Requires-Python: >=3.10
+Requires-Dist: fpsample
+Requires-Dist: gudhi
+Requires-Dist: multipers>=2.6.0
+Requires-Dist: numpy
+Requires-Dist: torch
+Provides-Extra: sklearn
+Requires-Dist: scikit-learn; extra == 'sklearn'
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: scikit-learn; extra == 'test'
+Description-Content-Type: text/markdown
+
+# The sublevel Flood bifiltration: towards scalable 2-parameter persistent homology
+
+[![PyPI - Version](https://img.shields.io/pypi/v/sublevel-flood?logo=python)](https://pypi.org/project/sublevel-flood/)
+
+This repository contains the code for the manuscript *Sublevel Flood bifiltration: scalable 2-parameter persistent homology*.
+The aim is to provide a way to compute 2-parameter persistent homology of large scale point clouds (up to $10^6$).
+To do so, it extends the [*Flood filtration*](https://proceedings.neurips.cc/paper_files/paper/2025/hash/aba03e6f25decb32bda9c5bf81c58305-Abstract-Conference.html) to the 2-parameter context.
+
+The Python package constructs a ```SimplexTreeMulti``` from the [```multipers```](https://github.com/DavidLapous/multipers) library.
+
+## Installation
+`sublevel-flood` can be installed via PyPI.
+
+```pip install sublevel-flood```
+
+## Usage
+
+The following example computes the sublevel Flood bifiltration (with 100 landmarks) of a point set in $\mathbb{R}^3$ distributed as a sphere with ambient noise, using CUDA. Then, it computes and plots its multiparameter module approximation with `multipers`.
+
+```python
+import slflood
+from slflood.data import noisy_sphere_data
+
+import matplotlib.pyplot as plt
+import multipers
+
+n_pts = 10000
+n_lms = 100
+dim = 3
+
+pts, lms_idx, fun = noisy_sphere_data(n_pts, n_lms, dim,
+                                      bandwidth=0.2)
+
+slf = slflood.slflood_bifiltration(pts.cuda(), lms_idx.cuda(), fun.cuda())
+mma = multipers.module_approximation(slf)
+mma.plot(box=[[0, 0], [1, 1]])
+plt.show()
+```
+
+![](https://raw.githubusercontent.com/MClemot/SLFlood/main/images/noisy_sphere_mma.png)
+
+## Citation
+```bibtex
+@inproceedings{clemot2026sublevel,
+      title={The sublevel {Flood} bifiltration: towards scalable 2-parameter persistent homology}, 
+      author={Cl{\'e}mot, Matt{\'e}o and Digne, Julie and Tierny, Julien},
+      year={2026},
+      booktitle={NeurIPS},
+}
+```

sublevel_flood-0.1/README.md

@@ -0,0 +1,50 @@
+# The sublevel Flood bifiltration: towards scalable 2-parameter persistent homology
+
+[![PyPI - Version](https://img.shields.io/pypi/v/sublevel-flood?logo=python)](https://pypi.org/project/sublevel-flood/)
+
+This repository contains the code for the manuscript *Sublevel Flood bifiltration: scalable 2-parameter persistent homology*.
+The aim is to provide a way to compute 2-parameter persistent homology of large scale point clouds (up to $10^6$).
+To do so, it extends the [*Flood filtration*](https://proceedings.neurips.cc/paper_files/paper/2025/hash/aba03e6f25decb32bda9c5bf81c58305-Abstract-Conference.html) to the 2-parameter context.
+
+The Python package constructs a ```SimplexTreeMulti``` from the [```multipers```](https://github.com/DavidLapous/multipers) library.
+
+## Installation
+`sublevel-flood` can be installed via PyPI.
+
+```pip install sublevel-flood```
+
+## Usage
+
+The following example computes the sublevel Flood bifiltration (with 100 landmarks) of a point set in $\mathbb{R}^3$ distributed as a sphere with ambient noise, using CUDA. Then, it computes and plots its multiparameter module approximation with `multipers`.
+
+```python
+import slflood
+from slflood.data import noisy_sphere_data
+
+import matplotlib.pyplot as plt
+import multipers
+
+n_pts = 10000
+n_lms = 100
+dim = 3
+
+pts, lms_idx, fun = noisy_sphere_data(n_pts, n_lms, dim,
+                                      bandwidth=0.2)
+
+slf = slflood.slflood_bifiltration(pts.cuda(), lms_idx.cuda(), fun.cuda())
+mma = multipers.module_approximation(slf)
+mma.plot(box=[[0, 0], [1, 1]])
+plt.show()
+```
+
+![](https://raw.githubusercontent.com/MClemot/SLFlood/main/images/noisy_sphere_mma.png)
+
+## Citation
+```bibtex
+@inproceedings{clemot2026sublevel,
+      title={The sublevel {Flood} bifiltration: towards scalable 2-parameter persistent homology}, 
+      author={Cl{\'e}mot, Matt{\'e}o and Digne, Julie and Tierny, Julien},
+      year={2026},
+      booktitle={NeurIPS},
+}
+```

sublevel_flood-0.1/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "sublevel-flood"
+version = "0.1"
+description = "A package that computes the sublevel Flood bifiltration, a scalable approach for using 2-parameter persistent homology with large-scale, low-dimensional point sets."
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "fpsample",
+    "gudhi",
+    "multipers >= 2.6.0",
+    "numpy",
+    "torch",
+]
+
+[project.urls]
+Homepage = "https://github.com/MClemot/SLFlood"
+
+[project.optional-dependencies]
+sklearn = ["scikit-learn"]
+test = ["pytest", "scikit-learn"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["slflood"]

sublevel_flood-0.1/slflood/__init__.py

@@ -0,0 +1,6 @@
+from .landmarks import generate_landmarks_FPS
+from .slflood import slflood_bifiltration
+
+__version__ = "0.1"
+
+__all__ = ["slflood_bifiltration", "generate_landmarks_FPS"]

sublevel_flood-0.1/slflood/data/__init__.py

@@ -0,0 +1,3 @@
+from .generate import noisy_sphere_data
+
+__all__ = ["noisy_sphere_data"]

sublevel_flood-0.1/slflood/data/generate.py

@@ -0,0 +1,41 @@
+from multipers.filtrations.density import KDE
+import torch
+from ..landmarks import generate_landmarks_FPS
+
+device = 'cpu'
+
+def generate_ambient_noise(n_pts, dim, radius=1., seed=None):
+    if seed is not None:
+        torch.manual_seed(seed)
+    directions = torch.normal(0, 1, size=(n_pts, dim), device=device)
+    directions /= torch.linalg.norm(directions, dim=-1, keepdim=True)
+    radii = torch.rand((n_pts, 1), device=device) ** (1. / dim)
+    return directions * radii * radius
+
+def generate_sphere(n, dim, device, seed=None):
+    if seed is not None:
+        torch.manual_seed(seed)
+    X = torch.normal(0, 1, size=(n,dim), device=device)
+    X = X / torch.norm(X, dim=-1, keepdim=True)
+    X += torch.normal(0, 0.01, size=(n,dim), device=device)
+    return X
+
+def generate_noisy_sphere(n_pts, dim, seed=None):
+    pts = generate_sphere(n_pts // 2, dim, device, seed=seed)
+    pts = torch.cat([pts, generate_ambient_noise(n_pts // 2, dim, seed=seed) * 1.5])
+    return pts
+
+def density(pts, bandwidth):
+    f = -KDE(bandwidth=bandwidth, return_log=False).fit(pts).score_samples(pts)
+    f -= f.min()
+    f /= f.max()
+    return f
+
+def sub_pts(pts, n_lms, start_idx=0):
+    return generate_landmarks_FPS(pts, n_lms, start_idx=start_idx, return_idxs=True)
+
+def noisy_sphere_data(n_pts, n_lms, dim, bandwidth=.1, seed=None):
+    pts = generate_noisy_sphere(n_pts, dim=dim, seed=seed)
+    fun = density(pts, bandwidth)
+    lms_idx = sub_pts(pts, n_lms)
+    return pts, lms_idx, fun

sublevel_flood-0.1/slflood/landmarks.py

@@ -0,0 +1,63 @@
+import fpsample
+import numpy as np
+import torch
+from typing import Union
+
+@torch.no_grad()
+def generate_landmarks_FPS(
+    points: torch.Tensor,
+    n_lms: int,
+    fps_h: Union[None, int] = None,
+    start_idx: Union[int, None] = None,
+    return_idxs: bool = False,
+) -> torch.Tensor:
+    """
+    Selects landmarks using Farthest-Point Sampling (bucket FPS).
+
+    This method implements a variant of Farthest-Point Sampling from
+    [here](https://dl.acm.org/doi/abs/10.1109/TCAD.2023.3274922).
+
+    Adapted from https://github.com/plus-rkwitt/flooder.
+
+    Args:
+        points (torch.Tensor):
+            A (P, d) tensor representing a point cloud. The tensor may reside on any
+            device (CPU or GPU) and be of any floating-point dtype.
+        n_lms (int):
+            The number of landmarks to sample (must be <= P and > 0).
+        fps_h (Union[None, int], optional):
+            h parameter (depth of kdtree) that is used for farthest point sampling to
+            select the landmarks. If None, then h is selected based on the size of the
+            point cloud. Defaults to None.
+        start_idx (int | None, optional):
+            If provided, the sampling starts from this index in the point cloud. If not,
+            the start index will be randomly picked from the point cloud.
+        return_idxs (bool, optional):
+            If true, returns the index of the selected landmarks. Defaults to False.
+
+    Returns:
+        torch.Tensor:
+            A (n_l, d) tensor containing a subset of the input `points`, representing the sampled landmarks,
+            or a (n_l) tensor of integers containing the indices of the sampled landmarks.
+    """
+    if n_lms <= 0:
+        raise RuntimeError(
+            f"Number of landmarks ({n_lms}) must be positive"
+        )
+    n_pts = len(points)
+    n_lms = min(n_lms, n_pts)
+    if fps_h is None:
+        if n_pts > 200_000:
+            fps_h = 9
+        elif n_pts > 80_000:
+            fps_h = 7
+        else:
+            fps_h = 5
+
+    index_set = torch.tensor(
+        fpsample.bucket_fps_kdline_sampling(points.cpu(), n_lms, h=fps_h, start_idx=start_idx).astype(np.int64),
+        device=points.device,
+    )
+    if return_idxs:
+        return index_set
+    return points[index_set]

sublevel_flood-0.1/slflood/meb.py

@@ -0,0 +1,88 @@
+import numpy as np
+from itertools import combinations
+
+def get_circumspheres(S):
+    """
+    Computes the circumspheres of a batch of set of points
+
+    Parameters
+    ----------
+    S : (B, M, N) ndarray, where 1 <= M <= N + 1
+            The input points
+
+    Returns
+    -------
+    C, r : ((2) ndarray, float)
+            The center and the squared radius of the circumsphere
+    """
+
+    U = S[:,1:] - S[:,0,None]
+    B = np.linalg.norm(U, axis=2, keepdims=True)
+    U /= B
+    B /= 2
+    G = np.matmul(U, U.transpose(0,2,1))
+    try:
+        x = np.linalg.solve(G, B)[:,:,0]
+    except np.linalg.LinAlgError:
+        G += np.random.normal(size=G.shape, loc=0, scale=1e-10)
+        x = np.linalg.solve(G, B)[:,:,0]
+    C = np.einsum('bi,bij->bj', x, U)
+    r = np.linalg.norm(C, axis=1)
+    C += S[:,0]
+    return C, r
+
+
+def get_middles(p):
+    """
+    Computes the middles of a batch of pair of points
+    """
+    center = 0.5 * (p[:, 0] + p[:, 1])
+    radius = 0.5 * np.linalg.norm(p[:, 1] - p[:, 0], axis=-1)
+    return center, radius
+
+
+def contains_all(center, radius, pts, tol=1e-7):
+    """
+    Check if ball contains all 4 tetrahedron vertices.
+    """
+    d = np.linalg.norm(pts - center[:, None, :], axis=-1)
+    return np.all(d <= radius[:, None] + tol, axis=-1)
+
+
+def minimum_enclosing_balls(X:np.ndarray):
+    """
+    Compute the minimum enclosing balls of a batch of sets of points
+
+    Parameters
+    ----------
+    X : (B, N, D)
+        input
+
+    Returns
+    -------
+    centers,radius : ((2) ndarray, float)
+        centers and radii of the minimum enclosing balls
+    """
+    B = X.shape[0]
+    N = X.shape[1]
+    D = X.shape[2]
+
+    best_r = np.full(B, np.inf)
+    best_c = np.zeros((B, D))
+
+    for k in range(N,1,-1):
+        for idx in combinations(range(N), k):
+            subset = X[:, idx]
+
+            if k == 2:
+                c, r = get_middles(subset)
+            else:
+                c, r = get_circumspheres(subset)
+
+            valid = contains_all(c, r, X)
+            update = valid & (r < best_r)
+
+            best_r[update] = r[update]
+            best_c[update] = c[update]
+
+    return best_c, best_r

sublevel_flood-0.1/slflood/sklearn/__init__.py

@@ -0,0 +1,3 @@
+from .transformer import SublevelFloodBifiltration
+
+__all__ = ["SublevelFloodBifiltration"]

sublevel_flood-0.1/slflood/sklearn/transformer.py

@@ -0,0 +1,104 @@
+import slflood
+
+from multipers.filtrations.density import KDE, DTM
+import numpy as np
+try:
+    from sklearn.base import BaseEstimator, TransformerMixin
+except ImportError as e:
+    raise ImportError(
+        "slflood.sklearn requires scikit-learn, which can be installed with `pip install sublevel-flood[sklearn]`"
+    ) from e
+import torch
+
+class SublevelFloodBifiltration(BaseEstimator, TransformerMixin):
+    """
+    Scikit-learn transformer computing the sublevel Flood bifiltration of point clouds.
+    For each input point cloud, a co-density function is estimated via kernel density or distance-to-measure.
+
+    Parameters
+    ----------
+    n_lms : int
+        Number of landmarks to sample (via farthest point sampling).
+
+    kde_bandwidth : float, optional
+        Bandwidth of the Gaussian kernel used to compute the density functions.
+        Mutually exclusive with ``dtm_mass``.
+
+    dtm_mass : float, optional
+        Mass parameter in (0, 1] used to compute the Distance-To-Measure (DTM) functions.
+        Mutually exclusive with ``kde_bandwidth``.
+
+    log_density : bool, default=False
+        If True and ``kde_bandwidth`` is set, uses the log-density instead of the density.
+
+    normalize : bool, default=True
+        Whether to normalize the density functions.
+
+    points_per_edge : int, default=20
+        Resolution of the sampling grid used on each simplex.
+
+    exact : bool, default=False
+        Whether to compute the exact (unmasked) or approximate (masked, faster) sublevel Flood bifiltration.
+
+    device : {'auto', 'cpu', 'cuda'}, default='auto'
+        Device used for computation. When 'auto', detects whether CUDA is available.
+
+    use_triton : bool, default=True
+        Whether to use Triton kernels when available.
+
+    n_jobs : int, default=-1
+        Currently unused.
+    """
+
+    def __init__(self, n_lms, kde_bandwidth=None, dtm_mass=None, log_density=False, normalize=True, points_per_edge=20, exact=False, device='auto', use_triton=True, n_jobs=-1):
+        self.n_lms = n_lms
+        self.kde_bandwidth = kde_bandwidth
+        self.dtm_mass = dtm_mass
+        self.log_density = log_density
+        self.normalize = normalize
+        self.points_per_edge = points_per_edge
+        self.exact = exact
+        self.device = device
+        self.use_triton = use_triton
+        self.n_jobs = n_jobs
+
+    def __get_device(self):
+        if self.device not in ['cpu', 'cuda', 'auto']:
+            raise ValueError(f"device must be 'cpu', 'cuda' or 'auto', got {self.device!r}")
+        if self.device == 'auto':
+            return 'cuda' if torch.cuda.is_available() else 'cpu'
+        if self.device == 'cuda' and not torch.cuda.is_available():
+            raise RuntimeError("device='cuda' was requested but CUDA is not available")
+        return self.device
+
+    def fit(self, X, y=None):
+        return self
+
+    def __density(self, X):
+        functions = []
+        for pts in X:
+            if self.kde_bandwidth is not None:
+                function = -KDE(bandwidth=self.kde_bandwidth, kernel="gaussian", return_log=self.log_density).fit(pts).score_samples(pts)
+            elif self.dtm_mass is not None:
+                function = DTM(masses=[self.dtm_mass]).fit(pts).score_samples(pts)[0]
+            else:
+                raise ValueError("exactly one parameter among kde_bandwidth and dtm_mass must be given")
+            if self.normalize:
+                function -= function.min()
+                function /= function.max()
+            functions.append(function)
+        return functions
+
+    def __transform(self, x, f, device):
+        x = torch.as_tensor(x, device=device)
+        f = torch.as_tensor(f, device=device)
+        return slflood.slflood_bifiltration(x, self.n_lms, f,
+                                            points_per_edge=self.points_per_edge,
+                                            exact=self.exact,
+                                            use_triton=self.use_triton)
+
+    def transform(self, X):
+        device = self.__get_device()
+        F = self.__density(X)
+
+        return [self.__transform(x, f, device) for (x,f) in zip(X, F)]

sublevel_flood-0.1/slflood/slflood.py

@@ -0,0 +1,258 @@
+"""Implementation of sublevel flood bifiltration."""
+
+from typing import Union
+
+import gudhi
+import multipers as mp
+import numbers
+import numpy as np
+import torch
+
+from .landmarks import generate_landmarks_FPS
+from .meb import minimum_enclosing_balls
+from .utils import generate_grid, mask_from_adjacency, pad_bigrades, reset, elapsed
+
+try:
+    from .triton_maxcummin import maxcummin
+    from .triton_mask_topk import compact_mask_indices
+    HAS_TRITON = True
+except ImportError:
+    HAS_TRITON = False
+
+@torch.no_grad()
+def near_validity(points: torch.Tensor,
+                  simplex_centers: torch.Tensor,
+                  simplex_radii: torch.Tensor) -> torch.Tensor:
+    centers_to_points = torch.cdist(simplex_centers, points)
+    return centers_to_points < simplex_radii.unsqueeze(1)
+
+
+@torch.no_grad()
+def vertex_bigrades(points: torch.Tensor,
+                    landmarks: torch.Tensor,
+                    function: torch.Tensor,
+                    cells: torch.Tensor,
+                    mask: torch.Tensor):
+    landmarks_to_points = torch.cdist(landmarks, points)
+    cummin = torch.cummin(landmarks_to_points, dim=1)[0]
+    critical = torch.ones_like(cummin, dtype=torch.int8)
+    critical[:, 1:] = cummin[:, 1:] != cummin[:, :-1]
+
+    critical_count = critical.sum(1)
+    critical_count_max = critical_count.max().item()
+    critical_indices = torch.argsort(critical, dim=1, descending=True, stable=True)[:,:critical_count_max]
+    bigrades_diameters = torch.gather(cummin, dim=1, index=critical_indices)
+    bigrades_functions = function[critical_indices]
+    bigrades = torch.stack((bigrades_diameters, bigrades_functions), dim=-1).cpu().to(torch.float32)
+
+    scc0 = [bigrades[k,:critical_count[k]] for k in range(len(landmarks))]
+
+    cell_mask = critical[cells].amax(dim=1)
+    mask = torch.max(mask, cell_mask)
+
+    return scc0, mask
+
+
+@torch.no_grad()
+def slflood_bifiltration(
+    points: torch.Tensor,
+    landmarks: Union[int, torch.Tensor],
+    function: torch.Tensor,
+    max_dimension: int = -1,
+    points_per_edge: int = 20,
+    near_radius: float = np.sqrt(2),
+    exact: bool = False,
+    batch_size: int = 32,
+    use_triton: bool = False,
+    fps_h: Union[None, int] = None,
+    start_idx: Union[int, None] = 0,
+    verbose = False,
+):
+    """
+    Constructs the sublevel Flood bifiltration from a set of points with scalar values and a set of landmarks.
+
+    Args:
+        points (torch.Tensor):
+            A (N, d) tensor containing the input point set.
+        landmarks (Union[int, torch.Tensor]):
+            Either an integer indicating the number of landmarks to randomly sample from `points` using FPS,
+            or a tensor of integers of shape (N_l) giving the indices of the entries of `points` to take as landmarks,
+            or a tensor of floats of shape (N_l, d) specifying an explicit set of landmarks.
+        function (torch.Tensor):
+            A (N) tensor containing the scalar values on the points.
+        max_dimension (int, optional):
+            The top dimension of the simplices to construct.
+            Defaults to -1 resulting in the dimension of the ambient space.
+        points_per_edge (int, optional):
+            Specifies resolution on simplices used for computing filtration values. Defaults to 20.
+        near_radius (float, optional):
+            Specifies the dilation factor of the minimum enclosing balls to obtain near points. Defaults to sqrt(2).
+        exact (bool, optional):
+            Whether to compute the exact (unmasked) or approximate (masked) sublevel Flood bifiltration.
+        batch_size (int, optional):
+            Size of simplex batches. Defaults to 32.
+        use_triton (bool, optional):
+            Whether to use Triton kernels if available.
+        fps_h (Union[None, int], optional):
+            h parameter (depth of kdtree) that is used for farthest point sampling to
+            select the landmarks. If None, then h is selected based on the size of the
+            point cloud. Defaults to None.
+        start_idx (int | None, optional):
+            If provided, FPS starts from this index in the point cloud. If not,
+            the start index will be randomly picked from the point cloud. Defaults to 0.
+
+    Returns:
+        mp.SimplexTreeMulti
+            Returns a k-critical mp.SimplexTreeMulti containing the sublevel Flood bifiltration.
+    """
+
+    reset()
+
+    if max_dimension == -1:
+        max_dimension = points.shape[1]
+    
+    if isinstance(landmarks, numbers.Integral):
+        if landmarks >= points.shape[0]:
+            landmarks = points
+        else:
+            landmarks = generate_landmarks_FPS(points, landmarks, fps_h, start_idx=start_idx)
+    elif landmarks.device != points.device:
+        raise RuntimeError(f"landmarks.device ({landmarks.device}) != points.device ({points.device})")
+
+    device = points.device
+    use_triton = use_triton and HAS_TRITON and device.type == 'cuda'
+    if use_triton:
+        points = points.float()
+        if torch.is_floating_point(landmarks):
+            landmarks = landmarks.float()
+    dtype = points.dtype
+
+    arg = torch.argsort(function)
+    ranks = torch.empty_like(arg)
+    ranks[arg] = torch.arange(points.size(0), device=device)
+    points = points[arg].contiguous()
+    function = function[arg].contiguous()
+
+    if not torch.is_floating_point(landmarks):
+        landmarks = points[ranks[landmarks]]
+
+    if verbose:
+        elapsed("SORT")
+
+    stree = gudhi.delaunay_complex.DelaunayComplex(landmarks.cpu()).create_simplex_tree(filtration=None)
+    cells_list = []
+    cells_idx  = dict()
+    simplices_list = [[] for _ in range(max_dimension-1)]
+    simplices_cocells = [[] for _ in range(max_dimension-1)]
+    for simplex, _ in stree.get_simplices():
+        dim = len(simplex)-1
+        if dim == max_dimension:
+            cells_idx[tuple(simplex)] = len(cells_list)
+            cells_list.append(simplex)
+        elif 0 < dim < max_dimension:
+            simplices_list[dim-1].append(simplex)
+            cocells = []
+            for cell,_ in stree.get_cofaces(simplex, max_dimension - dim):
+                cocells.append(cells_idx[tuple(cell)])
+            simplices_cocells[dim-1].append(cocells)
+
+    ret_stree = mp.SimplexTreeMulti(num_parameters=2, kcritical=True)
+
+    cells = torch.tensor(cells_list, device=device)
+    simplices_by_dim = [torch.tensor(l, device=device) for l in simplices_list]
+
+    if verbose:
+        elapsed("DELAUNAY")
+
+    simplex_vertices = landmarks[cells]
+    simplex_vertices_cpu = simplex_vertices.cpu().numpy().astype(np.float64, copy=False)
+    Lc, Lr = minimum_enclosing_balls(simplex_vertices_cpu)
+    simplex_radii = near_radius * torch.tensor(Lr, device=device, dtype=dtype)
+    simplex_centers = torch.tensor(Lc, device=device, dtype=dtype)
+
+    if verbose:
+        elapsed("BALLS")
+
+    mask = near_validity(points, simplex_centers, simplex_radii)
+    scc0, mask = vertex_bigrades(points, landmarks, function, cells, mask)
+    if exact:
+        mask = torch.ones_like(mask)
+    if verbose:
+        elapsed("MASK")
+
+    ret_stree.insert_batch(np.arange(len(scc0))[None,:], pad_bigrades(scc0))
+    if verbose:
+        elapsed("ASSIGN_0")
+
+    submasks = [mask_from_adjacency(cocells, mask) for cocells in simplices_cocells]
+    if verbose:
+        elapsed("SUBMASKS")
+
+    simplices_by_dim.append(cells)
+    submasks.append(mask)
+
+    for simplices, mask in zip(simplices_by_dim, submasks):
+        num_simplices, dim = simplices.shape
+
+        weights = generate_grid(points_per_edge, dim-1, device, dtype)
+        simplex_vertices = landmarks[simplices]
+        points_on_simplex = weights.unsqueeze(0) @ simplex_vertices
+
+        if not exact:
+            mask_counts = mask.sum(1, keepdim=False, dtype=torch.int32)
+            indices = torch.sort(mask_counts)[1]
+            mask = mask[indices].contiguous()
+            simplices = simplices[indices].contiguous()
+            points_on_simplex = points_on_simplex[indices].contiguous()
+
+            mask_counts_t = mask_counts[indices]
+            mask_counts = mask_counts[indices].cpu().numpy()
+
+        bigrades = []
+        for k in range(-(-num_simplices // batch_size)):
+            s1, s2 = k*batch_size, min((k+1)*batch_size, num_simplices)
+
+            # b_simplices = simplices[s1:s2]
+            # b_simplices_vertices = landmarks[b_simplices]
+            # b_sample = weights.unsqueeze(0) @ b_simplices_vertices
+            b_sample = points_on_simplex[s1:s2]
+
+            if not exact:
+                b_mask = mask[s1:s2]
+                b_mask_count_max = int(mask_counts[s2 - 1])
+                if use_triton:
+                    b_mask_indices = compact_mask_indices(b_mask, mask_counts_t[s1:s2], b_mask_count_max)
+                else:
+                    b_mask_indices = torch.argsort(b_mask, dim=1, descending=True, stable=True)[:, :b_mask_count_max]
+                b_mask_points = points[b_mask_indices]
+                b_mask_function = function[b_mask_indices]
+            else:
+                b_mask_points = points.repeat(s2-s1, 1, 1)
+                b_mask_function = function.repeat(s2-s1, 1)
+
+            if use_triton:
+                max_diameters = maxcummin(b_mask_points, b_sample)
+            else:
+                inter = torch.cdist(b_mask_points, b_sample)
+                cummin = inter.cummin(dim=1)[0]
+                max_diameters = torch.amax(cummin, dim=2)
+
+            critical = torch.ones_like(max_diameters, dtype=torch.int8)
+            critical[:, 1:] = max_diameters[:, 1:] != max_diameters[:, :-1]
+            critical_count = critical.sum(dim=1)
+            critical_indices = torch.argsort(critical, dim=1, descending=True, stable=True)
+            bigrade_diameters = torch.gather(max_diameters, dim=1, index=critical_indices)
+            bigrade_functions = torch.gather(b_mask_function, dim=1, index=critical_indices)
+
+            b_bigrades = torch.stack((bigrade_diameters, bigrade_functions), dim=-1).cpu().to(torch.float32)
+            critical_count = critical_count.cpu().numpy()
+
+            for i in range(len(b_bigrades)):
+                bigrades.append(b_bigrades[i, :critical_count[i]].numpy())
+
+        ret_stree.insert_batch(simplices.cpu().numpy().T, pad_bigrades(bigrades))
+
+    if verbose:
+        elapsed("LOOP")
+
+    return ret_stree

sublevel_flood-0.1/slflood/triton_mask_topk.py

@@ -0,0 +1,110 @@
+"""
+Triton replacement for:
+
+    b_mask_indices = torch.argsort(b_mask, dim=1, descending=True, stable=True)[:, :K]
+    or
+    b_mask_indices = torch.topk(b_mask, k=K, dim=1, sorted=False)[1]
+
+Algorithm: classic prefix-sum-based stream compaction, done as a single
+sequential pass per row over blocks of the N dimension, carrying a running
+(true_count_so_far, implicit false_count_so_far) across blocks:
+
+  For row-relative position j, let C[j] = number of True entries strictly
+  before j (exclusive prefix sum of the mask up to j).
+    - if mask[j] is True:  its destination column is C[j]
+    - if mask[j] is False: its destination column is true_count + (j - C[j])
+      where true_count is the row's TOTAL true count (known ahead of time -
+      reuse the mask_counts you already compute before the batch loop,
+      don't recompute here).
+
+  Only destinations < K are ever written (True entries always land below
+  K since true_count <= K by construction of K = max over the batch), so
+  most False-entry writes for well-covered rows never happen at all.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_N": 512}, num_warps=4),
+        triton.Config({"BLOCK_N": 1024}, num_warps=4),
+        triton.Config({"BLOCK_N": 1024}, num_warps=8),
+        triton.Config({"BLOCK_N": 2048}, num_warps=8),
+    ],
+    key=["N"],
+)
+@triton.jit(do_not_specialize=["K"])
+def _compact_kernel(
+    mask_ptr, counts_ptr, out_ptr,
+    stride_mb, stride_mn,
+    stride_cb,
+    stride_ob, stride_ok,
+    N, K,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    true_count = tl.load(counts_ptr + row * stride_cb)
+
+    carry_true = 0
+    n_blocks = tl.cdiv(N, BLOCK_N)
+
+    for blk in range(0, n_blocks):
+        n_offsets = blk * BLOCK_N + tl.arange(0, BLOCK_N)
+        n_valid = n_offsets < N
+
+        m = tl.load(
+            mask_ptr + row * stride_mb + n_offsets * stride_mn,
+            mask=n_valid, other=0,
+        ).to(tl.int32)
+
+        incl = tl.cumsum(m, axis=0)   # inclusive prefix sum within this block
+        excl = incl - m               # exclusive prefix sum within this block
+
+        global_true_rank = carry_true + excl                     # C[j], global
+        global_false_rank = n_offsets - global_true_rank          # j - C[j], global
+
+        dest = tl.where(m == 1, global_true_rank, true_count + global_false_rank)
+        store_mask = n_valid & (dest < K)
+
+        tl.store(out_ptr + row * stride_ob + dest * stride_ok, n_offsets, mask=store_mask)
+
+        block_true_count = tl.sum(tl.where(n_valid, m, 0), axis=0)
+        carry_true += block_true_count
+
+
+def compact_mask_indices(mask: torch.Tensor, mask_counts: torch.Tensor, k: int) -> torch.Tensor:
+    """
+    Drop-in replacement for:
+        b_mask_indices = torch.argsort(mask, dim=1, descending=True, stable=True)[:, :K]
+
+    mask:        (rows, N) bool, CUDA
+    mask_counts: (rows,)   int   - true count per row (mask.sum(1)).
+                 Pass the values you've already computed for batch sorting
+                 rather than recomputing them here - they're the same
+                 numbers.
+    k:           output width (b_mask_count_max for this batch)
+
+    Returns (rows, K) int64 tensor of indices, matching stable
+    descending-argsort semantics exactly: True positions in ascending
+    original order first, then enough False positions (also ascending) to
+    fill out to K columns.
+    """
+    assert mask.is_cuda
+    rows, N = mask.shape
+
+    mask_i32 = mask.to(torch.int32).contiguous()
+    counts_i32 = mask_counts.to(torch.int32).contiguous()
+    out = torch.empty((rows, k), device=mask.device, dtype=torch.int32)
+
+    grid = (rows,)
+    _compact_kernel[grid](
+        mask_i32, counts_i32, out,
+        mask_i32.stride(0), mask_i32.stride(1),
+        counts_i32.stride(0),
+        out.stride(0), out.stride(1),
+        N, k,
+    )
+    return out.to(torch.int64)

sublevel_flood-0.1/slflood/triton_maxcummin.py

@@ -0,0 +1,141 @@
+"""
+Triton replacement for:
+
+    inter          = torch.cdist(b_mask_points, b_sample)     # (B, M, G) - materialized
+    cummin         = inter.cummin(dim=1)[0]                   # (B, M, G) - materialized
+    max_diameters  = torch.amax(cummin, dim=2)                # (B, M)    - only this survives
+
+Algorithm (online scan, same structural pattern as flash-attention's
+online softmax):
+
+  For each batch row b and grid point g, maintain a running minimum
+  distance as witness points are added in order (m = 0, 1, 2, ...).
+  After adding point m, the "covering radius" at that prefix length is
+  max_g(running_min[g]). Only that (B, M) result is ever needed downstream
+  so the (B, M, G) intermediate never needs to exist in memory.
+
+  Grid points are split into blocks of BLOCK_G for parallelism (so a run
+  isn't limited to B independent programs, which would badly underuse a
+  GPU with far more SMs than typical batch sizes). Each (batch, g-block)
+  program computes its own partial max per m-step; a final cheap
+  torch.amax over the small g-block axis combines them.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_G": 256}, num_warps=4),
+        triton.Config({"BLOCK_G": 512}, num_warps=4),
+        triton.Config({"BLOCK_G": 512}, num_warps=8),
+        triton.Config({"BLOCK_G": 1024}, num_warps=8),
+        triton.Config({"BLOCK_G": 2048}, num_warps=8),
+    ],
+    key=["G", "D"],
+)
+@triton.jit(do_not_specialize=["M", "G"])
+def _flood_scan_kernel(
+    points_ptr, sample_ptr, out_ptr,
+    stride_pb, stride_pm, stride_pd,
+    stride_sb, stride_sg, stride_sd,
+    stride_ob, stride_om, stride_ogb,
+    M, G,
+    D: tl.constexpr,
+    D_PADDED: tl.constexpr,
+    BLOCK_G: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_g = tl.program_id(1)
+
+    g_offsets = pid_g * BLOCK_G + tl.arange(0, BLOCK_G)
+    g_mask = g_offsets < G
+
+    d_offsets = tl.arange(0, D_PADDED)
+    d_mask = d_offsets < D
+
+    # Load this program's slice of sample points once: (BLOCK_G, D_PADDED)
+    sample_ptrs = (
+        sample_ptr
+        + pid_b * stride_sb
+        + g_offsets[:, None] * stride_sg
+        + d_offsets[None, :] * stride_sd
+    )
+    load_mask = g_mask[:, None] & d_mask[None, :]
+    sample_block = tl.load(sample_ptrs, mask=load_mask, other=0.0)
+
+    running_min = tl.full([BLOCK_G], float("inf"), dtype=tl.float32)
+    out_base = out_ptr + pid_b * stride_ob + pid_g * stride_ogb
+
+    # Sequential scan over M - inherent to the cumulative-min structure.
+    # M is a runtime value (not tl.constexpr), so this compiles to an
+    # actual device-side loop rather than being unrolled at compile time
+    # (same pattern as the K-reduction loop in Triton's matmul tutorial).
+    for i in range(0, M):
+        point_ptrs = (
+            points_ptr
+            + pid_b * stride_pb
+            + i * stride_pm
+            + d_offsets * stride_pd
+        )
+        point_i = tl.load(point_ptrs, mask=d_mask, other=0.0)  # (D_PADDED,)
+
+        diff = sample_block - point_i[None, :]
+        sqdist = tl.sum(diff * diff, axis=1)  # (BLOCK_G,) - padded dims contribute 0
+        dist = tl.sqrt(sqdist)
+
+        running_min = tl.minimum(running_min, dist)
+
+        masked_min = tl.where(g_mask, running_min, float("-inf"))
+        local_max = tl.max(masked_min, axis=0)
+
+        tl.store(out_base + i * stride_om, local_max)
+
+
+def maxcummin(b_mask_points: torch.Tensor, b_sample: torch.Tensor) -> torch.Tensor:
+    """
+    Drop-in replacement for:
+        inter = torch.cdist(b_mask_points, b_sample)
+        cummin = inter.cummin(dim=1)[0]
+        max_diameters = torch.amax(cummin, dim=2)
+
+    b_mask_points: (B, M, D) float32, CUDA
+    b_sample:      (B, G, D) float32, CUDA
+    returns:       (B, M)    float32
+    """
+    assert b_mask_points.is_cuda and b_sample.is_cuda, "kernel is CUDA-only"
+    assert b_mask_points.dtype == torch.float32 and b_sample.dtype == torch.float32
+    B, M, D = b_mask_points.shape
+    B2, G, D2 = b_sample.shape
+    assert B == B2 and D == D2, "batch/dim mismatch between points and sample"
+
+    b_mask_points = b_mask_points.contiguous()
+    b_sample = b_sample.contiguous()
+
+    D_PADDED = triton.next_power_of_2(D)
+
+    def grid(meta):
+        return (B, triton.cdiv(G, meta["BLOCK_G"]))
+
+    # Sized for the smallest BLOCK_G in the autotune list so the buffer
+    # shape doesn't depend on which config gets picked; unused trailing
+    # columns stay -inf and are ignored by the final amax.
+    smallest_block_g = 256
+    max_g_blocks = triton.cdiv(G, smallest_block_g)
+    partial = torch.full(
+        (B, M, max_g_blocks), float("-inf"),
+        device=b_mask_points.device, dtype=torch.float32,
+    )
+
+    _flood_scan_kernel[grid](
+        b_mask_points, b_sample, partial,
+        b_mask_points.stride(0), b_mask_points.stride(1), b_mask_points.stride(2),
+        b_sample.stride(0), b_sample.stride(1), b_sample.stride(2),
+        partial.stride(0), partial.stride(1), partial.stride(2),
+        M, G,
+        D=D, D_PADDED=D_PADDED,
+    )
+
+    return partial.amax(dim=2)

sublevel_flood-0.1/slflood/utils.py

@@ -0,0 +1,94 @@
+import numpy as np
+import torch
+
+
+@torch.no_grad()
+def generate_grid(n: int, dim: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """Generates a grid of points on the unit simplex based on the number of points per edge.
+    Adapted from https://github.com/plus-rkwitt/flooder.
+
+    Args:
+        n (int):
+            Number of points per edge.
+        dim (int):
+            Dimension of the simplex.
+        device (torch.device):
+            Device to create the tensor on.
+        dtype (torch.dtype):
+            Dtype of the tensor.
+
+    Returns:
+        grid (torch.Tensor):
+            Tensor of shape (C, dim + 1), containing the grid points (coordinate weights).
+    """
+
+    combs = torch.combinations(torch.arange(n + dim - 1, device=torch.device('cpu')), r=dim)
+    padded = torch.cat(
+        [
+            torch.full((combs.shape[0], 1), -1, device=torch.device('cpu')),
+            combs,
+            torch.full((combs.shape[0], 1), n + dim - 1, device=torch.device('cpu')),
+        ],
+        dim=1,
+    )  # shape [C, dim + 2]
+    grid = torch.diff(padded, dim=1) - 1  # shape [C, dim + 1]
+    grid_float = torch.empty_like(grid, dtype=dtype)
+    torch.divide(grid, n - 1 , out=grid_float)
+    return grid_float.to(device)
+
+
+@torch.no_grad()
+def mask_from_adjacency(adjacency, mask):
+    """
+    adjacency: list of length n_k, adjacency[i] = list of cell indices adjacent to k-simplex i
+    mask:      torch.Tensor of shape (n_cells, *rest), binary
+
+    Returns: torch.Tensor of shape (n_k, *rest)
+    """
+    rest_shape = mask.shape[1:]
+    n_k = len(adjacency)
+
+    lengths = torch.tensor([len(a) for a in adjacency], dtype=torch.long, device=mask.device)
+    dest_idx = torch.repeat_interleave(torch.arange(n_k, device=mask.device), lengths)
+    src_idx = torch.tensor([c for adj in adjacency for c in adj], dtype=torch.long, device=mask.device)
+
+    values = mask[src_idx]  # (n_edges, *rest)
+
+    idx_scatter = dest_idx.view(-1, *([1] * len(rest_shape))).expand_as(values)
+    out = torch.zeros((n_k,) + rest_shape, dtype=mask.dtype, device=mask.device)
+    out.scatter_reduce_(0, idx_scatter, values, reduce="amax", include_self=True)
+
+    return out
+
+
+def pad_bigrades(bigrades_list):
+    """
+    bigrades_list: list of N arrays, each of shape (k_i, 2), k_i possibly different.
+    Returns: array of shape (N, max_k, 2), padded with (inf, inf).
+    """
+    N = len(bigrades_list)
+    max_k = max(arr.shape[0] for arr in bigrades_list)
+
+    padded = np.full((N, max_k, 2), np.inf, dtype=np.float64)
+    for i, arr in enumerate(bigrades_list):
+        k = arr.shape[0]
+        padded[i, :k, :] = arr
+
+    return padded
+
+
+import time
+
+t = time.time()
+
+def elapsed(s=None):
+    global t
+    if s is not None:
+        print("[timing]", s, int(1000*(time.time() - t)))
+    else:
+        print("[timing]", int(1000*(time.time() - t)))
+    t = time.time()
+
+def reset():
+    global t
+    t = time.time()

sublevel_flood-0.1/tests/test_sklearn.py

@@ -0,0 +1,17 @@
+import pytest
+pytest.importorskip("sklearn")
+
+from slflood.sklearn import SublevelFloodBifiltration
+
+from multipers.ml.mma import FilteredComplex2MMA
+import torch
+from sklearn.pipeline import Pipeline
+
+def test_sklearn():
+    X = torch.rand((10,10000,3))
+
+    pipeline = Pipeline([("slf", SublevelFloodBifiltration(100, 0.2, use_triton=True)),
+                         # ("mma", FilteredComplex2MMA(n_jobs=-1))
+                         ])
+
+    Y = pipeline.fit_transform(X)

sublevel_flood-0.1/tests/test_slflood.py

@@ -0,0 +1,49 @@
+import slflood
+import slflood.data
+
+import multipers as mp
+import pytest
+import torch
+
+# import matplotlib.pyplot as plt
+
+def run(n_pts, n_lms, dim, device):
+    pts, lms_idx, function = slflood.data.noisy_sphere_data(n_pts, n_lms, dim, 0.2, seed=0)
+    if device in ['cuda', 'triton']:
+        if not torch.cuda.is_available():
+            pytest.skip('CUDA not available')
+        else:
+            pts, lms_idx, function = pts.cuda(), lms_idx.cuda(), function.cuda()
+    ff = slflood.slflood_bifiltration(pts, lms_idx, function,
+                                      use_triton=device=='triton')
+    mma_ff = mp.module_approximation(ff)
+    # mma_ff.plot(box=[[0, 0], [1, 1]])
+    # plt.show()
+
+
+n_pts = 5000
+n_lms = 100
+
+
+def test_dim2_cpu():
+    run(n_pts, n_lms, 2, 'cpu')
+
+
+def test_dim2_cuda():
+    run(n_pts, n_lms, 2, 'cuda')
+
+
+def test_dim2_triton():
+    run(n_pts, n_lms, 2, 'triton')
+
+
+def test_dim3_cpu():
+    run(n_pts, n_lms, 3, 'cpu')
+
+
+def test_dim3_cuda():
+    run(n_pts, n_lms, 3, 'cuda')
+
+
+def test_dim3_triton():
+    run(n_pts, n_lms, 3, 'triton')