Guard compressed kernel allocations

Author: skilledwolf

README.md

@@ -98,6 +98,11 @@ Calling ``get_exchange_kernels_compressed(select=None)`` still builds the
 canonical symmetry-reduced list, so it avoids 5D materialization but not the
 underlying O(``nmax^4``) output scaling.
 
+The public compressed API also guards the numeric ``(nG, n_select)`` values
+array via ``compressed_limit_bytes`` (default 512 MiB). If you omit ``select``,
+``canonical_select_max_entries`` separately guards the number of canonical
+representatives before any numeric backend work begins.
+
 For the ``'laguerre'`` backend, dense Gauss-Legendre work tables are guarded by
 ``workspace_limit_bytes`` (default 512 MiB). Pass ``workspace_limit_bytes=None``
 to disable that backend-level guard.

src/quantumhall_matrixelements/__init__.py

@@ -19,11 +19,13 @@
 from numpy.typing import NDArray
 
 from ._materialize import (
+    DEFAULT_COMPRESSED_LIMIT_BYTES,
     DEFAULT_FULL_TENSOR_LIMIT_BYTES,
+    guard_compressed_values_allocation,
     guard_full_tensor_materialization,
     materialize_full_tensor,
 )
-from ._select import DEFAULT_CANONICAL_SELECT_MAX_ENTRIES
+from ._select import DEFAULT_CANONICAL_SELECT_MAX_ENTRIES, estimate_canonical_select_size
 from .diagnostic import get_exchange_kernels_opposite_field, get_form_factors_opposite_field
 from .exchange_hankel import get_exchange_kernels_hankel
 from .exchange_laguerre import (
@@ -144,6 +146,7 @@ def get_exchange_kernels_compressed(
     method: str | None = None,
     select: Iterable[Quad] | None = None,
     canonical_select_max_entries: int | None = DEFAULT_CANONICAL_SELECT_MAX_ENTRIES,
+    compressed_limit_bytes: float | int | None = DEFAULT_COMPRESSED_LIMIT_BYTES,
     **kwargs: Any,
 ) -> tuple[ComplexArray, list[Quad]]:
     """Return the compressed exchange-kernel representation ``(values, select_list)``.
@@ -153,8 +156,9 @@ def get_exchange_kernels_compressed(
 
     If ``select`` is omitted, the backend still constructs the canonical
     symmetry-reduced list, so the returned representation remains O(``nmax^4``)
-    in the number of stored entries. Pass an explicit ``select=...`` to compute
-    only the entries you need.
+    in the number of stored entries. ``compressed_limit_bytes`` caps the
+    resulting ``(nG, n_select)`` complex output array. Pass an explicit
+    ``select=...`` to compute only the entries you need.
     """
     chosen = (method or "laguerre").strip().lower()
     backend_fn: Any
@@ -175,13 +179,28 @@ def get_exchange_kernels_compressed(
     if G_magnitudes.shape != G_angles.shape:
         raise ValueError("G_magnitudes and G_angles must have the same shape.")
 
+    select_list: list[Quad] | None
+    if select is None:
+        select_list = None
+        n_select_est = estimate_canonical_select_size(int(nmax))
+    else:
+        select_list = [cast(Quad, tuple(int(x) for x in quad)) for quad in select]
+        n_select_est = len(select_list)
+
+    guard_compressed_values_allocation(
+        nG=int(G_magnitudes.size),
+        n_select=int(n_select_est),
+        compressed_limit_bytes=compressed_limit_bytes,
+        backend_name=f"{chosen} exchange kernels",
+    )
+
     return cast(
         tuple[ComplexArray, list[Quad]],
         backend_fn(
             G_magnitudes,
             G_angles,
             nmax,
-            select=select,
+            select=select_list,
             canonical_select_max_entries=canonical_select_max_entries,
             **kwargs,
         ),

src/quantumhall_matrixelements/_materialize.py

@@ -7,6 +7,8 @@
 
 # Default soft cap for allocating a full (nG, nmax, nmax, nmax, nmax) complex tensor.
 DEFAULT_FULL_TENSOR_LIMIT_BYTES = 512 * 1024 * 1024  # 512 MiB
+# Default soft cap for allocating compressed (nG, n_select) complex arrays.
+DEFAULT_COMPRESSED_LIMIT_BYTES = DEFAULT_FULL_TENSOR_LIMIT_BYTES
 # Default soft cap for dense backend work tables such as quadrature precomputes.
 DEFAULT_WORKSPACE_LIMIT_BYTES = DEFAULT_FULL_TENSOR_LIMIT_BYTES
 
@@ -90,6 +92,30 @@ def guard_full_tensor_materialization(
         )
 
 
+def guard_compressed_values_allocation(
+    *,
+    nG: int,
+    n_select: int,
+    compressed_limit_bytes: float | int | None = DEFAULT_COMPRESSED_LIMIT_BYTES,
+    backend_name: str = "exchange kernels",
+    dtype: Any = np.complex128,
+) -> None:
+    """Raise proactively if a compressed ``(nG, n_select)`` array would be too large."""
+    if compressed_limit_bytes is None:
+        return
+
+    est = float(nG) * float(n_select) * np.dtype(dtype).itemsize
+    if est > float(compressed_limit_bytes):
+        human_est = format_bytes(est)
+        human_lim = format_bytes(float(compressed_limit_bytes))
+        raise MemoryError(
+            f"Refusing to allocate compressed ({nG}, {n_select}) values array "
+            f"(~{human_est} > {human_lim}) for {backend_name}; pass a smaller "
+            "explicit select=..., split the G grid, increase compressed_limit_bytes, "
+            "or pass compressed_limit_bytes=None to disable this guard."
+        )
+
+
 # --------------------------------------------------------------------------- #
 # Helpers for controlled materialization / reconstruction                     #
 # --------------------------------------------------------------------------- #
@@ -142,6 +168,8 @@ def materialize_full_tensor(
 
 __all__ = [
     "guard_full_tensor_materialization",
+    "guard_compressed_values_allocation",
+    "DEFAULT_COMPRESSED_LIMIT_BYTES",
     "DEFAULT_FULL_TENSOR_LIMIT_BYTES",
     "DEFAULT_WORKSPACE_LIMIT_BYTES",
     "build_canonical_select",

src/quantumhall_matrixelements/_select.py

@@ -21,6 +21,12 @@
 DEFAULT_CANONICAL_SELECT_MAX_ENTRIES = 2_000_000
 
 
+def estimate_canonical_select_size(nmax: int) -> int:
+    """Return the size of the canonical symmetry-reduced select list."""
+    n_pairs = int(nmax) * int(nmax)
+    return (n_pairs * (n_pairs + 1)) // 2
+
+
 def _coerce_quad(nmax: int, quad: Sequence[int]) -> Quad:
     if len(quad) != 4:
         raise ValueError("select entries must be (n1, m1, n2, m2) tuples.")
@@ -45,8 +51,7 @@ def normalize_select(
         if canonical_select_max_entries is not None:
             # number of (n,m) pairs is nmax^2; canonical selection is the upper
             # triangle in that pair space.
-            n_pairs = int(nmax) * int(nmax)
-            n_select = (n_pairs * (n_pairs + 1)) // 2
+            n_select = estimate_canonical_select_size(nmax)
             if n_select > int(canonical_select_max_entries):
                 raise MemoryError(
                     f"Refusing to build canonical select list for nmax={nmax}: "
@@ -71,5 +76,6 @@ def normalize_select(
 
 __all__ = [
     "DEFAULT_CANONICAL_SELECT_MAX_ENTRIES",
+    "estimate_canonical_select_size",
     "normalize_select",
 ]

src/quantumhall_matrixelements/fock.py

@@ -7,6 +7,8 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from ._materialize import DEFAULT_COMPRESSED_LIMIT_BYTES, guard_compressed_values_allocation
+from ._select import estimate_canonical_select_size
 from .exchange_hankel import get_exchange_kernels_hankel
 from .exchange_laguerre import (
     QuadratureParams,
@@ -21,6 +23,7 @@
 _FAST_LAGUERRE_PASSTHROUGH_KEYS = {
     "adaptive_nquad",
     "canonical_select_max_entries",
+    "compressed_limit_bytes",
     "kappa",
     "nquad",
     "potential",
@@ -93,6 +96,37 @@ def apply(rho: ComplexArray, include_minus: bool = True) -> ComplexArray:
     return apply
 
 
+def _guard_compressed_kernel_values(
+    G_magnitudes: RealArray,
+    G_angles: RealArray,
+    nmax: int,
+    *,
+    select: Iterable[Quad] | None,
+    compressed_limit_bytes: float | int | None,
+    backend_name: str,
+) -> tuple[RealArray, RealArray, list[Quad] | None]:
+    G_magnitudes_arr = cast(RealArray, np.asarray(G_magnitudes, dtype=np.float64).ravel())
+    G_angles_arr = cast(RealArray, np.asarray(G_angles, dtype=np.float64).ravel())
+    if G_magnitudes_arr.shape != G_angles_arr.shape:
+        raise ValueError("G_magnitudes and G_angles must have the same shape.")
+
+    select_list: list[Quad] | None
+    if select is None:
+        select_list = None
+        n_select = estimate_canonical_select_size(int(nmax))
+    else:
+        select_list = [cast(Quad, tuple(int(x) for x in quad)) for quad in select]
+        n_select = len(select_list)
+
+    guard_compressed_values_allocation(
+        nG=int(G_magnitudes_arr.size),
+        n_select=int(n_select),
+        compressed_limit_bytes=compressed_limit_bytes,
+        backend_name=backend_name,
+    )
+    return G_magnitudes_arr, G_angles_arr, select_list
+
+
 def _build_compressed_constructor_hf(
     values: ComplexArray,
     select: list[Quad],
@@ -295,6 +329,10 @@ def get_fockmatrix_constructor(
     """
 
     chosen = (method or "laguerre").strip().lower()
+    backend_kwargs = dict(kwargs)
+    compressed_limit_bytes = backend_kwargs.pop(
+        "compressed_limit_bytes", DEFAULT_COMPRESSED_LIMIT_BYTES
+    )
     backend: Any
     if chosen in {"hankel", "hk"}:
         backend = get_exchange_kernels_hankel
@@ -317,14 +355,23 @@ def get_fockmatrix_constructor(
             "Use 'laguerre', 'ogata', or 'hankel'."
         )
 
+    G_magnitudes_arr, G_angles_arr, select_list = _guard_compressed_kernel_values(
+        G_magnitudes,
+        G_angles,
+        nmax,
+        select=select,
+        compressed_limit_bytes=compressed_limit_bytes,
+        backend_name=f"{chosen} exchange kernels",
+    )
+
     values, select_list = cast(
         tuple[ComplexArray, list[Quad]],
         backend(
-            G_magnitudes,
-            G_angles,
+            G_magnitudes_arr,
+            G_angles_arr,
             nmax,
-            select=select,
-            **kwargs,
+            select=select_list,
+            **backend_kwargs,
         ),
     )
 
@@ -349,6 +396,10 @@ def get_fockmatrix_constructor_hf(
         Σ^F_{n m}(G) = - Σ_{r,t} X_{m r n t}(G) ρ^*_{t r}(G).
     """
     chosen = (method or "laguerre").strip().lower()
+    backend_kwargs = dict(kwargs)
+    compressed_limit_bytes = backend_kwargs.pop(
+        "compressed_limit_bytes", DEFAULT_COMPRESSED_LIMIT_BYTES
+    )
     backend: Any
     if chosen in {"hankel", "hk"}:
         backend = get_exchange_kernels_hankel
@@ -362,14 +413,23 @@ def get_fockmatrix_constructor_hf(
             "Use 'laguerre', 'ogata', or 'hankel'."
         )
 
+    G_magnitudes_arr, G_angles_arr, select_list = _guard_compressed_kernel_values(
+        G_magnitudes,
+        G_angles,
+        nmax,
+        select=select,
+        compressed_limit_bytes=compressed_limit_bytes,
+        backend_name=f"{chosen} exchange kernels",
+    )
+
     values, select_list = cast(
         tuple[ComplexArray, list[Quad]],
         backend(
-            G_magnitudes,
-            G_angles,
+            G_magnitudes_arr,
+            G_angles_arr,
             nmax,
-            select=select,
-            **kwargs,
+            select=select_list,
+            **backend_kwargs,
         ),
     )
     return _build_compressed_constructor_hf(values, select_list, nmax)

tests/test_api_guards.py

@@ -66,6 +66,38 @@ def test_canonical_select_guard_is_bypassed_for_explicit_select():
     assert select_list == [(0, 0, 0, 0)]
 
 
+def test_compressed_values_guard_raises_before_backend_work():
+    Gs = np.array([0.0, 1.0], dtype=float)
+    thetas = np.array([0.0, 0.1], dtype=float)
+
+    with pytest.raises(MemoryError):
+        get_exchange_kernels_compressed(
+            Gs,
+            thetas,
+            2,
+            method="ogata",
+            select=[(0, 0, 0, 0)],
+            compressed_limit_bytes=1,
+        )
+
+
+def test_compressed_values_limit_none_bypasses_guard():
+    Gs = np.array([0.0, 1.0], dtype=float)
+    thetas = np.array([0.0, 0.1], dtype=float)
+    select = [(0, 0, 0, 0)]
+
+    values, select_list = get_exchange_kernels_compressed(
+        Gs,
+        thetas,
+        2,
+        method="ogata",
+        select=select,
+        compressed_limit_bytes=None,
+    )
+    assert values.shape == (2, 1)
+    assert select_list == select
+
+
 def test_laguerre_workspace_guard_raises_for_compressed_call():
     Gs = np.array([12.0], dtype=float)
     thetas = np.array([0.0], dtype=float)

tests/test_fock_constructor.py

@@ -99,3 +99,21 @@ def test_fock_constructor_laguerre_workspace_guard_uses_fast_path():
             nquad=80,
             workspace_limit_bytes=1,
         )
+
+
+def test_fock_constructor_laguerre_accepts_compressed_limit_on_fast_path():
+    Gs = np.array([0.0, 1.2])
+    thetas = np.array([0.0, 0.3])
+    rho = _random_hermitian_rho(len(Gs), 3, seed=7)
+
+    ref = get_fockmatrix_constructor(Gs, thetas, 3, method="laguerre", nquad=200)(rho)
+    out = get_fockmatrix_constructor(
+        Gs,
+        thetas,
+        3,
+        method="laguerre",
+        nquad=200,
+        compressed_limit_bytes=1,
+    )(rho)
+
+    assert np.allclose(out, ref, rtol=1e-10, atol=1e-12)