diffract-core 0.2.1__py3-none-any.whl

diffract/core/compute/extensions/rmt.py

@@ -0,0 +1,150 @@
+"""Random matrix theory distributions for the spectral kernels.
+
+Implements the closed-form cumulative distribution function of the
+Marchenko-Pastur law and the quantile function of the Tracy-Widom law for
+beta = 1 (GOE). The Marchenko-Pastur CDF integrates the density through the
+standard antiderivative of ``sqrt((x - a)(b - x)) / x``. The Tracy-Widom
+distribution is computed from its Painleve II representation: the
+Hastings-McLeod solution ``q'' = s q + 2 q^3`` with ``q(s) ~ Ai(s)`` as
+``s -> +inf`` is integrated numerically together with the auxiliary
+integrals that give ``F2(s) = exp(-int (x - s) q^2 dx)`` and
+``F1(s)^2 = F2(s) exp(-int q dx)``; quantiles invert the resulting CDF with
+a monotone PCHIP spline built once per process.
+
+References:
+    - Marchenko, V. A. and Pastur, L. A. "Distribution of eigenvalues for
+      some sets of random matrices". Mathematics of the USSR-Sbornik (1967).
+    - Tracy, C. A. and Widom, H. "Level-spacing distributions and the Airy
+      kernel". Communications in Mathematical Physics (1994).
+    - Tracy, C. A. and Widom, H. "On orthogonal and symplectic matrix
+      ensembles". Communications in Mathematical Physics (1996).
+    - Hastings, S. P. and McLeod, J. B. "A boundary value problem associated
+      with the second Painleve transcendent and the Korteweg-de Vries
+      equation". Archive for Rational Mechanics and Analysis (1980).
+    - Bejan, A. "Largest eigenvalues and sample covariance matrices".
+      M.Sc. dissertation, The University of Warwick (2005).
+"""
+
+from __future__ import annotations
+
+from functools import cache
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.integrate import quad, solve_ivp
+from scipy.interpolate import PchipInterpolator
+from scipy.special import airy
+
+_PAINLEVE_S_START = 10.0
+_PAINLEVE_S_END = -5.0
+_PAINLEVE_RTOL = 1e-12
+_PAINLEVE_ATOL = 1e-18
+_TW_GRID_STEP = 0.005
+
+
+def marchenko_pastur_cdf(
+    x: NDArray[np.floating[Any]] | float,
+    ratio: float,
+    sigma: float = 1.0,
+) -> NDArray[np.float64]:
+    """Evaluate the Marchenko-Pastur CDF for beta = 1.
+
+    Args:
+        x: Value or array of values to evaluate the CDF at.
+        ratio: Dimension ratio lambda in (0, 1] of the underlying matrix.
+        sigma: Standard deviation of the matrix entries.
+
+    Returns:
+        CDF values as a float64 array of the same shape as ``x``.
+
+    Raises:
+        ValueError: If ``ratio`` is outside (0, 1] or ``sigma`` is not positive.
+    """
+    if not 0 < ratio <= 1:
+        raise ValueError(f"ratio must be in (0, 1], got {ratio}")
+    if sigma <= 0:
+        raise ValueError(f"sigma must be positive, got {sigma}")
+
+    sqrt_ratio = np.sqrt(ratio)
+    edge_low = (1 - sqrt_ratio) ** 2
+    edge_high = (1 + sqrt_ratio) ** 2
+
+    z = np.asarray(x, dtype=np.float64) / sigma**2
+    result = np.zeros_like(z)
+    result[z >= edge_high] = 1.0
+
+    inside = (z > edge_low) & (z < edge_high)
+    if np.any(inside):
+        zi = z[inside]
+        root = np.sqrt((zi - edge_low) * (edge_high - zi))
+        span = edge_high - edge_low
+        first_arcsin = np.arcsin((2 * zi - edge_low - edge_high) / span)
+        product = edge_low * edge_high
+        second_arcsin = np.arcsin(
+            ((edge_low + edge_high) * zi - 2 * product) / (zi * span)
+        )
+        result[inside] = (
+            root
+            + (1 + ratio) * first_arcsin
+            - (1 - ratio) * second_arcsin
+            + np.pi * ratio
+        ) / (2 * np.pi * ratio)
+
+    return np.clip(result, 0.0, 1.0)
+
+
+@cache
+def _tw_beta1_cdf_spline() -> PchipInterpolator:
+    s_start = _PAINLEVE_S_START
+    ai_start, ai_prime_start, *_ = airy(s_start)
+    airy_tail_sq = quad(lambda t: airy(t)[0] ** 2, s_start, np.inf)[0]
+    airy_tail_weighted = quad(
+        lambda t: (t - s_start) * airy(t)[0] ** 2, s_start, np.inf
+    )[0]
+    airy_tail = quad(lambda t: airy(t)[0], s_start, np.inf)[0]
+
+    def rhs(s: float, y: NDArray[np.float64]) -> list[float]:
+        q, q_prime, _, k_val, _ = y
+        return [q_prime, s * q + 2 * q**3, -k_val, -(q**2), -q]
+
+    solution = solve_ivp(
+        rhs,
+        [s_start, _PAINLEVE_S_END],
+        [ai_start, ai_prime_start, airy_tail_weighted, airy_tail_sq, airy_tail],
+        method="DOP853",
+        rtol=_PAINLEVE_RTOL,
+        atol=_PAINLEVE_ATOL,
+        dense_output=True,
+    )
+    if solution.status != 0:
+        raise RuntimeError(f"Painleve II integration failed: {solution.message}")
+
+    s_grid = np.arange(_PAINLEVE_S_END, s_start, _TW_GRID_STEP)
+    states = solution.sol(s_grid)
+    cdf = np.exp(-(states[2] + states[4]) / 2)
+
+    increasing = np.concatenate(([True], np.diff(cdf) > 0))
+    return PchipInterpolator(cdf[increasing], s_grid[increasing], extrapolate=False)
+
+
+def tracy_widom_ppf(q: float) -> float:
+    """Evaluate the Tracy-Widom (beta = 1) quantile function.
+
+    Args:
+        q: Lower-tail probability inside the computed CDF range.
+
+    Returns:
+        The quantile corresponding to ``q``.
+
+    Raises:
+        ValueError: If ``q`` falls outside the computed CDF range.
+    """
+    spline = _tw_beta1_cdf_spline()
+    q_low, q_high = spline.x[0], spline.x[-1]
+    if not q_low <= q <= q_high:
+        raise ValueError(
+            f"q must be within the computed CDF range "
+            f"[{q_low:.2e}, {q_high:.10f}], got {q}"
+        )
+    return float(spline(q))

diffract/core/compute/extensions/utils.py

@@ -0,0 +1,36 @@
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+import diffract.core.utils.imports as import_utils
+
+if not import_utils.is_available("torch"):
+
+    def torch_cuda_wrapper(*_args: Any, **_kwargs: Any) -> None:
+        """Stub when torch is not available."""
+        raise ImportError
+
+else:
+    torch = import_utils.require("torch")
+
+    def torch_cuda_wrapper(
+        array: NDArray[np.floating[Any]],
+        function: Callable[[torch.Tensor], torch.Tensor],
+    ) -> Any:
+        """Execute function on CUDA tensor and return to CPU."""
+        torch.cuda.empty_cache()
+        with torch.no_grad():
+            tensor = torch.Tensor(array).to("cuda")
+            output = function(tensor)
+            if isinstance(output, torch.Tensor):
+                output = output.to("cpu").float().numpy()
+            else:
+                output_container_type = type(output)
+
+                output_ = [item.to("cpu").float().numpy() for item in output]
+                output = output_container_type(output_)
+
+        del tensor
+        return output

diffract/core/compute/field_signature.py

@@ -0,0 +1,183 @@
+"""Utilities to describe stored fields (shape/kind/dtype) for plotting/inspection."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+
+from diffract.core.compute.execution.enums import KernelApplyLevel
+from diffract.core.compute.registry import KernelRegistry
+from diffract.core.storage.interface import IStorageManager
+from diffract.core.storage.metadata import ValueMetadata, infer_value_metadata
+from diffract.core.utils.exceptions import format_exception_message
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class FieldSignature:
+    """Metadata describing a stored field's characteristics."""
+
+    field: str
+    kind: str
+    dtype: str | None
+    shape: tuple[int, ...] | None
+    ndim: int | None
+    is_numeric: bool
+    apply_level: KernelApplyLevel | None = None
+    source: str | None = None  # e.g., "metadata", "inferred"
+    available: bool = True
+
+
+def _signature_from_meta(
+    field: str,
+    meta: dict[str, Any],
+    *,
+    apply_level: KernelApplyLevel | None = None,
+    source: str = "metadata",
+) -> FieldSignature:
+    vm = ValueMetadata.from_jsonable(meta)
+    return FieldSignature(
+        field=field,
+        kind=vm.kind,
+        dtype=vm.dtype,
+        shape=vm.shape,
+        ndim=vm.ndim,
+        is_numeric=vm.is_numeric,
+        apply_level=apply_level,
+        source=source,
+    )
+
+
+def _signature_from_value(
+    field: str,
+    value: Any,
+    *,
+    apply_level: KernelApplyLevel | None = None,
+    source: str = "value",
+) -> FieldSignature:
+    vm = infer_value_metadata(value)
+    return FieldSignature(
+        field=field,
+        kind=vm.kind,
+        dtype=vm.dtype,
+        shape=vm.shape,
+        ndim=vm.ndim,
+        is_numeric=vm.is_numeric,
+        apply_level=apply_level,
+        source=source,
+    )
+
+
+def collect_field_signatures(
+    storage: IStorageManager,
+    field_names: Iterable[str] | None = None,
+    *,
+    sample_limit: int = 8,
+    apply_level_hint: KernelApplyLevel | None = None,
+) -> dict[str, FieldSignature]:
+    """Collect signatures for fields stored in a storage backend.
+
+    Sampling is limited to avoid loading all parameters.
+    """
+    signatures: dict[str, FieldSignature] = {}
+
+    fields = list(field_names) if field_names is not None else storage.list_fields()
+
+    for field in fields:
+        if field in signatures:
+            continue
+
+        try:
+            obj_ids = storage.list_objs_has_field(field)
+        except Exception as e:  # noqa: BLE001
+            logger.debug(
+                "Failed to list objects for field '%s': %s",
+                field,
+                format_exception_message(e),
+                exc_info=True,
+            )
+            obj_ids = []
+
+        obj_ids = obj_ids[:sample_limit]
+        signature: FieldSignature | None = None
+
+        for obj_uid in obj_ids:
+            meta = storage.get_field_metadata(obj_uid, field)
+            if meta is not None:
+                signature = _signature_from_meta(
+                    field, meta, apply_level=apply_level_hint, source="metadata"
+                )
+                break
+
+            try:
+                value = storage.get_field(obj_uid, field)
+            except Exception as e:  # noqa: BLE001
+                logger.debug(
+                    "Failed to read field '%s' for object '%s': %s",
+                    field,
+                    obj_uid,
+                    format_exception_message(e),
+                    exc_info=True,
+                )
+                continue
+
+            signature = _signature_from_value(
+                field, value, apply_level=apply_level_hint, source="value"
+            )
+            break
+
+        if signature is not None:
+            signatures[field] = signature
+
+    return signatures
+
+
+def _registry_apply_levels(registry: KernelRegistry) -> dict[str, KernelApplyLevel]:
+    mapping: dict[str, KernelApplyLevel] = {}
+    for kernel_name in registry.list_kernels():
+        apply_level = registry.get_kernel_apply_level(kernel_name)
+        for field in registry.get_fields_kernel_produce(kernel_name):
+            mapping[field] = apply_level
+    return mapping
+
+
+def collect_field_catalog(
+    storage: IStorageManager,
+    registry: KernelRegistry,
+    field_names: Iterable[str] | None = None,
+    *,
+    sample_limit: int = 8,
+) -> dict[str, FieldSignature]:
+    """Merge observed signatures with registry hints.
+
+    Includes computable-only fields.
+    """
+    observed = collect_field_signatures(
+        storage, field_names=field_names, sample_limit=sample_limit
+    )
+    apply_levels = _registry_apply_levels(registry)
+
+    # Enrich observed with apply_level if known.
+    for field, level in apply_levels.items():
+        if field in observed and observed[field].apply_level is None:
+            observed[field].apply_level = level
+
+    # Add registry-only fields (not yet computed).
+    for field, level in apply_levels.items():
+        if field not in observed:
+            observed[field] = FieldSignature(
+                field=field,
+                kind="object",
+                dtype=None,
+                shape=None,
+                ndim=None,
+                is_numeric=False,
+                apply_level=level,
+                source="registry",
+                available=False,
+            )
+
+    return observed

diffract/core/compute/kernels/__init__.py

@@ -0,0 +1,48 @@
+from .alignment import (
+    avg_vector_agreement,
+    max_vector_agreement,
+    overlap,
+    vector_agreement,
+)
+from .heavy_tailed import (
+    expon_p_value,
+    exponential_fit,
+    ht_concentration,
+    ht_presence,
+    ht_scale,
+    pl_p_value,
+    power_law_fit,
+    tpl_p_value,
+    truncated_power_law_fit,
+)
+from .marchenko_pastur import (
+    marchenko_pastur_fit,
+    mp_concentration,
+    mp_ks,
+    mp_num_spikes,
+    mp_presence,
+    mp_sval_max,
+)
+from .mat_decomposition import esd, esd_max, esd_min, svd
+from .mat_properties import (
+    aspect_ratio,
+    greater_dim,
+    lower_dim,
+    weights_rand,
+    weights_std,
+)
+from .model_quality import pl_alpha_weighted, rand_distance
+from .norms import (
+    alpha_norm,
+    frob_norm,
+    l1_norm,
+    l2_norm,
+    log_norm,
+    log_spectral_norm,
+    model_alpha_norm,
+    param_norm,
+    prod_frob_norm,
+    prod_spectral_norm,
+)
+from .ranks import effective_rank, hard_rank, mp_soft_rank, stable_rank
+from .tracy_widom import tw_esd_bound, tw_num_spikes

diffract/core/compute/kernels/alignment.py

@@ -0,0 +1,106 @@
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from diffract.core.compute.decorator import kernel
+from diffract.core.compute.execution.enums import (
+    KernelApplyLevel,
+    KernelRestrictions,
+)
+
+
+@kernel(
+    name="l_overlap",
+    require_fields=("weights_lsvs", "lower_dim"),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+    restrictions=KernelRestrictions.BINARY,
+)
+@kernel(
+    name="r_overlap",
+    require_fields=("weights_rsvs", "lower_dim"),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+    restrictions=KernelRestrictions.BINARY,
+)
+def overlap(
+    weights_svs: tuple[NDArray[np.floating[Any]], NDArray[np.floating[Any]]],
+    lower_dim: tuple[int, int],
+    *,
+    absolute: bool = False,
+) -> NDArray[np.floating[Any]]:
+    """Compute overlap matrix between two sets of singular vectors."""
+    svs, svs_other = weights_svs
+
+    l_dim, l_dim_other = lower_dim
+    if l_dim != l_dim_other:
+        msg = "l_dim and l_dim_other should be equal"
+        raise ValueError(msg)
+
+    result = (svs.T @ svs_other)[:l_dim, :l_dim]
+
+    if absolute:
+        result = np.abs(result)
+
+    return cast("NDArray[np.floating[Any]]", result)
+
+
+@kernel(
+    name="l_agreement",
+    require_fields=("l_overlap",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+@kernel(
+    name="r_agreement",
+    require_fields=("r_overlap",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+def vector_agreement(overlap: NDArray[np.floating[Any]]) -> NDArray[np.floating[Any]]:
+    """Return the diagonal of an overlap matrix (per-component agreement)."""
+    result = np.diag(overlap)
+    return cast("NDArray[np.floating[Any]]", result)
+
+
+@kernel(
+    name="max_l_agreement",
+    require_fields=("l_overlap",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+@kernel(
+    name="max_r_agreement",
+    require_fields=("r_overlap",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+def max_vector_agreement(
+    overlap: NDArray[np.floating[Any]],
+) -> NDArray[np.floating[Any]]:
+    """Return max absolute overlap per row (best agreement per component)."""
+    result = np.max(np.abs(overlap), axis=1)
+    return cast("NDArray[np.floating[Any]]", result)
+
+
+@kernel(
+    name="avg_max_l_agreement",
+    require_fields=("max_l_agreement",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+@kernel(
+    name="avg_max_r_agreement",
+    require_fields=("max_r_agreement",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+@kernel(
+    name="avg_l_agreement",
+    require_fields=("l_agreement",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+@kernel(
+    name="avg_r_agreement",
+    require_fields=("r_agreement",),
+    apply_level=KernelApplyLevel.CROSS_MODEL,
+)
+def avg_vector_agreement(
+    agreement: NDArray[np.floating[Any]],
+) -> NDArray[np.floating[Any]]:
+    """Return the mean of per-component agreement values."""
+    result = np.mean(agreement)
+    return cast("NDArray[np.floating[Any]]", result)