Allow per-state discrete transition prior_weight

Generalize prior_weight in estimate_stationary_state_transition to
accept a scalar or per-state array. Rows with prior_weight[i] == 0 use
the legacy fixed-count Dirichlet prior, while rows with prior_weight[i]
> 0 use the data-adaptive scaling. This lets frozen structural rows
(e.g. very sticky Local/No-Spike states) coexist with learnable rows
that want size-invariant regularization in the same EM update.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: edeno

src/non_local_detector/discrete_state_transitions.py

@@ -354,7 +354,7 @@ def estimate_stationary_state_transition(
     acausal_posterior: np.ndarray,
     stickiness: float = 0.0,
     concentration: float = 1.0,
-    prior_weight: float = 0.0,
+    prior_weight: float | np.ndarray = 0.0,
 ) -> np.ndarray:
     """Estimate the stationary state transition model.
 
@@ -366,20 +366,38 @@ def estimate_stationary_state_transition(
     acausal_posterior : np.ndarray, shape (n_time, n_states)
     stickiness : float, optional
     concentration : float, optional
-    prior_weight : float, optional
+    prior_weight : float or np.ndarray, shape (n_states,), optional
         Dimensionless weight for data-adaptive prior scaling. When > 0,
-        the effective prior pseudo-counts are scaled by the expected
-        transition count per state, making the prior influence approximately
-        invariant to the number of time bins. When 0.0 (default), the
-        fixed-count prior from ``concentration`` and ``stickiness`` is used
-        directly (legacy behavior).
+        the effective prior pseudo-counts for that row are scaled by the
+        expected transition count N_i, making the prior influence
+        approximately invariant to the number of time bins.
+
+        Can be a scalar (same weight for all rows) or an array with one
+        value per state. Rows with ``prior_weight[i] == 0`` use the legacy
+        fixed-count prior from ``concentration`` and ``stickiness``
+        directly. This allows mixing frozen rows (e.g., structural states
+        with very high stickiness) with data-adaptive rows in the same
+        call. When 0.0 (default), all rows use legacy behavior.
 
     Returns
     -------
     new_transition_matrix : np.ndarray, shape (n_states, n_states)
     """
-    if prior_weight < 0:
-        raise ValueError(f"prior_weight must be non-negative, got {prior_weight}")
+    n_states = acausal_posterior.shape[1]
+
+    # Normalize prior_weight to shape (n_states,)
+    prior_weight_arr = np.atleast_1d(np.asarray(prior_weight, dtype=float))
+    if prior_weight_arr.size == 1:
+        prior_weight_arr = np.full(n_states, prior_weight_arr.item())
+    if prior_weight_arr.shape != (n_states,):
+        raise ValueError(
+            f"prior_weight must be a scalar or array of shape ({n_states},), "
+            f"got shape {prior_weight_arr.shape}"
+        )
+    if np.any(prior_weight_arr < 0):
+        raise ValueError(
+            f"prior_weight must be non-negative, got {prior_weight_arr}"
+        )
 
     # p(x_t, x_{t+1} | O_{1:T})
     joint_distribution = estimate_joint_distribution(
@@ -389,12 +407,14 @@ def estimate_stationary_state_transition(
         acausal_posterior,
     )
 
-    n_states = acausal_posterior.shape[1]
     joint_sum = joint_distribution.sum(axis=0)  # (n_states, n_states)
 
     alpha = get_transition_prior(concentration, stickiness, n_states)
 
-    if prior_weight > 0.0:
+    # Legacy prior for rows with prior_weight[i] == 0
+    legacy_prior = alpha - 1.0  # (n_states, n_states)
+
+    if np.any(prior_weight_arr > 0):
         # Data-adaptive prior: scale pseudo-counts by expected transitions
         # so regularization strength is invariant to temporal resolution.
         alpha_shape = alpha - 1.0  # (n_states, n_states)
@@ -407,12 +427,21 @@ def estimate_stationary_state_transition(
         # N_i = expected transitions from state i
         N_i = joint_sum.sum(axis=-1, keepdims=True)  # (n_states, 1)
 
-        # alpha_eff - 1 = prior_weight * N_i * tilde_alpha
-        effective_prior = prior_weight * N_i * tilde_alpha
-        new_transition_matrix = joint_sum + effective_prior
+        # Per-row adaptive prior: prior_weight[i] * N_i * tilde_alpha[i]
+        adaptive_prior = prior_weight_arr[:, np.newaxis] * N_i * tilde_alpha
+
+        # Use adaptive prior for rows where prior_weight > 0, legacy otherwise
+        use_adaptive = prior_weight_arr > 0  # (n_states,)
+        effective_prior = np.where(
+            use_adaptive[:, np.newaxis],
+            adaptive_prior,
+            legacy_prior,
+        )
     else:
-        # Legacy fixed-count prior
-        new_transition_matrix = joint_sum + alpha - 1.0
+        # All rows use legacy fixed-count prior
+        effective_prior = legacy_prior
+
+    new_transition_matrix = joint_sum + effective_prior
 
     # Normalize rows to get transition probabilities.
     # Guard against all-zero rows (e.g., unvisited states) to avoid NaN.
@@ -607,7 +636,7 @@ def _estimate_discrete_transition(
     transition_concentration: float,
     transition_stickiness: float | np.ndarray,
     transition_regularization: float,
-    transition_prior_weight: float = 0.0,
+    transition_prior_weight: float | np.ndarray = 0.0,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Estimate the discrete transition matrix (stationary or non-stationary).
 

src/non_local_detector/models/base.py

@@ -206,7 +206,7 @@ def __init__(
         state_names: StateNames = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ) -> None:
         """
         Initialize the _DetectorBase class.
@@ -2099,7 +2099,7 @@ def __init__(
         state_names: StateNames = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ) -> None:
         """
         Initialize the ClusterlessDetector class.
@@ -2934,7 +2934,7 @@ def __init__(
         state_names: StateNames = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ) -> None:
         """
         Initialize the SortedSpikesDetector class.

src/non_local_detector/models/cont_frag_model.py

@@ -103,7 +103,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.cont_frag_defaults(),
@@ -250,7 +250,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.cont_frag_defaults(),

src/non_local_detector/models/decoder.py

@@ -85,7 +85,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.decoder_defaults(),
@@ -183,7 +183,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.decoder_defaults(),

src/non_local_detector/models/multienvironment_model.py

@@ -104,7 +104,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.multi_environment_defaults(),
@@ -220,7 +220,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.multi_environment_defaults(),

src/non_local_detector/models/non_local_model.py

@@ -118,7 +118,7 @@ def __init__(
         no_spike_rate: float = 1e-10,
         non_local_position_penalty: float = 0.0,
         non_local_penalty_sigma: float = 1.0,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.non_local_defaults(),
@@ -251,7 +251,7 @@ def __init__(
         no_spike_rate: float = 1e-10,
         non_local_position_penalty: float = 0.0,
         non_local_penalty_sigma: float = 1.0,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.non_local_defaults(),

src/non_local_detector/models/nospike_cont_frag_model.py

@@ -108,7 +108,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.nospike_cont_frag_defaults(),
@@ -229,7 +229,7 @@ def __init__(
         state_names: StateNames | None = None,
         sampling_frequency: float = 500.0,
         no_spike_rate: float = 1e-10,
-        discrete_transition_prior_weight: float = 0.0,
+        discrete_transition_prior_weight: float | np.ndarray = 0.0,
     ):
         params = _initialize_params(
             _ModelDefaults.nospike_cont_frag_defaults(),

src/non_local_detector/tests/test_discrete_transitions_estimation.py

@@ -429,6 +429,115 @@ def test_prior_weight_handles_unvisited_state(self):
             "Unvisited state should reflect sticky prior, not uniform"
         )
 
+    def test_per_state_prior_weight_scalar_equivalence(self, posterior_data):
+        """Passing scalar prior_weight should match passing uniform array."""
+        post = posterior_data
+        n_states = post["n_states"]
+
+        scalar_result = estimate_stationary_state_transition(
+            causal_posterior=post["causal_posterior"],
+            predictive_distribution=post["predictive_distribution"],
+            acausal_posterior=post["acausal_posterior"],
+            transition_matrix=post["transition_matrix"],
+            concentration=1.0,
+            stickiness=2.0,
+            prior_weight=0.1,
+        )
+        array_result = estimate_stationary_state_transition(
+            causal_posterior=post["causal_posterior"],
+            predictive_distribution=post["predictive_distribution"],
+            acausal_posterior=post["acausal_posterior"],
+            transition_matrix=post["transition_matrix"],
+            concentration=1.0,
+            stickiness=2.0,
+            prior_weight=np.full(n_states, 0.1),
+        )
+
+        np.testing.assert_allclose(scalar_result, array_result, atol=1e-12)
+
+    def test_per_state_prior_weight_mixed_modes(self, posterior_data):
+        """Rows with prior_weight=0 use legacy; rows with >0 use data-adaptive."""
+        post = posterior_data
+        n_states = post["n_states"]
+
+        # State 0 uses legacy path (strong fixed prior), others use adaptive
+        prior_weight = np.array([0.0] + [0.1] * (n_states - 1))
+        sticky = np.array([1e6] + [2.0] * (n_states - 1))
+
+        result = estimate_stationary_state_transition(
+            causal_posterior=post["causal_posterior"],
+            predictive_distribution=post["predictive_distribution"],
+            acausal_posterior=post["acausal_posterior"],
+            transition_matrix=post["transition_matrix"],
+            concentration=1.0,
+            stickiness=sticky,
+            prior_weight=prior_weight,
+        )
+
+        assert_stochastic_matrix(result)
+        # State 0 (frozen) should have near-1 diagonal due to huge stickiness
+        assert result[0, 0] > 0.99
+
+    def test_per_state_prior_weight_row_0_frozen_matches_legacy(self, posterior_data):
+        """For frozen row (prior_weight=0), result should match pure legacy call."""
+        post = posterior_data
+        n_states = post["n_states"]
+
+        sticky = np.array([1e6] + [2.0] * (n_states - 1))
+
+        legacy = estimate_stationary_state_transition(
+            causal_posterior=post["causal_posterior"],
+            predictive_distribution=post["predictive_distribution"],
+            acausal_posterior=post["acausal_posterior"],
+            transition_matrix=post["transition_matrix"],
+            concentration=1.0,
+            stickiness=sticky,
+            prior_weight=0.0,
+        )
+        mixed = estimate_stationary_state_transition(
+            causal_posterior=post["causal_posterior"],
+            predictive_distribution=post["predictive_distribution"],
+            acausal_posterior=post["acausal_posterior"],
+            transition_matrix=post["transition_matrix"],
+            concentration=1.0,
+            stickiness=sticky,
+            prior_weight=np.array([0.0] + [0.1] * (n_states - 1)),
+        )
+
+        # Row 0 (legacy) should match in both calls
+        np.testing.assert_allclose(mixed[0], legacy[0], atol=1e-12)
+
+    def test_per_state_prior_weight_negative_raises(self, posterior_data):
+        """Negative per-state prior_weight should raise ValueError."""
+        post = posterior_data
+        n_states = post["n_states"]
+
+        with pytest.raises(ValueError, match="non-negative"):
+            estimate_stationary_state_transition(
+                causal_posterior=post["causal_posterior"],
+                predictive_distribution=post["predictive_distribution"],
+                acausal_posterior=post["acausal_posterior"],
+                transition_matrix=post["transition_matrix"],
+                concentration=1.0,
+                stickiness=2.0,
+                prior_weight=np.array([0.1, -0.1] + [0.1] * (n_states - 2)),
+            )
+
+    def test_per_state_prior_weight_wrong_shape_raises(self, posterior_data):
+        """prior_weight array with wrong length should raise ValueError."""
+        post = posterior_data
+
+        with pytest.raises(ValueError, match="prior_weight"):
+            estimate_stationary_state_transition(
+                causal_posterior=post["causal_posterior"],
+                predictive_distribution=post["predictive_distribution"],
+                acausal_posterior=post["acausal_posterior"],
+                transition_matrix=post["transition_matrix"],
+                concentration=1.0,
+                stickiness=2.0,
+                prior_weight=np.array([0.1, 0.1]),  # wrong length
+            )
+
 
 @pytest.mark.unit
 class TestEstimateDiscreteTransition: