From d180fccad93fcac0f58303736cad0b28cda6e791 Mon Sep 17 00:00:00 2001 From: Negin Sobhani Date: Thu, 16 Jul 2026 20:24:49 -0600 Subject: [PATCH] Add HealDAv2 video data assimilation model wrapper Wraps physicsnemo VideoHealDA (8-frame 48h video-DA, dit-5B). Implements the pressure-level conventional obs normalization from healda main (conv-plevel expanded channel vocabulary), microwave sounder prep, IR PCA footprint encoding (IASI/CrIS-FSR/AIRS), and the 50-dim unified metadata featurization as a pure-torch port. load_default_package is stubbed pending checkpoint conversion and HF package publication. --- earth2studio/models/da/__init__.py | 1 + earth2studio/models/da/healda_v2.py | 1097 +++++++++++++++++++++ earth2studio/models/da/healda_v2_utils.py | 724 ++++++++++++++ test/conftest.py | 1 + test/models/da/test_da_healda_v2.py | 608 ++++++++++++ 5 files changed, 2431 insertions(+) create mode 100644 earth2studio/models/da/healda_v2.py create mode 100644 earth2studio/models/da/healda_v2_utils.py create mode 100644 test/models/da/test_da_healda_v2.py diff --git a/earth2studio/models/da/__init__.py b/earth2studio/models/da/__init__.py index d40fee02f..1fd9e72ac 100644 --- a/earth2studio/models/da/__init__.py +++ b/earth2studio/models/da/__init__.py @@ -15,5 +15,6 @@ # limitations under the License. from .healda import HealDA +from .healda_v2 import HealDAv2 from .interp import InterpEquirectangular from .sda_stormcast import StormCastSDA diff --git a/earth2studio/models/da/healda_v2.py b/earth2studio/models/da/healda_v2.py new file mode 100644 index 000000000..657fd4615 --- /dev/null +++ b/earth2studio/models/da/healda_v2.py @@ -0,0 +1,1097 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime as dt +from collections import OrderedDict +from collections.abc import Generator +from typing import Any, cast + +import numpy as np +import pandas as pd +import torch +import xarray as xr +from loguru import logger + +from earth2studio.models.auto import AutoModelMixin, Package +from earth2studio.models.da.base import AssimilationModel +from earth2studio.models.da.healda import E2S_CHANNELS, ERA5_CHANNELS +from earth2studio.models.da.healda_v2_utils import ( + CONV_CHANNELS, + CONV_GPS_CHANNELS, + CONV_GPS_LEVEL2_CHANNELS, + CONV_UV_CHANNELS, + CONV_UV_IN_SITU_TYPES, + CONV_VAR_CHANNEL, + IR_BT_MAX_VALID, + IR_BT_MIN_VALID, + PLATFORM_NAME_TO_ID, + SENSOR_OFFSET, + PCACodec, + QCLimits, + build_conv_plevel_channel_stats, + build_raw_to_local_lut, + compute_unified_metadata, + conv_plevel_local_channel_lut, + get_global_channel_id, + nearest_pressure_level_index, +) +from earth2studio.utils.imports import ( + OptionalDependencyFailure, + check_optional_dependencies, +) +from earth2studio.utils.type import CoordSystem, FrameSchema + +try: + import cupy as cp +except ImportError: + cp = None # type: ignore[assignment] + +try: + import cudf +except ImportError: + cudf = None # type: ignore[assignment, misc] + +try: + import earth2grid + from physicsnemo.experimental.models.healda import ( + VideoHealDA as _VideoHealDAModel, + ) + from physicsnemo.experimental.models.healda import ( + prepare_obs_context, + ) +except ImportError: + OptionalDependencyFailure("da-healda") + earth2grid = None + _VideoHealDAModel = None + prepare_obs_context = None + +# 8-frame video window at 6-hour spacing; each frame ingests its own +# (valid - 3h, valid + 3h] observation context, matching the training loader's +# end-aligned DA windows. +N_WINDOW = 8 +WINDOW_STEP_HOURS = 6 +FRAME_CONTEXT_HOURS = 3 + +# Microwave sounders consumed directly and infrared sounders consumed through +# PCA compression. UFSObsSat variable names for the raw IR sensors differ from +# the internal sensor names. +MW_SENSORS = ("atms", "mhs", "amsua", "amsub") +IR_PCA_SENSORS = ("iasi-pca", "cris-fsr-pca", "airs-pca") +IR_PCA_UFS_VARIABLE: dict[str, str] = { + "iasi-pca": "iasi", + "cris-fsr-pca": "crisfsr", + "airs-pca": "airs", +} +SAT_UFS_VARIABLES = (*MW_SENSORS, *IR_PCA_UFS_VARIABLE.values()) + +SENSOR_PLATFORMS: dict[str, list[str]] = { + "atms": ["npp", "n20"], + "mhs": ["metop-a", "metop-b", "metop-c", "n18", "n19"], + "amsua": ["metop-a", "metop-b", "metop-c", "n15", "n16", "n17", "n18", "n19"], + "amsub": ["n15", "n16", "n17"], + "iasi-pca": ["metop-a", "metop-b", "metop-c"], + "cris-fsr-pca": ["npp", "n20"], + "airs-pca": ["aqua"], +} + +# Footprint identity columns for grouping long-format IR rows into soundings +_FP_COLS = ["satellite", "time", "lat", "lon", "scan_angle", "satellite_za", "solza"] + +# Unified per-observation schema produced by the per-sensor prep methods +_OBS_FRAME_DTYPES = { + "lat": "float32", + "lon": "float32", + "obs_time_ns": "datetime64[ns]", + "observation": "float64", + "global_channel": "int64", + "global_platform": "int64", + "obs_type": "int64", + "height": "float32", + "pressure": "float32", + "scan_angle": "float32", + "sat_zenith_angle": "float32", + "sol_zenith_angle": "float32", +} + + +def _obs_frame(data: dict) -> pd.DataFrame: + return pd.DataFrame(data).astype(_OBS_FRAME_DTYPES, copy=False) + + +@check_optional_dependencies() +class HealDAv2(torch.nn.Module, AutoModelMixin): + """HealDA-v2 video data assimilation model for global weather analysis from + sparse observations on a HEALPix grid. + + HealDA-v2 is a stateless assimilation model that jointly produces an + 8-frame, 48-hour window of global weather analyses (6-hour spacing) from + conventional and satellite observations. The final frame, valid at the + request time, is the present-time analysis. Each frame ingests its own + ±3-hour observation context, so the full observation window spans + (request_time - 45h, request_time + 3h]. + + Compared to v1, the model consumes infrared hyperspectral sounders (IASI, + CrIS-FSR, AIRS) compressed to 32 PCA latents per footprint, and normalizes + vertically structured conventional observations with pressure-level + dependent statistics using an expanded per-level channel vocabulary + ("conv-plevel"). + + The model accepts observation DataFrames (from + :py:class:`earth2studio.data.UFSObsConv` and + :py:class:`earth2studio.data.UFSObsSat`) and produces a global analysis on + the HEALPix level-6 padded XY grid with ERA5-compatible variables. + + Parameters + ---------- + model : torch.nn.Module + The underlying VideoHealDA neural network + condition : torch.Tensor + Static conditioning fields (orography, land fraction) on the HEALPix + grid of size [1, n_static, 1, npix] or [1, n_static, time_length, npix] + era5_mean : torch.Tensor + ERA5 per-channel mean for output denormalization [1, out_variables, 1, 1] + era5_std : torch.Tensor + ERA5 per-channel std for output denormalization [1, out_variables, 1, 1] + channel_stats : pd.DataFrame + Per-global-channel statistics with columns ``Global_Channel_ID``, + ``mean``, ``stddev``, ``min_valid``, ``max_valid`` covering the + microwave, PCA and conv-plevel channels + raw_to_local : dict[str, np.ndarray] + Per-microwave-sensor lookup tables mapping raw (GSI) channel ids to + local channels (see + :func:`earth2studio.models.da.healda_v2_utils.build_raw_to_local_lut`) + codecs : dict[str, PCACodec] + PCA codecs keyed by PCA sensor name (``iasi-pca``, ``cris-fsr-pca``, + ``airs-pca``) + lat_lon : bool, optional + If True the model output is regridded from the native HEALPix grid to a + regular equiangular lat-lon grid using ``earth2grid``. If False the raw + HEALPix output is returned with an ``npix`` dimension, by default False + output_resolution : tuple[int, int], optional + ``(nlat, nlon)`` size of the output lat-lon grid. Only used when + ``lat_lon=True``, by default ``(181, 360)`` (1° resolution) + + Badges + ------ + region:global class:da product:wind product:temp product:atmos product:sat + product:insitu year:2026 gpu:40gb + """ + + def __init__( + self, + model: torch.nn.Module, + condition: torch.Tensor, + era5_mean: torch.Tensor, + era5_std: torch.Tensor, + channel_stats: pd.DataFrame, + raw_to_local: dict[str, np.ndarray], + codecs: dict[str, "PCACodec"], + lat_lon: bool = False, + output_resolution: tuple[int, int] = (181, 360), + ) -> None: + super().__init__() + self._model = model + self.register_buffer("condition", condition) + self.register_buffer("_era5_mean", era5_mean) + self.register_buffer("_era5_std", era5_std) + self.register_buffer("device_buffer", torch.empty(0)) + self._channel_stats = channel_stats[ + ["Global_Channel_ID", "mean", "stddev", "min_valid", "max_valid"] + ] + self._raw_to_local = raw_to_local + self._codecs = codecs + self._lat_lon = lat_lon + self._plevel_lut = conv_plevel_local_channel_lut() + + # Model geometry: observations are assigned to pixels on the backbone + # (level_model) grid, output is on the finer level_in grid. + self._time_length = int(model.time_length) + self._npix_out = int(model.npix) + self._npix_model = 12 * 4 ** int(model.level_model) + self._grid = earth2grid.healpix.Grid( + int(model.level_model), pixel_order=earth2grid.healpix.HEALPIX_PAD_XY + ) + + # Setup lat-lon regridder when requested + if self._lat_lon: + nlat, nlon = output_resolution + self._output_lat = np.linspace(90, -90, nlat) + self._output_lon = np.linspace(0, 360, nlon, endpoint=False) + out_grid = earth2grid.healpix.Grid( + int(np.round(np.emath.logn(4, self._npix_out / 12))), + pixel_order=earth2grid.healpix.HEALPIX_PAD_XY, + ) + ll_grid = earth2grid.latlon.equiangular_lat_lon_grid(nlat, nlon) + self._regridder = earth2grid.get_regridder(out_grid, ll_grid) + else: + self._output_lat = None + self._output_lon = None + self._regridder = None + + @property + def device(self) -> torch.device: + return self.device_buffer.device + + def init_coords(self) -> None: + """Initialization coords (not required)""" + return None + + def input_coords(self) -> tuple[FrameSchema, FrameSchema]: + """Input coordinate system specifying required DataFrame fields. + + Returns two FrameSchemas: one for conventional observations and one for + satellite observations. When calling the model, either may be ``None`` + but not both. + + Returns + ------- + tuple[FrameSchema, FrameSchema] + (conventional_schema, satellite_schema) describing the expected + columns for each observation DataFrame + """ + conv_schema = FrameSchema( + { + "time": np.empty(0, dtype="datetime64[ns]"), + "lat": np.empty(0, dtype=np.float32), + "lon": np.empty(0, dtype=np.float32), + "observation": np.empty(0, dtype=np.float32), + "variable": np.array(list(CONV_VAR_CHANNEL.keys()), dtype=str), + "type": np.empty(0, dtype=np.uint16), + "elev": np.empty(0, dtype=np.float32), + "pres": np.empty(0, dtype=np.float32), + } + ) + sat_schema = FrameSchema( + { + "time": np.empty(0, dtype="datetime64[ns]"), + "lat": np.empty(0, dtype=np.float32), + "lon": np.empty(0, dtype=np.float32), + "observation": np.empty(0, dtype=np.float32), + "variable": np.array(list(SAT_UFS_VARIABLES), dtype=str), + "sensor_index": np.empty(0, dtype=np.uint16), + "satellite": np.empty(0, dtype=str), + "scan_angle": np.empty(0, dtype=np.float32), + "satellite_za": np.empty(0, dtype=np.float32), + "solza": np.empty(0, dtype=np.float32), + } + ) + return conv_schema, sat_schema + + def output_coords( + self, + input_coords: tuple[CoordSystem, CoordSystem], + request_time: np.ndarray | None = None, + **kwargs: Any, + ) -> tuple[CoordSystem]: + """Output coordinate system for the HealDA-v2 analysis window. + + Parameters + ---------- + input_coords : tuple[CoordSystem] + Input coordinate system + request_time : np.ndarray | None, optional + Analysis valid time, by default None + + Returns + ------- + tuple[CoordSystem] + Coordinate system with time, lead_time, variable, and lat/lon or + npix dimensions. The ``lead_time`` axis holds the 8 frame offsets + relative to the analysis time (-42h ... 0h); the analysis frame is + ``lead_time == 0``. + """ + if request_time is None: + request_time = np.array([np.datetime64("NaT")], dtype="datetime64[ns]") + + lead_time = np.array( + [ + np.timedelta64(-(self._time_length - 1 - g) * WINDOW_STEP_HOURS, "h") + for g in range(self._time_length) + ], + dtype="timedelta64[ns]", + ) + + if self._lat_lon: + return ( + CoordSystem( + OrderedDict( + { + "time": request_time, + "lead_time": lead_time, + "variable": np.array(E2S_CHANNELS, dtype=str), + "lat": self._output_lat, + "lon": self._output_lon, + } + ) + ), + ) + + return ( + CoordSystem( + OrderedDict( + { + "time": request_time, + "lead_time": lead_time, + "variable": np.array(E2S_CHANNELS, dtype=str), + "npix": np.arange(self._npix_out), + } + ) + ), + ) + + @classmethod + def load_default_package(cls) -> Package: + """Load the default HealDA-v2 model package. + + Raises + ------ + NotImplementedError + The HealDA-v2 model package has not been published yet + """ + raise NotImplementedError( + "The HealDA-v2 model package has not been published yet. Provide a " + "local or remote Package with the expected layout to load_model." + ) + + @classmethod + @check_optional_dependencies() + def load_model( + cls, + package: Package, + lat_lon: bool = False, + output_resolution: tuple[int, int] = (181, 360), + ) -> AssimilationModel: + """Load HealDA-v2 model from package. + + The package is expected to contain:: + + healda_v2.mdlus # VideoHealDA checkpoint + static/condition_hpx6_padxy.npy # [1, 2, 1, npix] conditioning + stats/channel_table.parquet # global channel stats + stats/conv_normalizations_by_level.csv # per-level conv stats + stats/era5_13_levels_stats.csv # output denormalization + stats/normalizations/{sensor}_normalizations.csv # MW raw ids + codecs/{iasi,cris-fsr,airs}-pca.pt # PCA codecs + + Parameters + ---------- + package : Package + Package containing model checkpoint and statistics + lat_lon : bool, optional + If True the output is regridded to a regular lat-lon grid, + by default False + output_resolution : tuple[int, int], optional + ``(nlat, nlon)`` size of the output lat-lon grid. Only used when + ``lat_lon=True``, by default ``(181, 360)`` + + Returns + ------- + AssimilationModel + Loaded HealDA-v2 assimilation model + """ + model = _VideoHealDAModel.from_checkpoint(package.resolve("healda_v2.mdlus")) + model.eval() + + condition = torch.from_numpy( + np.load(package.resolve("static/condition_hpx6_padxy.npy")) + ) + + # Global channel table: microwave + PCA sensor stats keyed by global id + channel_table = pd.read_parquet(package.resolve("stats/channel_table.parquet")) + channel_stats = channel_table[ + ["Global_Channel_ID", "mean", "stddev", "min_valid", "max_valid"] + ] + + # Append the expanded conv-plevel stats built from the per-level CSV, + # with base conv rows providing QC bounds and the surface fallback. + level_stats = pd.read_csv( + package.resolve("stats/conv_normalizations_by_level.csv") + ) + conv_offset = SENSOR_OFFSET["conv"] + base_conv = channel_stats[ + (channel_stats["Global_Channel_ID"] >= conv_offset) + & (channel_stats["Global_Channel_ID"] < conv_offset + len(CONV_CHANNELS)) + ] + plevel_offset = SENSOR_OFFSET["conv-plevel"] + channel_stats = pd.concat( + [ + channel_stats[channel_stats["Global_Channel_ID"] < plevel_offset], + build_conv_plevel_channel_stats(level_stats, base_conv), + ], + ignore_index=True, + ) + + # Microwave raw (GSI) channel id -> local channel lookup tables + raw_to_local: dict[str, np.ndarray] = {} + for sensor in MW_SENSORS: + df = pd.read_csv( + package.resolve(f"stats/normalizations/{sensor}_normalizations.csv") + ) + df = df[df["Platform_ID"] == -1] + raw_to_local[sensor] = build_raw_to_local_lut( + df["Raw_Channel_ID"].to_numpy() + ) + + # PCA codecs for the compressed infrared sounders + codecs = { + sensor: PCACodec.load(package.resolve(f"codecs/{sensor}.pt")) + for sensor in IR_PCA_SENSORS + } + + # ERA5 output normalization stats + stats = pd.read_csv(package.resolve("stats/era5_13_levels_stats.csv")) + level = stats["level"].astype(int) + channel = np.where( + level.eq(-1), stats["variable"], stats["variable"] + level.astype(str) + ) + ordered = stats.assign(channel=channel).set_index("channel").loc[ERA5_CHANNELS] + era5_mean = torch.from_numpy(ordered["mean"].to_numpy(dtype=np.float32)) + era5_std = torch.from_numpy(ordered["std"].to_numpy(dtype=np.float32)) + + return cls( + model=model, + condition=condition, + era5_mean=era5_mean.view(1, -1, 1, 1), + era5_std=era5_std.view(1, -1, 1, 1), + channel_stats=channel_stats, + raw_to_local=raw_to_local, + codecs=codecs, + lat_lon=lat_lon, + output_resolution=output_resolution, + ) + + # ------------------------------------------------------------------ + # Per-sensor preprocessing + # ------------------------------------------------------------------ + + def prep_conv(self, df: pd.DataFrame) -> pd.DataFrame: + """Standardize and QC a conventional DataFrame into the unified schema. + + Applies the HealDA-v2 conventional QC (height/pressure physical bounds, + in-situ-only UV, GPS bending-angle only) and maps each observation to + its pressure-level expanded ``conv-plevel`` global channel. + + Parameters + ---------- + df : pd.DataFrame + Raw conventional observation DataFrame from UFSObsConv + + Returns + ------- + pd.DataFrame + Standardized DataFrame with unified column schema + """ + unknown_vars = set(df["variable"].unique()) - set(CONV_VAR_CHANNEL.keys()) + if unknown_vars: + raise ValueError(f"Unknown conventional variable(s): {unknown_vars}") + + base_local = df["variable"].map(CONV_VAR_CHANNEL).to_numpy().astype(np.int64) + observation = df["observation"].to_numpy().astype(np.float64) + # Earth2Studio provides pressure-like values in Pa; the model was + # trained on hPa. + is_pres_obs = df["variable"].to_numpy() == "pres" + observation[is_pres_obs] /= 100.0 + height = df["elev"].to_numpy().astype(np.float32) + pressure = df["pres"].to_numpy().astype(np.float32) / np.float32(100.0) + obs_type = df["type"].fillna(0).to_numpy().astype(np.int64) + + # Physical bounds QC. NaN height/pressure rows are dropped, which also + # covers observation sources lacking those coordinates. + is_gps = np.isin(base_local, CONV_GPS_CHANNELS) + min_pressure = np.where( + is_gps, QCLimits.PRESSURE_MIN_GPS, QCLimits.PRESSURE_MIN_DEFAULT + ) + keep = ( + np.isfinite(height) + & (height >= QCLimits.HEIGHT_MIN) + & (height <= QCLimits.HEIGHT_MAX) + & np.isfinite(pressure) + & (pressure >= min_pressure) + & (pressure <= QCLimits.PRESSURE_MAX) + ) + # Satellite-derived winds are excluded; only in-situ UV types are kept + is_uv = np.isin(base_local, CONV_UV_CHANNELS) + keep &= ~is_uv | np.isin(obs_type, CONV_UV_IN_SITU_TYPES) + # GPS level-2 retrievals (gps_t / gps_q) are excluded + keep &= ~np.isin(base_local, CONV_GPS_LEVEL2_CHANNELS) + + base_local = base_local[keep] + pressure = pressure[keep] + + # Pressure-level channel expansion: (base channel, level bin) -> + # conv-plevel local channel -> global channel id + level_idx = nearest_pressure_level_index(pressure) + expanded_local = self._plevel_lut[base_local, level_idx].astype(np.int64) + global_channel = SENSOR_OFFSET["conv-plevel"] + expanded_local + + platform = np.array( + [ + PLATFORM_NAME_TO_ID[CONV_CHANNELS[channel].platform] + for channel in base_local + ], + dtype=np.int64, + ) + df = df.loc[keep] + return _obs_frame( + { + "lat": df["lat"].to_numpy().astype(np.float32), + "lon": df["lon"].to_numpy().astype(np.float32), + "obs_time_ns": df["time"].to_numpy().astype("datetime64[ns]"), + "observation": observation[keep], + "global_channel": global_channel, + "global_platform": platform, + "obs_type": obs_type[keep], + "height": height[keep], + "pressure": pressure, + "scan_angle": np.float32(np.nan), + "sat_zenith_angle": np.float32(np.nan), + "sol_zenith_angle": np.float32(np.nan), + } + ) + + def prep_mw(self, df: pd.DataFrame, sensor: str) -> pd.DataFrame: + """Standardize a microwave sounder DataFrame into the unified schema. + + Parameters + ---------- + df : pd.DataFrame + Raw satellite observation DataFrame from UFSObsSat for one sensor + sensor : str + Sensor name (atms, mhs, amsua, amsub) + + Returns + ------- + pd.DataFrame + Standardized DataFrame with unified column schema + """ + unknown = set(df["satellite"].unique()) - set(SENSOR_PLATFORMS[sensor]) + if unknown: + raise ValueError(f"Unknown satellite platform(s) for {sensor}: {unknown}") + + global_channel = get_global_channel_id( + sensor, + df["sensor_index"].to_numpy().astype(np.int64), + self._raw_to_local[sensor], + ) + # Raw ids unknown to the LUT map below the sensor offset; drop them + keep = global_channel >= SENSOR_OFFSET[sensor] + df = df.loc[keep] + return _obs_frame( + { + "lat": df["lat"].to_numpy().astype(np.float32), + "lon": df["lon"].to_numpy().astype(np.float32), + "obs_time_ns": df["time"].to_numpy().astype("datetime64[ns]"), + "observation": df["observation"].to_numpy().astype(np.float64), + "global_channel": global_channel[keep], + "global_platform": df["satellite"] + .map(PLATFORM_NAME_TO_ID) + .to_numpy() + .astype(np.int64), + "obs_type": np.int64(0), + "height": np.float32(np.nan), + "pressure": np.float32(np.nan), + "scan_angle": df["scan_angle"].to_numpy().astype(np.float32), + "sat_zenith_angle": df["satellite_za"].to_numpy().astype(np.float32), + "sol_zenith_angle": df["solza"].to_numpy().astype(np.float32), + } + ) + + def prep_ir_pca(self, df: pd.DataFrame, sensor: str) -> pd.DataFrame: + """Group long-format IR rows into footprints and PCA-encode them. + + Each footprint's brightness temperature spectrum is projected to 32 + latent observations. ``sensor_index`` is the GSI channel id; the + codec's ``sensor_chan`` to global-channel mapping is the source of + truth (AIRS has sparse channel ids). + + Parameters + ---------- + df : pd.DataFrame + Raw satellite observation DataFrame from UFSObsSat for one raw IR + sensor + sensor : str + PCA sensor name (iasi-pca, cris-fsr-pca, airs-pca) + + Returns + ------- + pd.DataFrame + Standardized DataFrame with one row per (footprint, latent) + """ + unknown = set(df["satellite"].unique()) - set(SENSOR_PLATFORMS[sensor]) + if unknown: + raise ValueError(f"Unknown satellite platform(s) for {sensor}: {unknown}") + + codec = self._codecs[sensor] + latent_width = codec.n_latent + df = df.reset_index(drop=True) + + footprint_id = df.groupby(_FP_COLS, sort=False).ngroup().to_numpy() + sensor_channel = df["sensor_index"].to_numpy().astype(np.int64) + channel_lookup = { + int(sensor_chan): int(gcid) + for sensor_chan, gcid in zip( + codec.sensor_chan.cpu().numpy(), + codec.channel_gcids.cpu().numpy(), + ) + } + global_channel = np.array( + [channel_lookup.get(int(ch), -1) for ch in sensor_channel], + dtype=np.int64, + ) + bt, first_row_idx = codec.preprocess( + global_channel, + df["observation"].to_numpy(), + footprint_id, + IR_BT_MIN_VALID, + IR_BT_MAX_VALID, + ) + latent_observation = codec.encode(torch.from_numpy(bt)).numpy() + + n_footprints = len(first_row_idx) + footprint_rows = df.iloc[first_row_idx] + + def repeat_per_latent(values: Any) -> np.ndarray: + return np.repeat(np.asarray(values), latent_width) + + return _obs_frame( + { + "lat": repeat_per_latent(footprint_rows["lat"]).astype(np.float32), + "lon": repeat_per_latent(footprint_rows["lon"]).astype(np.float32), + "obs_time_ns": repeat_per_latent( + footprint_rows["time"].to_numpy().astype("datetime64[ns]") + ), + "observation": latent_observation.ravel().astype(np.float64), + "global_channel": np.tile( + np.arange(latent_width, dtype=np.int64) + SENSOR_OFFSET[sensor], + n_footprints, + ), + "global_platform": repeat_per_latent( + footprint_rows["satellite"] + .map(PLATFORM_NAME_TO_ID) + .to_numpy() + .astype(np.int64) + ), + "obs_type": np.int64(0), + "height": np.float32(np.nan), + "pressure": np.float32(np.nan), + "scan_angle": repeat_per_latent(footprint_rows["scan_angle"]).astype( + np.float32 + ), + "sat_zenith_angle": repeat_per_latent( + footprint_rows["satellite_za"] + ).astype(np.float32), + "sol_zenith_angle": repeat_per_latent(footprint_rows["solza"]).astype( + np.float32 + ), + } + ) + + def _frame_parts( + self, conv_df: pd.DataFrame | None, sat_df: pd.DataFrame | None + ) -> list[pd.DataFrame]: + """Per-sensor preprocessed DataFrames for a single 6-hour frame.""" + parts: list[pd.DataFrame] = [] + if conv_df is not None and len(conv_df) > 0: + parts.append(self.prep_conv(conv_df)) + if sat_df is not None and len(sat_df) > 0: + for sensor in MW_SENSORS: + rows = sat_df[sat_df["variable"] == sensor] + if len(rows) > 0: + parts.append(self.prep_mw(rows, sensor)) + for sensor in IR_PCA_SENSORS: + if sensor not in self._codecs: + continue + rows = sat_df[sat_df["variable"] == IR_PCA_UFS_VARIABLE[sensor]] + if len(rows) > 0: + parts.append(self.prep_ir_pca(rows, sensor)) + return [part for part in parts if len(part) > 0] + + # ------------------------------------------------------------------ + # Window assembly and forward + # ------------------------------------------------------------------ + + @staticmethod + def _frame_valid_times( + request_time: np.datetime64, time_length: int + ) -> list[pd.Timestamp]: + """Valid time of each window frame, ending at ``request_time``.""" + analysis_time = pd.Timestamp(request_time) + return [ + analysis_time + - pd.Timedelta(hours=WINDOW_STEP_HOURS * (time_length - 1 - g)) + for g in range(time_length) + ] + + @staticmethod + def _slice_frame( + df: pd.DataFrame | None, valid_time: pd.Timestamp + ) -> pd.DataFrame | None: + """Observations in the frame's end-aligned (valid - 3h, valid + 3h] window.""" + if df is None or len(df) == 0: + return None + times = pd.to_datetime(df["time"]) + lo = valid_time - pd.Timedelta(hours=FRAME_CONTEXT_HOURS) + hi = valid_time + pd.Timedelta(hours=FRAME_CONTEXT_HOURS) + frame_df = df[(times > lo) & (times <= hi)] + return frame_df if len(frame_df) > 0 else None + + def filter_and_normalize( + self, + conv_obs: pd.DataFrame | None, + sat_obs: pd.DataFrame | None, + request_time: np.datetime64, + ) -> pd.DataFrame: + """Preprocess, QC, and normalize observations for the full window. + + Buckets raw observations into the 8 window frames, converts each + frame's rows into the unified per-observation schema, applies + per-channel valid-range QC, and z-score normalizes with the global + channel statistics. + + Parameters + ---------- + conv_obs : pd.DataFrame | None + Raw conventional observation DataFrame (or ``None``) + sat_obs : pd.DataFrame | None + Raw satellite observation DataFrame (or ``None``) + request_time : np.datetime64 + Analysis valid time (final frame) + + Returns + ------- + pd.DataFrame + Normalized unified observation DataFrame with additional ``frame`` + and ``target_sec`` columns; may be empty + """ + frame_times = self._frame_valid_times(request_time, self._time_length) + parts: list[pd.DataFrame] = [] + for frame_idx, valid_time in enumerate(frame_times): + frame_conv = self._slice_frame(conv_obs, valid_time) + frame_sat = self._slice_frame(sat_obs, valid_time) + for part in self._frame_parts(frame_conv, frame_sat): + part = part.copy() + part["frame"] = np.int64(frame_idx) + part["target_sec"] = np.int64(self._datetime64_to_epoch_sec(valid_time)) + parts.append(part) + + if not parts: + return pd.DataFrame(columns=[*_OBS_FRAME_DTYPES, "frame", "target_sec"]) + + obs = pd.concat(parts, ignore_index=True) + obs = obs.merge( + self._channel_stats, + left_on="global_channel", + right_on="Global_Channel_ID", + how="left", + ) + observation = obs["observation"].to_numpy() + valid = ( + np.isfinite(observation) + & (observation >= obs["min_valid"].to_numpy()) + & (observation <= obs["max_valid"].to_numpy()) + ) + obs = obs[valid].reset_index(drop=True) + obs["observation"] = (obs["observation"] - obs["mean"]) / obs["stddev"] + return obs.drop( + columns=["Global_Channel_ID", "mean", "stddev", "min_valid", "max_valid"] + ) + + @staticmethod + def _datetime64_to_epoch_sec(t: np.datetime64 | pd.Timestamp) -> int: + """Convert a time to integer UTC epoch seconds.""" + t_dt = pd.Timestamp(t).to_pydatetime() + return int( + dt.datetime( + t_dt.year, + t_dt.month, + t_dt.day, + t_dt.hour, + t_dt.minute, + t_dt.second, + tzinfo=dt.timezone.utc, + ).timestamp() + ) + + def build_input( + self, obs: pd.DataFrame, request_time: np.datetime64 + ) -> dict[str, Any]: + """Convert the normalized observation DataFrame into model inputs. + + Computes the 50-dim metadata features, assigns each observation to its + (frame, pixel) bucket on the backbone HEALPix grid, and packs the + observations for the ragged pixel cross-attention. + + Parameters + ---------- + obs : pd.DataFrame + Output of :meth:`filter_and_normalize` + request_time : np.datetime64 + Analysis valid time (final frame) + + Returns + ------- + dict[str, Any] + Dictionary with ``condition``, ``second_of_day``, ``day_of_year`` + tensors and the packed ``obs_ctx`` + """ + device = self.device + + def col(name: str, dtype: torch.dtype) -> torch.Tensor: + return torch.as_tensor(obs[name].to_numpy(), dtype=dtype, device=device) + + lon = col("lon", torch.float32) + lat = col("lat", torch.float32) + time_ns = torch.as_tensor( + obs["obs_time_ns"].to_numpy().astype("datetime64[ns]").astype(np.int64), + device=device, + ) + float_metadata = compute_unified_metadata( + col("target_sec", torch.int64), + time=time_ns, + lon=lon, + lat=lat, + height=col("height", torch.float32), + pressure=col("pressure", torch.float32), + scan_angle=col("scan_angle", torch.float32), + sat_zenith_angle=col("sat_zenith_angle", torch.float32), + sol_zenith_angle=col("sol_zenith_angle", torch.float32), + ) + + pix = self._grid.ang2pix(lon, lat).long() + frame_idx = col("frame", torch.long) + flat_idx = (frame_idx * self._npix_model + pix).int() + obs_ctx = prepare_obs_context( + obs=col("observation", torch.float32), + float_metadata=float_metadata, + obs_type=col("obs_type", torch.long), + channel=col("global_channel", torch.long), + platform=col("global_platform", torch.long), + flat_idx=flat_idx, + total_pixels=self._time_length * self._npix_model, + ) + + # Per-frame calendar features, each shaped [1, time_length] + frame_times = self._frame_valid_times(request_time, self._time_length) + second_of_day = [ + float((t - t.normalize()).total_seconds()) for t in frame_times + ] + day_of_year = [ + float((t - t.normalize().replace(month=1, day=1)).total_seconds() / 86400.0) + for t in frame_times + ] + + condition = self.condition + if condition.shape[2] == 1 and self._time_length > 1: + condition = condition.expand(-1, -1, self._time_length, -1).contiguous() + + return { + "condition": condition, + "second_of_day": torch.tensor( + [second_of_day], dtype=torch.float32, device=device + ), + "day_of_year": torch.tensor( + [day_of_year], dtype=torch.float32, device=device + ), + "obs_ctx": obs_ctx, + } + + @torch.inference_mode() + def _forward(self, inputs: dict[str, Any]) -> torch.Tensor: + noise_labels = torch.zeros(1, device=self.device) + autocast = self.device.type == "cuda" + with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=autocast): + prediction = self._model( + inputs["condition"], + noise_labels, + inputs["second_of_day"], + inputs["day_of_year"], + inputs["obs_ctx"], + ) + # Denormalize: prediction is [batch, channels, time, npix] + return self._era5_std * prediction.float() + self._era5_mean + + # ------------------------------------------------------------------ + # Call / generator + # ------------------------------------------------------------------ + + def __call__( + self, + conv_obs: pd.DataFrame | None = None, + sat_obs: pd.DataFrame | None = None, + ) -> xr.DataArray: + """Run HealDA-v2 inference from conventional and/or satellite observations. + + At least one of the two observation DataFrames must be provided. Each + DataFrame must carry a single ``request_time`` entry in its ``.attrs``; + it is the valid time of the final window frame (the analysis). The + DataFrames should cover the full observation window + ``(request_time - 45h, request_time + 3h]``. + + Parameters + ---------- + conv_obs : pd.DataFrame | None, optional + Conventional observation DataFrame from + :py:class:`earth2studio.data.UFSObsConv`, by default None + sat_obs : pd.DataFrame | None, optional + Satellite observation DataFrame from + :py:class:`earth2studio.data.UFSObsSat`, by default None + + Returns + ------- + xr.DataArray + Global analysis window with dimensions + [time, lead_time, variable, npix] (or lat/lon). The analysis frame + is ``lead_time == 0``. Data is on the same device as the model + (cupy array for GPU, numpy for CPU). + + Raises + ------ + ValueError + If both *conv_obs* and *sat_obs* are ``None``, if ``request_time`` + is missing, or if more than one request time is provided + """ + if conv_obs is None and sat_obs is None: + raise ValueError("At least one of conv_obs or sat_obs must be provided.") + + request_time = None + for df in (conv_obs, sat_obs): + if df is not None: + request_time = df.attrs.get("request_time", None) + if request_time is not None: + break + if request_time is None: + raise ValueError( + "Observation DataFrame must have 'request_time' in attrs. " + "This is typically set by earth2studio data sources." + ) + + if isinstance(request_time, np.ndarray): + request_time = request_time.astype("datetime64[ns]") + else: + request_time = np.array( + [np.datetime64(request_time, "ns")], dtype="datetime64[ns]" + ) + if len(request_time) != 1: + raise ValueError( + "HealDAv2 assimilates a single analysis time per call; got " + f"{len(request_time)} request times." + ) + + # Convert cudf to pandas if needed + if cudf is not None: + if isinstance(conv_obs, cudf.DataFrame): + conv_obs = conv_obs.to_pandas() + if isinstance(sat_obs, cudf.DataFrame): + sat_obs = sat_obs.to_pandas() + + analysis_time = cast(np.datetime64, request_time[0]) + obs = self.filter_and_normalize(conv_obs, sat_obs, analysis_time) + + (output_coords,) = self.output_coords( + self.input_coords(), request_time=request_time + ) + if len(obs) == 0: + logger.warning("No observations after filtering, returning empty analysis") + return self._empty_output(output_coords) + + inputs = self.build_input(obs, analysis_time) + prediction = self._forward(inputs) + return self.build_output(prediction, output_coords) + + def create_generator(self) -> Generator[ + xr.DataArray, + tuple[pd.DataFrame | None, pd.DataFrame | None], + None, + ]: + """Creates a generator which accepts collection of input observations + and yields the output global assimilated data. + + Yields + ------ + xr.DataArray + Global analysis window on the HEALPix grid + + Receives + -------- + tuple[pd.DataFrame | None, pd.DataFrame | None] + A ``(conv_obs, sat_obs)`` tuple sent via ``generator.send()``. + Either element may be ``None`` but not both. + """ + inputs = yield None # type: ignore[misc] + try: + while True: + conv_obs, sat_obs = inputs if inputs is not None else (None, None) + da = self.__call__(conv_obs, sat_obs) + inputs = yield da + except GeneratorExit: + logger.debug("HealDAv2 generator clean up complete.") + + # ------------------------------------------------------------------ + # Output assembly + # ------------------------------------------------------------------ + + def build_output( + self, + prediction: torch.Tensor, + output_coords: CoordSystem, + ) -> xr.DataArray: + """Convert model output tensor to xarray DataArray. + + Parameters + ---------- + prediction : torch.Tensor + Model output [batch, variable, time_length, npix] + output_coords : CoordSystem + Output coordinate system + + Returns + ------- + xr.DataArray + Analysis window, either on HEALPix (npix) or lat-lon grid + """ + # [1, variable, time_length, npix] -> [time, lead_time, variable, npix] + out = prediction.permute(0, 2, 1, 3).contiguous() + if self._lat_lon and self._regridder is not None: + out = self._regridder(out.double()) + + if self.device.type == "cuda" and cp is not None: + data = cp.asarray(out) + else: + data = out.cpu().numpy() + + if self._lat_lon: + return xr.DataArray( + data=data, + dims=["time", "lead_time", "variable", "lat", "lon"], + coords=output_coords, + ) + + return xr.DataArray( + data=data, + dims=["time", "lead_time", "variable", "npix"], + coords=output_coords, + ) + + def _empty_output(self, output_coords: CoordSystem) -> xr.DataArray: + """Return an empty (NaN-filled) DataArray.""" + dims = list(output_coords.keys()) + shape = tuple(len(v) for v in output_coords.values()) + data = torch.full(shape, float("nan"), dtype=torch.float32, device=self.device) + if self.device.type == "cuda" and cp is not None: + data_np = cp.asarray(data) + else: + data_np = data.cpu().numpy() + return xr.DataArray(data=data_np, dims=dims, coords=output_coords) diff --git a/earth2studio/models/da/healda_v2_utils.py b/earth2studio/models/da/healda_v2_utils.py new file mode 100644 index 000000000..e15f36ea5 --- /dev/null +++ b/earth2studio/models/da/healda_v2_utils.py @@ -0,0 +1,724 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Self-contained helpers for the HealDA-v2 (video-DA) assimilation model. + +Ports the observation vocabulary, pressure-level conventional-obs +normalization, PCA infrared codec, and the 50-dim unified observation metadata +featurization from the HealDA training repository so the Earth2Studio wrapper +has no dependency on the internal ``healda`` package. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +# --------------------------------------------------------------------------- +# Sensor vocabulary +# +# The global channel id space concatenates each sensor's local channels in this +# fixed order. Offsets must match the channel table the model was trained with. +# --------------------------------------------------------------------------- + +# Sensor name -> number of local channels, in canonical (offset) order. +SENSOR_CHANNELS: dict[str, int] = { + "atms": 22, + "mhs": 5, + "amsua": 15, + "amsub": 5, + "iasi": 175, + "cris-fsr": 100, + "conv": 8, + "iasi-pca": 32, + "cris-fsr-pca": 32, + "airs": 117, + "airs-pca": 32, + "conv-plevel": 92, +} + +SENSOR_OFFSET: dict[str, int] = {} +_offset = 0 +for _name, _nchan in SENSOR_CHANNELS.items(): + SENSOR_OFFSET[_name] = _offset + _offset += _nchan + +SENSOR_NAME_TO_ID: dict[str, int] = { + name: idx for idx, name in enumerate(SENSOR_CHANNELS.keys()) +} + +# Valid brightness-temperature range for the raw infrared sounders whose +# spectra are PCA-compressed (values outside are treated as missing). +IR_BT_MIN_VALID = 150.0 +IR_BT_MAX_VALID = 350.0 + +PLATFORM_NAME_TO_ID: dict[str, int] = { + "aqua": 0, + "aura": 1, + "f10": 2, + "f11": 3, + "f13": 4, + "f14": 5, + "f15": 6, + "g08": 7, + "g10": 8, + "g11": 9, + "g12": 10, + "m08": 11, + "m09": 12, + "m10": 13, + "metop-a": 14, + "metop-b": 15, + "metop-c": 16, + "n11": 17, + "n12": 18, + "n14": 19, + "n15": 20, + "n16": 21, + "n17": 22, + "n18": 23, + "n19": 24, + "n20": 25, + "npp": 26, + "gps": 27, + "ps": 28, + "q": 29, + "t": 30, + "uv": 31, +} + + +# --------------------------------------------------------------------------- +# Conventional observation channels and QC limits +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ConvChannel: + """Conventional sensor channel definition.""" + + name: str + platform: str + min_valid: float + max_valid: float + + +CONV_CHANNELS: list[ConvChannel] = [ + ConvChannel("gps_angle", "gps", 0.0, 0.1), + ConvChannel("gps_t", "gps", 150.0, 350.0), + ConvChannel("gps_q", "gps", 0.0, 1.0), + ConvChannel("ps", "ps", 500.0, 1100.0), + ConvChannel("q", "q", 0.0, 1.0), + ConvChannel("t", "t", 150.0, 350.0), + ConvChannel("u", "uv", -100.0, 100.0), + ConvChannel("v", "uv", -100.0, 100.0), +] + +CONV_CHANNEL_NAMES: list[str] = [c.name for c in CONV_CHANNELS] +CONV_GPS_CHANNELS: list[int] = [ + i for i, c in enumerate(CONV_CHANNELS) if c.platform == "gps" +] +CONV_GPS_LEVEL2_CHANNELS: list[int] = [ + i for i, c in enumerate(CONV_CHANNELS) if c.name in ("gps_t", "gps_q") +] +CONV_UV_CHANNELS: list[int] = [ + i for i, c in enumerate(CONV_CHANNELS) if c.platform == "uv" +] +CONV_UV_IN_SITU_TYPES: list[int] = [ + 220, + 221, + 229, + 230, + 231, + 232, + 233, + 234, + 235, + 280, + 282, +] + +# Earth2Studio conventional variable name -> base conv local channel id +CONV_VAR_CHANNEL: dict[str, int] = { + "gps": CONV_CHANNEL_NAMES.index("gps_angle"), + "gps_t": CONV_CHANNEL_NAMES.index("gps_t"), + "gps_q": CONV_CHANNEL_NAMES.index("gps_q"), + "pres": CONV_CHANNEL_NAMES.index("ps"), + "q": CONV_CHANNEL_NAMES.index("q"), + "t": CONV_CHANNEL_NAMES.index("t"), + "u": CONV_CHANNEL_NAMES.index("u"), + "v": CONV_CHANNEL_NAMES.index("v"), +} + + +class QCLimits: + """Conventional observation QC filtering limits.""" + + # Height limits (meters) + HEIGHT_MIN = 0.0 + HEIGHT_MAX = 60000.0 + # Pressure limits (hPa). The HealDA-v2 training configuration lowers the + # non-GPS floor from the historical 200 hPa to 1 hPa. + PRESSURE_MIN_GPS = 0.5 + PRESSURE_MIN_DEFAULT = 1.0 + PRESSURE_MAX = 1100.0 + + +# --------------------------------------------------------------------------- +# Pressure-level ("conv-plevel") channel expansion +# +# The vertically structured conventional channels (all except surface pressure) +# are expanded into one channel per ERA5 pressure level; normalization stats are +# then per (base channel, level). Surface pressure keeps a single channel with +# its base normalization. +# --------------------------------------------------------------------------- + +PRESSURE_LEVELS_HPA = np.array( + [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50], + dtype=np.int16, +) +CONV_PLEVEL_SURFACE_LOCAL_CHANNEL: int = CONV_CHANNEL_NAMES.index("ps") +CONV_PLEVEL_BASE_LOCAL_CHANNELS = np.array( + [ + local_channel + for local_channel in range(len(CONV_CHANNELS)) + if local_channel != CONV_PLEVEL_SURFACE_LOCAL_CHANNEL + ], + dtype=np.int16, +) +CONV_PLEVEL_N_LEVELS: int = PRESSURE_LEVELS_HPA.size +CONV_PLEVEL_N_VERTICAL_CHANNELS: int = CONV_PLEVEL_BASE_LOCAL_CHANNELS.size +CONV_PLEVEL_SURFACE_EXPANDED_CHANNEL: int = ( + CONV_PLEVEL_N_VERTICAL_CHANNELS * CONV_PLEVEL_N_LEVELS +) +CONV_PLEVEL_N_CHANNELS: int = CONV_PLEVEL_SURFACE_EXPANDED_CHANNEL + 1 + +_PRESSURE_LEVELS_ASC = np.sort(PRESSURE_LEVELS_HPA) +_PRESSURE_LEVEL_EDGES = ( + _PRESSURE_LEVELS_ASC[:-1].astype(np.float32) + + _PRESSURE_LEVELS_ASC[1:].astype(np.float32) +) / 2.0 + + +def nearest_pressure_level_index(pressure: np.ndarray) -> np.ndarray: + """Return indices into :data:`PRESSURE_LEVELS_HPA` for nearest-level binning. + + Values outside the 50-1000 hPa span are clipped to the nearest endpoint by + the fixed midpoint thresholds. + + Parameters + ---------- + pressure : np.ndarray + Observation pressures in hPa + + Returns + ------- + np.ndarray + Index of the nearest pressure level for each observation + """ + pressure = np.asarray(pressure, dtype=np.float32) + idx: np.ndarray = np.zeros(pressure.size, dtype=np.int8) + for edge in _PRESSURE_LEVEL_EDGES: + idx += pressure > edge + return (PRESSURE_LEVELS_HPA.size - 1 - idx).astype(np.int8, copy=False) + + +def conv_plevel_local_channel_lut() -> np.ndarray: + """Map (base conv local channel, pressure-level index) to expanded local id. + + The first axis is the original 8-channel conv vocabulary. The second axis is + the index returned by :func:`nearest_pressure_level_index`. Surface pressure + maps to the same final expanded channel for every pressure bin. + + Returns + ------- + np.ndarray + Lookup table of shape [len(CONV_CHANNELS), CONV_PLEVEL_N_LEVELS] + """ + n_conv_channels = len(CONV_CHANNELS) + base_to_group: np.ndarray = np.full(n_conv_channels, -1, dtype=np.int16) + for group, local_channel in enumerate(CONV_PLEVEL_BASE_LOCAL_CHANNELS): + base_to_group[local_channel] = group + + lut: np.ndarray = np.empty((n_conv_channels, CONV_PLEVEL_N_LEVELS), dtype=np.uint16) + for local_channel in range(n_conv_channels): + if local_channel == CONV_PLEVEL_SURFACE_LOCAL_CHANNEL: + lut[local_channel, :] = CONV_PLEVEL_SURFACE_EXPANDED_CHANNEL + continue + lut[local_channel, :] = base_to_group[ + local_channel + ] * CONV_PLEVEL_N_LEVELS + np.arange(CONV_PLEVEL_N_LEVELS, dtype=np.uint16) + return lut + + +def build_conv_plevel_channel_stats( + level_stats: pd.DataFrame, + base_conv_stats: pd.DataFrame, +) -> pd.DataFrame: + """Build the 92-row ``conv-plevel`` channel stats table. + + For pressure-expanded rows, matching level stats override the base conv + normalization; missing (channel, level) rows fall back to base channel + stats. The surface pressure row always uses the base ``ps`` normalization. + QC ``min_valid``/``max_valid`` are inherited from the base channel. + + Parameters + ---------- + level_stats : pd.DataFrame + Per-level normalization stats with columns ``Global_Channel_ID`` (base + conv global id), ``Level_hPa``, ``obs_mean``, ``obs_std`` + base_conv_stats : pd.DataFrame + Base conv channel stats with columns ``Global_Channel_ID``, ``mean``, + ``stddev`` for the 8 base conv channels + + Returns + ------- + pd.DataFrame + Stats table with columns ``Global_Channel_ID``, ``mean``, ``stddev``, + ``min_valid``, ``max_valid`` for the 92 conv-plevel channels + """ + conv_offset = SENSOR_OFFSET["conv"] + conv_plevel_offset = SENSOR_OFFSET["conv-plevel"] + + base_norms = { + int(row["Global_Channel_ID"]): (float(row["mean"]), float(row["stddev"])) + for _, row in base_conv_stats.iterrows() + } + level_norms = { + (int(row["Global_Channel_ID"]), int(row["Level_hPa"])): ( + float(row["obs_mean"]), + float(row["obs_std"]), + ) + for _, row in level_stats.iterrows() + } + + rows = [] + expanded_local = 0 + for base_local in CONV_PLEVEL_BASE_LOCAL_CHANNELS: + channel = CONV_CHANNELS[int(base_local)] + base_gid = conv_offset + int(base_local) + base_mean, base_stddev = base_norms.get(base_gid, (0.0, 1.0)) + for level_hpa in PRESSURE_LEVELS_HPA: + mean, stddev = level_norms.get( + (base_gid, int(level_hpa)), (base_mean, base_stddev) + ) + rows.append( + { + "Global_Channel_ID": conv_plevel_offset + expanded_local, + "mean": mean, + "stddev": stddev, + "min_valid": channel.min_valid, + "max_valid": channel.max_valid, + } + ) + expanded_local += 1 + + surface = CONV_CHANNELS[CONV_PLEVEL_SURFACE_LOCAL_CHANNEL] + surface_gid = conv_offset + CONV_PLEVEL_SURFACE_LOCAL_CHANNEL + surface_mean, surface_stddev = base_norms.get(surface_gid, (0.0, 1.0)) + rows.append( + { + "Global_Channel_ID": conv_plevel_offset + expanded_local, + "mean": surface_mean, + "stddev": surface_stddev, + "min_valid": surface.min_valid, + "max_valid": surface.max_valid, + } + ) + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Raw channel id -> local channel lookup for microwave sounders +# --------------------------------------------------------------------------- + + +def build_raw_to_local_lut(raw_channel_ids: np.ndarray) -> np.ndarray: + """Build a 1-indexed raw-channel-id to local-channel lookup table. + + Index 0 of the LUT (and any raw id not present) maps to 0, which callers + interpret as local channel -1 after subtracting one. + + Parameters + ---------- + raw_channel_ids : np.ndarray + Sorted unique raw (GSI) channel ids for a sensor + + Returns + ------- + np.ndarray + Lookup table such that ``lut[raw_id] - 1`` is the local channel + """ + raw_ids = np.unique(np.asarray(raw_channel_ids).ravel()) + lut: np.ndarray = np.zeros(int(raw_ids.max()) + 1, dtype=np.int64) + for local_idx, raw in enumerate(raw_ids, start=1): + lut[int(raw)] = local_idx + return lut + + +def get_global_channel_id( + sensor: str, raw_channel_ids: np.ndarray, raw_to_local: np.ndarray +) -> np.ndarray: + """Map per-sensor raw channel ids to unified global channel ids. + + Parameters + ---------- + sensor : str + Sensor name (key of :data:`SENSOR_OFFSET`) + raw_channel_ids : np.ndarray + Raw (GSI) channel ids + raw_to_local : np.ndarray + Lookup table from :func:`build_raw_to_local_lut` + + Returns + ------- + np.ndarray + Global channel ids (int64); raw ids unknown to the LUT map to + ``SENSOR_OFFSET[sensor] - 1`` and should be filtered by callers + """ + raw_channel_ids = np.asarray(raw_channel_ids) + safe_ids = np.minimum(raw_channel_ids, len(raw_to_local) - 1) + local_channels = raw_to_local[safe_ids] - 1 + return (local_channels + SENSOR_OFFSET[sensor]).astype(np.int64) + + +# --------------------------------------------------------------------------- +# PCA codec for compressed infrared sounders (AIRS / IASI / CrIS-FSR) +# --------------------------------------------------------------------------- + +_STD_FLOOR = 1e-6 + + +class PCACodec(torch.nn.Module): + """Deployable PCA codec for compressed IR sensors. + + The codec standardizes brightness temperatures per channel and projects + onto a fixed set of PCA components:: + + encode: z = ((bt - bt_mean) / bt_std) @ components (N, C) -> (N, k) + + Invalid channels (NaN) are zeroed after standardization so they contribute + nothing to the latent projection. A saved codec is a single ``.pt`` state + dict; channel count and latent dimension are recovered from buffer shapes. + + Parameters + ---------- + bt_mean : torch.Tensor + Per-channel brightness temperature mean [n_channels] + bt_std : torch.Tensor + Per-channel brightness temperature std [n_channels] + components : torch.Tensor + PCA projection matrix [n_channels, n_latent] + channel_gcids : torch.Tensor | None, optional + Global channel id of each codec column, by default None + sensor_chan : torch.Tensor | None, optional + Raw (GSI) sensor channel id of each codec column, by default None + """ + + def __init__( + self, + bt_mean: torch.Tensor, + bt_std: torch.Tensor, + components: torch.Tensor, + channel_gcids: torch.Tensor | None = None, + sensor_chan: torch.Tensor | None = None, + ): + super().__init__() + bt_mean = torch.as_tensor(bt_mean, dtype=torch.float32) + bt_std = torch.as_tensor(bt_std, dtype=torch.float32).clamp_min(_STD_FLOOR) + components = torch.as_tensor(components, dtype=torch.float32) + if bt_mean.ndim != 1 or bt_std.shape != bt_mean.shape: + raise ValueError("bt_mean and bt_std must be 1D with matching length") + if components.ndim != 2 or components.shape[0] != bt_mean.shape[0]: + raise ValueError("components must have shape (n_channels, n_latent)") + + n_channels = bt_mean.shape[0] + if channel_gcids is None: + channel_gcids = torch.arange(n_channels, dtype=torch.int32) + if sensor_chan is None: + sensor_chan = torch.arange(n_channels, dtype=torch.int32) + + self.register_buffer("bt_mean", bt_mean) + self.register_buffer("bt_std", bt_std) + self.register_buffer("components", components) + self.register_buffer( + "channel_gcids", torch.as_tensor(channel_gcids, dtype=torch.int32) + ) + self.register_buffer( + "sensor_chan", torch.as_tensor(sensor_chan, dtype=torch.int32) + ) + + @property + def n_channels(self) -> int: + return int(self.bt_mean.shape[0]) + + @property + def n_latent(self) -> int: + return int(self.components.shape[1]) + + def standardize( + self, bt: torch.Tensor, valid: torch.Tensor | None = None + ) -> torch.Tensor: + if valid is None: + valid = ~torch.isnan(bt) + z = (bt - self.bt_mean) / self.bt_std + return z.masked_fill(~valid, 0.0) + + def encode( + self, bt: torch.Tensor, valid: torch.Tensor | None = None + ) -> torch.Tensor: + return self.standardize(bt, valid) @ self.components + + def preprocess( + self, + channel_gcid: np.ndarray, + obs: np.ndarray, + footprint_id: np.ndarray, + bt_min: float = float("-inf"), + bt_max: float = float("inf"), + ) -> tuple[np.ndarray, np.ndarray]: + """Shape long per-(footprint, channel) BT rows into the wide + ``(n_fp, n_channels)`` matrix that :meth:`encode` expects. + + Maps each row's global channel id to a codec column (dropping channels + the codec does not model), groups rows by footprint, and marks missing + or out-of-range BT as NaN (which :meth:`encode` masks to zero). + + Parameters + ---------- + channel_gcid : np.ndarray + Global channel ids matching ``self.channel_gcids`` [n_rows] + obs : np.ndarray + Brightness temperatures [n_rows] + footprint_id : np.ndarray + Group id identifying the sounding per row [n_rows] + bt_min : float, optional + Minimum valid BT; entries below are treated as missing + bt_max : float, optional + Maximum valid BT; entries above are treated as missing + + Returns + ------- + tuple[np.ndarray, np.ndarray] + ``bt`` wide BT matrix [n_fp, n_channels] (NaN for missing) and + ``first_row_idx`` [n_fp] index of the first input row of each + footprint, for gathering footprint-level metadata + """ + gcid = np.asarray(channel_gcid) + lut: np.ndarray = np.full(int(self.channel_gcids.max()) + 1, -1, dtype=np.int64) + lut[self.channel_gcids.numpy()] = np.arange(self.n_channels) + col = lut[np.minimum(gcid, len(lut) - 1)] + keep_idx = np.where(col >= 0)[0] + _, first_in_kept, inv = np.unique( + np.asarray(footprint_id)[keep_idx], return_index=True, return_inverse=True + ) + first_row_idx = keep_idx[first_in_kept] + + bt: np.ndarray = np.full( + (len(first_row_idx), self.n_channels), np.nan, dtype=np.float32 + ) + bt[inv, col[keep_idx]] = np.asarray(obs)[keep_idx].astype(np.float32) + bt[(bt < bt_min) | (bt > bt_max)] = np.nan + return bt, first_row_idx + + @classmethod + def load(cls, path: str | Path, map_location: str = "cpu") -> PCACodec: + """Load a codec from a single ``.pt`` state-dict file. + + Parameters + ---------- + path : str | Path + Path to the saved codec + map_location : str, optional + Torch map location, by default "cpu" + + Returns + ------- + PCACodec + Loaded codec in eval mode + """ + state = torch.load(path, map_location=map_location, weights_only=True) + codec = cls( + bt_mean=state["bt_mean"], + bt_std=state["bt_std"], + components=state["components"], + channel_gcids=state.get("channel_gcids"), + sensor_chan=state.get("sensor_chan"), + ) + codec.load_state_dict(state) + codec.eval() + return codec + + +# --------------------------------------------------------------------------- +# 50-dim unified observation metadata (features v2) +# --------------------------------------------------------------------------- + +N_METADATA_FEATURES = 50 +# Conv/sat-split layout: +# [0:10) SHARED: LST fourier(2)[0:4) + dt [dt,dt2][4:6) + dt fourier(1)[6:8) +# + lat [sin,cos][8:10) +# [10:30) SAT-private: scan fourier(3)[10:16) + sat_zen fourier(4)[16:24) +# + sol_zen fourier(3)[24:30) -- zero for conv rows +# [30:50) CONV-private: height fourier(5)[30:40) + pressure fourier(5)[40:50) +# -- zero for sat rows +# A row is conv iff height is not NaN; exactly one private block is non-zero. + + +def fourier_features(x_norm: torch.Tensor, num_freqs: int) -> torch.Tensor: + """Sin/cos Fourier features for frequencies ``1..num_freqs``. + + Parameters + ---------- + x_norm : torch.Tensor + Input angles (radians when pre-multiplied by 2*pi) [n] + num_freqs : int + Number of frequencies + + Returns + ------- + torch.Tensor + Features [n, 2 * num_freqs] laid out as [sin(1x)..sin(kx), cos(1x)..cos(kx)] + """ + freqs = torch.arange(1, num_freqs + 1, device=x_norm.device, dtype=x_norm.dtype) + x_expanded = x_norm.unsqueeze(-1) * freqs + return torch.cat([torch.sin(x_expanded), torch.cos(x_expanded)], dim=-1) + + +def local_solar_time(lon_deg: torch.Tensor, abs_time_ns: torch.Tensor) -> torch.Tensor: + """Local solar time in hours from longitude and absolute time. + + Parameters + ---------- + lon_deg : torch.Tensor + Longitude in degrees [n] + abs_time_ns : torch.Tensor + Observation times as epoch nanoseconds (int64) [n] + + Returns + ------- + torch.Tensor + Local solar time in [0, 24) hours [n] + """ + sec_of_day = (abs_time_ns // 1_000_000_000) % 86400 + utc_hours = sec_of_day.float() / 3600.0 + return (utc_hours + lon_deg / 15.0) % 24.0 + + +def compute_unified_metadata( + target_time_sec: torch.Tensor, + time: torch.Tensor, + lon: torch.Tensor, + lat: torch.Tensor, + height: torch.Tensor, + pressure: torch.Tensor, + scan_angle: torch.Tensor, + sat_zenith_angle: torch.Tensor, + sol_zenith_angle: torch.Tensor, +) -> torch.Tensor: + """Compute the 50-dim unified observation metadata features (v2). + + Conv/sat specialization: height validity determines which private block + fills feature slots 10-29 (satellite) or 30-49 (conventional). + + Parameters + ---------- + target_time_sec : torch.Tensor + Analysis (frame) time as epoch seconds (int64) [n] + time : torch.Tensor + Observation times as epoch nanoseconds (int64) [n] + lon : torch.Tensor + Longitude in degrees [n] + lat : torch.Tensor + Latitude in degrees [n] + height : torch.Tensor + Height in meters (NaN for satellite obs) [n] + pressure : torch.Tensor + Pressure in hPa (NaN for satellite obs) [n] + scan_angle : torch.Tensor + Scan angle in degrees (NaN for conventional obs) [n] + sat_zenith_angle : torch.Tensor + Satellite zenith angle in degrees (NaN for conventional obs) [n] + sol_zenith_angle : torch.Tensor + Solar zenith angle in degrees (NaN for conventional obs) [n] + + Returns + ------- + torch.Tensor + Metadata features [n, 50] + """ + n_obs = lon.shape[0] + for name, tensor in [ + ("target_time_sec", target_time_sec), + ("time", time), + ("lat", lat), + ("height", height), + ("pressure", pressure), + ("scan_angle", scan_angle), + ("sat_zenith_angle", sat_zenith_angle), + ("sol_zenith_angle", sol_zenith_angle), + ]: + if tensor.shape[0] != n_obs: + raise ValueError(f"{name} has length {tensor.shape[0]}, expected {n_obs}") + + out = torch.zeros( + n_obs, N_METADATA_FEATURES, dtype=torch.float32, device=lon.device + ) + if n_obs == 0: + return out + + is_conv = ~torch.isnan(height) + two_pi = 2 * math.pi + + # Shared: local solar time fourier(2) -> [0:4) + lst = local_solar_time(lon, time) + out[:, 0:4] = fourier_features(lst / 24.0 * two_pi, 2) + + # Shared: relative time polynomial -> [4:6) + target_time_ns = target_time_sec * 1_000_000_000 + dt_days = (time - target_time_ns).float() * 1e-9 / 86400.0 + out[:, 4] = dt_days + out[:, 5] = dt_days**2 + + # Shared: relative time fourier(1) -> [6:8) + out[:, 6:8] = fourier_features(dt_days, 1) + + # Shared: latitude -> [8:10) + lat_rad = torch.deg2rad(lat) + out[:, 8] = torch.sin(lat_rad) + out[:, 9] = torch.cos(lat_rad) + + # Sat-private [10:30): scan fourier(3) + sat_zen fourier(4) + sol_zen fourier(3) + is_sat = ~is_conv + if is_sat.any(): + s = is_sat + out[s, 10:16] = fourier_features(scan_angle[s] / 50.0 * two_pi, 3) + out[s, 16:24] = fourier_features(sat_zenith_angle[s] / 90.0 * two_pi, 4) + out[s, 24:30] = fourier_features(sol_zenith_angle[s] / 180.0 * two_pi, 3) + + # Conv-private [30:50): height fourier(5) + pressure fourier(5) + if is_conv.any(): + c = is_conv + h_norm = torch.clamp(height[c] / 60000.0, 0.0, 1.0) + out[c, 30:40] = fourier_features(h_norm * two_pi, 5) + p_norm = torch.clamp(pressure[c] / 1100.0, 0.0, 1.0) + out[c, 40:50] = fourier_features(p_norm * two_pi, 5) + + return out diff --git a/test/conftest.py b/test/conftest.py index 4c672321e..d4c3d9c56 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -109,6 +109,7 @@ def _package_available(package: str) -> bool: ], # Model da tests "test/models/da/test_da_healda.py": ["da-healda"], + "test/models/da/test_da_healda_v2.py": ["da-healda"], "test/models/da/test_da_interp.py": ["da-interp"], "test/models/da/test_da_sda_stormcast.py": ["da-stormcast"], # Serve tests diff --git a/test/models/da/test_da_healda_v2.py b/test/models/da/test_da_healda_v2.py new file mode 100644 index 000000000..a183b900c --- /dev/null +++ b/test/models/da/test_da_healda_v2.py @@ -0,0 +1,608 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest +import torch +import xarray as xr + +from earth2studio.models.da.healda import E2S_CHANNELS +from earth2studio.models.da.healda_v2 import ( + N_WINDOW, + HealDAv2, +) +from earth2studio.models.da.healda_v2_utils import ( + CONV_PLEVEL_SURFACE_EXPANDED_CHANNEL, + PLATFORM_NAME_TO_ID, + PRESSURE_LEVELS_HPA, + SENSOR_OFFSET, + PCACodec, + build_conv_plevel_channel_stats, + build_raw_to_local_lut, + compute_unified_metadata, + nearest_pressure_level_index, +) + +# ---------- Constants ---------- + +NVAR = len(E2S_CHANNELS) # 74 +LEVEL_MODEL = 1 +NPIX_MODEL = 12 * 4**LEVEL_MODEL # 48 +NPIX_OUT = 192 +NLAT = 5 +NLON = 10 + +REQUEST_TIME = np.array([np.datetime64("2024-01-01T12:00:00")]) + +# Expanded conv-plevel channel bookkeeping: vertical channel groups follow the +# base conv order with surface pressure removed: +# gps_angle -> 0, gps_t -> 1, gps_q -> 2, q -> 3, t -> 4, u -> 5, v -> 6 +T_GROUP = 4 +U_GROUP = 5 +CONV_PLEVEL_OFFSET = SENSOR_OFFSET["conv-plevel"] + +# Synthetic level stats used for the t channel at 600 hPa +T_600_MEAN = 250.0 +T_600_STD = 10.0 + + +# ---------- Mock neural network / grid ---------- + + +class PhooVideoHealDAModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.in_channels = 2 + self.out_channels = NVAR + self.npix = NPIX_OUT + self.level_model = LEVEL_MODEL + self.time_length = N_WINDOW + self.last_call: dict = {} + + def forward(self, x, t, second_of_day, day_of_year, obs_ctx, class_labels=None): + self.last_call = { + "x": x, + "t": t, + "second_of_day": second_of_day, + "day_of_year": day_of_year, + "obs_ctx": obs_ctx, + } + return torch.zeros( + x.shape[0], self.out_channels, self.time_length, self.npix, device=x.device + ) + + +class MockGrid: + def ang2pix(self, lon, lat): + return torch.zeros(lon.shape[0], dtype=torch.long, device=lon.device) + + +class MockRegridder: + def __init__(self, nlat, nlon): + self.nlat = nlat + self.nlon = nlon + + def __call__(self, x): + return torch.randn( + *x.shape[:-1], self.nlat, self.nlon, dtype=x.dtype, device=x.device + ) + + +# ---------- Fixtures / builders ---------- + + +def _build_channel_stats() -> pd.DataFrame: + """Channel stats for atms, the PCA sensors, and conv-plevel.""" + conv_offset = SENSOR_OFFSET["conv"] + base_conv = pd.DataFrame( + { + "Global_Channel_ID": np.arange(8) + conv_offset, + "mean": [0.05, 250.0, 0.005, 950.0, 0.005, 260.0, 0.0, 0.0], + "stddev": [0.02, 20.0, 0.002, 40.0, 0.002, 20.0, 5.0, 5.0], + } + ) + # min/max valid for the base conv channels (order matches CONV_CHANNELS) + base_conv["min_valid"] = [0.0, 150.0, 0.0, 500.0, 0.0, 150.0, -100.0, -100.0] + base_conv["max_valid"] = [0.1, 350.0, 1.0, 1100.0, 1.0, 350.0, 100.0, 100.0] + + level_stats = pd.DataFrame( + { + "Global_Channel_ID": [conv_offset + 5], # base conv channel "t" + "Level_hPa": [600], + "obs_mean": [T_600_MEAN], + "obs_std": [T_600_STD], + } + ) + plevel = build_conv_plevel_channel_stats(level_stats, base_conv) + + atms = pd.DataFrame( + { + "Global_Channel_ID": np.arange(22) + SENSOR_OFFSET["atms"], + "mean": 250.0, + "stddev": 50.0, + "min_valid": 0.0, + "max_valid": 400.0, + } + ) + pca_parts = [ + pd.DataFrame( + { + "Global_Channel_ID": np.arange(32) + SENSOR_OFFSET[sensor], + "mean": 0.0, + "stddev": 1.0, + "min_valid": -np.inf, + "max_valid": np.inf, + } + ) + for sensor in ("iasi-pca", "cris-fsr-pca", "airs-pca") + ] + return pd.concat([atms, *pca_parts, plevel], ignore_index=True) + + +def _build_codec() -> PCACodec: + """Tiny synthetic AIRS codec: 3 channels -> 2 latents.""" + airs_offset = SENSOR_OFFSET["airs"] + return PCACodec( + bt_mean=torch.tensor([250.0, 260.0, 270.0]), + bt_std=torch.tensor([10.0, 10.0, 10.0]), + components=torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]), + channel_gcids=torch.tensor([airs_offset, airs_offset + 1, airs_offset + 2]), + sensor_chan=torch.tensor([1, 5, 9]), # sparse GSI channel ids + ) + + +def _build_model(device="cpu", lat_lon=False): + with patch("earth2studio.models.da.healda_v2.earth2grid") as mock_e2g: + mock_e2g.healpix.Grid.return_value = MockGrid() + mock_e2g.healpix.HEALPIX_PAD_XY = 0 + if lat_lon: + mock_e2g.get_regridder.return_value = MockRegridder(NLAT, NLON) + model = HealDAv2( + model=PhooVideoHealDAModel(), + condition=torch.zeros(1, 2, 1, NPIX_OUT), + era5_mean=torch.zeros(1, NVAR, 1, 1), + era5_std=torch.ones(1, NVAR, 1, 1), + channel_stats=_build_channel_stats(), + raw_to_local={"atms": build_raw_to_local_lut(np.arange(1, 23))}, + codecs={"airs-pca": _build_codec()}, + lat_lon=lat_lon, + output_resolution=(NLAT, NLON), + ) + model._grid = MockGrid() + return model.to(device) + + +def _build_raw_conv_df( + n_obs=10, + request_time=None, + variable="t", + observation=260.0, + obs_type=120, + elev=100.0, + pres_pa=50000.0, +): + if request_time is None: + request_time = REQUEST_TIME + t = request_time[0].astype("datetime64[ns]") + df = pd.DataFrame( + { + "time": np.full(n_obs, t), + "lat": np.random.uniform(-90, 90, n_obs).astype(np.float32), + "lon": np.random.uniform(0, 360, n_obs).astype(np.float32), + "observation": np.full(n_obs, observation, dtype=np.float32), + "variable": variable, + "type": np.full(n_obs, obs_type, dtype=np.uint16), + "elev": np.full(n_obs, elev, dtype=np.float32), + "pres": np.full(n_obs, pres_pa, dtype=np.float32), + } + ) + df.attrs = {"request_time": request_time} + return df + + +def _build_raw_sat_df(n_obs=10, request_time=None, sensor="atms", satellite="n20"): + if request_time is None: + request_time = REQUEST_TIME + t = request_time[0].astype("datetime64[ns]") + df = pd.DataFrame( + { + "time": np.full(n_obs, t), + "lat": np.random.uniform(-90, 90, n_obs).astype(np.float32), + "lon": np.random.uniform(0, 360, n_obs).astype(np.float32), + "observation": np.random.uniform(200, 300, n_obs).astype(np.float32), + "variable": sensor, + "sensor_index": np.ones(n_obs, dtype=np.uint16), + "satellite": satellite, + "scan_angle": np.zeros(n_obs, dtype=np.float32), + "satellite_za": np.full(n_obs, 30.0, dtype=np.float32), + "solza": np.full(n_obs, 45.0, dtype=np.float32), + } + ) + df.attrs = {"request_time": request_time} + return df + + +def _mock_forward(inputs): + return torch.zeros(1, NVAR, N_WINDOW, NPIX_OUT) + + +# ---------- Utils tests ---------- + + +def test_nearest_pressure_level_index(): + pressures = np.array([1000.0, 612.0, 25.0, 1200.0, 962.6, 962.4]) + idx = nearest_pressure_level_index(pressures) + levels = PRESSURE_LEVELS_HPA[idx] + # 612 -> 600; 25 clips to 50; 1200 clips to 1000; 962.5 is the 1000/925 + # midpoint: strictly above -> 1000, at/below -> 925 + assert list(levels) == [1000, 600, 50, 1000, 1000, 925] + + +def test_build_conv_plevel_channel_stats(): + stats = _build_channel_stats() + plevel = stats[stats["Global_Channel_ID"] >= CONV_PLEVEL_OFFSET].reset_index( + drop=True + ) + assert len(plevel) == 92 + + # t @ 600 hPa (group 4, level index 4) uses the level stats + level_600 = int(np.where(PRESSURE_LEVELS_HPA == 600)[0][0]) + t_600 = plevel.iloc[T_GROUP * 13 + level_600] + assert t_600["mean"] == T_600_MEAN + assert t_600["stddev"] == T_600_STD + # t @ 500 hPa has no level stats -> falls back to the base t normalization + level_500 = int(np.where(PRESSURE_LEVELS_HPA == 500)[0][0]) + t_500 = plevel.iloc[T_GROUP * 13 + level_500] + assert t_500["mean"] == 260.0 + assert t_500["stddev"] == 20.0 + # QC bounds inherited from the base channel + assert t_600["min_valid"] == 150.0 + assert t_600["max_valid"] == 350.0 + # Surface pressure row uses the base ps normalization and bounds + ps = plevel.iloc[CONV_PLEVEL_SURFACE_EXPANDED_CHANNEL] + assert ps["mean"] == 950.0 + assert ps["min_valid"] == 500.0 + + +def test_compute_unified_metadata_branches(): + n = 2 + target = torch.full((n,), 1700000000, dtype=torch.int64) + time_ns = target * 1_000_000_000 + lon = torch.tensor([0.0, 90.0]) + lat = torch.tensor([30.0, -45.0]) + # Row 0 is conventional (finite height), row 1 satellite (NaN height) + height = torch.tensor([100.0, float("nan")]) + pressure = torch.tensor([500.0, float("nan")]) + scan = torch.tensor([float("nan"), 10.0]) + sat_za = torch.tensor([float("nan"), 30.0]) + sol_za = torch.tensor([float("nan"), 45.0]) + + meta = compute_unified_metadata( + target, time_ns, lon, lat, height, pressure, scan, sat_za, sol_za + ) + assert meta.shape == (n, 50) + # Shared latitude features + assert torch.allclose(meta[0, 8], torch.sin(torch.deg2rad(lat[0]))) + assert torch.allclose(meta[1, 9], torch.cos(torch.deg2rad(lat[1]))) + # Zero relative time -> dt features are [0, 0, sin 0, cos 0] + assert torch.allclose(meta[:, 4:7], torch.zeros(n, 3)) + assert torch.allclose(meta[:, 7], torch.ones(n)) + # Branch exclusivity: conv rows zero the sat block and vice versa + assert torch.all(meta[0, 10:30] == 0) + assert torch.any(meta[0, 30:50] != 0) + assert torch.any(meta[1, 10:30] != 0) + assert torch.all(meta[1, 30:50] == 0) + + +# ---------- Preprocessing tests ---------- + + +def test_prep_conv_plevel_channel_and_normalization(): + model = _build_model() + # t observation at 612 hPa -> binned to 600 hPa + df = _build_raw_conv_df(n_obs=4, observation=265.0, pres_pa=61200.0) + obs = model.filter_and_normalize(df, None, REQUEST_TIME[0]) + + assert len(obs) == 4 + level_600 = int(np.where(PRESSURE_LEVELS_HPA == 600)[0][0]) + expected_gid = CONV_PLEVEL_OFFSET + T_GROUP * 13 + level_600 + assert (obs["global_channel"] == expected_gid).all() + assert (obs["global_platform"] == PLATFORM_NAME_TO_ID["t"]).all() + # z-scored with the 600 hPa level stats + np.testing.assert_allclose( + obs["observation"].to_numpy(), (265.0 - T_600_MEAN) / T_600_STD, rtol=1e-6 + ) + # Pressure converted Pa -> hPa for the metadata path + np.testing.assert_allclose(obs["pressure"].to_numpy(), 612.0, rtol=1e-6) + # Final frame of the window + assert (obs["frame"] == N_WINDOW - 1).all() + + +def test_prep_conv_qc(): + model = _build_model() + + # NaN pressure -> dropped + df = _build_raw_conv_df(n_obs=3, pres_pa=np.nan) + assert len(model.prep_conv(df)) == 0 + + # Non-GPS pressure floor is 1 hPa: 0.8 hPa dropped, GPS 0.8 hPa kept + df = _build_raw_conv_df(n_obs=3, pres_pa=80.0) # 0.8 hPa + assert len(model.prep_conv(df)) == 0 + df = _build_raw_conv_df(n_obs=3, variable="gps", observation=0.05, pres_pa=80.0) + assert len(model.prep_conv(df)) == 3 + + # GPS level-2 retrievals dropped + df = _build_raw_conv_df(n_obs=3, variable="gps_t", observation=250.0) + assert len(model.prep_conv(df)) == 0 + + # Satellite-derived UV dropped, in-situ UV kept + df = _build_raw_conv_df(n_obs=3, variable="u", observation=10.0, obs_type=240) + assert len(model.prep_conv(df)) == 0 + df = _build_raw_conv_df(n_obs=3, variable="u", observation=10.0, obs_type=220) + assert len(model.prep_conv(df)) == 3 + + # Height out of physical bounds dropped + df = _build_raw_conv_df(n_obs=3, elev=70000.0) + assert len(model.prep_conv(df)) == 0 + + # Unknown variable raises + df = _build_raw_conv_df(n_obs=3, variable="bogus") + with pytest.raises(ValueError, match="Unknown conventional"): + model.prep_conv(df) + + +def test_prep_conv_surface_pressure(): + model = _build_model() + # Station pressure: observation in Pa -> hPa, valid range (500, 1100) + df = _build_raw_conv_df(n_obs=2, variable="pres", observation=90000.0) + obs = model.filter_and_normalize(df, None, REQUEST_TIME[0]) + assert len(obs) == 2 + expected_gid = CONV_PLEVEL_OFFSET + CONV_PLEVEL_SURFACE_EXPANDED_CHANNEL + assert (obs["global_channel"] == expected_gid).all() + np.testing.assert_allclose( + obs["observation"].to_numpy(), (900.0 - 950.0) / 40.0, rtol=1e-6 + ) + + # 300 hPa station pressure violates the (500, 1100) valid range + df = _build_raw_conv_df(n_obs=2, variable="pres", observation=30000.0) + obs = model.filter_and_normalize(df, None, REQUEST_TIME[0]) + assert len(obs) == 0 + + +def test_prep_mw(): + model = _build_model() + df = _build_raw_sat_df(n_obs=5, sensor="atms") + out = model.prep_mw(df, "atms") + assert len(out) == 5 + # Raw channel 1 -> local 0 -> global SENSOR_OFFSET["atms"] + assert (out["global_channel"] == SENSOR_OFFSET["atms"]).all() + assert (out["global_platform"] == PLATFORM_NAME_TO_ID["n20"]).all() + assert np.isnan(out["height"]).all() + + with pytest.raises(ValueError, match="Unknown satellite"): + model.prep_mw(_build_raw_sat_df(n_obs=2, satellite="bogus"), "atms") + + +def test_prep_ir_pca(): + model = _build_model() + codec = _build_codec() + # Two footprints x three channels in long format (channels 1, 5, 9) + t = REQUEST_TIME[0].astype("datetime64[ns]") + rows = [] + for fp, (latv, lonv) in enumerate([(10.0, 20.0), (30.0, 40.0)]): + for chan, bt in [(1, 260.0), (5, 270.0), (9, 280.0)]: + rows.append( + { + "time": t, + "lat": latv, + "lon": lonv, + "observation": bt, + "variable": "airs", + "sensor_index": chan, + "satellite": "aqua", + "scan_angle": 5.0, + "satellite_za": 20.0, + "solza": 60.0, + } + ) + df = pd.DataFrame(rows) + + out = model.prep_ir_pca(df, "airs-pca") + # 2 footprints x 2 latents + assert len(out) == 4 + assert list(out["global_channel"][:2]) == [ + SENSOR_OFFSET["airs-pca"], + SENSOR_OFFSET["airs-pca"] + 1, + ] + assert (out["global_platform"] == PLATFORM_NAME_TO_ID["aqua"]).all() + # Hand-computed latents: standardized bt = [1, 1, 1] @ components + expected = codec.encode(torch.tensor([[260.0, 270.0, 280.0]])).numpy().ravel() + np.testing.assert_allclose(out["observation"].to_numpy()[:2], expected, rtol=1e-6) + # Footprint metadata repeated per latent + np.testing.assert_allclose(out["lat"].to_numpy(), [10.0, 10.0, 30.0, 30.0]) + + +def test_frame_bucketing(): + model = _build_model() + analysis_time = REQUEST_TIME[0] + # Frame g valid time = analysis - 6h * (7 - g); window is (valid-3h, valid+3h] + frame3_valid = pd.Timestamp(analysis_time) - pd.Timedelta(hours=6 * 4) + + df = _build_raw_conv_df(n_obs=4) + df["time"] = np.array( + [ + frame3_valid + pd.Timedelta(hours=3), # inclusive end -> frame 3 + frame3_valid - pd.Timedelta(hours=3), # exclusive start -> frame 2 + frame3_valid, + pd.Timestamp(analysis_time) + pd.Timedelta(hours=4), # outside window + ], + dtype="datetime64[ns]", + ) + obs = model.filter_and_normalize(df, None, analysis_time) + frames = obs.sort_values("obs_time_ns")["frame"].to_list() + assert len(obs) == 3 + assert frames == [2, 3, 3] + # Each frame's target_sec matches its valid time + target = obs[obs["frame"] == 3]["target_sec"].unique() + assert len(target) == 1 + assert target[0] == int(frame3_valid.timestamp()) + + +# ---------- build_input / call tests ---------- + + +def test_build_input_obs_ctx(): + model = _build_model() + conv_df = _build_raw_conv_df(15) + sat_df = _build_raw_sat_df(10) + obs = model.filter_and_normalize(conv_df, sat_df, REQUEST_TIME[0]) + inputs = model.build_input(obs, REQUEST_TIME[0]) + + obs_ctx = inputs["obs_ctx"] + total_pixels = N_WINDOW * NPIX_MODEL + assert obs_ctx.cu_seqlens_k.numel() == total_pixels + 1 + assert int(obs_ctx.cu_seqlens_k[-1]) == len(obs) + assert obs_ctx.obs.shape[0] == len(obs) + assert obs_ctx.float_metadata.shape == (len(obs), 50) + assert inputs["second_of_day"].shape == (1, N_WINDOW) + assert inputs["day_of_year"].shape == (1, N_WINDOW) + assert inputs["condition"].shape == (1, 2, N_WINDOW, NPIX_OUT) + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda:0", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="cuda missing" + ), + ), + ], +) +@pytest.mark.parametrize("lat_lon", [False, True]) +def test_healda_v2_call(device, lat_lon): + model = _build_model(device=device, lat_lon=lat_lon) + conv_df = _build_raw_conv_df(15) + sat_df = _build_raw_sat_df(10) + + out = model(conv_df, sat_df) + + assert isinstance(out, xr.DataArray) + if lat_lon: + assert out.dims == ("time", "lead_time", "variable", "lat", "lon") + assert out.shape == (1, N_WINDOW, NVAR, NLAT, NLON) + else: + assert out.dims == ("time", "lead_time", "variable", "npix") + assert out.shape == (1, N_WINDOW, NVAR, NPIX_OUT) + assert np.all(out.coords["time"].values == REQUEST_TIME) + # Analysis frame is the final lead_time == 0 + assert out.coords["lead_time"].values[-1] == np.timedelta64(0, "ns") + assert out.coords["lead_time"].values[0] == np.timedelta64(-42, "h").astype( + "timedelta64[ns]" + ) + + +def test_healda_v2_call_missing_request_time(): + model = _build_model() + df = _build_raw_conv_df(5) + df.attrs = {} + with pytest.raises(ValueError, match="request_time"): + model(df) + + +def test_healda_v2_call_multiple_request_times(): + model = _build_model() + request_time = np.array( + [np.datetime64("2024-01-01T12:00:00"), np.datetime64("2024-01-01T18:00:00")] + ) + df = _build_raw_conv_df(5, request_time=request_time) + with pytest.raises(ValueError, match="single analysis time"): + model(df) + + +def test_healda_v2_call_no_obs(): + model = _build_model() + with pytest.raises(ValueError, match="At least one"): + model(None, None) + + +def test_healda_v2_call_empty_obs(): + model = _build_model() + df = _build_raw_conv_df(5, observation=9999.0) # violates t valid range + out = model(df) + assert isinstance(out, xr.DataArray) + assert out.shape == (1, N_WINDOW, NVAR, NPIX_OUT) + assert np.all(np.isnan(out.values)) + + +def test_healda_v2_generator(): + model = _build_model() + conv_df = _build_raw_conv_df(10) + sat_df = _build_raw_sat_df(10) + + gen = model.create_generator() + assert gen.send(None) is None + + with patch.object(model, "_forward", _mock_forward): + da = gen.send((conv_df, None)) + assert da.shape == (1, N_WINDOW, NVAR, NPIX_OUT) + da = gen.send((None, sat_df)) + assert da.shape == (1, N_WINDOW, NVAR, NPIX_OUT) + da = gen.send((conv_df, sat_df)) + assert da.shape == (1, N_WINDOW, NVAR, NPIX_OUT) + + with pytest.raises(ValueError, match="At least one"): + gen.send((None, None)) + gen.close() + + +# ---------- Coords tests ---------- + + +def test_healda_v2_init_coords(): + model = _build_model() + assert model.init_coords() is None + + +def test_healda_v2_input_coords(): + model = _build_model() + conv_schema, sat_schema = model.input_coords() + for field in ("time", "lat", "lon", "observation", "variable", "elev", "pres"): + assert field in conv_schema + for field in ("time", "lat", "lon", "observation", "sensor_index", "satellite"): + assert field in sat_schema + # IR sensors are part of the sat schema variables + sat_vars = set(sat_schema["variable"]) + assert {"iasi", "crisfsr", "airs"} <= sat_vars + + +def test_healda_v2_output_coords(): + model = _build_model() + (coords,) = model.output_coords(model.input_coords(), request_time=REQUEST_TIME) + assert list(coords.keys()) == ["time", "lead_time", "variable", "npix"] + assert len(coords["lead_time"]) == N_WINDOW + assert len(coords["variable"]) == NVAR + assert len(coords["npix"]) == NPIX_OUT + + +def test_healda_v2_default_package_unpublished(): + with pytest.raises(NotImplementedError, match="not been published"): + HealDAv2.load_default_package()