canns 0.14.3__py3-none-any.whl → 0.15.1__py3-none-any.whl

canns/analyzer/data/asa/embedding.py

@@ -6,7 +6,6 @@ import numpy as np
 from scipy.ndimage import gaussian_filter1d
 
 from .config import DataLoadError, ProcessingError, SpikeEmbeddingConfig
-from .filters import _gaussian_filter1d
 
 
 def embed_spike_trains(spike_trains, config: SpikeEmbeddingConfig | None = None, **kwargs):
@@ -31,10 +30,9 @@ def embed_spike_trains(spike_trains, config: SpikeEmbeddingConfig | None = None,
     Returns
     -------
     tuple
-        ``(spikes_bin, xx, yy, tt)`` where:
-        - ``spikes_bin`` is a (T, N) binned spike matrix.
-        - ``xx``, ``yy``, ``tt`` are position/time arrays when ``speed_filter=True``,
-          otherwise ``None``.
+        ``(spikes_bin, xx, yy, tt)``. ``spikes_bin`` is a (T, N) binned spike matrix.
+        ``xx``, ``yy``, ``tt`` are position/time arrays when ``speed_filter=True``,
+        otherwise ``None``.
 
     Examples
     --------
@@ -263,8 +261,8 @@ def _load_pos(t, x, y, res=100000, dt=1000):
     xx = np.interp(tt, t, x)
     yy = np.interp(tt, t, y)
 
-    xxs = _gaussian_filter1d(xx - np.min(xx), sigma=100)
-    yys = _gaussian_filter1d(yy - np.min(yy), sigma=100)
+    xxs = gaussian_filter1d(xx - np.min(xx), sigma=100)
+    yys = gaussian_filter1d(yy - np.min(yy), sigma=100)
     dx = (xxs[1:] - xxs[:-1]) * 100
     dy = (yys[1:] - yys[:-1]) * 100
     speed = np.sqrt(dx**2 + dy**2) / 0.01

canns/analyzer/data/asa/fr.py

@@ -1,18 +1,11 @@
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass
 
 import numpy as np
 
 from ...visualization.core import PlotConfig, finalize_figure
-
-
-def _ensure_parent_dir(save_path: str | None) -> None:
-    if save_path:
-        parent = os.path.dirname(save_path)
-        if parent:
-            os.makedirs(parent, exist_ok=True)
+from .utils import _ensure_parent_dir
 
 
 def _slice_range(r: tuple[int, int] | None, length: int) -> slice:

canns/analyzer/data/asa/path.py

@@ -240,6 +240,76 @@ def interp_coords_to_full_1d(idx_map: np.ndarray, coords1: np.ndarray, T_full: i
     return np.mod(out, 2 * np.pi)[:, None]
 
 
+def _align_activity_to_coords(
+    coords: np.ndarray,
+    activity: np.ndarray,
+    times: np.ndarray | None = None,
+    *,
+    label: str = "activity",
+    auto_filter: bool = True,
+) -> np.ndarray:
+    """
+    Align activity to coords by optional time indices and validate lengths.
+
+    Parameters
+    ----------
+    coords : ndarray
+        Decoded coordinates array.
+    activity : ndarray
+        Activity matrix (firing rate or spikes).
+    times : ndarray, optional
+        Optional time indices to align activity to coords when coords are computed
+        on a subset of timepoints.
+    label : str
+        Label for error messages (default: "activity").
+    auto_filter : bool
+        If True and lengths mismatch, auto-filter activity with activity>0 to mimic
+        decode filtering.
+
+    Returns
+    -------
+    ndarray
+        Aligned activity array.
+
+    Raises
+    ------
+    ValueError
+        If activity length doesn't match coords length after alignment attempts.
+    """
+    coords = np.asarray(coords)
+    activity = np.asarray(activity)
+
+    if times is not None:
+        times = np.asarray(times)
+        try:
+            activity = activity[times]
+        except Exception as exc:
+            raise ValueError(
+                f"Failed to index {label} with `times`. Ensure `times` indexes the original time axis."
+            ) from exc
+
+    if activity.shape[0] != coords.shape[0]:
+        # Try to reproduce decode's zero-spike filtering if lengths mismatch.
+        if auto_filter and times is None and activity.ndim == 2:
+            mask = np.sum(activity > 0, axis=1) >= 1
+            if mask.sum() == coords.shape[0]:
+                activity = activity[mask]
+            else:
+                raise ValueError(
+                    f"{label} length must match coords length. Got {activity.shape[0]} vs {coords.shape[0]}. "
+                    "If coords are computed on a subset of timepoints (e.g., decode['times']), pass "
+                    "`times=decoding['times']` or slice the activity accordingly."
+                )
+        else:
+            raise ValueError(
+                f"{label} length must match coords length. Got {activity.shape[0]} vs {coords.shape[0]}. "
+                "If coords are computed on a subset of timepoints (e.g., decode['times']), pass "
+                "`times=decoding['times']` or slice the activity accordingly."
+            )
+
+    return activity
+
+
 def align_coords_to_position_2d(
     t_full: np.ndarray,
     x_full: np.ndarray,

canns/analyzer/data/asa/plotting.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 from typing import Any
 
 import matplotlib.pyplot as plt
@@ -24,31 +23,7 @@ from ...visualization.core import (
 from ...visualization.core.jupyter_utils import display_animation_in_jupyter, is_jupyter_environment
 from .config import CANN2DPlotConfig, ProcessingError, SpikeEmbeddingConfig
 from .embedding import embed_spike_trains
-
-
-def _ensure_plot_config(
-    config: PlotConfig | None,
-    factory,
-    *,
-    kwargs: dict[str, Any] | None = None,
-    **defaults: Any,
-) -> PlotConfig:
-    if config is None:
-        defaults.update({"kwargs": kwargs or {}})
-        return factory(**defaults)
-
-    if kwargs:
-        config_kwargs = config.kwargs or {}
-        config_kwargs.update(kwargs)
-        config.kwargs = config_kwargs
-    return config
-
-
-def _ensure_parent_dir(save_path: str | None) -> None:
-    if save_path:
-        parent = os.path.dirname(save_path)
-        if parent:
-            os.makedirs(parent, exist_ok=True)
+from .utils import _ensure_parent_dir, _ensure_plot_config
 
 
 def _render_torus_frame(frame_index: int, frame_data: dict[str, Any]) -> np.ndarray:
@@ -403,7 +378,7 @@ def plot_path_compare_1d(
     return fig, axes
 
 
-def plot_cohomap(
+def plot_cohomap_scatter(
     decoding_result: dict[str, Any],
     position_data: dict[str, Any],
     config: PlotConfig | None = None,
@@ -453,7 +428,7 @@ def plot_cohomap(
         >>> # Decode coordinates
         >>> decoding = decode_circular_coordinates(persistence_result, spike_data)
         >>> # Visualize with trajectory data
-        >>> fig = plot_cohomap(
+        >>> fig = plot_cohomap_scatter(
         ...     decoding,
         ...     position_data={'x': xx, 'y': yy},
         ...     save_path='cohomap.png',
@@ -518,7 +493,7 @@ def plot_cohomap(
     return fig
 
 
-def plot_cohomap_multi(
+def plot_cohomap_scatter_multi(
     decoding_result: dict,
     position_data: dict,
     config: PlotConfig | None = None,
@@ -560,7 +535,7 @@ def plot_cohomap_multi(
 
     Examples
     --------
-    >>> fig = plot_cohomap_multi(decoding, {"x": xx, "y": yy}, show=False)  # doctest: +SKIP
+    >>> fig = plot_cohomap_scatter_multi(decoding, {"x": xx, "y": yy}, show=False)  # doctest: +SKIP
     """
     config = _ensure_plot_config(
         config,

canns/analyzer/data/asa/utils.py

@@ -0,0 +1,160 @@
+"""Shared utility functions for ASA (Attractor State Analysis) modules."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+from scipy.ndimage import gaussian_filter
+
+from ...visualization.core import PlotConfig
+from .path import find_coords_matrix, find_times_box
+
+
+def _ensure_plot_config(
+    config: PlotConfig | None,
+    factory,
+    *args,
+    kwargs: dict | None = None,
+    **defaults,
+) -> PlotConfig:
+    """Ensure a PlotConfig exists, creating one from factory if needed.
+
+    Args:
+        config: Optional existing PlotConfig.
+        factory: Factory function to create PlotConfig if config is None.
+        *args: Positional arguments for factory.
+        kwargs: Optional dict to merge into config.kwargs.
+        **defaults: Keyword arguments for factory.
+
+    Returns:
+        PlotConfig instance.
+    """
+    if config is None:
+        if kwargs:
+            defaults.update({"kwargs": kwargs})
+        return factory(*args, **defaults)
+
+    # If config exists and kwargs provided, merge them
+    if kwargs:
+        config_kwargs = config.kwargs or {}
+        config_kwargs.update(kwargs)
+        config.kwargs = config_kwargs
+    return config
+
+
+def _ensure_parent_dir(save_path: str | None) -> None:
+    """Create parent directory for save_path if it doesn't exist.
+
+    Args:
+        save_path: Optional file path. If provided, creates parent directory.
+    """
+    if save_path:
+        parent = os.path.dirname(save_path)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+
+
+def _circmean(x: np.ndarray) -> float:
+    """Compute circular mean of angles.
+
+    Args:
+        x: Array of angles in radians.
+
+    Returns:
+        Circular mean in radians.
+    """
+    return float(np.arctan2(np.mean(np.sin(x)), np.mean(np.cos(x))))
+
+
+def _smooth_circular_map(
+    mtot: np.ndarray,
+    smooth_sigma: float,
+    *,
+    fill_nan: bool = False,
+    fill_sigma: float | None = None,
+    fill_min_weight: float = 1e-3,
+) -> np.ndarray:
+    """Smooth a circular phase map using Gaussian filtering in sin/cos space.
+
+    Args:
+        mtot: Phase map array (angles in radians).
+        smooth_sigma: Gaussian smoothing sigma.
+        fill_nan: Whether to fill NaN values using weighted interpolation.
+        fill_sigma: Sigma for NaN filling (defaults to smooth_sigma).
+        fill_min_weight: Minimum weight threshold for valid interpolation.
+
+    Returns:
+        Smoothed phase map.
+    """
+    mtot = np.asarray(mtot, dtype=float)
+    nans = np.isnan(mtot)
+    mask = (~nans).astype(float)
+    sintot = np.sin(mtot)
+    costot = np.cos(mtot)
+    sintot[nans] = 0.0
+    costot[nans] = 0.0
+
+    if fill_nan:
+        if fill_sigma is None:
+            fill_sigma = smooth_sigma if smooth_sigma and smooth_sigma > 0 else 1.0
+        weight = gaussian_filter(mask, fill_sigma)
+        sintot = gaussian_filter(sintot * mask, fill_sigma)
+        costot = gaussian_filter(costot * mask, fill_sigma)
+        min_weight = max(float(fill_min_weight), 0.0)
+        valid = weight > min_weight
+        sintot = np.divide(sintot, weight, out=np.zeros_like(sintot), where=valid)
+        costot = np.divide(costot, weight, out=np.zeros_like(costot), where=valid)
+        mtot = np.arctan2(sintot, costot)
+        if fill_min_weight > 0:
+            mtot[~valid] = np.nan
+        return mtot
+
+    if smooth_sigma and smooth_sigma > 0:
+        sintot = gaussian_filter(sintot, smooth_sigma)
+        costot = gaussian_filter(costot, smooth_sigma)
+    mtot = np.arctan2(sintot, costot)
+    mtot[nans] = np.nan
+    return mtot
+
+
+def _extract_coords_and_times(
+    decoding_result: dict[str, Any],
+    coords_key: str | None = None,
+) -> tuple[np.ndarray, np.ndarray | None]:
+    """Extract coordinates and time indices from decoding result.
+
+    Args:
+        decoding_result: Dictionary containing decoding results.
+        coords_key: Optional key for coordinates (defaults to 'coordsbox' or auto-detect).
+
+    Returns:
+        Tuple of (coords, times_box) where coords is (T, 2) and times_box is optional.
+    """
+    if coords_key is not None:
+        if coords_key not in decoding_result:
+            raise KeyError(f"coords_key '{coords_key}' not found in decoding_result.")
+        coords = np.asarray(decoding_result[coords_key])
+    elif "coordsbox" in decoding_result:
+        coords = np.asarray(decoding_result["coordsbox"])
+    else:
+        coords, _ = find_coords_matrix(decoding_result)
+
+    times_box, _ = find_times_box(decoding_result)
+    return coords, times_box
+
+
+def _phase_map_valid_fraction(phase_map: np.ndarray) -> float:
+    """Calculate fraction of valid (non-NaN) values in phase map.
+
+    Args:
+        phase_map: Phase map array.
+
+    Returns:
+        Fraction of valid values (0.0 to 1.0).
+    """
+    valid = np.isfinite(phase_map)
+    if valid.size == 0:
+        return 0.0
+    return float(np.mean(valid))

canns/analyzer/data/cell_classification/__init__.py

@@ -11,6 +11,9 @@ Vollan, Gardner, Moser & Moser (Nature, 2025)
 __version__ = "0.1.0"
 
 from .core import (  # noqa: F401
+    BTNAnalyzer,
+    BTNConfig,
+    BTNResult,
     GridnessAnalyzer,
     GridnessResult,
     HDCellResult,
@@ -46,6 +49,8 @@ from .utils import (  # noqa: F401
 )
 from .visualization import (  # noqa: F401
     plot_autocorrelogram,
+    plot_btn_autocorr_summary,
+    plot_btn_distance_matrix,
     plot_grid_score_histogram,
     plot_gridness_analysis,
     plot_hd_analysis,
@@ -57,6 +62,9 @@ from .visualization import (  # noqa: F401
 __all__ = [
     "GridnessAnalyzer",
     "GridnessResult",
+    "BTNAnalyzer",
+    "BTNConfig",
+    "BTNResult",
     "HeadDirectionAnalyzer",
     "HDCellResult",
     "compute_2d_autocorrelation",
@@ -94,4 +102,6 @@ __all__ = [
     "plot_polar_tuning",
     "plot_temporal_autocorr",
     "plot_hd_analysis",
+    "plot_btn_autocorr_summary",
+    "plot_btn_distance_matrix",
 ]

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

@@ -1,5 +1,6 @@
 """Core analysis modules."""
 
+from .btn import BTNAnalyzer, BTNConfig, BTNResult
 from .grid_cells import GridnessAnalyzer, GridnessResult, compute_2d_autocorrelation
 from .grid_modules_leiden import identify_grid_modules_and_stats
 from .head_direction import HDCellResult, HeadDirectionAnalyzer
@@ -15,6 +16,9 @@ __all__ = [
     "GridnessAnalyzer",
     "compute_2d_autocorrelation",
     "GridnessResult",
+    "BTNAnalyzer",
+    "BTNConfig",
+    "BTNResult",
     "HeadDirectionAnalyzer",
     "HDCellResult",
     "compute_rate_map",

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

@@ -0,0 +1,272 @@
+"""
+BTN (Bursty/Theta/Non-bursty) classification.
+
+Workflow:
+1) Compute ISI autocorrelograms from spike times.
+2) Normalize and smooth each autocorr curve.
+3) Compute cosine distance between curves.
+4) Cluster with Tomato using a manual kNN graph and density weights.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+from scipy.ndimage import gaussian_filter1d
+from scipy.spatial.distance import pdist, squareform
+
+try:
+    from gudhi.clustering.tomato import Tomato
+except Exception as exc:  # pragma: no cover - optional dependency
+    Tomato = None
+    _TOMATO_IMPORT_ERROR = exc
+else:
+    _TOMATO_IMPORT_ERROR = None
+
+
+@dataclass
+class BTNConfig:
+    """Configuration for BTN clustering."""
+
+    maxt: float = 0.2
+    res: float = 1e-3
+    smooth_sigma: float = 4.0
+    nbs: int = 80
+    n_clusters: int = 4
+    metric: str = "cosine"
+    b_one: bool = True
+    b_log: bool = False
+
+
+@dataclass
+class BTNResult:
+    """Result container for BTN clustering."""
+
+    labels: np.ndarray
+    n_clusters: int
+    cluster_sizes: np.ndarray
+    btn_labels: np.ndarray | None
+    mapping: dict[int, str] | None
+    intermediates: dict[str, np.ndarray] | None
+
+
+class BTNAnalyzer:
+    """Analyzer that clusters neurons into BTN groups using Tomato."""
+
+    def __init__(self, config: BTNConfig | None = None):
+        self.config = config or BTNConfig()
+
+    def classify_btn(
+        self,
+        spike_data: dict[str, Any],
+        *,
+        mapping: dict[int, str] | None = None,
+        return_intermediates: bool = False,
+        plot_diagram: bool = False,
+    ) -> BTNResult:
+        """Cluster neurons into BTN classes using ISI autocorr + Tomato.
+
+        Parameters
+        ----------
+        spike_data : dict
+            ASA-style dict with keys ``spike`` and ``t`` (and optionally x/y).
+        mapping : dict, optional
+            Optional mapping from cluster id to BTN label string.
+        return_intermediates : bool
+            If True, include intermediate arrays in the result.
+        plot_diagram : bool
+            If True, call Tomato.plot_diagram() for visual inspection.
+        """
+        _require_tomato()
+        spikes = _extract_spike_times(spike_data)
+        acorr, bin_times = _isi_autocorr(
+            spikes,
+            maxt=self.config.maxt,
+            res=self.config.res,
+            b_one=self.config.b_one,
+            b_log=self.config.b_log,
+        )
+        acorr_norm = _normalize_autocorr(acorr)
+        acorr_smooth = gaussian_filter1d(acorr_norm, sigma=self.config.smooth_sigma, axis=1)
+
+        dist = squareform(pdist(acorr_smooth, metric=self.config.metric))
+        num_nodes = dist.shape[0]
+        order = np.argsort(dist, axis=1)
+        if num_nodes > 1:
+            nbs_max = num_nodes - 1
+            nbs = max(1, min(int(self.config.nbs), nbs_max))
+            knn_indices = order[:, 1 : nbs + 1]
+        else:
+            nbs = 1
+            knn_indices = order[:, :1]
+        knn_dists = dist[np.arange(dist.shape[0])[:, None], knn_indices]
+        weights = np.sum(np.exp(-knn_dists), axis=1)
+
+        t = Tomato(graph_type="manual", density_type="manual", metric="precomputed")
+        t.fit(knn_indices, weights=weights)
+        if plot_diagram:
+            t.plot_diagram()
+        if self.config.n_clusters is not None:
+            t.n_clusters_ = int(self.config.n_clusters)
+        labels = np.asarray(t.labels_, dtype=int)
+        if labels.size:
+            if np.any(labels < 0):
+                valid = labels[labels >= 0]
+                cluster_sizes = np.bincount(valid) if valid.size else np.array([])
+            else:
+                cluster_sizes = np.bincount(labels)
+        else:
+            cluster_sizes = np.array([])
+
+        btn_labels = None
+        if mapping is not None:
+            btn_labels = np.array([mapping.get(int(c), "unknown") for c in labels], dtype=object)
+
+        intermediates = None
+        if return_intermediates:
+            intermediates = {
+                "acorr": acorr,
+                "acorr_norm": acorr_norm,
+                "acorr_smooth": acorr_smooth,
+                "distance_matrix": dist,
+                "knn_indices": knn_indices,
+                "knn_dists": knn_dists,
+                "density_weights": weights,
+                "bin_times": bin_times,
+            }
+
+        if self.config.n_clusters is not None:
+            n_clusters = int(self.config.n_clusters)
+        else:
+            if labels.size:
+                valid_labels = labels[labels >= 0]
+                n_clusters = int(np.unique(valid_labels).size)
+            else:
+                n_clusters = 0
+
+        return BTNResult(
+            labels=labels,
+            n_clusters=n_clusters,
+            cluster_sizes=cluster_sizes,
+            btn_labels=btn_labels,
+            mapping=mapping,
+            intermediates=intermediates,
+        )
+
+
+def _require_tomato() -> None:
+    if Tomato is None:
+        raise ImportError(
+            "Tomato clustering requires gudhi. Install with: pip install gudhi"
+        ) from _TOMATO_IMPORT_ERROR
+
+
+def _extract_spike_times(spike_data: dict[str, Any]) -> dict[int, np.ndarray]:
+    if not isinstance(spike_data, dict) or "spike" not in spike_data or "t" not in spike_data:
+        raise ValueError("spike_data must be a dict with keys 'spike' and 't'")
+
+    spike_raw = spike_data["spike"]
+    t = np.asarray(spike_data["t"])
+    if (
+        isinstance(spike_raw, np.ndarray)
+        and spike_raw.ndim == 2
+        and spike_raw.shape[0] == t.shape[0]
+    ):
+        raise ValueError(
+            "BTN expects spike times per neuron, but got a binned spike matrix (T x N)."
+        )
+    if (
+        hasattr(spike_raw, "item")
+        and callable(spike_raw.item)
+        and np.asarray(spike_raw).shape == ()
+    ):
+        spikes_all = spike_raw[()]
+    elif isinstance(spike_raw, dict):
+        spikes_all = spike_raw
+    elif isinstance(spike_raw, (list, np.ndarray)):
+        spikes_all = spike_raw
+    else:
+        spikes_all = spike_raw
+
+    min_time0 = np.min(t)
+    max_time0 = np.max(t)
+
+    spikes: dict[int, np.ndarray] = {}
+    if isinstance(spikes_all, dict):
+        for i, key in enumerate(spikes_all.keys()):
+            s = np.asarray(spikes_all[key])
+            spikes[i] = s[(s >= min_time0) & (s < max_time0)]
+    else:
+        cell_inds = np.arange(len(spikes_all))
+        for i, m in enumerate(cell_inds):
+            s = np.asarray(spikes_all[m]) if len(spikes_all[m]) > 0 else np.array([])
+            if s.size > 0:
+                spikes[i] = s[(s >= min_time0) & (s < max_time0)]
+            else:
+                spikes[i] = np.array([])
+
+    return spikes
+
+
+def _isi_autocorr(
+    spk: dict[int, np.ndarray],
+    *,
+    maxt: float = 0.2,
+    res: float = 1e-3,
+    b_one: bool = True,
+    b_log: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute ISI autocorrelogram for each neuron.
+
+    Returns
+    -------
+    acorr : ndarray
+        Shape (N, num_bins).
+    bin_times : ndarray
+        Bin edges used for histogramming.
+    """
+    if b_log:
+        num_bins = 100
+        bin_times = np.ones(num_bins + 1) * 10
+        bin_times = np.power(bin_times, np.linspace(np.log10(0.005), np.log10(maxt), num_bins + 1))
+        bin_times = np.unique(np.concatenate((-bin_times, bin_times)))
+        num_bins = len(bin_times)
+    elif b_one:
+        num_bins = int(maxt / res) + 1
+        bin_times = np.linspace(0, maxt, num_bins)
+    else:
+        num_bins = int(2 * maxt / res) + 1
+        bin_times = np.linspace(-maxt, maxt, num_bins)
+
+    num_neurons = len(spk)
+    acorr = np.zeros((num_neurons, len(bin_times) - 1), dtype=int)
+
+    maxt = maxt - 1e-5
+    mint = -maxt
+    if b_one:
+        mint = -1e-5
+
+    for i, n in enumerate(spk):
+        spike_times = np.asarray(spk[n])
+        for ss in spike_times:
+            stemp = spike_times[(spike_times < ss + maxt) & (spike_times > ss + mint)]
+            dd = stemp - ss
+            if b_one:
+                dd = dd[dd >= 0]
+            bins = np.digitize(dd, bin_times) - 1
+            bins = bins[(bins >= 0) & (bins < num_bins - 1)]
+            if bins.size:
+                acorr[i, :] += np.bincount(bins, minlength=num_bins)[:-1]
+
+    return acorr, bin_times
+
+
+def _normalize_autocorr(acorr: np.ndarray) -> np.ndarray:
+    acorr = acorr.astype(float, copy=True)
+    denom = acorr[:, 0].copy()
+    valid = denom > 0
+    acorr[valid, :] = acorr[valid, :] / denom[valid, None]
+    acorr[:, 0] = 0.0
+    return acorr

canns/analyzer/data/cell_classification/visualization/__init__.py

@@ -1,5 +1,6 @@
 """Visualization modules."""
 
+from .btn_plots import plot_btn_autocorr_summary, plot_btn_distance_matrix
 from .grid_plots import (
     plot_autocorrelogram,
     plot_grid_score_histogram,
@@ -16,4 +17,6 @@ __all__ = [
     "plot_polar_tuning",
     "plot_temporal_autocorr",
     "plot_hd_analysis",
+    "plot_btn_distance_matrix",
+    "plot_btn_autocorr_summary",
 ]