search

exoplanet-dev / exoplanet

Fast & scalable MCMC for all your exoplanet needs!
https://docs.exoplanet.codes
MIT License
206 stars 52 forks source link

Implementing support for pymc v4 #271

Closed dfm closed 6 months ago

dfm commented 2 years ago

The major change is that this will now use pymc v4 if it is installed, but should continue to work with pymc3 in a more or less completely backwards compatible way.

From the user perspective, I think the biggest changes will be:

  1. My plan is to deprecated pymc3-ext since nearly everything is implemented in pymc now.
  2. The custom distributions are all functions now instead of distributions with transforms since that was a more straightforward mode for implementing both pymc and pymc3 support, since the handling of distributions has completely changed.
  3. ... anything else?

Still to do:

saimn commented 1 year ago

Hi @dfm, First, thanks for the great package (and the very good docs!) :) Is there anything I could do to help finalizing the compatibility with PyMC4 ? I have some code based on exoplanet and which I converted to PyMC4 using this branch (and exoplanet-core's main branch) some time ago, hoping I would not need to write a compat layer for PyMC3, and it works well (I don't use all exoplanet features though). So it seems you already did most of the work here, what else is needed to get it merged ? Would be happy to help if I can, instead of spending some time to make my code compatible with PyMC3 ;)

dfm commented 1 year ago

@saimn — Thanks for the kind words! Unfortunately, the situation here is not great. I was nearly ready to merge this (I still needed to run through the case studies to make sure everything was working there, but things were looking good!) when the PyMC project announced their plans for the (now released) v5 of PyMC which will require yet another overhaul of the back end (probably small and feasible, but an absolute nightmare for dependency management, etc.). And the near-term plan of the project is to completely remove the C back end in favor of numba, which would require a full rewrite of exoplanet-core. Given the volatility of the PyMC project, I have a lot of trouble mustering the enthusiasm or even capacity to keep rewriting all of exoplanet every few months. This might be less true for other high-level use cases of PyMC, but the low-level interfaces that exoplanet relies on are consistently broken every new release! I'm very happy to revisit this once things settle down, but for now, my priorities are on maintaining support for PyMC3 (which is still being widely used in a lot of projects), and working on the pure-JAX version of exoplanet (still early days and poorly documented, but what I'm most excited about!).

Sorry that I can't give a more positive response, and I'm always happy to chat more about this!

saimn commented 1 year ago

Thanks for the quick answer. Yeah I understand the pain to maintain the compatibility with the various versions, and the PyTensor fork is clearly not helping... I was interested by moving to PyMC4 to test their new JAX backend, which I haven't done yet but it seems promising. Pure-JAX is another very interesting option I need to look at, but for now I need to finalize a proof of concept so I guess I will go back to PyMC3 for that.

dfm commented 1 year ago

@saimn — Yeah I think that's what I'd recommend for now if that's ok! Please let me know if you run into issues and I'm happy to help.

vandalt commented 12 months ago

Hi @dfm! With exoplanet-dev/exoplanetcore, exoplanet-dev/celerite2 and exoplanet-dev/pymc-ext now (or soon) supporting PyMC v5, is there anything I can do to help exoplanet support it as well? I'm guessing the latest state of this PR would be a good starting point?

dfm commented 12 months ago

@vandalt — Yeah, I think that starting from this PR is probably the right way to go because I made some progress handling the backwards incompatibilities. If you're keen, I'd be very happy for you to take over!