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

canns/analyzer/data/asa/embedding.py

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+from typing import Any
+
+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):
+    """
+    Load and preprocess spike train data from npz file.
+
+    This function converts raw spike times into a time-binned spike matrix,
+    optionally applying Gaussian smoothing and filtering based on animal movement speed.
+
+    Parameters
+    ----------
+    spike_trains : dict
+        Dictionary containing ``'spike'`` and ``'t'``, and optionally ``'x'``/``'y'``.
+        ``'spike'`` can be a dict of neuron->spike_times, a list/array of arrays, or
+        a numpy object array from ``np.load``.
+    config : SpikeEmbeddingConfig, optional
+        Configuration object controlling binning, smoothing, and speed filtering.
+    **kwargs : Any
+        Legacy keyword parameters (``res``, ``dt``, ``sigma``, ``smooth0``, ``speed0``,
+        ``min_speed``). Prefer ``config`` in new code.
+
+    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``.
+
+    Examples
+    --------
+    >>> from canns.analyzer.data import SpikeEmbeddingConfig, embed_spike_trains
+    >>> cfg = SpikeEmbeddingConfig(smooth=False, speed_filter=False)
+    >>> spikes, xx, yy, tt = embed_spike_trains(mock_data, config=cfg)  # doctest: +SKIP
+    >>> spikes.ndim
+    2
+    """
+    # Handle backward compatibility and configuration
+    if config is None:
+        config = SpikeEmbeddingConfig(
+            res=kwargs.get("res", 100000),
+            dt=kwargs.get("dt", 1000),
+            sigma=kwargs.get("sigma", 5000),
+            smooth=kwargs.get("smooth0", True),
+            speed_filter=kwargs.get("speed0", True),
+            min_speed=kwargs.get("min_speed", 2.5),
+        )
+
+    try:
+        # Step 1: Extract and filter spike data
+        spikes_filtered = _extract_spike_data(spike_trains, config)
+
+        # Step 2: Create time bins
+        time_bins = _create_time_bins(spike_trains["t"], config)
+
+        # Step 3: Bin spike data
+        spikes_bin = _bin_spike_data(spikes_filtered, time_bins, config)
+
+        # Step 4: Apply temporal smoothing if requested
+        if config.smooth:
+            spikes_bin = _apply_temporal_smoothing(spikes_bin, config)
+
+        # Step 5: Apply speed filtering if requested
+        if config.speed_filter:
+            return _apply_speed_filtering(spikes_bin, spike_trains, config)
+
+        return spikes_bin, None, None, None
+
+    except Exception as e:
+        raise ProcessingError(f"Failed to embed spike trains: {e}") from e
+
+
+def _extract_spike_data(
+    spike_trains: dict[str, Any], config: SpikeEmbeddingConfig
+) -> dict[int, np.ndarray]:
+    """Extract and filter spike data within time window."""
+    try:
+        # Handle different spike data formats
+        spike_data = spike_trains["spike"]
+        if hasattr(spike_data, "item") and callable(spike_data.item):
+            # numpy array with .item() method (from npz file)
+            spikes_all = spike_data[()]
+        elif isinstance(spike_data, dict):
+            # Already a dictionary
+            spikes_all = spike_data
+        elif isinstance(spike_data, (list, np.ndarray)):
+            # List or array format
+            spikes_all = spike_data
+        else:
+            # Try direct access
+            spikes_all = spike_data
+
+        t = spike_trains["t"]
+
+        min_time0 = np.min(t)
+        max_time0 = np.max(t)
+
+        # Extract spike intervals for each cell
+        if isinstance(spikes_all, dict):
+            # Dictionary format
+            spikes = {}
+            for i, key in enumerate(spikes_all.keys()):
+                s = np.array(spikes_all[key])
+                spikes[i] = s[(s >= min_time0) & (s < max_time0)]
+        else:
+            # List/array format
+            cell_inds = np.arange(len(spikes_all))
+            spikes = {}
+
+            for i, m in enumerate(cell_inds):
+                s = np.array(spikes_all[m]) if len(spikes_all[m]) > 0 else np.array([])
+                # Filter spikes within time window
+                if len(s) > 0:
+                    spikes[i] = s[(s >= min_time0) & (s < max_time0)]
+                else:
+                    spikes[i] = np.array([])
+
+        return spikes
+
+    except KeyError as e:
+        raise DataLoadError(f"Missing required data key: {e}") from e
+    except Exception as e:
+        raise ProcessingError(f"Error extracting spike data: {e}") from e
+
+
+def _create_time_bins(t: np.ndarray, config: SpikeEmbeddingConfig) -> np.ndarray:
+    """Create time bins for spike discretization."""
+    min_time0 = np.min(t)
+    max_time0 = np.max(t)
+
+    min_time = min_time0 * config.res
+    max_time = max_time0 * config.res
+
+    return np.arange(np.floor(min_time), np.ceil(max_time) + 1, config.dt)
+
+
+def _bin_spike_data(
+    spikes: dict[int, np.ndarray], time_bins: np.ndarray, config: SpikeEmbeddingConfig
+) -> np.ndarray:
+    """Convert spike times to binned spike matrix."""
+    min_time = time_bins[0]
+    max_time = time_bins[-1]
+
+    spikes_bin = np.zeros((len(time_bins), len(spikes)), dtype=int)
+
+    for n in spikes:
+        spike_times = np.array(spikes[n] * config.res - min_time, dtype=int)
+        # Filter valid spike times
+        spike_times = spike_times[(spike_times < (max_time - min_time)) & (spike_times > 0)]
+        spike_times = np.array(spike_times / config.dt, int)
+
+        # Bin spikes
+        for j in spike_times:
+            if j < len(time_bins):
+                spikes_bin[j, n] += 1
+
+    return spikes_bin
+
+
+def _apply_temporal_smoothing(spikes_bin: np.ndarray, config: SpikeEmbeddingConfig) -> np.ndarray:
+    """Apply Gaussian temporal smoothing to spike matrix."""
+    # Calculate smoothing parameters (legacy implementation used custom kernel)
+    # Current implementation uses scipy's gaussian_filter1d for better performance
+
+    # Apply smoothing (simplified version - could be further optimized)
+    smoothed = np.zeros((spikes_bin.shape[0], spikes_bin.shape[1]))
+
+    # Use scipy's gaussian_filter1d for better performance
+
+    sigma_bins = config.sigma / config.dt
+
+    for n in range(spikes_bin.shape[1]):
+        smoothed[:, n] = gaussian_filter1d(
+            spikes_bin[:, n].astype(float), sigma=sigma_bins, mode="constant"
+        )
+
+    # Normalize
+    normalization_factor = 1 / np.sqrt(2 * np.pi * (config.sigma / config.res) ** 2)
+    return smoothed * normalization_factor
+
+
+def _apply_speed_filtering(
+    spikes_bin: np.ndarray, spike_trains: dict[str, Any], config: SpikeEmbeddingConfig
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Apply speed-based filtering to spike data."""
+    try:
+        xx, yy, tt_pos, speed = _load_pos(
+            spike_trains["t"], spike_trains["x"], spike_trains["y"], res=config.res, dt=config.dt
+        )
+
+        valid = speed > config.min_speed
+
+        return (spikes_bin[valid, :], xx[valid], yy[valid], tt_pos[valid])
+
+    except KeyError as e:
+        raise DataLoadError(f"Missing position data for speed filtering: {e}") from e
+    except Exception as e:
+        raise ProcessingError(f"Error in speed filtering: {e}") from e
+
+
+def _load_pos(t, x, y, res=100000, dt=1000):
+    """
+    Compute animal position and speed from spike data file.
+
+    Interpolates animal positions to match spike time bins and computes smoothed velocity vectors and speed.
+
+    Parameters:
+        t (ndarray): Time points of the spikes (in seconds).
+        x (ndarray): X coordinates of the animal's position.
+        y (ndarray): Y coordinates of the animal's position.
+        res (int): Time scaling factor to align with spike resolution.
+        dt (int): Temporal bin size in microseconds.
+
+    Returns:
+        xx (ndarray): Interpolated x positions.
+        yy (ndarray): Interpolated y positions.
+        tt (ndarray): Corresponding time points (in seconds).
+        speed (ndarray): Speed at each time point (in cm/s).
+    """
+
+    min_time0 = np.min(t)
+    max_time0 = np.max(t)
+
+    times = np.where((t >= min_time0) & (t < max_time0))
+    x = x[times]
+    y = y[times]
+    t = t[times]
+
+    min_time = min_time0 * res
+    max_time = max_time0 * res
+
+    tt = np.arange(np.floor(min_time), np.ceil(max_time) + 1, dt) / res
+
+    idt = np.concatenate(([0], np.digitize(t[1:-1], tt[:]) - 1, [len(tt) + 1]))
+    idtt = np.digitize(np.arange(len(tt)), idt) - 1
+
+    idx = np.concatenate((np.unique(idtt), [np.max(idtt) + 1]))
+    divisor = np.bincount(idtt)
+    steps = 1.0 / divisor[divisor > 0]
+    N = np.max(divisor)
+    ranges = np.multiply(np.arange(N)[np.newaxis, :], steps[:, np.newaxis])
+    ranges[ranges >= 1] = np.nan
+
+    rangesx = x[idx[:-1], np.newaxis] + np.multiply(
+        ranges, (x[idx[1:]] - x[idx[:-1]])[:, np.newaxis]
+    )
+    xx = rangesx[~np.isnan(ranges)]
+
+    rangesy = y[idx[:-1], np.newaxis] + np.multiply(
+        ranges, (y[idx[1:]] - y[idx[:-1]])[:, np.newaxis]
+    )
+    yy = rangesy[~np.isnan(ranges)]
+
+    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
+    speed = np.concatenate(([speed[0]], speed))
+    return xx, yy, tt, speed

canns/analyzer/data/asa/filters.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+import numbers
+
+import numpy as np
+from numpy.exceptions import AxisError
+from scipy.ndimage import _nd_image, _ni_support
+from scipy.ndimage._filters import _invalid_origin
+
+
+def _gaussian_filter1d(
+    input,
+    sigma,
+    axis=-1,
+    order=0,
+    output=None,
+    mode="reflect",
+    cval=0.0,
+    truncate=4.0,
+    *,
+    radius=None,
+):
+    """1-D Gaussian filter.
+
+    Parameters
+    ----------
+    %(input)s
+    sigma : scalar
+        standard deviation for Gaussian kernel
+    %(axis)s
+    order : int, optional
+        An order of 0 corresponds to convolution with a Gaussian
+        kernel. A positive order corresponds to convolution with
+        that derivative of a Gaussian.
+    %(output)s
+    %(mode_reflect)s
+    %(cval)s
+    truncate : float, optional
+        Truncate the filter at this many standard deviations.
+        Default is 4.0.
+    radius : None or int, optional
+        Radius of the Gaussian kernel. If specified, the size of
+        the kernel will be ``2*radius + 1``, and `truncate` is ignored.
+        Default is None.
+
+    Returns
+    -------
+    gaussian_filter1d : ndarray
+
+    Notes
+    -----
+    The Gaussian kernel will have size ``2*radius + 1`` along each axis. If
+    `radius` is None, a default ``radius = round(truncate * sigma)`` will be
+    used.
+
+    Examples
+    --------
+    >>> from scipy.ndimage import gaussian_filter1d
+    >>> import numpy as np
+    >>> gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 1)
+    array([ 1.42704095,  2.06782203,  3.        ,  3.93217797,  4.57295905])
+    >>> _gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 4)
+    array([ 2.91948343,  2.95023502,  3.        ,  3.04976498,  3.08051657])
+    >>> import matplotlib.pyplot as plt
+    >>> rng = np.random.default_rng()
+    >>> x = rng.standard_normal(101).cumsum()
+    >>> y3 = _gaussian_filter1d(x, 3)
+    >>> y6 = _gaussian_filter1d(x, 6)
+    >>> plt.plot(x, 'k', label='original data')
+    >>> plt.plot(y3, '--', label='filtered, sigma=3')
+    >>> plt.plot(y6, ':', label='filtered, sigma=6')
+    >>> plt.legend()
+    >>> plt.grid()
+    >>> plt.show()
+
+    """
+    sd = float(sigma)
+    # make the radius of the filter equal to truncate standard deviations
+    lw = int(truncate * sd + 0.5)
+    if radius is not None:
+        lw = radius
+    if not isinstance(lw, numbers.Integral) or lw < 0:
+        raise ValueError("Radius must be a nonnegative integer.")
+    # Since we are calling correlate, not convolve, revert the kernel
+    weights = _gaussian_kernel1d(sigma, order, lw)[::-1]
+    return _correlate1d(input, weights, axis, output, mode, cval, 0)
+
+
+def _gaussian_kernel1d(sigma, order, radius):
+    """
+    Computes a 1-D Gaussian convolution kernel.
+    """
+    if order < 0:
+        raise ValueError("order must be non-negative")
+    exponent_range = np.arange(order + 1)
+    sigma2 = sigma * sigma
+    x = np.arange(-radius, radius + 1)
+    phi_x = np.exp(-0.5 / sigma2 * x**2)
+    phi_x = phi_x / phi_x.sum()
+
+    if order == 0:
+        return phi_x
+    else:
+        # f(x) = q(x) * phi(x) = q(x) * exp(p(x))
+        # f'(x) = (q'(x) + q(x) * p'(x)) * phi(x)
+        # p'(x) = -1 / sigma ** 2
+        # Implement q'(x) + q(x) * p'(x) as a matrix operator and apply to the
+        # coefficients of q(x)
+        q = np.zeros(order + 1)
+        q[0] = 1
+        D = np.diag(exponent_range[1:], 1)  # D @ q(x) = q'(x)
+        P = np.diag(np.ones(order) / -sigma2, -1)  # P @ q(x) = q(x) * p'(x)
+        Q_deriv = D + P
+        for _ in range(order):
+            q = Q_deriv.dot(q)
+        q = (x[:, None] ** exponent_range).dot(q)
+        return q * phi_x
+
+
+def _correlate1d(input, weights, axis=-1, output=None, mode="reflect", cval=0.0, origin=0):
+    """Calculate a 1-D correlation along the given axis.
+
+    The lines of the array along the given axis are correlated with the
+    given weights.
+
+    Parameters
+    ----------
+    %(input)s
+    weights : array
+        1-D sequence of numbers.
+    %(axis)s
+    %(output)s
+    %(mode_reflect)s
+    %(cval)s
+    %(origin)s
+
+    Returns
+    -------
+    result : ndarray
+        Correlation result. Has the same shape as `input`.
+
+    Examples
+    --------
+    >>> from scipy.ndimage import correlate1d
+    >>> correlate1d([2, 8, 0, 4, 1, 9, 9, 0], weights=[1, 3])
+    array([ 8, 26,  8, 12,  7, 28, 36,  9])
+    """
+    input = np.asarray(input)
+    weights = np.asarray(weights)
+    complex_input = input.dtype.kind == "c"
+    complex_weights = weights.dtype.kind == "c"
+    if complex_input or complex_weights:
+        if complex_weights:
+            weights = weights.conj()
+            weights = weights.astype(np.complex128, copy=False)
+        kwargs = dict(axis=axis, mode=mode, origin=origin)
+        output = _ni_support._get_output(output, input, complex_output=True)
+        return _complex_via_real_components(_correlate1d, input, weights, output, cval, **kwargs)
+
+    output = _ni_support._get_output(output, input)
+    weights = np.asarray(weights, dtype=np.float64)
+    if weights.ndim != 1 or weights.shape[0] < 1:
+        raise RuntimeError("no filter weights given")
+    if not weights.flags.contiguous:
+        weights = weights.copy()
+    axis = _normalize_axis_index(axis, input.ndim)
+    if _invalid_origin(origin, len(weights)):
+        raise ValueError(
+            "Invalid origin; origin must satisfy "
+            "-(len(weights) // 2) <= origin <= "
+            "(len(weights)-1) // 2"
+        )
+    mode = _ni_support._extend_mode_to_code(mode)
+    _nd_image.correlate1d(input, weights, axis, output, mode, cval, origin)
+    return output
+
+
+def _complex_via_real_components(func, input, weights, output, cval, **kwargs):
+    """Complex convolution via a linear combination of real convolutions."""
+    complex_input = input.dtype.kind == "c"
+    complex_weights = weights.dtype.kind == "c"
+    if complex_input and complex_weights:
+        # real component of the output
+        func(input.real, weights.real, output=output.real, cval=np.real(cval), **kwargs)
+        output.real -= func(input.imag, weights.imag, output=None, cval=np.imag(cval), **kwargs)
+        # imaginary component of the output
+        func(input.real, weights.imag, output=output.imag, cval=np.real(cval), **kwargs)
+        output.imag += func(input.imag, weights.real, output=None, cval=np.imag(cval), **kwargs)
+    elif complex_input:
+        func(input.real, weights, output=output.real, cval=np.real(cval), **kwargs)
+        func(input.imag, weights, output=output.imag, cval=np.imag(cval), **kwargs)
+    else:
+        if np.iscomplexobj(cval):
+            raise ValueError("Cannot provide a complex-valued cval when the input is real.")
+        func(input, weights.real, output=output.real, cval=cval, **kwargs)
+        func(input, weights.imag, output=output.imag, cval=cval, **kwargs)
+    return output
+
+
+def _normalize_axis_index(axis, ndim):
+    # Check if `axis` is in the correct range and normalize it
+    if axis < -ndim or axis >= ndim:
+        msg = f"axis {axis} is out of bounds for array of dimension {ndim}"
+        raise AxisError(msg)
+
+    if axis < 0:
+        axis = axis + ndim
+    return axis