model/fuxi_ocean.py

"""Executable tile-streamed FuXi-Ocean reconstruction."""

import math

import numpy as np
import torch
from torch import nn
import torch.distributed as dist


FORMAT_VERSION = "fuxi_ocean_sampled_tiles_v2"
GLOBAL_SHAPE = (2160, 4320)
HISTORY_STEPS = 4
OCEAN_CHANNELS = 105
ATMOSPHERE_CHANNELS = 5
DEPTHS = (0, 2, 4, 6, 10, 15, 20, 30, 40, 50, 60, 70, 80, 90, 100, 125,
           150, 200, 300, 400, 500, 600, 800, 1000, 1250, 1500)


def enumerate_global_tiles(global_shape, tile_shape, overlap=0):
    """Enumerate all windows and non-overlapping ownership crops in row-major order."""
    height, width = map(int, global_shape)
    tile_h, tile_w = map(int, tile_shape)
    overlap = int(overlap)
    if min(height, width, tile_h, tile_w) <= 0 or overlap < 0 or overlap >= min(tile_h, tile_w):
        raise ValueError("invalid global/tile shape or overlap")

    def starts(size, tile):
        values = list(range(0, max(size - tile, 0) + 1, tile - overlap))
        last = max(size - tile, 0)
        if not values or values[-1] != last:
            values.append(last)
        return values

    ys, xs = starts(height, tile_h), starts(width, tile_w)
    records = []
    for row, y0 in enumerate(ys):
        for column, x0 in enumerate(xs):
            y1, x1 = min(y0 + tile_h, height), min(x0 + tile_w, width)
            top = 0 if row == 0 else (ys[row - 1] + tile_h + y0) // 2 - y0
            bottom = y1 - y0 if row == len(ys) - 1 else (y1 + ys[row + 1]) // 2 - y0
            left = 0 if column == 0 else (xs[column - 1] + tile_w + x0) // 2 - x0
            right = x1 - x0 if column == len(xs) - 1 else (x1 + xs[column + 1]) // 2 - x0
            records.append((y0, y1, x0, x1, top, bottom, left, right))
    return np.asarray(records, dtype=np.int32)


def assemble_tiles(tiles, tile_records, global_shape, require_complete=True):
    """Crop overlap by ownership records and assemble a dense field for verification/use."""
    tiles = np.asarray(tiles)
    records = np.asarray(tile_records)
    height, width = map(int, global_shape)
    output = np.zeros((*tiles.shape[1:-2], height, width), dtype=tiles.dtype)
    coverage = np.zeros((height, width), dtype=np.uint8)
    for tile, record in zip(tiles, records):
        y0, _y1, x0, _x1, top, bottom, left, right = map(int, record)
        output[..., y0 + top:y0 + bottom, x0 + left:x0 + right] = tile[..., top:bottom, left:right]
        coverage[y0 + top:y0 + bottom, x0 + left:x0 + right] = 1
    is_complete = bool(coverage.all())
    if require_complete and not is_complete:
        raise ValueError("tile set does not completely cover the global grid")
    return output, coverage


def ocean_channel_depth_indices():
    """Channel order: S26, T26, U26, V26, SSH."""
    return np.asarray(list(range(26)) * 4 + [0], dtype=np.int64)


def channel_mask(depth_mask):
    indices = torch.as_tensor(ocean_channel_depth_indices(), device=depth_mask.device)
    return depth_mask.index_select(1, indices)


class ChannelLayerNorm(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.norm = nn.LayerNorm(channels)

    def forward(self, x):
        return self.norm(x.movedim(1, -1)).movedim(-1, 1)


class PriorNetwork(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.spatial = nn.Conv2d(3, embed_dim, 1)
        self.temporal = nn.Sequential(nn.Linear(6, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim))

    def forward(self, coordinates, bathymetry, time_features):
        spatial = self.spatial(torch.cat((coordinates, bathymetry), dim=1))
        hour, day, step = time_features.unbind(dim=1)
        temporal = torch.stack((torch.sin(2 * math.pi * hour / 24), torch.cos(2 * math.pi * hour / 24),
                                torch.sin(2 * math.pi * day / 366), torch.cos(2 * math.pi * day / 366),
                                torch.sin(step / 40), torch.cos(step / 40)), dim=1)
        return spatial, self.temporal(temporal)


class SharedPatchEncoder(nn.Module):
    def __init__(self, in_channels, embed_dim, patch_size, levels):
        super().__init__()
        layers = []
        channels = in_channels
        self.skip_channels = []
        for level in range(levels):
            out = embed_dim * (2 ** level)
            layers += [nn.Conv2d(channels, out, patch_size, patch_size), ChannelLayerNorm(out), nn.GELU()]
            self.skip_channels.append(out)
            channels = out
        self.layers = nn.ModuleList(layers)
        self.levels = levels

    def forward(self, x, prior):
        skips = []
        for level in range(self.levels):
            conv, norm, activation = self.layers[level * 3:level * 3 + 3]
            x = conv(x)
            local_prior = torch.nn.functional.interpolate(prior, size=x.shape[-2:], mode="bilinear", align_corners=False)
            repeats = math.ceil(x.shape[1] / local_prior.shape[1])
            modulation = local_prior.repeat(1, repeats, 1, 1)[:, :x.shape[1]]
            x = activation(norm(x * (1.0 + torch.tanh(modulation))))
            skips.append(x)
        return x, skips


class AdaLNBlock(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
        self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 4))
        self.attention = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim))

    def forward(self, x, context):
        scale1, shift1, scale2, shift2 = self.modulation(context).chunk(4, dim=-1)
        q = self.norm1(x) * (1 + scale1[:, None]) + shift1[:, None]
        x = x + self.attention(q, q, q, need_weights=False)[0]
        q = self.norm2(x) * (1 + scale2[:, None]) + shift2[:, None]
        return x + self.ffn(q)


class FuXiOcean(nn.Module):
    def __init__(self, embed_dim=24, patch_size=2, encoder_levels=2, attention_layers=2, attention_heads=4):
        super().__init__()
        self.prior = PriorNetwork(embed_dim)
        self.encoder = SharedPatchEncoder(OCEAN_CHANNELS + ATMOSPHERE_CHANNELS, embed_dim, patch_size, encoder_levels)
        latent = embed_dim * 2 ** (encoder_levels - 1)
        self.fusion = nn.Sequential(nn.Conv2d(HISTORY_STEPS * latent, latent, 1), ChannelLayerNorm(latent), nn.GELU())
        self.blocks = nn.ModuleList(AdaLNBlock(latent, attention_heads) for _ in range(attention_layers))
        self.context_projection = nn.Linear(embed_dim, latent)
        decoder = []
        for level in reversed(range(encoder_levels)):
            skip = embed_dim * 2 ** level
            out = embed_dim * 2 ** max(level - 1, 0)
            decoder.append(nn.Sequential(nn.ConvTranspose2d(latent + skip, out, patch_size, patch_size),
                                         ChannelLayerNorm(out), nn.GELU()))
            latent = out
        self.decoder = nn.ModuleList(decoder)
        self.output = nn.Conv2d(latent, OCEAN_CHANNELS, 1)

    def forward(self, ocean_history, atmosphere, coordinates, bathymetry, depth_mask, time_features):
        if ocean_history.ndim != 5 or ocean_history.shape[1:3] != (HISTORY_STEPS, OCEAN_CHANNELS):
            raise ValueError("ocean_history must be [B,4,105,tile_h,tile_w]")
        spatial_prior, temporal_prior = self.prior(coordinates, bathymetry, time_features)
        encoded, all_skips = [], []
        for index in range(HISTORY_STEPS):
            state = torch.cat((ocean_history[:, index], atmosphere), dim=1)
            latent, skips = self.encoder(state, spatial_prior)
            encoded.append(latent)
            all_skips.append(skips)
        x = self.fusion(torch.cat(encoded, dim=1))
        shape = x.shape
        tokens = x.flatten(2).transpose(1, 2)
        context = self.context_projection(temporal_prior)
        for block in self.blocks:
            tokens = block(tokens, context)
        x = tokens.transpose(1, 2).reshape(shape)
        for decoder_index, decoder in enumerate(self.decoder):
            level = len(self.decoder) - 1 - decoder_index
            skip = torch.stack([values[level] for values in all_skips]).mean(dim=0)
            x = decoder(torch.cat((x, skip), dim=1))
        prediction = self.output(x)
        return prediction * channel_mask(depth_mask).to(prediction.dtype)


def latitude_weighted_charbonnier(prediction, target, latitude_deg, valid_mask, epsilon=1e-6):
    weights = torch.cos(torch.deg2rad(latitude_deg)).clamp_min(0)[:, None, :, None]
    mask = valid_mask.to(prediction.dtype)
    numerator = (torch.sqrt((prediction - target) ** 2 + epsilon ** 2) * weights * mask).sum()
    denominator = (weights * mask).sum()
    if dist.is_available() and dist.is_initialized():
        global_numerator = numerator.detach().clone()
        global_denominator = denominator.detach().clone()
        dist.all_reduce(global_numerator, op=dist.ReduceOp.SUM)
        dist.all_reduce(global_denominator, op=dist.ReduceOp.SUM)
        # DDP averages gradients, so this yields the gradient of the global ratio.
        return numerator * dist.get_world_size() / global_denominator.clamp_min(1.0)
    return numerator / denominator.clamp_min(1.0)


def weighted_error_sums(prediction, target, latitude_deg, valid_mask):
    weights = np.cos(np.deg2rad(latitude_deg))[:, None, :, None] * valid_mask
    error = prediction - target
    return ((weights * error ** 2).sum(axis=(0, 2, 3)),
            (weights * error).sum(axis=(0, 2, 3)), weights.sum(axis=(0, 2, 3)))