Atelier 1 - Effet du support

import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import VBox, HBox, IntText, Button, Output
from numpy.fft import fft2, ifft2, fftshift
from numpy.random import Generator, PCG64

# --- Configuration centralisée ---
# Regrouper les paramètres modifiables au même endroit facilite les ajustements.
CONFIG = {
    "size": 500,
    "range_short_factor": 10.0,
    "range_long_factor": 250.0,
    "sill": 1.0,
    "seed_short": 42,
    "seed_long": 24,
    "plot_vmin": 0,
    "plot_vmax": 10,
    "dtype": np.float32  # Utiliser float32 réduit l'usage mémoire de 50%
}

# --- Fonctions ---
def spherical_covariance_fft(size: int, range_val: float, sill: float = 1.0) -> np.ndarray:
    """
    Calcule la covariance sphérique sur une grille étendue pour la simulation FFT.

    Args:
        size (int): Taille de la grille de simulation finale.
        range_val (float): Portée du modèle de covariance.
        sill (float, optional): Le palier (variance). Défaut à 1.0.

    Returns:
        np.ndarray: Modèle de covariance centré via fftshift.
    """
    extended_size = 2 * size
    x = np.arange(-extended_size // 2, extended_size // 2)
    X, Y = np.meshgrid(x, x)
    h = np.sqrt(X**2 + Y**2).astype(CONFIG["dtype"])
    
    # Normalisation de la distance pour le calcul
    h_norm = h / range_val
    
    # Formule de la covariance sphérique
    cov = sill * (1 - 1.5 * h_norm + 0.5 * h_norm**3)
    
    # La covariance est nulle au-delà de la portée
    cov[h > range_val] = 0
    
    # fftshift centre la covariance pour la transformée de Fourier
    return fftshift(cov)

def fftma_simulation(
    size: int, range_val: float, rng: Generator, sill: float = 1.0
) -> np.ndarray:
    """
    Génère un champ gaussien 2D en utilisant la méthode FFT-MA.

    Args:
        size (int): Taille du champ à générer.
        range_val (float): Portée du variogramme.
        rng (Generator): Générateur de nombres aléatoires de NumPy.
        sill (float, optional): Le palier (variance). Défaut à 1.0.

    Returns:
        np.ndarray: Le champ gaussien simulé.
    """
    extended_size = 2 * size
    cov_model = spherical_covariance_fft(size, range_val, sill)
    
    # Utiliser np.abs() est une sécurité contre les imprécisions numériques
    cov_fft = np.sqrt(np.abs(fft2(cov_model)))
    
    # Génération du bruit blanc avec le dtype spécifié
    white_noise = rng.normal(size=(extended_size, extended_size)).astype(CONFIG["dtype"])
    white_fft = fft2(white_noise)
    
    # Convolution dans le domaine fréquentiel
    z_fft = cov_fft * white_fft
    z_ext = np.real(ifft2(z_fft))
    
    # Extraction de la partie centrale pour éviter les artefacts de bord
    start = extended_size // 4
    end = start + size
    return z_ext[start:end, start:end]

def gaussian_to_lognormal(field: np.ndarray) -> np.ndarray:
    """Convertit un champ gaussien en champ log-normal."""
    return np.exp(field)

def match_histogram(reference: np.ndarray, target: np.ndarray) -> np.ndarray:
    """Force l'histogramme du champ 'target' à correspondre à celui de 'reference'."""
    flat_ref = np.sort(reference.ravel())
    sorted_idx = np.argsort(target.ravel())
    result = np.zeros_like(target.ravel(), dtype=CONFIG["dtype"])
    result[sorted_idx] = flat_ref
    return result.reshape(target.shape)

def aggregate(field: np.ndarray, block_size: int) -> np.ndarray:
    """Agrége un champ par moyenne sur des blocs de taille block_size x block_size."""
    s = field.shape[0]
    # Assure que la taille du champ est un multiple de la taille du bloc
    trimmed_size = s - (s % block_size)
    trimmed = field[:trimmed_size, :trimmed_size]
    
    # Redimensionnement efficace pour calculer la moyenne par bloc
    reshaped = trimmed.reshape(
        trimmed_size // block_size, block_size,
        trimmed_size // block_size, block_size
    )
    return reshaped.mean(axis=(1, 3))

# --- Préparation des données initiales ---
# Utilisation du générateur de nombres aléatoires moderne de NumPy
rng_short = np.random.default_rng(CONFIG["seed_short"])
rng_long = np.random.default_rng(CONFIG["seed_long"])

size = CONFIG["size"]
range_short = CONFIG["range_short_factor"]
range_long = CONFIG["range_long_factor"]

gauss_short = fftma_simulation(size, range_short, rng_short, CONFIG["sill"])
gauss_long = fftma_simulation(size, range_long, rng_long, CONFIG["sill"])

lognorm_short = gaussian_to_lognormal(gauss_short)
lognorm_long = gaussian_to_lognormal(gauss_long)
lognorm_long = match_histogram(lognorm_short, lognorm_long)

# --- Affichage avec bouton ---
def plot_fields(support: int):
    """Affiche les champs, histogrammes et fonctions de répartition."""
    fig, axes = plt.subplots(2, 2, figsize=(8, 6), constrained_layout=True)
    
    agg_short = aggregate(lognorm_short, support)
    agg_long = aggregate(lognorm_long, support)
    
    # Dictionnaire pour éviter la répétition dans le traçage des images
    plot_data = {
        "short": {"data": agg_short, "ax": axes[0, 0], "title": "Portée courte"},
        "long":  {"data": agg_long, "ax": axes[0, 1], "title": "Portée longue"},
    }

    # Boucle pour tracer les deux cartes
    for key, val in plot_data.items():
        im = val["ax"].imshow(val["data"], cmap='viridis', 
                              vmin=CONFIG["plot_vmin"], vmax=CONFIG["plot_vmax"])
        val["ax"].set_title(f'{val["title"]} – Support {support}×{support}', fontsize=12)
        val["ax"].axis('off')
        fig.colorbar(im, ax=val["ax"])
    
    # Calculs pour les graphiques de distribution
    mean_val = agg_short.mean()
    bins = np.linspace(CONFIG["plot_vmin"], CONFIG["plot_vmax"], 11)

    # Histogrammes
    axes[1, 0].hist(agg_short.ravel(), bins=bins, alpha=0.7, color='steelblue', label='Portée courte', ec='black')
    axes[1, 0].hist(agg_long.ravel(), bins=bins, alpha=0.7, color='darkorange', label='Portée longue', ec='black')
    axes[1, 0].set_title('Histogrammes')

    # Fonctions de répartition empirique (ECDF)
    for data, color, label in zip([agg_short.ravel(), agg_long.ravel()],
                                  ['steelblue', 'darkorange'],
                                  ['Portée courte', 'Portée longue']):
        sorted_data = np.sort(data)
        ecdf = np.arange(1, len(sorted_data) + 1) / len(sorted_data)
        axes[1, 1].plot(sorted_data, ecdf, color=color, lw=2, label=label)
    axes[1, 1].set_title('Fonctions de répartition')

    # Paramètres communs aux deux graphiques du bas
    for ax in [axes[1, 0], axes[1, 1]]:
        ax.axvline(mean_val, color='black', linestyle='--', linewidth=2, label=f'Moyenne = {mean_val:.2f}')
        ax.set_xlim([CONFIG["plot_vmin"], CONFIG["plot_vmax"]])
        ax.grid(True)
        ax.legend()
    axes[1,1].set_ylim([0, 1])

    plt.show()

# --- Widgets ---
support_input = IntText(value=1, description="Support (m):", min=1)
calc_button = Button(description="Calculer", button_style="success")
output = Output()

def on_calculate_clicked(b):
    with output:
        output.clear_output(wait=True) # wait=True pour un affichage plus fluide
        support = support_input.value
        if support > 0:
            plot_fields(support)

calc_button.on_click(on_calculate_clicked)

# Affichage final
display(VBox([HBox([support_input, calc_button]), output]))
# Lancer un premier calcul pour l'affichage initial
on_calculate_clicked(None)