Atelier 2 - Krigeage Simple vs Ordinaire (1D)

import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist
import ipywidgets as widgets
from IPython.display import display

# --- Grille ---
x_grid = np.linspace(0, 400, 400)

# --- Modèles de covariance ---
def exponential_cov(h, range_, sill): return sill * np.exp(-h / range_)
def gaussian_cov(h, range_, sill): return sill * np.exp(-(h / range_)**2)
def spherical_cov(h, range_, sill):
    cov = np.zeros_like(h)
    hr = h / range_
    mask = h <= range_
    cov[mask] = sill * (1 - 1.5 * hr[mask] + 0.5 * hr[mask]**3)
    return cov
def nugget_cov(h, sill): return sill * (h == 0)

# --- Matrice de covariance combinée ---
def compute_covariance_matrix(x, model1, model2=None):
    h = cdist(x[:, None], x[:, None])
    cov1 = model1(h)
    if model2:
        cov2 = model2(h)
        return cov1 + cov2
    return cov1

# --- Simulation d’un champ de référence ---
def simulate_field(x, cov_model1, cov_model2=None, mean=0):
    C = compute_covariance_matrix(x, cov_model1, cov_model2)
    L = np.linalg.cholesky(C + 1e-10 * np.eye(len(x)))  # stabilité numérique
    z = np.random.randn(len(x))
    return mean + L @ z

# --- Krigeage simple ---
def simple_kriging(x_grid, x_samples, y_samples, cov_func, mean):
    d_grid_samples = cdist(x_grid[:, None], x_samples[:, None])
    d_samples = cdist(x_samples[:, None], x_samples[:, None])
    C = cov_func(d_samples)
    c0 = cov_func(d_grid_samples)
    weights = np.linalg.solve(C, c0.T)
    y_est = mean + weights.T @ (y_samples - mean)
    sigma2 = cov_func(np.zeros(1))[0] - np.sum(weights * c0.T, axis=0)
    return y_est, sigma2, weights

# --- Krigeage ordinaire ---
def ordinary_kriging(x_grid, x_samples, y_samples, cov_func):
    d_grid_samples = cdist(x_grid[:, None], x_samples[:, None])
    d_samples = cdist(x_samples[:, None], x_samples[:, None])
    n = len(x_samples)
    C = cov_func(d_samples)
    C_aug = np.zeros((n + 1, n + 1))
    C_aug[:n, :n] = C
    C_aug[n, :n] = 1
    C_aug[:n, n] = 1

    y_est, sigma2, weights_list = [], [], []
    for i in range(len(x_grid)):
        c0 = cov_func(d_grid_samples[i])
        c_aug = np.append(c0, 1)
        sol = np.linalg.solve(C_aug, c_aug)
        weights, lam = sol[:-1], sol[-1]
        y_est.append(weights @ y_samples)
        sigma2.append(cov_func(np.zeros(1))[0] - weights @ c0 - lam)
        weights_list.append(weights)
    return np.array(y_est), np.array(sigma2), np.array(weights_list)

# --- Variables globales pour garder la simulation ---
cached_sim = {
    'params': None,
    'field': None
}

# --- Mise à jour graphique ---
def update(mean_ref, sill1, range1, model1_name,
           sill2, range2, model2_name, mean_ks, x_pred):

    models = {
        'Exponentiel': lambda r, s: lambda h: exponential_cov(h, r, s),
        'Gaussien': lambda r, s: lambda h: gaussian_cov(h, r, s),
        'Sphérique': lambda r, s: lambda h: spherical_cov(h, r, s),
        'Pépite': lambda _, s: lambda h: nugget_cov(h, s),
        'Aucun': lambda r, s: lambda h: np.zeros_like(h)
    }

    cov1 = models[model1_name](range1, sill1)
    cov2 = None if model2_name == "Aucun" else models[model2_name](range2, sill2)
    cov_func = lambda h: cov1(h) + (cov2(h) if cov2 else 0)

    current_params = (mean_ref, sill1, range1, model1_name, sill2, range2, model2_name)

    # Générer ou réutiliser le champ simulé
    if cached_sim['params'] != current_params:
        cached_sim['field'] = simulate_field(x_grid, cov1, cov2, mean_ref)
        cached_sim['params'] = current_params

    y_ref = cached_sim['field']

    # Échantillonnage fixe
    x_samples = np.array([30, 90, 145, 225, 300, 360])
    y_samples = np.interp(x_samples, x_grid, y_ref)

    # Estimations
    y_ks, var_ks, weights_ks = simple_kriging(x_grid, x_samples, y_samples, cov_func, mean_ks)
    y_ko, var_ko, weights_ko = ordinary_kriging(x_grid, x_samples, y_samples, cov_func)

    # Affichage
    fig, axs = plt.subplots(1, 2, figsize=(14, 5))

    axs[0].plot(x_grid, y_ref, '--', color='gray', label='Champ simulé (réf)')
    axs[0].plot(x_grid, y_ks, label='Krigeage simple', color='red')
    axs[0].fill_between(x_grid, y_ks - np.sqrt(var_ks), y_ks + np.sqrt(var_ks), color='red', alpha=0.2)
    axs[0].plot(x_grid, y_ko, label='Krigeage ordinaire', color='blue')
    axs[0].fill_between(x_grid, y_ko - np.sqrt(var_ko), y_ko + np.sqrt(var_ko), color='blue', alpha=0.2)
    axs[0].scatter(x_samples, y_samples, color='black', zorder=10, label='Données')
    axs[0].axvline(x_pred, color='green', linestyle='--', label=f'Estimation à x = {x_pred:.1f}')
    axs[0].set_title("Estimation et champ de référence", fontsize=13)
    axs[0].set_xlabel("Position")
    axs[0].set_ylabel("Valeur")
    axs[0].legend()
    axs[0].grid(True)

    idx_pred = np.argmin(np.abs(x_grid - x_pred))
    axs[1].stem(x_samples, weights_ks[:, idx_pred], basefmt=" ", linefmt='r-', markerfmt='ro', label='KS')
    axs[1].stem(x_samples, weights_ko[idx_pred], basefmt=" ", linefmt='b-', markerfmt='bo', label='KO')
    axs[1].set_title(f"Poids pour x = {x_pred:.1f}", fontsize=13)
    axs[1].set_xlabel("Position des données")
    axs[1].set_ylabel("Poids")
    axs[1].legend()
    axs[1].grid(True)

    plt.tight_layout()
    plt.show()

# --- Widgets et affichage ---

style = {'description_width': '150px'}

mean_ref = widgets.FloatSlider(value=0, min=-2, max=2, step=0.1, description="Moyenne champ", style=style)
sill1 = widgets.FloatSlider(value=1.0, min=0, max=2.0, step=0.1, description="Sill structure 1", style=style)
range1 = widgets.FloatSlider(value=20, min=1, max=50, step=1, description="Portée structure 1", style=style)
model1_name = widgets.Dropdown(options=['Exponentiel', 'Gaussien', 'Sphérique', 'Pépite'], value='Exponentiel',
                               description="Type structure 1", style=style)

sill2 = widgets.FloatSlider(value=0.5, min=0, max=2.0, step=0.1, description="Sill structure 2", style=style)
range2 = widgets.FloatSlider(value=10, min=1, max=50, step=1, description="Portée structure 2", style=style)
model2_name = widgets.Dropdown(options=['Aucun', 'Exponentiel', 'Gaussien', 'Sphérique', 'Pépite'], value='Aucun',
                               description="Type structure 2", style=style)

mean_ks = widgets.FloatSlider(value=0, min=-2, max=2, step=0.1, description="Moyenne (KS)", style=style)
x_pred = widgets.FloatSlider(value=50, min=0, max=400, step=1, description="x à estimer", style=style)

group_field = widgets.VBox([
    widgets.HTML("<h4>🔷 Paramètres du champ de référence</h4>"),
    mean_ref
])

group_variogram = widgets.VBox([
    widgets.HTML("<h4>🔶 Paramètres du variogramme</h4>"),
    sill1, range1, model1_name,
    sill2, range2, model2_name
])

group_kriging = widgets.VBox([
    widgets.HTML("<h4>🔸 Paramètres du krigeage</h4>"),
    mean_ks, x_pred
])

ui = widgets.VBox([group_field, group_variogram, group_kriging])

out = widgets.interactive_output(update, {
    'mean_ref': mean_ref,
    'sill1': sill1,
    'range1': range1,
    'model1_name': model1_name,
    'sill2': sill2,
    'range2': range2,
    'model2_name': model2_name,
    'mean_ks': mean_ks,
    'x_pred': x_pred
})

display(ui, out)