v0.5.23 PR

v0.5.23

Author: cwindolf

pyproject.toml

@@ -75,7 +75,7 @@ line-length = 88
 
 [tool.ruff.lint]
 extend-select = ["E", "F", "I", "N", "B", "SIM", "NPY", "RUF", "PTH", "PD", "FURB"]
-ignore = ["F722", "N812", "SIM102", "SIM103", "SIM105", "SIM108", "E501", "N806", "N803", "N815", "RUF005", "SIM300", "C408", "TRY401"]
+ignore = ["F722", "N812", "SIM102", "SIM103", "SIM105", "SIM108", "SIM114", "E501", "N806", "N803", "N815", "RUF005", "SIM300", "C408", "TRY401"]
 
 [tool.ruff.lint.per-file-ignores]
 "**/__init__.py" = ["F401", "F403"]
@@ -96,5 +96,4 @@ replace-imports-with-any = ["hdbscan"]
 
 [tool.ty.rules]
 "unsupported-base" = "ignore"
-"invalid-key" = "ignore"
 "possibly-unresolved-reference" = "error"

src/dartsort/__init__.py

@@ -76,7 +76,12 @@
     get_logger,
     set_log_level,
 )
-from .util.motion import MotionInfo, get_motion_info, try_load_motion_info
+from .util.motion import (
+    MotionInfo,
+    detect_for_motion,
+    get_motion_info,
+    try_load_motion_info,
+)
 from .util.noise_util import EmbeddedNoise
 from .util.preprocess_util import preprocess
 from .util.py_util import databag, ensure_path

src/dartsort/clustering/agglomerate.py

@@ -1156,7 +1156,7 @@ def _calc_coentropy(
         u = cands[i]
         q = resps[i]
         log_q = np.log(q)
-        np.nan_to_num(log_q, copy=False, neginf=0.0)
+        np.nan_to_num(log_q, copy=False, neginf=0.0, posinf=np.inf)
         dh = q * log_q
 
         ui0 = u[0]

src/dartsort/clustering/cluster_util.py

@@ -1,19 +1,23 @@
 from typing import cast
 
 import h5py
+import numba
 import numpy as np
 from scipy.cluster.hierarchy import fcluster, linkage
 from scipy.spatial import KDTree
 
-from ..util import data_util, waveform_util
+from ..util import waveform_util
 from ..util.data_util import (
     DARTsortSorting,
     apply_label_remapping_in_place,
+    count_not_sorted,
     mean_by_label_1d,
     pos_int_unique_and_counts,
+    yield_masked_chunks,
 )
 from ..util.logging_util import get_logger
 from ..util.motion import MotionInfo
+from ..util.py_util import databag
 
 logger = get_logger(__name__)
 
@@ -509,7 +513,7 @@ def get_main_channel_pcs(
     with h5py.File(sorting.parent_h5_path, "r", locking=False) as h5:
         feats_dset = h5[dataset_name]
         channel_index = cast(h5py.Dataset, h5["channel_index"])[:]
-        for ixs, feats in data_util.yield_masked_chunks(
+        for ixs, feats in yield_masked_chunks(
             mask, feats_dset, show_progress=show_progress, desc_prefix="Main channel"
         ):
             feats = feats[:, :rank]
@@ -576,3 +580,102 @@ def decrumb(
     if flatten:
         sorting = sorting.flatten(in_place=in_place)
     return sorting
+
+
+@databag
+class ViolationCounts:
+    unit_ids: np.ndarray
+    spike_counts: np.ndarray
+    """Same shape as unit_ids (flat)"""
+    viol_counts: np.ndarray
+    """Indexed by pair of ids (not flat)"""
+
+
+def violation_matrix(
+    st: DARTsortSorting, *, censor_ms: float = 0.25, viol_ms: float = 1.0
+) -> ViolationCounts:
+    """Count ACG and CCG violations within viol_ms
+
+    Times within censor_ms of each other are ignored in the violation
+    count. The censorship is right-exclusive, so that if censor_ms is 0,
+    exact duplicates are counted; if censor_ms corresponds to 10 samples,
+    9-sample viols are excluded and 10-sample viols are counted.
+    """
+    assert st.labels is not None
+    censor_samples = int(censor_ms * (st.sampling_frequency / 1000.0))
+    viol_samples = int(viol_ms * (st.sampling_frequency / 1000.0))
+
+    unit_ids, spike_counts, _ = pos_int_unique_and_counts(st.labels)
+    nu = (unit_ids.max() + 1).item() if unit_ids.size else 0
+    if not nu or (viol_samples < max(0, censor_samples)):
+        # nothing can be counted, but keep the matrix shape consistent
+        return ViolationCounts(
+            unit_ids=unit_ids,
+            spike_counts=spike_counts,
+            viol_counts=np.zeros((nu, nu), dtype=np.int64),
+        )
+
+    labels = st.labels
+    times = st.times_samples
+    if count_not_sorted(times) > 0:
+        tsort = np.argsort(times, kind="stable")
+        labels = labels[tsort]
+        times = times[tsort]
+
+    # count in chunks, per thread buffer; counts are ti<=tj
+    n = times.size
+    nchunks = max(1, numba.get_num_threads())
+    nchunks = min(nchunks, max(n, 1))
+    starts = (np.arange(nchunks + 1) * n) // nchunks
+    viol_counts = np.zeros((nchunks, nu, nu), dtype=np.int64)
+
+    _violation_count_matrix(
+        times, labels, censor_samples, viol_samples, starts, viol_counts
+    )
+
+    viol_counts = viol_counts.sum(axis=0)
+    viol_diag = np.diagonal(viol_counts).copy()
+    viol_counts += viol_counts.T
+    np.fill_diagonal(viol_counts, viol_diag)
+
+    return ViolationCounts(
+        unit_ids=unit_ids,
+        spike_counts=spike_counts,
+        viol_counts=viol_counts,
+    )
+
+
+@numba.njit(nogil=True, parallel=True)
+def _violation_count_matrix(
+    times: np.ndarray,
+    labels: np.ndarray,
+    censor_samples: int,
+    viol_samples: int,
+    starts: np.ndarray,
+    counts: np.ndarray,
+):
+    n = times.shape[0]
+
+    # parallelize over chunks
+    for c in numba.prange(starts.shape[0] - 1):  # ty: ignore[not-iterable]
+        out = counts[c]
+
+        for i in range(starts[c], starts[c + 1]):
+            li = labels[i]
+            if li < 0:
+                continue
+
+            ti = times[i]
+            first = ti + censor_samples
+            last = ti + viol_samples
+
+            # make sure to read js past the chunk end!
+            for j in range(i + 1, n):
+                if times[j] < first:
+                    continue
+                if times[j] > last:  # be inclusive here i suppose
+                    break
+                lj = labels[j]
+                if lj < 0:
+                    continue
+                out[li, lj] += 1

src/dartsort/clustering/density.py

@@ -188,7 +188,7 @@ def get_smoothed_density_ratio(
     else:
         dens_ = dens[0]
         dens_ /= dens[1]
-        np.nan_to_num(dens_, out=dens_)  # type: ignore
+        np.nan_to_num(dens_, out=dens_, posinf=np.inf, neginf=-np.inf)  # type: ignore
         return dens_
 
 
@@ -407,7 +407,7 @@ def kdt_density(
             desc=f"KDTdens[{n_jobs}]",
         ):
             density[i0:i1] = dens
-    np.nan_to_num(density, copy=False)
+    np.nan_to_num(density, posinf=np.inf, neginf=-np.inf, copy=False)
     return density
 
 

src/dartsort/clustering/merge.py

@@ -358,21 +358,47 @@ def get_deconv_resid_decrease_iter(
     )
 
 
-def combine_templates(template_data_a, template_data_b):
+def combine_templates(template_data_a: TemplateData, template_data_b: TemplateData):
     rgeom = template_data_a.registered_geom
     if rgeom is not None:
+        assert template_data_b.registered_geom is not None
         if not np.array_equal(rgeom, template_data_b.registered_geom):
             raise ValueError(
                 f"Template data had different registered geoms: "
                 f"{template_data_a.registered_geom=} {template_data_b.registered_geom=}"
             )
 
+    ta = template_data_a.templates
+    tb = template_data_b.templates
+    assert ta.shape[2] == tb.shape[2]
+    if ta.shape[1] > tb.shape[1]:
+        i0 = (
+            template_data_a.trough_offset_samples
+            - template_data_b.trough_offset_samples
+        )
+        assert i0 >= 0
+        i1 = i0 + template_data_b.spike_length_samples
+        ta = ta[:, i0:i1]
+        trough_offset_samples = template_data_b.trough_offset_samples
+        sampling_frequency = template_data_b.sampling_frequency
+    elif tb.shape[1] > ta.shape[1]:
+        i0 = (
+            template_data_b.trough_offset_samples
+            - template_data_a.trough_offset_samples
+        )
+        assert i0 >= 0
+        i1 = i0 + template_data_a.spike_length_samples
+        tb = tb[:, i0:i1]
+        trough_offset_samples = template_data_a.trough_offset_samples
+        sampling_frequency = template_data_a.sampling_frequency
+    else:
+        trough_offset_samples = template_data_a.trough_offset_samples
+        sampling_frequency = template_data_a.sampling_frequency
+
     ids_a = template_data_a.unit_ids
     ids_b = template_data_b.unit_ids + ids_a.max() + 1
     unit_ids = np.concatenate((ids_a, ids_b))
-    templates = np.concatenate(
-        (template_data_a.templates, template_data_b.templates), axis=0
-    )
+    templates = np.concatenate((ta, tb), axis=0)
     spike_counts = np.concatenate(
         (template_data_a.spike_counts, template_data_b.spike_counts)
     )
@@ -389,8 +415,8 @@ def combine_templates(template_data_a, template_data_b):
         spike_counts=spike_counts,
         registered_geom=rgeom,
         spike_counts_by_channel=spike_counts_by_channel,
-        trough_offset_samples=template_data_a.trough_offset_samples,
-        sampling_frequency=template_data_a.sampling_frequency,
+        trough_offset_samples=trough_offset_samples,
+        sampling_frequency=sampling_frequency,
     )
 
     cross_mask = np.logical_and(

src/dartsort/clustering/mixture.py

@@ -5855,31 +5855,40 @@ def _noise_factors(*, noise, obs_ix, miss_near_ix, cache_prefix):
         jmnixv = jmnix[jmnixvix]
         ncoi = joixv.numel()
         ncmi = jmnixv.numel()
+        assert ncoi > 0
 
+        # -- observed-only stuff
+        # calculate observed-only stuff
         jCoo = noise.marginal_covariance(
             channels=joixv, cache_prefix=cache_prefix, cache_key=j
         )
-        jCom = noise.offdiag_covariance(
-            channels_left=joixv, channels_right=jmnixv, device=dev
-        )
-        jCom = jCom.to_dense().to(device=dev)
-
         jL = jCoo.cholesky(upper=False)  # C = LL'
         jLinv = jL.inverse().to_dense()
         jCooinv = jLinv.T @ jLinv  # Cinv = Linv' Linv
         jlogdet = 2.0 * jL.to_dense().diagonal(dim1=-2, dim2=-1).log().sum()
-        jCooinvCom = jCooinv @ jCom
 
+        # set observed-only stuff
         logdet[j] = jlogdet
         # fancy inds to front! love that.
-        jCooinv = jCooinv.view(rank, ncoi, rank, ncoi)
-        Cooinv[j, :, joixvix[:, None], :, joixvix[None]] = jCooinv.permute(1, 3, 0, 2)
+        jCooinv_out = jCooinv.view(rank, ncoi, rank, ncoi).permute(1, 3, 0, 2)
+        Cooinv[j, :, joixvix[:, None], :, joixvix[None]] = jCooinv_out
+        jLinv_out = jLinv.view(rank, ncoi, rank, ncoi).permute(1, 3, 0, 2)
+        Linv[j, :, joixvix[:, None], :, joixvix[None]] = jLinv_out
+
+        # -- observed-missing stuff
+        if ncmi == 0:
+            # i am a lonely island
+            # importantly, that means my Com is 0. so everything is 0 below.
+            continue
+        jCom = noise.offdiag_covariance(
+            channels_left=joixv, channels_right=jmnixv, device=dev
+        )
+        jCom = jCom.to_dense().to(device=dev)
+        jCooinvCom = jCooinv @ jCom
         jCooinvCom = jCooinvCom.view(rank, ncoi, rank, ncmi)
         CooinvCom[j, :, joixvix[:, None], :, jmnixvix[None]] = jCooinvCom.permute(
             1, 3, 0, 2
         )
-        jLinv = jLinv.view(rank, ncoi, rank, ncoi)
-        Linv[j, :, joixvix[:, None], :, joixvix[None]] = jLinv.permute(1, 3, 0, 2)
 
     obsdim = rank * nc_obs
     missdim = rank * nc_miss_near
@@ -6739,7 +6748,7 @@ def _update_lut_mean_batch(
 
     # constplogdet. add in the signal-rank-0-only terms.
     lut_params.constplogdet[i0:i1] = neighb_cov.nobs[nn].mul_(LOG_2PI)  # type: ignore  # ty: ignore[x]
-    lut_params.constplogdet[i0:i1] += neighb_cov.b.logdet[nn]
+    lut_params.constplogdet[i0:i1] += neighb_cov.b.logdet[nn]  # ty: ignore[not-subscriptable]
     if pnoid:
         assert lut_params.constplogdet[i0:i1].isfinite().all()  # type: ignore  # ty: ignore[x]
 
@@ -6799,7 +6808,7 @@ def _update_lut_ppca_batch(
     assert lut_params.Tpad is not None
     lut_params.Tpad[i0:i1, :, :-1] = T  # type: ignore  # ty: ignore[x]
     cap_logdet = L.diagonal(dim1=-2, dim2=-1).log().sum(dim=1).mul_(2.0)
-    lut_params.constplogdet[i0:i1] += cap_logdet
+    lut_params.constplogdet[i0:i1] += cap_logdet  # ty: ignore[not-subscriptable]
     if pnoid:
         assert lut_params.b.constplogdet[i0:i1].isfinite().all()
 

src/dartsort/config.py

@@ -205,6 +205,7 @@ class DARTsortUserConfig:
     """Upsampling of templates during matching to allow for temporal aliasing of waveforms."""
 
     # -- final merge step
+    #TODO name this more prominently, clarify flags for dedup ms, decouple dedup and agg
     agg_kind: Literal["none", "template_distance", "qda"] = "qda"
     """Final distance or GMM-based merge type."""