Add constrained VaR

[nabenabe0928]

Contributor Agreements

Please read the contributor agreements and if you agree, please click the checkbox below.

• I agree to the contributor agreements.

[nabenabe0928]

from __future__ import annotations

from typing import Any

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import optuna
import optunahub


CONSTRAINTS_KEY = "constraints"


def numpy_objective(X: np.ndarray) -> np.ndarray:
    C = np.asarray([[2.0, -12.2, 21.2, -6.4, -4.7, 6.2], [1.0, -11, 43.3, -74.8, 56.9, -10]])
    X_poly = np.zeros_like(X)
    for i in range(C.shape[1]):
        X_poly = X * (X_poly + C[:, i])
    X0X1 = X[..., 0] * X[..., 1]
    out = np.sum(X_poly, axis=-1)
    out += X0X1 * (-4.1 - 0.1 * X0X1 + 0.4 * X[..., 0] + 0.4 * X[..., 1])
    return out


def constraint_func0(X: np.ndarray) -> np.ndarray:
    return (X[..., 0] - 1.5)**4 + (X[..., 1] - 1.5)**4 - 10.125


def constraint_func1(X: np.ndarray) -> np.ndarray:
    return (X[..., 0] - 2.5)**3 - (X[..., 1] + 1.5)**3 + 15.75


def objective(trial: optuna.Trial) -> float:
    x = trial.suggest_float("x", -1, 4)
    y = trial.suggest_float("y", -1, 4)
    X = np.array([x, y])
    f = numpy_objective(X).item()
    c0 = constraint_func0(X).item()
    c1 = constraint_func1(X).item()
    trial.set_user_attr(CONSTRAINTS_KEY, (c0, c1))
    return f


def constraints(trial: optuna.trial.FrozenTrial) -> tuple[float, float]:
    return trial.user_attrs[CONSTRAINTS_KEY]


def optimize(n_trials: int, seed: int) -> optuna.Study:
    RobustGPSampler = optunahub.load_local_module("samplers/value_at_risk", registry_root="./package").RobustGPSampler
    sampler = RobustGPSampler(seed=seed, uniform_input_noise_rads={"x": 0.5, "y": 0.5}, constraints_func=constraints)
    study = optuna.create_study(sampler=sampler)
    study.optimize(objective, n_trials=n_trials)
    return study


def plot_landscape(ax: plt.Axes) -> plt.Axes:
    dx = np.linspace(-1, 4, 100)
    XY = np.stack(np.meshgrid(dx, dx), axis=-1)
    Z = numpy_objective(XY)
    Z = np.log((Z - Z.min()) / (Z.max() - Z.min()) + 1e-12)
    ax.contourf(XY[..., 0], XY[..., 1], Z, levels=50, vmin=np.percentile(Z, 10), vmax=np.percentile(Z, 90))
    c = ((constraint_func0(XY) <= 0) & (constraint_func1(XY) <= 0)).astype(float)
    ax.pcolormesh(XY[..., 0], XY[..., 1], c, alpha=0.2)
    ax.set_xlim(dx[0], dx[-1])
    ax.set_ylim(dx[0], dx[-1])
    return ax


def main(n_trials: int, seed: int) -> None:
    plt.rcParams["font.family"] = "Times New Roman"
    plt.rcParams["font.size"] = 18
    _, ax = plt.subplots()
    ax = plot_landscape(ax)
    study = optimize(n_trials=n_trials, seed=seed)
    robust_params = study.sampler.get_robust_trial(study).params
    XY = np.asarray([[t.params["x"], t.params["y"]] for t in study.trials])
    ax.scatter(XY[:, 0], XY[:, 1], c="darkred", ec="white")
    x, y = robust_params["x"], robust_params["y"]
    ax.scatter(x, y, c="yellow", ec="black", label="Best trial")
    r = mpatches.Rectangle(xy=(x - 0.5, y - 0.5), width=1, height=1, ec="cyan", lw=5, fill=False)
    ax.add_patch(r)
    plt.savefig(f"example-var{seed}.png", bbox_inches="tight")


if __name__ == "__main__":
    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument("--n-trials", type=int, default=30)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()
    main(n_trials=args.n_trials, seed=args.seed)

[nabenabe0928]

The current version is apparently weighting the raw constraint boundary too much.

[nabenabe0928]

Due to the nature of the acquisition function, which is just a posterior mean of VaR in our case, the convergence is observed, but the acquisition function now captures the promising region very well:)

[c-bata]

@.not522 Could you review this PR? Let me reassign another reviewer later.

[msakai]

LGTM!