feat: Add Stochastic Decomposition Layer and Fix Dependencies (#188)

* feat: Add Stochastic Decomposition Layer and Fix Dependencies

* refactor: Move test_stochastic_decomposition.py per review

Author: zer-art

graph_weather/data/nnja_ai.py

@@ -15,7 +15,7 @@
 from torch.utils.data import Dataset
 
 try:
-    from nnja import DataCatalog
+    from nnja_ai import DataCatalog
 except ImportError:
     raise ImportError("NNJA-AI library not installed. Install with: " "`pip install nnja-ai`")
 

graph_weather/models/__init__.py

@@ -12,3 +12,4 @@
 from .layers.decoder import Decoder
 from .layers.encoder import Encoder
 from .layers.processor import Processor
+from .layers.stochastic_decomposition import StochasticDecompositionLayer

graph_weather/models/layers/stochastic_decomposition.py

@@ -0,0 +1,68 @@
+"""Stochastic Decomposition Layers to inject controllable noise into feature maps
+
+In the original paper the Stochastic Decomposition Layer (SDL) is described as
+
+The SDL decomposes the intermediate feature map into a deterministic component (the input)
+and a stochastic component (the noise). This decomposition allows the model to separate
+the signal processing from the uncertainty quantification.
+The stochastic component is generated by modulating random Gaussian noise with a learned
+style vector derived from a latent control variable. This architecture enables the
+ensemble generation process to be explicitly controlled by the latent variable, rather
+than relying on implicit randomness.
+
+The SDL operation is defined as:
+Output = x + (alpha * Style(z) * epsilon)
+
+Where x is the deterministic input, z is the latent control vector, and alpha is a
+learnable channel-wise scaling factor that determines the magnitude of the stochastic
+perturbation.
+
+"""
+
+import torch
+import torch.nn as nn
+
+
+class StochasticDecompositionLayer(nn.Module):
+    """Stochastic Decomposition Layer for controllable probabilistic outputs."""
+
+    def __init__(self, input_dim: int, latent_dim: int):
+        """Initialize the Stochastic Decomposition Layer.
+
+        Args:
+            input_dim: Number of channels in the input feature map
+            latent_dim: Dimension of the latent control vector
+        """
+        super().__init__()
+        self.input_dim = input_dim
+        self.latent_dim = latent_dim
+
+        self.alpha = nn.Parameter(torch.zeros(1, input_dim, 1))
+
+        self.style_net = nn.Linear(latent_dim, input_dim)
+
+    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        """Apply stochastic decomposition to input features.
+
+        Args:
+            x: Input features [Batch, Channels, *Spatial]
+            z: Latent control vector [Batch, Latent_Dim]
+
+        Returns:
+            Output features with injected stochasticity
+        """
+        if x.size(1) != self.input_dim:
+            raise ValueError(f"Expected {self.input_dim} channels, got {x.size(1)}")
+        epsilon = torch.randn_like(x)
+
+        style = self.style_net(z)  # [B, C]
+
+        spatial_dims = x.dim() - 2
+        for _ in range(spatial_dims):
+            style = style.unsqueeze(-1)
+
+        alpha_broadcast = self.alpha
+        while alpha_broadcast.dim() < x.dim():
+            alpha_broadcast = alpha_broadcast.unsqueeze(-1)
+
+        return x + (alpha_broadcast * style * epsilon)

tests/models/layers/test_stochastic_decomposition.py

@@ -0,0 +1,91 @@
+"""Tests for the Stochastic Decomposition Layer."""
+
+import pytest
+import torch
+
+from graph_weather.models.layers.stochastic_decomposition import (
+    StochasticDecompositionLayer,
+)
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 32, 10),
+        (2, 32, 16, 16),
+        (2, 32, 8, 16, 16),
+    ],
+)
+def test_sdl_shapes(shape):
+    """Ensure SDL handles arbitrary spatial/temporal dimensions via broadcasting"""
+    batch, channels = shape[0], shape[1]
+    latent_dim = 16
+
+    x = torch.randn(*shape)
+    z = torch.randn(batch, latent_dim)
+
+    model = StochasticDecompositionLayer(input_dim=channels, latent_dim=latent_dim)
+    out = model(x, z)
+
+    assert out.shape == shape
+    assert not torch.isnan(out).any()
+
+
+def test_initialization_is_deterministic():
+    """Alpha initialized to 0 should imply Identity function initially"""
+    x = torch.randn(2, 64, 32, 32)
+    z = torch.randn(2, 16)
+
+    model = StochasticDecompositionLayer(input_dim=64, latent_dim=16)
+
+    assert torch.allclose(model.alpha, torch.zeros_like(model.alpha))
+
+    out = model(x, z)
+    assert torch.allclose(out, x, atol=1e-6)
+
+
+def test_reproducibility():
+    """Fixed seed + fixed latent = fixed output"""
+    x = torch.randn(2, 16, 10)
+    z = torch.randn(2, 8)
+
+    model = StochasticDecompositionLayer(16, 8)
+
+    with torch.no_grad():
+        model.alpha.fill_(0.5)
+
+    torch.manual_seed(42)
+    out1 = model(x, z)
+
+    torch.manual_seed(42)
+    out2 = model(x, z)
+
+    assert torch.equal(out1, out2)
+
+
+def test_gradient_flow():
+    """Test that gradients flow correctly through the layer."""
+    x = torch.randn(2, 16, 10, requires_grad=True)
+    z = torch.randn(2, 8, requires_grad=True)
+
+    model = StochasticDecompositionLayer(16, 8)
+    with torch.no_grad():
+        model.alpha.fill_(0.1)
+
+    out = model(x, z)
+    loss = out.sum()
+    loss.backward()
+
+    assert model.style_net.weight.grad is not None
+    assert model.alpha.grad is not None
+    assert x.grad is not None
+
+
+def test_channel_mismatch_error():
+    """Test that channel mismatch raises ValueError."""
+    x = torch.randn(2, 32, 10)
+    z = torch.randn(2, 8)
+    model = StochasticDecompositionLayer(input_dim=16, latent_dim=8)
+
+    with pytest.raises(ValueError):
+        model(x, z)