rlouf / mcx

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.
https://rlouf.github.io/mcx
Apache License 2.0
325 stars 17 forks source link

Add pareto distribution #82

Open tblazina opened 3 years ago

tblazina commented 3 years ago

I'm not familiar really with this distribution and am a bit confused with all the different ways it is parameterized depending on which library you look at, for example Stan and PyMC3 both use a shape and scale parameter but the jax.scipy.stats implementation uses a parameter b as well as a loc and scale parameter. I guess the use of the b parameter stems from the jax.random.pareto function which I seems to be similar to the Numpy implementation where it is the "The Lomax or Pareto II distribution is a shifted Pareto distribution". I am not sure which would be preferable to use, some guidance/input would be appreciated. 🙏

rlouf commented 3 years ago

If we note $m$ the scale parameter and $b$ the shape parameter, my understanding is that JAX implemnts:

$$P(x) = \frac{b m^b}{(x-loc)^{b+1}}$$

While on the other hand PyMC3 implements

$$P(x) = \frac{b m^{b}}{x^{b+1}}$$

What I would do is leave scale and shape as the first two arguments when initializing the distribution (same defaults as PyMC3) and add a keyword argument loc = 0. So it would have the signature Pareto(shape, scale, loc=0). What do you think?

tblazina commented 3 years ago

Sounds reasonable to me 👍 Would you prefer using b and m as the parameter names or rather shape and scale? I personally like shape and scale better but not sure it there is some reason that being consistent with the jax implementation would be preferable?

rlouf commented 3 years ago

shape and scale make more sense to me, and it is better to have an API close to PyMC3's. You can keep loc.

rlouf commented 3 years ago

Hi @tblazina Looks like the tests are not passing :( Are you still planning on working on this?

tblazina commented 3 years ago

Hi @rlouf - sorry about that, this fell by the wayside but I would plan on finishing this yes! I'll try to get to it in the next few days and if I don't think I can get around to doing it I will let you know!