diffract-core 0.2.1__py3-none-any.whl

diffract/core/compute/kernels/heavy_tailed.py

@@ -0,0 +1,394 @@
+import io
+import warnings
+from contextlib import redirect_stderr, redirect_stdout
+from operator import itemgetter
+from typing import Any, Literal, cast
+
+import numpy as np
+import powerlaw as pl
+from numpy.typing import NDArray
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.compute.decorator import kernel
+
+# Availability is probed without importing the accelerated module itself:
+# importing it initializes the taichi runtime, which must not happen (and
+# must not be able to fail) during `import diffract`.
+_DIFFRACT_FIT_AVAILABLE = import_utils.is_available("taichi")
+
+_fit_class: Any = None
+_DIFFRACT_FIT_BROKEN = False
+
+# The statistical floor of the accelerated fitter: its xmin candidates
+# must leave a tail of MIN_TAIL_SIZE (50) points, per Clauset et al., so
+# below ~2x that it cannot fit at all. From this size up it is also
+# faster than the `powerlaw` library at every tail size once warm.
+AUTO_DIFFRACT_MIN_SIZE = 100
+
+FitMethod = Literal["auto", "powerlaw", "diffract"]
+
+
+def _accelerated_fit_ready() -> bool:
+    """True when the accelerated implementation can actually initialize.
+
+    taichi may be importable yet fail at runtime initialization; `auto`
+    must degrade to the powerlaw library then, not raise.
+    """
+    global _DIFFRACT_FIT_BROKEN  # noqa: PLW0603
+    if _DIFFRACT_FIT_BROKEN or not _DIFFRACT_FIT_AVAILABLE:
+        return False
+    try:
+        _get_fit_class()
+    except ModuleNotFoundError:
+        _DIFFRACT_FIT_BROKEN = True
+        return False
+    return True
+
+
+def _resolve_fit_method(fit_method: FitMethod, data_size: int) -> str:
+    """Resolve `auto` to an implementation by data size and availability.
+
+    `auto` never raises over the accelerated path: explicit
+    fit_method="diffract" surfaces initialization errors instead.
+    """
+    if fit_method != "auto":
+        return fit_method
+    if data_size < AUTO_DIFFRACT_MIN_SIZE:
+        return "powerlaw"
+    return "diffract" if _accelerated_fit_ready() else "powerlaw"
+
+
+def _get_fit_class() -> Any:
+    """Import the accelerated Fit implementation on first use."""
+    global _fit_class  # noqa: PLW0603
+    if _fit_class is None:
+        try:
+            module = import_utils.get_module(
+                "diffract.core.compute.extensions.power_law"
+            )
+            _fit_class = None if module is None else module.Fit
+        except Exception as e:  # optional accelerator must not crash callers
+            raise ModuleNotFoundError(
+                "The accelerated 'diffract' fit implementation failed to "
+                f"initialize: {e}. Use fit_method='powerlaw' instead."
+            ) from e
+        if _fit_class is None:
+            raise ModuleNotFoundError(
+                "The accelerated 'diffract' fit implementation needs taichi. "
+                'Install it with: pip install "diffract-core[taichi]", '
+                "or use fit_method='powerlaw'."
+            )
+    return _fit_class
+
+
+# region general
+
+
+@kernel(name="pl_concentration", require_fields=("esd", "pl_esd_xmin"))
+@kernel(name="tpl_concentration", require_fields=("esd", "tpl_esd_xmin"))
+@kernel(name="expon_concentration", require_fields=("esd", "expon_esd_xmin"))
+def ht_concentration(esd: NDArray[np.floating[Any]], ht_esd_xmin: float) -> float:
+    """Compute heavy-tailed concentration (tail size / total size)."""
+    tail_size = cast("int", np.sum(esd >= ht_esd_xmin).item())
+
+    return tail_size / esd.size
+
+
+@kernel(name="pl_presence", require_fields=("esd_min", "esd_max", "pl_esd_xmin"))
+@kernel(name="tpl_presence", require_fields=("esd_min", "esd_max", "tpl_esd_xmin"))
+@kernel(name="expon_presence", require_fields=("esd_min", "esd_max", "expon_esd_xmin"))
+def ht_presence(esd_min: float, esd_max: float, ht_esd_xmin: float) -> float:
+    """Compute heavy-tailed presence (tail width / total width)."""
+    esd_width = esd_max - esd_min
+    ht_width = esd_max - ht_esd_xmin
+
+    return ht_width / esd_width if esd_width > 0 else float("nan")
+
+
+@kernel(name="tpl_scale", require_fields=("esd_max", "tpl_lambda"))
+@kernel(name="expon_scale", require_fields=("esd_max", "expon_lambda"))
+def ht_scale(esd_max: float, ht_lambda: float) -> float:
+    """Compute heavy-tailed scale (max * lambda)."""
+    return esd_max * ht_lambda
+
+
+# endregion general
+
+# region power_law
+
+
+def power_law_fit_powerlaw_implementation(
+    data: NDArray[np.floating[Any]],
+) -> tuple[float, float, float]:
+    """Fit power law using powerlaw library."""
+    f = io.StringIO()
+    with redirect_stdout(f), redirect_stderr(f), warnings.catch_warnings():
+        warnings.simplefilter(action="ignore", category=RuntimeWarning)
+
+        distribution = "power_law"
+        fit = pl.Fit(
+            data,
+            distribution=distribution,
+            xmin_distribution=distribution,
+            discrete=False,
+        )
+
+    fit_result = getattr(fit, distribution)
+    pl_alpha = fit_result.alpha
+    pl_esd_xmin = fit_result.xmin
+    pl_ks = fit_result.KS()
+
+    return pl_alpha, pl_esd_xmin, pl_ks
+
+
+def power_law_fit_diffract_implementation(
+    data: NDArray[np.floating[Any]],
+) -> tuple[float, float, float]:
+    """Fit power law using diffract's accelerated Taichi implementation."""
+    fit = _get_fit_class()(data, "power_law")
+    try:
+        params = fit.fit_params()
+    finally:
+        fit.close()
+
+    pl_alpha, pl_esd_xmin, pl_ks = itemgetter("pl_alpha", "xmin", "ks_distance")(params)
+
+    return pl_alpha, pl_esd_xmin, pl_ks
+
+
+@kernel(produce_fields=("pl_alpha", "pl_esd_xmin", "pl_ks"))
+def power_law_fit(
+    esd: NDArray[np.floating[Any]],
+    *,
+    fit_method: FitMethod = "auto",
+) -> tuple[float, float, float]:
+    """Fit power law distribution to ESD data.
+
+    `auto` uses the accelerated implementation when the taichi extra is
+    installed and the ESD is large enough for it to fit at all.
+    """
+    match _resolve_fit_method(fit_method, esd.size):
+        case "powerlaw":
+            return power_law_fit_powerlaw_implementation(esd)
+        case "diffract":
+            return power_law_fit_diffract_implementation(esd)
+        case _:
+            msg = (
+                f"Fit method {fit_method!r} not implemented, choose method "
+                "from ('auto', 'powerlaw', 'diffract')."
+            )
+            raise ValueError(msg)
+
+
+def pl_p_value(
+    esd: NDArray[np.floating[Any]], pl_alpha: float, pl_esd_xmin: float, pl_ks: float
+) -> float:
+    """Compute p-value for power law fit (requires the taichi extra)."""
+    fit = _get_fit_class()(esd, "power_law")
+    try:
+        fit.set_params(
+            xmin=pl_esd_xmin,
+            pl_alpha=pl_alpha,
+            ks_distance=pl_ks,
+            tail_size=cast("int", np.sum(esd >= pl_esd_xmin).item()),
+        )
+        _, p_value = fit.p_value_test()
+    finally:
+        fit.close()
+
+    return p_value
+
+
+# endregion power_law
+
+# region truncated_power_law
+
+
+def truncated_power_law_fit_powerlaw_implementation(
+    data: NDArray[np.floating[Any]],
+) -> tuple[float, float, float, float]:
+    """Fit truncated power law using powerlaw library."""
+    f = io.StringIO()
+    with redirect_stdout(f), redirect_stderr(f), warnings.catch_warnings():
+        warnings.simplefilter(action="ignore", category=RuntimeWarning)
+
+        distribution = "truncated_power_law"
+        fit = pl.Fit(
+            data,
+            distribution=distribution,
+            xmin_distribution=distribution,
+            discrete=False,
+        )
+
+    fit_result = getattr(fit, distribution)
+    tpl_alpha = fit_result.alpha
+    tpl_lambda = fit_result.Lambda
+    tpl_esd_xmin = fit_result.xmin
+    tpl_ks = fit_result.KS()
+
+    return tpl_alpha, tpl_lambda, tpl_esd_xmin, tpl_ks
+
+
+def truncated_power_law_fit_diffract_implementation(
+    data: NDArray[np.floating[Any]],
+) -> tuple[float, float, float, float]:
+    """Fit truncated power law using diffract's accelerated Taichi implementation."""
+    fit = _get_fit_class()(data, "truncated_power_law")
+    try:
+        params = fit.fit_params()
+    finally:
+        fit.close()
+
+    tpl_alpha, tpl_lambda, tpl_esd_xmin, tpl_ks = itemgetter(
+        "pl_alpha", "expon_lambda", "xmin", "ks_distance"
+    )(params)
+
+    return tpl_alpha, tpl_lambda, tpl_esd_xmin, tpl_ks
+
+
+@kernel(produce_fields=("tpl_alpha", "tpl_lambda", "tpl_esd_xmin", "tpl_ks"))
+def truncated_power_law_fit(
+    esd: NDArray[np.floating[Any]],
+    *,
+    fit_method: FitMethod = "auto",
+) -> tuple[float, float, float, float]:
+    """Fit truncated power law distribution to ESD data.
+
+    `auto` uses the accelerated implementation when the taichi extra is
+    installed and the ESD is large enough for it to fit at all.
+    """
+    match _resolve_fit_method(fit_method, esd.size):
+        case "powerlaw":
+            return truncated_power_law_fit_powerlaw_implementation(esd)
+        case "diffract":
+            return truncated_power_law_fit_diffract_implementation(esd)
+        case _:
+            msg = (
+                f"Fit method {fit_method!r} not implemented, choose method "
+                "from ('auto', 'powerlaw', 'diffract')."
+            )
+            raise ValueError(msg)
+
+
+def tpl_p_value(
+    esd: NDArray[np.floating[Any]],
+    tpl_alpha: float,
+    tpl_lambda: float,
+    tpl_esd_xmin: float,
+    tpl_ks: float,
+) -> float:
+    """Compute p-value for truncated power law fit (requires the taichi extra)."""
+    fit = _get_fit_class()(esd, "truncated_power_law")
+    try:
+        fit.set_params(
+            xmin=tpl_esd_xmin,
+            pl_alpha=tpl_alpha,
+            expon_lambda=tpl_lambda,
+            ks_distance=tpl_ks,
+            tail_size=cast("int", np.sum(esd >= tpl_esd_xmin).item()),
+        )
+        _, p_value = fit.p_value_test()
+    finally:
+        fit.close()
+
+    return p_value
+
+
+# endregion truncated_power_law
+
+# region exponential
+
+
+def exponential_fit_powerlaw_implementation(
+    data: NDArray[np.floating[Any]],
+) -> tuple[float, float, float]:
+    """Fit exponential distribution using powerlaw library."""
+    f = io.StringIO()
+    with redirect_stdout(f), redirect_stderr(f), warnings.catch_warnings():
+        warnings.simplefilter(action="ignore", category=RuntimeWarning)
+
+        distribution = "exponential"
+        fit = pl.Fit(
+            data,
+            distribution=distribution,
+            xmin_distribution=distribution,
+            discrete=False,
+        )
+
+    fit_result = getattr(fit, distribution)
+    expon_lambda = fit_result.Lambda
+    expon_esd_xmin = fit_result.xmin
+    expon_ks = fit_result.KS()
+
+    return expon_lambda, expon_esd_xmin, expon_ks
+
+
+def exponential_fit_diffract_implementation(
+    data: NDArray[np.floating[Any]],
+) -> tuple[float, float, float]:
+    """Fit exponential distribution using the accelerated implementation."""
+    fit = _get_fit_class()(data, "exponential")
+    try:
+        params = fit.fit_params()
+    finally:
+        fit.close()
+
+    expon_lambda, expon_esd_xmin, expon_ks = itemgetter(
+        "expon_lambda", "xmin", "ks_distance"
+    )(params)
+
+    return expon_lambda, expon_esd_xmin, expon_ks
+
+
+@kernel(produce_fields=("expon_lambda", "expon_esd_xmin", "expon_ks"))
+def exponential_fit(
+    esd: NDArray[np.floating[Any]],
+    *,
+    fit_method: FitMethod = "auto",
+) -> tuple[float, float, float]:
+    """Fit exponential distribution to ESD data.
+
+    `auto` uses the accelerated implementation when the taichi extra is
+    installed and the ESD is large enough for it to fit at all.
+    """
+    match _resolve_fit_method(fit_method, esd.size):
+        case "powerlaw":
+            return exponential_fit_powerlaw_implementation(esd)
+        case "diffract":
+            return exponential_fit_diffract_implementation(esd)
+        case _:
+            msg = (
+                f"Fit method {fit_method!r} not implemented, choose method "
+                "from ('auto', 'powerlaw', 'diffract')."
+            )
+            raise ValueError(msg)
+
+
+def expon_p_value(
+    esd: NDArray[np.floating[Any]],
+    expon_lambda: float,
+    expon_esd_xmin: float,
+    expon_ks: float,
+) -> float:
+    """Compute p-value for exponential fit (requires the taichi extra)."""
+    fit = _get_fit_class()(esd, "exponential")
+    try:
+        fit.set_params(
+            xmin=expon_esd_xmin,
+            expon_lambda=expon_lambda,
+            ks_distance=expon_ks,
+            tail_size=cast("int", np.sum(esd >= expon_esd_xmin).item()),
+        )
+        _, p_value = fit.p_value_test()
+    finally:
+        fit.close()
+
+    return p_value
+
+
+# The p-value kernels are only executable with the accelerated implementation,
+# so they join the registry only when it is importable.
+if _DIFFRACT_FIT_AVAILABLE:
+    kernel(pl_p_value)
+    kernel(tpl_p_value)
+    kernel(expon_p_value)

diffract/core/compute/kernels/marchenko_pastur.py

@@ -0,0 +1,99 @@
+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 marchenko_pastur_cdf
+
+
+@kernel(produce_fields=("mp_esd_max", "mp_esd_min", "mp_bulk_std"))
+def marchenko_pastur_fit(
+    esd_rand: NDArray[np.floating[Any]],
+    esd_max: float,
+    weights_std: NDArray[np.floating[Any]],
+    aspect_ratio: float,
+    lower_dim: int,
+) -> tuple[float, float, float]:
+    """Fit Marchenko-Pastur distribution to random ESD."""
+    sigma_2_fit = weights_std**2
+    eig_max_fit = sigma_2_fit * (1 + (1 / np.sqrt(aspect_ratio))) ** 2
+
+    eigenvals_bleeding_out = esd_rand[esd_rand > eig_max_fit]
+    if eigenvals_bleeding_out.size > 0:
+        sigma_2_corrected = sigma_2_fit - (np.sum(eigenvals_bleeding_out) / lower_dim)
+        eig_max_corrected = sigma_2_corrected * (1 + (1 / np.sqrt(aspect_ratio))) ** 2
+        mp_esd_max = eig_max_corrected
+        mp_bulk_std = np.sqrt(sigma_2_corrected)
+    else:
+        mp_esd_max = eig_max_fit
+        mp_bulk_std = weights_std
+
+    mp_esd_max = min(mp_esd_max, esd_max)
+
+    mp_esd_min = (mp_bulk_std**2) * (1 - (1 / np.sqrt(aspect_ratio))) ** 2
+
+    return mp_esd_max, mp_esd_min, mp_bulk_std
+
+
+@kernel
+def mp_sval_max(mp_esd_max: float, greater_dim: int) -> float:
+    """Compute maximum singular value from MP ESD max."""
+    return np.sqrt(mp_esd_max * greater_dim)
+
+
+@kernel
+def mp_ks(
+    aspect_ratio: float,
+    mp_bulk_std: float,
+    esd: NDArray[np.floating[Any]],
+    mp_esd_max: float,
+    mp_esd_min: float,
+) -> float:
+    """Compute Kolmogorov-Smirnov distance for Marchenko-Pastur fit."""
+    ratio = 1 / (aspect_ratio + 1e-8)
+    sigma = mp_bulk_std  # we consider the fitted bulk std as the standard deviation
+
+    esd_mask = (mp_esd_min < esd) & (esd < mp_esd_max)
+    if np.sum(esd_mask).astype(int).item() == 0:
+        ks_distance = 1
+    else:
+        esd_filtered = esd[esd_mask]
+        model_cdf = marchenko_pastur_cdf(esd_filtered, ratio, sigma)
+
+        empirical_cdf = np.arange(esd_filtered.size) / esd_filtered.size
+
+        ks_distance = np.max(np.abs(empirical_cdf - model_cdf))
+
+    return ks_distance
+
+
+@kernel
+def mp_concentration(
+    esd: NDArray[np.floating[Any]], mp_esd_max: float, mp_esd_min: float
+) -> float:
+    """Compute MP bulk concentration (bulk size / total size)."""
+    mask = (mp_esd_min <= esd) & (esd <= mp_esd_max)
+    bulk_size = cast("int", np.sum(mask.astype(int)).item())
+
+    return bulk_size / esd.size
+
+
+@kernel
+def mp_presence(
+    esd_min: float,
+    esd_max: float,
+    mp_esd_max: float,
+    mp_esd_min: float,
+) -> float:
+    """Compute MP bulk presence (bulk width / total width)."""
+    esd_width = esd_max - esd_min
+    mp_width = mp_esd_max - mp_esd_min
+
+    return mp_width / esd_width
+
+
+@kernel
+def mp_num_spikes(esd: NDArray[np.floating[Any]], mp_esd_max: float) -> int:
+    """Count number of spikes above MP bulk maximum."""
+    return cast("int", np.sum((esd > mp_esd_max).astype(int)).item())

diffract/core/compute/kernels/mat_decomposition.py

@@ -0,0 +1,101 @@
+import functools
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.compute.decorator import kernel
+
+torch = import_utils.optional_import("torch")
+TORCH_AVAILABLE = torch is not None
+if TORCH_AVAILABLE:
+    from diffract.core.compute.extensions.utils import torch_cuda_wrapper
+
+    IS_CUDA_AVAILABLE = bool(torch.cuda.is_available())
+else:
+    IS_CUDA_AVAILABLE = False
+
+
+@kernel(
+    name="weights_svd",
+    require_fields=("weights",),
+    produce_fields=(
+        "weights_lsvs",
+        "weights_svals",
+        "weights_rsvs",
+    ),
+)
+@kernel(
+    name="weights_rand_svd",
+    require_fields=("weights_rand",),
+    produce_fields=(
+        "weights_rand_lsvs",
+        "weights_rand_svals",
+        "weights_rand_rsvs",
+    ),
+)
+def svd(
+    mat: NDArray[np.floating[Any]], *, allow_cuda: bool = True
+) -> tuple[
+    NDArray[np.floating[Any]], NDArray[np.floating[Any]], NDArray[np.floating[Any]]
+]:
+    """Compute SVD decomposition of matrix."""
+    if IS_CUDA_AVAILABLE and allow_cuda:
+        svd_fn = functools.partial(torch.linalg.svd, full_matrices=False)
+        u, svals, vt = torch_cuda_wrapper(mat, svd_fn)
+    else:
+        u, svals, vt = np.linalg.svd(mat, full_matrices=False)
+
+    svals = np.abs(svals.real).flatten()
+    svals_argsort_index = np.argsort(svals)
+
+    lsvs_sorted = u[:, svals_argsort_index]
+    svals_sorted = svals[svals_argsort_index]
+    rsvs_sorted = vt.T[:, svals_argsort_index]
+
+    return lsvs_sorted, svals_sorted, rsvs_sorted
+
+
+@kernel(name="esd", require_fields=("weights_svals", "greater_dim"))
+@kernel(name="esd_rand", require_fields=("weights_rand_svals", "greater_dim"))
+def esd(
+    svals: NDArray[np.floating[Any]], greater_dim: int
+) -> NDArray[np.floating[Any]]:
+    """Compute empirical spectral distribution from singular values."""
+    result = np.square(svals) / greater_dim
+    return cast("NDArray[np.floating[Any]]", result)
+
+
+# region utils
+
+
+@kernel(name="max_weights_sval", require_fields=("weights_svals",))
+@kernel(name="max_weights_rand_sval", require_fields=("weights_rand_svals",))
+def max_sval(svals: NDArray[np.floating[Any]]) -> float:
+    """Get maximum singular value."""
+    return cast("float", svals[-1].item())
+
+
+@kernel(name="min_weights_sval", require_fields=("weights_svals",))
+@kernel(name="min_weights_rand_sval", require_fields=("weights_rand_svals",))
+def min_sval(svals: NDArray[np.floating[Any]]) -> float:
+    """Get minimum singular value."""
+    return cast("float", svals[0].item())
+
+
+@kernel(name="esd_max", require_fields=("esd",))
+@kernel(name="esd_rand_max", require_fields=("esd_rand",))
+def esd_max(esd: NDArray[np.floating[Any]]) -> float:
+    """Get maximum ESD value."""
+    return cast("float", esd[-1].item())
+
+
+@kernel(name="esd_min", require_fields=("esd",))
+@kernel(name="esd_rand_min", require_fields=("esd_rand",))
+def esd_min(esd: NDArray[np.floating[Any]]) -> float:
+    """Get minimum ESD value."""
+    return cast("float", esd[0].item())
+
+
+# endregion utils

diffract/core/compute/kernels/mat_properties.py

@@ -0,0 +1,53 @@
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from diffract.core.compute.decorator import kernel
+
+
+@kernel
+def greater_dim(weights: NDArray[int]) -> int:
+    """Get greater dimension from shape."""
+    return max(weights.shape)
+
+
+@kernel
+def lower_dim(weights: NDArray[int]) -> int:
+    """Get lower dimension from shape."""
+    return min(weights.shape)
+
+
+@kernel
+def aspect_ratio(greater_dim: int, lower_dim: int) -> float:
+    """Compute aspect ratio (greater_dim / lower_dim)."""
+    return greater_dim / lower_dim
+
+
+@kernel
+def weights_std(weights: NDArray[np.floating[Any]]) -> float:
+    """Compute standard deviation of weights."""
+    return cast("float", weights.std())
+
+
+@kernel
+def weights_rand(
+    weights: NDArray[np.floating[Any]],
+    *,
+    n_randomise_iterations: int = 1,
+    seed: int = 42,
+) -> NDArray[np.floating[Any]]:
+    """Generate randomized version of weights by shuffling."""
+    rng = np.random.default_rng(None if seed == -1 else seed)
+
+    result = weights.copy()
+
+    shape = result.shape
+    index = np.arange(np.prod(shape))
+
+    result = result.reshape(-1)
+    for _ in range(n_randomise_iterations):
+        rng.shuffle(index)
+    result = result[index]
+
+    return result.reshape(shape)

diffract/core/compute/kernels/model_quality.py

@@ -0,0 +1,27 @@
+import math
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy import stats
+
+from diffract.core.compute.decorator import kernel
+
+
+@kernel
+def pl_alpha_weighted(esd_max: float, pl_alpha: float) -> float:
+    """Compute power law alpha weighted by log ESD max."""
+    return pl_alpha * math.log(esd_max)
+
+
+@kernel
+def rand_distance(
+    esd: NDArray[np.floating[Any]], esd_rand: NDArray[np.floating[Any]]
+) -> float:
+    """Compute Jensen-Shannon divergence between ESD and randomized ESD."""
+    avg = 0.5 * (esd + esd_rand)
+    divergence = 0.5 * (
+        cast("float", stats.entropy(esd, avg))
+        + cast("float", stats.entropy(esd_rand, avg))
+    )
+    return math.sqrt(divergence)