20337451080

Committed 18 Dec 2025 12:47PM UTC coverage: 41.778% (-51.9%) from 93.727%

Build # 20337451080

Build Type

Pull #917

github

Committed by

web-flow

Commit Message

Merge e8c4714e0 into d0ff6e647

Pull Request Pull Request #917: gh-915: Fail tests if array backend not found

Run Details

82 of 216 branches covered (37.96%)

Branch coverage included in aggregate %.

604 of 1426 relevant lines covered (42.36%)

0.42 hits per line

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

56.1

/glass/arraytools.py

"""Module for array utilities."""

from __future__ import annotations

import itertools
from functools import partial
from typing import TYPE_CHECKING

import array_api_compat
import array_api_extra as xpx

import glass._array_api_utils as _utils

if TYPE_CHECKING:
    from types import ModuleType
    from typing import Unpack

    from glass._types import AnyArray, FloatArray, IntArray


def broadcast_first(
    *arrays: FloatArray,
) -> tuple[FloatArray, ...]:
    """
    Broadcast arrays, treating the first axis as common.

    Parameters
    ----------
    arrays
        The arrays to broadcast.

    Returns
    -------
        The broadcasted arrays.

    """
    xp = array_api_compat.array_namespace(*arrays, use_compat=False)

    arrays = tuple(xp.moveaxis(a, 0, -1) if a.ndim else a for a in arrays)
    arrays = xp.broadcast_arrays(*arrays)
    return tuple(xp.moveaxis(a, -1, 0) if a.ndim else a for a in arrays)


def broadcast_leading_axes(
    *args: tuple[
        float | FloatArray,
        int,
    ],
    xp: ModuleType | None = None,
) -> tuple[
    tuple[int, ...],
    Unpack[tuple[FloatArray, ...]],
]:
    """
    Broadcast all but the last N axes.

    Parameters
    ----------
    args
        The arrays and the number of axes to keep.

    Returns
    -------
        The shape of the broadcast dimensions, and all input arrays
        with leading axes matching that shape.

    Examples
    --------
    Broadcast all dimensions of ``a``, all except the last dimension of
    ``b``, and all except the last two dimensions of ``c``.

    >>> import numpy as np
    >>> a = 0
    >>> b = np.zeros((4, 10))
    >>> c = np.zeros((3, 1, 5, 6))
    >>> dims, a, b, c = broadcast_leading_axes((a, 0), (b, 1), (c, 2))
    >>> dims
    (3, 4)
    >>> a.shape
    (3, 4)
    >>> b.shape
    (3, 4, 10)
    >>> c.shape
    (3, 4, 5, 6)

    """
    if xp is None:
        xp = array_api_compat.array_namespace(
            *[arg[0] for arg in args],
            use_compat=False,
        )

    shapes, trails = [], []
    for a, n in args:
        a_arr = xp.asarray(a)
        s = a_arr.shape
        i = len(s) - n
        shapes.append(s[:i])
        trails.append(s[i:])
    dims = xpx.broadcast_shapes(*shapes)
    arrs = (
        xp.broadcast_to(xp.asarray(a), dims + t)
        for (a, _), t in zip(args, trails, strict=False)
    )
    return (dims, *arrs)


def ndinterp(  # noqa: PLR0913
    x: float | FloatArray,
    xq: FloatArray,
    fq: FloatArray,
    axis: int = -1,
    left: float | None = None,
    right: float | None = None,
    period: float | None = None,
) -> FloatArray:
    """
    Interpolate multi-dimensional array over axis.

    Parameters
    ----------
    x
        The x-coordinates.
    xq
        The x-coordinates of the data points.
    fq
        The function values corresponding to the x-coordinates in *xq*.
    axis
        The axis to interpolate over.
    left
        The value to return for x < xq[0].
    right
        The value to return for x > xq[-1].
    period
        The period of the function, used for interpolating periodic data.

    Returns
    -------
        The interpolated array.

    """
    xp = array_api_compat.array_namespace(x, xq, fq, use_compat=False)
    uxpx = _utils.XPAdditions(xp)

    return uxpx.apply_along_axis(
        partial(uxpx.interp, x, xq),
        axis,
        fq,
        left=left,
        right=right,
        period=period,
    )


def trapezoid_product(
    f: tuple[FloatArray, FloatArray],
    *ff: tuple[FloatArray, FloatArray],
    axis: int = -1,
) -> float | FloatArray:
    """
    Trapezoidal rule for a product of functions.

    Parameters
    ----------
    f
        The first function.
    ff
        The other functions.
    axis
        The axis along which to integrate.

    Returns
    -------
        The integral of the product of the functions.

    """
    # Flatten ff into a 1D tuple of all ff inputs and then expand to get the namespace
    xp = array_api_compat.array_namespace(
        *f,
        *tuple(itertools.chain(*ff)),
        use_compat=False,
    )
    uxpx = _utils.XPAdditions(xp)

    x: FloatArray
    x, _ = f
    for x_, _ in ff:
        x = uxpx.union1d(
            x[(x >= x_[0]) & (x <= x_[-1])],
            x_[(x_ >= x[0]) & (x_ <= x[-1])],
        )
    y = uxpx.interp(x, *f)
    for f_ in ff:
        y *= uxpx.interp(x, *f_)
    return uxpx.trapezoid(y, x, axis=axis)


def cumulative_trapezoid(
    f: IntArray | FloatArray,
    x: IntArray | FloatArray,
) -> AnyArray:
    """
    Cumulative trapezoidal rule along last axis.

    Parameters
    ----------
    f
        The function values.
    x
        The x-coordinates.

    Returns
    -------
        The cumulative integral of the function.

    """
    xp = array_api_compat.array_namespace(f, x, use_compat=False)

    f = xp.asarray(f, dtype=xp.float64)
    x = xp.asarray(x, dtype=xp.float64)

    # Compute the cumulative trapezoid without mutating any arrays
    return xp.cumulative_sum(
        (f[..., 1:] + f[..., :-1]) * 0.5 * xp.diff(x),
        axis=-1,
        include_initial=True,
    )

1	"""Module for array utilities."""
2
3	from __future__ import annotations
4
5	import itertools
6	from functools import partial
7	from typing import TYPE_CHECKING
8
9	import array_api_compat
10	import array_api_extra as xpx
11
12	import glass._array_api_utils as _utils
13
14	if TYPE_CHECKING:
15	from types import ModuleType
16	from typing import Unpack
17
18	from glass._types import AnyArray, FloatArray, IntArray
19
20
21	def broadcast_first(	1✔
22	*arrays: FloatArray,
23	) -> tuple[FloatArray, ...]:
24	"""
25	Broadcast arrays, treating the first axis as common.
26
27	Parameters
28	----------
29	arrays
30	The arrays to broadcast.
31
32	Returns
33	-------
34	The broadcasted arrays.
35
36	"""
37	xp = array_api_compat.array_namespace(*arrays, use_compat=False)	×
38
39	arrays = tuple(xp.moveaxis(a, 0, -1) if a.ndim else a for a in arrays)	×
40	arrays = xp.broadcast_arrays(*arrays)	×
41	return tuple(xp.moveaxis(a, -1, 0) if a.ndim else a for a in arrays)	×
42
43
44	def broadcast_leading_axes(	1✔
45	*args: tuple[
46	float \| FloatArray,
47	int,
48	],
49	xp: ModuleType \| None = None,
50	) -> tuple[
51	tuple[int, ...],
52	Unpack[tuple[FloatArray, ...]],
53	]:
54	"""
55	Broadcast all but the last N axes.
56
57	Parameters
58	----------
59	args
60	The arrays and the number of axes to keep.
61
62	Returns
63	-------
64	The shape of the broadcast dimensions, and all input arrays
65	with leading axes matching that shape.
66
67	Examples
68	--------
69	Broadcast all dimensions of ``a``, all except the last dimension of
70	``b``, and all except the last two dimensions of ``c``.
71
72	>>> import numpy as np
73	>>> a = 0
74	>>> b = np.zeros((4, 10))
75	>>> c = np.zeros((3, 1, 5, 6))
76	>>> dims, a, b, c = broadcast_leading_axes((a, 0), (b, 1), (c, 2))
77	>>> dims
78	(3, 4)
79	>>> a.shape
80	(3, 4)
81	>>> b.shape
82	(3, 4, 10)
83	>>> c.shape
84	(3, 4, 5, 6)
85
86	"""
87	if xp is None:	1✔
88	xp = array_api_compat.array_namespace(	1✔
89	*[arg[0] for arg in args],
90	use_compat=False,
91	)
92
93	shapes, trails = [], []	1✔
94	for a, n in args:	1✔
95	a_arr = xp.asarray(a)	1✔
96	s = a_arr.shape	1✔
97	i = len(s) - n	1✔
98	shapes.append(s[:i])	1✔
99	trails.append(s[i:])	1✔
100	dims = xpx.broadcast_shapes(*shapes)	1✔
101	arrs = (	1✔
102	xp.broadcast_to(xp.asarray(a), dims + t)
103	for (a, _), t in zip(args, trails, strict=False)
104	)
105	return (dims, *arrs)	1✔
106
107
108	def ndinterp( # noqa: PLR0913	1✔
109	x: float \| FloatArray,
110	xq: FloatArray,
111	fq: FloatArray,
112	axis: int = -1,
113	left: float \| None = None,
114	right: float \| None = None,
115	period: float \| None = None,
116	) -> FloatArray:
117	"""
118	Interpolate multi-dimensional array over axis.
119
120	Parameters
121	----------
122	x
123	The x-coordinates.
124	xq
125	The x-coordinates of the data points.
126	fq
127	The function values corresponding to the x-coordinates in xq.
128	axis
129	The axis to interpolate over.
130	left
131	The value to return for x < xq[0].
132	right
133	The value to return for x > xq[-1].
134	period
135	The period of the function, used for interpolating periodic data.
136
137	Returns
138	-------
139	The interpolated array.
140
141	"""
142	xp = array_api_compat.array_namespace(x, xq, fq, use_compat=False)	×
143	uxpx = _utils.XPAdditions(xp)	×
144
145	return uxpx.apply_along_axis(	×
146	partial(uxpx.interp, x, xq),
147	axis,
148	fq,
149	left=left,
150	right=right,
151	period=period,
152	)
153
154
155	def trapezoid_product(	1✔
156	f: tuple[FloatArray, FloatArray],
157	*ff: tuple[FloatArray, FloatArray],
158	axis: int = -1,
159	) -> float \| FloatArray:
160	"""
161	Trapezoidal rule for a product of functions.
162
163	Parameters
164	----------
165	f
166	The first function.
167	ff
168	The other functions.
169	axis
170	The axis along which to integrate.
171
172	Returns
173	-------
174	The integral of the product of the functions.
175
176	"""
177	# Flatten ff into a 1D tuple of all ff inputs and then expand to get the namespace
178	xp = array_api_compat.array_namespace(	×
179	*f,
180	tuple(itertools.chain(ff)),
181	use_compat=False,
182	)
183	uxpx = _utils.XPAdditions(xp)	×
184
185	x: FloatArray
186	x, _ = f	×
187	for x_, _ in ff:	×
188	x = uxpx.union1d(	×
189	x[(x >= x_[0]) & (x <= x_[-1])],
190	x_[(x_ >= x[0]) & (x_ <= x[-1])],
191	)
192	y = uxpx.interp(x, *f)	×
193	for f_ in ff:	×
194	y = uxpx.interp(x, f_)	×
195	return uxpx.trapezoid(y, x, axis=axis)	×
196
197
198	def cumulative_trapezoid(	1✔
199	f: IntArray \| FloatArray,
200	x: IntArray \| FloatArray,
201	) -> AnyArray:
202	"""
203	Cumulative trapezoidal rule along last axis.
204
205	Parameters
206	----------
207	f
208	The function values.
209	x
210	The x-coordinates.
211
212	Returns
213	-------
214	The cumulative integral of the function.
215
216	"""
217	xp = array_api_compat.array_namespace(f, x, use_compat=False)	1✔
218
219	f = xp.asarray(f, dtype=xp.float64)	1✔
220	x = xp.asarray(x, dtype=xp.float64)	1✔
221
222	# Compute the cumulative trapezoid without mutating any arrays
223	return xp.cumulative_sum(	1✔
224	(f[..., 1:] + f[..., :-1]) * 0.5 * xp.diff(x),
225	axis=-1,
226	include_initial=True,
227	)

glass-dev / glass / 20337451080

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