Source code for denia.models

"""Surrogates and acquisition functions - the classical, reproducible core.

Two surrogates cover the practical range of tabular campaigns:

* :class:`EnsembleSurrogate` (default) - random-forest ensemble; uncertainty
  from across-tree dispersion. Robust to mixed/sparse features and scales to
  thousands of rows (the regime of literature-compiled datasets like OCM).
* :class:`GPSurrogate` - Gaussian process (Matérn 5/2); preferable for small
  in-house campaigns (tens to a few hundred rows).

Acquisition: expected improvement (EI) and upper confidence bound (UCB), plus
greedy batch selection with a diversity penalty so a recommended batch does
not collapse onto near-duplicates of the same candidate.
"""

from __future__ import annotations

import math
from typing import Tuple

import numpy as np

try:  # prefer scipy if available; keep a dependency-free fallback
    from scipy.stats import norm

    def norm_cdf(x):  # noqa: F811
        return norm.cdf(x)

    def norm_pdf(x):  # noqa: F811
        return norm.pdf(x)

except Exception:  # pragma: no cover
    def norm_cdf(x):  # noqa: F811
        return 0.5 * (1.0 + np.vectorize(math.erf)(np.asarray(x) / np.sqrt(2.0)))

    def norm_pdf(x):  # noqa: F811
        return np.exp(-0.5 * np.asarray(x) ** 2) / np.sqrt(2 * np.pi)

from sklearn.ensemble import RandomForestRegressor
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import ConstantKernel, Matern, WhiteKernel




class EnsembleSurrogate:
    """Random-forest surrogate with across-tree predictive dispersion."""

    def __init__(self, n_estimators: int = 200, random_state: int = 0):
        self.model = RandomForestRegressor(
            n_estimators=n_estimators, random_state=random_state, n_jobs=-1,
            min_samples_leaf=2,
        )

    def fit(self, X: np.ndarray, y: np.ndarray) -> "EnsembleSurrogate":
        mask = np.isfinite(y)
        self.model.fit(X[mask], y[mask])
        return self

    def predict(self, X: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        preds = np.stack([t.predict(X) for t in self.model.estimators_], axis=0)
        return preds.mean(axis=0), preds.std(axis=0) + 1e-9





class GPSurrogate:
    """Gaussian-process surrogate (Matérn 5/2 + noise) for small campaigns."""

    def __init__(self, random_state: int = 0):
        kernel = ConstantKernel(1.0) * Matern(nu=2.5) + WhiteKernel(1e-3)
        self.model = GaussianProcessRegressor(
            kernel=kernel, normalize_y=True, random_state=random_state, alpha=1e-8,
        )

    def fit(self, X: np.ndarray, y: np.ndarray) -> "GPSurrogate":
        mask = np.isfinite(y)
        self.model.fit(X[mask], y[mask])
        return self

    def predict(self, X: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        mu, std = self.model.predict(X, return_std=True)
        return mu, std + 1e-9



# --------------------------------------------------------------------------- acquisition
def expected_improvement(mu: np.ndarray, std: np.ndarray, best: float, xi: float = 0.01) -> np.ndarray:
    imp = mu - best - xi
    z = imp / std
    return imp * norm_cdf(z) + std * norm_pdf(z)


def upper_confidence_bound(mu: np.ndarray, std: np.ndarray, beta: float = 2.0) -> np.ndarray:
    return mu + beta * std




def select_batch(
    scores: np.ndarray,
    X: np.ndarray,
    batch_size: int,
    diversity: float = 0.3,
) -> np.ndarray:
    """Greedy batch selection with a squared-exponential diversity penalty.

    After each pick, candidate scores are down-weighted by proximity to the
    already-selected points (feature-space distance), preventing batches of
    near-duplicates - the practical failure mode of naive top-k selection on
    literature datasets full of replicated conditions.
    """
    scores = scores.astype(float).copy()
    n = len(scores)
    batch_size = min(batch_size, n)
    chosen = []
    scale = np.median(np.linalg.norm(X - X.mean(axis=0), axis=1)) + 1e-9
    for _ in range(batch_size):
        i = int(np.nanargmax(scores))
        chosen.append(i)
        scores[i] = -np.inf
        if diversity > 0:
            d = np.linalg.norm(X - X[i], axis=1) / scale
            active = np.isfinite(scores)
            scores[active] = scores[active] - diversity * np.exp(-(d[active] ** 2)) * np.abs(scores[active])
    return np.array(chosen)