21400098791

Committed 27 Jan 2026 02:01PM UTC coverage: 93.959% (-0.2%) from 94.174%

Build # 21400098791

Build Type

push

github

Committed by

web-flow

Commit Message

gh-969: determine `xp` when function called in XPAdditions (#994)

Run Details

211 of 213 branches covered (99.06%)

Branch coverage included in aggregate %.

78 of 78 new or added lines in 5 files covered. (100.0%)

3 existing lines in 1 file now uncovered.

1360 of 1459 relevant lines covered (93.21%)

5.07 hits per line

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

96.77

/glass/algorithm.py

"""Module for algorithms."""

from __future__ import annotations

import warnings
from typing import TYPE_CHECKING

import array_api_compat
import array_api_extra as xpx

if TYPE_CHECKING:
    from glass._types import FloatArray


def nnls(
    a: FloatArray,
    b: FloatArray,
    *,
    tol: float = 0.0,
    maxiter: int | None = None,
) -> FloatArray:
    """
    Compute a non-negative least squares solution.

    Implementation of the algorithm due to [Lawson95]_ as described by
    [Bro97]_.

    Parameters
    ----------
    a
        The matrix.
    b
        The vector.
    tol
        The tolerance for convergence.
    maxiter
        The maximum number of iterations.

    Returns
    -------
        The non-negative least squares solution.

    Raises
    ------
    ValueError
        If ``a`` is not a matrix.
    ValueError
        If ``b`` is not a vector.
    ValueError
        If the shapes of ``a`` and ``b`` do not match.

    """
    xp = array_api_compat.array_namespace(a, b, use_compat=False)

    a = xp.asarray(a)
    b = xp.asarray(b)

    if a.ndim != 2:
        msg = "input `a` is not a matrix"
        raise ValueError(msg)
    if b.ndim != 1:
        msg = "input `b` is not a vector"
        raise ValueError(msg)
    if a.shape[0] != b.shape[0]:
        msg = "the shapes of `a` and `b` do not match"
        raise ValueError(msg)

    _, n = a.shape

    if maxiter is None:
        maxiter = 3 * n

    index = xp.arange(n)
    q = xp.full(n, fill_value=False)
    x = xp.zeros(n)
    for _ in range(maxiter):
        if xp.all(q):
            break
        # The sum product over the last axis of arg1 and the second-to-last axis of arg2
        w = xp.sum((b - a @ x)[..., None] * a, axis=-2)

        m = int(index[~q][xp.argmax(w[~q])])
        if w[m] <= tol:
            break
        q = xpx.at(q)[m].set(True)
        while True:
            # Use `xp.task`` here instead of `a[:,q]` to mask the inner arrays, because
            # array-api requires a masking index to be the sole index, which would
            # return a 1-D array. However, we want to maintain the shape of `a`,
            # i.e. `[[a11],[a12],...]` rather than `[a11,a12,...]`
            aq = xp.take(a, xp.nonzero(q)[0], axis=1)
            xq = x[q]
            sq = xp.linalg.solve(aq.T @ aq, b @ aq)
            t = sq <= 0
            if not xp.any(t):
                break
            alpha = -xp.min(xq[t] / (xq[t] - sq[t]))
            x = xpx.at(x)[q].add(alpha * (sq - xq))
            q = xpx.at(q)[x <= 0].set(False)
        x = xpx.at(x)[q].set(sq)
        x = xpx.at(x)[~q].set(0)
    return x


def cov_clip(
    cov: FloatArray,
    rtol: float | None = None,
) -> FloatArray:
    """
    Covariance matrix from clipping non-positive eigenvalues.

    The relative tolerance *rtol* is defined as for
    :func:`~array_api.linalg.matrix_rank`.

    Parameters
    ----------
    cov
        A symmetric matrix (or a stack of matrices).
    rtol
        An optional relative tolerance for eigenvalues to be considered
        positive.

    Returns
    -------
        Covariance matrix with negative eigenvalues clipped.

    """
    xp = cov.__array_namespace__()

    # Hermitian eigendecomposition
    w, v = xp.linalg.eigh(cov)

    # get tolerance if not given
    if rtol is None:
        rtol = max(v.shape[-2], v.shape[-1]) * xp.finfo(w.dtype).eps

    # clip negative diagonal values
    w = xp.clip(w, rtol * xp.max(w, axis=-1, keepdims=True), None)

    # put matrix back together
    # enforce symmetry
    v = xp.sqrt(w[..., None, :]) * v
    return xp.matmul(v, xp.matrix_transpose(v))


def nearcorr(
    a: FloatArray,
    *,
    tol: float | None = None,
    niter: int = 100,
) -> FloatArray:
    """
    Compute the nearest correlation matrix.

    Returns the nearest correlation matrix using the alternating
    projections algorithm of [Higham02]_.

    Parameters
    ----------
    a
        Square matrix (or a stack of square matrices).
    tol
        Tolerance for convergence. Default is dimension times machine
        epsilon.
    niter
        Maximum number of iterations.

    Returns
    -------
        Nearest correlation matrix.

    """
    xp = a.__array_namespace__()

    # shorthand for Frobenius norm
    frob = xp.linalg.matrix_norm

    # get size of the covariance matrix and flatten leading dimensions
    *dim, m, n = a.shape
    if m != n:
        msg = "non-square matrix"
        raise ValueError(msg)

    # default tolerance
    if tol is None:
        tol = n * xp.finfo(a.dtype).eps

    # current result, flatten leading dimensions
    y = xp.reshape(a, (-1, n, n))

    # initial correction is zero
    ds = xp.zeros_like(a)

    # store identity matrix
    diag = xp.eye(n)

    # find the nearest correlation matrix
    for _ in range(niter):
        # apply Dykstra's correction to current result
        r = y - ds

        # project onto positive semi-definite matrices
        x = cov_clip(r)

        # compute Dykstra's correction
        ds = x - r

        # project onto matrices with unit diagonal
        y = (1 - diag) * x + diag

        # check for convergence
        if xp.all(frob(y - x) <= tol * frob(y)):
            break

    else:
        # nearest correlation matrix was not found
        warnings.warn(
            f"Nearest correlation matrix not found in {niter} iterations. "
            "The result may be invalid. Please run with a larger `niter` value, "
            "or run the function again on the returned result.",
            stacklevel=2,
        )

    # return result in original shape
    return xp.reshape(y, (*dim, n, n))


def cov_nearest(
    cov: FloatArray,
    tol: float | None = None,
    niter: int = 100,
) -> FloatArray:
    """
    Covariance matrix from nearest correlation matrix.

    Divides *cov* along rows and columns by the square root of the
    diagonal, then computes the nearest valid correlation matrix using
    :func:`glass.nearcorr`, before scaling rows and columns back.  The
    diagonal of the input is hence unchanged.

    Parameters
    ----------
    cov
        A square matrix (or a stack of matrices).
    tol
        Tolerance for convergence, see :func:`glass.nearcorr`.
    niter
        Maximum number of iterations.

    Returns
    -------
        Covariance matrix from nearest correlation matrix.

    """
    xp = cov.__array_namespace__()

    # get the diagonal
    diag = xp.linalg.diagonal(cov)

    # cannot fix negative diagonal
    if xp.any(diag < 0):
        msg = "negative values on the diagonal"
        raise ValueError(msg)

    # store the normalisation of the matrix
    norm = xp.sqrt(diag)
    norm = norm[..., None, :] * norm[..., :, None]

    # find nearest correlation matrix
    corr = cov / xp.where(norm > 0, norm, 1.0)
    return nearcorr(corr, niter=niter, tol=tol) * norm

1	"""Module for algorithms."""
2
3	from __future__ import annotations
4
5	import warnings
6	from typing import TYPE_CHECKING
7
8	import array_api_compat
9	import array_api_extra as xpx
10
11	if TYPE_CHECKING:
12	from glass._types import FloatArray
13
14
15	def nnls(	6✔
16	a: FloatArray,
17	b: FloatArray,
18	*,
19	tol: float = 0.0,
20	maxiter: int \| None = None,
21	) -> FloatArray:
22	"""
23	Compute a non-negative least squares solution.
24
25	Implementation of the algorithm due to [Lawson95]_ as described by
26	[Bro97]_.
27
28	Parameters
29	----------
30	a
31	The matrix.
32	b
33	The vector.
34	tol
35	The tolerance for convergence.
36	maxiter
37	The maximum number of iterations.
38
39	Returns
40	-------
41	The non-negative least squares solution.
42
43	Raises
44	------
45	ValueError
46	If ``a`` is not a matrix.
47	ValueError
48	If ``b`` is not a vector.
49	ValueError
50	If the shapes of ``a`` and ``b`` do not match.
51
52	"""
53	xp = array_api_compat.array_namespace(a, b, use_compat=False)	5✔
54
55	a = xp.asarray(a)	5✔
56	b = xp.asarray(b)	5✔
57
58	if a.ndim != 2:	5✔
59	msg = "input `a` is not a matrix"	5✔
60	raise ValueError(msg)	5✔
61	if b.ndim != 1:	5✔
62	msg = "input `b` is not a vector"	5✔
63	raise ValueError(msg)	5✔
64	if a.shape[0] != b.shape[0]:	5✔
65	msg = "the shapes of `a` and `b` do not match"	5✔
66	raise ValueError(msg)	5✔
67
68	_, n = a.shape	5✔
69
70	if maxiter is None:	5✔
71	maxiter = 3 * n	5✔
72
73	index = xp.arange(n)	5✔
74	q = xp.full(n, fill_value=False)	5✔
75	x = xp.zeros(n)	5✔
76	for _ in range(maxiter):	5✔
77	if xp.all(q):	5✔
78	break	5✔
79	# The sum product over the last axis of arg1 and the second-to-last axis of arg2
80	w = xp.sum((b - a @ x)[..., None] * a, axis=-2)	5✔
81
82	m = int(index[~q][xp.argmax(w[~q])])	5✔
83	if w[m] <= tol:	5✔
84	break	5✔
85	q = xpx.at(q)[m].set(True)	5✔
86	while True:	5✔
87	# Use `xp.task`` here instead of `a[:,q]` to mask the inner arrays, because
88	# array-api requires a masking index to be the sole index, which would
89	# return a 1-D array. However, we want to maintain the shape of `a`,
90	# i.e. `[[a11],[a12],...]` rather than `[a11,a12,...]`
91	aq = xp.take(a, xp.nonzero(q)[0], axis=1)	5✔
92	xq = x[q]	5✔
93	sq = xp.linalg.solve(aq.T @ aq, b @ aq)	5✔
94	t = sq <= 0	5✔
95	if not xp.any(t):	5✔
96	break	5✔
UNCOV 97	alpha = -xp.min(xq[t] / (xq[t] - sq[t]))	×
UNCOV 98	x = xpx.at(x)[q].add(alpha * (sq - xq))	×
UNCOV 99	q = xpx.at(q)[x <= 0].set(False)	×
100	x = xpx.at(x)[q].set(sq)	5✔
101	x = xpx.at(x)[~q].set(0)	5✔
102	return x	5✔
103
104
105	def cov_clip(	6✔
106	cov: FloatArray,
107	rtol: float \| None = None,
108	) -> FloatArray:
109	"""
110	Covariance matrix from clipping non-positive eigenvalues.
111
112	The relative tolerance rtol is defined as for
113	:func:`~array_api.linalg.matrix_rank`.
114
115	Parameters
116	----------
117	cov
118	A symmetric matrix (or a stack of matrices).
119	rtol
120	An optional relative tolerance for eigenvalues to be considered
121	positive.
122
123	Returns
124	-------
125	Covariance matrix with negative eigenvalues clipped.
126
127	"""
128	xp = cov.__array_namespace__()	5✔
129
130	# Hermitian eigendecomposition
131	w, v = xp.linalg.eigh(cov)	5✔
132
133	# get tolerance if not given
134	if rtol is None:	5✔
135	rtol = max(v.shape[-2], v.shape[-1]) * xp.finfo(w.dtype).eps	5✔
136
137	# clip negative diagonal values
138	w = xp.clip(w, rtol * xp.max(w, axis=-1, keepdims=True), None)	5✔
139
140	# put matrix back together
141	# enforce symmetry
142	v = xp.sqrt(w[..., None, :]) * v	5✔
143	return xp.matmul(v, xp.matrix_transpose(v))	5✔
144
145
146	def nearcorr(	6✔
147	a: FloatArray,
148	*,
149	tol: float \| None = None,
150	niter: int = 100,
151	) -> FloatArray:
152	"""
153	Compute the nearest correlation matrix.
154
155	Returns the nearest correlation matrix using the alternating
156	projections algorithm of [Higham02]_.
157
158	Parameters
159	----------
160	a
161	Square matrix (or a stack of square matrices).
162	tol
163	Tolerance for convergence. Default is dimension times machine
164	epsilon.
165	niter
166	Maximum number of iterations.
167
168	Returns
169	-------
170	Nearest correlation matrix.
171
172	"""
173	xp = a.__array_namespace__()	5✔
174
175	# shorthand for Frobenius norm
176	frob = xp.linalg.matrix_norm	5✔
177
178	# get size of the covariance matrix and flatten leading dimensions
179	*dim, m, n = a.shape	5✔
180	if m != n:	5✔
181	msg = "non-square matrix"	5✔
182	raise ValueError(msg)	5✔
183
184	# default tolerance
185	if tol is None:	5✔
186	tol = n * xp.finfo(a.dtype).eps	5✔
187
188	# current result, flatten leading dimensions
189	y = xp.reshape(a, (-1, n, n))	5✔
190
191	# initial correction is zero
192	ds = xp.zeros_like(a)	5✔
193
194	# store identity matrix
195	diag = xp.eye(n)	5✔
196
197	# find the nearest correlation matrix
198	for _ in range(niter):	5✔
199	# apply Dykstra's correction to current result
200	r = y - ds	5✔
201
202	# project onto positive semi-definite matrices
203	x = cov_clip(r)	5✔
204
205	# compute Dykstra's correction
206	ds = x - r	5✔
207
208	# project onto matrices with unit diagonal
209	y = (1 - diag) * x + diag	5✔
210
211	# check for convergence
212	if xp.all(frob(y - x) <= tol * frob(y)):	5✔
213	break	5✔
214
215	else:
216	# nearest correlation matrix was not found
217	warnings.warn(	5✔
218	f"Nearest correlation matrix not found in {niter} iterations. "
219	"The result may be invalid. Please run with a larger `niter` value, "
220	"or run the function again on the returned result.",
221	stacklevel=2,
222	)
223
224	# return result in original shape
225	return xp.reshape(y, (*dim, n, n))	5✔
226
227
228	def cov_nearest(	6✔
229	cov: FloatArray,
230	tol: float \| None = None,
231	niter: int = 100,
232	) -> FloatArray:
233	"""
234	Covariance matrix from nearest correlation matrix.
235
236	Divides cov along rows and columns by the square root of the
237	diagonal, then computes the nearest valid correlation matrix using
238	:func:`glass.nearcorr`, before scaling rows and columns back. The
239	diagonal of the input is hence unchanged.
240
241	Parameters
242	----------
243	cov
244	A square matrix (or a stack of matrices).
245	tol
246	Tolerance for convergence, see :func:`glass.nearcorr`.
247	niter
248	Maximum number of iterations.
249
250	Returns
251	-------
252	Covariance matrix from nearest correlation matrix.
253
254	"""
255	xp = cov.__array_namespace__()	5✔
256
257	# get the diagonal
258	diag = xp.linalg.diagonal(cov)	5✔
259
260	# cannot fix negative diagonal
261	if xp.any(diag < 0):	5✔
262	msg = "negative values on the diagonal"	5✔
263	raise ValueError(msg)	5✔
264
265	# store the normalisation of the matrix
266	norm = xp.sqrt(diag)	5✔
267	norm = norm[..., None, :] * norm[..., :, None]	5✔
268
269	# find nearest correlation matrix
270	corr = cov / xp.where(norm > 0, norm, 1.0)	5✔
271	return nearcorr(corr, niter=niter, tol=tol) * norm	5✔

glass-dev / glass / 21400098791

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