volatility_forecast_engine.py

"""Array-oriented primitives for constituent volatility forecast evaluation."""

from __future__ import annotations

from collections.abc import Mapping

import numpy as np
import pandas as pd


PRIOR_ARRAY_COLUMNS = (
    "prior_source",
    "peer_count",
    "peer_mean",
    "peer_observed_variance",
    "peer_mean_estimation_variance",
    "estimated_within_group_variance",
    "empirical_own_weight",
    "empirical_estimate",
)


def leave_one_out_group_priors_array(
    estimates: np.ndarray,
    estimation_variances: np.ndarray,
    group_codes: np.ndarray,
    group_count: int,
    min_group_peers: int,
) -> dict[str, np.ndarray]:
    """Array equivalent of leave_one_out_group_priors.

    Invalid rows remain NaN. Groups with too few peers fall back to the full
    usable universe, matching the DataFrame implementation.
    """
    estimates = np.asarray(estimates, dtype=float)
    estimation_variances = np.asarray(estimation_variances, dtype=float)
    group_codes = np.asarray(group_codes, dtype=np.int64)
    if not (
        estimates.shape == estimation_variances.shape == group_codes.shape
    ):
        raise ValueError("prior input arrays must have identical shapes")

    size = estimates.size
    numeric = {
        column: np.full(size, np.nan)
        for column in PRIOR_ARRAY_COLUMNS
        if column != "prior_source"
    }
    sources = np.full(size, None, dtype=object)
    usable = (
        np.isfinite(estimates)
        & np.isfinite(estimation_variances)
        & (estimation_variances >= 0)
        & (group_codes >= 0)
    )
    if not np.any(usable):
        return {"prior_source": sources, **numeric}

    usable_positions = np.flatnonzero(usable)
    usable_groups = group_codes[usable]
    usable_estimates = estimates[usable]
    usable_uncertainties = estimation_variances[usable]
    counts = np.bincount(usable_groups, minlength=group_count).astype(float)
    sums = np.bincount(
        usable_groups, weights=usable_estimates, minlength=group_count
    )
    sums_squared = np.bincount(
        usable_groups, weights=usable_estimates**2, minlength=group_count
    )
    uncertainty_sums = np.bincount(
        usable_groups, weights=usable_uncertainties, minlength=group_count
    )
    total_count = float(usable_positions.size)
    total_sum = float(usable_estimates.sum())
    total_sum_squared = float(np.dot(usable_estimates, usable_estimates))
    total_uncertainty = float(usable_uncertainties.sum())

    peer_counts = counts[usable_groups] - 1.0
    peer_sums = sums[usable_groups] - usable_estimates
    peer_sums_squared = sums_squared[usable_groups] - usable_estimates**2
    peer_uncertainty_sums = (
        uncertainty_sums[usable_groups] - usable_uncertainties
    )
    use_universe = peer_counts < min_group_peers
    peer_counts[use_universe] = total_count - 1.0
    peer_sums[use_universe] = total_sum - usable_estimates[use_universe]
    peer_sums_squared[use_universe] = (
        total_sum_squared - usable_estimates[use_universe] ** 2
    )
    peer_uncertainty_sums[use_universe] = (
        total_uncertainty - usable_uncertainties[use_universe]
    )
    enough = peer_counts >= min_group_peers
    positions = usable_positions[enough]
    if positions.size == 0:
        return {"prior_source": sources, **numeric}

    count = peer_counts[enough]
    estimate_sum = peer_sums[enough]
    estimate_sum_squared = peer_sums_squared[enough]
    uncertainty_sum = peer_uncertainty_sums[enough]
    peer_mean = estimate_sum / count
    observed_variance = np.full(count.size, np.nan)
    multiple = count > 1
    observed_variance[multiple] = (
        estimate_sum_squared[multiple]
        - estimate_sum[multiple] ** 2 / count[multiple]
    ) / (count[multiple] - 1.0)
    mean_uncertainty = uncertainty_sum / count
    signal_variance = np.where(
        np.isfinite(observed_variance),
        np.maximum(0.0, observed_variance - mean_uncertainty),
        0.0,
    )
    own_uncertainty = estimation_variances[positions]
    denominator = signal_variance + own_uncertainty
    own_weight = np.divide(
        signal_variance,
        denominator,
        out=np.zeros_like(signal_variance),
        where=denominator > 0,
    )

    sources[positions] = np.where(use_universe[enough], "universe", "sector")
    numeric["peer_count"][positions] = count
    numeric["peer_mean"][positions] = peer_mean
    numeric["peer_observed_variance"][positions] = observed_variance
    numeric["peer_mean_estimation_variance"][positions] = mean_uncertainty
    numeric["estimated_within_group_variance"][positions] = signal_variance
    numeric["empirical_own_weight"][positions] = own_weight
    numeric["empirical_estimate"][positions] = (
        own_weight * estimates[positions] + (1.0 - own_weight) * peer_mean
    )
    return {"prior_source": sources, **numeric}


class ArrayResultBuilder:
    """Collect equally sized column arrays and build one DataFrame at the end."""

    def __init__(self) -> None:
        self._parts: dict[str, list[np.ndarray]] = {}
        self._rows = 0

    def append(self, columns: Mapping[str, object], rows: int) -> None:
        if rows <= 0:
            return
        if self._parts and set(columns) != set(self._parts):
            missing = set(self._parts) - set(columns)
            extra = set(columns) - set(self._parts)
            raise ValueError(
                f"result columns changed; missing={sorted(missing)}, "
                f"extra={sorted(extra)}"
            )
        for column, values in columns.items():
            if np.isscalar(values) or values is None:
                array = np.full(rows, values)
            else:
                array = np.asarray(values)
                if array.ndim != 1 or array.size != rows:
                    raise ValueError(
                        f"column {column!r} has shape {array.shape}; "
                        f"expected ({rows},)"
                    )
            self._parts.setdefault(column, []).append(array)
        self._rows += rows

    def to_frame(self) -> pd.DataFrame:
        if not self._rows:
            return pd.DataFrame()
        return pd.DataFrame(
            {column: np.concatenate(parts) for column, parts in self._parts.items()}
        )