Switching Linear Dynamical Systems support - #294
Conversation
|
@anushri10 I had codex update your #279 to re-synch with main; here is what it did. If you'd prefer going a bit more manually to update #279 , that is totally fine! Feel free to ignore this. BUT, if you find it helpful, feel free to work off of this branch instead---if you choose this route, do please read the core bits of code very carefully and think about how it aligns with the choices you previously made in #279. The AI is not yet perfect, and sometimes makes annoying choices.
|
DanWaxman
left a comment
DanWaxman
left a comment
There was a problem hiding this comment.
This is a great start! I didn't read carefully, but had a bunch of high-level comments below.
| ResamplingDifferentiableMethod = Literal["stop_gradient", "straight_through", "soft"] | ||
| FilterEmissionOrder = Literal["zeroth", "first", "second"] | ||
| FilterStateOrder = Literal["zeroth", "first", "second"] | ||
| RBPFProposal = Literal["prior", "optimal"] |
There was a problem hiding this comment.
Is this just a bootstrap proposal? Can we call it that if so? It's more in line with our other nomenclature.
| The filter samples the discrete regime path while analytically | ||
| marginalizing the conditionally linear-Gaussian continuous state with |
There was a problem hiding this comment.
Needs to specify what the discrete regime is
| `SwitchingLinearGaussianObservation`. | ||
|
|
||
| Does not support missing observations (data cannot contain NaNs). | ||
|
|
| """ | ||
|
|
||
| n_particles: int = 1_000 | ||
| proposal: RBPFProposal = "optimal" |
There was a problem hiding this comment.
I'm sure this has been thought about, but let's make sure this is a good default choice
| class RBPFPosterior(NamedTuple): | ||
| """Normalized dynestyx view of a cd-dynamax SLDS RBPF result. | ||
|
|
||
| Attributes: | ||
| marginal_loglik: Particle estimate of the total marginal log | ||
| likelihood. | ||
| weights: Normalized particle weights with shape | ||
| ``(time, num_particles)``. | ||
| regimes: Discrete regime particles with shape | ||
| ``(time, num_particles)``. | ||
| means: Conditional continuous-state means with shape | ||
| ``(time, num_particles, state_dim)``. | ||
| covariances: Conditional continuous-state covariances with shape | ||
| ``(time, num_particles, state_dim, state_dim)``. | ||
| filtered_means: Mixture mean of the continuous state. | ||
| filtered_covariances: Mixture covariance of the continuous state, | ||
| including both within-particle and between-particle uncertainty. | ||
| filtered_regime_probs: Filtered categorical probabilities over | ||
| regimes. | ||
| """ |
There was a problem hiding this comment.
Not sure this is the right place for this class, but I don't feel super confident about this
| from numpyro.distributions import constraints | ||
|
|
||
|
|
||
| class MixedStateDistribution(dist.Distribution): |
There was a problem hiding this comment.
I'd prefer a different name for this (to me "mixed distribution" probably evokes a measure like
| particle_key, state_key = jr.split(key) | ||
| particle_indices = dist.Categorical(logits=self.log_weights).sample( | ||
| particle_key, sample_shape | ||
| ) | ||
| component_samples = dist.MultivariateNormal( | ||
| self.continuous_locs, | ||
| covariance_matrix=self.continuous_covariances, | ||
| ).sample(state_key, sample_shape) | ||
| state_indices = jnp.broadcast_to( | ||
| particle_indices[..., None, None], | ||
| component_samples.shape[:-2] + (1, self.event_shape[0] - 1), | ||
| ) | ||
| states = jnp.take_along_axis(component_samples, state_indices, axis=-2)[ | ||
| ..., 0, : | ||
| ] | ||
| regimes = jnp.take_along_axis( | ||
| jnp.broadcast_to( | ||
| self.regimes, | ||
| sample_shape + self.regimes.shape, | ||
| ), | ||
| particle_indices[..., None], | ||
| axis=-1, | ||
| )[..., 0] | ||
| return jnp.concatenate( | ||
| (regimes[..., None].astype(states.dtype), states), | ||
| axis=-1, | ||
| ) |
There was a problem hiding this comment.
Haven't read carefully, but it seems surprising this is so complicated
| SwitchingLinearGaussianObservation, | ||
| SwitchingLinearGaussianStateEvolution, |
There was a problem hiding this comment.
Why is this necessary, don't they both inherit from numpyro.distributions.Distribution?
| def test_slds_rbpf_supports_shared_model_in_plate(): | ||
| obs_times, obs_values = _make_observations() | ||
| batched_observations = jnp.stack((obs_values, obs_values)) | ||
|
|
||
| def plated_model(): | ||
| with dsx.plate("trajectories", 2): | ||
| _slds_model(obs_times, batched_observations) | ||
|
|
||
| config = RBPFConfig( | ||
| n_particles=32, | ||
| proposal="optimal", | ||
| crn_seed=jr.PRNGKey(1), | ||
| ) | ||
| with trace() as tr, seed(rng_seed=jr.PRNGKey(2)), dsx.Filter(config): | ||
| plated_model() | ||
|
|
||
| marginal_loglik = tr["f_marginal_loglik"]["value"] | ||
| assert marginal_loglik.shape == (2,) | ||
| assert jnp.isfinite(marginal_loglik).all() | ||
|
|
There was a problem hiding this comment.
Could we also have a test with multiple levels of plating? This is where lots of our plating machinery is more fragile and breaks, IME.
| # Temporary commit pin while the SLDS RBPF marginal-likelihood change is | ||
| # awaiting an upstream cd-dynamax release. | ||
| "cd-dynamax @ git+https://github.com/anushri10/cd_dynamax.git@6295031bf19162735564c3b3c7cad697ea9176d2", |
There was a problem hiding this comment.
This is fine for testing, but let's remember to change this before pushing to main.
2 participants
Supersedes #279.
Currently, 1-shot un-edited codex from PR #279.