nstat.extras.em.dynamax_bridge

EM-trained linear-Gaussian state-space models via Dynamax.

This module wraps dynamax.linear_gaussian_ssm.LinearGaussianSSM behind a thin nstat-style API so users can fit EM-trained state-space models without nstat owning the EM code itself.

Scope

All fit/inference routines return plain NumPy arrays (NOT pytrees) so callers stay decoupled from JAX.

Install

pip install nstat-toolbox[dynamax]

Pulls Dynamax (~50 MB) + JAX (~200 MB).

Functions

cmgf_poisson_filter(observations, ...)

Run a Conditional-Moments Gaussian Filter for Poisson observations.

cmgf_poisson_smoother(observations, ...)

Run a Conditional-Moments Gaussian Smoother for Poisson observations.

fit_hybrid_em(poisson_observations, ...[, ...])

Fit a hybrid Poisson + Gaussian SSM via EM (mPPCO_EM equivalent).

fit_hybrid_em_best_of(poisson_observations, ...)

Hybrid (Poisson + Gaussian) multi-restart EM with predictive-LL selection.

fit_linear_gaussian_em(observations, ...[, ...])

Fit a linear-Gaussian state-space model via Dynamax EM.

fit_point_process_em(observations, state_dim, *)

Fit a Poisson-LGSSM via Laplace-approximated EM (PP_EM equivalent).

fit_point_process_em_best_of(observations, ...)

Run fit_point_process_em() with several seeds and pick the best.

hybrid_predictive_ll(poisson_observations, ...)

True one-step-ahead predictive log-likelihood of a hybrid SSM.

point_process_predictive_ll(observations, ...)

True one-step-ahead predictive log-likelihood of a Poisson-LGSSM.

Classes

CMGFPoissonFilterResult(state_means, ...)

Output of a CMGF Poisson filter or smoother.

HybridEMResult(transition_matrix, ...)

Output of a Poisson + Gaussian SSM EM fit (mPPCO_EM equivalent).

LinearGaussianEMResult(transition_matrix, ...)

Output of a linear-Gaussian SSM EM fit.

MultiRestartResult(best_result, best_seed, ...)

Best-of-N restart EM result + per-restart diagnostics.

PointProcessEMResult(transition_matrix, ...)

Output of a Poisson-LGSSM EM fit (PP_EM equivalent).

PredictiveLogLik(total, per_timestep, poisson)

True one-step-ahead predictive log-likelihood of a state-space fit.
class nstat.extras.em.dynamax_bridge.LinearGaussianEMResult(transition_matrix, observation_matrix, transition_covariance, observation_covariance, initial_state_mean, initial_state_covariance, log_likelihoods, n_iter)[source]

Bases: object

Output of a linear-Gaussian SSM EM fit.

Variables:

  • transition_matrix (numpy.ndarray) – Learned \(A\), shape (state_dim, state_dim).

  • observation_matrix (numpy.ndarray) – Learned \(C\), shape (emission_dim, state_dim).

  • transition_covariance (numpy.ndarray) – Learned \(Q\), shape (state_dim, state_dim).

  • observation_covariance (numpy.ndarray) – Learned \(R\), shape (emission_dim, emission_dim).

  • initial_state_mean (numpy.ndarray) – Learned \(\hat x_0\), shape (state_dim,).

  • initial_state_covariance (numpy.ndarray) – Learned \(P_0\), shape (state_dim, state_dim).

  • log_likelihoods (numpy.ndarray) – EM log-likelihood trace, shape (n_iter,). Last value is the final log-likelihood at convergence.

  • n_iter (int) – Number of EM iterations actually run.

Parameters:
transition_matrix: ndarray
observation_matrix: ndarray
transition_covariance: ndarray
observation_covariance: ndarray
initial_state_mean: ndarray
initial_state_covariance: ndarray
log_likelihoods: ndarray
n_iter: int
nstat.extras.em.dynamax_bridge.fit_linear_gaussian_em(observations, state_dim, *, n_iter=50, seed=0)[source]

Fit a linear-Gaussian state-space model via Dynamax EM.

Counterpart to MATLAB nSTAT’s KF_EM family.

Parameters:

  • observations (ndarray) – Observation time series, shape (T, emission_dim).

  • state_dim (int) – Dimensionality of the latent state.

  • n_iter (int) – Number of EM iterations. Default 50.

  • seed (int) – Seed for JAX PRNG (used for parameter initialization).

Return type:

LinearGaussianEMResult

Notes

Dynamax internally represents parameters as nested JAX pytrees; this wrapper unpacks them into plain NumPy arrays so callers don’t need to know about pytrees or JAX.

class nstat.extras.em.dynamax_bridge.CMGFPoissonFilterResult(state_means, state_covariances, marginal_log_likelihood)[source]

Bases: object

Output of a CMGF Poisson filter or smoother.

Variables:

  • state_means (numpy.ndarray) – Posterior state means, shape (T, state_dim). These are filtered means when produced by cmgf_poisson_filter(), smoothed (forward + backward pass) when produced by cmgf_poisson_smoother().

  • state_covariances (numpy.ndarray) – Posterior state covariances, shape (T, state_dim, state_dim).

  • marginal_log_likelihood (float) – Sum of log p(y_t | y_{1:t-1}) (filter) or log p(y_{1:T}) (smoother) under the CMGF Gaussian approximation. Useful for model-comparison-style scalar diagnostics.

Parameters:
state_means: ndarray
state_covariances: ndarray
marginal_log_likelihood: float
nstat.extras.em.dynamax_bridge.cmgf_poisson_filter(observations, transition_matrix, observation_matrix, transition_covariance, initial_state_mean, initial_state_covariance)[source]

Run a Conditional-Moments Gaussian Filter for Poisson observations.

Counterpart to MATLAB nSTAT’s PPDecodeFilter / PP_filter families — point-process Kalman-like filtering for spike-count observations on a continuous linear-Gaussian latent state.

The model is:

\[\begin{split}x_t &= A x_{t-1} + w_t, \quad w_t \sim \mathcal{N}(0, Q) \\ y_t &\sim \text{Poisson}(\exp(C x_t))\end{split}\]
Parameters:

  • observations (ndarray) – Spike-count observations, shape (T, emission_dim). Must be non-negative integers (or float castable to non-negative int).

  • transition_matrix (ndarray) – \(A\), shape (state_dim, state_dim).

  • observation_matrix (ndarray) – \(C\), shape (emission_dim, state_dim).

  • transition_covariance (ndarray) – \(Q\), shape (state_dim, state_dim).

  • initial_state_mean (ndarray) – \(\hat{x}_0\), shape (state_dim,).

  • initial_state_covariance (ndarray) – \(P_0\), shape (state_dim, state_dim).

Return type:

CMGFPoissonFilterResult

Notes

Internally uses Dynamax’s EKF integration scheme. For UKF / GHKF variants drop down to dynamax.generalized_gaussian_ssm.inference directly.

nstat.extras.em.dynamax_bridge.cmgf_poisson_smoother(observations, transition_matrix, observation_matrix, transition_covariance, initial_state_mean, initial_state_covariance)[source]

Run a Conditional-Moments Gaussian Smoother for Poisson observations.

Forward CMGF filter + backward RTS-style smoother under the Gaussian approximation to the Poisson likelihood. Counterpart to MATLAB nSTAT’s PP_fixedIntervalSmoother.

See cmgf_poisson_filter() for the model specification and parameter shapes — this function is identical in signature, just returns smoothed (rather than filtered) state means + covariances.

Returns:

state_means and state_covariances are the smoothed posteriors \(p(x_t | y_{1:T})\).

Return type:

CMGFPoissonFilterResult

Parameters:
class nstat.extras.em.dynamax_bridge.PointProcessEMResult(transition_matrix, observation_matrix, transition_covariance, initial_state_mean, initial_state_covariance, marginal_log_likelihoods, n_iter)[source]

Bases: object

Output of a Poisson-LGSSM EM fit (PP_EM equivalent).

Variables:

  • transition_matrix (numpy.ndarray) – Learned \(A\), shape (state_dim, state_dim).

  • observation_matrix (numpy.ndarray) – Learned \(C\), shape (emission_dim, state_dim).

  • transition_covariance (numpy.ndarray) – Learned \(Q\), shape (state_dim, state_dim).

  • initial_state_mean (numpy.ndarray) – Learned \(\hat x_0\), shape (state_dim,).

  • initial_state_covariance (numpy.ndarray) – Learned \(P_0\), shape (state_dim, state_dim).

  • marginal_log_likelihoods (numpy.ndarray) – EM marginal log-likelihood trace under the CMGF Gaussian approximation, shape (n_iter,).

  • n_iter (int) – Number of EM iterations actually run.

Parameters:
transition_matrix: ndarray
observation_matrix: ndarray
transition_covariance: ndarray
initial_state_mean: ndarray
initial_state_covariance: ndarray
marginal_log_likelihoods: ndarray
n_iter: int
nstat.extras.em.dynamax_bridge.fit_point_process_em(observations, state_dim, *, n_iter=30, n_newton_iter=5, seed=0, init='random', ridge_lambda=0.0)[source]

Fit a Poisson-LGSSM via Laplace-approximated EM (PP_EM equivalent).

Counterpart to MATLAB nSTAT’s PP_EM family. Model:

\[\begin{split}x_t &= A x_{t-1} + w_t, \quad w_t \sim \mathcal{N}(0, Q) \\ y_t &\sim \text{Poisson}(\exp(C x_t))\end{split}\]

Algorithm (Smith & Brown 2003 PPLDS):

  1. Initialize A, C, Q, x0, P0.

  2. Loop n_iter times:

    • E-step: Dynamax CMGF smoother (EKF integration) on the Poisson-LGSSM with current parameters. Returns smoothed means \(\mu_t\), covariances \(\Sigma_t\), and the marginal log-likelihood under the Gaussian approximation.

    • M-step (dynamics): closed-form A, Q, x0, P0 from smoothed moments.

    • M-step (loadings): per-row Newton-Raphson on the expected Poisson log-likelihood with Laplace approximation for \(\mathbb{E}[\exp(c_i' x_t)]\).

Parameters:

  • observations (ndarray) – Spike-count time series, shape (T, emission_dim), integer- valued (or float castable to non-negative int).

  • state_dim (int) – Latent state dimensionality.

  • n_iter (int) – Outer EM iterations. Default 30.

  • n_newton_iter (int) – Inner Newton-Raphson iterations per M-step for C. Default 5.

  • seed (int) – RNG seed for parameter initialization.

  • init (str)

  • ridge_lambda (float)

Return type:

PointProcessEMResult

Warning

Experimental. A Poisson LDS has a gauge freedom \((A, C, x) \to (T A T^{-1}, C T^{-1}, T x)\) that leaves the observable log-rate \(C x\) (hence the likelihood) invariant for any invertible \(T\) (the full \(GL(d)\) group). This implementation pins that gauge to a canonical form — a cheap diagonal scale pin during the loop, then a single whiten + SVD-rotate + sign-fix after convergence (see _canonicalize_gauge()) — so the returned C satisfies \(C^\top C = \mathrm{diag}(S^2)\) and is a unique, seed-stable representative. What remains is local-optima multiplicity: EM may converge to genuinely different likelihoods across seeds, so prefer the predictions (the fitted log-rate / firing rate exp(C x)) over a single fit’s raw A / C. Bit-exact parity with MATLAB PP_EM would additionally require multi-restart model selection.

Notes

The marginal-log-likelihood trace is the time-varying-R Gaussian smoother surrogate, not the true Poisson marginal likelihood; it is not guaranteed monotonic. The E-step uses an IRLS pseudo-observation linearization + a time-varying-R RTS smoother that exposes the lag-one cross-covariances the A/Q M-step needs.

class nstat.extras.em.dynamax_bridge.HybridEMResult(transition_matrix, poisson_observation_matrix, gaussian_observation_matrix, transition_covariance, gaussian_observation_covariance, initial_state_mean, initial_state_covariance, marginal_log_likelihoods, n_iter)[source]

Bases: object

Output of a Poisson + Gaussian SSM EM fit (mPPCO_EM equivalent).

A single latent state drives both observation channels.

Variables:

  • transition_matrix (numpy.ndarray) – Learned \(A\).

  • poisson_observation_matrix (numpy.ndarray) – Learned \(C_p\), shape (poisson_emission_dim, state_dim).

  • gaussian_observation_matrix (numpy.ndarray) – Learned \(C_g\), shape (gaussian_emission_dim, state_dim).

  • transition_covariance (numpy.ndarray) – Learned \(Q\).

  • gaussian_observation_covariance (numpy.ndarray) – Learned \(R\).

  • initial_state_mean (numpy.ndarray) – \(\hat x_0\).

  • initial_state_covariance (numpy.ndarray) – \(P_0\).

  • marginal_log_likelihoods (numpy.ndarray) – EM trace under the Laplace-Gaussian approximation.

  • n_iter (int) – Iterations actually run.

Parameters:
transition_matrix: ndarray
poisson_observation_matrix: ndarray
gaussian_observation_matrix: ndarray
transition_covariance: ndarray
gaussian_observation_covariance: ndarray
initial_state_mean: ndarray
initial_state_covariance: ndarray
marginal_log_likelihoods: ndarray
n_iter: int
nstat.extras.em.dynamax_bridge.fit_hybrid_em(poisson_observations, gaussian_observations, state_dim, *, n_iter=30, n_newton_iter=3, seed=0, ridge_lambda=0.0)[source]

Fit a hybrid Poisson + Gaussian SSM via EM (mPPCO_EM equivalent).

Counterpart to MATLAB nSTAT’s mPPCO_EM family. Both observation channels share a single linear-Gaussian latent state:

\[\begin{split}x_t &= A x_{t-1} + w_t, \quad w_t \sim \mathcal{N}(0, Q) \\ y^{(p)}_t &\sim \text{Poisson}(\exp(C_p x_t)) \\ y^{(g)}_t &\sim \mathcal{N}(C_g x_t, R)\end{split}\]

Algorithm:

  1. Initialize all parameters.

  2. Loop n_iter times:

    • E-step: Laplace-approximate the Poisson channel into IRLS Gaussian pseudo-observations linearized around the previous smoothed means. Stack with the real Gaussian channel. Run dynamax.linear_gaussian_ssm.lgssm_smoother() on the augmented LG-SSM. The cross-channel information sharing that’s the whole point of mPPCO_EM happens inside the smoother — the Gaussian channel constrains the latent state, which tightens the pseudo-obs for the Poisson channel on the next iteration.

    • M-step (dynamics): closed-form A, Q, x0, P0 from smoothed moments (shared with PP_EM).

    • M-step (Poisson loadings): per-row Newton-Raphson for C_p (shared with PP_EM).

    • M-step (Gaussian loadings): closed-form least squares for C_g; closed-form residual cov for R (with the latent-uncertainty trace correction).

    • Gauge pin (in-loop): a cheap diagonal scale normalization keeps |C_p| / |C_g| finite without reshaping the optimization landscape each step.

  3. After convergence: pin the full PLDS gauge once (whiten + SVD-rotate + sign-fix on the stacked [C_p; C_g]) so both emission channels share a unique, seed-stable latent frame (see _canonicalize_gauge()).

Parameters:

  • poisson_observations (ndarray) – Spike-count series, shape (T, p_emission_dim).

  • gaussian_observations (ndarray) – Continuous-valued series, shape (T, g_emission_dim).

  • state_dim (int) – Latent state dimensionality.

  • n_iter (int) – Outer EM iterations.

  • n_newton_iter (int) – Inner Newton-Raphson iterations for C_p per M-step.

  • seed (int) – RNG seed for initialization.

  • ridge_lambda (float)

Return type:

HybridEMResult

Warning

Experimental. Like fit_point_process_em(), the shared latent is gauge-free up to \(GL(d)\); this fit pins it to a canonical form (computed from the stacked [C_p; C_g]), so the returned loadings are seed-stable up to local-optima multiplicity. The identifiable, observation-space outputs — the Gaussian noise R and the fitted Poisson rates — are the most reliable; treat a single fit’s raw C_p / C_g with the same care.

Notes

The E-step uses a fixed-R approximation: the IRLS pseudo-obs variance 1/λ_t varies in t, but Dynamax’s lgssm_smoother accepts only a fixed emission covariance. We use the time-averaged 1/mean(λ) per channel as the substitute. Real- data fits with sharply non-stationary Poisson rates may want full time-varying R; that requires a custom Kalman smoother (deferred to a future release).

class nstat.extras.em.dynamax_bridge.PredictiveLogLik(total, per_timestep, poisson, gaussian=None)[source]

Bases: object

True one-step-ahead predictive log-likelihood of a state-space fit.

Unlike the marginal_log_likelihoods trace reported by the EM trainers — a Gaussian-smoother surrogate that re-linearizes each iteration and is not a valid objective — this is the genuine predictive log-likelihood of the observations under the fitted parameters. At each step the latent state is predicted causally from the past only (one-step-ahead), and the true Poisson (and, for the hybrid, Gaussian) likelihood of the actual observation is scored under that predictive distribution.

Use it to (a) confirm EM actually improved the fit, (b) compare models / EM restarts on equal footing, and (c) score held-out data (pass a test segment together with the train-fitted parameters).

Variables:

  • total (float) – Summed predictive log-likelihood over all time steps and channels (nats). Higher is better.

  • per_timestep (numpy.ndarray) – Per-time-step predictive log-likelihood, shape (T,) (summed across neurons / channels).

  • poisson (float) – Poisson-channel contribution to total.

  • gaussian (float | None) – Gaussian-channel contribution (None for the point-process case).

Parameters:
total: float
per_timestep: ndarray
poisson: float
gaussian: float | None = None
nstat.extras.em.dynamax_bridge.point_process_predictive_ll(observations, transition_matrix, observation_matrix, transition_covariance, initial_state_mean, initial_state_covariance, *, n_quad=15)[source]

True one-step-ahead predictive log-likelihood of a Poisson-LGSSM.

The honest convergence / quality diagnostic for fit_point_process_em() — replaces the surrogate Gaussian-smoother trace. Runs a causal forward filter and scores the actual spike counts under the one-step-ahead predictive state via Gauss-Hermite quadrature of the Poisson likelihood (integrating over the latent uncertainty, not just plugging in the mean rate).

Pure NumPy — does not require dynamax. Pass a held-out segment plus the train-fitted parameters for a proper held-out score.

Parameters:

  • observations (ndarray) – Spike counts, shape (T, emission_dim) (1-D is reshaped).

  • transition_matrix (ndarray) – Fitted \(A\), \(C\), \(Q\).

  • observation_matrix (ndarray) – Fitted \(A\), \(C\), \(Q\).

  • transition_covariance (ndarray) – Fitted \(A\), \(C\), \(Q\).

  • initial_state_mean (ndarray) – Fitted \(\hat x_0\), \(P_0\).

  • initial_state_covariance (ndarray) – Fitted \(\hat x_0\), \(P_0\).

  • n_quad (int) – Gauss-Hermite nodes for the Poisson marginal (default 15; increase for sharper accuracy at high firing rates).

Return type:

PredictiveLogLik

nstat.extras.em.dynamax_bridge.hybrid_predictive_ll(poisson_observations, gaussian_observations, transition_matrix, poisson_observation_matrix, gaussian_observation_matrix, transition_covariance, gaussian_observation_covariance, initial_state_mean, initial_state_covariance, *, n_quad=15)[source]

True one-step-ahead predictive log-likelihood of a hybrid SSM.

The fit_hybrid_em() counterpart of point_process_predictive_ll(). Both channels share the causal forward filter; the Poisson channel is scored by Gauss-Hermite quadrature, the Gaussian channel by its exact multivariate-normal predictive density. total = poisson + gaussian.

Pure NumPy — does not require dynamax.

Parameters:

  • poisson_observations (ndarray) – Spike counts (T, p_dim) and continuous signal (T, g_dim).

  • gaussian_observations (ndarray) – Spike counts (T, p_dim) and continuous signal (T, g_dim).

  • transition_matrix (ndarray) – Fitted \(A\), \(C_p\), \(C_g\), \(Q\), \(R\), \(\hat x_0\), \(P_0\).

  • poisson_observation_matrix (ndarray) – Fitted \(A\), \(C_p\), \(C_g\), \(Q\), \(R\), \(\hat x_0\), \(P_0\).

  • gaussian_observation_matrix (ndarray) – Fitted \(A\), \(C_p\), \(C_g\), \(Q\), \(R\), \(\hat x_0\), \(P_0\).

  • transition_covariance (ndarray) – Fitted \(A\), \(C_p\), \(C_g\), \(Q\), \(R\), \(\hat x_0\), \(P_0\).

  • gaussian_observation_covariance (ndarray) – Fitted \(A\), \(C_p\), \(C_g\), \(Q\), \(R\), \(\hat x_0\), \(P_0\).

  • initial_state_mean (ndarray) – Fitted \(A\), \(C_p\), \(C_g\), \(Q\), \(R\), \(\hat x_0\), \(P_0\).

  • initial_state_covariance (ndarray) – Fitted \(A\), \(C_p\), \(C_g\), \(Q\), \(R\), \(\hat x_0\), \(P_0\).

  • n_quad (int) – Gauss-Hermite nodes for the Poisson marginal (default 15).

Return type:

PredictiveLogLik

class nstat.extras.em.dynamax_bridge.MultiRestartResult(best_result, best_seed, best_predictive_ll, all_seeds, all_predictive_lls)[source]

Bases: object

Best-of-N restart EM result + per-restart diagnostics.

Tier 0.1 (canonical gauge) made each fit a unique seed-stable representative of its equivalence class; Tier 0.2 added a true held-out predictive log-likelihood diagnostic. This dataclass is the natural composition of the two: run EM with several seeds, score each on held-out data, return the best. It is also the mitigation for the weak-observability A → 0 collapse exposed by Tier 0.2: even when some seeds collapse, the predictive-LL selector picks one that didn’t.

Variables:

  • best_result (object) – PointProcessEMResult or HybridEMResult from the seed with the highest held-out predictive log-likelihood.

  • best_seed (int) – Seed that produced best_result.

  • best_predictive_ll (float) – Held-out predictive log-likelihood total for best_result.

  • all_seeds (numpy.ndarray) – Seeds used, in order, shape (n_restarts,).

  • all_predictive_lls (numpy.ndarray) – Held-out predictive log-likelihood for each restart, same order as all_seeds. NaN if a restart raised an exception or produced non-finite parameters.

Parameters:
best_result: object
best_seed: int
best_predictive_ll: float
all_seeds: ndarray
all_predictive_lls: ndarray
nstat.extras.em.dynamax_bridge.fit_point_process_em_best_of(observations, state_dim, *, n_restarts=8, holdout_fraction=0.2, n_iter=30, n_newton_iter=5, base_seed=0, n_quad=15, init='random', ridge_lambda=0.0)[source]

Run fit_point_process_em() with several seeds and pick the best.

Splits observations into a training segment (first 1 - holdout_fraction of time) and a test segment (the remainder); fits PP_EM on the train segment for each of n_restarts random seeds (base_seed … base_seed + n_restarts - 1); scores each fit on the test segment with point_process_predictive_ll(); returns the best.

This is the recommended workflow for PP_EM on real data. Single-fit PP_EM under weak observability (few neurons, small loadings) can converge to a degenerate solution (A → 0, inflated C/Q) whose held-out predictive log-likelihood is worse than a constant-rate baseline — see the observability caveat in fit_point_process_em() and docs/extras/em_dynamax.md. Multi-restart selection on held-out predictive LL is the production-quality mitigation.

Parameters:
Return type:

MultiRestartResult

nstat.extras.em.dynamax_bridge.fit_hybrid_em_best_of(poisson_observations, gaussian_observations, state_dim, *, n_restarts=8, holdout_fraction=0.2, n_iter=30, n_newton_iter=3, base_seed=0, n_quad=15, ridge_lambda=0.0)[source]

Hybrid (Poisson + Gaussian) multi-restart EM with predictive-LL selection.

The fit_hybrid_em() counterpart of fit_point_process_em_best_of(). Both observation channels are split by the same trailing-time holdout; scoring uses hybrid_predictive_ll() (Poisson + Gaussian total).

Parameters:
Return type:

MultiRestartResult