Aifs crps

[loliverhennigh]

Earth2Studio Pull Request

Description

Duplicate here,
#500

[greptile-apps]

Greptile Overview

Greptile Summary

This PR introduces AIFSENS (AIFS Ensemble), a new ensemble forecasting model variant to the Earth2Studio framework. The addition includes a complete model implementation in earth2studio/models/px/aifsens.py with sophisticated atmospheric and temporal forcing calculations, comprehensive test coverage, updated documentation across multiple files, and a new optional dependency group in pyproject.toml. The AIFSENS model extends the existing AIFS capabilities by providing ensemble forecasting functionality with CRPS evaluation support, requiring different dependency versions (anemoi-models 0.5.1 vs 0.3.1) and including conflict constraints to prevent incompatible installations. The implementation follows established patterns in the codebase for prognostic models while handling the additional complexity of ensemble member management and specialized variable indexing schemes.

Critical Issue: The file earth2studio/lexicon/base.py contains an unresolved Git merge conflict with conflict markers that will cause syntax errors and prevent the code from running.

Important Files Changed

Filename	Score	Overview
earth2studio/lexicon/base.py	1/5	Contains unresolved Git merge conflict with conflict markers preventing code execution
earth2studio/models/px/aifsens.py	3/5	New AIFSENS model implementation with complex logic and hardcoded indices
pyproject.toml	4/5	Adds aifsens dependency group with version conflicts requiring careful management
test/models/px/test_aifsens.py	4/5	Comprehensive test suite for new AIFSENS model functionality
earth2studio/models/px/init.py	5/5	Simple import addition for AIFSENS model exposure
docs/modules/models.rst	5/5	Documentation update adding AIFSENS to model list
docs/userguide/about/install.md	5/5	Installation documentation for new aifsens extra dependency
docs/userguide/support/troubleshooting.md	5/5	Updated troubleshooting docs for Flash Attention issues
test/models/test_auto_models.py	5/5	Adds AIFSENS to automated model download testing
README.md	5/5	Updates model table to include AIFS Ensemble entry

Confidence score: 1/5

• This PR cannot be merged safely due to the unresolved Git merge conflict that will cause immediate syntax errors
• Score reflects the critical merge conflict in base.py that makes the code non-functional, plus complexity concerns in the new model implementation with hardcoded indices and external dependencies
• The merge conflict in earth2studio/lexicon/base.py must be resolved before this PR can be considered for merging

Sequence Diagram

sequenceDiagram
    participant User
    participant AIFSENS
    participant Package
    participant TorchModel
    participant InterpolationMatrix
    participant DataSources

    User->>AIFSENS: "load_default_package()"
    AIFSENS->>Package: "Create Package with HuggingFace URL"
    Package-->>AIFSENS: "Package instance"
    AIFSENS-->>User: "Package"

    User->>AIFSENS: "load_model(package)"
    AIFSENS->>Package: "resolve('aifs-ens-crps-1.0.ckpt')"
    Package-->>AIFSENS: "model_path"
    AIFSENS->>TorchModel: "torch.load(model_path)"
    TorchModel-->>AIFSENS: "loaded_model"
    AIFSENS->>Package: "Load interpolation matrices"
    Package-->>AIFSENS: "interpolation_matrix, inverse_matrix"
    AIFSENS->>AIFSENS: "Create AIFSENS instance with matrices"
    AIFSENS-->>User: "AIFSENS model"

    User->>DataSources: "fetch_data(time, variables, lead_time)"
    DataSources-->>User: "input_tensor, coords"

    User->>AIFSENS: "__call__(input_tensor, coords)"
    AIFSENS->>AIFSENS: "_prepare_input(x, coords)"
    AIFSENS->>InterpolationMatrix: "Interpolate to model grid"
    InterpolationMatrix-->>AIFSENS: "interpolated_data"
    AIFSENS->>AIFSENS: "Add temporal features (cos/sin time, zenith angle)"
    AIFSENS->>TorchModel: "model.predict_step(x, fcstep)"
    TorchModel-->>AIFSENS: "prediction_output"
    AIFSENS->>AIFSENS: "_prepare_output(y, coords)"
    AIFSENS->>InterpolationMatrix: "Interpolate back to lat/lon grid"
    InterpolationMatrix-->>AIFSENS: "final_output"
    AIFSENS-->>User: "output_tensor, output_coords"

    User->>AIFSENS: "create_iterator(x, coords)"
    AIFSENS->>AIFSENS: "_default_generator(x, coords)"
    loop Time Integration Steps
        AIFSENS->>AIFSENS: "_update_input(x, coords)"
        AIFSENS->>TorchModel: "model.predict_step(x, fcstep)"
        TorchModel-->>AIFSENS: "prediction"
        AIFSENS->>AIFSENS: "_prepare_output(prediction, coords)"
        AIFSENS-->>User: "yield output_tensor, coords"
    end

Loading

_{10 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Greptile Overview

Greptile Summary

This PR introduces the AIFS Ensemble (AIFSENS) model to Earth2Studio, a new ensemble forecasting variant of ECMWF's AIFS weather prediction system. The changes add a complete model implementation with proper documentation, testing infrastructure, and dependency management. The implementation includes complex functionality for grid interpolation, time-dependent feature calculations, and ensemble forecasting capabilities.

The core addition is the new earth2studio/models/px/aifsens.py file that implements the AIFSENS class with AutoModelMixin integration, supporting model loading from external packages and providing iterator-based time integration. Supporting changes include documentation updates in README.md and installation guides, test coverage in test_aifsens.py, and dependency management through a new 'aifsens' optional group in pyproject.toml with version-pinned requirements.

Important Files Changed

Filename	Score	Overview
earth2studio/lexicon/base.py	1/5	Contains unresolved Git merge conflict with conflict markers preventing code execution
earth2studio/models/px/aifsens.py	4/5	New AIFS ensemble model implementation with complex grid interpolation and forecasting logic
pyproject.toml	4/5	Adds new 'aifsens' dependency group with pinned versions and conflict rules
test/models/px/test_aifsens.py	4/5	Comprehensive test suite for AIFSENS model with mock implementations and integration tests
earth2studio/models/px/init.py	5/5	Simple import addition for AIFSENS class following existing patterns
docs/userguide/about/install.md	5/5	Documentation update adding installation instructions for AIFS Ensemble
docs/userguide/support/troubleshooting.md	5/5	Minor documentation clarification for Flash Attention build time issues
docs/modules/models.rst	5/5	Adds AIFSENS to autosummary documentation list
test/models/test_auto_models.py	5/5	Includes AIFSENS in automated model testing framework
README.md	5/5	Adds AIFS Ensemble model to prognostic models table

Confidence score: 1/5

• This PR contains a critical merge conflict that will cause immediate runtime failures when deployed
• Score reflects the presence of unresolved Git conflict markers in earth2studio/lexicon/base.py that prevent Python from parsing the dictionary syntax
• Pay close attention to earth2studio/lexicon/base.py which requires immediate conflict resolution before this PR can be merged

Sequence Diagram

sequenceDiagram
    participant User
    participant Package as "Package System"
    participant AIFSENS as "AIFSENS Model"
    participant HuggingFace as "HuggingFace Hub"
    participant ECMWF as "ECMWF Registry"
    participant Model as "Anemoi Model"
    participant Interpolation as "Interpolation Matrix"
    participant Output as "Output Tensor"

    User->>Package: "load_default_package()"
    Package->>HuggingFace: "Download aifs-ens-crps-1.0.ckpt"
    HuggingFace-->>Package: "Model checkpoint"
    Package->>ECMWF: "Download interpolation matrices"
    ECMWF-->>Package: "Interpolation data"
    Package-->>User: "Package ready"

    User->>AIFSENS: "load_model(package)"
    AIFSENS->>Package: "resolve('aifs-ens-crps-1.0.ckpt')"
    Package-->>AIFSENS: "Model path"
    AIFSENS->>Model: "torch.load(model_path)"
    Model-->>AIFSENS: "Loaded model"
    AIFSENS->>Package: "Load interpolation matrices"
    Package-->>AIFSENS: "Matrix data"
    AIFSENS->>Interpolation: "Create sparse CSR tensors"
    Interpolation-->>AIFSENS: "Interpolation matrices"
    AIFSENS-->>User: "Initialized AIFSENS model"

    User->>AIFSENS: "__call__(x, coords)"
    AIFSENS->>AIFSENS: "_prepare_input(x, coords)"
    AIFSENS->>Interpolation: "Apply interpolation matrix"
    Interpolation-->>AIFSENS: "Interpolated input"
    AIFSENS->>AIFSENS: "Add time-based forcings"
    AIFSENS->>Model: "predict_step(x, fcstep)"
    Model-->>AIFSENS: "Model predictions"
    AIFSENS->>AIFSENS: "_prepare_output(y, coords)"
    AIFSENS->>Interpolation: "Apply inverse interpolation"
    Interpolation-->>AIFSENS: "Final grid output"
    AIFSENS->>Output: "Format output tensor"
    Output-->>User: "Forecast predictions"

Loading

_{10 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Greptile Overview

Greptile Summary

This pull request adds support for the AIFSENS (AIFS Ensemble) model to Earth2Studio, expanding the library's weather forecasting capabilities with ensemble-based probabilistic predictions. The changes introduce a new prognostic model class that extends the existing AIFS framework with ensemble sampling capabilities, utilizing updated ECMWF checkpoints optimized for CRPS (Continuous Ranked Probability Score) evaluation.

The implementation includes comprehensive infrastructure updates: new optional dependencies (aifsens) with version-specific anemoi packages, extensive documentation covering installation and troubleshooting, and complete test coverage. The AIFSENS model incorporates advanced features including astronomical calculations for solar forcing, grid interpolation matrices for coordinate transformations, and sophisticated tensor operations for ensemble data handling. The model operates on a different checkpoint (aifs-ens-crps-1.0.ckpt) and implements specialized variable filtering to exclude generated forcing indices (92-101 range) that aren't relevant for ensemble forecasting workflows.

Important Files Changed

Filename	Score	Overview
earth2studio/lexicon/base.py	0/5	Contains unresolved Git merge conflicts with duplicate variable definitions that prevent compilation
earth2studio/models/px/aifsens.py	3/5	New AIFSENS ensemble model implementation with complex interpolation, astronomical calculations, and security concerns
test/models/px/test_aifsens.py	4/5	Comprehensive test suite for AIFSENS with mock classes and integration tests
pyproject.toml	4/5	Adds aifsens optional dependencies with proper conflict management between aifs variants
docs/userguide/about/install.md	5/5	Documentation update adding AIFSENS installation instructions following existing patterns
docs/userguide/support/troubleshooting.md	5/5	Minor clarification that Flash Attention issues affect both AIFS model variants
earth2studio/models/px/init.py	5/5	Straightforward import addition for AIFSENS model following existing patterns
docs/modules/models.rst	5/5	Documentation update adding AIFSENS to prognostic models autosummary list
README.md	5/5	Simple documentation addition to prognostic models comparison table
test/models/test_auto_models.py	5/5	Adds AIFSENS to automated model download test suite

Confidence score: 1/5

• This PR has critical issues that will prevent compilation and deployment due to unresolved merge conflicts in a core lexicon file
• Score reflects the unresolved Git merge conflict in earth2studio/lexicon/base.py that introduces duplicate dictionary keys and will cause immediate runtime failures, plus security concerns with torch.load(weights_only=False) in the AIFSENS implementation
• The earth2studio/lexicon/base.py file requires immediate attention to resolve merge conflicts before this PR can be merged safely

Sequence Diagram

sequenceDiagram
    participant User
    participant AIFSENS
    participant Package
    participant Model
    participant Random
    participant Interpolation
    participant Output

    User->>AIFSENS: "load_default_package()"
    AIFSENS->>Package: "create package from hf://ecmwf/aifs-ens-1.0"
    Package-->>AIFSENS: "package instance"
    
    User->>AIFSENS: "load_model(package)"
    AIFSENS->>Package: "resolve('aifs-ens-crps-1.0.ckpt')"
    Package-->>AIFSENS: "model path"
    AIFSENS->>Package: "load interpolation matrices"
    Package-->>AIFSENS: "interpolation data"
    AIFSENS-->>User: "AIFSENS model instance"
    
    User->>Random: "create random data source"
    Random-->>User: "data source instance"
    
    User->>Random: "fetch_data(time, variables, lead_time)"
    Random-->>User: "input tensor and coordinates"
    
    User->>AIFSENS: "__call__(input_tensor, coordinates)"
    AIFSENS->>AIFSENS: "_prepare_input(x, coords)"
    AIFSENS->>Interpolation: "interpolation_matrix @ x"
    Interpolation-->>AIFSENS: "interpolated input"
    AIFSENS->>AIFSENS: "add forcing terms (cos/sin julian day, local time, zenith angle)"
    AIFSENS->>Model: "predict_step(x, fcstep)"
    Model-->>AIFSENS: "model prediction"
    AIFSENS->>AIFSENS: "_prepare_output(y, coords)"
    AIFSENS->>Interpolation: "inverse_interpolation_matrix @ y"
    Interpolation-->>AIFSENS: "output on lat/lon grid"
    AIFSENS-->>User: "forecast output and coordinates"

Loading

_{10 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}