gh-1806: Implementation of Doubly Truncated Power Law and Lower Truncated Power Law (#1807)

* implementation of DoublyTruncatedPowerLaw

* implementation of LowerTruncatedPowerLaw

* chore: mathematical description in docstrings

* chore: mathematical details of `LowerTruncatedPowerLaw`

* chore: Fix bug in DoublyTruncatedPowerLaw cdf and icdf calculation

* chore: Refactor mean and variance calculation by using kth-moment in DoublyTruncatedPowerLaw

* chore: Refactor mean and variance calculation in LowerTruncatedPowerLaw

* chore: masking in icdf of LowerTruncatedPowerLaw

* chore: entropy of LowerTruncatedPowerLaw

* chore: `lax.sqaure` replaced with `jnp.sqaure`

* chore: moments and entropy were extra and removed

* chore: unit tests

* fix: nan gradients fixed, values still diverging

* Updated UpperTruncatedPowerLaw with adequate derivations, including fixing the discontinuity point for alpha equals minus one and correct tangents for lower and upper bounds.

* Changed constrains of alpha of LowerTruncatedPowerLaw to the smaller minus one and also changed equation to keep equational uniformity UpperTruncatedPowerLaw to and improved stability of the calculation by integrating formula parts inside log to prevent NaN values due to negative values of some parts that could balance out in the end.

* chore: code and docstring formated

* chore: equation refactor and simplified

* chore: equation refactor and simplified

* chore: use numpy arrays and numpy constants

* chore: high precision computation enable for powerlaws

* chore: `__name__` attribute calls removed

* chore: powerlaws shifted with truncated distributions

* chore: spelling mistakes fixed with code spell checker pre-commit hook

* fix typo: perforance->perforamce->performance

* chore: explicit enabling/disabling of 64bit floating point numbers

* chore: disable everytime and enable x64 for power laws

* chore: disable x64 for every test

* chore: linked explanation in comments for disabling x64 for future reference for devs

* chore: high precision test handeled efficiently for DoublyTruncatedPowerLaw

* chore: high precision exception handled in test_log_prob_gradient

---------

Co-authored-by: David Ziegler <[email protected]>

Author: Qazalbash

.github/workflows/ci.yml

@@ -72,6 +72,9 @@ jobs:
     - name: Test with pytest
       run: |
         CI=1 pytest -vs -k "not test_example" --durations=100 --ignore=test/infer/ --ignore=test/contrib/
+    - name: Test x64
+      run: |
+        JAX_ENABLE_X64=1 pytest -vs test/test_distributions.py -k powerLaw
 
 
   test-inference:

.pre-commit-config.yaml

@@ -21,3 +21,11 @@ repos:
       - id: check-yaml
       - id: check-added-large-files
         exclude: notebooks/*
+
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.3.0
+    hooks:
+      - id: codespell
+        stages: [commit, commit-msg]
+        args:
+          [--ignore-words-list, "Teh,aas", --check-filenames, --skip, "*.ipynb"]

README.md

@@ -34,5 +34,5 @@ Design Choices:
 Future Work:
 
 - Right now the jax, jaxlib, and numpyro versions are manually specified, so they have to be updated every NumPyro release. There are two ways forward for this:
-    1. If there is a CI/CD in place to build and push images to a repository like Dockerhub, then the jax, jaxlib, and numpyro versions can be passed in as environment variables (for example, if something like [Drone CI](http://plugins.drone.io/drone-plugins/drone-docker/) is used). If implemented this way, the jax/jaxlib/numpyro versions will be ephemereal (not stored in source code).
+    1. If there is a CI/CD in place to build and push images to a repository like Dockerhub, then the jax, jaxlib, and numpyro versions can be passed in as environment variables (for example, if something like [Drone CI](http://plugins.drone.io/drone-plugins/drone-docker/) is used). If implemented this way, the jax/jaxlib/numpyro versions will be ephemeral (not stored in source code).
     2. Alternative, one can create a Python script that will modify the Dockerfiles upon release accordingly (using a hook of some sort).

docker/README.md

@@ -662,6 +662,14 @@ VonMises
 Truncated Distributions
 -----------------------
 
+DoublyTruncatedPowerLaw
+^^^^^^^^^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.truncated.DoublyTruncatedPowerLaw
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 LeftTruncatedDistribution
 ^^^^^^^^^^^^^^^^^^^^^^^^^
 .. autoclass:: numpyro.distributions.truncated.LeftTruncatedDistribution
@@ -670,6 +678,14 @@ LeftTruncatedDistribution
     :show-inheritance:
     :member-order: bysource
 
+LowerTruncatedPowerLaw
+^^^^^^^^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.truncated.LowerTruncatedPowerLaw
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 RightTruncatedDistribution
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
 .. autoclass:: numpyro.distributions.truncated.RightTruncatedDistribution

docs/source/distributions.rst

@@ -9,7 +9,7 @@
 
 All models have discrete latent variables. Under the hood, we enumerate over
 (marginalize out) those discrete latent sites in inference. Those models have different
-complexity so they are great refererences for those who are new to Pyro/NumPyro
+complexity so they are great references for those who are new to Pyro/NumPyro
 enumeration mechanism. We recommend readers compare the implementations with the
 corresponding plate diagrams in [1] to see how concise a Pyro/NumPyro program is.
 

examples/annotation.py

@@ -103,7 +103,7 @@ def run_inference(model, args, rng_key, y):
 
 
 def main(args):
-    # generate artifical dataset
+    # generate artificial dataset
     num_data = args.num_data
     rng_key = jax.random.PRNGKey(0)
     t = jnp.arange(0, num_data)

examples/ar2.py

@@ -147,7 +147,7 @@ def predict(model, args, samples, rng_key, y, n_seasons):
 
 
 def main(args):
-    # generate artifical dataset
+    # generate artificial dataset
     rng_key, _ = random.split(random.PRNGKey(0))
     T = args.T
     t = jnp.linspace(0, T + args.future, (T + args.future) * N_POINTS_PER_UNIT)

examples/holt_winters.py

@@ -57,7 +57,7 @@
 dimensions of the age, space and time variables. This allows us to efficiently broadcast arrays
 in the likelihood.
 
-As written above, the model includes a lot of centred random effects. The NUTS alogrithm benefits
+As written above, the model includes a lot of centred random effects. The NUTS algorithm benefits
 from a non-centred reparamatrisation to overcome difficult posterior geometries [2]. Rather than
 manually writing out the non-centred parametrisation, we make use of the NumPyro's automatic
 reparametrisation in :class:`~numpyro.infer.reparam.LocScaleReparam`.

examples/mortality.py

@@ -95,7 +95,9 @@
 from numpyro.distributions.mixtures import Mixture, MixtureGeneral, MixtureSameFamily
 from numpyro.distributions.transforms import biject_to
 from numpyro.distributions.truncated import (
+    DoublyTruncatedPowerLaw,
     LeftTruncatedDistribution,
+    LowerTruncatedPowerLaw,
     RightTruncatedDistribution,
     TruncatedCauchy,
     TruncatedDistribution,
@@ -122,6 +124,7 @@
     "Binomial",
     "BinomialLogits",
     "BinomialProbs",
+    "CAR",
     "Categorical",
     "CategoricalLogits",
     "CategoricalProbs",
@@ -132,9 +135,10 @@
     "DirichletMultinomial",
     "DiscreteUniform",
     "Distribution",
+    "DoublyTruncatedPowerLaw",
     "EulerMaruyama",
-    "Exponential",
     "ExpandedDistribution",
+    "Exponential",
     "FoldedDistribution",
     "Gamma",
     "GammaPoisson",
@@ -152,29 +156,29 @@
     "Independent",
     "InverseGamma",
     "Kumaraswamy",
-    "LKJ",
-    "LKJCholesky",
     "Laplace",
     "LeftTruncatedDistribution",
+    "LKJ",
+    "LKJCholesky",
     "Logistic",
     "LogNormal",
     "LogUniform",
-    "MatrixNormal",
+    "LowerTruncatedPowerLaw",
+    "LowRankMultivariateNormal",
     "MaskedDistribution",
+    "MatrixNormal",
     "Mixture",
-    "MixtureSameFamily",
     "MixtureGeneral",
+    "MixtureSameFamily",
     "Multinomial",
     "MultinomialLogits",
     "MultinomialProbs",
     "MultivariateNormal",
-    "CAR",
     "MultivariateStudentT",
-    "LowRankMultivariateNormal",
-    "Normal",
-    "NegativeBinomialProbs",
-    "NegativeBinomialLogits",
     "NegativeBinomial2",
+    "NegativeBinomialLogits",
+    "NegativeBinomialProbs",
+    "Normal",
     "OrderedLogistic",
     "Pareto",
     "Poisson",
@@ -199,7 +203,7 @@
     "Wishart",
     "WishartCholesky",
     "ZeroInflatedDistribution",
-    "ZeroInflatedPoisson",
     "ZeroInflatedNegativeBinomial2",
+    "ZeroInflatedPoisson",
     "ZeroSumNormal",
 ]