18593169061

Committed 17 Oct 2025 12:48PM UTC coverage: 94.216%. First build

Build # 18593169061

Build Type

Pull #677

github

Committed by

web-flow

Commit Message

Merge c5562afc0 into f3875abce

Pull Request Pull Request #677: gh-410 and gh-676: Port straightforward function in observations and remove UnifiedGenerator

Run Details

201 of 203 branches covered (99.01%)

Branch coverage included in aggregate %.

106 of 115 new or added lines in 5 files covered. (92.17%)

1428 of 1526 relevant lines covered (93.58%)

7.47 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

85.71

/glass/_array_api_utils.py

"""
Array API Utilities for glass.
============================

This module provides utility functions and classes for working with multiple array
backends in the glass project, including NumPy, JAX, and array-api-strict. It includes
functions for importing backends, determining array namespaces, dispatching random
number generators, and providing missing functionality for array-api-strict through the
XPAdditions class.

Classes and functions in this module help ensure consistent behavior and compatibility
across different array libraries, and provide wrappers for common operations such as
integration, interpolation, and linear algebra.

"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, TypeAlias

import array_api_strict
import numpy as np
import numpy.random

if TYPE_CHECKING:
    from collections.abc import Callable
    from types import ModuleType

    from array_api_strict._array_object import Array as AArray
    from jaxtyping import Array as JAXArray
    from numpy.typing import DTypeLike, NDArray

    import glass.jax

    Size: TypeAlias = int | tuple[int, ...] | None

    AnyArray: TypeAlias = NDArray[Any] | JAXArray
    ComplexArray: TypeAlias = NDArray[np.complex128] | JAXArray
    DoubleArray: TypeAlias = NDArray[np.double] | JAXArray
    FloatArray: TypeAlias = NDArray[np.float64] | JAXArray
    IntArray: TypeAlias = NDArray[np.int_] | JAXArray


def import_numpy(backend: str, function_name: str) -> ModuleType:
    """
    Import the NumPy module, raising a helpful error if NumPy is not installed.

    Parameters
    ----------
    backend : str
        The name of the backend requested by the user.
    function_name : str
        The name of the function which is not implemented in the user's chosen backend.

    Returns
    -------
    ModuleType
        The NumPy module.

    Raises
    ------
    ModuleNotFoundError
        If NumPy is not found in the user's environment.

    Notes
    -----
    This is useful for explaining to the user why NumPy is required when their chosen
    backend does not implement a needed function.
    """
    try:
        import numpy as np  # noqa: PLC0415

    except ModuleNotFoundError as err:
        msg = (
            "numpy is required here as "
            + backend
            + " does not implement "
            + function_name
        )
        raise ModuleNotFoundError(msg) from err
    else:
        return np


def get_namespace(*arrays: AnyArray) -> ModuleType:
    """
    Return the array library (namespace) of input arrays if they all belong to the same
    library.

    Parameters
    ----------
    *arrays : AnyArray
        Arrays whose namespace is to be determined.

    Returns
    -------
    ModuleType
        The array namespace module.

    Raises
    ------
    ValueError
        If input arrays do not all belong to the same array library.
    """
    namespace = arrays[0].__array_namespace__()
    if any(
        array.__array_namespace__() != namespace
        for array in arrays
        if array is not None
    ):
        msg = "input arrays should belong to the same array library"
        raise ValueError(msg)

    return namespace


def rng_dispatcher(
    array: AnyArray,
) -> np.random.Generator | glass.jax.Generator | Generator:
    """
    Dispatch a random number generator based on the provided array's backend.

    Parameters
    ----------
    array : AnyArray
        The array whose backend determines the RNG.

    Returns
    -------
    np.random.Generator | glass.jax.Generator | Generator
        The appropriate random number generator for the array's backend.

    Raises
    ------
    NotImplementedError
        If the array backend is not supported.
    """
    backend = array.__array_namespace__().__name__
    if backend == "jax.numpy":
        import glass.jax  # noqa: PLC0415

        return glass.jax.Generator(seed=42)
    if backend == "numpy":
        return np.random.default_rng()
    if backend == "array_api_strict":
        return Generator(seed=42)
    msg = "the array backend in not supported"
    raise NotImplementedError(msg)


class Generator:
    """
    NumPy random number generator returning array_api_strict Array.

    This class wraps NumPy's random number generator and returns arrays compatible
    with array_api_strict.
    """

    __slots__ = ("rng",)

    def __init__(
        self,
        seed: int | bool | NDArray[np.int_ | np.bool] | None = None,  # noqa: FBT001
    ) -> None:
        """
        Initialize the Generator.

        Parameters
        ----------
        seed : int | bool | NDArray[np.int_ | np.bool] | None, optional
            Seed for the random number generator.
        """
        self.rng = numpy.random.default_rng(seed=seed)  # type: ignore[arg-type]

    def random(
        self,
        size: Size = None,
        dtype: DTypeLike | None = np.float64,
        out: NDArray[Any] | None = None,
    ) -> AArray:
        """
        Return random floats in the half-open interval [0.0, 1.0).

        Parameters
        ----------
        size : Size, optional
            Output shape.
        dtype : DTypeLike | None, optional
            Desired data type.
        out : NDArray[Any] | None, optional
            Optional output array.

        Returns
        -------
        AArray
            Array of random floats.
        """
        return array_api_strict.asarray(self.rng.random(size, dtype, out))  # type: ignore[arg-type]

    def normal(
        self,
        loc: float | NDArray[np.floating] = 0.0,
        scale: float | NDArray[np.floating] = 1.0,
        size: Size = None,
    ) -> AArray:
        """
        Draw samples from a Normal distribution (mean=loc, stdev=scale).

        Parameters
        ----------
        loc : float | NDArray[np.floating], optional
            Mean of the distribution.
        scale : float | NDArray[np.floating], optional
            Standard deviation of the distribution.
        size : Size, optional
            Output shape.

        Returns
        -------
        AArray
            Array of samples from the normal distribution.
        """
        return array_api_strict.asarray(self.rng.normal(loc, scale, size))

    def poisson(self, lam: float | NDArray[np.floating], size: Size = None) -> AArray:
        """
        Draw samples from a Poisson distribution.

        Parameters
        ----------
        lam : float | NDArray[np.floating]
            Expected number of events.
        size : Size, optional
            Output shape.

        Returns
        -------
        AArray
            Array of samples from the Poisson distribution.
        """
        return array_api_strict.asarray(self.rng.poisson(lam, size))

    def standard_normal(
        self,
        size: Size = None,
        dtype: DTypeLike | None = np.float64,
        out: NDArray[Any] | None = None,
    ) -> AArray:
        """
        Draw samples from a standard Normal distribution (mean=0, stdev=1).

        Parameters
        ----------
        size : Size, optional
            Output shape.
        dtype : DTypeLike | None, optional
            Desired data type.
        out : NDArray[Any] | None, optional
            Optional output array.

        Returns
        -------
        AArray
            Array of samples from the standard normal distribution.
        """
        return array_api_strict.asarray(self.rng.standard_normal(size, dtype, out))  # type: ignore[arg-type]

    def uniform(
        self,
        low: float | NDArray[np.floating] = 0.0,
        high: float | NDArray[np.floating] = 1.0,
        size: Size = None,
    ) -> AArray:
        """
        Draw samples from a Uniform distribution.

        Parameters
        ----------
        low : float | NDArray[np.floating], optional
            Lower bound of the distribution.
        high : float | NDArray[np.floating], optional
            Upper bound of the distribution.
        size : Size, optional
            Output shape.

        Returns
        -------
        AArray
            Array of samples from the uniform distribution.
        """
        return array_api_strict.asarray(self.rng.uniform(low, high, size))


class XPAdditions:
    """
    Additional functions missing from both array-api-strict and array-api-extra.

    This class provides wrappers for common array operations such as integration,
    interpolation, and linear algebra, ensuring compatibility across NumPy, JAX,
    and array-api-strict backends.

    This is intended as a temporary solution. See https://github.com/glass-dev/glass/issues/645
    for details.
    """

    xp: ModuleType
    backend: str

    def __init__(self, xp: ModuleType) -> None:
        """
        Initialize XPAdditions with the given array namespace.

        Parameters
        ----------
        xp : ModuleType
            The array namespace module.
        """
        self.xp = xp
        self.backend = xp.__name__

    def trapezoid(
        self, y: AnyArray, x: AnyArray = None, dx: float = 1.0, axis: int = -1
    ) -> AnyArray:
        """
        Integrate along the given axis using the composite trapezoidal rule.

        Parameters
        ----------
        y : AnyArray
            Input array to integrate.
        x : AnyArray, optional
            Sample points corresponding to y.
        dx : float, optional
            Spacing between sample points.
        axis : int, optional
            Axis along which to integrate.

        Returns
        -------
        AnyArray
            Integrated result.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/646
        """
        self.backend = self.xp.__name__
        if self.backend == "jax.numpy":
            import glass.jax  # noqa: PLC0415

            return glass.jax.trapezoid(y, x=x, dx=dx, axis=axis)
        if self.backend == "numpy":
            return self.xp.trapezoid(y, x=x, dx=dx, axis=axis)
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "trapezoid")

            # Using design principle of scipy (i.e. copy, use np, copy back)
            y_np = np.asarray(y, copy=True)
            x_np = np.asarray(x, copy=True)
            result_np = np.trapezoid(y_np, x_np, dx=dx, axis=axis)
            return self.xp.asarray(result_np, copy=True)

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def union1d(self, ar1: AnyArray, ar2: AnyArray) -> AnyArray:
        """
        Compute the set union of two 1D arrays.

        Parameters
        ----------
        ar1 : AnyArray
            First input array.
        ar2 : AnyArray
            Second input array.

        Returns
        -------
        AnyArray
            The union of the two arrays.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/647
        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.union1d(ar1, ar2)
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "union1d")

            # Using design principle of scipy (i.e. copy, use np, copy back)
            ar1_np = np.asarray(ar1, copy=True)
            ar2_np = np.asarray(ar2, copy=True)
            result_np = np.union1d(ar1_np, ar2_np)
            return self.xp.asarray(result_np, copy=True)

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def interp(  # noqa: PLR0913
        self,
        x: AnyArray,
        x_points: AnyArray,
        y_points: AnyArray,
        left: float | None = None,
        right: float | None = None,
        period: float | None = None,
    ) -> AnyArray:
        """
        One-dimensional linear interpolation for monotonically increasing sample points.

        Parameters
        ----------
        x : AnyArray
            The x-coordinates at which to evaluate the interpolated values.
        x_points : AnyArray
            The x-coordinates of the data points.
        y_points : AnyArray
            The y-coordinates of the data points.
        left : float | None, optional
            Value to return for x < x_points[0].
        right : float | None, optional
            Value to return for x > x_points[-1].
        period : float | None, optional
            Period for periodic interpolation.

        Returns
        -------
        AnyArray
            Interpolated values.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/650
        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.interp(
                x, x_points, y_points, left=left, right=right, period=period
            )
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "interp")

            # Using design principle of scipy (i.e. copy, use np, copy back)
            x_np = np.asarray(x, copy=True)
            x_points_np = np.asarray(x_points, copy=True)
            y_points_np = np.asarray(y_points, copy=True)
            result_np = np.interp(
                x_np, x_points_np, y_points_np, left=left, right=right, period=period
            )
            return self.xp.asarray(result_np, copy=True)

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def gradient(self, f: AnyArray) -> AnyArray:
        """
        Return the gradient of an N-dimensional array.

        Parameters
        ----------
        f : AnyArray
            Input array.

        Returns
        -------
        AnyArray
            Gradient of the input array.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/648
        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.gradient(f)
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "gradient")

            # Using design principle of scipy (i.e. copy, use np, copy back)
            f_np = np.asarray(f, copy=True)
            result_np = np.gradient(f_np)
            return self.xp.asarray(result_np, copy=True)

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def linalg_lstsq(
        self, a: AnyArray, b: AnyArray, rcond: float | None = None
    ) -> tuple[AnyArray, AnyArray, int, AnyArray]:
        """
        Solve a linear least squares problem.

        Parameters
        ----------
        a : AnyArray
            Coefficient matrix.
        b : AnyArray
            Ordinate or "dependent variable" values.
        rcond : float | None, optional
            Cut-off ratio for small singular values.

        Returns
        -------
        x : {(N,), (N, K)} AnyArray
            Least-squares solution. If b is two-dimensional, the solutions are in the K
            columns of x.

        residuals : {(1,), (K,), (0,)} AnyArray
            Sums of squared residuals: Squared Euclidean 2-norm for each column in b - a
            @ x. If the rank of a is < N or M <= N, this is an empty array. If b is
            1-dimensional, this is a (1,) shape array. Otherwise the shape is (K,).

        rank : int
            Rank of matrix a.

        s : (min(M, N),) AnyArray
            Singular values of a.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/649
        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.linalg.lstsq(a, b, rcond=rcond)  # type: ignore[no-any-return]
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "linalg.lstsq")

            # Using design principle of scipy (i.e. copy, use np, copy back)
            a_np = np.asarray(a, copy=True)
            b_np = np.asarray(b, copy=True)
            result_np = np.linalg.lstsq(a_np, b_np, rcond=rcond)
            return tuple(self.xp.asarray(res, copy=True) for res in result_np)

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def einsum(self, subscripts: str, *operands: AnyArray) -> AnyArray:
        """
        Evaluate the Einstein summation convention on the operands.

        Parameters
        ----------
        subscripts : str
            Specifies the subscripts for summation.
        *operands : AnyArray
            Arrays to be summed.

        Returns
        -------
        AnyArray
            Result of the Einstein summation.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/657
        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.einsum(subscripts, *operands)
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "einsum")

            # Using design principle of scipy (i.e. copy, use np, copy back)
            operands_np = (np.asarray(op, copy=True) for op in operands)
            result_np = np.einsum(subscripts, *operands_np)
            return self.xp.asarray(result_np, copy=True)

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def apply_along_axis(
        self,
        func1d: Callable[..., Any],
        axis: int,
        arr: AnyArray,
        *args: object,
        **kwargs: object,
    ) -> AnyArray:
        """
        Apply a function to 1-D slices along the given axis.

        Parameters
        ----------
        func1d : Callable[..., Any]
            Function to apply to 1-D slices.
        axis : int
            Axis along which to apply the function.
        arr : AnyArray
            Input array.
        *args : object
            Additional positional arguments to pass to func1d.
        **kwargs : object
            Additional keyword arguments to pass to func1d.

        Returns
        -------
        AnyArray
            Result of applying the function along the axis.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/651

        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.apply_along_axis(func1d, axis, arr, *args, **kwargs)
        if self.backend == "array_api_strict":
            # Import here to prevent users relying on numpy unless in this instance
            np = import_numpy(self.backend, "apply_along_axis")

            return self.xp.asarray(
                np.apply_along_axis(func1d, axis, arr, *args, **kwargs), copy=True
            )

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def vectorize(
        self,
        pyfunc: Callable[..., Any],
        otypes: tuple[type[float]],
    ) -> Callable[..., Any]:
        """
        Returns an object that acts like pyfunc, but takes arrays as input.

        Parameters
        ----------
        pyfunc : Callable[..., Any]
            Python function to vectorize.
        otypes : tuple[type[float]]
            Output types.

        Returns
        -------
        Callable[..., Any]
            Vectorized function.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Notes
        -----
        See https://github.com/glass-dev/glass/issues/671
        """
        if self.backend == "numpy":
            return self.xp.vectorize(pyfunc, otypes=otypes)  # type: ignore[no-any-return]
        if self.backend in {"array_api_strict", "jax.numpy"}:
            # Import here to prevent users relying on numpy unless in this instance
            np = import_numpy(self.backend, "vectorize")

            return np.vectorize(pyfunc, otypes=otypes)  # type: ignore[no-any-return]

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def radians(self, deg_arr: AnyArray) -> AnyArray:
        """
        Convert angles from degrees to radians.

        Parameters
        ----------
        deg_arr : AnyArray
            Array of angles in degrees.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Returns
        -------
        AnyArray
            Array of angles in radians.
        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.radians(deg_arr)
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "radians")

            return self.xp.asarray(np.radians(deg_arr))

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

    def degrees(self, deg_arr: AnyArray) -> AnyArray:
        """
        Convert angles from radians to degrees.

        Parameters
        ----------
        deg_arr : AnyArray
            Array of angles in radians.

        Raises
        ------
        NotImplementedError
            If the array backend is not supported.

        Returns
        -------
        AnyArray
            Array of angles in degrees.
        """
        if self.backend in {"numpy", "jax.numpy"}:
            return self.xp.degrees(deg_arr)
        if self.backend == "array_api_strict":
            np = import_numpy(self.backend, "degrees")

            return self.xp.asarray(np.degrees(deg_arr))

        msg = "the array backend in not supported"
        raise NotImplementedError(msg)

1	"""
2	Array API Utilities for glass.
3	============================
4
5	This module provides utility functions and classes for working with multiple array
6	backends in the glass project, including NumPy, JAX, and array-api-strict. It includes
7	functions for importing backends, determining array namespaces, dispatching random
8	number generators, and providing missing functionality for array-api-strict through the
9	XPAdditions class.
10
11	Classes and functions in this module help ensure consistent behavior and compatibility
12	across different array libraries, and provide wrappers for common operations such as
13	integration, interpolation, and linear algebra.
14
15	"""
16
17	from __future__ import annotations	8✔
18
19	from typing import TYPE_CHECKING, Any, TypeAlias	8✔
20
21	import array_api_strict	8✔
22	import numpy as np	8✔
23	import numpy.random	8✔
24
25	if TYPE_CHECKING:
26	from collections.abc import Callable
27	from types import ModuleType
28
29	from array_api_strict._array_object import Array as AArray
30	from jaxtyping import Array as JAXArray
31	from numpy.typing import DTypeLike, NDArray
32
33	import glass.jax
34
35	Size: TypeAlias = int \| tuple[int, ...] \| None
36
37	AnyArray: TypeAlias = NDArray[Any] \| JAXArray
38	ComplexArray: TypeAlias = NDArray[np.complex128] \| JAXArray
39	DoubleArray: TypeAlias = NDArray[np.double] \| JAXArray
40	FloatArray: TypeAlias = NDArray[np.float64] \| JAXArray
41	IntArray: TypeAlias = NDArray[np.int_] \| JAXArray
42
43
44	def import_numpy(backend: str, function_name: str) -> ModuleType:	8✔
45	"""
46	Import the NumPy module, raising a helpful error if NumPy is not installed.
47
48	Parameters
49	----------
50	backend : str
51	The name of the backend requested by the user.
52	function_name : str
53	The name of the function which is not implemented in the user's chosen backend.
54
55	Returns
56	-------
57	ModuleType
58	The NumPy module.
59
60	Raises
61	------
62	ModuleNotFoundError
63	If NumPy is not found in the user's environment.
64
65	Notes
66	-----
67	This is useful for explaining to the user why NumPy is required when their chosen
68	backend does not implement a needed function.
69	"""
70	try:	8✔
71	import numpy as np # noqa: PLC0415	8✔
72
NEW 73	except ModuleNotFoundError as err:	×
NEW 74	msg = (	×
75	"numpy is required here as "
76	+ backend
77	+ " does not implement "
78	+ function_name
79	)
NEW 80	raise ModuleNotFoundError(msg) from err	×
81	else:
82	return np	8✔
83
84
85	def get_namespace(*arrays: AnyArray) -> ModuleType:	8✔
86	"""
87	Return the array library (namespace) of input arrays if they all belong to the same
88	library.
89
90	Parameters
91	----------
92	*arrays : AnyArray
93	Arrays whose namespace is to be determined.
94
95	Returns
96	-------
97	ModuleType
98	The array namespace module.
99
100	Raises
101	------
102	ValueError
103	If input arrays do not all belong to the same array library.
104	"""
105	namespace = arrays[0].__array_namespace__()	8✔
106	if any(	8✔
107	array.__array_namespace__() != namespace
108	for array in arrays
109	if array is not None
110	):
111	msg = "input arrays should belong to the same array library"	8✔
112	raise ValueError(msg)	8✔
113
114	return namespace	8✔
115
116
117	def rng_dispatcher(	8✔
118	array: AnyArray,
119	) -> np.random.Generator \| glass.jax.Generator \| Generator:
120	"""
121	Dispatch a random number generator based on the provided array's backend.
122
123	Parameters
124	----------
125	array : AnyArray
126	The array whose backend determines the RNG.
127
128	Returns
129	-------
130	np.random.Generator \| glass.jax.Generator \| Generator
131	The appropriate random number generator for the array's backend.
132
133	Raises
134	------
135	NotImplementedError
136	If the array backend is not supported.
137	"""
138	backend = array.__array_namespace__().__name__	8✔
139	if backend == "jax.numpy":	8✔
140	import glass.jax # noqa: PLC0415	8✔
141
142	return glass.jax.Generator(seed=42)	8✔
143	if backend == "numpy":	8✔
144	return np.random.default_rng()	8✔
145	if backend == "array_api_strict":	8✔
146	return Generator(seed=42)	8✔
147	msg = "the array backend in not supported"	×
148	raise NotImplementedError(msg)	×
149
150
151	class Generator:	8✔
152	"""
153	NumPy random number generator returning array_api_strict Array.
154
155	This class wraps NumPy's random number generator and returns arrays compatible
156	with array_api_strict.
157	"""
158
159	__slots__ = ("rng",)	8✔
160
161	def __init__(	8✔
162	self,
163	seed: int \| bool \| NDArray[np.int_ \| np.bool] \| None = None, # noqa: FBT001
164	) -> None:
165	"""
166	Initialize the Generator.
167
168	Parameters
169	----------
170	seed : int \| bool \| NDArray[np.int_ \| np.bool] \| None, optional
171	Seed for the random number generator.
172	"""
173	self.rng = numpy.random.default_rng(seed=seed) # type: ignore[arg-type]	8✔
174
175	def random(	8✔
176	self,
177	size: Size = None,
178	dtype: DTypeLike \| None = np.float64,
179	out: NDArray[Any] \| None = None,
180	) -> AArray:
181	"""
182	Return random floats in the half-open interval [0.0, 1.0).
183
184	Parameters
185	----------
186	size : Size, optional
187	Output shape.
188	dtype : DTypeLike \| None, optional
189	Desired data type.
190	out : NDArray[Any] \| None, optional
191	Optional output array.
192
193	Returns
194	-------
195	AArray
196	Array of random floats.
197	"""
198	return array_api_strict.asarray(self.rng.random(size, dtype, out)) # type: ignore[arg-type]	8✔
199
200	def normal(	8✔
201	self,
202	loc: float \| NDArray[np.floating] = 0.0,
203	scale: float \| NDArray[np.floating] = 1.0,
204	size: Size = None,
205	) -> AArray:
206	"""
207	Draw samples from a Normal distribution (mean=loc, stdev=scale).
208
209	Parameters
210	----------
211	loc : float \| NDArray[np.floating], optional
212	Mean of the distribution.
213	scale : float \| NDArray[np.floating], optional
214	Standard deviation of the distribution.
215	size : Size, optional
216	Output shape.
217
218	Returns
219	-------
220	AArray
221	Array of samples from the normal distribution.
222	"""
223	return array_api_strict.asarray(self.rng.normal(loc, scale, size))	8✔
224
225	def poisson(self, lam: float \| NDArray[np.floating], size: Size = None) -> AArray:	8✔
226	"""
227	Draw samples from a Poisson distribution.
228
229	Parameters
230	----------
231	lam : float \| NDArray[np.floating]
232	Expected number of events.
233	size : Size, optional
234	Output shape.
235
236	Returns
237	-------
238	AArray
239	Array of samples from the Poisson distribution.
240	"""
241	return array_api_strict.asarray(self.rng.poisson(lam, size))	8✔
242
243	def standard_normal(	8✔
244	self,
245	size: Size = None,
246	dtype: DTypeLike \| None = np.float64,
247	out: NDArray[Any] \| None = None,
248	) -> AArray:
249	"""
250	Draw samples from a standard Normal distribution (mean=0, stdev=1).
251
252	Parameters
253	----------
254	size : Size, optional
255	Output shape.
256	dtype : DTypeLike \| None, optional
257	Desired data type.
258	out : NDArray[Any] \| None, optional
259	Optional output array.
260
261	Returns
262	-------
263	AArray
264	Array of samples from the standard normal distribution.
265	"""
266	return array_api_strict.asarray(self.rng.standard_normal(size, dtype, out)) # type: ignore[arg-type]	8✔
267
268	def uniform(	8✔
269	self,
270	low: float \| NDArray[np.floating] = 0.0,
271	high: float \| NDArray[np.floating] = 1.0,
272	size: Size = None,
273	) -> AArray:
274	"""
275	Draw samples from a Uniform distribution.
276
277	Parameters
278	----------
279	low : float \| NDArray[np.floating], optional
280	Lower bound of the distribution.
281	high : float \| NDArray[np.floating], optional
282	Upper bound of the distribution.
283	size : Size, optional
284	Output shape.
285
286	Returns
287	-------
288	AArray
289	Array of samples from the uniform distribution.
290	"""
291	return array_api_strict.asarray(self.rng.uniform(low, high, size))	8✔
292
293
294	class XPAdditions:	8✔
295	"""
296	Additional functions missing from both array-api-strict and array-api-extra.
297
298	This class provides wrappers for common array operations such as integration,
299	interpolation, and linear algebra, ensuring compatibility across NumPy, JAX,
300	and array-api-strict backends.
301
302	This is intended as a temporary solution. See https://github.com/glass-dev/glass/issues/645
303	for details.
304	"""
305
306	xp: ModuleType	8✔
307	backend: str	8✔
308
309	def __init__(self, xp: ModuleType) -> None:	8✔
310	"""
311	Initialize XPAdditions with the given array namespace.
312
313	Parameters
314	----------
315	xp : ModuleType
316	The array namespace module.
317	"""
318	self.xp = xp	8✔
319	self.backend = xp.__name__	8✔
320
321	def trapezoid(	8✔
322	self, y: AnyArray, x: AnyArray = None, dx: float = 1.0, axis: int = -1
323	) -> AnyArray:
324	"""
325	Integrate along the given axis using the composite trapezoidal rule.
326
327	Parameters
328	----------
329	y : AnyArray
330	Input array to integrate.
331	x : AnyArray, optional
332	Sample points corresponding to y.
333	dx : float, optional
334	Spacing between sample points.
335	axis : int, optional
336	Axis along which to integrate.
337
338	Returns
339	-------
340	AnyArray
341	Integrated result.
342
343	Raises
344	------
345	NotImplementedError
346	If the array backend is not supported.
347
348	Notes
349	-----
350	See https://github.com/glass-dev/glass/issues/646
351	"""
352	self.backend = self.xp.__name__	8✔
353	if self.backend == "jax.numpy":	8✔
354	import glass.jax # noqa: PLC0415	8✔
355
356	return glass.jax.trapezoid(y, x=x, dx=dx, axis=axis)	8✔
357	if self.backend == "numpy":	8✔
358	return self.xp.trapezoid(y, x=x, dx=dx, axis=axis)	8✔
359	if self.backend == "array_api_strict":	8✔
360	np = import_numpy(self.backend, "trapezoid")	8✔
361
362	# Using design principle of scipy (i.e. copy, use np, copy back)
363	y_np = np.asarray(y, copy=True)	8✔
364	x_np = np.asarray(x, copy=True)	8✔
365	result_np = np.trapezoid(y_np, x_np, dx=dx, axis=axis)	8✔
366	return self.xp.asarray(result_np, copy=True)	8✔
367
368	msg = "the array backend in not supported"	×
369	raise NotImplementedError(msg)	×
370
371	def union1d(self, ar1: AnyArray, ar2: AnyArray) -> AnyArray:	8✔
372	"""
373	Compute the set union of two 1D arrays.
374
375	Parameters
376	----------
377	ar1 : AnyArray
378	First input array.
379	ar2 : AnyArray
380	Second input array.
381
382	Returns
383	-------
384	AnyArray
385	The union of the two arrays.
386
387	Raises
388	------
389	NotImplementedError
390	If the array backend is not supported.
391
392	Notes
393	-----
394	See https://github.com/glass-dev/glass/issues/647
395	"""
396	if self.backend in {"numpy", "jax.numpy"}:	8✔
397	return self.xp.union1d(ar1, ar2)	8✔
398	if self.backend == "array_api_strict":	8✔
399	np = import_numpy(self.backend, "union1d")	8✔
400
401	# Using design principle of scipy (i.e. copy, use np, copy back)
402	ar1_np = np.asarray(ar1, copy=True)	8✔
403	ar2_np = np.asarray(ar2, copy=True)	8✔
404	result_np = np.union1d(ar1_np, ar2_np)	8✔
405	return self.xp.asarray(result_np, copy=True)	8✔
406
407	msg = "the array backend in not supported"	×
408	raise NotImplementedError(msg)	×
409
410	def interp( # noqa: PLR0913	8✔
411	self,
412	x: AnyArray,
413	x_points: AnyArray,
414	y_points: AnyArray,
415	left: float \| None = None,
416	right: float \| None = None,
417	period: float \| None = None,
418	) -> AnyArray:
419	"""
420	One-dimensional linear interpolation for monotonically increasing sample points.
421
422	Parameters
423	----------
424	x : AnyArray
425	The x-coordinates at which to evaluate the interpolated values.
426	x_points : AnyArray
427	The x-coordinates of the data points.
428	y_points : AnyArray
429	The y-coordinates of the data points.
430	left : float \| None, optional
431	Value to return for x < x_points[0].
432	right : float \| None, optional
433	Value to return for x > x_points[-1].
434	period : float \| None, optional
435	Period for periodic interpolation.
436
437	Returns
438	-------
439	AnyArray
440	Interpolated values.
441
442	Raises
443	------
444	NotImplementedError
445	If the array backend is not supported.
446
447	Notes
448	-----
449	See https://github.com/glass-dev/glass/issues/650
450	"""
451	if self.backend in {"numpy", "jax.numpy"}:	8✔
452	return self.xp.interp(	8✔
453	x, x_points, y_points, left=left, right=right, period=period
454	)
455	if self.backend == "array_api_strict":	8✔
456	np = import_numpy(self.backend, "interp")	8✔
457
458	# Using design principle of scipy (i.e. copy, use np, copy back)
459	x_np = np.asarray(x, copy=True)	8✔
460	x_points_np = np.asarray(x_points, copy=True)	8✔
461	y_points_np = np.asarray(y_points, copy=True)	8✔
462	result_np = np.interp(	8✔
463	x_np, x_points_np, y_points_np, left=left, right=right, period=period
464	)
465	return self.xp.asarray(result_np, copy=True)	8✔
466
467	msg = "the array backend in not supported"	×
468	raise NotImplementedError(msg)	×
469
470	def gradient(self, f: AnyArray) -> AnyArray:	8✔
471	"""
472	Return the gradient of an N-dimensional array.
473
474	Parameters
475	----------
476	f : AnyArray
477	Input array.
478
479	Returns
480	-------
481	AnyArray
482	Gradient of the input array.
483
484	Raises
485	------
486	NotImplementedError
487	If the array backend is not supported.
488
489	Notes
490	-----
491	See https://github.com/glass-dev/glass/issues/648
492	"""
493	if self.backend in {"numpy", "jax.numpy"}:	8✔
494	return self.xp.gradient(f)	8✔
495	if self.backend == "array_api_strict":	8✔
496	np = import_numpy(self.backend, "gradient")	8✔
497
498	# Using design principle of scipy (i.e. copy, use np, copy back)
499	f_np = np.asarray(f, copy=True)	8✔
500	result_np = np.gradient(f_np)	8✔
501	return self.xp.asarray(result_np, copy=True)	8✔
502
503	msg = "the array backend in not supported"	×
504	raise NotImplementedError(msg)	×
505
506	def linalg_lstsq(	8✔
507	self, a: AnyArray, b: AnyArray, rcond: float \| None = None
508	) -> tuple[AnyArray, AnyArray, int, AnyArray]:
509	"""
510	Solve a linear least squares problem.
511
512	Parameters
513	----------
514	a : AnyArray
515	Coefficient matrix.
516	b : AnyArray
517	Ordinate or "dependent variable" values.
518	rcond : float \| None, optional
519	Cut-off ratio for small singular values.
520
521	Returns
522	-------
523	x : {(N,), (N, K)} AnyArray
524	Least-squares solution. If b is two-dimensional, the solutions are in the K
525	columns of x.
526
527	residuals : {(1,), (K,), (0,)} AnyArray
528	Sums of squared residuals: Squared Euclidean 2-norm for each column in b - a
529	@ x. If the rank of a is < N or M <= N, this is an empty array. If b is
530	1-dimensional, this is a (1,) shape array. Otherwise the shape is (K,).
531
532	rank : int
533	Rank of matrix a.
534
535	s : (min(M, N),) AnyArray
536	Singular values of a.
537
538	Raises
539	------
540	NotImplementedError
541	If the array backend is not supported.
542
543	Notes
544	-----
545	See https://github.com/glass-dev/glass/issues/649
546	"""
547	if self.backend in {"numpy", "jax.numpy"}:	8✔
548	return self.xp.linalg.lstsq(a, b, rcond=rcond) # type: ignore[no-any-return]	8✔
549	if self.backend == "array_api_strict":	8✔
550	np = import_numpy(self.backend, "linalg.lstsq")	8✔
551
552	# Using design principle of scipy (i.e. copy, use np, copy back)
553	a_np = np.asarray(a, copy=True)	8✔
554	b_np = np.asarray(b, copy=True)	8✔
555	result_np = np.linalg.lstsq(a_np, b_np, rcond=rcond)	8✔
556	return tuple(self.xp.asarray(res, copy=True) for res in result_np)	8✔
557
558	msg = "the array backend in not supported"	×
559	raise NotImplementedError(msg)	×
560
561	def einsum(self, subscripts: str, *operands: AnyArray) -> AnyArray:	8✔
562	"""
563	Evaluate the Einstein summation convention on the operands.
564
565	Parameters
566	----------
567	subscripts : str
568	Specifies the subscripts for summation.
569	*operands : AnyArray
570	Arrays to be summed.
571
572	Returns
573	-------
574	AnyArray
575	Result of the Einstein summation.
576
577	Raises
578	------
579	NotImplementedError
580	If the array backend is not supported.
581
582	Notes
583	-----
584	See https://github.com/glass-dev/glass/issues/657
585	"""
586	if self.backend in {"numpy", "jax.numpy"}:	8✔
587	return self.xp.einsum(subscripts, *operands)	8✔
588	if self.backend == "array_api_strict":	8✔
589	np = import_numpy(self.backend, "einsum")	8✔
590
591	# Using design principle of scipy (i.e. copy, use np, copy back)
592	operands_np = (np.asarray(op, copy=True) for op in operands)	8✔
593	result_np = np.einsum(subscripts, *operands_np)	8✔
594	return self.xp.asarray(result_np, copy=True)	8✔
595
596	msg = "the array backend in not supported"	×
597	raise NotImplementedError(msg)	×
598
599	def apply_along_axis(	8✔
600	self,
601	func1d: Callable[..., Any],
602	axis: int,
603	arr: AnyArray,
604	*args: object,
605	**kwargs: object,
606	) -> AnyArray:
607	"""
608	Apply a function to 1-D slices along the given axis.
609
610	Parameters
611	----------
612	func1d : Callable[..., Any]
613	Function to apply to 1-D slices.
614	axis : int
615	Axis along which to apply the function.
616	arr : AnyArray
617	Input array.
618	*args : object
619	Additional positional arguments to pass to func1d.
620	**kwargs : object
621	Additional keyword arguments to pass to func1d.
622
623	Returns
624	-------
625	AnyArray
626	Result of applying the function along the axis.
627
628	Raises
629	------
630	NotImplementedError
631	If the array backend is not supported.
632
633	Notes
634	-----
635	See https://github.com/glass-dev/glass/issues/651
636
637	"""
638	if self.backend in {"numpy", "jax.numpy"}:	8✔
639	return self.xp.apply_along_axis(func1d, axis, arr, args, *kwargs)	8✔
640	if self.backend == "array_api_strict":	8✔
641	# Import here to prevent users relying on numpy unless in this instance
642	np = import_numpy(self.backend, "apply_along_axis")	8✔
643
644	return self.xp.asarray(	8✔
645	np.apply_along_axis(func1d, axis, arr, args, *kwargs), copy=True
646	)
647
NEW 648	msg = "the array backend in not supported"	×
NEW 649	raise NotImplementedError(msg)	×
650
651	def vectorize(	8✔
652	self,
653	pyfunc: Callable[..., Any],
654	otypes: tuple[type[float]],
655	) -> Callable[..., Any]:
656	"""
657	Returns an object that acts like pyfunc, but takes arrays as input.
658
659	Parameters
660	----------
661	pyfunc : Callable[..., Any]
662	Python function to vectorize.
663	otypes : tuple[type[float]]
664	Output types.
665
666	Returns
667	-------
668	Callable[..., Any]
669	Vectorized function.
670
671	Raises
672	------
673	NotImplementedError
674	If the array backend is not supported.
675
676	Notes
677	-----
678	See https://github.com/glass-dev/glass/issues/671
679	"""
680	if self.backend == "numpy":	8✔
681	return self.xp.vectorize(pyfunc, otypes=otypes) # type: ignore[no-any-return]	8✔
682	if self.backend in {"array_api_strict", "jax.numpy"}:	8✔
683	# Import here to prevent users relying on numpy unless in this instance
684	np = import_numpy(self.backend, "vectorize")	8✔
685
686	return np.vectorize(pyfunc, otypes=otypes) # type: ignore[no-any-return]	8✔
687
NEW 688	msg = "the array backend in not supported"	×
NEW 689	raise NotImplementedError(msg)	×
690
691	def radians(self, deg_arr: AnyArray) -> AnyArray:	8✔
692	"""
693	Convert angles from degrees to radians.
694
695	Parameters
696	----------
697	deg_arr : AnyArray
698	Array of angles in degrees.
699
700	Raises
701	------
702	NotImplementedError
703	If the array backend is not supported.
704
705	Returns
706	-------
707	AnyArray
708	Array of angles in radians.
709	"""
710	if self.backend in {"numpy", "jax.numpy"}:	8✔
711	return self.xp.radians(deg_arr)	8✔
712	if self.backend == "array_api_strict":	8✔
713	np = import_numpy(self.backend, "radians")	8✔
714
715	return self.xp.asarray(np.radians(deg_arr))	8✔
716
NEW 717	msg = "the array backend in not supported"	×
NEW 718	raise NotImplementedError(msg)	×
719
720	def degrees(self, deg_arr: AnyArray) -> AnyArray:	8✔
721	"""
722	Convert angles from radians to degrees.
723
724	Parameters
725	----------
726	deg_arr : AnyArray
727	Array of angles in radians.
728
729	Raises
730	------
731	NotImplementedError
732	If the array backend is not supported.
733
734	Returns
735	-------
736	AnyArray
737	Array of angles in degrees.
738	"""
739	if self.backend in {"numpy", "jax.numpy"}:	8✔
740	return self.xp.degrees(deg_arr)	8✔
741	if self.backend == "array_api_strict":	8✔
742	np = import_numpy(self.backend, "degrees")	8✔
743
744	return self.xp.asarray(np.degrees(deg_arr))	8✔
745
746	msg = "the array backend in not supported"	×
747	raise NotImplementedError(msg)	×

glass-dev / glass / 18593169061

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous