CorrDiffCMIP6 integration

[gertln]

Earth2Studio Pull Request

Summary

CorrDiffCMIP6 integration for CMIP6 downscaling: model wrapper + exports + numbered example + tests

Changes

• Add CorrDiffCMIP6 CMIP6-specific wrapper implementation in earth2studio/models/dx/corrdiff.py
  • CMIP6 preprocessing (time features + channel handling) and postprocessing
  • Optional multi-sample CPU streaming and optional tqdm progress for ensembles
• Export CorrDiffCMIP6 via earth2studio/models/dx/__init__.py
• Add numbered example examples/20_corrdiff_cmip6.py
• Extend tests in test/models/dx/test_corrdiff.py with focused CorrDiffCMIP6 coverage
  • preprocessing + required time handling
  • multi-sample streaming
  • postprocess crop/clamp/flip behavior

Notes / Follow-ups

• CorrDiffCMIP6.load_default_package() uses a placeholder URI until the final package URI is configured.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[gertln]

Plots

The example python examples/20_corrdiff_cmip6.py writes two figures:

• 20_corrdiff_cmip6_quicklook.png: 2×2 map comparison of coarse CMIP6 input vs CorrDiffCMIP6 downscaled output for
  • temperature (CMIP6 tas vs CorrDiff t2m)
  • 10 m wind speed (ws10m = sqrt(u²+v²)), computed from CMIP6 uas/vas and from downscaled u10m/v10m
• 20_corrdiff_cmip6_ws10m_ensemble_3x3.png: 3×3 grid showing 7 stochastic members of downscaled ws10m plus ensemble mean and ensemble std.

Single-sample comparison

Ensemble ws10m (members/mean/std)

[greptile-apps]

Disclaimer: This is AI-generated, please review response for accuracy

Greptile Summary

Adds CorrDiffCMIP6, a CMIP6-specialized variant of the CorrDiff downscaling model with time-dependent preprocessing (solar zenith angle, hour-of-day features) and CMIP6-specific handling (sea-air-ice cover smoothing, non-negativity constraints, relaxed grid validation).

Key changes:

• New CorrDiffCMIP6 class (864 lines) extending CorrDiff with time-aware preprocessing and optional CPU streaming for multi-sample inference
• create_time_window_wrapper() helper method for wrapping datasources with the model's time window metadata
• Comprehensive example (20_corrdiff_cmip6.py) demonstrating usage with CMIP6 datasources
• Test coverage for preprocessing validation, time handling, CPU streaming, and postprocessing constraints
• Export updates in __init__.py and CHANGELOG entries

Implementation notes:

• The channel reordering logic intentionally replicates a training quirk where hour-of-day values were placed in the coslat position but normalized with coslat statistics
• load_default_package() placeholder URI raises NotImplementedError (acknowledged in PR description)
• Example file calls the unimplemented load_default_package() method, preventing execution until package URI is configured

Confidence Score: 4/5

• Safe to merge after fixing example to avoid calling unimplemented method
• The implementation is well-tested and documented, but the example file calls load_default_package() which raises NotImplementedError, preventing the example from running. Once this is addressed (either by commenting out the call or providing a fallback), the PR should be ready to merge.
• examples/20_corrdiff_cmip6.py requires updating to avoid calling the unimplemented load_default_package() method

Important Files Changed

Filename	Overview
earth2studio/models/dx/corrdiff.py	Added CorrDiffCMIP6 class (864 lines) with time-dependent preprocessing, CPU streaming, and CMIP6-specific handling
examples/20_corrdiff_cmip6.py	Comprehensive example demonstrating CorrDiffCMIP6 usage with CMIP6 datasources and visualization
test/models/dx/test_corrdiff.py	Added 269 lines of CorrDiffCMIP6 tests covering preprocessing, sampling, postprocessing, and edge cases

[greptile-apps]

Additional Comments (1)

1. examples/20_corrdiff_cmip6.py, line 62 (link)

   syntax: calling load_default_package() will raise NotImplementedError - the example won't run as written

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}