OneScience-Group/CRAI-ClimateExtremes

1"""Independent PyTorch implementation of the CRAI reconstruction architecture."""2 3import torch4from torch import nn5import torch.nn.functional as F6 7 8class PartialConv2d(nn.Module):9    """Convolution normalized by valid pixels, with the corresponding mask update."""10 11    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):12        super().__init__()13        padding = kernel_size // 214        self.input_conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)15        self.kernel_area = float(kernel_size * kernel_size)16        self.stride = stride17        self.padding = padding18        self.kernel_size = kernel_size19        self.register_buffer("mask_kernel", torch.ones(1, 1, kernel_size, kernel_size))20 21    def forward(self, values, mask):22        if mask.shape[1] != 1:23            mask = mask.amax(dim=1, keepdim=True)24        with torch.no_grad():25            valid_count = F.conv2d(26                mask, self.mask_kernel, stride=self.stride, padding=self.padding27            )28            new_mask = (valid_count > 0).to(values.dtype)29            scale = self.kernel_area / valid_count.clamp_min(1.0)30        raw = self.input_conv(values * mask)31        if self.input_conv.bias is not None:32            bias = self.input_conv.bias.view(1, -1, 1, 1)33            raw = (raw - bias) * scale + bias34        else:35            raw = raw * scale36        return raw * new_mask, new_mask37 38 39class PConvBlock(nn.Module):40    def __init__(self, in_channels, out_channels, stride=1):41        super().__init__()42        self.partial = PartialConv2d(in_channels, out_channels, 3, stride)43        self.norm = nn.GroupNorm(1, out_channels)44 45    def forward(self, values, mask):46        values, mask = self.partial(values, mask)47        return F.leaky_relu(self.norm(values), 0.2), mask48 49 50class CRAIClimateExtremes(nn.Module):51    """Partial-convolution U-Net producing bounded extreme-index percentages."""52 53    def __init__(self, base_channels=8):54        super().__init__()55        c = base_channels56        self.enc1 = PConvBlock(1, c)57        self.enc2 = PConvBlock(c, 2 * c, stride=2)58        self.enc3 = PConvBlock(2 * c, 4 * c, stride=2)59        self.bottleneck = PConvBlock(4 * c, 8 * c, stride=2)60        self.dec3 = PConvBlock(8 * c + 4 * c, 4 * c)61        self.dec2 = PConvBlock(4 * c + 2 * c, 2 * c)62        self.dec1 = PConvBlock(2 * c + c, c)63        self.output = nn.Conv2d(c, 1, 1)64 65    @staticmethod66    def _upsample(values, mask, size):67        return (68            F.interpolate(values, size=size, mode="bilinear", align_corners=False),69            F.interpolate(mask, size=size, mode="nearest"),70        )71 72    def forward(self, inputs):73        if inputs.ndim != 4 or inputs.shape[1] != 2:74            raise ValueError("expected [B, 2, H, W] containing index and valid mask")75        values = inputs[:, :1] / 100.076        mask = (inputs[:, 1:2] > 0.5).to(values.dtype)77        e1, m1 = self.enc1(values, mask)78        e2, m2 = self.enc2(e1, m1)79        e3, m3 = self.enc3(e2, m2)80        x, m = self.bottleneck(e3, m3)81        x, m = self._upsample(x, m, e3.shape[-2:])82        x, m = self.dec3(torch.cat((x, e3), 1), torch.maximum(m, m3))83        x, m = self._upsample(x, m, e2.shape[-2:])84        x, m = self.dec2(torch.cat((x, e2), 1), torch.maximum(m, m2))85        x, m = self._upsample(x, m, e1.shape[-2:])86        x, _ = self.dec1(torch.cat((x, e1), 1), torch.maximum(m, m1))87        return 100.0 * torch.sigmoid(self.output(x))88 89 90if __name__ == "__main__":91    model = CRAIClimateExtremes()92    output = model(torch.rand(2, 2, 144, 192))93    print(tuple(output.shape), float(output.min()), float(output.max()))94