gh-1806: Implementation of Doubly Truncated Power Law and Lower Truncated Power Law

[Qazalbash]

This PR contains the implementation of Power Law distributions, that fixes #1806. There are two different types of power law implementations based on their truncation.

1. DoublyTruncatedPowerLaw
2. LowerTruncatedPowerLaw

They are named in such way to avoid confusion with TwoSidedTruncated* and LeftTruncated* distributions.

[Qazalbash]

@fehiepsi Power Law is not in SciPy. Can you guide me for the Unit Tests?

[fehiepsi]

Hi @Qazalbash, you could add if/else statements to skip some tests (like some other distributions). Do you have reference for this distribution? I couldn't find it on wikipedia.

[Qazalbash]

@fehiepsi unfortunately I couldn't find much on internet too. I have seen many papers in astrophysics using these distributions. Therefore, I did the derivations. If you like to see, I can share them. This way we can cross check it too.

[fehiepsi]

I think we would need to let users know how to use the distribution. Could you add references / formulas in the docstring? I can double-check the formulas - though I hope that the tests will cover them.

[Qazalbash]

test_sample_gradient and test_log_prob_gradient are failing only in case of alpha=-1.0. The gradient is coming out to be jnp.nan. I can not understand why, so I have pushed the changes for review.

[fehiepsi]

I think you need double where. See this note

[Qazalbash]

Your suggestion worked and the gradient is some float now! But values are still diverging. I have rechecked my derivations and implementation a couple of times. You can double-check it!

[InfinityMod]

@Qazalbash Please see my pull request in your repository. Probably this can solve the problems.

[Qazalbash]

@InfinityMod, I will look into it.

[Qazalbash]

@InfinityMod The only problem left is the failing of test cases, can you look into it?

[InfinityMod]

Hi, please see #1860 to solve the long testing times and the problems significantly when increasing the numerical accuracy to 64-bit. @Qazalbash, @fehiepsi, looking forward to the review and the validation.

[InfinityMod]

@Qazalbash, could you merge #1860 into this pull request?
I think we are then good to go.

[Qazalbash]

Can you change the base of your PR to Qazalbash:powerlaw-dist?

[InfinityMod]

I'm not sure if it can still be merged into the master branch if I change the base right now.
Another solution would be to wait until the merge request went through for #1860 and then change for #1807 the code of util.py (see: changelog). It should be much effort, as it's only one line to change. Then, changes align with each other.

[Qazalbash]

We should wait till #1860 gets merged to master then we will merge the master with my branch.

[fehiepsi]

i just merged the other PR. Thanks both!

[Qazalbash]

@fehiepsi I have merged the master into this branch.

[InfinityMod]

Ah, testing went through now. No stopping anymore.
However, some tests failed. But that's really due to the computation accuracy of the implemented (other) functions.
Any suggestions on what to do? I think we can't improve the accuracy for all these tests just because they failed on x64 computational accuracy..

[fehiepsi]

for some reasons, float64 is enabled for other distributions, not just power law ones. I will think a bit more on why it is the case.

[Qazalbash]

@fehiepsi @InfinityMod Finally all test cases are passing. I fixed them by disabling the higher precession at the start of every test because previously we were enabling it and not disabling it after the test, which let other tests and distributions be tested on x64 precession.

My thought process!

I tried to fix them in the following fashion,

Switch between enable and disable (f1da2d5)

if jax_dist in [dist.LowerTruncatedPowerLaw, dist.DoublyTruncatedPowerLaw]:
    numpyro.enable_x64()
else:
    numpyro.enable_x64(False)

Disable before x64 in tests for power laws (42ed59d)

numpyro.enable_x64(False)
if jax_dist in [dist.LowerTruncatedPowerLaw, dist.DoublyTruncatedPowerLaw]:
    numpyro.enable_x64()

Why no sandwich?

The next obvious thing was to sandwich tests between enable and disable, like,

if jax_dist in [dist.LowerTruncatedPowerLaw, dist.DoublyTruncatedPowerLaw]:
    numpyro.enable_x64()

# tests

if jax_dist in [dist.LowerTruncatedPowerLaw, dist.DoublyTruncatedPowerLaw]:
    numpyro.enable_x64(False)

This is okay if every test comes to the last line. Many tests were using pytest.skip which stops it from coming to an end and we have to add an x64 disable statement at those points with a concern that we might have missed any point where we could place it.

Disable every time (2671edb)

The best way I could figure out was to disable it at the start of every test and it worked!

[fehiepsi]

Glad to merge this awesome work! Thanks for adding codespell. 💯

[InfinityMod]

Awesome, glad it finally worked out!