diffract-core 0.2.1__py3-none-any.whl

diffract/core/compute/kernels/norms.py

@@ -0,0 +1,88 @@
+import math
+from sys import float_info
+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
+
+
+@kernel(name="pl_alpha_norm", require_fields=("esd", "pl_alpha"))
+@kernel(name="tpl_alpha_norm", require_fields=("esd", "tpl_alpha"))
+def alpha_norm(esd: NDArray[np.floating[Any]], alpha: float) -> float:
+    """Compute alpha norm: sum(esd^alpha)."""
+    return cast("float", np.sum(esd**alpha).item())
+
+
+@kernel(
+    name="model_pl_alpha_norm",
+    require_fields=("pl_alpha_norm",),
+    apply_level=KernelApplyLevel.IN_MODEL,
+)
+@kernel(
+    name="model_tpl_alpha_norm",
+    require_fields=("tpl_alpha_norm",),
+    apply_level=KernelApplyLevel.IN_MODEL,
+)
+def model_alpha_norm(alpha_norm: tuple[float, ...]) -> float:
+    """Compute mean log10 alpha norm across model parameters."""
+    alpha_norm_arr: NDArray[np.floating[Any]] = np.array(alpha_norm)
+    log_alpha_norm = np.where(
+        alpha_norm_arr == 0, float_info.min_10_exp, np.log10(alpha_norm_arr)
+    )
+    return cast("float", np.mean(log_alpha_norm).item())
+
+
+@kernel
+def frob_norm(weights_svals: NDArray[np.floating[Any]]) -> float:
+    """Compute Frobenius norm from singular values."""
+    return math.sqrt(cast("float", np.sum(weights_svals**2).item()))
+
+
+@kernel
+def l1_norm(weights_svals: NDArray[np.floating[Any]]) -> float:
+    """Compute L1 norm (sum of singular values)."""
+    return cast("float", weights_svals.sum().item())
+
+
+@kernel
+def l2_norm(max_weights_sval: float) -> float:
+    """Compute L2 (spectral) norm (maximum singular value)."""
+    return max_weights_sval
+
+
+@kernel(apply_level=KernelApplyLevel.IN_MODEL)
+def log_norm(frob_norm: tuple[float, ...]) -> float:
+    """Compute mean squared log Frobenius norm across model parameters."""
+    frob_norm_arr: NDArray[np.floating[Any]] = np.array(frob_norm)
+    return cast("float", np.mean(np.log(frob_norm_arr) ** 2).item())
+
+
+@kernel(apply_level=KernelApplyLevel.IN_MODEL)
+def log_spectral_norm(l2_norm: tuple[float, ...]) -> float:
+    """Compute mean squared log spectral norm across model parameters."""
+    l2_norm_arr: NDArray[np.floating[Any]] = np.array(l2_norm)
+    return cast("float", np.mean(np.log(l2_norm_arr) ** 2).item())
+
+
+@kernel(apply_level=KernelApplyLevel.IN_MODEL)
+def param_norm(frob_norm: tuple[float, ...]) -> float:
+    """Compute sum of squared Frobenius norms across model parameters."""
+    frob_norm_arr: NDArray[np.floating[Any]] = np.array(frob_norm)
+    return cast("float", np.sum(frob_norm_arr**2).item())
+
+
+@kernel(apply_level=KernelApplyLevel.IN_MODEL)
+def prod_frob_norm(frob_norm: tuple[float, ...]) -> float:
+    """Compute product of Frobenius norms across model parameters."""
+    frob_norm_arr: NDArray[np.floating[Any]] = np.array(frob_norm)
+    return cast("float", np.prod(frob_norm_arr).item())
+
+
+@kernel(apply_level=KernelApplyLevel.IN_MODEL)
+def prod_spectral_norm(l2_norm: tuple[float, ...]) -> float:
+    """Compute product of spectral norms across model parameters."""
+    l2_norm_arr: NDArray[np.floating[Any]] = np.array(l2_norm)
+    return cast("float", np.prod(l2_norm_arr).item())

diffract/core/compute/kernels/ranks.py

@@ -0,0 +1,35 @@
+import math
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from diffract.core.compute.decorator import kernel
+
+EPS = 1e-16
+
+
+@kernel
+def effective_rank(weights_svals: NDArray[np.floating[Any]]) -> float:
+    """Compute effective rank from singular values using entropy."""
+    probabilities = weights_svals / cast("float", weights_svals.sum().item() + EPS)
+    entropy = cast("float", -np.sum(probabilities * np.log(probabilities + EPS)).item())
+    return math.exp(entropy)
+
+
+@kernel
+def hard_rank(esd: NDArray[np.floating[Any]], *, threshold: float = 1e-5) -> int:
+    """Count eigenvalues above threshold."""
+    return cast("int", np.sum((esd > threshold).astype(int)).item())
+
+
+@kernel
+def mp_soft_rank(mp_esd_max: float, esd_max: float) -> float:
+    """Compute soft rank from MP bulk maximum relative to ESD maximum."""
+    return mp_esd_max / (esd_max + EPS)
+
+
+@kernel
+def stable_rank(frob_norm: float, l2_norm: float) -> float:
+    """Compute stable rank (Frobenius norm / spectral norm)^2."""
+    return (frob_norm / (l2_norm + EPS)) ** 2

diffract/core/compute/kernels/tracy_widom.py

@@ -0,0 +1,36 @@
+import math
+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.extensions.rmt import tracy_widom_ppf
+
+
+@kernel
+def tw_esd_bound(
+    greater_dim: int,
+    lower_dim: int,
+    mp_bulk_std: float,
+    *,
+    p_value_threshold: float = 0.005,
+) -> float:
+    """Compute Tracy-Widom bound for ESD spike detection."""
+    g_dim, l_dim = greater_dim, lower_dim
+
+    loc = math.sqrt(g_dim - 1) + math.sqrt(l_dim)
+    inv_loc = 1 / math.sqrt(g_dim - 1) + 1 / math.sqrt(l_dim)
+    scale = (mp_bulk_std**2) / g_dim
+
+    mu: float = scale * (loc**2)
+    sigma: float = scale * (loc) * (inv_loc ** (1 / 3))
+    pure_twd_bound = tracy_widom_ppf(1 - p_value_threshold)
+
+    return pure_twd_bound * sigma + mu
+
+
+@kernel
+def tw_num_spikes(esd: NDArray[np.floating[Any]], tw_esd_bound: float) -> int:
+    """Count number of spikes above Tracy-Widom bound."""
+    return cast("int", np.sum((esd > tw_esd_bound).astype(int)).item())

diffract/core/compute/metadata.py

@@ -0,0 +1,67 @@
+"""Kernel metadata and information structures."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from .config import KernelConfig
+from .execution import KernelApplyLevel, KernelExecutionProtocol, KernelRestrictions
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+@dataclass(frozen=True)
+class KernelInfo:
+    """Optional human-readable kernel documentation fields.
+
+    Contains optional metadata for kernel documentation and introspection.
+    Used to provide additional context about kernel purpose and usage.
+
+    Attributes:
+        summary: Brief description of kernel functionality.
+        notes: Additional implementation notes or usage guidelines.
+    """
+
+    summary: str | None = None
+    notes: str | None = None
+
+
+@dataclass
+class KernelMetadata:
+    """Complete metadata container for registered kernels.
+
+    Stores all information required for kernel registration, execution,
+    and dependency resolution within the kernel registry system.
+
+    Attributes:
+        name: Unique kernel identifier.
+        require_fields: Tuple of field names required as input.
+        produce_fields: Tuple of field names produced as output.
+        implementation: Callable implementing the kernel logic.
+        apply_level: Level at which kernel is applied (parameter/model/cross-model).
+        execution_protocol: Execution strategy (sequential/parallel).
+        restrictions: Optional execution restrictions.
+        config: Configuration parameters with defaults.
+        info: Optional documentation metadata.
+    """
+
+    name: str
+    require_fields: tuple[str, ...]
+    produce_fields: tuple[str, ...]
+    implementation: Callable[..., Any]
+    apply_level: KernelApplyLevel
+    execution_protocol: KernelExecutionProtocol | None
+    restrictions: KernelRestrictions | None
+    config: KernelConfig
+    info: KernelInfo
+
+    def __repr__(self) -> str:  # pragma: no cover - debug helper
+        """Return string representation for debugging and logging."""
+        conf = ", ".join(f"{k}={v}" for k, v in self.config.as_dict().items())
+        restrictions = f",{self.restrictions}" if self.restrictions else ""
+        return (
+            f"{self.name} ({', '.join(self.require_fields)}; {conf}): "
+            f"{self.apply_level}{restrictions}"
+        )