Support PyMC v5

[vandalt]

This builds upon #271, where pretty much all the hard work required to update from PyMC3 had already been done, and implements support of PyMC v5. Most of the changes were replacing aesara by pytensor plus a few API changes.

Some things left to address:

• [x] Should the joss directory be updated? I'm guessing the answer is no
• [ ] Update requirement version for exoplanet-core, pymc-ext and celerite2 once pymc >= 5 support is released there. This is required for the tests to pass.
• [ ] Is it OK to remove warning filters in utils.docs_setup()? I left TODO comments to discuss this.
• [ ] There is a reference to pytensor.opt.Assert (which used to be theano.opt.Assert) which no longer exists. See comment in the code
• [ ] For the Docs: is the lowest compatible Python version 3.6? I have not tested this but I saw that 3.8 is the lowest version test by nox (and the lowest still getting security updates I think)
• [ ] Updating the case studies. Not technically required before merging into main, but probably worth waiting in case it exposes bugs or deprecation warnings.

[dfm]

@vandalt — Thanks!! 🚀

It will take me a week or so to go through this in detail, but here are some initial comments.

First, can you take a look at the merge conflicts with main? You could rebase if you want, but I'd also be fine with the merge commit (merge from main, fix the conflicts, then commit).

Should the joss directory be updated? I'm guessing the answer is no

Agreed. I think we should leave that as is.

Update requirement version for exoplanet-core, pymc-ext and celerite2 once pymc >= 5 support is released there. This is required for the tests to pass.

I'll try to get those released ASAP. Might be worth adding a test against the GitHub versions of those in the meantime as well though? What do you think?

Is it OK to remove warning filters in utils.docs_setup()? I left TODO comments to discuss this.

This is fine with me as long as pytensor doesn't spew so many warnings!

For the Docs: is the lowest compatible Python version 3.6? I have not tested this but I saw that 3.8 is the lowest version test by nox (and the lowest still getting security updates I think)

I'm happy to update to 3.8 or 3.9 as the lower bound!

Updating the case studies. Not technically required before merging into main, but probably worth waiting in case it exposes bugs or deprecation warnings.

I think this is a big one, but I don't think we should worry about solving it before merging this PR. Instead, we should open a separate issue to track that.

[vandalt]

Thanks!

I'll try to get those released ASAP. Might be worth adding a test against the GitHub versions of those in the meantime as well though? What do you think?

Yes, I can add that!

This is fine with me as long as pytensor doesn't spew so many warnings!

I'll try removing and make sure there are not too many warnings when running the tests with pytest -s or the tutorials

I think this [case studies update] is a big one, but I don't think we should worry about solving it before merging this PR. Instead, we should open a separate issue to track that.

I agree it's a big one, but I think trying a quick update+run for each notebooks (without necessarily fixing everything) would be worthwhile. So far the pRV one helped uncover a problem with passing observed=ecc to eccentricity priors (basically having observed pointing to a prior or Deterministic parameter no longer works. We need to use pm.Potential). I added a test for this and will push it soon, hopefully with a fix

[vandalt]

@dfm another point I forgot to raise is that unit_disk and angles are defined both here and in exoplanet-dev/pymc-ext. Should one of them be removed (with a deprecation warning at first)?

[vandalt]

Hi @dfm!

I finally took time to look at the remaining issues with this PR.

• I just finished a first attempt at updating the case studies. I'll send a PR on that repo later today.
• While doing that, I noticed (as mentioned above) that PyMC 5 no longer supports observed to depend on other random variables/model parameters. This required an update of the eccentricity prior to use pm.Potential() when ecc is a derived parameter.
• I rebased onto main (not super experienced with rebase/merged conflict, hopefully I did not mess the commit history too much!)

I think the only remaining thing would be to update the celerite2, pymc-ext and exoplanet-core requirements. Let me know if it's best to wait for release candidates or if I should test against the git versions for now.

Thanks!

[dfm]

Thanks for all your work on this, @vandalt!!

I'll work on getting those other packages released ASAP, but it probably won't be until next week. Can you try to remember to ping me if you don't see any motion by Monday afternoon or Tuesday morning next week?

[dfm]

@vandalt — I did manage to publish new releases of all of the dependencies yesterday and re-ran the tests. Any guesses for why the PyMC3 tests are failing now?

[vandalt]

@dfm

Yes sorry, I had forgotten to re-run the tests with PyMC3 locally. Most errors were related to differences in how to access the shape of a variable in PyMC3 vs 5.

For the Kipping eccentricity prior, I had added a test for something that was not implemented in the PyMC3 version (truncated distribution with observed eccentricity), so I modified the test to expect the NotImplementedError.

For the vaneylen19 eccentricity prior, I'm not 100% sure what caused the issue, but reverting the PyMC3 version of the prior closer to its original form fixed the issue on my local setup. This comes at the cost of slightly more if USING_PYMC3 conditions, but at least now the tests pass for both versions!

[dfm]

Great - thanks!! I'll look into debugging the docs build errors because those can be tricky to track down.

[vandalt]

It seems the docs thing was not too tricky (and 100% my fault :sweat_smile:): I had moved exoplanet-core from the "main" dependencies to the pymc and pymc3 extra groups, because the required versions were now different. I just added exoplanet-core to the docs extra as well.

[dfm]

@vandalt — I pushed a few changes to fix some merge conflicts and get the tests passing, but now I think we're good to go here!! I'm going to merge and we can fix issues as they come up in future PRs. Thanks again for all your work getting this over the finish line!! I (and the whole community) really really appreciate everything that you've done here!!! 🚀