Added tutorials on PDEs

[MatthieuDarcy]

This is a proof of concept on how to use dynestyx to solve inverse problems involving time dependent PDEs.

Discretization is done manually using a finite difference scheme, yielding an ODE for which the dynestyx machinery can be applied.

Limitations:
In the heat equation initial value inference problem, I wasn't able to use the usual dist.MultivariateNormal when using MCMC directly on the solver. Instead, I had to use a clunky dist.Normal(initial_condition(mu), ic_noise).to_event(1) combined with a reparametrization to ensure that MCMC worked at all. This is not in line with the usual dynestyx syntax and is much slower.

[DanWaxman]

Thanks Matthieu, this looks great!!

Discretization is done manually using a finite difference scheme, yielding an ODE for which the dynestyx machinery can be applied.

I don't want to put the cart before the horse -- but do you think we can package at least a few discretization schemes automatically (and/or, are there libraries we can import to do this)? It would be wonderful to have something like PDEDynamicalModel(...) and a handler PDEDiscretizer which transforms PDEDynamicalModel -> DynamicalModel. Not necessarily asking for that to be in this PR, but would be useful to get a wider vision & roadmap.

In the heat equation initial value inference problem, I wasn't able to use the usual dist.MultivariateNormal when using MCMC directly on the solver. Instead, I had to use a clunky dist.Normal(initial_condition(mu), ic_noise).to_event(1) combined with a reparametrization to ensure that MCMC worked at all. This is not in line with the usual dynestyx syntax and is much slower.

I'm curious why you say this is much slower (I'd expect it to be faster actually, though somewhat-negligibly-so). I think this is an interesting and valuable example where reparameterization is really necessary, though. But the corresponding loc-scale reparameterization is also possible for the MultivariateNormal -- it would look something like this

from numpyro.distributions.transforms import LowerCholeskyAffine
from numpyro.infer.reparam import TransformReparam

# define the initial condition a function of mu and x.
initial_condition = lambda mu : jnp.exp(- (x - mu)**2/(2*0.1**2)) * jnp.sin(jnp.pi*x)
ic_noise = 1e-2 # the noise on the initial condition
obs_noise = 1e-2 # the noise on the observations


def heat_equation_model(mu=None, obs_times=None, obs_values=None, predict_times=None):
    mu = numpyro.sample("mu", dist.Uniform(0, 1), obs=mu)

    # Create the dynamical model with sampled mu
    # Reparameterize the initial condition
    standard_x0 = dist.MultivariateNormal(
        loc=jnp.zeros(n_interior),
        scale_tril=jnp.eye(n_interior),
    )

    initial_distribution = dist.TransformedDistribution(
        standard_x0,
        LowerCholeskyAffine(
            loc=initial_condition(mu),
            scale_tril=ic_noise * jnp.eye(n_interior),
        ),
    )

    dynamics = DynamicalModel(
        initial_condition=initial_distribution,
        state_evolution=ContinuousTimeStateEvolution(
            drift=lambda x, u, t: A @ x,
        ),
        observation_model=LinearGaussianObservation(
            H=jnp.eye(n_interior), R=obs_noise**2 * jnp.eye(n_interior)
        ),
    )

    return dsx.sample("f", dynamics, obs_times=obs_times, obs_values=obs_values, predict_times=predict_times)

This gets pretty similar results (and similar timing, if not slower) on my machine to the existing results.

Some other comments:

• I'd like to avoid the HTML(FuncAnimation(...).to_jshtml()) pattern if possible -- it really bloats the filesize a ton (each of these notebooks is like 7MB), which is somewhat unsustainable. Not sure I have a great solution for this, to be fair.
• Super minor, but I'd prefer to keep the naming schemes in snake_case.
• Perhaps we should consider a bit coarser time discretization -- I think the default is like $\delta_t = 10^{-4}$, which may make these notebooks unnecessarily slow.

[MatthieuDarcy]

Re: adding discretization schemes. This would be very interesting and there are several candidate packages in the jax ecosystem we could consider.

I have now changed the reparametrization to use MultivariateNormal, which produces the same results as the previous but is more similar to other tutorials. This is slower than running the KF (which is what I meant originally, sorry if unclear) by a factor of 2x.

I have removed the HTML(FuncAnimation(...).to_jshtml()). It looks nice but is probably not worth the size.

Regarding naming schemes, did you mean the names of the notebooks themselves? If so then fixed.

Time discretization: I used the ImplicitEuler which requires an adaptive time step. Nonetheless, I reduced to dt = 1e-2 and reduced the tolerance so hopefully it runs a bit faster. Right now it takes 2-3 minutes to run each mcmc on my machine. If this is too slow, I will reduce the space discretization.

[DanWaxman]

Re: adding discretization schemes. This would be very interesting and there are several candidate packages in the jax ecosystem we could consider.

Cool! Looking forward to chatting more about this!

I have now changed the reparametrization to use MultivariateNormal, which produces the same results as the previous but is more similar to other tutorials. This is slower than running the KF (which is what I meant originally, sorry if unclear) by a factor of 2x.

Cool! And ah okay, yeah that makes sense.

I have removed the HTML(FuncAnimation(...).to_jshtml()). It looks nice but is probably not worth the size.

Thanks!

Regarding naming schemes, did you mean the names of the notebooks themselves? If so then fixed.

Yep!

Time discretization: I used the ImplicitEuler which requires an adaptive time step. Nonetheless, I reduced to dt = 1e-2 and reduced the tolerance so hopefully it runs a bit faster. Right now it takes 2-3 minutes to run each mcmc on my machine. If this is too slow, I will reduce the space discretization.

Well, this is true, but ContinuousTime(En)KF is also integrating and will be using a default solver. (I wouldn't be surprised if this is also a big part of the reason why the filtering path is a lot faster). I think 2-3 minutes is fine! I just find people more likely to play around with tutorials if they don't take super long to run (which I think 2-3 minutes falls into squarely), so prefer to reduce as much as practical such that the qualitative takeaway is the same.

I think this actually points to some places we should improve cd_dynamax/our integration, which hasn't really been built around ODEs at all. For example, the solver is non-configurable for the Kalman filter (which, tbh, feels more just like a bug on our end). The EnKF, on the other hand, will always solve as an SDE, which is rather wasteful here. More stuff to talk about when we meet, I suppose :)

[mattlevine22]

Interesting, great work!

• great to see that ODESimulator and Filtering (linear and non-linear) both work well for the two problems you set up!

• pretty cool that KF works with 0 state-noise; I might have stressed about the numerics, but glad to see it survived!

• I forget what choice EnKF makes for default state-noise covariances if not specified, and it was hard for me to figure this out quickly from our docs; I created an issue about this Clearer documentation for assumed diffusion in filtering algorithms #277.

• agreed that a big picture look at nice Discretizer interfaces seems worth it!

• also agreed that CD-KF filter should expose solver settings! I created an issue about this Continuous-Discrete Kalman Filter config should expose solver settings #278.

• I'd probably lean towards specialized state-evolution classes rather than a PDEDynamicalModel class since none of our other model classes are distinguished at such a high level.

Overall this seems great to me---I'd be down to just add these directly to our Deep Dives + Examples and merge this PR; then do further iterations after our discussion next week. @DanWaxman what say you?

[DanWaxman]

I’d prefer the solver thing gets fixed first, or at least commented on in the notebook, I think gives potentially the wrong conclusion (or maybe the right conclusion for the wrong reason). Other than that I’m happy… it becomes another thing we need to update in #255, but at this point, c’est la vie.

[mattlevine22]

I'd prefer to fast-track this PR and add a short comment in the notebooks explaining that the FilterConfig has a solver associated with it (and note that that the ContinuousKF has a fixed solver that cannot be changed, pending #278).

• (my concern is unnecessary blockage / slowdowns to/from Staging Branch for Numpyro-Free Usage #255). With the small comment, it keeps us moving, while not impeding other work. We should just come back and refine the solvers when we are able to prioritize that.

For the ContinuousTimeEnKF, @MatthieuDarcy you could specify the solver there if you want. Indeed, it is unfortunate that CD-Dynamax is solving an ODE (zero diffusion) using an SDE solver (it is expecting the general case). When #226 lands, these should be mostly fixed (would be nice for that PR to discretize ODEs as mappings of diracs by the way).

[MatthieuDarcy]

Thanks for the informative discussion, I wasn't aware of some of these subtleties.

I added a note on the use of the ContinuousTimeEnKFshowing how one can modify the solver used. UsingImplicitEuler is indeed much slower (~6 minutes on my machine vs ~2 minutes using default settings), so I left a commented block with default parameters and saying that one can run this instead as it yields very similar results in practice. Again, I can reduce spatial discretization if we want a faster tutorial.

The notebook on the linear heat equation is unchanged.

[mattlevine22]

1. For the heat equation example, I don't really understand what is going on with the initial condition transformation (and why that is needed / useful here). Relatedly, it would be better (if possible) not to re-simulate data before running the Filter here. In fact, if the model is really just an LTI system, might be better/easier to define it that way from the beginning.

2. Can you please add a similar solver note to the heat equation example (just say what default solver is being used, and that it currently cannot change it, but soon you will be able to)?

[MatthieuDarcy]

Regarding the reparametrization, I'm no expert on MCMC but my understanding is that this is a common issue in hierarchical Bayesian models. My understanding is that here $$u_0 \sim \mathcal{N}(\mu(\theta), I)$$ where $$u_0 \in \mathbb{R}^{62}$$. MCMC treat entries of this component as independent at the proposal stage and therefore gets stuck. Reparametrization is just doing $$u_0 = \mu(\theta) + \varepsilon$$ which somehow fixes this. I'm not competent enough to fully explain it, but it is necessary as otherwise MCMC will fail completely. Interestingly this is not an issue for the KF because it marginalizes over the latent state $$u_0$$ (might be worth pointing out).

Regarding defining things using LTI since the beginning, my issue was that LTI implicitly assumes everything is an SDE, hence you cannot run ImplicitEuler. Later on, I have to define LTI because KF assumes that the dynamics are linear SDEs. I have now written a small paragraph pointing this out as well as the fact that the KF solver cannot be customized. However I agree that regenerating the data is not necessary and I have removed that part.

[mattlevine22]

Ahhh, I think I understand, thanks for the explanation!

You're saying that MCMC over p(mu, u_0 | Y) is a funnel with bad geometry, whereas (mu, eps_0 | Y) has better geometry?

• Yes, can you add a note about this implementation challenge in joint-MCMC that goes away when using filter-marginalized MCMC?

Otherwise, I think I'm happy with it for now. @DanWaxman are you ok with merging this?

[DanWaxman]

You're saying that MCMC over p(mu, u_0 | Y) is a funnel with bad geometry, whereas (mu, eps_0 | Y) has better geometry?

Yeah. The MCMC gets stuck near $\mu = 0$ otherwise.

This may also pop up other places -- it could be worth it to make a little helper for what I suggested above (or try to contribute some multivariate normal reparameterization to numpyro) to make dynestyx users' lives easier.

Otherwise, I think I'm happy with it for now. @DanWaxman are you ok with merging this?

In principle, yes -- though these tutorials need to be linked in the mkdocs.yaml file + deep dives landing page. It should probably also go in deep_dives, though I'm not too pressed about organizing to that level (we'll probably re-organize in the coming weeks anyways).

[mattlevine22]

This may also pop up other places -- it could be worth it to make a little helper for what I suggested above (or try to contribute some multivariate normal reparameterization to numpyro) to make dynestyx users' lives easier.

Mind creating a short issue about this?

[MatthieuDarcy]

Added a comment explaining this and highlighting that this is not an issue when running the KF.

[DanWaxman]

I moved things to deep_dives/, fixed some capitalization, and added a hyperlink for when the heat equation notebook is referenced. Otherwise, looks good to me, thanks again @MatthieuDarcy ! Looking forward to chatting more about this.