27141040624

Committed 08 Jun 2026 01:29PM UTC coverage: 95.591% (-0.2%) from 95.749%

Build # 27141040624

Build Type

Pull #1114

github

Committed by

web-flow

Commit Message

Merge f0f9a9b3e into 3416cae3d

Pull Request Pull Request #1114: gh-1112: distribute redshifts over shells

Coverage Stats

208 of 210 branches covered (99.05%)

Branch coverage included in aggregate %.

27 of 28 new or added lines in 3 files covered. (96.43%)

3 existing lines in 1 file now uncovered.

1418 of 1491 relevant lines covered (95.1%)

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

glass-dev / glass / 27141040624

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