Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updated Likelihood Interface #29

Open
wants to merge 34 commits into
base: master
Choose a base branch
from
Open

Conversation

charlesknipp
Copy link
Collaborator

@charlesknipp charlesknipp commented Sep 9, 2024

This is a WIP PR which adds a suite of estimation tools for SSJ. While this doesn't implement anything on the block level (as of writing), this does introduce a number of conventions which may be useful for likelihood evaluation regardless of how it's structured. These features are listed below:

  • Probability distributions, each with a sampling method rand, a log likelihood method logpdf, and a method to check whether a given draw lies within the existing support in_support. This design may change, any feedback is appreciated (especially for the Prior class).
  • Shock objects (limited to AR, MA, and ARMA for now) and containers to build off of the dictionary design paradigm used throughout the module.

Lastly, I added a notebook which demonstrates these methods in action. Some of these methods are rather clunky (particularly when it comes to parameterization) but are completely generalizable and will at least get the ball rolling on future iterations.

@charlesknipp
Copy link
Collaborator Author

charlesknipp commented Sep 9, 2024

Credit to @bbardoczy for suggesting optimagic instead of scipy.optimize. This saved a ton of time allowing me to pass a dictionary as an argument to the objective function.

I still need to figure out whether the user should be allowed to pass a numpy array within a nested dictionary. For ARMA(p,q) shocks, their parameters phi and theta are arrays, so I'm not particularly sure how I want to handle this case. Any suggestions are welcome.

@charlesknipp
Copy link
Collaborator Author

I cleaned up the estimation notebook considerably. Right now, I limit the estimation to shock parameters to eliminate the need to recompute the Jacobian (this is a much bigger issue).

I also added a random walk MH algorithm at the bottom of the notebook. It seems finicky since the current log likelihood function references globals. This is definitely something I plan on addressing later this week.

@charlesknipp
Copy link
Collaborator Author

charlesknipp commented Sep 11, 2024

PyMC Integration

I added another notebook to try interfacing with pymc. It's not a very intuitive process, but can be disguised under the hood. There is also a ton of room for improvement in terms of optimization, which that requires integrating existing code with pytensor. So I'm not convinced this is the best path forward.

In my demo here the sampler takes about 8 minutes, and is significantly slower than the simple/stupid sampler in my other notebook. This is likely because of the "black box" inference of the modeling language.

Why Use a PPL?

The advantage to using a PPL (probablistic programming language) like pymc is clear when applying the same model to a number of different solvers. This also paves the way for using more robust and efficient inference algorithms like HMC.

I'm not too familiar with the Bayesian landscape in Python, so definitely let me know if there is anything more robust.

@charlesknipp
Copy link
Collaborator Author

@bbardoczy I figured out the problem with parameter inference in PyMC. Apparently each draw gets allocated to a zero dimensional numpy array as opposed to a scalar. This causes AccumulatedDerivative to error when multiplying, even though we are still technically multiplying a scalar.

I can't change PyMC's behavior since it's technically not doing anything wrong, and in most cases is helpful for sampling efficiency. The problem is, I don't want to change simple_displacement.py since it may override expected behavior. I may need a little help understanding exactly what this design choice was.

@bbardoczy
Copy link
Collaborator

I see. My guess is that we could safely extend all the scalar operations of the AccumulatedDerivative class to also work with zero dimensional numpy arrays. You could try that and see if it breaks any of tests. There are some direct test of AccumulatedDerivatce in test_displacement_handlers.py. You should also check that the notebooks still run properly. They all use SimpleBlocks and SolvedBlocks.

@charlesknipp
Copy link
Collaborator Author

What's New

More efficient means of automatically recomputing Jacobians between MCMC iterations. Blocks which vary with respect to a certain parameter will remain as their original blocks, whereas the rest (whose Jacobians stay constant between iterations) are replaced with a JacobianBlock.

While the constructor itself works in this demo, I have yet to unit test the updated estimation scheme.

@charlesknipp
Copy link
Collaborator Author

Interface Changes

  • no more pymc
  • added sampling interface with custom distribution module
  • model reduction scheme to intelligently re-evaluate Jacobians
  • added new shocks such as stacked shocks and news shocks
  • updated unit test to run both the manual likelihood calculation as well as the new interface

I still have a little house cleaning left (updating requirements.txt and removing unused code), but for now this is pretty stable.

@charlesknipp charlesknipp marked this pull request as ready for review January 24, 2025 22:25
@charlesknipp
Copy link
Collaborator Author

There is one issue with the addition of stacked shocks. __matmul__ doesn't know how to properly broadcast to a list of numpy arrays. I overcame this by introducing the functions stacked_responses and get_responses, which use check the type of the generated response before the matmul.

I think there are better ways around this, but nothing I could get working before the week's end. Even when the broadcast works, how should it be stored to preserve the array of intended operations?

It's not a huge issue since estimation works. Although I think there may be some complications for both solve_impulse functions

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants