feat: Dagum distribution

.github/workflows/ci.yml

@@ -79,7 +79,7 @@ jobs:
         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
+        JAX_ENABLE_X64=1 pytest -vs test/test_distributions.py -k "powerLaw or Dagum"
     - name: Test tracer leak
       if: matrix.python-version == '3.10'
       env:

docs/source/distributions.rst

@@ -144,6 +144,14 @@ CirculantNormal
     :show-inheritance:
     :member-order: bysource
 
+Dagum
+^^^^^
+.. autoclass:: numpyro.distributions.continuous.Dagum
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 Dirichlet
 ^^^^^^^^^
 .. autoclass:: numpyro.distributions.continuous.Dirichlet

numpyro/distributions/__init__.py

@@ -20,6 +20,7 @@
     Cauchy,
     Chi2,
     CirculantNormal,
+    Dagum,
     Dirichlet,
     EulerMaruyama,
     Exponential,
@@ -134,6 +135,7 @@
     "Cauchy",
     "Chi2",
     "CirculantNormal",
+    "Dagum",
     "Delta",
     "Dirichlet",
     "DirichletMultinomial",

numpyro/distributions/continuous.py

@@ -25,6 +25,9 @@
 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 # POSSIBILITY OF SUCH DAMAGE.
 
+
+from typing import Optional
+
 import numpy as np
 
 from jax import lax, vmap
@@ -3212,3 +3215,105 @@ def entropy(self):
         (n,) = self.event_shape
         log_abs_det_jacobian = 2 * jnp.log(2) * ((n - 1) // 2) - jnp.log(n) * n
         return self.base_dist.entropy() + log_abs_det_jacobian / 2
+
+
+class Dagum(Distribution):
+    arg_constraints = {
+        "p": constraints.positive,
+        "a": constraints.positive,
+        "b": constraints.positive,
+    }
+    support = constraints.positive
+    reparametrized_params = ["p", "a", "b"]
+
+    def __init__(
+        self,
+        p: ArrayLike,
+        a: ArrayLike,
+        b: ArrayLike,
+        *,
+        validate_args: Optional[bool] = None,
+    ) -> None:
+        r"""The Dagum distribution (or Mielke Beta-Kappa distribution) is a continuous
+        probability distribution defined over positive real numbers. If :math:`p`,
+        :math:`a` and :math:`b` are shape values, then Dagum distribution is defined as,
+
+        .. math::
+
+            f(x\mid p,a,b):=\frac{ap}{x}
+            \left(\frac{(x/b)^{ap}}{\left((x/b)^{a}+1\right)^{p+1}}\right)
+
+        :param p: shape parameter :math:`p>0`.
+        :param a: shape parameter :math:`a>0`.
+        :param b: scale parameter :math:`b>0`.
+
+        **References:**
+
+        1. Wikipedia. (n.d.). Dagum distribution. Retrieved March 31, 2025, from
+           https://en.wikipedia.org/wiki/Dagum_distribution
+        """
+        self.p, self.a, self.b = promote_shapes(p, a, b)
+        batch_shape = lax.broadcast_shapes(jnp.shape(p), jnp.shape(a), jnp.shape(b))
+        super().__init__(batch_shape=batch_shape, validate_args=validate_args)
+
+    @validate_sample
+    def log_prob(self, value: ArrayLike) -> ArrayLike:
+        a_ln_x_m_ln_b = xlogy(self.a, value) - xlogy(self.a, self.b)
+        return (
+            jnp.log(self.a)
+            + jnp.log(self.p)
+            - jnp.log(value)
+            + self.p * a_ln_x_m_ln_b
+            - (self.p + 1.0) * nn.softplus(a_ln_x_m_ln_b)
+        )
+
+    def cdf(self, value: ArrayLike) -> ArrayLike:
+        return jnp.exp(
+            -self.p * nn.softplus(xlogy(self.a, self.b) - xlogy(self.a, value))
+        )
+
+    def icdf(self, q: ArrayLike) -> ArrayLike:
+        q_root_p = jnp.power(q, -jnp.reciprocal(self.p))
+        return self.b * jnp.power(q_root_p - 1.0, -jnp.reciprocal(self.a))
+
+    def sample(self, key, sample_shape: tuple[int, ...] = ()) -> jnp.ndarray:
+        assert is_prng_key(key)
+        return self.icdf(random.uniform(key, shape=self.shape(sample_shape)))
+
+    @property
+    def mean(self) -> ArrayLike:
+        safe_a = jnp.where(self.a > 1.0, self.a, 2.0)
+        return jnp.where(
+            self.a > 1.0,
+            self.b
+            * jnp.exp(
+                gammaln(1.0 - 1.0 / safe_a)
+                + gammaln(self.p + 1.0 / safe_a)
+                - gammaln(self.p)
+            ),
+            jnp.inf,
+        )
+
+    @property
+    def variance(self) -> ArrayLike:
+        safe_a = jnp.where(self.a > 2.0, self.a, 3.0)
+        return jnp.where(
+            self.a > 2.0,
+            jnp.square(self.b)
+            * (
+                jnp.exp(
+                    gammaln(1.0 - 2.0 / safe_a)
+                    + gammaln(self.p + 2.0 / safe_a)
+                    - gammaln(self.p)
+                )
+                - jnp.exp(
+                    2.0
+                    * (
+                        gammaln(1.0 - 1.0 / safe_a)
+                        + gammaln(self.p + 1.0 / safe_a)
+                        - gammaln(self.p)
+                    )
+                )
+            ),
+            jnp.inf,
+        )

test/test_distributions.py

@@ -959,6 +959,11 @@ def get_sp_dist(jax_dist):
     T(dist.Levy, 0.0, 1.0),
     T(dist.Levy, 0.0, np.array([1.0, 2.0, 10.0])),
     T(dist.Levy, np.array([1.0, 2.0, 10.0]), np.pi),
+    T(dist.Dagum, 1.0, 1.0, 1.0),
+    T(dist.Dagum, 3.0, 4.0, 5.0),
+    T(dist.Dagum, 2.0, np.array([1.0, 2.0, 10.0]), 4.0),
+    T(dist.Dagum, 2.0, 3.0, np.array([0.5, 2.0, 1.0])),
+    T(dist.Dagum, np.array([5.0, 2.0, 10.0]), 3.0, 5.0),
 ]
 
 DIRECTIONAL = [
@@ -1390,10 +1395,10 @@ def test_sample_gradient(jax_dist, sp_dist, params):
     }.get(jax_dist.__name__, [])
 
     if (
-        jax_dist in [dist.DoublyTruncatedPowerLaw]
+        jax_dist in [dist.DoublyTruncatedPowerLaw, dist.Dagum]
         and jnp.result_type(float) == jnp.float32
     ):
-        pytest.skip("DoublyTruncatedPowerLaw is tested with x64 only.")
+        pytest.skip(f"{jax_dist.__name__} is tested with x64 only.")
 
     dist_args = [
         p
@@ -1856,10 +1861,10 @@ def test_log_prob_gradient(jax_dist, sp_dist, params):
     if jax_dist is _ImproperWrapper:
         pytest.skip("no param for ImproperUniform to test for log_prob gradient")
     if (
-        jax_dist in [dist.DoublyTruncatedPowerLaw]
+        jax_dist in [dist.DoublyTruncatedPowerLaw, dist.Dagum]
         and jnp.result_type(float) == jnp.float32
     ):
-        pytest.skip("DoublyTruncatedPowerLaw is tested with x64 only.")
+        pytest.skip(f"{jax_dist.__name__} is tested with x64 only.")
 
     rng_key = random.PRNGKey(0)
     value = jax_dist(*params).sample(rng_key)
@@ -1940,6 +1945,8 @@ def test_mean_var(jax_dist, sp_dist, params):
         pytest.skip(
             f"{jax_dist.__name__} distribution does not has mean/var implemented"
         )
+    if jax_dist in [dist.Dagum] and jnp.result_type(float) == jnp.float32:
+        pytest.skip(f"{jax_dist.__name__} is tested with x64 only.")
 
     n = (
         20000