Connection NumPyro Forecasting

[juanitorduz]

Hi! I recently released a very lightweight JAX + NumPyro port of the Pyro Forecasting module, see https://github.com/juanitorduz/numpyro_forecast . My idea is to just do the port and present some examples. There is no intention to provide a user API for general forecasting models (like the Nixtla packages), but rather to provide a set of helper functions and a simple functional and OOP skeleton for implementing NumPyro Forecasting models.

I shared this new package in the NumPyro discourse: https://forum.pyro.ai/t/numpyro-forecast/9072 and I was refered to dynestyx :) I am not an expert in dynamical systems or in custom samplers; I am essentially a (Bayesian) forecasting practitioner. I wonder if dynestyx can also be used for forecasting models like the ones in https://juanitorduz.github.io/numpyro_forecast/docs/examples/, or if I could integrate' dynestyx' samplers into the NumPyro Forecast models. Any tips or hints are much appreciated. I am catching up with your tutorials and examples :)

[mattlevine22]

Hi @juanitorduz, thanks for looking into dynestyx! What you've been building in numpyro_forecast looks very cool, and very much in line with our goals of dynestyx (namely, to have a comprehensive Bayesian workflow + probabilistic evaluation framework specifically catered to timeseries/dynamics that can be paired with our modeling/inference toolkit).

A few thoughts / hints that I can think of:

• We try to adhere to "separation of concerns" as much as possible. To that end, our DynamicalModel has no notion of predict itself; instead it simply encodes the particulars of the recursive structure of a given dynamic model one is working with.
• Notions of prediction occur through interpretations of our dsx.sample(dynamical_model, ...) statement (e.g., wrapping the model in with Simulator(...):...dsx.sample(... predict_times=[...])

Some concrete directions:

1. Use numpyro_forecast to evaluate a dynestyx model

• Ultimately, I expect this will involve turning a code block like section 6.6 of this tutorial into a nf.ForecastModel class.

2. Use Dynestyx to fit your forecast models

• At a high-level, I'd expect many of your models to be definable as DynamicalModel classes (paired with a Simulator and/or filter, since your ForecastModels seem to require these sorts of methods within them). Once you do that, you could probably use your own SVI/MCMC fit pipeline or ours. Either way, the key benefit would be that your model would be running a non-linear filtering algorithm to connect the noisy data to the underlying model (rather than directly feeding noisy data into a predictive model).

I'd be more than happy to chat offline as well to discuss opportunities to collaborate.

Best,
Matt

[juanitorduz]

Thanks @mattlevine22 ! This is super insightful!

Let me spend a bit of time understanding the internals of dynestyx so I can fully grasp your comments above (they make a lot of sense, but I want to map them to code :) ). After that, having an offline chat would be super useful 💪

Best,
Juan