soc-pipeline¶

def fit_clauset_powerlaw(
    x_data: np.ndarray,
    name: str = "values",
    discrete: bool = False,
    xmin_method: str = "ks",
    min_samples: int = 100,
) -> FitResult:
    """Fit a power-law to the tail of x_data using the Clauset 2009 method.

    Args:
        x_data: 1-D array of positive observed sizes/durations.
        name: Caller-provided label for logging / round-tripping.
        discrete: True for integer event counts, False for continuous sizes.
        xmin_method: Currently only 'ks' (Kolmogorov-Smirnov minimization).
        min_samples: Minimum sample size; below this, the function returns a
            FitResult with an error message and unset fields.

    Returns:
        FitResult dataclass.

    Notes:
        Requires the `powerlaw` package (Alstott et al. 2014). If the package
        is not installed, FitResult.error is set and other fields are None.
    """
    try:
        import powerlaw  # type: ignore
    except Exception as exc:  # pragma: no cover - import-time only
        return FitResult(name=name, error=f"powerlaw missing: {exc}")

    if xmin_method != "ks":
        return FitResult(name=name, error=f"xmin_method {xmin_method} not supported")

    x_data = np.asarray(x_data, dtype=float)
    x_data = x_data[np.isfinite(x_data) & (x_data > 0)]
    n_total = int(len(x_data))
    if n_total < min_samples:
        return FitResult(
            name=name,
            n_total=n_total,
            error=f"too few values: {n_total} < {min_samples}",
        )

    fit = powerlaw.Fit(x_data, discrete=discrete, xmin_distance="D", verbose=False)
    alpha = float(fit.power_law.alpha)
    sigma = float(fit.power_law.sigma)
    xmin = float(fit.power_law.xmin)
    n_tail = int(np.sum(x_data >= xmin))
    try:
        ks = float(fit.power_law.D)
    except Exception:
        ks = None

    try:
        R_ln, p_ln = fit.distribution_compare(
            "power_law", "lognormal", normalized_ratio=True
        )
    except Exception:
        R_ln, p_ln = (None, None)
    try:
        R_exp, p_exp = fit.distribution_compare(
            "power_law", "exponential", normalized_ratio=True
        )
    except Exception:
        R_exp, p_exp = (None, None)

    rejects = False
    if R_exp is not None and R_exp < 0:
        rejects = True
    if R_ln is not None and R_ln < 0:
        rejects = True

    if R_ln is None:
        winner = "inconclusive"
    elif R_ln > 0 and (p_ln is None or p_ln < 0.1):
        winner = "power_law"
    elif R_ln < 0 and (p_ln is None or p_ln < 0.1):
        winner = "lognormal"
    else:
        winner = "inconclusive"

    return FitResult(
        alpha=alpha,
        xmin=xmin,
        sigma=sigma,
        n_total=n_total,
        n_tail=n_tail,
        ks_statistic=ks,
        vs_lognormal_R=None if R_ln is None else float(R_ln),
        vs_lognormal_p=None if p_ln is None else float(p_ln),
        vs_exponential_R=None if R_exp is None else float(R_exp),
        vs_exponential_p=None if p_exp is None else float(p_exp),
        vs_powerlaw_lognormal_winner=winner,
        rejects_power_law=bool(rejects),
        name=name,
    )

def bootstrap_ci(
    x_data: np.ndarray,
    n_boot: int = 200,
    seed: int = 42,
    discrete: bool = False,
    ci_pct: tuple[float, float] = (2.5, 97.5),
    min_samples: int = 200,
) -> BootstrapResult:
    """Bootstrap a CI on the Clauset 2009 alpha estimator.

    Args:
        x_data: 1-D positive sample.
        n_boot: Number of resamples (default 200; >=200 recommended).
        seed: RNG seed for reproducibility.
        discrete: Pass-through to powerlaw.Fit.
        ci_pct: Lower/upper percentile pair for the CI (default 95% CI).
        min_samples: Minimum sample size to attempt bootstrap.

    Returns:
        BootstrapResult dataclass. If fewer than 20 fits succeed, returns
        a result with error set.
    """
    try:
        import powerlaw  # type: ignore
    except Exception as exc:
        return BootstrapResult(error=f"powerlaw missing: {exc}")

    x_data = np.asarray(x_data, dtype=float)
    x_data = x_data[np.isfinite(x_data) & (x_data > 0)]
    if len(x_data) < min_samples:
        return BootstrapResult(error=f"too few values: {len(x_data)} < {min_samples}")
    rng = np.random.default_rng(seed)
    n = len(x_data)
    alphas: list[float] = []
    for _ in range(n_boot):
        sample = rng.choice(x_data, size=n, replace=True)
        try:
            f = powerlaw.Fit(sample, discrete=discrete, xmin_distance="D", verbose=False)
            alphas.append(float(f.power_law.alpha))
        except Exception:
            continue
    if len(alphas) < 20:
        return BootstrapResult(
            n_boot_succeeded=len(alphas),
            error=f"only {len(alphas)}/{n_boot} fits succeeded",
        )
    arr = np.asarray(alphas)
    return BootstrapResult(
        alpha_mean=float(arr.mean()),
        alpha_median=float(np.median(arr)),
        alpha_std=float(arr.std()),
        ci_low=float(np.percentile(arr, ci_pct[0])),
        ci_high=float(np.percentile(arr, ci_pct[1])),
        n_boot_succeeded=int(len(arr)),
    )

def vuong_lr_test(
    x_data: np.ndarray,
    vs: Distribution = "lognormal",
    discrete: bool = False,
    alpha_threshold: float = 0.1,
) -> LRResult:
    """Run Vuong 1989 LR test on power-law vs alternative.

    Args:
        x_data: Positive 1-D sample.
        vs: Alternative distribution ('lognormal' / 'exponential' /
            'stretched_exponential' / 'truncated_power_law').
        discrete: Pass-through to powerlaw.Fit.
        alpha_threshold: p-value cutoff for declaring a winner.

    Returns:
        LRResult dataclass.
    """
    try:
        import powerlaw  # type: ignore
    except Exception as exc:
        return LRResult(vs=vs, error=f"powerlaw missing: {exc}")

    x_data = np.asarray(x_data, dtype=float)
    x_data = x_data[np.isfinite(x_data) & (x_data > 0)]
    if len(x_data) < 100:
        return LRResult(vs=vs, error=f"too few values: {len(x_data)}")
    try:
        fit = powerlaw.Fit(x_data, discrete=discrete, xmin_distance="D", verbose=False)
        R, p = fit.distribution_compare("power_law", vs, normalized_ratio=True)
    except Exception as exc:
        return LRResult(vs=vs, error=f"lr test failed: {exc}")

    winner = "inconclusive"
    if R is not None and p is not None and p < alpha_threshold:
        winner = "power_law" if R > 0 else vs

    return LRResult(vs=vs, R=float(R) if R is not None else None,
                    p=float(p) if p is not None else None, winner=winner)

def synthetic_null(
    kind: NullKind | None = None,
    n: int = 20_000,
    seed: int = 42,
) -> dict[str, NullCase] | NullCase:
    """Run the pipeline on one or all synthetic nulls.

    Args:
        kind: If None, run all three null kinds (gaussian_walk, exponential,
            poisson_iat) and return a dict keyed by name. If a specific kind,
            return a single NullCase.
        n: Sample size per null.
        seed: RNG seed.

    Returns:
        Either a dict[str, NullCase] (when kind is None) or a single NullCase.

    Notes:
        Use the dict form for the standard validation suite. The returned
        cases each carry a `correctly_rejected` flag — pipeline is healthy
        iff all three reject power-law.
    """
    rng = np.random.default_rng(seed)
    if kind is not None:
        sample = _generate(kind, n, rng)
        fit = fit_clauset_powerlaw(sample, name=kind)
        return NullCase(name=kind, fit=fit, correctly_rejected=bool(fit.rejects_power_law))

    out: dict[str, NullCase] = {}
    for k in ("gaussian_walk", "exponential", "poisson_iat"):
        sample = _generate(k, n, rng)  # type: ignore[arg-type]
        fit = fit_clauset_powerlaw(sample, name=k)
        out[k] = NullCase(
            name=k,  # type: ignore[arg-type]
            fit=fit,
            correctly_rejected=bool(fit.rejects_power_law),
        )
    return out

def shape_normalized_collapse(
    samples: dict[str, tuple[np.ndarray, float]],
    s_star_percentile: float = 99.0,
    n_points: int = 120,
) -> CollapseResult:
    """Compute shape-normalized collapse curves across systems.

    Args:
        samples: dict mapping system name -> (event_sizes, known_alpha).
            event_sizes: 1-D positive array. known_alpha: scaling exponent
            from Clauset fit (or literature value).
        s_star_percentile: Percentile used to choose system-specific cutoff s*.
            99 means s* = 99th percentile (robust to outliers).
        n_points: Number of log-spaced CCDF grid points per system.

    Returns:
        CollapseResult with one SystemCurve per system.
    """
    out: dict[str, SystemCurve] = {}
    alpha_min, alpha_max = float("inf"), float("-inf")
    for name, (vals, alpha) in samples.items():
        vals = np.asarray(vals, dtype=float)
        vals = vals[np.isfinite(vals) & (vals > 0)]
        if len(vals) < 100:
            continue
        grid, ccdf = empirical_ccdf(vals, n_points=n_points)
        if grid is None or ccdf is None:
            continue
        s_star = float(np.percentile(vals, s_star_percentile))
        x_rescaled = grid / s_star
        y_rescaled = (s_star ** (alpha - 1)) * ccdf
        out[name] = SystemCurve(
            name=name,
            alpha=float(alpha),
            s_star=s_star,
            n=int(len(vals)),
            x_rescaled=np.asarray(x_rescaled),
            y_rescaled=np.asarray(y_rescaled),
            raw_grid=np.asarray(grid),
            raw_ccdf=np.asarray(ccdf),
        )
        alpha_min = min(alpha_min, float(alpha))
        alpha_max = max(alpha_max, float(alpha))

    return CollapseResult(
        systems=out,
        alpha_range=(alpha_min, alpha_max) if out else None,
        n_systems=len(out),
        notes=(
            "Strict alpha collapse fails across different observables. "
            "Functional-form collapse (power-law tail with exponential cutoff) succeeds."
        ),
    )

def fit_omori_p(
    dts_sec: np.ndarray,
    min_sec: float = 300.0,
    max_sec: float = 30 * 86400,
    n_bins: int = 24,
    c_grid_days: tuple[float, ...] = (0.001, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5),
) -> OmoriResult:
    """Omori-Utsu fit on a stack of aftershock times after a main shock.

    Args:
        dts_sec: Delay (seconds) of each aftershock after its main shock.
        min_sec: Lower edge of fit window (default 5 min).
        max_sec: Upper edge of fit window (default 30 days).
        n_bins: Number of log-spaced time bins.
        c_grid_days: Grid for the c (time-shift) parameter; best by R^2 wins.

    Returns:
        OmoriResult dataclass.
    """
    dts = np.asarray(dts_sec, dtype=float)
    if len(dts) < 100:
        return OmoriResult(error=f"too few aftershocks ({len(dts)})")

    dts = dts[(dts >= min_sec) & (dts <= max_sec)]
    if len(dts) == 0:
        return OmoriResult(error="no aftershocks in [min_sec, max_sec]")

    bins = np.logspace(math.log10(min_sec), math.log10(max_sec), n_bins + 1)
    centers = np.sqrt(bins[:-1] * bins[1:])
    widths = bins[1:] - bins[:-1]
    counts, _ = np.histogram(dts, bins=bins)
    rate_per_sec = counts / widths
    valid = counts >= 3
    if valid.sum() < 6:
        return OmoriResult(error=f"too few valid bins ({int(valid.sum())})")

    t_sec = centers[valid]
    r = rate_per_sec[valid]
    w = np.sqrt(counts[valid].astype(float))

    best: dict[str, Any] | None = None
    for c_day in c_grid_days:
        c_sec = c_day * 86400
        x = np.log10(t_sec + c_sec)
        y = np.log10(r)
        W = w ** 2
        Sw = np.sum(W)
        Sx = np.sum(W * x)
        Sy = np.sum(W * y)
        Sxx = np.sum(W * x * x)
        Sxy = np.sum(W * x * y)
        denom = Sw * Sxx - Sx * Sx
        if denom == 0:
            continue
        slope = (Sw * Sxy - Sx * Sy) / denom
        intercept = (Sy - slope * Sx) / Sw
        p_val = -slope
        logK = intercept
        pred_y = intercept + slope * x
        ss_res = float(np.sum(W * (y - pred_y) ** 2))
        ss_tot = float(np.sum(W * (y - np.average(y, weights=W)) ** 2))
        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else None
        residuals = y - pred_y
        dof = max(len(y) - 2, 1)
        mse = float(np.sum(W * residuals ** 2) / dof)
        var_slope = mse * Sw / denom
        sigma_p = float(math.sqrt(abs(var_slope)))
        if best is None or (r2 is not None and r2 > best["R2"]):
            best = {
                "c_days": c_day,
                "p": float(p_val),
                "p_sigma": sigma_p,
                "logK": float(logK),
                "R2": float(r2) if r2 is not None else None,
            }

    if best is None:
        return OmoriResult(error="no valid Omori fit")

    return OmoriResult(
        p=best["p"],
        p_sigma=best["p_sigma"],
        c_days=best["c_days"],
        logK=best["logK"],
        R2=best["R2"],
        n_aftershocks_in_fit=int(np.sum(counts[valid])),
        n_bins_used=int(valid.sum()),
        t_range_days=(float(t_sec.min() / 86400), float(t_sec.max() / 86400)),
    )