Expand predictive output options and update tests

Added 'predictive_posterior' as a valid output option for model results, updated output normalization logic, and clarified documentation for output choices and memory usage. Refactored and extended tests to verify correct behavior for new and existing output combinations, and improved code formatting and assertion messages for clarity and consistency.

Author: edeno

src/non_local_detector/likelihoods/clusterless_gmm.py

@@ -16,10 +16,8 @@
 from non_local_detector.environment import Environment
 from non_local_detector.likelihoods.common import (
     EPS,
-    LOG_EPS,
     get_position_at_time,
     get_spike_time_bin_ind,
-    safe_divide,
 )
 from non_local_detector.likelihoods.gmm import GaussianMixtureModel
 
@@ -326,7 +324,6 @@ def fit_clusterless_gmm_encoding_model(
         unit="electrode",
         disable=disable_progress_bar,
     ):
-
         # Clip to encoding window
         in_bounds = np.logical_and(
             elect_times >= position_time[0], elect_times <= position_time[-1]
@@ -483,7 +480,6 @@ def predict_clusterless_gmm_log_likelihood(
         unit="electrode",
         disable=disable_progress_bar,
     ):
-
         # Clip to decoding window
         in_bounds = np.logical_and(elect_times >= time[0], elect_times <= time[-1])
         elect_times = elect_times[in_bounds]

src/non_local_detector/models/base.py

@@ -62,14 +62,21 @@
 }
 
 # Valid options for return_outputs parameter
-VALID_OUTPUTS: set[str] = {"filter", "predictive", "log_likelihood", "all"}
+VALID_OUTPUTS: set[str] = {
+    "filter",
+    "predictive",
+    "predictive_posterior",
+    "log_likelihood",
+    "all",
+}
 
 # Mapping of single string options to sets of outputs
 OUTPUT_INCLUDES: dict[str, set[str]] = {
     "filter": {"filter"},
-    "predictive": {"predictive"},
+    "predictive": {"predictive", "predictive_posterior"},
+    "predictive_posterior": {"predictive_posterior"},
     "log_likelihood": {"log_likelihood"},
-    "all": {"filter", "predictive", "log_likelihood"},
+    "all": {"filter", "predictive", "predictive_posterior", "log_likelihood"},
 }
 
 
@@ -86,7 +93,8 @@ def _normalize_return_outputs(
     Returns
     -------
     set of str
-        Normalized set containing any of: 'filter', 'predictive', 'log_likelihood'
+        Normalized set containing any of: 'filter', 'predictive',
+        'predictive_posterior', 'log_likelihood'
 
     Raises
     ------
@@ -2233,22 +2241,24 @@ def predict(
             Options:
             - None: smoother only (default, minimal memory)
             - 'filter': filtered (causal) posterior and state probabilities
-            - 'predictive': one-step-ahead predictive state distributions
+            - 'predictive': both aggregated and full predictive distributions
+            - 'predictive_posterior': only full predictive posterior (state bins)
             - 'log_likelihood': per-timepoint log likelihoods
             - 'all': all outputs above
-            - List/set: e.g., ['filter', 'predictive'] for multiple outputs
+            - List/set: e.g., ['filter', 'log_likelihood'] for multiple outputs
 
             The smoother (acausal_posterior, acausal_state_probabilities) and
             marginal_log_likelihood are ALWAYS included.
 
             When to use each output:
             - 'filter': Online/causal decoding, debugging forward pass
-            - 'predictive': Model evaluation, predictive checks
+            - 'predictive': Model evaluation, predictive checks (includes both formats)
+            - 'predictive_posterior': When you only need full distribution, not aggregated
             - 'log_likelihood': Diagnostics, per-timepoint metrics, model comparison
 
-            Memory warning: 'log_likelihood' and 'filter' can be very large
-            (~400 GB for 1M timepoints × 100k spatial bins). Only request
-            what you need for your analysis.
+            Memory warning: 'log_likelihood', 'filter', 'predictive', and
+            'predictive_posterior' can be very large (~400 GB for 1M timepoints × 100k
+            spatial bins). Use None for minimal memory (smoother only).
         save_log_likelihood_to_results : bool, optional
             DEPRECATED. Use return_outputs='log_likelihood' instead.
             Whether to save the log likelihood to the results, by default None.
@@ -2276,8 +2286,9 @@ def predict(
                 Filtered discrete state probabilities
             - predictive_state_probabilities : (n_time, n_states) - if 'predictive'
                 One-step-ahead predictive distributions over discrete states
-            - predictive_posterior : (n_time, n_state_bins) - if 'predictive'
+            - predictive_posterior : (n_time, n_state_bins) - if 'predictive_posterior'
                 One-step-ahead predictive distributions over state bins
+                (Warning: can be very large, ~same size as causal_posterior)
             - log_likelihood : (n_time, n_state_bins) - if 'log_likelihood'
                 Per-timepoint observation log likelihoods
 
@@ -2406,7 +2417,9 @@ def predict(
                 else None
             ),
             predictive_posterior=(
-                predictive_posterior if "predictive" in requested_outputs else None
+                predictive_posterior
+                if "predictive_posterior" in requested_outputs
+                else None
             ),
         )
 
@@ -3087,7 +3100,9 @@ def predict(
                 else None
             ),
             predictive_posterior=(
-                predictive_posterior if "predictive" in requested_outputs else None
+                predictive_posterior
+                if "predictive_posterior" in requested_outputs
+                else None
             ),
         )
 

src/non_local_detector/tests/likelihoods/test_clusterless_kde_optimization.py

@@ -37,8 +37,9 @@ def test_numerical_equivalence_high_dim(self):
             result_original = kde_distance(eval_points, samples, std)
             result_vectorized = kde_distance_vectorized(eval_points, samples, std)
 
-            assert jnp.allclose(result_original, result_vectorized, rtol=1e-5, atol=1e-8), \
-                f"Failed for {n_features}D"
+            assert jnp.allclose(
+                result_original, result_vectorized, rtol=1e-5, atol=1e-8
+            ), f"Failed for {n_features}D"
 
     def test_numerical_stability(self):
         """Test numerical stability with small std values."""
@@ -198,8 +199,12 @@ def test_jit_compilation_caching(self):
 
         # Warmup call
         result_warmup = estimate_log_joint_mark_intensity(
-            dec_features, enc_features, waveform_stds,
-            occupancy, mean_rate, position_distance
+            dec_features,
+            enc_features,
+            waveform_stds,
+            occupancy,
+            mean_rate,
+            position_distance,
         )
         result_warmup.block_until_ready()
 
@@ -208,24 +213,37 @@ def test_jit_compilation_caching(self):
         for _ in range(10):
             start = time.perf_counter()
             result = estimate_log_joint_mark_intensity(
-                dec_features, enc_features, waveform_stds,
-                occupancy, mean_rate, position_distance
+                dec_features,
+                enc_features,
+                waveform_stds,
+                occupancy,
+                mean_rate,
+                position_distance,
             )
             result.block_until_ready()
             times.append(time.perf_counter() - start)
 
         # All calls should be fast (< 10ms) after compilation
         avg_time = np.mean(times)
-        assert avg_time < 0.01, \
+        assert avg_time < 0.01, (
             f"Average call time ({avg_time:.4f}s) too slow, JIT may not be working"
+        )
 
         # Results should be consistent
         result1 = estimate_log_joint_mark_intensity(
-            dec_features, enc_features, waveform_stds,
-            occupancy, mean_rate, position_distance
+            dec_features,
+            enc_features,
+            waveform_stds,
+            occupancy,
+            mean_rate,
+            position_distance,
         )
         result2 = estimate_log_joint_mark_intensity(
-            dec_features, enc_features, waveform_stds,
-            occupancy, mean_rate, position_distance
+            dec_features,
+            enc_features,
+            waveform_stds,
+            occupancy,
+            mean_rate,
+            position_distance,
         )
         assert jnp.allclose(result1, result2)

src/non_local_detector/tests/likelihoods/test_enc_tile_size.py

@@ -148,9 +148,9 @@ def test_enc_tile_size_with_pos_tile_size(enc_tile_size, pos_tile_size):
 
     # Should match
     max_diff = np.max(np.abs(result_baseline - result_both_tiling))
-    assert np.allclose(
-        result_baseline, result_both_tiling, rtol=1e-5, atol=1e-7
-    ), f"enc_tile_size={enc_tile_size}, pos_tile_size={pos_tile_size}: Max diff = {max_diff}"
+    assert np.allclose(result_baseline, result_both_tiling, rtol=1e-5, atol=1e-7), (
+        f"enc_tile_size={enc_tile_size}, pos_tile_size={pos_tile_size}: Max diff = {max_diff}"
+    )
 
 
 @pytest.mark.unit

src/non_local_detector/tests/likelihoods/test_gmm_kde_convergence.py

@@ -170,20 +170,24 @@ def normalize_per_time(ll):
         peak_gmm = np.argmax(ll_gmm, axis=1)
         peak_agreement = np.mean(peak_kde == peak_gmm)
 
-        print(f"{n_comp:4d}       | {corr:11.4f} | {mse:16.6f} | {peak_agreement:14.1%}")
+        print(
+            f"{n_comp:4d}       | {corr:11.4f} | {mse:16.6f} | {peak_agreement:14.1%}"
+        )
 
     # Verify convergence trend
     print("\n=== Convergence Analysis ===")
     print(f"Correlation improvement: {correlations[0]:.4f} → {correlations[-1]:.4f}")
     print(f"MSE improvement: {mse_values[0]:.6f} → {mse_values[-1]:.6f}")
 
     # Key assertion: correlation should increase with more components
-    assert correlations[-1] > correlations[0], \
+    assert correlations[-1] > correlations[0], (
         f"Correlation should increase: {correlations[0]:.3f} → {correlations[-1]:.3f}"
+    )
 
     # MSE should decrease
-    assert mse_values[-1] < mse_values[0], \
+    assert mse_values[-1] < mse_values[0], (
         f"MSE should decrease: {mse_values[0]:.4f} → {mse_values[-1]:.4f}"
+    )
 
 
 def test_mathematical_formula_consistency(convergence_test_data):
@@ -219,12 +223,17 @@ def test_mathematical_formula_consistency(convergence_test_data):
     gmm_formula = np.log(mean_rate) + np.log(marginal_density) - np.log(occupancy)
 
     print("\n=== Formula Verification ===")
-    print(f"KDE: log({mean_rate} * {marginal_density} / {occupancy}) = {kde_formula:.6f}")
-    print(f"GMM: log({mean_rate}) + log({marginal_density}) - log({occupancy}) = {gmm_formula:.6f}")
+    print(
+        f"KDE: log({mean_rate} * {marginal_density} / {occupancy}) = {kde_formula:.6f}"
+    )
+    print(
+        f"GMM: log({mean_rate}) + log({marginal_density}) - log({occupancy}) = {gmm_formula:.6f}"
+    )
     print(f"Difference: {abs(kde_formula - gmm_formula):.10f}")
 
-    assert np.isclose(kde_formula, gmm_formula, rtol=1e-10), \
+    assert np.isclose(kde_formula, gmm_formula, rtol=1e-10), (
         "KDE and GMM formulas should be mathematically identical"
+    )
 
 
 def test_ground_process_intensity_calculation(convergence_test_data):
@@ -298,18 +307,25 @@ def test_segment_sum_correctness(convergence_test_data):
     segment_ids = jnp.array([0, 0, 1, 1, 2])
 
     # KDE way
-    result_kde = jax.ops.segment_sum(values, segment_ids, num_segments=3, indices_are_sorted=True)
+    result_kde = jax.ops.segment_sum(
+        values, segment_ids, num_segments=3, indices_are_sorted=True
+    )
 
     # GMM way (from jax.ops import segment_sum)
     from jax.ops import segment_sum
-    result_gmm = segment_sum(values, segment_ids, num_segments=3, indices_are_sorted=True)
+
+    result_gmm = segment_sum(
+        values, segment_ids, num_segments=3, indices_are_sorted=True
+    )
 
     print("\n=== segment_sum Verification ===")
     print(f"Input: values={values}, segment_ids={segment_ids}")
     print(f"KDE result: {result_kde}")
     print(f"GMM result: {result_gmm}")
 
-    assert jnp.allclose(result_kde, result_gmm), "segment_sum implementations should match"
+    assert jnp.allclose(result_kde, result_gmm), (
+        "segment_sum implementations should match"
+    )
 
 
 def test_log_space_operations(convergence_test_data):

src/non_local_detector/tests/likelihoods/test_kde_gmm_comparison.py

@@ -66,9 +66,7 @@ def shared_simulation_data():
 
     # Decoding period spikes (subset of encoding spikes for simplicity)
     decoding_spike_times = [times[:20] for times in encoding_spike_times]
-    decoding_spike_features = [
-        feats[:20] for feats in encoding_spike_features
-    ]
+    decoding_spike_features = [feats[:20] for feats in encoding_spike_features]
 
     # Create and fit environment
     environment = Environment(position_range=[(0, 10), (-3, 3)])
@@ -159,9 +157,7 @@ def test_kde_end_to_end_pipeline(shared_simulation_data):
         encoding_positions=kde_encoding["encoding_positions"],
         environment=data["environment"],
         mean_rates=jnp.asarray(kde_encoding["mean_rates"]),
-        summed_ground_process_intensity=kde_encoding[
-            "summed_ground_process_intensity"
-        ],
+        summed_ground_process_intensity=kde_encoding["summed_ground_process_intensity"],
         position_std=jnp.asarray(kde_encoding["position_std"]),
         waveform_std=jnp.asarray(kde_encoding["waveform_std"]),
         is_local=False,
@@ -191,9 +187,7 @@ def test_kde_end_to_end_pipeline(shared_simulation_data):
         encoding_positions=kde_encoding["encoding_positions"],
         environment=data["environment"],
         mean_rates=jnp.asarray(kde_encoding["mean_rates"]),
-        summed_ground_process_intensity=kde_encoding[
-            "summed_ground_process_intensity"
-        ],
+        summed_ground_process_intensity=kde_encoding["summed_ground_process_intensity"],
         position_std=jnp.asarray(kde_encoding["position_std"]),
         waveform_std=jnp.asarray(kde_encoding["waveform_std"]),
         is_local=True,
@@ -404,9 +398,7 @@ def test_api_consistency_predict_functions(shared_simulation_data):
         encoding_positions=kde_encoding["encoding_positions"],
         environment=data["environment"],
         mean_rates=kde_encoding["mean_rates"],
-        summed_ground_process_intensity=kde_encoding[
-            "summed_ground_process_intensity"
-        ],
+        summed_ground_process_intensity=kde_encoding["summed_ground_process_intensity"],
         position_std=kde_encoding["position_std"],
         waveform_std=kde_encoding["waveform_std"],
         is_local=False,
@@ -454,7 +446,10 @@ def test_kde_gmm_output_shape_consistency(shared_simulation_data):
 
     kde_enc = fit_clusterless_kde_encoding_model(**common_params, position_std=1.0)
     gmm_enc = fit_clusterless_gmm_encoding_model(
-        **common_params, gmm_components_occupancy=4, gmm_components_gpi=4, gmm_components_joint=8
+        **common_params,
+        gmm_components_occupancy=4,
+        gmm_components_gpi=4,
+        gmm_components_joint=8,
     )
 
     # Predict with both
@@ -522,7 +517,10 @@ def test_both_support_local_and_nonlocal_modes(shared_simulation_data):
     # Fit models
     kde_enc = fit_clusterless_kde_encoding_model(**common_params, position_std=1.0)
     gmm_enc = fit_clusterless_gmm_encoding_model(
-        **common_params, gmm_components_occupancy=4, gmm_components_gpi=4, gmm_components_joint=8
+        **common_params,
+        gmm_components_occupancy=4,
+        gmm_components_gpi=4,
+        gmm_components_joint=8,
     )
 
     time = jnp.asarray(data["time"])