gpu/single-gpu/home_benjamin_hdbasin/code/hdbasin/objectives.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Optional

import numpy as np


Array = np.ndarray


@dataclass
class Objective:
    name: str
    dim: int
    loss: Callable[[Array], float]
    grad: Optional[Callable[[Array], Array]]
    optimum_x: Optional[Array]
    optimum_loss: Optional[float]
    metadata: dict

    def evaluate(self, x: Array, rng: Optional[np.random.Generator] = None) -> tuple[float, Optional[Array]]:
        loss = float(self.loss(np.asarray(x, dtype=float)))
        grad = None if self.grad is None else np.asarray(self.grad(np.asarray(x, dtype=float)), dtype=float)
        return loss, grad


def _orthonormal(rng: np.random.Generator, d: int, h: int) -> Array:
    q, _ = np.linalg.qr(rng.normal(size=(d, h)))
    return q[:, :h]


def convex_quadratic(dim: int, seed: int) -> Objective:
    rng = np.random.default_rng(seed)
    q, _ = np.linalg.qr(rng.normal(size=(dim, dim)))
    eigs = np.geomspace(0.5, 20.0, dim)
    a = q @ np.diag(eigs) @ q.T
    x_star = rng.uniform(0.15, 0.85, size=dim)

    def loss(x: Array) -> float:
        dx = x - x_star
        return 0.5 * float(dx @ a @ dx)

    def grad(x: Array) -> Array:
        return a @ (x - x_star)

    return Objective("convex_quadratic", dim, loss, grad, x_star, 0.0, {"condition": float(eigs[-1] / eigs[0])})


def rotated_active(dim: int, active_dim: int, seed: int, eps: float = 0.02) -> Objective:
    rng = np.random.default_rng(seed)
    u = _orthonormal(rng, dim, active_dim)
    p = u @ u.T
    z_star = rng.uniform(-0.4, 0.4, size=active_dim)
    x_anchor = np.clip(u @ z_star + 0.5, 0.0, 1.0)
    lambdas = np.geomspace(0.5, 8.0, active_dim)

    def loss(x: Array) -> float:
        z = u.T @ (x - x_anchor)
        inactive = (np.eye(dim) - p) @ (x - x_anchor)
        return float(np.sum(lambdas * z * z) + eps * np.dot(inactive, inactive))

    def grad(x: Array) -> Array:
        z = u.T @ (x - x_anchor)
        inactive = (np.eye(dim) - p) @ (x - x_anchor)
        return 2.0 * u @ (lambdas * z) + 2.0 * eps * inactive

    return Objective(
        "rotated_active",
        dim,
        loss,
        grad,
        x_anchor,
        0.0,
        {"true_active_dim": active_dim, "true_U": u, "eps": eps},
    )


def smooth_multibasin(dim: int, active_dim: int, seed: int) -> Objective:
    rng = np.random.default_rng(seed)
    u = _orthonormal(rng, dim, active_dim)
    centers = rng.uniform(-0.65, 0.65, size=(4, active_dim))
    offsets = np.array([0.25, 0.08, 0.0, 0.16])
    temperature = 0.06
    p_perp = np.eye(dim) - u @ u.T

    def basin_terms(x: Array) -> Array:
        z = u.T @ (x - 0.5)
        return offsets + np.sum((z[None, :] - centers) ** 2, axis=1)

    def loss(x: Array) -> float:
        terms = basin_terms(x)
        m = float(np.min(terms))
        soft = m - temperature * np.log(np.sum(np.exp(-(terms - m) / temperature)))
        inactive = p_perp @ (x - 0.5)
        return float(soft + 0.03 * np.dot(inactive, inactive))

    def grad(x: Array) -> Array:
        z = u.T @ (x - 0.5)
        terms = basin_terms(x)
        weights = np.exp(-(terms - np.min(terms)) / temperature)
        weights = weights / np.sum(weights)
        grad_z = np.sum(weights[:, None] * 2.0 * (z[None, :] - centers), axis=0)
        return u @ grad_z + 0.06 * (p_perp @ (x - 0.5))

    # approximate optimum by center of lowest-offset basin clipped into box
    best = int(np.argmin(offsets))
    x_star = np.clip(0.5 + u @ centers[best], 0.0, 1.0)
    return Objective("smooth_multibasin", dim, loss, grad, x_star, float(loss(x_star)), {"true_active_dim": active_dim})


def highdim_multiwell(dim: int, seed: int, active_dim: int | None = None, n_wells: int = 7) -> Objective:
    rng = np.random.default_rng(seed)
    active_dim = min(dim, active_dim or min(6, max(3, dim // 10)))
    u = _orthonormal(rng, dim, active_dim)
    p_perp = np.eye(dim) - u @ u.T
    centers = rng.uniform(-0.7, 0.7, size=(n_wells, active_dim))
    offsets = np.sort(rng.uniform(0.0, 0.35, size=n_wells))
    rng.shuffle(offsets)
    best = int(np.argmin(offsets))
    offsets[best] = 0.0
    widths = rng.uniform(0.06, 0.18, size=n_wells)
    temperature = 0.045
    ripple = rng.normal(size=(active_dim, min(4, active_dim)))
    ripple = ripple / (np.linalg.norm(ripple, axis=0, keepdims=True) + 1e-9)

    def basin_terms(x: Array) -> Array:
        z = u.T @ (x - 0.5)
        scaled = (z[None, :] - centers) / widths[:, None]
        return offsets + np.sum(scaled * scaled, axis=1) / active_dim

    def loss(x: Array) -> float:
        x = np.asarray(x, dtype=float)
        z = u.T @ (x - 0.5)
        terms = basin_terms(x)
        m = float(np.min(terms))
        soft_min = m - temperature * np.log(np.sum(np.exp(-(terms - m) / temperature)))
        inactive = p_perp @ (x - 0.5)
        ripples = 0.012 * float(np.sum(np.sin(9.0 * (ripple.T @ z))))
        return float(soft_min + 0.015 * np.dot(inactive, inactive) + ripples)

    def grad(x: Array) -> Array:
        x = np.asarray(x, dtype=float)
        z = u.T @ (x - 0.5)
        terms = basin_terms(x)
        weights = np.exp(-(terms - np.min(terms)) / temperature)
        weights = weights / np.sum(weights)
        grad_z = np.zeros(active_dim)
        for w, center, width in zip(weights, centers, widths):
            grad_z += w * 2.0 * (z - center) / (active_dim * width * width)
        grad_z += 0.108 * (ripple @ np.cos(9.0 * (ripple.T @ z)))
        return u @ grad_z + 0.03 * (p_perp @ (x - 0.5))

    x_star = np.clip(0.5 + u @ centers[best], 0.0, 1.0)
    return Objective(
        f"highdim_multiwell_d{dim}",
        dim,
        loss,
        grad,
        x_star,
        float(loss(x_star)),
        {
            "true_active_dim": active_dim,
            "true_U": u,
            "n_wells": n_wells,
            "well_offsets": offsets.tolist(),
        },
    )


def saddle_rejection(dim: int, seed: int) -> Objective:
    _ = seed

    def loss(x: Array) -> float:
        y = x - 0.5
        tail = 0.1 * float(np.sum(y[2:] ** 2)) if dim > 2 else 0.0
        return float(y[0] ** 2 - y[1] ** 2 + tail)

    def grad(x: Array) -> Array:
        y = x - 0.5
        g = np.zeros(dim)
        g[0] = 2.0 * y[0]
        g[1] = -2.0 * y[1]
        if dim > 2:
            g[2:] = 0.2 * y[2:]
        return g

    x_star = np.full(dim, 0.5)
    x_star[1] = 1.0
    return Objective("saddle_rejection", dim, loss, grad, x_star, float(loss(x_star)), {})


def boundary_minimum(dim: int, seed: int) -> Objective:
    _ = seed

    def loss(x: Array) -> float:
        return float(np.dot(x, x))

    def grad(x: Array) -> Array:
        return 2.0 * x

    return Objective("boundary_minimum", dim, loss, grad, np.zeros(dim), 0.0, {})


def noisy_objective(base: Objective, sigma: float, seed: int) -> Objective:
    rng = np.random.default_rng(seed)

    def loss(x: Array) -> float:
        return float(base.loss(x) + rng.normal(0.0, sigma))

    return Objective(
        f"noisy_{base.name}_sigma_{sigma:g}",
        base.dim,
        loss,
        None,
        base.optimum_x,
        base.optimum_loss,
        {**base.metadata, "noise_sigma": sigma, "base": base.name},
    )


def synthetic_suite(seed: int) -> list[Objective]:
    return [
        convex_quadratic(8, seed),
        rotated_active(20, 3, seed + 11),
        smooth_multibasin(12, 3, seed + 17),
        saddle_rejection(8, seed + 23),
        boundary_minimum(8, seed + 29),
        noisy_objective(convex_quadratic(8, seed + 31), 0.05, seed + 37),
    ]