18598858404

Committed 17 Oct 2025 04:31PM UTC coverage: 94.23% (-0.2%) from 94.467%

Build # 18598858404

Build Type

Pull #677

github

Committed by

web-flow

Commit Message

Merge a52607df1 into c5b0599c8

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 %.

123 of 132 new or added lines in 5 files covered. (93.18%)

20 existing lines in 2 files now uncovered.

1432 of 1530 relevant lines covered (93.59%)

7.47 hits per line

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

86.03

/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

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

    import numpy as np
    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 | AArray
    ComplexArray: TypeAlias = NDArray[np.complex128] | JAXArray | AArray
    DoubleArray: TypeAlias = NDArray[np.double] | JAXArray | AArray
    FloatArray: TypeAlias = NDArray[np.float64] | JAXArray | AArray
    IntArray: TypeAlias = NDArray[np.int_] | JAXArray | AArray


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  # noqa: ICN001, PLC0415

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


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.
    """
    xp = get_namespace(array)
    backend = xp.__name__

    if backend == "jax.numpy":
        import glass.jax  # noqa: PLC0415

        return glass.jax.Generator(seed=42)

    if backend == "numpy":
        return xp.random.default_rng()  # type: ignore[no-any-return]

    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__ = ("axp", "nxp", "rng")

    def __init__(
        self,
        seed: int | bool | AArray | 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.
        """
        import array_api_strict  # noqa: PLC0415
        import numpy as np  # noqa: PLC0415

        self.axp = array_api_strict
        self.nxp = np
        self.rng = self.nxp.random.default_rng(seed=seed)

    def random(
        self,
        size: Size = None,
        dtype: DTypeLike | None = None,
        out: AArray | 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.
        """
        dtype = dtype if dtype is not None else self.nxp.float64
        return self.axp.asarray(self.rng.random(size, dtype, out))  # type: ignore[arg-type]

    def normal(
        self,
        loc: float | AArray = 0.0,
        scale: float | AArray = 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 self.axp.asarray(self.rng.normal(loc, scale, size))

    def poisson(self, lam: float | AArray, 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 self.axp.asarray(self.rng.poisson(lam, size))

    def standard_normal(
        self,
        size: Size = None,
        dtype: DTypeLike | None = None,
        out: AArray | 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.
        """
        dtype = dtype if dtype is not None else self.nxp.float64
        return self.axp.asarray(self.rng.standard_normal(size, dtype, out))  # type: ignore[arg-type]

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

glass-dev / glass / 18598858404

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