Integrated Mean Squared Error (IMSE) Example

This example demonstrates how to use the Integrated Mean Squared Error (IMSE) acquisition function provided in smt-optim. IMSE is often used in Stepwise Uncertainty Reduction (SUR) to globally minimize surrogate uncertainty.

import numpy as np
import matplotlib.pyplot as plt
from smt.surrogate_models import KRG
from smt_optim.acquisition_functions.integrated_variance_reduction import integrated_variance_reduction

np.random.seed(42)

# Define a 1D test function
def f(x):
    return (x * 6 - 2)**2 * np.sin(x * 12 - 4)

# Training data (3 points)
xt = np.array([[0.0], [0.5], [1.0]])
yt = f(xt)

# Train the Kriging model
sm = KRG(theta0=[1e-2], print_global=False)
sm.set_training_values(xt, yt)
sm.train()

# Define Monte-Carlo integration points over the domain [0, 1]
integration_points = np.linspace(0, 1, 100).reshape(-1, 1)

# Evaluate IMSE for candidate points across the domain
candidate_points = np.linspace(0, 1, 50).reshape(-1, 1)
imse_values = integrated_variance_reduction(
    sm, 
    candidate_points, 
    integration_points, 
    inv_block=True
)

# Plotting the result
fig, ax1 = plt.subplots(figsize=(8, 5))

# Plot the surrogate mean and variance
x_plot = np.linspace(0, 1, 200).reshape(-1, 1)
y_mean = sm.predict_values(x_plot)
y_var = sm.predict_variances(x_plot)
y_std = np.sqrt(np.maximum(y_var, 0))

ax1.plot(x_plot, y_mean, 'b-', label='Surrogate Mean')
ax1.fill_between(x_plot.flatten(), 
                 (y_mean - 3 * y_std).flatten(), 
                 (y_mean + 3 * y_std).flatten(), 
                 color='b', alpha=0.2, label='99% CI')
ax1.plot(xt, yt, 'ro', label='Training Points')
ax1.set_xlabel('x')
ax1.set_ylabel('Objective Function', color='b')
ax1.tick_params(axis='y', labelcolor='b')
ax1.legend(loc='upper left')

# Plot IMSE on a secondary axis
ax2 = ax1.twinx()
ax2.plot(candidate_points, imse_values, 'g--', label='IMSE')
ax2.set_ylabel('IMSE Value', color='g')
ax2.tick_params(axis='y', labelcolor='g')
ax2.legend(loc='upper right')

plt.title('Integrated Mean Squared Error (IMSE) Acquisition')
plt.show()

3D IMSE Convergence (Space-Filling)

This second example demonstrates how IMSE effectively explores a 3D domain space by sequentially adding points that reduce the global variance. We track the Normalized Shannon Entropy (1D projection) to show how space filling improves across iterations.

from scipy.stats import entropy
from smt.sampling_methods import Random, LHS
from scipy.optimize import minimize

np.random.seed(42)

# Define a simple 3D objective function
def f_3d(x):
    return np.sum((x - 0.5)**2, axis=1, keepdims=True)

n_runs = 2
n_start = 10
n_end = 25
n_iters = n_end - n_start

xlimits = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]])
bounds = [(0, 1), (0, 1), (0, 1)]
bins = np.linspace(0, 1, 11) # 10 bins for 1D entropy

all_runs_entropies = np.zeros((n_runs, n_iters + 1))

# We MUST use LHS for stable integration points for IMSE accuracy
integration_points = LHS(xlimits=xlimits, criterion='ese', seed=42)(250)

def get_avg_1d_entropy(points):
    entropies = []
    for dim in range(points.shape[1]):
        counts, _ = np.histogram(points[:, dim], bins=bins)
        pk = (counts + 1e-6) / np.sum(counts + 1e-6)
        # Normalize by the number of bins (10) so the max entropy is exactly 1
        entropies.append(entropy(pk, base=len(pk)))
    return np.mean(entropies)

for run in range(n_runs):
    print(f"Run {run+1}/{n_runs}...")
    # 1. Initialize Random DoE with a deterministic seed for this run
    sampling = Random(xlimits=xlimits, seed=777+run)
    xt = sampling(n_start)
    yt = f_3d(xt)
    
    all_runs_entropies[run, 0] = get_avg_1d_entropy(xt)
    
    for i in range(n_iters):
        # Train surrogate
        sm = KRG(print_global=False)
        sm.set_training_values(xt, yt)
        sm.train()
        
        # Optimize IMSE
        def obj(x):
            return -integrated_variance_reduction(
                sm, np.atleast_2d(x), integration_points=integration_points, inv_block=True
            )[0, 0]
            
        best_x = None
        best_val = np.inf
        
        # Multi-start optimization (3 random starts for speed)
        starts = LHS(xlimits=xlimits,criterion ="ese", seed=run * 100 + i)(3)
        for x0 in starts:
            res = minimize(obj, x0=x0, bounds=bounds, method='L-BFGS-B')
            if res.fun < best_val:
                best_val = res.fun
                best_x = res.x
                
        # Enrich
        xt = np.vstack((xt, best_x))
        yt = np.vstack((yt, f_3d(np.array([best_x]))))
        
        # Record entropy
        all_runs_entropies[run, i + 1] = get_avg_1d_entropy(xt)

# Average over runs
avg_entropies = np.mean(all_runs_entropies, axis=0)

# Plotting the convergence
plt.figure(figsize=(8, 5))
plt.plot(range(n_start, n_end + 1), avg_entropies, marker='o', linestyle='-', color='purple', label=f'Avg 1D Entropy ({n_runs} runs)')
plt.xlabel('Number of Points in DoE')
plt.ylabel('Normalized Shannon Entropy (10 bins)')
plt.title('IMSE 3D Space-Filling Convergence')
plt.grid(True)
plt.legend()
plt.show()

Kriging (KRG) IMSE with Adaptive Sampling Animation

This section demonstrates how IMSE can be used with a Kriging model. We use the Forrester function. The animation shows 10 iterations of adaptive sampling, in which the Design of Experiments is sequentially enriched at the point that minimizes the IMSE.

from IPython.display import HTML
from matplotlib.animation import FuncAnimation
import warnings
warnings.filterwarnings("ignore")

# Define Multi-Fidelity Forrester Function
def f_hf(x):
    return (x * 6 - 2)**2 * np.sin(x * 12 - 4)


xt_hf = np.array([[0.0], [0.5], [1.0]])
yt_hf = f_hf(xt_hf)

# Precompute true functions for plotting
x_plot = np.linspace(0, 1, 100).reshape(-1, 1)
y_true_hf = f_hf(x_plot)
integration_points = np.linspace(0, 1, 75).reshape(-1, 1)
candidate_points = np.linspace(0, 1, 150).reshape(-1, 1)

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8), gridspec_kw={'height_ratios': [2, 1]})
plt.close() # Prevent static plot from showing

# Animation update function
def update(frame):
    global xt_hf, yt_hf
    
    # 1. Train KRG Model
    sm = KRG(theta0=[1e-2], print_global=False)
    sm.set_training_values(xt_hf, yt_hf) 
    sm.train()
    
    # 2. Predict Mean and Variance
    y_mean = sm.predict_values(x_plot)
    y_var = sm.predict_variances(x_plot)
    y_std = np.sqrt(np.maximum(y_var, 0))
    
    # 3. Compute IMSE
    imse_vals = integrated_variance_reduction(
        sm, candidate_points, integration_points=integration_points, inv_block=True
    )
    
    # 4. Clear axes and plot updated data
    ax1.clear()
    ax2.clear()
    
    ax1.plot(x_plot, y_true_hf, 'k--', label='True HF')
    ax1.plot(x_plot, y_mean, 'b-', label='MFK Mean')
    ax1.fill_between(x_plot.flatten(), 
                     (y_mean - 3 * y_std).flatten(), 
                     (y_mean + 3 * y_std).flatten(), 
                     color='b', alpha=0.2, label='99% CI')
    ax1.plot(xt_hf, yt_hf, 'ro', markersize=8, label='HF Points')
    
    ax1.set_title(f'MFK Adaptive Sampling - Iteration {frame}')
    ax1.set_ylabel('Objective')
    ax1.legend(loc='upper right')
    
    ax2.plot(candidate_points, imse_vals, 'g-', label='IMSE')
    ax2.set_ylabel('IMSE', color='g')
    ax2.set_xlabel('x')
    ax2.tick_params(axis='y', labelcolor='g')
    
    # Find the next point to add (min IMSE)
    best_idx = np.argmax(imse_vals)
    next_x = candidate_points[best_idx]
    next_y = f_hf(next_x)
    
    ax2.axvline(next_x, color='r', linestyle='--', label='Next Point')
    ax2.legend(loc='upper right')
    
    # Enrich HF design for the *next* frame
    if frame < 10:
        xt_hf = np.vstack((xt_hf, next_x))
        yt_hf = np.vstack((yt_hf, next_y))

# Create animation (6 frames: initial + 5 additions)
anim = FuncAnimation(fig, update, frames=11, interval=1500, repeat_delay=3000)
HTML(anim.to_jshtml())