| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| import torch |
| import torch.nn as nn |
|
|
| from model.prithvi_wxc_official import PrithviWxC as _PrithviWxC |
|
|
|
|
| def _identity_scalers(n, device="cpu"): |
| return torch.zeros(n, device=device), torch.ones(n, device=device) |
|
|
|
|
| class PrithviWxC(nn.Module): |
| """ |
| Config-driven Prithvi WxC wrapper. |
| |
| The official model is an encoder-decoder vision Transformer (Hiera + MaxViT, |
| alternating local/global attention) that maps an input state history to a |
| target state. This wrapper builds the official network with identity |
| scalers so that a small YAML config can be exercised; loading the 2.3B |
| checkpoints published on Hugging Face is possible by calling |
| ``load_state_dict`` with the official checkpoint state (the buffer names, |
| e.g. ``input_scalers_mu``, are identical). |
| |
| Args: |
| in_channels: number of input/output parameters. |
| input_size_time: number of input timestamps (paper uses 2). |
| in_channels_static: number of static input channels (>= 1). |
| n_lats_px / n_lons_px: grid size in pixels. |
| patch_size_px: patch/token size in pixels. |
| mask_unit_size_px: mask unit size in pixels. |
| mask_ratio_inputs: input masking ratio (0..1, 0 disables). |
| embed_dim / n_blocks_encoder / n_blocks_decoder / mlp_multiplier / n_heads: |
| transformer architecture hyper-parameters (paper: 2560 / 25 / 5 / 4 / 16). |
| dropout / drop_path / parameter_dropout: regularization ratios. |
| residual: one of "none" / "temporal" / "climate". |
| masking_mode: "global" / "local" / "both". |
| positional_encoding: "absolute" or "fourier". |
| encoder_shifting / decoder_shifting: Swin-style shifting. |
| """ |
|
|
| def __init__( |
| self, |
| in_channels=4, |
| input_size_time=2, |
| in_channels_static=4, |
| n_lats_px=32, |
| n_lons_px=64, |
| patch_size_px=(2, 2), |
| mask_unit_size_px=(8, 8), |
| mask_ratio_inputs=0.0, |
| embed_dim=16, |
| n_blocks_encoder=1, |
| n_blocks_decoder=1, |
| mlp_multiplier=4.0, |
| n_heads=4, |
| dropout=0.0, |
| drop_path=0.0, |
| parameter_dropout=0.0, |
| residual="none", |
| masking_mode="global", |
| positional_encoding="absolute", |
| encoder_shifting=False, |
| decoder_shifting=False, |
| ): |
| super().__init__() |
| self.in_channels = int(in_channels) |
| self.input_size_time = int(input_size_time) |
| self.in_channels_static = int(in_channels_static) |
| self.n_lats_px = int(n_lats_px) |
| self.n_lons_px = int(n_lons_px) |
| self.residual = residual |
| self.positional_encoding = positional_encoding |
|
|
| mu_in, sigma_in = _identity_scalers(self.in_channels) |
| mu_s, sigma_s = _identity_scalers(self.in_channels_static) |
|
|
| self.model = _PrithviWxC( |
| in_channels=self.in_channels, |
| input_size_time=self.input_size_time, |
| in_channels_static=self.in_channels_static, |
| input_scalers_mu=mu_in, |
| input_scalers_sigma=sigma_in, |
| input_scalers_epsilon=1e-6, |
| static_input_scalers_mu=mu_s, |
| static_input_scalers_sigma=sigma_s, |
| static_input_scalers_epsilon=1e-6, |
| output_scalers=sigma_in, |
| n_lats_px=self.n_lats_px, |
| n_lons_px=self.n_lons_px, |
| patch_size_px=tuple(patch_size_px), |
| mask_unit_size_px=tuple(mask_unit_size_px), |
| mask_ratio_inputs=float(mask_ratio_inputs), |
| mask_ratio_targets=0.0, |
| embed_dim=int(embed_dim), |
| n_blocks_encoder=int(n_blocks_encoder), |
| n_blocks_decoder=int(n_blocks_decoder), |
| mlp_multiplier=int(mlp_multiplier), |
| n_heads=int(n_heads), |
| dropout=float(dropout), |
| drop_path=float(drop_path), |
| parameter_dropout=float(parameter_dropout), |
| residual=residual, |
| masking_mode=masking_mode, |
| positional_encoding=positional_encoding, |
| encoder_shifting=bool(encoder_shifting), |
| decoder_shifting=bool(decoder_shifting), |
| ) |
|
|
| def forward(self, x, static, input_time=None, lead_time=None, climate=None): |
| """ |
| Args: |
| x: Tensor of shape [batch, time, parameter, lat, lon]. |
| static: Tensor of shape [batch, static_channel, lat, lon]. |
| input_time: Tensor of shape [batch] (optional, default zeros). |
| lead_time: Tensor of shape [batch] (optional, default zeros). |
| climate: Optional Tensor of shape [batch, parameter, lat, lon]. |
| Returns: |
| Tensor of shape [batch, parameter, lat, lon]. |
| """ |
| if input_time is None: |
| input_time = torch.zeros(x.shape[0], device=x.device) |
| if lead_time is None: |
| lead_time = torch.zeros(x.shape[0], device=x.device) |
| batch = { |
| "x": x, |
| "y": x[:, -1], |
| "static": static, |
| "input_time": input_time, |
| "lead_time": lead_time, |
| } |
| if climate is not None: |
| batch["climate"] = climate |
| return self.model(batch) |
