feat: initial implementation of DGMR solar radiation nowcasting models

- Add the model architecture
- Include pre-trained model weights for DGMR_SO & Generator_only models
- Add inference pipeline and sample data for testing
- Configure project structure with requirements and documentation

.gitattributes

@@ -0,0 +1,5 @@
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__/

README.md

@@ -0,0 +1,73 @@
+# DGMR Solar Radiation Nowcasting
+
+A deep learning model for solar radiation nowcasting using modified [Deep Generative Model of Rainfall (DGMR)](https://www.nature.com/articles/s41586-021-03854-z) architecture with Solar radiation Output (DGMR-SO). The model predicts clearsky index and converts it to solar radiation for up to 36 time steps ahead.
+
+## Overview
+
+This repository implements two model variants for solar radiation forecasting:
+- **DGMR_SO**: Full Deep Generative Models with one generator and two discriminators during the training stage
+- **Generator_only**: Only one generator during the training stage
+
+The model uses multiple input sources:
+- **Himawari satellite data**: Clearsky index calculated from Himawari satellite data
+- **WRF Prediction**: Clearsky index from WRF's solar irradiation prediction
+- **Topography**: Static topographical features
+- **Time features**: Temporal sin/cos encoding for day and hour
+
+## Installation
+
+1. Clone the repository:
+```bash
+git clone <repository-url>
+cd DGMR_SolRad
+```
+
+2. Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+## Requirements
+
+- Python 3.x
+- PyTorch 2.4.0
+- NumPy 1.26.4
+- einops 0.8.0
+
+## Usage
+
+### Basic Inference
+
+Run solar radiation prediction using the pre-trained models:
+
+```bash
+python inference.py --model-type DGMR_SO --basetime 202504131100
+```
+
+### Command Line Arguments
+
+- `--model-type`: Choose between `DGMR_SO` or `Generator_only` (default: `DGMR_SO`)
+- `--basetime`: Timestamp for input data in format YYYYMMDDHHMM (default: `202504131100`)
+
+### Example
+
+```bash
+# Using DGMR_SO model
+python inference.py --model-type DGMR_SO --basetime 202504131100
+
+# Using Generator-only model
+python inference.py --model-type Generator_only --basetime 202507151200
+```
+
+## Sample Data
+
+The repository includes sample data files:
+- `sample_202504131100.npz`
+- `sample_202504161200.npz`
+- `sample_202507151200.npz`
+
+## Model Weights
+
+Pre-trained weights are available for both models:
+- `model_weights/DGMR_SO/ft36/weights.ckpt`
+- `model_weights/Generator_only/ft36/weights.ckpt`

inference.py

@@ -0,0 +1,75 @@
+import torch
+import numpy as np
+import argparse
+
+from model_architect.inference_model import Predictor
+
+
+def data_loading(BASETIME, device):
+    data_npz = np.load(f'./sample_data/sample_{BASETIME}.npz')
+
+    inputs = {}
+    for key in data_npz:
+        inputs[key] = torch.from_numpy(data_npz[key]).to(device)
+
+    return inputs
+
+
+def model_loading(model_type, device):
+    if model_type == 'DGMR_SO':
+        ckpt_path = './model_weights/DGMR_SO/ft36/weights.ckpt'
+    elif model_type == 'Generator_only':
+        ckpt_path = './model_weights/Generator_only/ft36/weights.ckpt'
+
+    model = Predictor(
+        model_type=model_type,
+    )
+
+    ckpt = torch.load(ckpt_path, weights_only=True)
+    model.load_state_dict(ckpt['generator_state_dict'])
+    model.eval()
+    model.to(device)
+
+    return model
+
+
+def arg_parse():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--model-type',
+        type=str,
+        default='DGMR_SO',
+        choices=[
+            'Generator_only',
+            'DGMR_SO'])
+    parser.add_argument('--basetime', type=str, default='202504131100')
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = arg_parse()
+    model_type = args.model_type
+    BASETIME = args.basetime
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    inputs = data_loading(BASETIME, device)
+    model = model_loading(model_type, device)
+
+    # prediction
+    with torch.no_grad():
+        pred_clr_idx = model(
+            inputs['Himawari'],
+            inputs['WRF'],
+            inputs['topo'],
+            inputs['time_feat'],
+            pred_step=36,
+        )
+    pred_clr_idx = pred_clr_idx.squeeze(2).clamp(0, 1)
+
+    # transform clearsky index to solar radiation
+    pred_srad = pred_clr_idx * inputs['clearsky']  # dim: (1, 36, 512, 512)
+
+    # save prediction
+    np.save(f'./pred_{BASETIME}_{model_type}.npy', pred_srad.cpu().numpy())
+    print('Done')

model_architect/DGMR_SO/discriminator.py

@@ -0,0 +1,143 @@
+import torch
+import torch.nn as nn
+from ..components.common import DBlock
+import einops
+
+
+class TemporalDiscriminator(nn.Module):
+    def __init__(self, in_channel: int, base_c: int = 24):
+        super().__init__()
+        self.in_channel = in_channel
+
+        # (N, C, D, H, W)
+        self.down_sample = nn.AvgPool3d(
+            kernel_size=(1, 2, 2), 
+            stride=(1, 2, 2)
+        )
+
+        # (N, D, C, H, W)
+        self.space_to_depth = nn.PixelUnshuffle(downscale_factor=2)
+
+        # in_channel, out_channel
+        # Conv3D -> (N, C, D, H, W)
+        chn = base_c * 2 * in_channel
+        self.d3_1 = DBlock(in_channel=in_channel * 4,
+                           out_channel=chn,
+                           conv_type='3d', apply_relu=False, apply_down=True)
+
+        self.d3_2 = DBlock(in_channel=chn,
+                           out_channel=2 * chn,
+                           conv_type='3d', apply_relu=True, apply_down=True)
+
+        self.Dlist = nn.ModuleList()
+        for i in range(3):
+            chn = chn * 2
+            self.Dlist.append(
+                DBlock(in_channel=chn,
+                       out_channel=2 * chn,
+                       conv_type='2d', apply_relu=True, apply_down=True)
+            )
+
+        self.last_D = DBlock(in_channel=2 * chn,
+                             out_channel=2 * chn,
+                             conv_type='2d', apply_relu=True, apply_down=False)
+
+        self.fc = nn.Linear(2 * chn, 1)
+        self.relu = nn.ReLU()
+        # TODO: close bn
+        # self.bn = nn.BatchNorm1d(2*chn)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.down_sample(x)
+        x = self.space_to_depth(x)
+
+        # go through the 3D Block
+        # from (N, D, C, H, W) -> (N, C, D, H, W)
+        x = torch.permute(x, dims=(0, 2, 1, 3, 4))
+        x = self.d3_1(x)
+        x = self.d3_2(x)
+        # go through 2D Block, permute -> (N, D, C, H, W)
+        x = torch.permute(x, dims=(0, 2, 1, 3, 4))
+        n, d, c, h, w = list(x.size())
+        ####
+        fea = einops.rearrange(x, "n d c h w -> (n d) c h w")
+        for dd in self.Dlist:
+            fea = dd(fea)
+
+        fea = self.last_D(fea)
+
+        fea = torch.sum(self.relu(fea), dim=[2, 3])
+        # fea = self.bn(fea)
+        fea = self.fc(fea)
+
+        y = torch.reshape(fea, (n, d, 1))  # dims -> (N, D, 1)
+        y = torch.sum(y, keepdim=True, dim=1)  # dims -> (N, 1, 1)
+
+        return y
+
+
+class SpatialDiscriminator(nn.Module):
+    def __init__(self, in_channel: int, base_c: int = 24):
+        super().__init__()
+        self.in_channel = in_channel
+
+        # (N, C, H, W)
+        self.down_sample = nn.AvgPool2d(kernel_size=(2, 2), stride=(2, 2))
+
+        # (N, C, H, W)
+        self.space_to_depth = nn.PixelUnshuffle(downscale_factor=2)
+
+        # first Dblock doesn't apply relu
+        chn = base_c * in_channel
+        self.d1 = DBlock(in_channel=in_channel * 4,
+                         out_channel=chn * 2,
+                         conv_type='2d', apply_relu=False, apply_down=True)
+
+        self.Dlist = nn.ModuleList()
+        for i in range(4):
+            chn = chn * 2
+            self.Dlist.append(
+                DBlock(in_channel=chn,
+                       out_channel=2 * chn,
+                       conv_type='2d', apply_relu=True, apply_down=True)
+            )
+
+        self.last_D = DBlock(in_channel=2 * chn,
+                             out_channel=2 * chn,
+                             conv_type='2d', apply_relu=True, apply_down=False)
+
+        self.fc = nn.Linear(2 * chn, 1)
+        self.relu = nn.ReLU()
+        # TODO: close BN
+        # self.bn = nn.BatchNorm1d(2*chn)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # note: input dims -> (N, D, C, H, W)
+        # randomly pick up 8 out of 18
+        perm = torch.randperm(x.shape[1])
+        random_idx = perm[:8]
+
+        fea = x[:, random_idx, :, :, :]
+
+        n, d, c, h, w = list(fea.size())
+
+        fea = einops.rearrange(fea, "n d c h w -> (n d) c h w")
+        fea = self.down_sample(fea)
+        fea = self.space_to_depth(fea)
+
+        # apply DBlock
+        fea = self.d1(fea)
+        for dd in self.Dlist:
+            fea = dd(fea)
+
+        fea = self.last_D(fea)
+
+        # sum
+        fea = torch.sum(self.relu(fea), dim=[2, 3])
+        # fea = self.bn(fea)
+        fea = self.fc(fea)
+
+        y = torch.reshape(fea, (n, d, 1))  # dims -> (N, D, 1)
+        y = torch.sum(y, keepdim=True, dim=1)  # dims -> (N, 1, 1)
+
+        return y

model_architect/DGMR_SO/generator.py

@@ -0,0 +1,274 @@
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn.utils.parametrizations import spectral_norm
+from torch.distributions import normal
+import einops
+
+from ..components.ConvGRU import ConvGRU
+from ..components.common import GBlock, Up_GBlock, LBlock, AttentionLayer, DBlock
+
+
+class Sampler(nn.Module):
+    def __init__(self, in_channels, base_channels=24, up_step=4):
+        """
+        up_step should be the same as down_step in context-condition-stack
+
+        """
+        super().__init__()
+        base_c = base_channels
+
+        self.up_steps = up_step
+        self.convgru_list = nn.ModuleList()
+        self.conv1x1_list = nn.ModuleList()
+        self.gblock_list = nn.ModuleList()
+        self.upg_list = nn.ModuleList()
+
+        for i in range(self.up_steps):
+            # different scale
+            chs1 = base_c * 2**(self.up_steps - i + 1) * in_channels
+            chs2 = base_c * 2**(self.up_steps - i) * in_channels
+            # convgru
+            self.convgru_list.append(
+                ConvGRU(chs1, chs2, 3)
+            )
+            # conv1x1
+            self.conv1x1_list.append(
+                spectral_norm(
+                    nn.Conv2d(
+                        in_channels=chs2,
+                        out_channels=chs1,
+                        kernel_size=(
+                            1,
+                            1))))
+            # GBlock
+            self.gblock_list.append(
+                GBlock(in_channel=chs1, out_channel=chs1)
+            )
+            # upgblock
+            self.upg_list.append(
+                Up_GBlock(in_channel=chs1)
+            )
+
+            # output
+            # TODO: close Batch
+            # self.bn = nn.BatchNorm2d(chs2)
+            self.relu = nn.ReLU()
+            self.last_conv1x1 = spectral_norm(
+                nn.Conv2d(in_channels=chs2,
+                          out_channels=4,
+                          kernel_size=(1, 1))
+            )
+            self.depth_to_space = nn.PixelShuffle(upscale_factor=2)
+
+    def forward(self, latents, img_feat, init_states, pred_step):
+        """
+        latent dim -> (N, C, W, H)
+        img_feat dim -> (N, D, C, W, H)
+        init_states dim -> (N, C, W, H)
+        """
+        # expand time dims at axis=1
+        latents = torch.unsqueeze(latents, dim=1)
+
+        # repeat batch_size
+        if latents.shape[0] == 1:
+            # expand batch
+            latents = einops.repeat(
+                latents,
+                "b d c h w -> (repeat b) d c h w",
+                repeat=init_states[0].shape[0])
+        # repeat time step
+        latents = einops.repeat(
+            latents, "b d c h w -> b (repeat d) c h w", repeat=pred_step
+        )
+        # TODO: add img feas
+        # seq_out = latents + img_feat
+        seq_out = torch.cat((img_feat, latents), dim=2)
+        # init_states should be reversed
+        # scale up step
+        for i in range(self.up_steps):
+            seq_out = self.convgru_list[i](
+                seq_out, init_states[(self.up_steps - 1) - i])
+            # seq_out output shape -> (D, N, C, W, H)
+
+            # forloop time step
+            seq_out = [self.conv1x1_list[i](h) for h in seq_out]
+            seq_out = [self.gblock_list[i](h) for h in seq_out]
+            seq_out = [self.upg_list[i](h) for h in seq_out]
+            # output: seq_out list dim -> D * [N, C, H, W]
+            # should stack at dim == 1 to become (N, D, C, H, W)
+            seq_out = torch.stack(seq_out, dim=1)
+
+        # final output
+        # forloop time step
+        output = []
+        for t in range(seq_out.shape[1]):
+            y = seq_out[:, t, :, :, :]
+            # y = self.bn(y)
+            y = self.relu(y)
+            y = self.last_conv1x1(y)
+            y = self.depth_to_space(y)
+            output.append(y)
+
+        output = torch.stack(output, dim=1)
+
+        return output
+
+class LatentConditionStack(nn.Module):
+    def __init__(self, out_channels, down_step, sigma, attn=True):
+        """
+        in_shape dims -> e.g. (8, 8) -> (H, W)
+        x) base_c is 1/96 of out_channels
+        base_c is set to 4
+        """
+        super().__init__()
+
+        self.down_step = down_step
+        self.base_c = out_channels // 96
+        if self.base_c < 4:
+            self.base_c = 4
+
+        self.out_channels = out_channels
+        self.attn = attn
+
+        # define the distribution
+        self.dist = normal.Normal(loc=0.0, scale=sigma)
+
+        self.conv3x3 = spectral_norm(
+            nn.Conv2d(
+                in_channels=self.base_c,
+                out_channels=self.base_c,
+                kernel_size=(3, 3),
+                padding=1
+            )
+        )
+
+        cc = self.base_c
+        self.l1 = LBlock(cc, cc * 3)
+        self.l2 = LBlock(cc * 3, cc * 6)
+        self.l3 = LBlock(cc * 6, cc * 24)
+        if self.attn:
+            self.attn = AttentionLayer(cc * 24, cc * 24)
+        self.l4 = LBlock(cc * 24, self.out_channels)
+
+    def forward(self, x, batch_size=1, z=None):
+        """
+        x shape -> (batch_size, time, c, width, height)
+        """
+        width = x.shape[3]
+        height = x.shape[4]
+        # shape after downstep
+        s_w = width // (2 * 2**self.down_step)
+        s_h = height // (2 * 2**self.down_step)
+
+        in_shape = [self.base_c] + [s_w, s_h]
+
+        target_shape = [batch_size] + in_shape
+        if z is None:
+            z = self.dist.sample(target_shape)
+            z = z.type_as(x)
+
+        # first conv
+        z = self.conv3x3(z)
+
+        # Lblock
+        z = self.l1(z)
+        z = self.l2(z)
+        z = self.l3(z)
+        if self.attn:
+            z = self.attn(z)
+
+        z = self.l4(z)
+
+        return z
+
+# TODO: modification(Change the amount of parameters)
+
+
+class ContextConditionStack(nn.Module):
+    def __init__(self,
+                 in_channels: int = 1,
+                 base_channels: int = 24,
+                 down_step: int = 4,
+                 prev_step: int = 4):
+        """
+        base_channels: e.g. 24 -> output_channel: 384
+        output_channel: base_c*in_c*2**(down_step-2) * prev_step
+        down_step: int
+        prev_step: int
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.down_step = down_step
+        self.prev_step = prev_step
+        ###
+        base_c = base_channels
+        in_c = in_channels
+
+        # different scales channels
+        chs = [4 * in_c] + [base_c * in_c * 2 **
+                            (i + 1) for i in range(down_step)]
+
+        self.space_to_depth = nn.PixelUnshuffle(downscale_factor=2)
+        self.Dlist = nn.ModuleList()
+        self.convList = nn.ModuleList()
+        for i in range(down_step):
+            self.Dlist.append(
+                DBlock(in_channel=chs[i],
+                       out_channel=chs[i + 1],
+                       apply_relu=True, apply_down=True)
+            )
+
+            self.convList.append(
+                spectral_norm(
+                    nn.Conv2d(in_channels=prev_step * chs[i + 1],
+                              out_channels=prev_step * chs[i + 1] // 4,
+                              kernel_size=(3, 3),
+                              padding=1)
+                )
+            )
+
+        # ReLU
+        self.relu = nn.ReLU()
+
+    def forward(self,
+                x: torch.Tensor) -> Tuple[torch.Tensor,
+                                          torch.Tensor,
+                                          torch.Tensor,
+                                          torch.Tensor]:
+        """
+        ## input dims -> (N, D, C, H, W)
+        """
+        x = self.space_to_depth(x)
+        tsteps = x.shape[1]
+        assert tsteps == self.prev_step
+
+        # different feature index represent different scale
+        # features
+        # [scale1 -> [t1, t2, t3, t4], scale2 -> [t1, t2, t3, t4], scale3 -> [....]]
+        features = [[] for i in range(tsteps)]
+
+        for st in range(tsteps):
+            in_x = x[:, st, :, :, :]
+            # in_x -> (Batch(N), C, H, W)
+            for scale in range(self.down_step):
+                in_x = self.Dlist[scale](in_x)
+                features[scale].append(in_x)
+
+        out_scale = []
+        for i, cc in enumerate(self.convList):
+            # after stacking, dims -> (Batch, Time, C, H, W)
+            # and mixing layer is to concat Time, C
+            stacked = self._mixing_layer(torch.stack(features[i], dim=1))
+            out = self.relu(cc(stacked))
+            out_scale.append(out)
+
+        return out_scale
+
+    def _mixing_layer(self, x):
+        # conver from (N, Time, C, H, W) -> (N, Time*C, H, W)
+        # Then apply Conv2d
+        stacked = einops.rearrange(x, "b t c h w -> b (t c) h w")
+
+        return stacked

model_architect/DGMR_SO/img_extractor.py

@@ -0,0 +1,49 @@
+import torch.nn as nn
+
+from ..components.common import DBlock
+
+
+class ImageExtractor(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            apply_down_flag,
+            down_step=4):
+        """
+        in_c -> 1
+        x) base_c is 1/96 of out_channels
+        base_c is set to 4
+        """
+        super().__init__()
+        self.down_step = down_step
+
+        self.base_c = out_channels // 96
+        if self.base_c < 4:
+            self.base_c = 4
+        cc = self.base_c
+
+        self.space_to_depth = nn.PixelUnshuffle(downscale_factor=2)
+
+        chs = [in_channels * 4, cc * 3, cc * 6, cc * 24, out_channels]
+        self.DList = nn.ModuleList()
+        for i in range(down_step):
+            self.DList.append(
+                DBlock(
+                    in_channel=chs[i],
+                    out_channel=chs[i + 1],
+                    conv_type='2d',
+                    apply_down=apply_down_flag[i]
+                ),
+            )
+
+    def forward(self, x):
+        """
+        x
+        """
+        y = self.space_to_depth(x)
+        # forloop ImageExtractor
+        for i in range(self.down_step):
+            y = self.DList[i](y)
+
+        return y

model_architect/DGMR_SO/model.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+from .discriminator import TemporalDiscriminator, SpatialDiscriminator
+from .generator import Sampler, ContextConditionStack, LatentConditionStack
+from .img_extractor import ImageExtractor
+
+
+class Generator(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            base_channels,
+            down_step,
+            prev_step,
+            sigma
+    ):
+        super().__init__()
+        out_channels = base_channels * \
+            2**(down_step - 2) * prev_step * in_channels
+
+        self.latentStack = LatentConditionStack(
+            out_channels=out_channels,
+            down_step=down_step,
+            sigma=sigma
+        )
+
+        self.contextStack = ContextConditionStack(
+            in_channels=in_channels,
+            base_channels=base_channels,
+            down_step=down_step,
+            prev_step=prev_step
+        )
+
+        self.sampler = Sampler(
+            in_channels=in_channels,
+            base_channels=base_channels,
+            up_step=down_step
+        )
+
+        self.encode_time = nn.Linear(
+            4, base_channels * 2**(down_step) * in_channels)
+
+        self.topo_extractor = ImageExtractor(
+            in_channels=1,
+            out_channels=base_channels * 2**(down_step - 1) * in_channels,
+            apply_down_flag=[True, True, True, True],
+            down_step=down_step
+        )
+
+        self.nwp_extractor = ImageExtractor(
+            in_channels=1,  # TODO: fixed now
+            out_channels=base_channels * 2**(down_step - 1) * in_channels,
+            apply_down_flag=[False, True, False, True],
+            down_step=down_step
+        )
+
+    def forward(self, x, x2, topo, datetime_feat, pred_step=36):
+        """
+        x: input seq -> dims (N, D, C, H, W)
+        x2: input seq (WRF) -> dims (N, D, C, H, W)
+        topo: topography -> dims (N, 1, H=512, W=512)
+        datetime_feat -> dims (N, D, 4)
+        """
+        context_inits = self.contextStack(x)
+        batch_size = context_inits[0].shape[0]
+        zlatent = self.latentStack(x, batch_size=batch_size)
+
+        # topo feature
+        topo_feat = self.topo_extractor(topo)
+        # encode time feature
+        time_feat = self.encode_time(datetime_feat)
+        # extract nwp feature
+        nwp_feat = []
+        # forloop x2
+        for i in range(x2.shape[1]):
+            nwp_ = self.nwp_extractor(x2[:, i, ...])
+            # concat topo and nwp feature
+            concat_feat = torch.cat((nwp_, topo_feat), dim=1)
+            nwp_feat.append(concat_feat)
+        nwp_feat = torch.stack(nwp_feat, dim=1)
+        fuse_feat = nwp_feat + time_feat.unsqueeze(-1).unsqueeze(-1)
+
+        pred = self.sampler(zlatent, fuse_feat, context_inits, pred_step)
+
+        return pred
+
+
+class Discriminator(nn.Module):
+    def __init__(self, in_channels, base_channels):
+        super().__init__()
+        self.spatial = SpatialDiscriminator(
+            in_channel=in_channels, base_c=base_channels)
+        self.temporal = TemporalDiscriminator(
+            in_channel=in_channels, base_c=base_channels)
+
+    def forward(self, x, y):
+        """
+        x -> dims (N, D, C, H, W) e.g. input_frames
+        y -> dims (N, D, C, H, W) e.g. output_grames
+        """
+        spatial_out = self.spatial(y)
+        temporal_out = self.temporal(torch.cat([x, y], dim=1))
+
+        dis_out = torch.cat([spatial_out, temporal_out], dim=1)
+
+        return dis_out

model_architect/Generator_only/generator_clr_idx_wrf_topot.py

@@ -0,0 +1,282 @@
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn.utils.parametrizations import spectral_norm
+import einops
+
+from ..components.ConvGRU import ConvGRUCell
+from ..components.common import GBlock, Up_GBlock, DBlock
+
+
+class Sampler(nn.Module):
+    def __init__(self, in_channels, base_channels=24, up_step=4):
+        """
+        up_step should be the same as down_step in context-condition-stack
+
+        """
+        super().__init__()
+        base_c = base_channels
+
+        self.up_steps = up_step
+        self.convgru_list = nn.ModuleList()
+        self.conv1x1_list = nn.ModuleList()
+        self.gblock_list = nn.ModuleList()
+        self.upg_list = nn.ModuleList()
+
+        # image extractor
+        self.img_extractor = ImageExtractor(
+            in_channels=2,
+            out_channels=base_c * 2**(self.up_steps) * in_channels,
+            apply_down_flag=[True, True, True, True],
+            down_step=self.up_steps
+        )
+
+        self.nwp_extractor = ImageExtractor(
+            in_channels=1,
+            out_channels=base_c * 2**(self.up_steps) * in_channels,
+            apply_down_flag=[False, True, False, True],
+            down_step=self.up_steps
+        )
+
+        self.encode_time = nn.Linear(
+            4, base_c * 2**(self.up_steps + 1) * in_channels)
+
+        for i in range(self.up_steps):
+            # different scale
+            chs1 = base_c * 2**(self.up_steps - i + 1) * in_channels
+            chs2 = base_c * 2**(self.up_steps - i) * in_channels
+            # convgru
+            self.convgru_list.append(
+                ConvGRUCell(chs1, chs2, 3)
+            )
+            # conv1x1
+            self.conv1x1_list.append(
+                spectral_norm(
+                    nn.Conv2d(
+                        in_channels=chs2,
+                        out_channels=chs1,
+                        kernel_size=(
+                            1,
+                            1))))
+            # GBlock
+            self.gblock_list.append(
+                GBlock(in_channel=chs1, out_channel=chs1)
+            )
+            # upgblock
+            self.upg_list.append(
+                Up_GBlock(in_channel=chs1)
+            )
+
+            # output
+            # self.bn = nn.BatchNorm2d(chs2)
+            self.relu = nn.ReLU()
+            self.last_conv1x1 = spectral_norm(
+                nn.Conv2d(in_channels=chs2,
+                          out_channels=4,
+                          kernel_size=(1, 1))
+            )
+            self.depth_to_space = nn.PixelShuffle(upscale_factor=2)
+
+    def forward(
+            self,
+            input_img,
+            nwp_inputs,
+            topo,
+            time_feat,
+            init_states,
+            pred_step,
+            thres=None):
+        """
+        input_img dim  -> (N, Tstep, C, W, H) -> Tstep can be 1 or pred_steps
+        nwp_inputs dim -> (N, Tstep, C, W, H)
+        init_states dim -> [(N, C, W, H)-1, (N, C, W, H)-2, ...]
+        probs -> (tsteps)
+        """
+        hh = init_states
+        output = []
+        img_t_len = input_img.shape[1]
+
+        xx = None
+        for t in range(pred_step):
+            time_emb = self.encode_time(time_feat[:, t])  # time emb
+
+            # The 1st tstep image should be truth
+            # and 2nd tstep image apply schedule sampling
+            if t == 0:
+                input_ = input_img[:, 0, :, :, :]
+            else:
+                if thres is not None and img_t_len > 1:
+                    # use groud truth
+                    if torch.rand(1) < thres:
+                        input_ = input_img[:, t, :, :, :]
+                    else:
+                        input_ = xx
+                else:
+                    input_ = xx
+
+            nwp_in = nwp_inputs[:, t]
+            # image extractor -> extract T step image
+            # xx is the output of image extractor
+            input_ = torch.cat((input_, topo), dim=1)
+            xx = self.img_extractor(input_)
+            nwp_feat = self.nwp_extractor(nwp_in)
+            xx = torch.cat((xx, nwp_feat), dim=1)
+
+            xx = xx + time_emb[:, :, None, None]  # add time embedding
+
+            for up in range(self.up_steps):
+                # convGRU
+                # init_states should be reversed
+                h_index = (self.up_steps - 1) - up
+                xx, out_hh = self.convgru_list[up](xx, hh[h_index])
+                hh[h_index] = out_hh
+                # conv1x1
+                xx = self.conv1x1_list[up](xx)
+                # gblock
+                xx = self.gblock_list[up](xx)
+                # upg_list
+                xx = self.upg_list[up](xx)
+
+            # xx = self.bn(xx)
+            xx = self.relu(xx)
+            xx = self.last_conv1x1(xx)
+            xx = self.depth_to_space(xx)
+
+            # prediction
+            output.append(xx)
+
+        output = torch.stack(output, dim=1)
+
+        return output
+
+
+class ImageExtractor(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            apply_down_flag,
+            down_step=4):
+        """
+        in_c -> 1
+        x) base_c is 1/96 of out_channels
+        base_c is set to 4
+        """
+        super().__init__()
+        self.down_step = down_step
+
+        self.base_c = out_channels // 96
+        if self.base_c < 4:
+            self.base_c = 4
+        cc = self.base_c
+
+        self.space_to_depth = nn.PixelUnshuffle(downscale_factor=2)
+
+        chs = [in_channels * 4, cc * 3, cc * 6, cc * 24, out_channels]
+        self.DList = nn.ModuleList()
+        for i in range(down_step):
+            self.DList.append(
+                DBlock(
+                    in_channel=chs[i],
+                    out_channel=chs[i + 1],
+                    conv_type='2d',
+                    apply_down=apply_down_flag[i]
+                ),
+            )
+
+    def forward(self, x):
+        y = self.space_to_depth(x)
+        # forloop ImageExtractor
+        for i in range(self.down_step):
+            y = self.DList[i](y)
+
+        return y
+
+
+class ContextConditionStack(nn.Module):
+    def __init__(self,
+                 in_channels: int = 1,
+                 base_channels: int = 24,
+                 down_step: int = 4,
+                 prev_step: int = 4):
+        """
+        base_channels: e.g. 24 -> output_channel: 384
+        output_channel: base_c*in_c*2**(down_step-2) * prev_step
+        down_step: int
+        prev_step: int
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.down_step = down_step
+        self.prev_step = prev_step
+        ###
+        base_c = base_channels
+        in_c = in_channels
+
+        # different scales channels
+        chs = [4 * in_c] + [base_c * in_c * 2 **
+                            (i + 1) for i in range(down_step)]
+
+        self.space_to_depth = nn.PixelUnshuffle(downscale_factor=2)
+        self.Dlist = nn.ModuleList()
+        self.convList = nn.ModuleList()
+        for i in range(down_step):
+            self.Dlist.append(
+                DBlock(in_channel=chs[i],
+                       out_channel=chs[i + 1],
+                       apply_relu=True, apply_down=True)
+            )
+
+            self.convList.append(
+                spectral_norm(
+                    nn.Conv2d(in_channels=prev_step * chs[i + 1],
+                              out_channels=prev_step * chs[i + 1] // 4,
+                              kernel_size=(3, 3),
+                              padding=1)
+                )
+            )
+
+        # ReLU
+        self.relu = nn.ReLU()
+
+    def forward(self,
+                x: torch.Tensor) -> Tuple[torch.Tensor,
+                                          torch.Tensor,
+                                          torch.Tensor,
+                                          torch.Tensor]:
+        """
+        ## input dims -> (N, D, C, H, W)
+        """
+        x = self.space_to_depth(x)
+        tsteps = x.shape[1]
+        assert tsteps == self.prev_step
+
+        # different feature index represent different scale
+        # features
+        # [scale1 -> [t1, t2, t3, t4], scale2 -> [t1, t2, t3, t4], scale3 -> [....]]
+        features = [[] for i in range(tsteps)]
+
+        for st in range(tsteps):
+            in_x = x[:, st, :, :, :]
+            # in_x -> (Batch(N), C, H, W)
+            for scale in range(self.down_step):
+                in_x = self.Dlist[scale](in_x)
+                features[scale].append(in_x)
+
+        out_scale = []
+        for i, cc in enumerate(self.convList):
+            # after stacking, dims -> (Batch, Time, C, H, W)
+            # and mixing layer is to concat Time, C
+            stacked = self._mixing_layer(torch.stack(features[i], dim=1))
+            out = self.relu(cc(stacked))
+            out_scale.append(out)
+
+        return out_scale
+
+    def _mixing_layer(self, x):
+        # conver from (N, Time, C, H, W) -> (N, Time*C, H, W)
+        # Then apply Conv2d
+        stacked = einops.rearrange(x, "b t c h w -> b (t c) h w")
+
+        return stacked

model_architect/Generator_only/model_clr_idx.py

@@ -0,0 +1,57 @@
+import torch.nn as nn
+
+from .generator_clr_idx_wrf_topot import Sampler, ContextConditionStack
+
+
+class Generator(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        base_channels,
+        down_step,
+        prev_step
+    ):
+
+        super().__init__()
+        self.contextStack = ContextConditionStack(
+            in_channels=in_channels,
+            base_channels=base_channels,
+            down_step=down_step,
+            prev_step=prev_step
+        )
+
+        self.sampler = Sampler(
+            in_channels=in_channels,
+            base_channels=base_channels,
+            up_step=down_step,
+        )
+
+    def forward(
+        self,
+        x,
+        x2,
+        topo,
+        time_feat,
+        pred_step=12,
+        y=None,
+        thres=None
+    ):
+        """
+        x: input seq -> dims (N, T, C, H, W)
+        x2: input seq for WRF -> dims (N, T, C, H, W)
+        """
+        context_inits = self.contextStack(x)
+        if y is None:
+            y = x[:, -1:, :, :, :]
+
+        pred = self.sampler(
+            y,
+            x2,
+            topo,
+            time_feat,
+            context_inits,
+            pred_step,
+            thres
+        )
+
+        return pred

model_architect/components/ConvGRU.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import spectral_norm
+from torch.autograd import Variable
+
+
+class ConvGRUCell(torch.nn.Module):
+    def __init__(self, in_channel, out_channel, kernel_size=3):
+        super().__init__()
+        padding = kernel_size // 2
+        self.out_channel = out_channel
+
+        self.conv1 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=in_channel + out_channel,
+                out_channels=2 * out_channel,
+                kernel_size=kernel_size,
+                padding=padding
+            )
+        )
+
+        self.conv2 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=in_channel + out_channel,
+                out_channels=out_channel,
+                kernel_size=kernel_size,
+                padding=padding
+            )
+        )
+
+    def forward(self, x, h_st):
+        """
+        x -> dim (Batch, channels*2, width, height)
+        h_st -> dim (Batch, channels, width, height)
+        """
+        x_shape = x.shape
+        h_shape = h_st.shape
+        # resize width, height
+        if h_shape[2] > x_shape[2]:
+            w_l = 1
+        else:
+            w_l = 0
+
+        if h_shape[3] > x_shape[3]:
+            h_b = 1
+        else:
+            h_b = 0
+
+        x = F.pad(x, (0, h_b, w_l, 0), "reflect")
+
+        # print(x.shape, h_st.shape)
+        xx = torch.cat([x, h_st], dim=1)
+        xx = self.conv1(xx)
+        gamma, beta = torch.split(xx, self.out_channel, dim=1)
+
+        reset_gate = torch.sigmoid(gamma)
+        update_gate = torch.sigmoid(beta)
+
+        out = torch.cat([x, h_st * reset_gate], dim=1)
+        out = torch.tanh(self.conv2(out))
+
+        out = (1 - update_gate) * out + h_st * update_gate
+        new_st = out
+
+        return out, new_st
+
+
+class ConvGRU(torch.nn.Module):
+    def __init__(self, in_channel, out_channel, kernel_size):
+        super().__init__()
+
+        self.out_channel = out_channel
+        self.convgru_cell = ConvGRUCell(in_channel, out_channel, kernel_size)
+
+    def _get_init_state(self, batch_size, imd_w, imd_h, dtype):
+        state = Variable(
+            torch.zeros(
+                batch_size,
+                self.out_channel,
+                self.h,
+                self.w)).type(dtype)
+
+        return state
+
+    def forward(self, x_sequence, init_hidden=None):
+        """
+        Args:
+            x_sequence shape -> (batch_size, time, c, width, height)
+        Return:
+            outputs shape -> (time, batch_size, c, width, height)
+        """
+        seq_len = x_sequence.shape[1]
+
+        img_w = x_sequence.shape[3]
+        img_h = x_sequence.shape[4]
+
+        dtype = x_sequence.type()
+        if init_hidden is None:
+            hidden_state = self._get_init_state(
+                x_sequence.shape[0], img_w, img_h, dtype)
+        else:
+            hidden_state = init_hidden
+
+        out_list = []
+        for t in range(seq_len):
+            out, hidden_state = self.convgru_cell(
+                x_sequence[:, t, :, :, :], hidden_state)
+            out_list.append(out)
+
+        outputs = torch.stack(out_list, dim=0)
+
+        return outputs

model_architect/components/common.py

@@ -0,0 +1,261 @@
+from torch.nn.utils.parametrizations import spectral_norm
+import torch
+from torch.nn import functional as F
+import einops
+
+
+class GBlock(torch.nn.Module):
+    def __init__(self, in_channel: int, out_channel: int):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+        # TODO: close batch
+        # self.bn1 = torch.nn.BatchNorm2d(in_channel)
+        # self.bn2 = torch.nn.BatchNorm2d(out_channel)
+
+        self.relu = torch.nn.ReLU()
+        # conv1x1
+        self.conv1x1 = spectral_norm(
+            torch.nn.Conv2d(in_channel, out_channel, kernel_size=1)
+        )
+
+        self.conv3x3_1 = spectral_norm(
+            torch.nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1)
+        )
+        self.conv3x3_2 = spectral_norm(
+            torch.nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1)
+        )
+
+    def forward(self, x: torch.Tensor):
+        # if shape is different then applied
+        if x.shape[1] != self.out_channel:
+            res = self.conv1x1(x)
+        else:
+            res = x.clone()
+
+        # first
+        # x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv3x3_1(x)
+        # second
+        # x = self.bn2(x)
+        x = self.relu(x)
+        x = self.conv3x3_2(x)
+
+        y = x + res
+
+        return y
+
+
+class Up_GBlock(torch.nn.Module):
+    def __init__(self, in_channel: int):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = int(in_channel / 2)
+
+        # TODO: close batch
+        # self.bn1 = torch.nn.BatchNorm2d(in_channel)
+        # self.bn2 = torch.nn.BatchNorm2d(self.out_channel)
+        # self.bn2 = torch.nn.BatchNorm2d(in_channel)
+
+        self.relu = torch.nn.ReLU()
+        self.up = torch.nn.Upsample(scale_factor=2, mode='nearest')
+
+        self.conv1x1 = spectral_norm(
+            torch.nn.Conv2d(in_channel, self.out_channel, kernel_size=1)
+        )
+
+        self.conv3x3_1 = spectral_norm(
+            torch.nn.Conv2d(in_channel, in_channel, kernel_size=3, padding=1)
+        )
+        self.conv3x3_2 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channel,
+                self.out_channel,
+                kernel_size=3,
+                padding=1))
+
+    def forward(self, x):
+        res = self.up(x)
+        res = self.conv1x1(res)
+
+        # x = self.bn1(x)
+        x = self.relu(x)
+        x = self.up(x)
+        x = self.conv3x3_1(x)
+
+        # x = self.bn2(x)
+        x = self.relu(x)
+        x = self.conv3x3_2(x)
+
+        y = x + res
+
+        return y
+
+
+class DBlock(torch.nn.Module):
+    def __init__(
+            self,
+            in_channel: int,
+            out_channel: int,
+            conv_type='2d',
+            apply_relu=True,
+            apply_down=False):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.apply_relu = apply_relu
+        self.apply_down = apply_down
+
+        # construct layer
+        if conv_type == '2d':
+            self.avg_pool = torch.nn.AvgPool2d(kernel_size=2, stride=2)
+            conv = torch.nn.Conv2d
+        elif conv_type == '3d':
+            self.avg_pool = torch.nn.AvgPool3d(kernel_size=2, stride=2)
+            conv = torch.nn.Conv3d
+
+        self.relu = torch.nn.ReLU()
+        self.conv1x1 = spectral_norm(
+            conv(in_channel, out_channel, kernel_size=1)
+        )
+
+        self.conv3x3_1 = spectral_norm(
+            conv(in_channel, out_channel, kernel_size=3, padding=1)
+        )
+        self.conv3x3_2 = spectral_norm(
+            conv(out_channel, out_channel, kernel_size=3, padding=1)
+        )
+
+    def forward(self, x):
+        # Residual block
+        if x.shape[1] != self.out_channel:
+            res = self.conv1x1(x)
+        else:
+            res = x.clone()
+        if self.apply_down:
+            res = self.avg_pool(res)
+
+        ##
+        if self.apply_relu:
+            x = self.relu(x)
+        x = self.conv3x3_1(x)
+        x = self.relu(x)
+        x = self.conv3x3_2(x)
+        if self.apply_down:
+            x = self.avg_pool(x)
+
+        # connect
+        y = res + x
+
+        return y
+
+
+class LBlock(torch.nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+        self.relu = torch.nn.ReLU()
+        conv = torch.nn.Conv2d
+        self.conv1x1 = conv(
+            in_channel,
+            (out_channel - in_channel),
+            kernel_size=1)
+
+        self.conv3x3_1 = conv(in_channel, in_channel, kernel_size=3, padding=1)
+        self.conv3x3_2 = conv(
+            in_channel,
+            out_channel,
+            kernel_size=3,
+            padding=1)
+
+    def forward(self, x):
+        res = torch.cat([x, self.conv1x1(x)], dim=1)
+
+        x = self.relu(x)
+        x = self.conv3x3_1(x)
+        x = self.relu(x)
+        x = self.conv3x3_2(x)
+
+        y = x + res
+
+        return y
+
+# Attention Layer
+
+
+def attention_einsum(q, k, v):
+    """
+    Apply self-attention to tensors
+    """
+
+    # Reshape 3D tensor to 2D tensor with first dimension L = h x w
+    k = einops.rearrange(k, "h w c -> (h w) c")  # [h, w, c] -> [L, c]
+    v = einops.rearrange(v, "h w c -> (h w) c")  # [h, w, c] -> [L, c]
+
+    # Einstein summation corresponding to the query * key operation.
+    beta = F.softmax(torch.einsum("hwc, Lc->hwL", q, k), dim=-1)
+
+    # Einstein summation corresponding to the attention * value operation.
+    out = torch.einsum("hwL, Lc->hwc", beta, v)
+
+    return out
+
+
+class AttentionLayer(torch.nn.Module):
+    def __init__(self, in_channel, out_channel, ratio_kq=8, ratio_v=8):
+        super().__init__()
+
+        self.ratio_kq = ratio_kq
+        self.ratio_v = ratio_v
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+        # compute query, key, and value using 1x1 convolution
+        self.query = torch.nn.Conv2d(
+            in_channel,
+            out_channel // ratio_kq,
+            kernel_size=1,
+            bias=False
+        )
+
+        self.key = torch.nn.Conv2d(
+            in_channel,
+            out_channel // ratio_kq,
+            kernel_size=1,
+            bias=False
+        )
+
+        self.value = torch.nn.Conv2d(
+            in_channel,
+            out_channel // ratio_v,
+            kernel_size=1,
+            bias=False
+        )
+
+        self.conv = torch.nn.Conv2d(
+            out_channel // 8,
+            out_channel,
+            kernel_size=1,
+            bias=False
+        )
+
+        self.gamma = torch.nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        query = self.query(x)
+        key = self.key(x)
+        value = self.value(x)
+        # apply attention
+        out = []
+        for i in range(x.shape[0]):
+            out.append(attention_einsum(query[i], key[i], value[i]))
+
+        out = torch.stack(out, dim=0)
+        out = self.gamma * self.conv(out)
+        out = out + x  # skip connection
+
+        return out

model_architect/inference_model.py

@@ -0,0 +1,40 @@
+import torch.nn as nn
+
+from .DGMR_SO.model import Generator as DGMR_SO
+from .Generator_only.model_clr_idx import Generator as Generator_only
+
+
+class Predictor(nn.Module):
+    def __init__(
+        self,
+        model_type,
+    ):
+        super().__init__()
+
+        if model_type == 'DGMR_SO':
+            self.generator = DGMR_SO(
+                in_channels=1,
+                base_channels=24,
+                down_step=4,
+                prev_step=4,
+                sigma=1
+            )
+
+        elif model_type == 'Generator_only':
+            self.generator = Generator_only(
+                in_channels=1,
+                base_channels=24,
+                down_step=4,
+                prev_step=4,
+            )
+
+    def forward(self, x, x2, topo, datetime_feat, pred_step=36):
+        """
+        x: input seq -> dims (N, D, C, H, W)
+        x2: input seq (WRF) -> dims (N, D, C, H, W)
+        topo: topography -> dims (N, 1, H=512, W=512)
+        datetime_feat -> dims (N, D, 4)
+        """
+        pred = self.generator(x, x2, topo, datetime_feat, pred_step=pred_step)
+
+        return pred

model_weights/DGMR_SO/ft36/weights.ckpt

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

model_weights/Generator_only/ft36/weights.ckpt

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

requirements.txt

@@ -0,0 +1,3 @@
+numpy==1.26.4
+torch==2.4.0
+einops==0.8.0

sample_data/sample_202504131100.npz

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

sample_data/sample_202504161200.npz

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

sample_data/sample_202507151200.npz

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