Coverage for src/pyhrp/operators.py: 100%
87 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-10-14 19:39 +0000
1"""The covariance as an operator: the products HRP needs, without the matrix.
3The allocators never need the covariance as an explicit array. Risk parity needs
4quadratic forms on principal blocks, the Schur variant adds cross-block products
5and solves against a principal block, and Ward or single-linkage clustering
6needs the diagonal plus products with the full matrix. This module captures that
7as the structural protocol :class:`CovarianceOperator`, six members of the
8``SymmetricOperator`` of `cvx-linalg <https://github.com/Jebel-Quant/linalg>`_,
9whose dense, factor (diagonal-plus-low-rank) and Gram (returns-matrix) backends
10satisfy it without an adapter. A factor or Gram backend never forms the
11``n x n`` covariance, so the allocation runs in ``O(n k)`` memory.
13- CovarianceOperator: The structural protocol the allocators consume
14- DenseCovariance: The protocol over an explicit matrix (what a DataFrame is wrapped in)
15- as_operator: Turn a covariance DataFrame or an operator into an operator plus asset names
16"""
18from __future__ import annotations
20from collections.abc import Callable, Sequence
21from typing import Any, Protocol, runtime_checkable
23import numpy as np
24import polars as pl
26__all__ = ["CovarianceOperator", "DenseCovariance", "as_operator", "cg_min_norm"]
28# Below this reciprocal condition number a principal block is treated as singular
29# and solved by conjugate gradients for the minimum-norm solution.
30RCOND_MIN = 1e-12
32# Conjugate gradients stop once the residual is this small relative to the right-hand side.
33CG_RTOL = 1e-10
36@runtime_checkable
37class CovarianceOperator(Protocol):
38 """A symmetric covariance reached only through block products and a block solve.
40 Index sets are 1-D integer arrays of positions into ``range(n)``.
41 """
43 @property
44 def n(self) -> int:
45 """Number of assets."""
46 ...
48 @property
49 def diag(self) -> np.ndarray:
50 """The variances, ``diag(Sigma)``, as a length-``n`` vector."""
51 ...
53 def matvec(self, x: np.ndarray) -> np.ndarray:
54 """Return ``Sigma @ x``."""
55 ...
57 def block_matvec(self, rows: Any, cols: Any, v: np.ndarray) -> np.ndarray:
58 """Return ``Sigma[rows, cols] @ v``."""
59 ...
61 def solve_free(self, free: Any, rhs: np.ndarray) -> np.ndarray:
62 """Return ``Sigma[free, free]^{-1} @ rhs``."""
63 ...
65 def rcond_free(self, free: Any) -> float:
66 """Return the reciprocal condition number of ``Sigma[free, free]``."""
67 ...
70class DenseCovariance:
71 """:class:`CovarianceOperator` over an explicit symmetric matrix.
73 This is what a covariance DataFrame is wrapped in, so the DataFrame path and
74 the operator path run the same allocation code.
76 Args:
77 matrix: The ``n x n`` covariance.
79 Examples:
80 >>> import numpy as np
81 >>> from pyhrp.operators import DenseCovariance
82 >>> op = DenseCovariance(np.array([[4.0, 1.0], [1.0, 9.0]]))
83 >>> op.diag.tolist()
84 [4.0, 9.0]
85 >>> op.block_matvec(np.array([1]), np.array([0, 1]), np.array([1.0, 1.0])).tolist()
86 [10.0]
87 """
89 def __init__(self, matrix: np.ndarray) -> None:
90 """Store the matrix as a float64 array."""
91 self._a = np.asarray(matrix, dtype=np.float64)
93 @property
94 def n(self) -> int:
95 """Number of assets."""
96 return int(self._a.shape[0])
98 @property
99 def diag(self) -> np.ndarray:
100 """The variances, ``diag(Sigma)``."""
101 return np.diagonal(self._a).copy()
103 def matvec(self, x: np.ndarray) -> np.ndarray:
104 """Return ``Sigma @ x``."""
105 return np.asarray(self._a @ x)
107 def block_matvec(self, rows: Any, cols: Any, v: np.ndarray) -> np.ndarray:
108 """Return ``Sigma[rows, cols] @ v``."""
109 return np.asarray(self._a[np.ix_(rows, cols)] @ v)
111 def solve_free(self, free: Any, rhs: np.ndarray) -> np.ndarray:
112 """Return ``Sigma[free, free]^{-1} @ rhs`` by an LU solve on the block."""
113 return np.asarray(np.linalg.solve(self._a[np.ix_(free, free)], rhs))
115 def rcond_free(self, free: Any) -> float:
116 """Return the reciprocal 2-norm condition number of ``Sigma[free, free]``."""
117 eig = np.linalg.eigvalsh(self._a[np.ix_(free, free)])
118 if eig[-1] <= 0.0:
119 return 0.0
120 return max(float(eig[0]), 0.0) / float(eig[-1])
123def as_operator(
124 cov: pl.DataFrame | CovarianceOperator, assets: Sequence[str] | None = None
125) -> tuple[CovarianceOperator, list[str]]:
126 """Return the covariance as an operator together with the asset names.
128 A DataFrame is wrapped in :class:`DenseCovariance` and names its own assets
129 (its columns). An operator carries no names, so they come from *assets*, or
130 default to ``"0", "1", ...``.
132 Args:
133 cov: A square covariance DataFrame (columns are assets) or an operator.
134 assets: Asset names for an operator; must match the columns for a DataFrame.
136 Returns:
137 tuple[CovarianceOperator, list[str]]: The operator and the ``n`` asset names.
139 Raises:
140 TypeError: If *cov* is neither a DataFrame nor a :class:`CovarianceOperator`.
141 ValueError: If *assets* does not have one name per asset, or contradicts
142 the DataFrame's columns.
144 Examples:
145 >>> import polars as pl
146 >>> from pyhrp.operators import as_operator
147 >>> op, names = as_operator(pl.DataFrame({"A": [1.0, 0.0], "B": [0.0, 2.0]}))
148 >>> op.n, names
149 (2, ['A', 'B'])
150 """
151 if isinstance(cov, pl.DataFrame):
152 if assets is not None and list(assets) != cov.columns:
153 msg = f"assets {list(assets)} do not match the covariance columns {cov.columns}"
154 raise ValueError(msg)
155 return DenseCovariance(cov.to_numpy()), list(cov.columns)
156 if not isinstance(cov, CovarianceOperator):
157 msg = f"cov must be a polars DataFrame or a CovarianceOperator, got {type(cov).__name__}"
158 raise TypeError(msg)
159 names = [str(i) for i in range(cov.n)] if assets is None else list(assets)
160 if len(names) != cov.n:
161 msg = f"expected {cov.n} asset names, got {len(names)}"
162 raise ValueError(msg)
163 return cov, names
166def block(op: CovarianceOperator, rows: np.ndarray, cols: np.ndarray) -> np.ndarray:
167 """Materialise ``Sigma[rows, cols]`` from ``len(cols)`` products with the identity."""
168 return np.asarray(op.block_matvec(rows, cols, np.eye(cols.size)))
171def solve_block(op: CovarianceOperator, idx: np.ndarray, rhs: np.ndarray) -> np.ndarray:
172 """Solve against the principal block ``Sigma[idx, idx]``, robust to singularity.
174 Covariance blocks of collinear assets, or of more assets than observations
175 (a Gram operator with ``T < n``), are singular. A block the operator reports
176 as singular, or whose solve fails, is solved by conjugate gradients over
177 ``block_matvec`` (see :func:`cg_min_norm`), which never forms the block.
179 So is every block of an operator that can multiply but not solve, one whose
180 ``rcond_free`` or ``solve_free`` raises ``NotImplementedError`` -- such as a
181 cvx-linalg ``SumOperator`` blending the sample covariance with a shrinkage
182 target. The Schur allocator therefore needs no direct solver at all.
183 """
184 try:
185 if op.rcond_free(idx) >= RCOND_MIN:
186 return np.asarray(op.solve_free(idx, rhs))
187 except (np.linalg.LinAlgError, NotImplementedError):
188 pass
189 return cg_min_norm(_block_product(op, idx), rhs)
192def _block_product(op: CovarianceOperator, idx: np.ndarray) -> Callable[[np.ndarray], np.ndarray]:
193 """The product ``v -> Sigma[idx, idx] @ v`` for an iterative solve.
195 ``block_matvec`` gathers the block's storage on every call (the columns of a
196 Gram factor, the rows of factor loadings), which dominates a long CG run. An
197 operator that offers ``restricted(idx)``, as the cvx-linalg backends do, is
198 sliced once and then multiplied with its plain ``matvec``.
199 """
200 restricted = getattr(op, "restricted", None)
201 if callable(restricted):
202 sub = restricted(idx)
203 return lambda v: np.asarray(sub.matvec(v))
204 return lambda v: np.asarray(op.block_matvec(idx, idx, v))
207def cg_min_norm(apply: Callable[[np.ndarray], np.ndarray], b: np.ndarray, rtol: float = CG_RTOL) -> np.ndarray:
208 """Minimum-norm solution of ``A x = b`` for a symmetric PSD ``A``, by conjugate gradients.
210 The Schur right-hand sides are ``Sigma[R, L] w_L``, and for a PSD ``Sigma``
211 the columns of ``Sigma[R, L]`` lie in the range of ``Sigma[R, R]``, so the
212 system is consistent even when the block is singular. Started from zero, CG
213 then stays in that range and converges to the minimum-norm solution, the
214 one ``lstsq`` returns, using only products with ``A``.
216 Iteration stops when the residual falls below ``rtol * |b|``, when a search
217 direction has no curvature (it lies in the null space, so the residual left
218 is round-off outside the range), or after ``2 * len(b)`` steps. In exact
219 arithmetic CG terminates within ``rank(A)`` steps.
221 Args:
222 apply (Callable[[np.ndarray], np.ndarray]): The product ``v -> A @ v``.
223 b (np.ndarray): The right-hand side.
224 rtol (float): Relative residual tolerance.
226 Returns:
227 np.ndarray: The approximate minimum-norm solution.
229 Examples:
230 >>> import numpy as np
231 >>> from pyhrp.operators import cg_min_norm
232 >>> a = np.array([[1.0, 2.0], [2.0, 4.0]]) # singular
233 >>> x = cg_min_norm(lambda v: a @ v, np.array([1.0, 2.0]))
234 >>> np.allclose(x, np.linalg.pinv(a) @ [1.0, 2.0])
235 True
236 """
237 x = np.zeros_like(b, dtype=np.float64)
238 r = np.asarray(b, dtype=np.float64).copy()
239 p = r.copy()
240 rr = float(r @ r)
241 stop = (rtol**2) * rr
242 for _ in range(2 * b.size):
243 if rr <= stop:
244 break
245 ap = apply(p)
246 curvature = float(p @ ap)
247 if curvature <= 0.0:
248 break
249 alpha = rr / curvature
250 x += alpha * p
251 r -= alpha * ap
252 rr_next = float(r @ r)
253 p = r + (rr_next / rr) * p
254 rr = rr_next
255 return x