Added named_arrays.random.gamma() function to draw samples from the gamma distribution. (#106)

Author: roytsmart

named_arrays/_scalars/scalar_named_array_functions.py

@@ -625,6 +625,56 @@ def random_binomial(
     )
 
 
+@_implements(na.random.gamma)
+def random_gamma(
+    shape: float | na.AbstractScalarArray,
+    scale: float | u.Quantity | na.AbstractScalarArray = 1,
+    shape_random: None | dict[str, int] = None,
+    seed: None | int = None,
+) -> na.ScalarArray:
+    alpha = shape
+    theta = scale
+
+    try:
+        alpha = scalars._normalize(alpha)
+        theta = scalars._normalize(theta)
+    except na.ScalarTypeError:
+        return NotImplemented
+
+    if shape_random is None:
+        shape_random = dict()
+
+    shape_base = na.shape_broadcasted(alpha, theta)
+    shape = na.broadcast_shapes(shape_base, shape_random)
+
+    alpha = alpha.ndarray_aligned(shape)
+    theta = theta.ndarray_aligned(shape)
+
+    unit = na.unit(theta)
+
+    if unit is not None:
+        theta = theta.value
+
+    if seed is None:
+        func = np.random.gamma
+    else:
+        func = np.random.default_rng(seed).gamma
+
+    value = func(
+        shape=alpha,
+        scale=theta,
+        size=tuple(shape.values()),
+    )
+
+    if unit is not None:
+        value = value << unit
+
+    return na.ScalarArray(
+        ndarray=value,
+        axes=tuple(shape.keys()),
+    )
+
+
 def plt_plot_like(
         func: Callable,
         *args: na.AbstractScalarArray,

named_arrays/_scalars/uncertainties/uncertainties_named_array_functions.py

@@ -24,6 +24,7 @@
     na.random.normal,
     na.random.poisson,
     na.random.binomial,
+    na.random.gamma,
 )
 PLT_PLOT_LIKE_FUNCTIONS = named_arrays._scalars.scalar_named_array_functions.PLT_PLOT_LIKE_FUNCTIONS
 NDFILTER_FUNCTIONS = named_arrays._scalars.scalar_named_array_functions.NDFILTER_FUNCTIONS

named_arrays/_vectors/vector_named_array_functions.py

@@ -27,6 +27,7 @@
     na.random.normal,
     na.random.poisson,
     na.random.binomial,
+    na.random.gamma,
 )
 PLT_PLOT_LIKE_FUNCTIONS = named_arrays._scalars.scalar_named_array_functions.PLT_PLOT_LIKE_FUNCTIONS
 NDFILTER_FUNCTIONS = named_arrays._scalars.scalar_named_array_functions.NDFILTER_FUNCTIONS

named_arrays/random.py

@@ -10,6 +10,7 @@
     "normal",
     "poisson",
     "binomial",
+    "gamma",
 ]
 
 
@@ -19,6 +20,8 @@
 RandomWidthT = TypeVar("RandomWidthT", bound="float | complex | u.Quantity | na.AbstractArray")
 NumTrialsT = TypeVar("NumTrialsT", bound="int | na.AbstractArray")
 ProbabilityT = TypeVar("ProbabilityT", bound="float | na.AbstractArray")
+ShapeT = TypeVar("ShapeT", bound="float | na.AbstractArray")
+ScaleT = TypeVar("ScaleT", bound="float | u.Quantity | na.AbstractArray")
 
 
 def uniform(
@@ -164,3 +167,37 @@ def binomial(
         shape_random=shape_random,
         seed=seed,
     )
+
+
+def gamma(
+    shape: ShapeT,
+    scale: ScaleT = 1,
+    shape_random: None | dict[str, int] = None,
+    seed: None | int = None,
+) -> ShapeT | ScaleT:
+    """
+    Draw samples from a gamma distribution.
+
+    Parameters
+    ----------
+    shape
+        The shape parameter of the distribution.
+    scale
+        The scale parameter of the distribution.
+    shape_random
+        Additional dimensions to be broadcast against `shape` and `scale`.
+    seed
+        Optional seed for the random number generator,
+        can be provided for repeatability.
+
+    See Also
+    --------
+    :func:`numpy.random.gamma` : Equivalent numpy function
+    """
+    return na._named_array_function(
+        func=gamma,
+        shape=na.as_named_array(shape),
+        scale=scale,
+        shape_random=shape_random,
+        seed=seed,
+    )

named_arrays/tests/test_random.py

@@ -55,3 +55,54 @@ def test_binomial(
 
     assert np.all(result >= 0)
     assert np.all(result <= n)
+
+
+@pytest.mark.parametrize(
+    argnames="shape",
+    argvalues=[
+        0.5,
+        na.ScalarArray(0.51),
+        na.linspace(0.4, 0.5, axis="p", num=5),
+        na.UniformUncertainScalarArray(0.5, width=0.1),
+        na.Cartesian2dVectorArray(0.5, 0.6),
+    ],
+)
+@pytest.mark.parametrize(
+    argnames="scale",
+    argvalues=[
+        10,
+        (11 * u.photon).astype(int),
+        na.ScalarArray(12),
+        (na.arange(1, 10, axis="x") << u.photon).astype(int),
+        na.Cartesian2dVectorArray(10, 11),
+    ],
+)
+@pytest.mark.parametrize(
+    argnames="shape_random",
+    argvalues=[
+        None,
+        dict(_s=6),
+    ],
+)
+@pytest.mark.parametrize(
+    argnames="seed",
+    argvalues=[
+        None,
+        42,
+    ],
+)
+def test_gamma(
+    shape: float | na.AbstractScalar | na.AbstractVectorArray,
+    scale: float | u.Quantity | na.AbstractScalar | na.AbstractVectorArray,
+    shape_random: None | dict[str, int],
+    seed: None | int,
+):
+    result = na.random.gamma(
+        shape=shape,
+        scale=scale,
+        shape_random=shape_random,
+        seed=seed,
+    )
+
+    assert na.unit(result) == na.unit(scale)
+    assert np.all(result >= 0)