canns 0.12.7__py3-none-any.whl → 0.13.0__py3-none-any.whl

canns/analyzer/data/asa/config.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ...visualization import PlotConfig
+
+
+@dataclass
+class SpikeEmbeddingConfig:
+    """Configuration for spike train embedding.
+
+    Attributes
+    ----------
+    res : int
+        Time scaling factor that converts seconds to integer bins.
+    dt : int
+        Bin width in the same scaled units as ``res``.
+    sigma : int
+        Gaussian smoothing width (scaled units).
+    smooth : bool
+        Whether to apply temporal smoothing to spike counts.
+    speed_filter : bool
+        Whether to filter by animal speed (requires x/y/t in the input).
+    min_speed : float
+        Minimum speed threshold for ``speed_filter`` (cm/s by convention).
+
+    Examples
+    --------
+    >>> from canns.analyzer.data import SpikeEmbeddingConfig
+    >>> cfg = SpikeEmbeddingConfig(smooth=False, speed_filter=False)
+    >>> cfg.min_speed
+    2.5
+    """
+
+    res: int = 100000
+    dt: int = 1000
+    sigma: int = 5000
+    smooth: bool = True
+    speed_filter: bool = True
+    min_speed: float = 2.5
+
+
+@dataclass
+class TDAConfig:
+    """Configuration for Topological Data Analysis (TDA).
+
+    Attributes
+    ----------
+    dim : int
+        Target PCA dimension before TDA.
+    num_times : int
+        Downsampling stride in time.
+    active_times : int
+        Number of most active time points to keep.
+    k : int
+        Number of neighbors used in denoising.
+    n_points : int
+        Number of points sampled for persistent homology.
+    metric : str
+        Distance metric for point cloud (e.g., "cosine").
+    nbs : int
+        Number of neighbors for distance matrix construction.
+    maxdim : int
+        Maximum homology dimension for persistence.
+    coeff : int
+        Field coefficient for persistent homology.
+    show : bool
+        Whether to show barcode plots.
+    do_shuffle : bool
+        Whether to run shuffle analysis.
+    num_shuffles : int
+        Number of shuffles for null distribution.
+    progress_bar : bool
+        Whether to show progress bars.
+
+    Examples
+    --------
+    >>> from canns.analyzer.data import TDAConfig
+    >>> cfg = TDAConfig(maxdim=1, do_shuffle=False, show=False)
+    >>> cfg.maxdim
+    1
+    """
+
+    dim: int = 6
+    num_times: int = 5
+    active_times: int = 15000
+    k: int = 1000
+    n_points: int = 1200
+    metric: str = "cosine"
+    nbs: int = 800
+    maxdim: int = 1
+    coeff: int = 47
+    show: bool = True
+    do_shuffle: bool = False
+    num_shuffles: int = 1000
+    progress_bar: bool = True
+
+
+@dataclass
+class CANN2DPlotConfig(PlotConfig):
+    """Specialized PlotConfig for CANN2D visualizations.
+
+    Extends :class:`canns.analyzer.visualization.PlotConfig` with fields that
+    control 3D projection and torus animation parameters.
+
+    Examples
+    --------
+    >>> from canns.analyzer.data import CANN2DPlotConfig
+    >>> cfg = CANN2DPlotConfig.for_projection_3d(title="Projection")
+    >>> cfg.zlabel
+    'Component 3'
+    """
+
+    # 3D projection specific parameters
+    zlabel: str = "Component 3"
+    dpi: int = 300
+
+    # Torus animation specific parameters
+    numangsint: int = 51
+    r1: float = 1.5  # Major radius
+    r2: float = 1.0  # Minor radius
+    window_size: int = 300
+    frame_step: int = 5
+    n_frames: int = 20
+
+    @classmethod
+    def for_projection_3d(cls, **kwargs) -> CANN2DPlotConfig:
+        """Create configuration for 3D projection plots.
+
+        Examples
+        --------
+        >>> cfg = CANN2DPlotConfig.for_projection_3d(figsize=(6, 5))
+        >>> cfg.figsize
+        (6, 5)
+        """
+        defaults = {
+            "title": "3D Data Projection",
+            "xlabel": "Component 1",
+            "ylabel": "Component 2",
+            "zlabel": "Component 3",
+            "figsize": (10, 8),
+            "dpi": 300,
+        }
+        defaults.update(kwargs)
+        return cls.for_static_plot(**defaults)
+
+    @classmethod
+    def for_torus_animation(cls, **kwargs) -> CANN2DPlotConfig:
+        """Create configuration for 3D torus bump animations.
+
+        Examples
+        --------
+        >>> cfg = CANN2DPlotConfig.for_torus_animation(fps=10, n_frames=50)
+        >>> cfg.fps, cfg.n_frames
+        (10, 50)
+        """
+        defaults = {
+            "title": "3D Bump on Torus",
+            "figsize": (8, 8),
+            "fps": 5,
+            "repeat": True,
+            "show_progress_bar": True,
+            "numangsint": 51,
+            "r1": 1.5,
+            "r2": 1.0,
+            "window_size": 300,
+            "frame_step": 5,
+            "n_frames": 20,
+        }
+        defaults.update(kwargs)
+        time_steps = kwargs.get("time_steps_per_second", 1000)
+        config = cls.for_animation(time_steps, **defaults)
+        # Add torus-specific attributes
+        config.numangsint = defaults["numangsint"]
+        config.r1 = defaults["r1"]
+        config.r2 = defaults["r2"]
+        config.window_size = defaults["window_size"]
+        config.frame_step = defaults["frame_step"]
+        config.n_frames = defaults["n_frames"]
+        return config
+
+
+# ==================== Constants ====================
+
+
+class Constants:
+    """Constants used throughout CANN2D analysis.
+
+    Examples
+    --------
+    >>> from canns.analyzer.data import Constants
+    >>> Constants.DEFAULT_DPI
+    300
+    """
+
+    DEFAULT_FIGSIZE = (10, 8)
+    DEFAULT_DPI = 300
+    GAUSSIAN_SIGMA_FACTOR = 100
+    SPEED_CONVERSION_FACTOR = 100
+    TIME_CONVERSION_FACTOR = 0.01
+    MULTIPROCESSING_CORES = 4
+
+
+# ==================== Custom Exceptions ====================
+
+
+class CANN2DError(Exception):
+    """Base exception for CANN2D analysis errors.
+
+    Examples
+    --------
+    >>> try:  # doctest: +SKIP
+    ...     raise CANN2DError("boom")
+    ... except CANN2DError:
+    ...     pass
+    """
+
+    pass
+
+
+class DataLoadError(CANN2DError):
+    """Raised when data loading fails.
+
+    Examples
+    --------
+    >>> try:  # doctest: +SKIP
+    ...     raise DataLoadError("missing data")
+    ... except DataLoadError:
+    ...     pass
+    """
+
+    pass
+
+
+class ProcessingError(CANN2DError):
+    """Raised when data processing fails.
+
+    Examples
+    --------
+    >>> try:  # doctest: +SKIP
+    ...     raise ProcessingError("processing failed")
+    ... except ProcessingError:
+    ...     pass
+    """
+
+    pass

canns/analyzer/data/asa/decode.py

@@ -0,0 +1,448 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+from scipy.sparse.linalg import lsmr
+from sklearn import preprocessing
+
+from .config import SpikeEmbeddingConfig
+from .embedding import embed_spike_trains
+
+
+def decode_circular_coordinates(
+    persistence_result: dict[str, Any],
+    spike_data: dict[str, Any],
+    real_ground: bool = True,
+    real_of: bool = True,
+    save_path: str | None = None,
+) -> dict[str, Any]:
+    """
+    Decode circular coordinates (bump positions) from cohomology.
+
+    Parameters
+    ----------
+    persistence_result : dict
+        Output from :func:`canns.analyzer.data.tda_vis`, containing keys:
+        ``persistence``, ``indstemp``, ``movetimes``, ``n_points``.
+    spike_data : dict
+        Spike data dictionary containing ``'spike'``, ``'t'`` and optionally ``'x'``/``'y'``.
+    real_ground : bool
+        Whether x/y/t ground-truth exists (controls whether speed filtering is applied).
+    real_of : bool
+        Whether the experiment is open-field (controls box coordinate handling).
+    save_path : str, optional
+        Path to save decoding results. Defaults to ``Results/spikes_decoding.npz``.
+
+    Returns
+    -------
+    dict
+        Dictionary containing:
+        - ``coords``: decoded coordinates for all timepoints.
+        - ``coordsbox``: decoded coordinates for box timepoints.
+        - ``times``: time indices for ``coords``.
+        - ``times_box``: time indices for ``coordsbox``.
+        - ``centcosall`` / ``centsinall``: cosine/sine centroids.
+
+    Examples
+    --------
+    >>> from canns.analyzer.data import tda_vis, decode_circular_coordinates
+    >>> persistence = tda_vis(embed_spikes, config=tda_cfg)  # doctest: +SKIP
+    >>> decoding = decode_circular_coordinates(persistence, spike_data)  # doctest: +SKIP
+    >>> decoding["coords"].shape  # doctest: +SKIP
+    """
+    ph_classes = [0, 1]  # Decode the ith most persistent cohomology class
+    num_circ = len(ph_classes)
+    dec_tresh = 0.99
+    coeff = 47
+
+    # Extract persistence analysis results
+    persistence = persistence_result["persistence"]
+    indstemp = persistence_result["indstemp"]
+    movetimes = persistence_result["movetimes"]
+    n_points = persistence_result["n_points"]
+
+    diagrams = persistence["dgms"]  # the multiset describing the lives of the persistence classes
+    cocycles = persistence["cocycles"][1]  # the cocycle representatives for the 1-dim classes
+    dists_land = persistence["dperm2all"]  # the pairwise distance between the points
+    births1 = diagrams[1][:, 0]  # the time of birth for the 1-dim classes
+    deaths1 = diagrams[1][:, 1]  # the time of death for the 1-dim classes
+    deaths1[np.isinf(deaths1)] = 0
+    lives1 = deaths1 - births1  # the lifetime for the 1-dim classes
+    iMax = np.argsort(lives1)
+    coords1 = np.zeros((num_circ, len(indstemp)))
+    threshold = births1[iMax[-2]] + (deaths1[iMax[-2]] - births1[iMax[-2]]) * dec_tresh
+
+    for c in ph_classes:
+        cocycle = cocycles[iMax[-(c + 1)]]
+        f, verts = _get_coords(cocycle, threshold, len(indstemp), dists_land, coeff)
+        if len(verts) != len(indstemp):
+            raise ValueError(
+                "Circular coordinate reconstruction returned fewer vertices than sampled points. "
+                "Increase n_points/active_times or use denser data."
+            )
+        coords1[c, :] = f
+
+    # Whether the user-provided dataset has ground-truth x/y/t.
+    if real_ground:
+        sspikes, _, _, _ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=True)
+        )
+    else:
+        sspikes, _, _, _ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=False)
+        )
+
+    num_neurons = sspikes.shape[1]
+    centcosall = np.zeros((num_neurons, 2, n_points))
+    centsinall = np.zeros((num_neurons, 2, n_points))
+    dspk = preprocessing.scale(sspikes[movetimes[indstemp], :])
+
+    for neurid in range(num_neurons):
+        spktemp = dspk[:, neurid].copy()
+        centcosall[neurid, :, :] = np.multiply(np.cos(coords1[:, :] * 2 * np.pi), spktemp)
+        centsinall[neurid, :, :] = np.multiply(np.sin(coords1[:, :] * 2 * np.pi), spktemp)
+
+    # Whether the user-provided dataset has ground-truth x/y/t.
+    if real_ground:
+        sspikes, _, _, _ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=True)
+        )
+        spikes, __, __, __ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=False, speed_filter=True)
+        )
+    else:
+        sspikes, _, _, _ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=False)
+        )
+        spikes, _, _, _ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=False, speed_filter=False)
+        )
+
+    times = np.where(np.sum(spikes > 0, 1) >= 1)[0]
+    dspk = preprocessing.scale(sspikes)
+    sspikes = sspikes[times, :]
+    dspk = dspk[times, :]
+
+    a = np.zeros((len(sspikes[:, 0]), 2, num_neurons))
+    for n in range(num_neurons):
+        a[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centcosall[n, :, :], 1))
+
+    c = np.zeros((len(sspikes[:, 0]), 2, num_neurons))
+    for n in range(num_neurons):
+        c[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centsinall[n, :, :], 1))
+
+    mtot2 = np.sum(c, 2)
+    mtot1 = np.sum(a, 2)
+    coords = np.arctan2(mtot2, mtot1) % (2 * np.pi)
+
+    # Whether the dataset comes from a real open-field (OF) environment.
+    if real_of:
+        coordsbox = coords.copy()
+        times_box = times.copy()
+    else:
+        sspikes, _, _, _ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=True, speed_filter=True)
+        )
+        spikes, __, __, __ = embed_spike_trains(
+            spike_data, config=SpikeEmbeddingConfig(smooth=False, speed_filter=True)
+        )
+        dspk = preprocessing.scale(sspikes)
+        times_box = np.where(np.sum(spikes > 0, 1) >= 1)[0]
+        dspk = dspk[times_box, :]
+
+        a = np.zeros((len(times_box), 2, num_neurons))
+        for n in range(num_neurons):
+            a[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centcosall[n, :, :], 1))
+
+        c = np.zeros((len(times_box), 2, num_neurons))
+        for n in range(num_neurons):
+            c[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centsinall[n, :, :], 1))
+
+        mtot2 = np.sum(c, 2)
+        mtot1 = np.sum(a, 2)
+        coordsbox = np.arctan2(mtot2, mtot1) % (2 * np.pi)
+
+    # Prepare results dictionary
+    results = {
+        "coords": coords,
+        "coordsbox": coordsbox,
+        "times": times,
+        "times_box": times_box,
+        "centcosall": centcosall,
+        "centsinall": centsinall,
+    }
+
+    # Save results
+    if save_path is None:
+        os.makedirs("Results", exist_ok=True)
+        save_path = "Results/spikes_decoding.npz"
+
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    np.savez_compressed(save_path, **results)
+
+    return results
+
+
+def decode_circular_coordinates1(
+    persistence_result: dict[str, Any],
+    spike_data: dict[str, Any],
+    save_path: str | None = None,
+) -> dict[str, Any]:
+    """Legacy helper kept for backward compatibility."""
+    ph_classes = [0, 1]  # Decode the ith most persistent cohomology class
+    num_circ = len(ph_classes)
+    dec_tresh = 0.99
+    coeff = 47
+
+    # Extract persistence analysis results
+    persistence = persistence_result["persistence"]
+    indstemp = persistence_result["indstemp"]
+    movetimes = persistence_result["movetimes"]
+    n_points = persistence_result["n_points"]
+
+    diagrams = persistence["dgms"]  # the multiset describing the lives of the persistence classes
+    cocycles = persistence["cocycles"][1]  # the cocycle representatives for the 1-dim classes
+    dists_land = persistence["dperm2all"]  # the pairwise distance between the points
+    births1 = diagrams[1][:, 0]  # the time of birth for the 1-dim classes
+    deaths1 = diagrams[1][:, 1]  # the time of death for the 1-dim classes
+    deaths1[np.isinf(deaths1)] = 0
+    lives1 = deaths1 - births1  # the lifetime for the 1-dim classes
+    iMax = np.argsort(lives1)
+    coords1 = np.zeros((num_circ, len(indstemp)))
+    threshold = births1[iMax[-2]] + (deaths1[iMax[-2]] - births1[iMax[-2]]) * dec_tresh
+
+    for c in ph_classes:
+        cocycle = cocycles[iMax[-(c + 1)]]
+        f, verts = _get_coords(cocycle, threshold, len(indstemp), dists_land, coeff)
+        if len(verts) != len(indstemp):
+            raise ValueError(
+                "Circular coordinate reconstruction returned fewer vertices than sampled points. "
+                "Increase n_points/active_times or use denser data."
+            )
+        coords1[c, :] = f
+
+    sspikes = spike_data["spike"]
+    num_neurons = sspikes.shape[1]
+    centcosall = np.zeros((num_neurons, 2, n_points))
+    centsinall = np.zeros((num_neurons, 2, n_points))
+    dspk = preprocessing.scale(sspikes[movetimes[indstemp], :])
+
+    for neurid in range(num_neurons):
+        spktemp = dspk[:, neurid].copy()
+        centcosall[neurid, :, :] = np.multiply(np.cos(coords1[:, :] * 2 * np.pi), spktemp)
+        centsinall[neurid, :, :] = np.multiply(np.sin(coords1[:, :] * 2 * np.pi), spktemp)
+
+    times = np.where(np.sum(sspikes > 0, 1) >= 1)[0]
+    dspk = preprocessing.scale(sspikes)
+    sspikes = sspikes[times, :]
+    dspk = dspk[times, :]
+
+    a = np.zeros((len(sspikes[:, 0]), 2, num_neurons))
+    for n in range(num_neurons):
+        a[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centcosall[n, :, :], 1))
+
+    c = np.zeros((len(sspikes[:, 0]), 2, num_neurons))
+    for n in range(num_neurons):
+        c[:, :, n] = np.multiply(dspk[:, n : n + 1], np.sum(centsinall[n, :, :], 1))
+
+    mtot2 = np.sum(c, 2)
+    mtot1 = np.sum(a, 2)
+    coords = np.arctan2(mtot2, mtot1) % (2 * np.pi)
+
+    coordsbox = coords.copy()
+    times_box = times.copy()
+
+    # Prepare results dictionary
+    results = {
+        "coords": coords,
+        "coordsbox": coordsbox,
+        "times": times,
+        "times_box": times_box,
+        "centcosall": centcosall,
+        "centsinall": centsinall,
+    }
+
+    # Save results
+    if save_path is None:
+        os.makedirs("Results", exist_ok=True)
+        save_path = "Results/spikes_decoding.npz"
+
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    np.savez_compressed(save_path, **results)
+
+    return results
+
+
+def decode_circular_coordinates_multi(
+    persistence_result: dict,
+    spike_data: dict,
+    save_path: str | None = None,
+    num_circ: int = 2,  # Number of H1 cocycles/circular coordinates to decode
+) -> dict:
+    """Decode multiple circular coordinates from TDA persistence.
+
+    Parameters
+    ----------
+    persistence_result : dict
+        Output from :func:`canns.analyzer.data.tda_vis`, containing keys:
+        ``persistence``, ``indstemp``, ``movetimes``, ``n_points``.
+    spike_data : dict
+        Spike data dictionary containing ``'spike'``, ``'t'`` and optionally ``'x'``/``'y'``.
+    save_path : str, optional
+        Path to save decoding results. Defaults to ``Results/spikes_decoding.npz``.
+    num_circ : int
+        Number of H1 cocycles/circular coordinates to decode.
+
+    Returns
+    -------
+    dict
+        Dictionary with ``coords``, ``coordsbox``, ``times``, ``times_box`` and centroid terms.
+
+    Examples
+    --------
+    >>> decoding = decode_circular_coordinates_multi(persistence, spike_data, num_circ=2)  # doctest: +SKIP
+    >>> decoding["coords"].shape  # doctest: +SKIP
+    """
+    from sklearn import preprocessing
+
+    dec_tresh = 0.99
+    coeff = 47
+
+    persistence = persistence_result["persistence"]
+    indstemp = persistence_result["indstemp"]
+    movetimes = persistence_result["movetimes"]
+    n_points = persistence_result["n_points"]
+
+    diagrams = persistence["dgms"]
+    cocycles = persistence["cocycles"][1]
+    dists_land = persistence["dperm2all"]
+
+    births1 = diagrams[1][:, 0]
+    deaths1 = diagrams[1][:, 1]
+    deaths1[np.isinf(deaths1)] = 0
+    lives1 = deaths1 - births1
+
+    if lives1.size < num_circ or len(cocycles) < num_circ:
+        raise ValueError(
+            f"Requested num_circ={num_circ}, but only {lives1.size} H1 feature(s) are available. "
+            "This usually means the chosen time window is too short, the data are too sparse, "
+            "or the embedding parameters are not appropriate."
+        )
+
+    iMax = np.argsort(lives1)
+    coords1 = np.zeros((num_circ, len(indstemp)))
+
+    for i in range(num_circ):
+        idx = iMax[-(i + 1)]
+        threshold = births1[idx] + (deaths1[idx] - births1[idx]) * dec_tresh
+        cocycle = cocycles[idx]
+        f, verts = _get_coords(cocycle, threshold, len(indstemp), dists_land, coeff)
+        if len(verts) != len(indstemp):
+            raise ValueError(
+                "Circular coordinate reconstruction returned fewer vertices than sampled points. "
+                "Increase n_points/active_times or use denser data."
+            )
+        coords1[i, :] = f
+
+    sspikes = spike_data["spike"]
+    num_neurons = sspikes.shape[1]
+
+    centcosall = np.zeros((num_neurons, num_circ, n_points))
+    centsinall = np.zeros((num_neurons, num_circ, n_points))
+    dspk = preprocessing.scale(sspikes[movetimes[indstemp], :])
+
+    for n in range(num_neurons):
+        spktemp = dspk[:, n].copy()
+        centcosall[n, :, :] = np.multiply(np.cos(coords1 * 2 * np.pi), spktemp)
+        centsinall[n, :, :] = np.multiply(np.sin(coords1 * 2 * np.pi), spktemp)
+
+    times = np.where(np.sum(sspikes > 0, 1) >= 1)[0]
+    dspk = preprocessing.scale(sspikes)
+    sspikes = sspikes[times, :]
+    dspk = dspk[times, :]
+
+    a = np.zeros((len(sspikes), num_circ, num_neurons))
+    c = np.zeros((len(sspikes), num_circ, num_neurons))
+
+    for n in range(num_neurons):
+        a[:, :, n] = dspk[:, n : n + 1] * np.sum(centcosall[n, :, :], axis=1)
+        c[:, :, n] = dspk[:, n : n + 1] * np.sum(centsinall[n, :, :], axis=1)
+
+    mtot1 = np.sum(a, 2)
+    mtot2 = np.sum(c, 2)
+    coords = np.arctan2(mtot2, mtot1) % (2 * np.pi)
+
+    results = {
+        "coords": coords,
+        "coordsbox": coords.copy(),
+        "times": times,
+        "times_box": times.copy(),
+        "centcosall": centcosall,
+        "centsinall": centsinall,
+    }
+
+    if save_path is None:
+        os.makedirs("Results", exist_ok=True)
+        save_path = "Results/spikes_decoding.npz"
+
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        np.savez_compressed(save_path, **results)
+    return results
+
+
+def _get_coords(cocycle, threshold, num_sampled, dists, coeff):
+    """
+    Reconstruct circular coordinates from cocycle information.
+
+    Parameters:
+        cocycle (ndarray): Persistent cocycle representative.
+        threshold (float): Maximum allowable edge distance.
+        num_sampled (int): Number of sampled points.
+        dists (ndarray): Pairwise distance matrix.
+        coeff (int): Finite field modulus for cohomology.
+
+    Returns:
+        f (ndarray): Circular coordinate values (in [0,1]).
+        verts (ndarray): Indices of used vertices.
+    """
+    zint = np.where(coeff - cocycle[:, 2] < cocycle[:, 2])
+    cocycle[zint, 2] = cocycle[zint, 2] - coeff
+    d = np.zeros((num_sampled, num_sampled))
+    d[np.tril_indices(num_sampled)] = np.nan
+    d[cocycle[:, 1], cocycle[:, 0]] = cocycle[:, 2]
+    d[dists > threshold] = np.nan
+    d[dists == 0] = np.nan
+    edges = np.where(~np.isnan(d))
+    verts = np.array(np.unique(edges))
+    num_edges = np.shape(edges)[1]
+    num_verts = np.size(verts)
+    values = d[edges]
+    A = np.zeros((num_edges, num_verts), dtype=int)
+    v1 = np.zeros((num_edges, 2), dtype=int)
+    v2 = np.zeros((num_edges, 2), dtype=int)
+    for i in range(num_edges):
+        # Extract scalar indices from np.where results
+        idx1 = np.where(verts == edges[0][i])[0]
+        idx2 = np.where(verts == edges[1][i])[0]
+
+        # Handle case where np.where returns multiple matches (shouldn't happen in valid data)
+        if len(idx1) > 0:
+            v1[i, :] = [i, idx1[0]]
+        else:
+            raise ValueError(f"No vertex found for edge {edges[0][i]}")
+
+        if len(idx2) > 0:
+            v2[i, :] = [i, idx2[0]]
+        else:
+            raise ValueError(f"No vertex found for edge {edges[1][i]}")
+
+    A[v1[:, 0], v1[:, 1]] = -1
+    A[v2[:, 0], v2[:, 1]] = 1
+
+    L = np.ones((num_edges,))
+    Aw = A * np.sqrt(L[:, np.newaxis])
+    Bw = values * np.sqrt(L)
+    f = lsmr(Aw, Bw)[0] % 1
+    return f, verts