[備忘録] Optunaの探索結果をUMAPで可視化する

ハイパパラメータが3つ以上存在する場合，探索結果の可視化はなかなか難しいです．
今回はUMAPを利用して探索結果の可視化を行います．

import matplotlib.pyplot as plt

import numpy as np

import optuna

import umap


DIM = 6
SEED = 42
N_SAMPLES = 1000
N_GRID = 50
LB, UB = -5, 5


def func(X: np.ndarray) -> np.ndarray:
    return np.sum(X ** 2, axis=-1)


def objective(trial: optuna.Trial) -> float:
    X = np.asarray([trial.suggest_float(f"x{d}", LB, UB) for d in range(DIM)])
    return func(X)


def get_mapper():
    rng = np.random.RandomState(SEED)
    S = rng.random((N_SAMPLES, DIM)) * (UB - LB) + LB
    mapper = umap.UMAP(random_state=SEED).fit(S)

    S_trans = mapper.transform(S)
    x_bound = np.min(S_trans[:, 0]), np.max(S_trans[:, 0])
    y_bound = np.min(S_trans[:, 1]), np.max(S_trans[:, 1])
    return mapper, x_bound, y_bound


def plot_contour(fig, ax, mapper, x_bound, y_bound):
    X, Y = np.meshgrid(np.linspace(x_bound, N_GRID), np.linspace(y_bound, N_GRID))
    S_grid = np.stack([X.flatten(), Y.flatten()], axis=-1)
    F = func(mapper.inverse_transform(S_grid))

    level = np.linspace(0, 25 * DIM, N_GRID)
    cb = ax.contourf(X, Y, F.reshape(X.shape), level)
    fig.colorbar(cb)


study = optuna.create_study()
study.optimize(objective, n_trials=100)

mapper, x_bound, y_bound = get_mapper()
X = np.asarray([[t.params[f"x{d}"] for d in range(DIM)] for t in study.trials])
X_trans = mapper.transform(X)

fig, ax = plt.subplots()
plot_contour(fig, ax, mapper, x_bound, y_bound)
ax.scatter(X_trans[:, 0], X_trans[:, 1], color="red", s=1)
plt.show()