"""OneScience adapter for the official Microsoft Aurora implementation. The neural network is the flattened Microsoft Aurora implementation in this directory. This module owns only the contract between the flat ERA5 tensor used by OneScience and Aurora's structured batch object. """ from __future__ import annotations from datetime import datetime from importlib.util import find_spec from pathlib import Path from typing import Mapping, Sequence import numpy as np import torch from torch import nn _TORCHVISION_SCHEMA_LIBRARY: torch.library.Library | None = None def _operator_schema_exists(name: str) -> bool: try: torch._C._dispatch_has_kernel_for_dispatch_key(name, "Meta") except RuntimeError: return False return True def _register_missing_torchvision_schemas() -> None: """Allow timm imports when the site torchvision extension cannot be loaded.""" global _TORCHVISION_SCHEMA_LIBRARY spec = find_spec("torchvision") if spec is not None and spec.submodule_search_locations: package_dir = Path(next(iter(spec.submodule_search_locations))) for extension in sorted(package_dir.glob("_C*.so")): try: torch.ops.load_library(str(extension)) except (OSError, RuntimeError): continue break schemas = { "torchvision::nms": "nms(Tensor boxes, Tensor scores, float iou_threshold) -> Tensor", "torchvision::qnms": "qnms(Tensor boxes, Tensor scores, float iou_threshold) -> Tensor", } missing = [schema for name, schema in schemas.items() if not _operator_schema_exists(name)] if not missing: return library = torch.library.Library("torchvision", "FRAGMENT") for schema in missing: library.define(schema) _TORCHVISION_SCHEMA_LIBRARY = library _register_missing_torchvision_schemas() # The implementation modules are imported here as a single public model surface. Keeping the # official classes in their original modules preserves their state-dict names and checkpoint # compatibility while all callers use this file as the model entry point. from .aurora_area import area, expand_matrix, radius_earth from .aurora_batch import Batch, Metadata, interpolate, interpolate_numpy from .aurora_compat import * from .aurora_decoder import LinearPatchReconstruction, Perceiver3DDecoder from .aurora_encoder import Perceiver3DEncoder from .aurora_film import AdaptiveLayerNorm from .aurora_fourier import * from .aurora_insolation import insolation from .aurora_levelcond import LevelConditioned from .aurora_lora import * from .aurora_network import ( Aurora, Aurora12hPretrained, AuroraAirPollution, AuroraHighRes, AuroraPretrained, AuroraSmall, AuroraSmallPretrained, AuroraV1p5, AuroraV1p5Ensemble, AuroraWave, ) from .aurora_normalisation import * from .aurora_patchembed import LevelPatchEmbed from .aurora_perceiver import MLP, PerceiverAttention, PerceiverResampler from .aurora_posencoding import * from .aurora_rollout import rollout as aurora_rollout # Keep the official public API available from the unified model entry point. rollout = aurora_rollout from .aurora_swin3d import Swin3DTransformerBackbone, Swin3DTransformerBlock, WindowAttention from .aurora_util import * AURORA_SURFACE_VARS = ("2t", "10u", "10v", "msl") AURORA_STATIC_VARS = ("lsm", "z", "slt") AURORA_ATMOS_VARS = ("z", "u", "v", "t", "q") AURORA_LEVELS = (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000) ERA5_SURFACE_MAPPING = { "2t": "2m_temperature", "10u": "10m_u_component_of_wind", "10v": "10m_v_component_of_wind", "msl": "mean_sea_level_pressure", } ERA5_ATMOS_MAPPING = { "z": "geopotential_{level}", "u": "u_component_of_wind_{level}", "v": "v_component_of_wind_{level}", "t": "temperature_{level}", "q": "specific_humidity_{level}", } def _channel_order() -> tuple[str, ...]: channels = [ERA5_SURFACE_MAPPING[name] for name in AURORA_SURFACE_VARS] for name in AURORA_ATMOS_VARS: channels.extend(ERA5_ATMOS_MAPPING[name].format(level=level) for level in AURORA_LEVELS) return tuple(channels) DEFAULT_ERA5_CHANNEL_ORDER = _channel_order() def _as_float_tensor(value: torch.Tensor | np.ndarray, *, name: str) -> torch.Tensor: """Convert a field to float32 without changing its spatial dimensions.""" tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) if tensor.ndim != 2: raise ValueError(f"Static field {name!r} must have shape [H, W], got {tuple(tensor.shape)}") return tensor.to(dtype=torch.float32) def _parse_time(value: datetime | str) -> datetime: """Parse OneScience's ``YYYYmmddHH`` index and common ISO timestamps.""" if isinstance(value, datetime): return value text = str(value) for fmt in ("%Y%m%d%H", "%Y%m%d%H%M", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d %H:%M:%S"): try: return datetime.strptime(text, fmt) except ValueError: continue raise ValueError(f"Unsupported Aurora timestamp {value!r}") def _normalise_times(times: datetime | str | Sequence[datetime | str], batch_size: int) -> tuple[datetime, ...]: if isinstance(times, (datetime, str)): parsed = (_parse_time(times),) * batch_size else: parsed = tuple(_parse_time(value) for value in times) if len(parsed) != batch_size: raise ValueError(f"Expected {batch_size} timestamps, received {len(parsed)}") return parsed def _default_coordinates(height: int, width: int) -> tuple[torch.Tensor, torch.Tensor]: lat = torch.linspace(90.0, -90.0, height, dtype=torch.float32) lon = torch.arange(width, dtype=torch.float32) * (360.0 / width) return lat, lon class AuroraOneScience(nn.Module): """Aurora with a strict OneScience ERA5 tensor interface. Args: variant: ``"small"`` selects the official small architecture; ``"base"`` selects the official full 0.25-degree architecture. Both are randomly initialized unless a local checkpoint is explicitly supplied. static_path: Optional ``.npz`` containing ``lsm``, ``z``, ``slt``, ``lat`` and ``lon``. channel_order: Must equal the configured 69-channel ERA5 order exactly. load_checkpoint: If true, load ``checkpoint_path`` after constructing the network. """ def __init__( self, *, variant: str = "small", static_path: str | Path | None = None, channel_order: Sequence[str] = DEFAULT_ERA5_CHANNEL_ORDER, atmos_levels: Sequence[int | float] = AURORA_LEVELS, checkpoint_path: str | Path | None = None, load_checkpoint: bool = False, checkpoint_strict: bool = True, network_kwargs: Mapping[str, object] | None = None, ) -> None: super().__init__() self.channel_order = tuple(channel_order) if self.channel_order != DEFAULT_ERA5_CHANNEL_ORDER: raise ValueError( "Aurora requires the exact ERA5 channel order defined by " "DEFAULT_ERA5_CHANNEL_ORDER" ) self.atmos_levels = tuple(atmos_levels) if self.atmos_levels != AURORA_LEVELS: raise ValueError(f"Aurora requires pressure levels {AURORA_LEVELS}, got {self.atmos_levels}") kwargs = dict(network_kwargs or {}) kwargs.setdefault("max_history_size", 2) if variant == "small": self.network = AuroraSmallPretrained(**kwargs) elif variant in {"base", "full", "pretrained"}: self.network = AuroraPretrained(**kwargs) else: raise ValueError(f"Unknown Aurora variant {variant!r}; expected 'small' or 'base'") self.register_buffer("_static_lsm", torch.empty(0), persistent=False) self.register_buffer("_static_z", torch.empty(0), persistent=False) self.register_buffer("_static_slt", torch.empty(0), persistent=False) self.register_buffer("_lat", torch.empty(0), persistent=False) self.register_buffer("_lon", torch.empty(0), persistent=False) if static_path is not None: self.load_static_fields(static_path) if load_checkpoint: if checkpoint_path is None: raise ValueError("checkpoint_path is required when load_checkpoint=True") self.load_checkpoint_local(checkpoint_path, strict=checkpoint_strict) @property def aurora(self) -> nn.Module: """Expose the official network for activation-checkpointing and inspection.""" return self.network def load_static_fields(self, path: str | Path) -> None: """Load and validate project static fields from a compressed NumPy archive.""" static_path = Path(path).expanduser() with np.load(static_path) as data: required = set(AURORA_STATIC_VARS) | {"lat", "lon"} missing = sorted(required - set(data.files)) if missing: raise ValueError(f"{static_path}: missing static fields {missing}") fields = {name: _as_float_tensor(data[name], name=name) for name in AURORA_STATIC_VARS} lat = torch.as_tensor(data["lat"], dtype=torch.float32) lon = torch.as_tensor(data["lon"], dtype=torch.float32) self.set_static_fields(fields, lat=lat, lon=lon) def set_static_fields( self, static_vars: Mapping[str, torch.Tensor | np.ndarray], *, lat: torch.Tensor | np.ndarray | None = None, lon: torch.Tensor | np.ndarray | None = None, ) -> None: """Set static fields and coordinates without making them checkpoint parameters.""" missing = sorted(set(AURORA_STATIC_VARS) - set(static_vars)) if missing: raise ValueError(f"Missing Aurora static variables: {missing}") fields = {name: _as_float_tensor(static_vars[name], name=name) for name in AURORA_STATIC_VARS} shape = fields["lsm"].shape if any(value.shape != shape for value in fields.values()): raise ValueError("Aurora static fields must have identical [H, W] shapes") if lat is None or lon is None: default_lat, default_lon = _default_coordinates(*shape) lat = default_lat if lat is None else lat lon = default_lon if lon is None else lon lat = lat if isinstance(lat, torch.Tensor) else torch.as_tensor(lat) lon = lon if isinstance(lon, torch.Tensor) else torch.as_tensor(lon) lat = lat.to(dtype=torch.float32) lon = lon.to(dtype=torch.float32) if lat.shape != (shape[0],) or lon.shape != (shape[1],): raise ValueError(f"Coordinates must have shapes {(shape[0],)} and {(shape[1],)}") if not torch.all(lat[1:] < lat[:-1]) or not torch.all(lon[1:] > lon[:-1]): raise ValueError("Aurora coordinates must be decreasing latitude and increasing longitude") self._static_lsm = fields["lsm"] self._static_z = fields["z"] self._static_slt = fields["slt"] self._lat = lat self._lon = lon def load_checkpoint_local(self, path: str | Path, *, strict: bool = True) -> None: """Load an official Aurora checkpoint after model construction.""" self.network.load_checkpoint_local(str(Path(path).expanduser()), strict=strict) def configure_activation_checkpointing(self) -> None: """Delegate official activation checkpointing configuration.""" self.network.configure_activation_checkpointing() def _static_for_shape( self, shape: tuple[int, int], static_vars: Mapping[str, torch.Tensor | np.ndarray] | None, lat: torch.Tensor | np.ndarray | None, lon: torch.Tensor | np.ndarray | None, ) -> tuple[dict[str, torch.Tensor], torch.Tensor, torch.Tensor]: if static_vars is not None: missing = sorted(set(AURORA_STATIC_VARS) - set(static_vars)) if missing: raise ValueError(f"Missing Aurora static variables: {missing}") fields = {name: _as_float_tensor(static_vars[name], name=name) for name in AURORA_STATIC_VARS} lat_value = lat if lat is not None else self._lat lon_value = lon if lon is not None else self._lon else: fields = { "lsm": self._static_lsm, "z": self._static_z, "slt": self._static_slt, } lat_value = lat if lat is not None else self._lat lon_value = lon if lon is not None else self._lon if any(value.numel() == 0 for value in fields.values()): raise ValueError("Static fields are not configured; pass static_vars or static_path") if any(tuple(value.shape) != shape for value in fields.values()): raise ValueError(f"Static fields must match input grid {shape}") lat_tensor = lat_value if isinstance(lat_value, torch.Tensor) else torch.as_tensor(lat_value) lon_tensor = lon_value if isinstance(lon_value, torch.Tensor) else torch.as_tensor(lon_value) if lat_tensor.numel() == 0 or lon_tensor.numel() == 0: lat_tensor, lon_tensor = _default_coordinates(*shape) return fields, lat_tensor.to(torch.float32), lon_tensor.to(torch.float32) def tensor_to_batch( self, inputs: torch.Tensor, *, times: datetime | str | Sequence[datetime | str], static_vars: Mapping[str, torch.Tensor | np.ndarray] | None = None, lat: torch.Tensor | np.ndarray | None = None, lon: torch.Tensor | np.ndarray | None = None, ) -> Batch: """Convert ``[B, T, 69, H, W]`` ERA5 data to an official Aurora Batch.""" if inputs.ndim != 5: raise ValueError(f"Aurora input must have shape [B, T, C, H, W], got {tuple(inputs.shape)}") batch_size, history, channels, height, width = inputs.shape if channels != len(self.channel_order): raise ValueError(f"Expected {len(self.channel_order)} channels, got {channels}") if history < 1 or history > 2: raise ValueError(f"Aurora supports one or two history frames, got {history}") fields, lat_tensor, lon_tensor = self._static_for_shape( (height, width), static_vars, lat, lon ) channel_indices = {name: self.channel_order.index(name) for name in self.channel_order} surface = { name: inputs[:, :, channel_indices[ERA5_SURFACE_MAPPING[name]]] for name in AURORA_SURFACE_VARS } atmos = {} for name in AURORA_ATMOS_VARS: indices = [ channel_indices[ERA5_ATMOS_MAPPING[name].format(level=level)] for level in self.atmos_levels ] atmos[name] = torch.stack([inputs[:, :, index] for index in indices], dim=2) metadata = Metadata( lat=lat_tensor, lon=lon_tensor, time=_normalise_times(times, batch_size), atmos_levels=self.atmos_levels, ) return Batch(surface, fields, atmos, metadata) def batch_to_tensor(self, prediction: Batch) -> torch.Tensor: """Flatten a one-step Aurora prediction back to ``[B, 69, H, W]``.""" surface = [prediction.surf_vars[name][:, -1] for name in AURORA_SURFACE_VARS] atmos = [ prediction.atmos_vars[name][:, -1, level_index] for name in AURORA_ATMOS_VARS for level_index in range(len(self.atmos_levels)) ] return torch.stack(surface + atmos, dim=1) @staticmethod def crop_target(target: torch.Tensor, patch_size: int = 4) -> torch.Tensor: """Apply the same one-row latitude crop used by the official Batch.crop.""" if target.ndim not in {4, 5}: raise ValueError(f"Target must have [B,C,H,W] or [B,T,C,H,W], got {tuple(target.shape)}") height, width = target.shape[-2:] if width % patch_size != 0: raise ValueError("Target longitude size must be divisible by the Aurora patch size") if height % patch_size == 0: return target if height % patch_size == 1: return target[..., :-1, :] raise ValueError("Aurora permits at most one extra latitude row") def forward( self, inputs: torch.Tensor, *, times: datetime | str | Sequence[datetime | str], static_vars: Mapping[str, torch.Tensor | np.ndarray] | None = None, lat: torch.Tensor | np.ndarray | None = None, lon: torch.Tensor | np.ndarray | None = None, return_batch: bool = False, ) -> torch.Tensor | Batch: batch = self.tensor_to_batch(inputs, times=times, static_vars=static_vars, lat=lat, lon=lon) prediction = self.network(batch) return prediction if return_batch else self.batch_to_tensor(prediction) def rollout( self, inputs: torch.Tensor, *, times: datetime | str | Sequence[datetime | str], steps: int, static_vars: Mapping[str, torch.Tensor | np.ndarray] | None = None, lat: torch.Tensor | np.ndarray | None = None, lon: torch.Tensor | np.ndarray | None = None, ) -> torch.Tensor: """Run official autoregressive rollout and return ``[B, steps, 69, H, W]``.""" if steps < 1: raise ValueError(f"Rollout steps must be positive, got {steps}") batch = self.tensor_to_batch(inputs, times=times, static_vars=static_vars, lat=lat, lon=lon) predictions = [ self.batch_to_tensor(prediction) for prediction in aurora_rollout(self.network, batch, steps=steps) ] return torch.stack(predictions, dim=1) def build_aurora_model( config: Mapping[str, object], *, project_root: str | Path | None = None, load_pretrained: bool = False, ) -> AuroraOneScience: """Build the adapter from the project's YAML configuration mapping.""" model_cfg = config["model"] data_cfg = config["data"] root = Path(project_root or ".").expanduser().resolve() static_value = data_cfg.get("static_file") static_path = None if static_value: static_path = Path(str(static_value)).expanduser() if not static_path.is_absolute(): static_path = root / static_path checkpoint_value = model_cfg.get("checkpoint", {}).get("local_small") checkpoint_path = None if checkpoint_value: checkpoint_path = Path(str(checkpoint_value)).expanduser() if not checkpoint_path.is_absolute(): checkpoint_path = root / checkpoint_path return AuroraOneScience( variant=str(model_cfg.get("variant", "small")), static_path=static_path, channel_order=tuple(data_cfg.get("channel_order", DEFAULT_ERA5_CHANNEL_ORDER)), atmos_levels=tuple(model_cfg.get("atmos_levels", AURORA_LEVELS)), checkpoint_path=checkpoint_path, load_checkpoint=load_pretrained, )