Function to optimize prior under constraints

[AlexAndorra]

Ready for review 🍾

This PR introduces the pm.find_optim_prior function, that finds the optimal prior parameters of a distribution under some constraints.

For instance, sometimes, while building a model, we think "I need a Gamma that has 95% probability mass between 0.1 and 0.4". It may take a long time to find an appropriate combination of parameters that give a decent Gamma.
pm.find_optim_prior basically answers the question: "What are the alpha and beta parameters that fit this constraint?".

Notes:

• pm.find_optim_prior can be used with any PyMC distribution for which the log-cdf function is implemented.
• For now, it only works for two-parameter distributions, as there are exactly two constraints (a lower and an upper bound).
• When possible, it uses Aesara and gradients of the PyMC log CDF for more efficient computations. Otherwise, it falls back to Scipy's default.

Depending on what your PR does, here are a few things you might want to address in the description:

• Docstring
• Tests
• linting/style checks have been run
• Mentioned in RELEASE-NOTES.md

[codecov]

Codecov Report

Merging #5231 (d89e375) into main (600fe90) will increase coverage by 0.00%.
The diff coverage is 87.80%.

@@           Coverage Diff           @@
##             main    #5231   +/-   ##
=======================================
  Coverage   79.08%   79.08%           
=======================================
  Files          87       88    +1     
  Lines       14394    14435   +41     
=======================================
+ Hits        11383    11416   +33     
- Misses       3011     3019    +8     

Impacted Files	Coverage Δ
pymc/func_utils.py	87.50% <87.50%> (ø)
pymc/__init__.py	96.07% <100.00%> (+0.07%)	⬆️
pymc/parallel_sampling.py	86.33% <0.00%> (-1.00%)	⬇️

[ricardoV94]

Looks cool, left some suggestions below

[ferrine]

I would do the API differently to make that less verbose and elegant. Ideally, something like this:

intercept = pm.Normal.fit("intercept", lower=0.1, upper=0.9)
# That could (in future) be extended as
intercept = pm.Normal.fit("intercept", data=input_data)

[ricardoV94]

I would do the API differently to make that less verbose and elegant. Ideally, something like this:

intercept = pm.Normal.fit("intercept", lower=0.1, upper=0.9)
That could (in future) be extended as
intercept = pm.Normal.fit("intercept", data=input_data)

I don't agree. You still need to provide initial parameter guesses, lower/upper can be distribution parameters, and you might want to specify that some parameters should be fixed, as I mentioned above. All that is more difficult if you mix optimization goals with distribution initialization.

Also are you suggesting this would automatically add a new random variable to the model? I think this suggestion can create a lot of ambiguity, similar to how the .dist often confuses people. Specially because you need to play nice with all the additional keywords like dims, size, observed, transform and so on...

[ricardoV94]

Should this go into util.py? That's not a very discoverable file name but this is also a very small helper function.

[AlexAndorra]

I would do the API differently to make that less verbose and elegant. Ideally, something like this

@ferrine I agree with @ricardoV94 on this, at least for a first version, where we can try and see how people use that feature and what is missing (or not).

Should this go into util.py? That's not a very discoverable file name but this is also a very small helper function.

@ricardoV94 I actually put that into util.py originally, but got into issues when trying to make find_optim_prior accessible at the pm level 🤷‍♂️ It's not something I'm a master at, so if someone knows how to do that, I'm down to put it back into util.py

[ricardoV94]

You can import it explicitly in init.py I think.

[ricardoV94]

I suggest that instead of minimizing two distances: logcdf(lower) - np.log(alpha/2) and logcdf(upper) - np.log(1-alpha/2), we just minimize logdiffexp(logcdf(upper), logcdf(lower)) - np.log(mass), which among other things allows one to set the lower point to zero/a value where at the edge of the distribution support.

@aseyboldt do you see any loss in functionality by doing things like this?

[ricardoV94]

Perhaps add a test with a discrete distribution?

[AlexAndorra]

All green @ricardoV94 🥳 Let's make that Xmas miracle happen!

[ricardoV94]

I think the exception we want here is the one found in aesara.gradient.NullTypeGradError as well as NotImplementedError

Suggested change

	except Exception:
	except (NotImplementedError, NullTyppeGradError):

[AlexAndorra]

The exception that's thrown is a TypeError by the Aesara grad method (line 501) in aesara/gradient.py:

if cost is not None and cost.ndim != 0:
>           raise TypeError("cost must be a scalar.")

I added that error in the Except clause and tests pass locally. aesara.gradient.NullTypeGradError and NotImplementedError don't seem to be raised but I kept them in case they are by other cases we may have forgotten

[ricardoV94]

We shouldn't catch that TypeError. That means we produced a wrong input to aesara grad.

The other two exceptions are the ones that (should) appear when a grad is not implemented for an Op.

[AlexAndorra]

That means we should make these two exceptions appear then, shouldn't we? Because they are not raised right now -- only the TypeError is raised
(here is to my first GH comment of the year 🥂 )

[ricardoV94]

That means we should make these two exceptions appear then, shouldn't we? Because they are not raised right now -- only the TypeError is raised
(here is to my first GH comment of the year 🥂 )

Those two exceptions mean there is no grad implemented for some Op in the cdf which can very well happen and its a good reason to silently default to the scipy approximation. The TypeError, on the other hand, should not be catched, as that means we did something wrong.

In fact there was a point during this PR when it was always silently defaulting to the scipy approximation because we were passing two values to grad and suppressing the TypeError.

[ricardoV94]

That seems to be the case. Locally it is passing in float64/float32 and 1e-5 precision, so we don't need to have the separate test just for the Poisson anymore.

[AlexAndorra]

Ha ha damn, I just pushed without that change. Tests are indeed passing locally. Gonna refactor the tests and push again

[AlexAndorra]

Hot damn, tests are passing locally 🔥 Pushed!
Why does the symbolic gradien help so much with numerical errors?

[ricardoV94]

Because the logcdf uses gammaincc whose gradient is notoriously tricky. We somewhat recently added a numerically stable(r) implementation to Aesara

[ricardoV94]

aesara-devs/aesara#513

[ricardoV94]

See my reply above. We shouldn't catch the TypeError

[ricardoV94]

We should move the new tests to a file test_func_utils.py if func_utils.py is where the functions are going to live.

[AlexAndorra]

Good point, just did it

[ricardoV94]

Follow up PR (not to be sadistic with this one): we should add a code example in the docstrings

[AlexAndorra]

I'm not preaching for my choir here, but I actually should add that here. Don't merge in the meantime. Will ping when done

[AlexAndorra]

Done @ricardoV94 -- we can merge once tests pass