Refactor get_spike_time_bin_ind to common module

Moved the get_spike_time_bin_ind function to likelihoods/common.py and updated imports in clusterless_gmm.py, clusterless_kde.py, and clusterless_kde_log.py. Removed duplicate implementations from the affected files for consistency and code reuse. Minor docstring and code cleanup included.

Author: edeno

src/non_local_detector/likelihoods/clusterless_gmm.py

@@ -14,7 +14,12 @@
 from track_linearization import get_linearized_position  # type: ignore[import-untyped]
 
 from non_local_detector.environment import Environment
-from non_local_detector.likelihoods.common import EPS, get_position_at_time, safe_divide
+from non_local_detector.likelihoods.common import (
+    EPS,
+    get_position_at_time,
+    get_spike_time_bin_ind,
+    safe_divide,
+)
 from non_local_detector.likelihoods.gmm import GaussianMixtureModel
 
 # ---------------------------------------------------------------------
@@ -38,30 +43,6 @@ def _as_jnp(x) -> jnp.ndarray:
     return x if isinstance(x, jnp.ndarray) else jnp.asarray(x)
 
 
-def get_spike_time_bin_ind(
-    spike_times: jnp.ndarray, time_bin_edges: jnp.ndarray
-) -> jnp.ndarray:
-    """Map spike times to decoding time-bin indices on device.
-
-    Parameters
-    ----------
-    spike_times : jnp.ndarray, shape (n_spikes,)
-        Array of spike times to map to bins.
-    time_bin_edges : jnp.ndarray, shape (n_bins + 1,)
-        Edges of time bins defining intervals [t0, t1), ..., [t_{n-1}, tn].
-
-    Returns
-    -------
-    bin_indices : jnp.ndarray, shape (n_spikes,)
-        Index of time bin for each spike (0 to n_bins-1).
-    """
-    # Right-closed bins [t_i, t_{i+1}), except the last edge which is included
-    # Use JAX ops to keep everything on device
-    inds = jnp.searchsorted(time_bin_edges, spike_times, side="right") - 1
-    last = jnp.isclose(spike_times, time_bin_edges[-1])
-    return jnp.where(last, time_bin_edges.shape[0] - 2, inds).astype(jnp.int32)
-
-
 def _gmm_logp(gmm: GaussianMixtureModel, X: jnp.ndarray) -> jnp.ndarray:
     """Log density under a fitted GMM.
 

src/non_local_detector/likelihoods/clusterless_kde.py

@@ -12,25 +12,10 @@
     block_kde,
     gaussian_pdf,
     get_position_at_time,
+    get_spike_time_bin_ind,
 )
 
 
-def get_spike_time_bin_ind(spike_times: np.ndarray, time: np.ndarray) -> np.ndarray:
-    """Get the index of the time bin for each spike time.
-
-    Parameters
-    ----------
-    spike_times : np.ndarray, shape (n_spikes,)
-    time : np.ndarray, shape (n_time_bins,)
-        Bin edges.
-
-    Returns
-    -------
-    ind : np.ndarray, shape (n_spikes,)
-    """
-    return np.digitize(spike_times, time[1:-1])
-
-
 def kde_distance(
     eval_points: jnp.ndarray, samples: jnp.ndarray, std: jnp.ndarray
 ) -> jnp.ndarray:

src/non_local_detector/likelihoods/clusterless_kde_log.py

@@ -12,27 +12,12 @@
     block_kde,
     gaussian_pdf,
     get_position_at_time,
+    get_spike_time_bin_ind,
     log_gaussian_pdf,
     safe_log,
 )
 
 
-def get_spike_time_bin_ind(spike_times: np.ndarray, time: np.ndarray) -> np.ndarray:
-    """Get the index of the time bin for each spike time.
-
-    Parameters
-    ----------
-    spike_times : np.ndarray, shape (n_spikes,)
-    time : np.ndarray, shape (n_time_bins,)
-        Bin edges.
-
-    Returns
-    -------
-    ind : np.ndarray, shape (n_spikes,)
-    """
-    return np.digitize(spike_times, time[1:-1])
-
-
 @jax.jit
 def kde_distance(
     eval_points: jnp.ndarray, samples: jnp.ndarray, std: jnp.ndarray
@@ -123,6 +108,7 @@ def log_gaussian_per_dim(eval_dim, sample_dim, sigma):
     return jnp.sum(per_dim_log_distances, axis=0)
 
 
+@jax.jit
 def log_kde_distance_streaming(
     eval_points: jnp.ndarray,
     samples: jnp.ndarray,
@@ -152,6 +138,10 @@ def log_kde_distance_streaming(
 
     Notes
     -----
+    This function is JIT-compiled and will be specialized for each unique combination
+    of input shapes. The shape dimensions (n_dims, n_samp, n_eval) are traced during
+    compilation, so different shapes will result in separate compiled versions.
+
     Memory usage:
     - log_kde_distance (vmap): O(D×n_samp×n_eval) peak
     - log_kde_distance_streaming (fori_loop): O(n_samp×n_eval) peak
@@ -489,6 +479,7 @@ def process_pos_tile(
                 )(log_pos_tile, logK_mark_chunk.T)
 
                 # Update output with this tile
+                # fori_loop handles in-place updates efficiently when possible
                 return jax.lax.dynamic_update_slice(
                     log_marginal_chunk, log_marginal_tile, (0, pos_start)
                 )
@@ -838,15 +829,6 @@ def block_estimate_log_joint_mark_intensity(
     if n_decoding_spikes == 0:
         return jnp.full((0, n_position_bins), LOG_EPS)
 
-    # Create JIT-compiled update function with buffer donation
-    # donate_argnums=(0,) allows JAX to reuse the output buffer in-place
-    _update_block = jax.jit(
-        lambda out_array, block_result, start_idx: jax.lax.dynamic_update_slice(
-            out_array, block_result, (start_idx, 0)
-        ),
-        donate_argnums=(0,),
-    )
-
     out = jnp.zeros((n_decoding_spikes, n_position_bins))
     for start_ind in range(0, n_decoding_spikes, block_size):
         block_inds = slice(start_ind, start_ind + block_size)
@@ -865,7 +847,7 @@ def block_estimate_log_joint_mark_intensity(
             position_eval_points=position_eval_points,
             position_std=position_std,
         )
-        out = _update_block(out, block_result, start_ind)
+        out = jax.lax.dynamic_update_slice(out, block_result, (start_ind, 0))
 
     return jnp.clip(out, a_min=LOG_EPS, a_max=None)
 
@@ -1171,12 +1153,12 @@ def predict_clusterless_kde_log_likelihood(
                     enc_tile_size=enc_tile_size,
                     pos_tile_size=pos_tile_size,
                     use_streaming=use_streaming,
-                    encoding_positions=electrode_encoding_positions
-                    if use_streaming
-                    else None,
-                    position_eval_points=interior_place_bin_centers
-                    if use_streaming
-                    else None,
+                    encoding_positions=(
+                        electrode_encoding_positions if use_streaming else None
+                    ),
+                    position_eval_points=(
+                        interior_place_bin_centers if use_streaming else None
+                    ),
                     position_std=position_std if use_streaming else None,
                 ),
                 get_spike_time_bin_ind(electrode_spike_times, time),

src/non_local_detector/likelihoods/common.py

@@ -50,6 +50,22 @@ def get_position_at_time(
     return position_at_spike_times
 
 
+def get_spike_time_bin_ind(spike_times: np.ndarray, time: np.ndarray) -> np.ndarray:
+    """Get the index of the time bin for each spike time.
+
+    Parameters
+    ----------
+    spike_times : np.ndarray, shape (n_spikes,)
+    time : np.ndarray, shape (n_time_bins,)
+        Bin edges.
+
+    Returns
+    -------
+    ind : np.ndarray, shape (n_spikes,)
+    """
+    return np.digitize(spike_times, time[1:-1])
+
+
 @jax.jit
 def log_gaussian_pdf(
     x: jnp.ndarray, mean: jnp.ndarray, sigma: jnp.ndarray