Integrating score-based diffusion models with machine learning-enhanced localization for advanced data assimilation in geological carbon storage

Accurate characterization of subsurface heterogeneity is critical for geological carbon storage (GCS). This repository provides the implementation for a framework that integrates score-based diffusion models with machine learning-enhanced localization to improve data assimilation in channelized reservoirs. By training diffusion models on geostatistical realizations and using ML proxies to estimate localization coefficients from large super-ensembles (5000 members), the approach preserves up to 40% more ensemble variance than traditional distance-based methods while maintaining geological realism. The framework addresses ensemble collapse in the Ensemble Smoother with Multiple Data Assimilation (ESMDA), particularly critical for small ensemble sizes typical in operational settings.