Support HaarReparam for non-compartmental variables #2523
Conversation
| @@ -868,3 +849,77 @@ def init(self, state): | |||
| def step(self, state): | |||
| with poutine.block(hide_types=["observe"]): | |||
| super().step(state.copy()) | |||
|
|
|||
|
|
|||
| class _HaarSplitReparam: | |||
There was a problem hiding this comment.
@fehiepsi I found this class provides useful helpers to convert between user <--> aux coordinates, in addition to our usual .reparam() functionality. These extra helpers seem necessary for poutine.reparam() to play well with init_to_value(). That is, if we reparameterize a model, then a user's custom init_to_value() function will need to be converted to different coordinates. WDYT of generalizing these to be part of the Reparam interface, say as a pair of methods (.aux_to_user(), .user_to_aux()) or similar? (in follow-up PRs to Pyro and NumPyro).
There was a problem hiding this comment.
These extra helpers seem necessary for poutine.reparam() to play well with init_to_value().
I see. I agree that those helper methods would simplify much of the user code for this init strategy.
| "HeterogeneousSIRModel", | ||
| "OverdispersedSEIRModel", | ||
| "OverdispersedSIRModel", | ||
| "RegionalSIRModel", | ||
| "SimpleSEIRModel", | ||
| "SimpleSIRModel", | ||
| "SparseSIRModel", | ||
| "SuperspreadingSEIRModel", | ||
| "SuperspreadingSIRModel", | ||
| "UnknownStartSIRModel", |
There was a problem hiding this comment.
Removing these because they are inteded more as examples than as reusable components.
| dist.Uniform(-0.5, self.population + 0.5) | ||
| .mask(False).expand(shape).to_event()) | ||
| assert auxiliary.shape == shape, "particle plates are not supported" | ||
| # Split tenors into current state. |
| # Transform to Haar coordinates. | ||
| config = {} | ||
| for name, dim in self.dims.items(): | ||
| config[name] = HaarReparam(dim=dim, flip=True) |
There was a problem hiding this comment.
make flip an arg?
There was a problem hiding this comment.
I think that's unnecessary since we don't expose any plumbing to set that, and since _HaarSplitReparam is merely a minimal internal helper.
|
Thanks for reviewing! |
Addresses #2426
This follows up #2517 to support
HaarReparamandSplitReparamfor non-compartmental latent time series. After this PRhaar_full_masswill trigger learning of covariance among the low frequency components of all time series, hopefully improving mixing in long-duration datasets.Summary:
haarandhaar_full_masslogic into a_HaarSplitReparamclass.._sample_auxiliary()method.HeterogeneousRegionalSIRModel(ported from Add epidemiology tutorial with a regional SEIR model #2518) to exercise new code.RetoRtwhich is more standard.pytest_covbreakage.From @martinjankowiak's comment in #2517 (review):
I hope the refactoring in this PR helps to reduce complexity.
Tested
HeterogeneousRegionalSIRModelwith test