Add predictive_posterior output to filter/smoother

Extended core filter/smoother functions and model base classes to compute and return predictive_posterior (one-step-ahead predicted state probabilities over state bins). Updated xarray dataset conversion and output handling to include predictive_posterior. Adjusted tests to verify presence and correctness of predictive_posterior in results.

Author: edeno

src/non_local_detector/core.py

@@ -263,7 +263,14 @@ def chunked_filter_smoother(
     cache_log_likelihoods: bool = True,
     dtype: jnp.dtype = jnp.float32,
 ) -> tuple[
-    np.ndarray, np.ndarray, float, np.ndarray, np.ndarray, np.ndarray, np.ndarray
+    np.ndarray,
+    np.ndarray,
+    float,
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
 ]:
     """Filter and smooth the state probabilities in chunks.
 
@@ -302,8 +309,11 @@ def chunked_filter_smoother(
         Log likelihoods for each state at each time point
     causal_posterior : np.ndarray, shape (n_time, n_state_bins)
         Filtered state probabilities
+    predictive_posterior : np.ndarray, shape (n_time, n_state_bins)
+        One-step-ahead predicted state probabilities over state bins
     """
     causal_posterior = []
+    predictive_posterior = []
     predictive_state_probabilities = []
     causal_state_probabilities = []
     acausal_posterior = []
@@ -394,6 +404,7 @@ def chunked_filter_smoother(
         causal_state_probabilities.append(
             causal_posterior_chunk @ state_aggregation_matrix
         )
+        predictive_posterior.append(predicted_probs_chunk)
         predictive_state_probabilities.append(
             predicted_probs_chunk @ state_aggregation_matrix
         )
@@ -403,6 +414,7 @@ def chunked_filter_smoother(
     # Concatenate JAX arrays on device
     causal_posterior_jax = jnp.concatenate(causal_posterior)
     causal_state_probabilities_jax = jnp.concatenate(causal_state_probabilities)
+    predictive_posterior_jax = jnp.concatenate(predictive_posterior)
     predictive_state_probabilities_jax = jnp.concatenate(predictive_state_probabilities)
 
     # Backward pass: accumulate JAX arrays
@@ -439,6 +451,7 @@ def chunked_filter_smoother(
         np.asarray(predictive_state_probabilities_jax),
         log_likelihoods,  # Keep as original (may be None or NumPy)
         np.asarray(causal_posterior_jax),
+        np.asarray(predictive_posterior_jax),
     )
 
 
@@ -751,7 +764,14 @@ def chunked_filter_smoother_covariate_dependent(
     cache_log_likelihoods: bool = True,
     dtype: jnp.dtype = jnp.float32,
 ) -> tuple[
-    np.ndarray, np.ndarray, float, np.ndarray, np.ndarray, np.ndarray, np.ndarray
+    np.ndarray,
+    np.ndarray,
+    float,
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
 ]:
     """Filter and smooth the state probabilities in chunks with covariate dependent transitions.
 
@@ -790,8 +810,11 @@ def chunked_filter_smoother_covariate_dependent(
         Log likelihoods for each state at each time point
     causal_posterior : np.ndarray, shape (n_time, n_state_bins)
         Filtered state probabilities
+    predictive_posterior : np.ndarray, shape (n_time, n_state_bins)
+        One-step-ahead predicted state probabilities over state bins
     """
     causal_posterior = []
+    predictive_posterior = []
     predictive_state_probabilities = []
     causal_state_probabilities = []
     acausal_posterior = []
@@ -893,6 +916,7 @@ def chunked_filter_smoother_covariate_dependent(
         causal_state_probabilities.append(
             causal_posterior_chunk @ state_aggregation_matrix
         )
+        predictive_posterior.append(predicted_probs_chunk)
         predictive_state_probabilities.append(
             predicted_probs_chunk @ state_aggregation_matrix
         )
@@ -902,6 +926,7 @@ def chunked_filter_smoother_covariate_dependent(
     # Concatenate JAX arrays on device
     causal_posterior_jax = jnp.concatenate(causal_posterior)
     causal_state_probabilities_jax = jnp.concatenate(causal_state_probabilities)
+    predictive_posterior_jax = jnp.concatenate(predictive_posterior)
     predictive_state_probabilities_jax = jnp.concatenate(predictive_state_probabilities)
 
     # Backward pass: accumulate JAX arrays
@@ -946,6 +971,7 @@ def chunked_filter_smoother_covariate_dependent(
         np.asarray(predictive_state_probabilities_jax),
         log_likelihoods,  # Keep as original (may be None or NumPy)
         np.asarray(causal_posterior_jax),
+        np.asarray(predictive_posterior_jax),
     )
 
 

src/non_local_detector/models/base.py

@@ -106,7 +106,7 @@ def _normalize_return_outputs(
             )
         return OUTPUT_INCLUDES.get(return_outputs, {return_outputs})
 
-    if isinstance(return_outputs, (list, set)):
+    if isinstance(return_outputs, list | set):
         outputs_set = set(return_outputs)
         invalid = outputs_set - VALID_OUTPUTS
         if invalid:
@@ -1063,7 +1063,14 @@ def _predict(
         cache_likelihood: bool = True,
         n_chunks: int = 1,
     ) -> tuple[
-        np.ndarray, np.ndarray, float, np.ndarray, np.ndarray, np.ndarray, np.ndarray
+        np.ndarray,
+        np.ndarray,
+        float,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
     ]:
         """
         Compute the posterior probabilities.
@@ -1092,6 +1099,7 @@ def _predict(
         predictive_state_probabilities : np.ndarray, shape (n_time, n_states)
         log_likelihoods : np.ndarray, shape (n_time, n_state_bins)
         causal_posterior : np.ndarray, shape (n_time, n_state_bins)
+        predictive_posterior : np.ndarray, shape (n_time, n_state_bins)
         """
 
         logger.info("Computing posterior...")
@@ -1549,6 +1557,7 @@ def _convert_results_to_xarray(
         causal_posterior: np.ndarray | None = None,
         causal_state_probabilities: np.ndarray | None = None,
         predictive_state_probabilities: np.ndarray | None = None,
+        predictive_posterior: np.ndarray | None = None,
     ) -> xr.Dataset:
         """
         Convert the results to an xarray Dataset.
@@ -1571,6 +1580,8 @@ def _convert_results_to_xarray(
             Causal state probabilities, by default None.
         predictive_state_probabilities : np.ndarray, optional, shape (n_time, n_states)
             One-step-ahead predicted state probabilities, by default None.
+        predictive_posterior : np.ndarray, optional, shape (n_time, n_state_bins)
+            One-step-ahead predicted posterior probabilities over state bins, by default None.
 
         Returns
         -------
@@ -1698,6 +1709,14 @@ def _convert_results_to_xarray(
                 predictive_state_probabilities,
             )
 
+        if predictive_posterior is not None:
+            data_vars["predictive_posterior"] = (
+                ("time", "state_bins"),
+                self._create_masked_posterior(
+                    predictive_posterior, is_track_interior, n_total_bins
+                ),
+            )
+
         # Create Dataset with MultiIndex coordinates
         results = xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs)
 
@@ -2256,7 +2275,9 @@ def predict(
             - causal_state_probabilities : (n_time, n_states) - if 'filter'
                 Filtered discrete state probabilities
             - predictive_state_probabilities : (n_time, n_states) - if 'predictive'
-                One-step-ahead predictive distributions
+                One-step-ahead predictive distributions over discrete states
+            - predictive_posterior : (n_time, n_state_bins) - if 'predictive'
+                One-step-ahead predictive distributions over state bins
             - log_likelihood : (n_time, n_state_bins) - if 'log_likelihood'
                 Per-timepoint observation log likelihoods
 
@@ -2351,6 +2372,7 @@ def predict(
             predictive_state_probabilities,
             log_likelihood,
             causal_posterior,
+            predictive_posterior,
         ) = self._predict(
             time=time,
             log_likelihood_args=(
@@ -2383,6 +2405,9 @@ def predict(
                 if "predictive" in requested_outputs
                 else None
             ),
+            predictive_posterior=(
+                predictive_posterior if "predictive" in requested_outputs else None
+            ),
         )
 
     def estimate_parameters(
@@ -3029,6 +3054,7 @@ def predict(
             predictive_state_probabilities,
             log_likelihood,
             causal_posterior,
+            predictive_posterior,
         ) = self._predict(
             time=time,
             log_likelihood_args=(
@@ -3060,6 +3086,9 @@ def predict(
                 if "predictive" in requested_outputs
                 else None
             ),
+            predictive_posterior=(
+                predictive_posterior if "predictive" in requested_outputs else None
+            ),
         )
 
     def estimate_parameters(

src/non_local_detector/tests/core/test_chunked_parity.py

@@ -70,6 +70,7 @@ def log_likelihood_func(time_idx, *args):
             predictive_state_probabilities,  # one-step-ahead predictions
             _,
             causal_posterior,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -114,6 +115,7 @@ def log_likelihood_func(time_idx, *args):
             _,  # predictive_state_probabilities
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -155,6 +157,7 @@ def log_likelihood_func(time_idx, *args):
             _,  # predictive_state_probabilities
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -175,6 +178,7 @@ def log_likelihood_func(time_idx, *args):
             _,  # predictive_state_probabilities
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -211,7 +215,7 @@ def log_likelihood_func(time_idx, *args):
         total_log_like_std = log_marginals.sum()
 
         # Act - Chunked version
-        (_, _, log_like_chunked, _, _, _, _) = chunked_filter_smoother(
+        (_, _, log_like_chunked, _, _, _, _, _) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
             initial_distribution=np.array(init),
@@ -251,6 +255,7 @@ def log_likelihood_func(time_idx, *args):
             _,  # predictive_state_probabilities
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -286,6 +291,7 @@ def log_likelihood_func(time_idx, *args):
             _,  # predictive_state_probabilities
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -330,6 +336,7 @@ def log_likelihood_func(time_idx, *args):
             _,  # predictive_state_probabilities
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -369,6 +376,7 @@ def log_likelihood_func(time_idx, *args):
             _,
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -390,6 +398,7 @@ def log_likelihood_func(time_idx, *args):
             _,
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,
@@ -459,6 +468,7 @@ def log_likelihood_func(time_idx, *args):
             _,  # predictive_state_probabilities
             _,
             _,
+            _,  # predictive_posterior (not used in this test)
         ) = chunked_filter_smoother(
             time=time,
             state_ind=state_ind,

src/non_local_detector/tests/models/test_return_outputs.py

@@ -112,22 +112,31 @@ def test_return_predictive_string(self, simple_fitted_detector):
         # Should have smoother (always)
         assert "acausal_posterior" in results
 
-        # Should have predictive
+        # Should have predictive (both aggregated and full)
         assert "predictive_state_probabilities" in results
+        assert "predictive_posterior" in results
 
         # Should NOT have filter or log_likelihood
         assert "causal_posterior" not in results
         assert "log_likelihood" not in results
 
-        # Verify shape (aggregated to discrete states)
+        # Verify shape of aggregated version (discrete states)
         n_time = len(time)
         n_states = results.acausal_state_probabilities.shape[1]
         assert results.predictive_state_probabilities.shape == (n_time, n_states)
 
-        # Verify probabilities sum to 1
+        # Verify shape of full version (state bins)
+        n_state_bins = results.acausal_posterior.shape[1]
+        assert results.predictive_posterior.shape == (n_time, n_state_bins)
+
+        # Verify probabilities sum to 1 (aggregated version)
         predictive_sums = results.predictive_state_probabilities.sum(dim="states")
         assert np.allclose(predictive_sums.values, 1.0, atol=1e-10)
 
+        # Verify probabilities sum to 1 (full version)
+        predictive_posterior_sums = results.predictive_posterior.sum(dim="state_bins")
+        assert np.allclose(predictive_posterior_sums.values, 1.0, atol=1e-10)
+
     def test_return_log_likelihood_string(self, simple_fitted_detector):
         """Test return_outputs='log_likelihood' returns log likelihoods."""
         detector, spike_times, time, position = simple_fitted_detector
@@ -173,6 +182,7 @@ def test_return_all_string(self, simple_fitted_detector):
         assert "causal_posterior" in results
         assert "causal_state_probabilities" in results
         assert "predictive_state_probabilities" in results
+        assert "predictive_posterior" in results
         assert "log_likelihood" in results
         assert "marginal_log_likelihoods" in results.attrs
 
@@ -195,6 +205,7 @@ def test_return_multiple_outputs_list(self, simple_fitted_detector):
         assert "causal_posterior" in results
         assert "causal_state_probabilities" in results
         assert "predictive_state_probabilities" in results
+        assert "predictive_posterior" in results
 
         # Should NOT have log_likelihood
         assert "log_likelihood" not in results