canns 0.13.2__py3-none-any.whl → 0.14.0__py3-none-any.whl

canns/analyzer/data/cell_classification/core/grid_modules_leiden.py

@@ -0,0 +1,288 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+
+def _base_mask(shape: tuple[int, int], center_bins: int = 2) -> np.ndarray:
+    """Create a shared mask: center disk + outside circle (corners)."""
+    h, w = shape
+    cy = (h - 1) / 2.0
+    cx = (w - 1) / 2.0
+    yy, xx = np.ogrid[:h, :w]
+    rr = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2)
+
+    # outer circle radius: half-min dimension
+    outer_r = min(cy, cx)
+    mask_outer = rr > outer_r
+
+    mask_center = rr <= float(center_bins)
+    return mask_outer | mask_center
+
+
+def _safe_corr(a: np.ndarray, b: np.ndarray) -> float:
+    a = np.asarray(a).ravel()
+    b = np.asarray(b).ravel()
+    if a.size != b.size or a.size == 0:
+        return 0.0
+    a = a - a.mean()
+    b = b - b.mean()
+    da = np.sqrt((a * a).mean())
+    db = np.sqrt((b * b).mean())
+    if da <= 1e-12 or db <= 1e-12:
+        return 0.0
+    return float((a * b).mean() / (da * db))
+
+
+def _vectorize_autocorrs(
+    autocorrs: np.ndarray, center_bins: int = 2
+) -> tuple[np.ndarray, np.ndarray]:
+    """Vectorize autocorrs into point-cloud matrix X and return the shared mask."""
+    if autocorrs.ndim != 3:
+        raise ValueError(f"autocorrs must be (N,H,W), got {autocorrs.shape}")
+    N, H, W = autocorrs.shape
+    mask = _base_mask((H, W), center_bins=center_bins)
+    # Use shared mask only; replace NaN with 0 so dimensions match across cells.
+    ac = np.nan_to_num(autocorrs, nan=0.0, posinf=0.0, neginf=0.0)
+    X = ac[:, ~mask]  # (N, n_features)
+    return X.astype(np.float32, copy=False), mask
+
+
+def _build_knn_graph(X: np.ndarray, k: int = 30, metric: str = "manhattan"):
+    """Return an igraph graph built from kNN with edge weights."""
+    try:
+        from sklearn.neighbors import NearestNeighbors
+    except Exception as e:
+        raise ImportError(f"scikit-learn is required for kNN graph: {e}") from e
+
+    try:
+        import igraph as ig
+    except Exception as e:
+        raise ImportError(f"python-igraph is required for Leiden clustering: {e}") from e
+
+    N = X.shape[0]
+    k_eff = min(max(int(k), 1), max(N - 1, 1))
+
+    nbrs = NearestNeighbors(n_neighbors=k_eff + 1, metric=metric)
+    nbrs.fit(X)
+    dist, ind = nbrs.kneighbors(X, return_distance=True)
+
+    edges = []
+    weights = []
+    eps = 1e-6
+    seen = set()
+    for i in range(N):
+        for jj in range(1, k_eff + 1):
+            j = int(ind[i, jj])
+            if i == j:
+                continue
+            a, b = (i, j) if i < j else (j, i)
+            if (a, b) in seen:
+                continue
+            seen.add((a, b))
+            d = float(dist[i, jj])
+            w = 1.0 / (d + eps)
+            edges.append((a, b))
+            weights.append(w)
+
+    g = ig.Graph(n=N, edges=edges, directed=False)
+    g.es["weight"] = weights
+    return g
+
+
+def _leiden_membership(graph, resolution: float = 1.0) -> np.ndarray:
+    """Run Leiden and return membership array."""
+    try:
+        import leidenalg
+    except Exception as e:
+        raise ImportError(f"leidenalg is required for Leiden clustering: {e}") from e
+
+    part = leidenalg.find_partition(
+        graph,
+        leidenalg.RBConfigurationVertexPartition,
+        weights=graph.es["weight"] if "weight" in graph.es.attributes() else None,
+        resolution_parameter=float(resolution),
+    )
+    return np.asarray(part.membership, dtype=int)
+
+
+class _DSU:
+    def __init__(self, n: int):
+        self.p = list(range(n))
+        self.r = [0] * n
+
+    def find(self, a: int) -> int:
+        while self.p[a] != a:
+            self.p[a] = self.p[self.p[a]]
+            a = self.p[a]
+        return a
+
+    def union(self, a: int, b: int):
+        ra, rb = self.find(a), self.find(b)
+        if ra == rb:
+            return
+        if self.r[ra] < self.r[rb]:
+            ra, rb = rb, ra
+        self.p[rb] = ra
+        if self.r[ra] == self.r[rb]:
+            self.r[ra] += 1
+
+
+def identify_grid_modules_and_stats(
+    autocorrs: np.ndarray,
+    *,
+    gridness_analyzer,
+    center_bins: int = 2,
+    k: int = 30,
+    resolution: float = 1.0,
+    score_thr: float = 0.3,
+    consistency_thr: float = 0.5,
+    min_cells: int = 10,
+    merge_corr_thr: float = 0.7,
+    metric: str = "manhattan",
+) -> dict[str, Any]:
+    """Identify grid modules with Leiden clustering on autocorrelogram point cloud.
+
+    Parameters
+    ----------
+    autocorrs : np.ndarray
+        Array of shape (N, H, W).
+    gridness_analyzer :
+        An instance that provides compute_gridness_score(autocorr)->GridnessResult.
+    center_bins : int
+        Radius (in bins) to mask around center peak.
+    k : int
+        Neighbors for kNN graph.
+    resolution : float
+        Leiden resolution parameter.
+    score_thr, consistency_thr, min_cells, merge_corr_thr
+        Module acceptance and merging thresholds.
+
+    Returns
+    -------
+    dict with keys:
+        module_id (N,), cluster_id (N,), modules (list of dict), params
+    """
+    if autocorrs.ndim != 3:
+        raise ValueError(f"autocorrs must be (N,H,W). Got {autocorrs.shape}")
+    N, H, W = autocorrs.shape
+
+    X, mask = _vectorize_autocorrs(autocorrs, center_bins=int(center_bins))
+    g = _build_knn_graph(X, k=int(k), metric=str(metric))
+    cluster_id = _leiden_membership(g, resolution=float(resolution))
+
+    # group members
+    clusters: dict[int, np.ndarray] = {}
+    for cid in np.unique(cluster_id):
+        clusters[int(cid)] = np.where(cluster_id == cid)[0]
+
+    base_mask = mask  # shared
+    ac = np.nan_to_num(autocorrs, nan=0.0, posinf=0.0, neginf=0.0)
+
+    cluster_stats = {}
+    candidate_cids = []
+
+    # compute cluster metrics
+    for cid, idxs in clusters.items():
+        if idxs.size == 0:
+            continue
+        avg = ac[idxs].mean(axis=0)
+        med = np.median(ac[idxs], axis=0)
+
+        gr = gridness_analyzer.compute_gridness_score(med)
+        grid_score = float(getattr(gr, "score", getattr(gr, "grid_score", np.nan)))
+
+        flat_avg = avg[~base_mask]
+        cors = []
+        for i in idxs:
+            c = _safe_corr(flat_avg, ac[i][~base_mask])
+            cors.append(c)
+        consistency = float(np.median(cors)) if len(cors) else 0.0
+
+        cluster_stats[cid] = {
+            "cid": cid,
+            "size": int(idxs.size),
+            "grid_score": grid_score,
+            "consistency": consistency,
+            "gridness_result": gr,
+            "avg_autocorr": avg,
+        }
+
+        if (
+            (grid_score > float(score_thr))
+            and (consistency > float(consistency_thr))
+            and (idxs.size >= int(min_cells))
+        ):
+            candidate_cids.append(cid)
+
+    # Merge candidate clusters based on avg autocorr correlation
+    cand = list(candidate_cids)
+    dsu = _DSU(len(cand))
+    for i in range(len(cand)):
+        for j in range(i + 1, len(cand)):
+            ci, cj = cand[i], cand[j]
+            ai = cluster_stats[ci]["avg_autocorr"][~base_mask]
+            aj = cluster_stats[cj]["avg_autocorr"][~base_mask]
+            corr = _safe_corr(ai, aj)
+            if corr > float(merge_corr_thr):
+                dsu.union(i, j)
+
+    # Build merged modules
+    root_to_members: dict[int, list[int]] = {}
+    for idx, cid in enumerate(cand):
+        r = dsu.find(idx)
+        root_to_members.setdefault(r, []).append(cid)
+
+    modules = []
+    module_id = np.full((N,), -1, dtype=int)
+
+    # assign module ids in stable order
+    roots = sorted(root_to_members.keys())
+    for mid, r in enumerate(roots):
+        member_cids = root_to_members[r]
+        member_idxs = (
+            np.concatenate([clusters[c] for c in member_cids])
+            if member_cids
+            else np.array([], dtype=int)
+        )
+
+        # compute module-level stats (median over clusters)
+        gs = [cluster_stats[c]["grid_score"] for c in member_cids]
+        cs = [cluster_stats[c]["consistency"] for c in member_cids]
+        grid_score = float(np.median(gs)) if gs else float("nan")
+        consistency = float(np.median(cs)) if cs else float("nan")
+
+        module_id[member_idxs] = mid
+
+        modules.append(
+            {
+                "module_id": mid,
+                "clusters": member_cids,
+                "indices": member_idxs.astype(int),
+                "size": int(member_idxs.size),
+                "grid_score": grid_score,
+                "consistency": consistency,
+            }
+        )
+
+    n_grid_cells = int((module_id != -1).sum())
+    out = {
+        "cluster_id": cluster_id.astype(int),
+        "module_id": module_id.astype(int),
+        "n_units": int(N),
+        "n_grid_cells": int(n_grid_cells),
+        "n_modules": int(len(modules)),
+        "modules": modules,
+        "params": {
+            "center_bins": int(center_bins),
+            "k": int(k),
+            "resolution": float(resolution),
+            "score_thr": float(score_thr),
+            "consistency_thr": float(consistency_thr),
+            "min_cells": int(min_cells),
+            "merge_corr_thr": float(merge_corr_thr),
+            "metric": str(metric),
+        },
+    }
+    return out

canns/analyzer/data/cell_classification/core/head_direction.py

@@ -0,0 +1,347 @@
+"""
+Head Direction Cell Classification
+
+Implementation of head direction cell identification based on Mean Vector Length (MVL).
+
+Based on MATLAB code from the sweeps analysis pipeline.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from ..utils.circular_stats import circ_mean, circ_r, circ_rtest
+
+
+@dataclass
+class HDCellResult:
+    """
+    Results from head direction cell classification.
+
+    Attributes
+    ----------
+    is_hd : bool
+        Whether the cell is classified as a head direction cell
+    mvl_hd : float
+        Mean Vector Length for head direction tuning
+    preferred_direction : float
+        Preferred head direction in radians
+    mvl_theta : float or None
+        Mean Vector Length for theta phase tuning (if provided)
+    tuning_curve : tuple
+        Tuple of (bin_centers, firing_rates)
+    rayleigh_p : float
+        P-value from Rayleigh test for non-uniformity
+    """
+
+    is_hd: bool
+    mvl_hd: float
+    preferred_direction: float
+    mvl_theta: float | None
+    tuning_curve: tuple[np.ndarray, np.ndarray]
+    rayleigh_p: float
+
+
+class HeadDirectionAnalyzer:
+    """
+    Analyzer for classifying head direction cells based on directional tuning.
+
+    Head direction cells fire when the animal's head points in a specific direction.
+    Classification is based on the strength of directional tuning measured by
+    Mean Vector Length (MVL).
+
+    Parameters
+    ----------
+    mvl_hd_threshold : float, optional
+        MVL threshold for head direction. Default is 0.4 (strict).
+        Use 0.2 for looser threshold.
+    mvl_theta_threshold : float, optional
+        MVL threshold for theta phase modulation. Default is 0.3.
+    strict_mode : bool, optional
+        If True, requires both HD and theta criteria. Default is True.
+    n_bins : int, optional
+        Number of directional bins for tuning curve. Default is 60 (6° bins).
+
+    Examples
+    --------
+    >>> analyzer = HeadDirectionAnalyzer(mvl_hd_threshold=0.4, strict_mode=True)
+    >>> result = analyzer.classify_hd_cell(spike_times, head_directions, time_stamps)
+    >>> print(f"Is HD cell: {result.is_hd}")
+    >>> print(f"MVL: {result.mvl_hd:.3f}")
+    >>> print(f"Preferred direction: {np.rad2deg(result.preferred_direction):.1f}°")
+
+    Notes
+    -----
+    Based on MATLAB classification from fig2.m and plotSwsExample.m:
+    - Strict: MVL_hd > 0.4 AND MVL_theta > 0.3
+    - Loose: MVL_hd > 0.2 AND MVL_theta > 0.3
+
+    References
+    ----------
+    Classification thresholds follow standard conventions in head direction
+    cell literature and the CircStat toolbox.
+    """
+
+    def __init__(
+        self,
+        mvl_hd_threshold: float = 0.4,
+        mvl_theta_threshold: float = 0.3,
+        strict_mode: bool = True,
+        n_bins: int = 60,
+    ):
+        self.mvl_hd_threshold = mvl_hd_threshold
+        self.mvl_theta_threshold = mvl_theta_threshold
+        self.strict_mode = strict_mode
+        self.n_bins = n_bins
+
+    def classify_hd_cell(
+        self,
+        spike_times: np.ndarray,
+        head_directions: np.ndarray,
+        time_stamps: np.ndarray,
+        theta_phases: np.ndarray | None = None,
+    ) -> HDCellResult:
+        """
+        Classify a cell as head direction cell based on MVL thresholds.
+
+        Parameters
+        ----------
+        spike_times : np.ndarray
+            Spike times in seconds
+        head_directions : np.ndarray
+            Head direction at each time point (radians)
+        time_stamps : np.ndarray
+            Time stamps corresponding to head_directions (seconds)
+        theta_phases : np.ndarray, optional
+            Theta phase at each time point (radians). If None, theta
+            criterion is not checked.
+
+        Returns
+        -------
+        result : HDCellResult
+            Classification result with MVL, preferred direction, and tuning curve
+
+        Examples
+        --------
+        >>> # Simulate a head direction cell
+        >>> time_stamps = np.linspace(0, 100, 10000)
+        >>> head_directions = np.linspace(0, 20*np.pi, 10000) % (2*np.pi) - np.pi
+        >>> preferred_dir = 0.5
+        >>> spike_times = time_stamps[np.abs(head_directions - preferred_dir) < 0.3]
+        >>> result = analyzer.classify_hd_cell(spike_times, head_directions, time_stamps)
+        """
+        # Compute directional tuning curve
+        bin_centers, firing_rates, occupancy = self.compute_tuning_curve(
+            spike_times, head_directions, time_stamps
+        )
+
+        # Compute MVL for head direction
+        mvl_hd = self.compute_mvl(bin_centers, weights=firing_rates)
+
+        # Compute preferred direction
+        preferred_direction = circ_mean(bin_centers, w=firing_rates)
+
+        # Rayleigh test for non-uniformity
+        rayleigh_p = circ_rtest(bin_centers, w=firing_rates)
+
+        # Compute MVL for theta phase if provided
+        mvl_theta = None
+        if theta_phases is not None:
+            # Get theta phases at spike times
+            spike_theta = np.interp(spike_times, time_stamps, theta_phases)
+            mvl_theta = self.compute_mvl(spike_theta)
+
+        # Classification logic
+        is_hd = mvl_hd > self.mvl_hd_threshold
+
+        if self.strict_mode and theta_phases is not None:
+            # Strict mode: require both HD and theta criteria
+            is_hd = is_hd and (mvl_theta > self.mvl_theta_threshold)
+        elif self.strict_mode and theta_phases is None:
+            # If strict mode but no theta data, just use HD threshold
+            pass
+
+        # Create result
+        result = HDCellResult(
+            is_hd=is_hd,
+            mvl_hd=mvl_hd,
+            preferred_direction=preferred_direction,
+            mvl_theta=mvl_theta,
+            tuning_curve=(bin_centers, firing_rates),
+            rayleigh_p=rayleigh_p,
+        )
+
+        return result
+
+    def compute_tuning_curve(
+        self,
+        spike_times: np.ndarray,
+        head_directions: np.ndarray,
+        time_stamps: np.ndarray,
+        n_bins: int | None = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Compute directional tuning curve.
+
+        Parameters
+        ----------
+        spike_times : np.ndarray
+            Spike times in seconds
+        head_directions : np.ndarray
+            Head direction at each time point (radians)
+        time_stamps : np.ndarray
+            Time stamps corresponding to head_directions (seconds)
+        n_bins : int, optional
+            Number of bins. If None, uses self.n_bins.
+
+        Returns
+        -------
+        bin_centers : np.ndarray
+            Center of each directional bin (radians)
+        firing_rates : np.ndarray
+            Firing rate in each bin (Hz)
+        occupancy : np.ndarray
+            Time spent in each bin (seconds)
+
+        Examples
+        --------
+        >>> bins, rates, occ = analyzer.compute_tuning_curve(
+        ...     spike_times, head_directions, time_stamps
+        ... )
+        >>> # Plot polar tuning curve
+        >>> import matplotlib.pyplot as plt
+        >>> ax = plt.subplot(111, projection='polar')
+        >>> ax.plot(bins, rates)
+        """
+        if n_bins is None:
+            n_bins = self.n_bins
+
+        # Define bin edges
+        bin_edges = np.linspace(-np.pi, np.pi, n_bins + 1)
+        bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
+
+        # Compute occupancy (time spent in each bin)
+        dt = np.median(np.diff(time_stamps))  # Sampling interval
+        hd_binned = np.digitize(head_directions, bin_edges) - 1  # 0-indexed
+        hd_binned = np.clip(hd_binned, 0, n_bins - 1)  # Handle edge cases
+
+        occupancy = np.bincount(hd_binned, minlength=n_bins) * dt
+
+        # Count spikes in each bin
+        # Interpolate HD at spike times
+        spike_hd = np.interp(spike_times, time_stamps, head_directions)
+        spike_bins = np.digitize(spike_hd, bin_edges) - 1
+        spike_bins = np.clip(spike_bins, 0, n_bins - 1)
+
+        spike_counts = np.bincount(spike_bins, minlength=n_bins)
+
+        # Compute firing rates
+        firing_rates = np.zeros(n_bins)
+        valid = occupancy > 0
+        firing_rates[valid] = spike_counts[valid] / occupancy[valid]
+
+        return bin_centers, firing_rates, occupancy
+
+    def compute_mvl(self, angles: np.ndarray, weights: np.ndarray | None = None) -> float:
+        """
+        Compute Mean Vector Length (MVL).
+
+        The MVL is a measure of circular variance, ranging from 0 (uniform
+        distribution) to 1 (concentrated distribution).
+
+        Parameters
+        ----------
+        angles : np.ndarray
+            Angles in radians
+        weights : np.ndarray, optional
+            Weights for each angle (e.g., firing rates). If None, uniform weights.
+
+        Returns
+        -------
+        mvl : float
+            Mean vector length
+
+        Examples
+        --------
+        >>> # Concentrated distribution
+        >>> angles = np.random.normal(0, 0.1, 100)
+        >>> mvl = analyzer.compute_mvl(angles)
+        >>> print(f"MVL: {mvl:.3f}")  # Should be close to 1
+
+        >>> # Uniform distribution
+        >>> angles = np.random.uniform(-np.pi, np.pi, 100)
+        >>> mvl = analyzer.compute_mvl(angles)
+        >>> print(f"MVL: {mvl:.3f}")  # Should be close to 0
+
+        Notes
+        -----
+        Uses the circ_r function from circular statistics utilities.
+        """
+        return circ_r(angles, w=weights)
+
+
+if __name__ == "__main__":
+    print("Testing HeadDirectionAnalyzer...")
+
+    # Simulate a head direction cell
+    print("\nSimulating head direction cell...")
+    time_stamps = np.linspace(0, 100, 10000)  # 100 seconds, 100 Hz sampling
+    dt = time_stamps[1] - time_stamps[0]
+
+    # Animal rotates and explores
+    angular_velocity = 0.5  # rad/s average
+    head_directions = np.cumsum(np.random.randn(len(time_stamps)) * angular_velocity * dt)
+    head_directions = np.arctan2(
+        np.sin(head_directions), np.cos(head_directions)
+    )  # Wrap to [-π, π]
+
+    # Cell fires preferentially at 0.5 radians (~28°)
+    preferred_dir = 0.5
+    tuning_width = 0.5  # radians (~28° width)
+
+    # Generate spikes based on von Mises tuning
+    from scipy.stats import vonmises
+
+    firing_prob = vonmises.pdf(head_directions - preferred_dir, kappa=1 / tuning_width**2)
+    firing_prob = firing_prob / firing_prob.max() * 0.1  # Max 10% per bin
+
+    # Poisson spike generation
+    spike_mask = np.random.rand(len(time_stamps)) < firing_prob
+    spike_times = time_stamps[spike_mask]
+
+    print(f"Generated {len(spike_times)} spikes")
+    print(f"Mean firing rate: {len(spike_times) / time_stamps[-1]:.2f} Hz")
+
+    # Classify
+    print("\nClassifying cell...")
+    analyzer = HeadDirectionAnalyzer(mvl_hd_threshold=0.4, strict_mode=False)
+    result = analyzer.classify_hd_cell(spike_times, head_directions, time_stamps)
+
+    print("\nResults:")
+    print(f"  Is HD cell: {result.is_hd}")
+    print(f"  MVL: {result.mvl_hd:.3f}")
+    print(f"  Preferred direction: {np.rad2deg(result.preferred_direction):.1f}°")
+    print(f"  True preferred direction: {np.rad2deg(preferred_dir):.1f}°")
+    print(f"  Rayleigh test p-value: {result.rayleigh_p:.6f}")
+
+    # Check tuning curve
+    bin_centers, firing_rates = result.tuning_curve
+    max_rate_idx = np.argmax(firing_rates)
+    print(f"  Peak firing rate: {firing_rates[max_rate_idx]:.2f} Hz")
+    print(f"  Peak at direction: {np.rad2deg(bin_centers[max_rate_idx]):.1f}°")
+
+    # Test with non-directional cell
+    print("\n\nSimulating non-directional cell...")
+    # Random spikes (Poisson process, no directional tuning)
+    mean_rate = 5  # Hz
+    n_spikes = int(mean_rate * time_stamps[-1])
+    spike_times_random = np.sort(np.random.uniform(0, time_stamps[-1], n_spikes))
+
+    result_random = analyzer.classify_hd_cell(spike_times_random, head_directions, time_stamps)
+
+    print("Results for non-directional cell:")
+    print(f"  Is HD cell: {result_random.is_hd}")
+    print(f"  MVL: {result_random.mvl_hd:.3f}")
+    print(f"  Rayleigh test p-value: {result_random.rayleigh_p:.3f}")
+
+    print("\nHeadDirectionAnalyzer test completed!")