How to`fit_constrained` and handle missing values

[adoherty21]

Apologies if this is not the right place for this question.

I am evaluating dynamax for implementing simple, fast and efficient structural time series models for imputation/univariate forecasting and dynamic regression. I am impressed with the API, however, there are a couple of blockers for me using the library. I was wondering you would be able to provide some insight into:

1. How to constrain parameters before using fit_sgd. I have looked into the source code and have seen the from_unconstrained function, however, I am unsure how to use this to constrain a single parameter before fitting.
2. When there are missing values the inference/fitting doesn't work. Is handling missing data an improvement that will be considered?

The notebook I am using to investigate is here.

I would be happy to contribute if there is anything I can help with.

[gileshd]

Hi @adoherty21, thanks for the kind words.

If you are interesting in structural time series models you might be interested in sts-jax a sister library built on top of dynamax for precisely this use case!

In response to your questions:

1. A bijector which maps to the constrained form of your parameters can be specified in ParameterProperties. For instance in your notebook this is generated by lgssm.initialize and then passed as an argument into fit_sgd.

You can add a tfp.Bijector to a particular parameter in this props object in order to constrain this parameter when using fit_sgd. Let me know if that makes sense.

2. I am not sure about the roadmap for supporting missing values - there is some discussion of this and related issue in #99. Contributions on this front are welcome!

[adoherty21]

Thanks a lot @gileshd.

That makes sense for constraining.

Regarding missing values I will start to think about the implementation and make time to contribute.

[adoherty21]

Closed.