polystat.py

from __future__ import annotations
import math
from typing import Dict, List, Callable
import numpy as np
from numpy.polynomial import Polynomial

def _as_numpy_1d(x) -> np.ndarray:
    arr = np.asarray(x)
    return arr.ravel()

def _safe_quantile(x: np.ndarray, q: float) -> float:
    if x.size == 0:
        return np.nan
    return float(np.quantile(x, q))

def roots(p: Polynomial) -> np.ndarray:
    return p.roots()

def real_roots(z: np.ndarray, tol: float = 1e-10) -> np.ndarray:
    real = np.real(z[np.isfinite(z) & (np.abs(np.imag(z)) <= tol)])
    return np.sort(real)

def nonreal_roots(z: np.ndarray, tol: float = 1e-10) -> np.ndarray:
    return z[np.isfinite(z) & (np.abs(np.imag(z)) > tol)]

def real_gaps(rr: np.ndarray) -> np.ndarray:
    if rr.size <= 1:
        return np.array([], dtype=float)
    return np.diff(rr)

def radii(z: np.ndarray) -> np.ndarray:
    return np.abs(z)

def angles(z: np.ndarray) -> np.ndarray:
    ang = np.angle(z)
    return np.mod(ang, 2 * np.pi)

def angular_gaps(ang: np.ndarray) -> np.ndarray:
    a = np.sort(_as_numpy_1d(ang).astype(float))
    if a.size <= 1:
        return np.array([], dtype=float)
    diffs = np.diff(a)
    wrap = (a[0] + 2 * np.pi) - a[-1]
    return np.concatenate([diffs, [wrap]])

def nn_distances_complex(z: np.ndarray) -> np.ndarray:
    n = z.size
    if n < 2:
        return np.full(n, np.nan)
    out = np.empty(n, dtype=float)
    xy = np.column_stack([z.real, z.imag])
    for i in range(n):
        di = np.hypot(xy[:, 0] - xy[i, 0], xy[:, 1] - xy[i, 1])
        di[i] = np.inf
        out[i] = np.min(di)
    return out

def diameter_complex(z: np.ndarray) -> float:
    n = z.size
    if n < 2:
        return np.nan
    xy = np.column_stack([z.real, z.imag])
    dmax = 0.0
    for i in range(n - 1):
        di = np.hypot(xy[i+1:, 0] - xy[i, 0], xy[i+1:, 1] - xy[i, 1])
        dmax = max(dmax, float(np.max(di)))
    return dmax

def critical_points(p: Polynomial) -> np.ndarray:
    return p.deriv().roots()

def cp_to_root_nearest_distances(p: Polynomial, roots: np.ndarray) -> np.ndarray:
    c = p.deriv().roots()
    if roots.size == 0 or c.size == 0:
        return np.array([], dtype=float)
    out = np.empty(c.size, dtype=float)
    for i, ci in enumerate(c):
        d = np.abs(roots - ci)
        out[i] = float(np.min(d))
    return out

def basic_stats_1d(x: np.ndarray) -> Dict[str, float]:
    arr = _as_numpy_1d(x).astype(float)
    n = arr.size
    if n == 0:
        return dict(count=0, mean=np.nan, std=np.nan, var=np.nan,
                    min=np.nan, q25=np.nan, median=np.nan, q75=np.nan,
                    iqr=np.nan, max=np.nan, range=np.nan)
    mean = float(np.mean(arr))
    var = float(np.var(arr, ddof=1)) if n > 1 else np.nan
    std = math.sqrt(var) if np.isfinite(var) else np.nan
    mn = float(np.min(arr))
    q25 = _safe_quantile(arr, 0.25)
    med = _safe_quantile(arr, 0.50)
    q75 = _safe_quantile(arr, 0.75)
    mx = float(np.max(arr))
    iqr = q75 - q25 if np.all(np.isfinite([q75, q25])) else np.nan
    rng = mx - mn if np.all(np.isfinite([mx, mn])) else np.nan
    return dict(count=int(n), mean=mean, std=std, var=var,
                min=mn, q25=q25, median=med, q75=q75,
                iqr=iqr, max=mx, range=rng)

TargetFunc = Callable[[Polynomial, Dict, float, float], object]

TARGETS: Dict[str, Dict] = {}

def register_target(name: str, deps: List[str], func: TargetFunc):
    if name in TARGETS:
        raise ValueError(f"Target {name} already registered")
    TARGETS[name] = {"deps": deps, "func": func}

def resolve_target(name: str, p: Polynomial, cache: Dict, tol=1e-10, cluster_tol=1e-6):
    if name in cache:
        return cache[name]
    if name not in TARGETS:
        raise ValueError(f"Unknown target {name}")
    for dep in TARGETS[name]["deps"]:
        resolve_target(dep, p, cache, tol, cluster_tol)
    val = TARGETS[name]["func"](p, cache, tol, cluster_tol)
    cache[name] = val
    return val

def _register_stats_for_array(base: str, prefix: str):
    def stat_func(stat):
        return lambda p, c, tol, ct: basic_stats_1d(c[base])[stat]
    for stat in ["count", "mean", "std", "var", "min", "median", "max", "range"]:
        register_target(f"{prefix}_{stat}", [base], stat_func(stat))

register_target("roots", [], lambda p, c, tol, ct: roots(p))
register_target("real_roots", ["roots"], lambda p, c, tol, ct: real_roots(c["roots"], tol=tol))
register_target("nonreal_roots", ["roots"], lambda p, c, tol, ct: nonreal_roots(c["roots"], tol=tol))
register_target("critical_points", [], lambda p, c, tol, ct: critical_points(p))

# Counts
register_target("n_roots", ["roots"], lambda p, c, tol, ct: c["roots"].size)
register_target("n_real", ["real_roots"], lambda p, c, tol, ct: c["real_roots"].size)
register_target("n_nonreal", ["nonreal_roots"], lambda p, c, tol, ct: c["nonreal_roots"].size)

# Lists for secondary stats
register_target("real_gaps", ["real_roots"], lambda p, c, tol, ct: real_gaps(c["real_roots"]))
register_target("radii", ["roots"], lambda p, c, tol, ct: radii(c["roots"]))
register_target("angles", ["roots"], lambda p, c, tol, ct: angles(c["roots"]))
register_target("angular_gaps", ["angles"], lambda p, c, tol, ct: angular_gaps(c["angles"]))
register_target("nn_distances", ["roots"], lambda p, c, tol, ct: nn_distances_complex(c["roots"]))
register_target("cp_to_root_nn", ["roots", "critical_points"], lambda p, c, tol, ct: cp_to_root_nearest_distances(p, c["roots"]))

# Auto-register stats for each list
_register_stats_for_array("real_roots", "real")
_register_stats_for_array("real_gaps", "real_gap")
_register_stats_for_array("radii", "radius")
_register_stats_for_array("angles", "angle")
_register_stats_for_array("angular_gaps", "ang_gap")
_register_stats_for_array("nn_distances", "nn")
_register_stats_for_array("cp_to_root_nn", "cp_nn")

# Special geometry
register_target("diameter", ["roots"], lambda p, c, tol, ct: diameter_complex(c["roots"]))