fresh new main branch

.gitattributes

@@ -0,0 +1,6 @@
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.safetensor filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+veg/MultiTaskConvLSTM_veg_variables filter=lfs diff=lfs merge=lfs -text
+no_veg/MultiTaskConvLSTM_no_veg_variables filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,56 @@
+# MultiTask ConvLSTM for Precipitation Prediction
+
+This repository contains two MultiTask ConvLSTM models:
+- **veg/**: Model trained with vegetation input variables
+- **noveg/**: Model trained without vegetation input variables
+
+Both directories include:
+- `convlstm.py`: base ConvLSTM layers
+- `model.py`: MultiTask ConvLSTM model definition
+- `example_inference.py`: inference script
+- `data/`: example `.pth` files (test)
+
+These scripts are provided for reproducibility of the model architecture and workflow. 
+Exact runtime and performance may vary depending on hardware.
+
+## Example Data
+
+We provide a large test `.pth` files
+so you can immediately run the inference script without preprocessing.
+These files are already preprocessed and normalized from the ECWMF REA5 reanalysis data.
+
+Each `.pth` file loads as a list of batches:
+
+- `X_batch`: shape `(B, T_in, C_in, H*W)`  
+- `y_batch`: shape `(B, T_out, C_out, H*W)`  
+- `y_zero_batch`: shape `(B, T_out, C_out, H*W)`  
+
+with `H=81`, `W=97`. Inside `evaluate(...)`, these are reshaped to `(B, T, C, H, W)`.
+
+---
+
+## How to Use
+
+Ensure all files are in the correct directory then run the example_inference.py file.
+
+# 1 Get the repo
+git clone https://huggingface.co/<your-username>/MultiTaskConvLSTM
+cd MultiTaskConvLSTM
+
+# 2 Install minimal deps
+pip install -r requirements.txt
+
+# 3 Run inference (choose one variant)
+python veg/example_inference.py
+# or
+python noveg/example_inference.py
+
+
+## Citation If you use this model, please cite: > Lilly Horvath-Makkos (2025). [title] [journal] BibTeX:
+bibtex
+@article{horvathmakkos2025,
+  title={Title},
+  author={Horvath-Makkos, Lilly},
+  journal={Journal},
+  year={2025}
+}

no_veg/ConvLSTM.py

@@ -0,0 +1,184 @@
+#ConvLSTM definition
+
+
+import torch.nn as nn
+import torch
+
+
+class ConvLSTMCell(nn.Module):
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, bias):
+        #Input_um is the number of channels per input tensor, hidden_dim is the numer of channels of hidden state, bias is a booleam, wehther or not to add a bias
+
+        super(ConvLSTMCell, self).__init__()
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+
+        self.kernel_size = kernel_size
+        self.padding = (kernel_size[0])// 2, (kernel_size[1]) // 2
+        self.bias = bias
+
+        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,
+                              out_channels=4 * self.hidden_dim,
+                              kernel_size=self.kernel_size,
+                              padding=self.padding,
+                              bias=self.bias)
+
+    def forward(self, input_tensor, cur_state):
+        h_cur, c_cur = cur_state
+
+        combined = torch.cat([input_tensor, h_cur], dim=1)  # concatenate along channel axis
+
+        combined_conv = self.conv(combined)
+        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
+        i = torch.sigmoid(cc_i)
+        f = torch.sigmoid(cc_f)
+        o = torch.sigmoid(cc_o)
+        g = torch.tanh(cc_g)
+
+        c_next = f * c_cur + i * g
+        h_next = o * torch.tanh(c_next)
+
+        return h_next, c_next
+
+    def init_hidden(self, batch_size, image_size):
+        height, width = image_size
+        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
+                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))
+
+
+class ConvLSTM(nn.Module):
+
+    """
+
+    Parameters:
+        input_dim: Number of channels in input
+        hidden_dim: Number of hidden channels
+        kernel_size: Size of kernel in convolutions
+        num_layers: Number of LSTM layers stacked on each other
+        batch_first: Whether or not dimension 0 is the batch or not
+        bias: Bias or no bias in Convolution
+        return_all_layers: Return the list of computations for all layers
+        Note: Will do same padding.
+
+    Input:
+        A tensor of size B, T, C, H, W or T, B, C, H, W
+    Output:
+        A tuple of two lists of length num_layers (or length 1 if return_all_layers is False).
+            0 - layer_output_list is the list of lists of length T of each output
+            1 - last_state_list is the list of last states
+                    each element of the list is a tuple (h, c) for hidden state and memory
+    Example:
+        >> x = torch.rand((32, 10, 64, 128, 128))
+        >> convlstm = ConvLSTM(64, 16, 3, 1, True, True, False)
+        >> _, last_states = convlstm(x)
+        >> h = last_states[0][0]  # 0 for layer index, 0 for h index
+    """
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers,
+                 batch_first=False, bias=True, return_all_layers=False):
+        super(ConvLSTM, self).__init__()
+
+        self._check_kernel_size_consistency(kernel_size)
+
+        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers
+        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)
+        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)
+        if not len(kernel_size) == len(hidden_dim) == num_layers:
+            raise ValueError('Inconsistent list length.')
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+        self.batch_first = batch_first
+        self.bias = bias
+        self.return_all_layers = return_all_layers
+
+        cell_list = []
+        for i in range(0, self.num_layers):
+            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]
+            # print(f"Layer {i}: input_dim={cur_input_dim}, hidden_dim={self.hidden_dim[i]}")
+            cell_list.append(ConvLSTMCell(input_dim=cur_input_dim,
+                                          hidden_dim=self.hidden_dim[i],
+                                          kernel_size=self.kernel_size[i],
+                                          bias=self.bias))
+
+        self.cell_list = nn.ModuleList(cell_list)
+
+    def forward(self, input_tensor, hidden_state=None):
+        """
+
+        Parameters
+        ----------
+        input_tensor: todo
+            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)
+        hidden_state: todo
+            None. todo implement stateful
+
+        Returns
+        -------
+        last_state_list, layer_output
+        """
+        if not self.batch_first:
+            # (t, b, c, h, w) -> (b, t, c, h, w)
+            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)
+
+        b, _, _, h, w = input_tensor.size()
+
+        # Implement stateful ConvLSTM
+        if hidden_state is not None:
+            raise NotImplementedError()
+        else:
+            # Since the init is done in forward. Can send image size here
+            hidden_state = self._init_hidden(batch_size=b,
+                                             image_size=(h, w))
+
+        layer_output_list = []
+        last_state_list = []
+
+        seq_len = input_tensor.size(1)
+        cur_layer_input = input_tensor
+
+        for layer_idx in range(self.num_layers):
+
+            h, c = hidden_state[layer_idx]
+            output_inner = []
+            for t in range(seq_len):
+                # print(f"Layer {layer_idx}, Time {t}, Input shape: {cur_layer_input[:, t, :, :, :].shape}")
+                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :],
+                                                 cur_state=[h, c])
+                output_inner.append(h)
+
+            layer_output = torch.stack(output_inner, dim=1)
+            cur_layer_input = layer_output
+            
+            # print(f"ConvLSTM Layer {layer_idx} output shape: {cur_layer_input.shape}")
+
+            layer_output_list.append(layer_output)
+            last_state_list.append([h, c])
+
+        if not self.return_all_layers:
+            layer_output_list = layer_output_list[-1:]
+            last_state_list = last_state_list[-1:]
+
+        return layer_output_list, last_state_list
+
+    def _init_hidden(self, batch_size, image_size):
+        init_states = []
+        for i in range(self.num_layers):
+            init_states.append(self.cell_list[i].init_hidden(batch_size, image_size))
+        return init_states
+
+    @staticmethod
+    def _check_kernel_size_consistency(kernel_size):
+        if not (isinstance(kernel_size, tuple) or
+                (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):
+            raise ValueError('`kernel_size` must be tuple or list of tuples')
+
+    @staticmethod
+    def _extend_for_multilayer(param, num_layers):
+        if not isinstance(param, list):
+            param = [param] * num_layers
+        return param

no_veg/MultiTaskConvLSTM.py

@@ -0,0 +1,129 @@
+from ConvLSTM import ConvLSTM
+import torch
+import torch.nn as nn
+from collections import defaultdict
+
+#MLP definition
+class MLP_5D(nn.Module):
+    def __init__(self, height, width):
+        super(MLP_5D, self).__init__()
+        # Define the fully connected layers
+        self.fc1 = nn.Linear(64, 128)  # Input channels = 41, output features = 128
+        self.dropout1 = nn.Dropout(0.05)
+        self.fc2 = nn.Linear(128, 64)  # Output features = 64
+        self.dropout2 = nn.Dropout(0.05)
+        self.fc3 = nn.Linear(64, 1)    # Final output, reducing to 1 channel
+
+        self.height = height
+        self.width = width
+
+    def forward(self, x):
+        batch_size, timesteps, channels, height, width = x.shape
+        
+        # Ensure the input spatial dimensions match the expected height and width
+        assert height == self.height and width == self.width, "Height and width mismatch"
+        
+        # Reshape to (batch * timesteps * height * width, channels)
+        x = x.permute(0, 1, 3, 4, 2).reshape(-1, channels)
+        # print(x.shape)
+        
+        # Apply MLP (Fully connected layers)
+        x = self.fc1(x)
+        x = torch.nn.functional.softplus(x)
+        x = self.dropout1(x)
+        x = self.fc2(x)
+        x = torch.nn.functional.softplus(x)
+        x = self.dropout2(x)
+        x = self.fc3(x)
+        x = torch.nn.functional.softplus(x)
+        
+        # Reshape back to (batch, timesteps, 1, height, width)
+        x = x.view(batch_size, timesteps, self.height, self.width, 1).permute(0, 1, 4, 2, 3)
+
+        return x
+    
+  
+# MultiTask ConvLSTM definition
+
+class ConvLSTMNetwork(nn.Module):
+    def __init__(self, input_dim, hidden_dims, kernel_size, num_layers, output_channels, batch_first=True, pool_size=(2,2)):
+        super(ConvLSTMNetwork, self).__init__()
+        
+        # ConvLSTM module
+        self.convlstm = ConvLSTM(input_dim=input_dim,
+                                 hidden_dim=hidden_dims,
+                                 kernel_size=kernel_size,
+                                 num_layers=num_layers,
+                                 batch_first=batch_first,
+                                 bias=True,
+                                 return_all_layers=True)
+        
+        # Batch Normalization for each ConvLSTM layer's output
+        self.batch_norms = nn.ModuleList([
+            nn.BatchNorm3d(hidden_dim) for hidden_dim in hidden_dims
+        ])
+
+        # Final Conv3D layer for regression pathway
+        self.conv3d = nn.Conv3d(in_channels=hidden_dims[-1],
+                                out_channels=output_channels,
+                                kernel_size=(1, 3, 3),
+                                padding=(0, 1, 1))
+
+        # MLP for regression output: (B,T,C,H,W) -> (B,T,1,H,W)
+        self.mlp = MLP_5D(height=81, width=97)
+
+        # Classification head for pixel-level zero precipitation probability
+        # We'll produce (B,T,1,H,W) as well:
+        # The classification head takes (B,C,T,H,W) input. We'll reorder dimensions before applying it.
+        # Then apply Sigmoid to get probabilities between 0 and 1.
+        self.classification_head = nn.Sequential(
+            nn.Conv3d(output_channels, 1, kernel_size=(1,1,1)),  # from C to 1 channel
+            nn.Sigmoid()
+        )
+
+        self.activation_variance = defaultdict(list)
+
+    def forward(self, x):
+        """
+        x: (B, T, input_dim, H, W)
+        """
+        # Forward through ConvLSTM
+        layer_output_list, last_state_list = self.convlstm(x)
+        
+        # Apply batch norms
+        for i, output in enumerate(layer_output_list):
+            # output: (B, T, C, H, W)
+            output = output.permute(0, 2, 1, 3, 4)  # (B, C, T, H, W) for BatchNorm3d
+            output = self.batch_norms[i](output)
+            output = output.permute(0, 2, 1, 3, 4)  # back to (B, T, C, H, W)
+
+            #Track variance across spatial dimensions for hooks with activation tracking 
+            activation_variance = output.var(dim=(3, 4)).mean().item()
+            self.activation_variance[f"ConvLSTM_layer_{i}"].append(activation_variance)
+
+            layer_output_list[i] = output
+        
+        # Take output from the last ConvLSTM layer
+        final_output = layer_output_list[-1]  # (B, T, C, H, W)
+
+        # Pass through Conv3D: needs (B,C,T,H,W)
+        final_output = final_output.permute(0, 2, 1, 3, 4)  # (B,C,T,H,W)
+        final_output = self.conv3d(final_output)
+        # Now final_output: (B, output_channels, T, H, W)
+
+        # Return to (B,T,C,H,W) for MLP (regression)
+        final_output_t = final_output.permute(0, 2, 1, 3, 4)  # (B,T,C,H,W)
+
+        # Regression output
+        regression_output = self.mlp(final_output_t)  # (B,T,1,H,W)
+
+        # Classification output:
+        # The classification head is defined for (B,C,T,H,W), so reorder again
+        final_output_c = final_output  # still (B,output_channels,T,H,W)
+        classification_output = self.classification_head(final_output_c)
+        # classification_output: (B,1,T,H,W)
+
+        # Permute classification output to match (B,T,1,H,W) format
+        classification_output = classification_output.permute(0, 2, 1, 3, 4)  # (B,T,1,H,W)
+
+        return regression_output, classification_output

no_veg/MultiTaskConvLSTM_no_veg_variables.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82e89dda4a11281b4a9dbecf081951c304977d3f481f7f6024a2edd3261b02e9
+size 1317646

no_veg/data/normalized_test_data_no_veg_input.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5b76c968c3a80b260f8db30d5e9c219c241ac26fd44856c3c87008394dac8e9
+size 1644632048

no_veg/example_inference.py

@@ -0,0 +1,231 @@
+# example_inference
+import torch
+from MultiTaskConvLSTM import ConvLSTMNetwork
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
+import torch
+import toch.nn as nn
+from tqdm.auto import tqdm
+from utils import (
+    mse, mae, nash_sutcliffe_efficiency, r2_score, pearson_correlation,
+    spearman_correlation, percentage_error, percentage_bias,
+    kendall_tau, spatial_correlation
+)
+import torch.optim as optim
+
+
+device = 'cpu'
+
+height = 81
+width = 97
+
+set_lookback = 1
+set_forecast_horizon = 1
+
+#Define variables for evaluation
+batch_size = 16
+time_steps_out = set_forecast_horizon
+channels = 9
+
+#Variable names
+#Variable names
+variable_names = ['10 metre U wind component', '10 metre V wind component', '2 metre dewpoint temperature', '2 metre temperature', 'Total column rain water', 'Total precipitation', 'Time-integrated surface latent heat net flux']
+
+# Adjust input_dim and output_channels according to your data specifics
+model = ConvLSTMNetwork(
+    input_dim=9 * set_lookback, 
+    hidden_dims=[9, 32, 64], 
+    kernel_size=(3,3), 
+    num_layers=3, 
+    output_channels=64 * set_forecast_horizon, 
+    batch_first=True
+).to(device)
+
+# Define separate loss functions
+loss_fn = nn.MSELoss()       # For regression output
+bce_loss_fn = nn.BCELoss()       # For classification output
+
+optimizer = optim.AdamW(model.parameters(), lr = 0.005) 
+
+checkpoint = torch.load("MultiTaskConvLSTM_no_veg_variables")
+model.load_state_dict(checkpoint['model_state_dict'])
+
+# If you want to move the model to the GPU (optional, depending on your setup)
+model.to(device)  # Assuming you have a variable `device` for CUDA or CPU
+
+# Ensure that the model is in evaluation mode if you're using it for inference
+model.eval()
+
+print("Model loaded successfully")
+
+
+threshold = 0.1
+precip_index = 10
+
+def evaluate(model, test_loader, reg_loss_fn, class_loss_fn, device, variable_names, height, width):
+    """
+    Evaluate the model on the test set for both regression and classification tasks.
+    """
+    model.eval()  # Set the model to evaluation model
+
+    # input_to_true = {'zero_to_non_zero': 0, 'non_zero_to_zero': 0}
+    # input_to_pred_REG = {'zero_to_non_zero': 0, 'non_zero_to_zero': 0}
+    # input_to_pred_CLASS = {'zero_to_non_zero': 0, 'non_zero_to_zero': 0}
+
+    test_reg_loss = 0.0
+    test_class_loss = 0.0
+    test_total_loss = 0.0
+
+    y_true_reg = []  # List to store true values for regression
+    y_pred_reg = []  # List to store predicted values for regression
+
+    y_pred_reg2 = []
+
+    y_true_class = []  # List to store true values for classification
+    y_pred_class = []  # List to store predicted probabilities for classification
+
+    # Disable gradient computation
+    with torch.no_grad():
+        for X_test, y_test, y_zero_test in tqdm(test_loader, desc="Evaluating on Test Set"):
+            # Move the batch to the device
+            X_test, y_test, y_zero_test = X_test.to(device), y_test.to(device), y_zero_test.to(device)
+
+            # Reshape inputs and targets
+            batch_size, time_steps_in, channels_in, grid_points = X_test.shape
+            batch_size, time_steps_out, channels_out, grid_points = y_test.shape
+            X_test = X_test.view(batch_size, time_steps_in, channels_in, height, width)
+            y_test = y_test.view(batch_size, time_steps_out, channels_out, height, width)
+            y_zero_test = y_zero_test.view(batch_size, time_steps_out, channels_out, height, width)
+
+            # Forward pass
+            regression_output, classification_output = model(X_test)
+
+            classification_predictions = (classification_output > 0.7).float()
+
+            # Compute regression loss
+            reg_loss = reg_loss_fn(regression_output, y_test)
+
+            # Compute classification loss
+            class_loss = class_loss_fn(classification_output, y_zero_test)
+
+            # Total loss
+            total_loss = reg_loss + class_loss
+
+            regression_output2 = torch.where(classification_predictions == 0, regression_output, classification_predictions)
+
+            # Accumulate losses
+            test_reg_loss += reg_loss.item() * X_test.size(0)
+            test_class_loss += class_loss.item() * X_test.size(0)
+            test_total_loss += total_loss.item() * X_test.size(0)
+
+            # Collect true and predicted values for regression and classification
+            y_true_reg.append(y_test.cpu())
+            y_pred_reg.append(regression_output.cpu())
+            y_pred_reg2.append(regression_output2.cpu())
+            y_true_class.append(y_zero_test.cpu())
+            y_pred_class.append(classification_output.cpu())
+
+    # Normalize losses by the total dataset size
+    test_reg_loss /= len(test_loader)
+    test_class_loss /= len(test_loader)
+    test_total_loss /= len(test_loader)
+
+    print(f"Test Regression Loss: {test_reg_loss:.16f}")
+    print(f"Test Classification Loss: {test_class_loss:.16f}")
+    print(f"Test Total Loss: {test_total_loss:.16f}")
+
+    y_true_reg_flat = torch.cat(y_true_reg, dim=0).flatten()  # Keep as PyTorch tensor
+    y_pred_reg_flat = torch.cat(y_pred_reg, dim=0).flatten()  # Keep as PyTorch tensor
+    y_true_class_flat = torch.cat(y_true_class, dim=0).flatten()  # Keep as PyTorch tensor
+    y_pred_class_flat = torch.cat(y_pred_class, dim=0).flatten()  # Keep as PyTorch tensor
+
+    # Compute regression metrics
+    regression_metrics = {
+        "MSE": mse(y_true_reg_flat, y_pred_reg_flat),
+        "MAE": mae(y_true_reg_flat, y_pred_reg_flat),
+        "NSE": nash_sutcliffe_efficiency(y_true_reg_flat, y_pred_reg_flat),
+        "R2": r2_score(y_true_reg_flat, y_pred_reg_flat),
+        "Pearson": pearson_correlation(y_true_reg_flat, y_pred_reg_flat),
+        "Spearman": spearman_correlation(y_true_reg_flat, y_pred_reg_flat),
+        "NSE": nash_sutcliffe_efficiency(y_true_reg_flat, y_pred_reg_flat),
+        "Percentage Error": percentage_error(y_true_reg_flat, y_pred_reg_flat),
+        "Percentage Bias": percentage_bias(y_true_reg_flat, y_pred_reg_flat),
+        "Kendall Tau": kendall_tau(y_true_reg_flat, y_pred_reg_flat),
+        "Spatial Correlation": spatial_correlation(y_true_reg_flat, y_pred_reg_flat)}
+
+    print("\nRegression Metrics:")
+    for metric, value in regression_metrics.items():
+        print(f"{metric}: {value:.16f}")
+
+
+    # Compute classification metrics
+    classification_metrics = {
+        "Accuracy": accuracy_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "Precision": precision_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "Recall": recall_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "F1": f1_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "ROC-AUC": roc_auc_score(y_true_class_flat, y_pred_class_flat),
+    }
+
+    print("\nClassification Metrics:")
+    for metric, value in classification_metrics.items():
+        print(f"{metric}: {value:.16f}")
+
+    torch.save({
+        'y_true_reg': y_true_reg_flat,
+        'y_pred_reg': y_pred_reg_flat,
+        'y_true_class': y_true_class_flat,
+        'y_pred_class': y_pred_class_flat,
+    }, 'results')
+
+    return test_total_loss, regression_metrics, classification_metrics
+
+
+"""
+EXPECTED DATALOADER BATCH FORMAT (normalized_test_data):
+
+Each batch must be a tuple: (X_batch, y_batch, y_zero_batch)
+
+X_batch contains the previous hours variables. y_batch contains the next hour's precipitation.
+y_zero_batch contains the next hour's precipitation thresholded as 0 for precipiation <=0.1mm/h and 
+1 for precipitation >0.1mm.
+
+Shapes BEFORE reshaping inside `evaluate`:
+    X_batch:   (B, T_in, C_in, G)        # G = H*W = 81*97 = 7857
+    y_batch:   (B, T_out, C_out, G)
+    y_zero_batch: (B, T_out, C_out, G)   # binary 0/1 "zero-precip" targets
+
+    If your preprocessing produces (B,T, C, H, W), reshape to (B, T, C, H*W) before inference.
+
+DTypes:
+    X_batch, y_batch:       torch.float32
+    y_zero_batch:           torch.float32 (will be used with BCELoss)
+
+Reshaping done in 'evaluate':
+    X_test = X_batch.view(B, T_in, C_in, H, W)         -> (B, T_in, C_in, 81, 97)
+    y_test = y_batch.view(B, T_out, C_out, H, W)       -> (B, T_out, C_out, 81, 97)
+    y_zero_test = y_zero_batch.view(B, T_out, C_out, H, W)
+
+Model input:
+    model expects X_test shaped (B, T_in, input_dim, H, W)
+    where input_dim == 9 * set_lookback  (with set_lookback=1 -> input_dim=9)
+
+Notes:
+  • Make sure G == H*W (i.e., 7857 for 81x97).
+  • C_out for precipitation should be 1 (one target channel), and y_zero_batch
+    is the 0/1 mask for “zero precipitation” at each pixel & time.
+  • y_zero_batch should be probabilities/labels in {0,1} for BCELoss.
+"""
+
+normalized_test_data = torch.load("data/normalized_test_data_no_veg_input.pth")
+
+
+test_total_loss, regression_metrics, classification_metrics = evaluate(
+    model=model,
+    test_loader=normalized_test_data,
+    reg_loss_fn=loss_fn,
+    class_loss_fn=bce_loss_fn,
+    device=device,
+    variable_names=variable_names,
+    height=height,
+    width=width,
+)

no_veg/utils.py

@@ -0,0 +1,92 @@
+#Definition of evaluation metrics
+from scipy.stats import pearsonr, spearmanr
+import torch.nn.functional as F
+import torch
+from scipy.stats import kendalltau
+import scipy.stats as stats
+
+
+def nash_sutcliffe_efficiency(observed, predicted):
+    # Ensure inputs are tensors on the CPU
+    observed = observed.cpu()
+    predicted = predicted.cpu()
+
+    # Compute the numerator and denominator
+    numerator = torch.sum((observed - predicted) ** 2)
+    denominator = torch.sum((observed - torch.mean(observed)) ** 2)
+
+    # Calculate NSE
+    nse = 1 - (numerator / denominator)
+    return nse.item()
+
+def pearson_correlation(y_true, y_pred):
+    y_true = y_true.view(-1).cpu().numpy()  # Flatten and move to CPU
+    y_pred = y_pred.view(-1).cpu().numpy()  # Flatten and move to CPU
+    
+    return pearsonr(y_true, y_pred)[0]  # Return the correlation coefficient
+
+def spearman_correlation(y_true, y_pred):
+    y_true = y_true.view(-1).cpu().numpy()  # Flatten and move to CPU
+    y_pred = y_pred.view(-1).cpu().numpy()  # Flatten and move to CPU
+    
+    return spearmanr(y_true, y_pred).correlation  # Return the Spearman correlation
+
+def mse(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+    
+    return torch.mean((y_true - y_pred) ** 2).item()
+
+def mae(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+    
+    return torch.mean(torch.abs(y_true - y_pred)).item()
+
+def percentage_error(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+    
+    return 100 * torch.mean((y_pred - y_true) / (y_true + 1e-6)).item()
+
+def percentage_bias(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+
+    return 100 * torch.sum(y_pred - y_true) / (torch.sum(y_true) + 1e-6)
+
+def kendall_tau(y_true, y_pred):
+    y_true = y_true.view(-1).cpu().numpy()  # Flatten and move to CPU
+    y_pred = y_pred.view(-1).cpu().numpy()  # Flatten and move to CPU
+    
+    return kendalltau(y_true, y_pred).correlation  # Return the Kendall Tau
+
+def r2_score(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+
+    ss_total = torch.sum((y_true - torch.mean(y_true)) ** 2)
+    ss_residual = torch.sum((y_true - y_pred) ** 2)
+    
+    return 1 - (ss_residual / (ss_total + 1e-6)).item()
+
+def spatial_correlation(y_true, y_pred):
+    # Flatten the tensors to work with them
+    y_true_flat = y_true.view(-1).cpu()
+    y_pred_flat = y_pred.view(-1).cpu()
+
+    # Compute the numerator: sum(P * T)
+    numerator = torch.sum(y_pred_flat * y_true_flat)
+
+    # Compute the denominator: sqrt(sum(P^2) * sum(T^2))
+    denominator = torch.sqrt(torch.sum(y_pred_flat ** 2) * torch.sum(y_true_flat ** 2))
+
+    # Compute the correlation (add epsilon to avoid division by zero)
+    correlation = numerator / (denominator)
+
+    return correlation.item()

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.0
+numpy>=1.24
+scikit-learn>=1.3
+tqdm>=4.65

veg/ConvLSTM.py

@@ -0,0 +1,184 @@
+#ConvLSTM definition
+
+
+import torch.nn as nn
+import torch
+
+
+class ConvLSTMCell(nn.Module):
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, bias):
+        #Input_um is the number of channels per input tensor, hidden_dim is the numer of channels of hidden state, bias is a booleam, wehther or not to add a bias
+
+        super(ConvLSTMCell, self).__init__()
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+
+        self.kernel_size = kernel_size
+        self.padding = (kernel_size[0])// 2, (kernel_size[1]) // 2
+        self.bias = bias
+
+        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,
+                              out_channels=4 * self.hidden_dim,
+                              kernel_size=self.kernel_size,
+                              padding=self.padding,
+                              bias=self.bias)
+
+    def forward(self, input_tensor, cur_state):
+        h_cur, c_cur = cur_state
+
+        combined = torch.cat([input_tensor, h_cur], dim=1)  # concatenate along channel axis
+
+        combined_conv = self.conv(combined)
+        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
+        i = torch.sigmoid(cc_i)
+        f = torch.sigmoid(cc_f)
+        o = torch.sigmoid(cc_o)
+        g = torch.tanh(cc_g)
+
+        c_next = f * c_cur + i * g
+        h_next = o * torch.tanh(c_next)
+
+        return h_next, c_next
+
+    def init_hidden(self, batch_size, image_size):
+        height, width = image_size
+        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
+                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))
+
+
+class ConvLSTM(nn.Module):
+
+    """
+
+    Parameters:
+        input_dim: Number of channels in input
+        hidden_dim: Number of hidden channels
+        kernel_size: Size of kernel in convolutions
+        num_layers: Number of LSTM layers stacked on each other
+        batch_first: Whether or not dimension 0 is the batch or not
+        bias: Bias or no bias in Convolution
+        return_all_layers: Return the list of computations for all layers
+        Note: Will do same padding.
+
+    Input:
+        A tensor of size B, T, C, H, W or T, B, C, H, W
+    Output:
+        A tuple of two lists of length num_layers (or length 1 if return_all_layers is False).
+            0 - layer_output_list is the list of lists of length T of each output
+            1 - last_state_list is the list of last states
+                    each element of the list is a tuple (h, c) for hidden state and memory
+    Example:
+        >> x = torch.rand((32, 10, 64, 128, 128))
+        >> convlstm = ConvLSTM(64, 16, 3, 1, True, True, False)
+        >> _, last_states = convlstm(x)
+        >> h = last_states[0][0]  # 0 for layer index, 0 for h index
+    """
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers,
+                 batch_first=False, bias=True, return_all_layers=False):
+        super(ConvLSTM, self).__init__()
+
+        self._check_kernel_size_consistency(kernel_size)
+
+        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers
+        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)
+        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)
+        if not len(kernel_size) == len(hidden_dim) == num_layers:
+            raise ValueError('Inconsistent list length.')
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+        self.batch_first = batch_first
+        self.bias = bias
+        self.return_all_layers = return_all_layers
+
+        cell_list = []
+        for i in range(0, self.num_layers):
+            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]
+            # print(f"Layer {i}: input_dim={cur_input_dim}, hidden_dim={self.hidden_dim[i]}")
+            cell_list.append(ConvLSTMCell(input_dim=cur_input_dim,
+                                          hidden_dim=self.hidden_dim[i],
+                                          kernel_size=self.kernel_size[i],
+                                          bias=self.bias))
+
+        self.cell_list = nn.ModuleList(cell_list)
+
+    def forward(self, input_tensor, hidden_state=None):
+        """
+
+        Parameters
+        ----------
+        input_tensor: todo
+            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)
+        hidden_state: todo
+            None. todo implement stateful
+
+        Returns
+        -------
+        last_state_list, layer_output
+        """
+        if not self.batch_first:
+            # (t, b, c, h, w) -> (b, t, c, h, w)
+            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)
+
+        b, _, _, h, w = input_tensor.size()
+
+        # Implement stateful ConvLSTM
+        if hidden_state is not None:
+            raise NotImplementedError()
+        else:
+            # Since the init is done in forward. Can send image size here
+            hidden_state = self._init_hidden(batch_size=b,
+                                             image_size=(h, w))
+
+        layer_output_list = []
+        last_state_list = []
+
+        seq_len = input_tensor.size(1)
+        cur_layer_input = input_tensor
+
+        for layer_idx in range(self.num_layers):
+
+            h, c = hidden_state[layer_idx]
+            output_inner = []
+            for t in range(seq_len):
+                # print(f"Layer {layer_idx}, Time {t}, Input shape: {cur_layer_input[:, t, :, :, :].shape}")
+                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :],
+                                                 cur_state=[h, c])
+                output_inner.append(h)
+
+            layer_output = torch.stack(output_inner, dim=1)
+            cur_layer_input = layer_output
+            
+            # print(f"ConvLSTM Layer {layer_idx} output shape: {cur_layer_input.shape}")
+
+            layer_output_list.append(layer_output)
+            last_state_list.append([h, c])
+
+        if not self.return_all_layers:
+            layer_output_list = layer_output_list[-1:]
+            last_state_list = last_state_list[-1:]
+
+        return layer_output_list, last_state_list
+
+    def _init_hidden(self, batch_size, image_size):
+        init_states = []
+        for i in range(self.num_layers):
+            init_states.append(self.cell_list[i].init_hidden(batch_size, image_size))
+        return init_states
+
+    @staticmethod
+    def _check_kernel_size_consistency(kernel_size):
+        if not (isinstance(kernel_size, tuple) or
+                (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):
+            raise ValueError('`kernel_size` must be tuple or list of tuples')
+
+    @staticmethod
+    def _extend_for_multilayer(param, num_layers):
+        if not isinstance(param, list):
+            param = [param] * num_layers
+        return param

veg/MultiTaskConvLSTM.py

@@ -0,0 +1,129 @@
+from ConvLSTM import ConvLSTM
+import torch
+import torch.nn as nn
+from collections import defaultdict
+
+#MLP definition
+class MLP_5D(nn.Module):
+    def __init__(self, height, width):
+        super(MLP_5D, self).__init__()
+        # Define the fully connected layers
+        self.fc1 = nn.Linear(64, 128)  # Input channels = 41, output features = 128
+        self.dropout1 = nn.Dropout(0.05)
+        self.fc2 = nn.Linear(128, 64)  # Output features = 64
+        self.dropout2 = nn.Dropout(0.05)
+        self.fc3 = nn.Linear(64, 1)    # Final output, reducing to 1 channel
+
+        self.height = height
+        self.width = width
+
+    def forward(self, x):
+        batch_size, timesteps, channels, height, width = x.shape
+        
+        # Ensure the input spatial dimensions match the expected height and width
+        assert height == self.height and width == self.width, "Height and width mismatch"
+        
+        # Reshape to (batch * timesteps * height * width, channels)
+        x = x.permute(0, 1, 3, 4, 2).reshape(-1, channels)
+        # print(x.shape)
+        
+        # Apply MLP (Fully connected layers)
+        x = self.fc1(x)
+        x = torch.nn.functional.softplus(x)
+        x = self.dropout1(x)
+        x = self.fc2(x)
+        x = torch.nn.functional.softplus(x)
+        x = self.dropout2(x)
+        x = self.fc3(x)
+        x = torch.nn.functional.softplus(x)
+        
+        # Reshape back to (batch, timesteps, 1, height, width)
+        x = x.view(batch_size, timesteps, self.height, self.width, 1).permute(0, 1, 4, 2, 3)
+
+        return x
+    
+  
+# MultiTask ConvLSTM definition
+
+class ConvLSTMNetwork(nn.Module):
+    def __init__(self, input_dim, hidden_dims, kernel_size, num_layers, output_channels, batch_first=True, pool_size=(2,2)):
+        super(ConvLSTMNetwork, self).__init__()
+        
+        # ConvLSTM module
+        self.convlstm = ConvLSTM(input_dim=input_dim,
+                                 hidden_dim=hidden_dims,
+                                 kernel_size=kernel_size,
+                                 num_layers=num_layers,
+                                 batch_first=batch_first,
+                                 bias=True,
+                                 return_all_layers=True)
+        
+        # Batch Normalization for each ConvLSTM layer's output
+        self.batch_norms = nn.ModuleList([
+            nn.BatchNorm3d(hidden_dim) for hidden_dim in hidden_dims
+        ])
+
+        # Final Conv3D layer for regression pathway
+        self.conv3d = nn.Conv3d(in_channels=hidden_dims[-1],
+                                out_channels=output_channels,
+                                kernel_size=(1, 3, 3),
+                                padding=(0, 1, 1))
+
+        # MLP for regression output: (B,T,C,H,W) -> (B,T,1,H,W)
+        self.mlp = MLP_5D(height=81, width=97)
+
+        # Classification head for pixel-level zero precipitation probability
+        # We'll produce (B,T,1,H,W) as well:
+        # The classification head takes (B,C,T,H,W) input. We'll reorder dimensions before applying it.
+        # Then apply Sigmoid to get probabilities between 0 and 1.
+        self.classification_head = nn.Sequential(
+            nn.Conv3d(output_channels, 1, kernel_size=(1,1,1)),  # from C to 1 channel
+            nn.Sigmoid()
+        )
+
+        self.activation_variance = defaultdict(list)
+
+    def forward(self, x):
+        """
+        x: (B, T, input_dim, H, W)
+        """
+        # Forward through ConvLSTM
+        layer_output_list, last_state_list = self.convlstm(x)
+        
+        # Apply batch norms
+        for i, output in enumerate(layer_output_list):
+            # output: (B, T, C, H, W)
+            output = output.permute(0, 2, 1, 3, 4)  # (B, C, T, H, W) for BatchNorm3d
+            output = self.batch_norms[i](output)
+            output = output.permute(0, 2, 1, 3, 4)  # back to (B, T, C, H, W)
+
+            #Track variance across spatial dimensions for hooks with activation tracking 
+            activation_variance = output.var(dim=(3, 4)).mean().item()
+            self.activation_variance[f"ConvLSTM_layer_{i}"].append(activation_variance)
+
+            layer_output_list[i] = output
+        
+        # Take output from the last ConvLSTM layer
+        final_output = layer_output_list[-1]  # (B, T, C, H, W)
+
+        # Pass through Conv3D: needs (B,C,T,H,W)
+        final_output = final_output.permute(0, 2, 1, 3, 4)  # (B,C,T,H,W)
+        final_output = self.conv3d(final_output)
+        # Now final_output: (B, output_channels, T, H, W)
+
+        # Return to (B,T,C,H,W) for MLP (regression)
+        final_output_t = final_output.permute(0, 2, 1, 3, 4)  # (B,T,C,H,W)
+
+        # Regression output
+        regression_output = self.mlp(final_output_t)  # (B,T,1,H,W)
+
+        # Classification output:
+        # The classification head is defined for (B,C,T,H,W), so reorder again
+        final_output_c = final_output  # still (B,output_channels,T,H,W)
+        classification_output = self.classification_head(final_output_c)
+        # classification_output: (B,1,T,H,W)
+
+        # Permute classification output to match (B,T,1,H,W) format
+        classification_output = classification_output.permute(0, 2, 1, 3, 4)  # (B,T,1,H,W)
+
+        return regression_output, classification_output

veg/MultiTaskConvLSTM_veg_variables.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7e8b5e8b33db227257dd794dd6ba5ff10d4759aeed1f758550d4f3acb69cc26
+size 1383333

veg/data/normalized_test_data_veg_input.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:075c60cf83a2d7ef68720b72c77f41a034d6418b67085e016fff6a5bdfef878c
+size 2631223074

veg/example_inference.py

@@ -0,0 +1,230 @@
+# example_inference
+import torch
+from MultiTaskConvLSTM import ConvLSTMNetwork
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
+import torch
+import toch.nn as nn
+from tqdm.auto import tqdm
+from utils import (
+    mse, mae, nash_sutcliffe_efficiency, r2_score, pearson_correlation,
+    spearman_correlation, percentage_error, percentage_bias,
+    kendall_tau, spatial_correlation
+)
+import torch.optim as optim
+
+
+device = 'cpu'
+
+height = 81
+width = 97
+
+set_lookback = 1
+set_forecast_horizon = 1
+
+#Define variables for evaluation
+batch_size = 16
+time_steps_out = set_forecast_horizon
+channels = 14
+
+#Variable names
+variable_names = ['10 metre U wind component', '10 metre V wind component', '2 metre dewpoint temperature', '2 metre temperature', 'UV visible albedo for direct radiation (climatological)', 'Total column rain water', 'Volumetric soil water layer 1', 'Leaf area index, high vegetation', 'Leaf area index, low vegetation', 'Forecast surface roughness', 'Total precipitation', 'Time-integrated surface latent heat net flux', 'Evaporation']
+
+# Adjust input_dim and output_channels according to your data specifics
+model = ConvLSTMNetwork(
+    input_dim=14 * set_lookback, 
+    hidden_dims=[14, 32, 64], 
+    kernel_size=(3,3), 
+    num_layers=3, 
+    output_channels=64 * set_forecast_horizon, 
+    batch_first=True
+).to(device)
+
+# Define separate loss functions
+loss_fn = nn.MSELoss()       # For regression output
+bce_loss_fn = nn.BCELoss()       # For classification output
+
+optimizer = optim.AdamW(model.parameters(), lr = 0.005) 
+
+checkpoint = torch.load("MultiTaskConvLSTM_veg_variables")
+model.load_state_dict(checkpoint['model_state_dict'])
+
+# If you want to move the model to the GPU (optional, depending on your setup)
+model.to(device)  # Assuming you have a variable `device` for CUDA or CPU
+
+# Ensure that the model is in evaluation mode if you're using it for inference
+model.eval()
+
+print("Model loaded successfully")
+
+
+threshold = 0.1
+precip_index = 10
+
+def evaluate(model, test_loader, reg_loss_fn, class_loss_fn, device, variable_names, height, width):
+    """
+    Evaluate the model on the test set for both regression and classification tasks.
+    """
+    model.eval()  # Set the model to evaluation model
+
+    # input_to_true = {'zero_to_non_zero': 0, 'non_zero_to_zero': 0}
+    # input_to_pred_REG = {'zero_to_non_zero': 0, 'non_zero_to_zero': 0}
+    # input_to_pred_CLASS = {'zero_to_non_zero': 0, 'non_zero_to_zero': 0}
+
+    test_reg_loss = 0.0
+    test_class_loss = 0.0
+    test_total_loss = 0.0
+
+    y_true_reg = []  # List to store true values for regression
+    y_pred_reg = []  # List to store predicted values for regression
+
+    y_pred_reg2 = []
+
+    y_true_class = []  # List to store true values for classification
+    y_pred_class = []  # List to store predicted probabilities for classification
+
+    # Disable gradient computation
+    with torch.no_grad():
+        for X_test, y_test, y_zero_test in tqdm(test_loader, desc="Evaluating on Test Set"):
+            # Move the batch to the device
+            X_test, y_test, y_zero_test = X_test.to(device), y_test.to(device), y_zero_test.to(device)
+
+            # Reshape inputs and targets
+            batch_size, time_steps_in, channels_in, grid_points = X_test.shape
+            batch_size, time_steps_out, channels_out, grid_points = y_test.shape
+            X_test = X_test.view(batch_size, time_steps_in, channels_in, height, width)
+            y_test = y_test.view(batch_size, time_steps_out, channels_out, height, width)
+            y_zero_test = y_zero_test.view(batch_size, time_steps_out, channels_out, height, width)
+
+            # Forward pass
+            regression_output, classification_output = model(X_test)
+
+            classification_predictions = (classification_output > 0.7).float()
+
+            # Compute regression loss
+            reg_loss = reg_loss_fn(regression_output, y_test)
+
+            # Compute classification loss
+            class_loss = class_loss_fn(classification_output, y_zero_test)
+
+            # Total loss
+            total_loss = reg_loss + class_loss
+
+            regression_output2 = torch.where(classification_predictions == 0, regression_output, classification_predictions)
+
+            # Accumulate losses
+            test_reg_loss += reg_loss.item() * X_test.size(0)
+            test_class_loss += class_loss.item() * X_test.size(0)
+            test_total_loss += total_loss.item() * X_test.size(0)
+
+            # Collect true and predicted values for regression and classification
+            y_true_reg.append(y_test.cpu())
+            y_pred_reg.append(regression_output.cpu())
+            y_pred_reg2.append(regression_output2.cpu())
+            y_true_class.append(y_zero_test.cpu())
+            y_pred_class.append(classification_output.cpu())
+
+    # Normalize losses by the total dataset size
+    test_reg_loss /= len(test_loader)
+    test_class_loss /= len(test_loader)
+    test_total_loss /= len(test_loader)
+
+    print(f"Test Regression Loss: {test_reg_loss:.16f}")
+    print(f"Test Classification Loss: {test_class_loss:.16f}")
+    print(f"Test Total Loss: {test_total_loss:.16f}")
+
+    y_true_reg_flat = torch.cat(y_true_reg, dim=0).flatten()  # Keep as PyTorch tensor
+    y_pred_reg_flat = torch.cat(y_pred_reg, dim=0).flatten()  # Keep as PyTorch tensor
+    y_true_class_flat = torch.cat(y_true_class, dim=0).flatten()  # Keep as PyTorch tensor
+    y_pred_class_flat = torch.cat(y_pred_class, dim=0).flatten()  # Keep as PyTorch tensor
+
+    # Compute regression metrics
+    regression_metrics = {
+        "MSE": mse(y_true_reg_flat, y_pred_reg_flat),
+        "MAE": mae(y_true_reg_flat, y_pred_reg_flat),
+        "NSE": nash_sutcliffe_efficiency(y_true_reg_flat, y_pred_reg_flat),
+        "R2": r2_score(y_true_reg_flat, y_pred_reg_flat),
+        "Pearson": pearson_correlation(y_true_reg_flat, y_pred_reg_flat),
+        "Spearman": spearman_correlation(y_true_reg_flat, y_pred_reg_flat),
+        "NSE": nash_sutcliffe_efficiency(y_true_reg_flat, y_pred_reg_flat),
+        "Percentage Error": percentage_error(y_true_reg_flat, y_pred_reg_flat),
+        "Percentage Bias": percentage_bias(y_true_reg_flat, y_pred_reg_flat),
+        "Kendall Tau": kendall_tau(y_true_reg_flat, y_pred_reg_flat),
+        "Spatial Correlation": spatial_correlation(y_true_reg_flat, y_pred_reg_flat)}
+
+    print("\nRegression Metrics:")
+    for metric, value in regression_metrics.items():
+        print(f"{metric}: {value:.16f}")
+
+
+    # Compute classification metrics
+    classification_metrics = {
+        "Accuracy": accuracy_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "Precision": precision_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "Recall": recall_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "F1": f1_score(y_true_class_flat, (y_pred_class_flat > 0.7)),
+        "ROC-AUC": roc_auc_score(y_true_class_flat, y_pred_class_flat),
+    }
+
+    print("\nClassification Metrics:")
+    for metric, value in classification_metrics.items():
+        print(f"{metric}: {value:.16f}")
+
+    torch.save({
+        'y_true_reg': y_true_reg_flat,
+        'y_pred_reg': y_pred_reg_flat,
+        'y_true_class': y_true_class_flat,
+        'y_pred_class': y_pred_class_flat,
+    }, 'results')
+
+    return test_total_loss, regression_metrics, classification_metrics
+
+
+"""
+EXPECTED DATALOADER BATCH FORMAT (normalized_test_data):
+
+Each batch must be a tuple: (X_batch, y_batch, y_zero_batch)
+
+X_batch contains the previous hours variables. y_batch contains the next hour's precipitation.
+y_zero_batch contains the next hour's precipitation thresholded as 0 for precipiation <=0.1mm/h and 
+1 for precipitation >0.1mm.
+
+Shapes BEFORE reshaping inside `evaluate`:
+    X_batch:   (B, T_in, C_in, G)        # G = H*W = 81*97 = 7857
+    y_batch:   (B, T_out, C_out, G)
+    y_zero_batch: (B, T_out, C_out, G)   # binary 0/1 "zero-precip" targets
+
+    If your preprocessing produces (B,T, C, H, W), reshape to (B, T, C, H*W) before inference.
+
+DTypes:
+    X_batch, y_batch:       torch.float32
+    y_zero_batch:           torch.float32 (will be used with BCELoss)
+
+Reshaping done in 'evaluate':
+    X_test = X_batch.view(B, T_in, C_in, H, W)         -> (B, T_in, C_in, 81, 97)
+    y_test = y_batch.view(B, T_out, C_out, H, W)       -> (B, T_out, C_out, 81, 97)
+    y_zero_test = y_zero_batch.view(B, T_out, C_out, H, W)
+
+Model input:
+    model expects X_test shaped (B, T_in, input_dim, H, W)
+    where input_dim == 9 * set_lookback  (with set_lookback=1 -> input_dim=9)
+
+Notes:
+  • Make sure G == H*W (i.e., 7857 for 81x97).
+  • C_out for precipitation should be 1 (one target channel), and y_zero_batch
+    is the 0/1 mask for “zero precipitation” at each pixel & time.
+  • y_zero_batch should be probabilities/labels in {0,1} for BCELoss.
+"""
+
+normalized_test_data = torch.load("data/normalized_test_data_veg_input.pth")
+
+
+test_total_loss, regression_metrics, classification_metrics = evaluate(
+    model=model,
+    test_loader=normalized_test_data,
+    reg_loss_fn=loss_fn,
+    class_loss_fn=bce_loss_fn,
+    device=device,
+    variable_names=variable_names,
+    height=height,
+    width=width,
+)

veg/utils.py

@@ -0,0 +1,92 @@
+#Definition of evaluation metrics
+from scipy.stats import pearsonr, spearmanr
+import torch.nn.functional as F
+import torch
+from scipy.stats import kendalltau
+import scipy.stats as stats
+
+
+def nash_sutcliffe_efficiency(observed, predicted):
+    # Ensure inputs are tensors on the CPU
+    observed = observed.cpu()
+    predicted = predicted.cpu()
+
+    # Compute the numerator and denominator
+    numerator = torch.sum((observed - predicted) ** 2)
+    denominator = torch.sum((observed - torch.mean(observed)) ** 2)
+
+    # Calculate NSE
+    nse = 1 - (numerator / denominator)
+    return nse.item()
+
+def pearson_correlation(y_true, y_pred):
+    y_true = y_true.view(-1).cpu().numpy()  # Flatten and move to CPU
+    y_pred = y_pred.view(-1).cpu().numpy()  # Flatten and move to CPU
+    
+    return pearsonr(y_true, y_pred)[0]  # Return the correlation coefficient
+
+def spearman_correlation(y_true, y_pred):
+    y_true = y_true.view(-1).cpu().numpy()  # Flatten and move to CPU
+    y_pred = y_pred.view(-1).cpu().numpy()  # Flatten and move to CPU
+    
+    return spearmanr(y_true, y_pred).correlation  # Return the Spearman correlation
+
+def mse(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+    
+    return torch.mean((y_true - y_pred) ** 2).item()
+
+def mae(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+    
+    return torch.mean(torch.abs(y_true - y_pred)).item()
+
+def percentage_error(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+    
+    return 100 * torch.mean((y_pred - y_true) / (y_true + 1e-6)).item()
+
+def percentage_bias(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+
+    return 100 * torch.sum(y_pred - y_true) / (torch.sum(y_true) + 1e-6)
+
+def kendall_tau(y_true, y_pred):
+    y_true = y_true.view(-1).cpu().numpy()  # Flatten and move to CPU
+    y_pred = y_pred.view(-1).cpu().numpy()  # Flatten and move to CPU
+    
+    return kendalltau(y_true, y_pred).correlation  # Return the Kendall Tau
+
+def r2_score(y_true, y_pred):
+    # Ensure inputs are tensors on the CPU
+    y_true = y_true.cpu()
+    y_pred = y_pred.cpu()
+
+    ss_total = torch.sum((y_true - torch.mean(y_true)) ** 2)
+    ss_residual = torch.sum((y_true - y_pred) ** 2)
+    
+    return 1 - (ss_residual / (ss_total + 1e-6)).item()
+
+def spatial_correlation(y_true, y_pred):
+    # Flatten the tensors to work with them
+    y_true_flat = y_true.view(-1).cpu()
+    y_pred_flat = y_pred.view(-1).cpu()
+
+    # Compute the numerator: sum(P * T)
+    numerator = torch.sum(y_pred_flat * y_true_flat)
+
+    # Compute the denominator: sqrt(sum(P^2) * sum(T^2))
+    denominator = torch.sqrt(torch.sum(y_pred_flat ** 2) * torch.sum(y_true_flat ** 2))
+
+    # Compute the correlation (add epsilon to avoid division by zero)
+    correlation = numerator / (denominator)
+
+    return correlation.item()