feat: Add Stochastic Decomposition Layer and Fix Dependencies

[zer-art]

Pull Request

Description

This PR implements the Stochastic Decomposition Layer (SDL) as a new standalone module, addressing issue #185.

The SDL enables the injection of controllable noise into feature maps, facilitating uncertainty quantification by separating deterministic signals from stochastic perturbations. This implementation is based on Section 2.2 of the paper "Controllable Probabilistic Forecasting with Stochastic Decomposition Layers".

Important

Dependency Fix: This PR also includes a critical fix for the nnja_ai dependency. The main package __init__.py imports nnja_ai via SensorDataset, but the existing code referenced the wrong package name (nnja). This was causing the entire package to crash on import, blocking the execution of the new tests.

Key Changes:

• New Feature: Added StochasticDecompositionLayer in graph_weather/models/layers/stochastic_decomposition.py.
  • Implements the operation: $Output = x + (\alpha \cdot Style(z) \cdot \epsilon)$
  • Includes a learnable alpha (initialized to 0) and a style network to modulate noise based on a latent control vector $z$.
• Dependency Fix: Updated graph_weather/data/nnja_ai.py to fix an import error (from nnja import ... -> from nnja_ai import ...).
• Exports: Exposed the new layer in graph_weather/models/__init__.py.

Fixes #185

How Has This Been Tested?

I have added a comprehensive test suite in tests/test_stochastic_decomposition.py to verify the module's correctness and stability.

Test Coverage:

• Broadcast Support: Verified that the layer correctly handles and broadcasts over arbitrary spatial/temporal dimensions (tested on 3D, 4D, and 5D tensors).
• Initialization: Verified that the layer acts as an identity function at initialization (alpha=0), ensuring stability at the start of training.
• Reproducibility: Confirmed that fixed random seeds and fixed latent vectors produce identical outputs.
• Gradient Flow: Verified that gradients propagate correctly through the stochastic components (alpha, style_net) and inputs.
• Error Handling: Validated that dimension mismatch errors are raised correctly.

To reproduce, run:

pytest tests/test_stochastic_decomposition.py

• Yes

If your changes affect data processing, have you plotted any changes? i.e. have you done a quick sanity check?

• Yes

Checklist:

• My code follows OCF's coding style guidelines
• I have performed a self-review of my own code
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• I have checked my code and corrected any misspellings

[zer-art]

I noticed the pre-commit.ci check is failing.

I ran the pre-commit hooks locally specifically on my added files (stochastic_decomposition.py, test_stochastic_decomposition.py and etc ), and they all Passed.

The failures appear to be stemming from existing legacy files (e.g., formatting issues in aurora/ or cafa/). I avoided auto-formatting those unrelated files to keep this PR clean and focused solely on the SDL implementation.

[jacobbieker]

Looks really good! Nice and concise. Thanks for this. If you can move the test to the right subfolder, we can then merge this change.

[jacobbieker]

LGTM