example.py

from __future__ import annotations

import numpy as np
from sklearn.neighbors import KernelDensity


lower, upper = -5, 5

# Objective function for analysis
def sphere(x: np.ndarray) -> np.ndarray:
    coef = np.array([1, 5])
    x *= coef
    return np.sum(x ** 2, axis=-1)


def collect_observations(n_samples: int = 1000, dim: int = 2) -> tuple[np.ndarray, np.ndarray]:
    # Define the search space and collect observations using random search
    n_samples, dim = 1000, 2
    X = np.random.random((n_samples, dim)) * (upper - lower) + lower
    return X, sphere(X)


class PEDANOVAEvaluator:
    def __init__(self, X_train: np.ndarray, y_train: np.ndarray) -> None:
        # Define the top quantile and compute the local/global PDFs for each dimension
        # Note that we are using Eq. (14) to calculate the KDEs rather tahn Eq. (15), which is a quicker version
        quantile = 0.1
        (n_samples, dim) = X_train.shape
        X_sorted = X_train[np.argsort(y_train)]
        self.local_kdes = {
            f"x{d}": KernelDensity(kernel="gaussian").fit(X_sorted[:int(quantile * n_samples), d, None])
            for d in range(dim)
        }
        self.global_kdes = {
            f"x{d}": KernelDensity(kernel="gaussian").fit(X_sorted[:, d, None])
            for d in range(dim)
        }

    def __call__(self, X: np.ndrray) -> list[float]:
        hpi_list = []
        # Compute the global HPI of each dimension based on Eq. (16)
        for d in range(X.shape[-1]):
            log_pdf_global = self.global_kdes[f"x{d}"].score_samples(X[:, d, None])
            log_pdf_local = self.local_kdes[f"x{d}"].score_samples(X[:, d, None])
            # Equivalent to `(pdf_local / pdf_global - 1) ** 2`
            integrand = np.expm1(log_pdf_local - log_pdf_global) ** 2
            hpi_list.append(float(np.exp(log_pdf_global) @ integrand))

        hpi_sum = sum(hpi_list)
        return [hpi / hpi_sum for hpi in hpi_list]


X, y = collect_observations()
evaluator = PEDANOVAEvaluator(X, y)
# Compute the global HPI of each dimension based on Eq. (16)
dX = np.broadcast_to(np.linspace(lower, upper, 100)[:, None], (100, X.shape[-1]))
# Print the global HPI
print(evaluator(dX))