"""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)))