diff --git a/.agents/rules/e2s-013-assimilation-models.mdc b/.agents/rules/e2s-013-assimilation-models.mdc index 62c90618c..5dff13f35 100644 --- a/.agents/rules/e2s-013-assimilation-models.mdc +++ b/.agents/rules/e2s-013-assimilation-models.mdc @@ -572,6 +572,6 @@ class MyDAModel(torch.nn.Module, AutoModelMixin): - **Register `device_buffer`** to track device via a `device` property - **Do NOT use `@batch_func` or `@batch_coords`** - these are px/dx conventions only - **Do NOT use `@torch.inference_mode()`** if the forward pass requires gradients (e.g., DPS guidance); document the reason if omitted -- **Documentation**: Add the DA model to [models.rst](mdc:docs/modules/models.rst) in the `earth2studio.models.da` section, maintaining alphabetical order +- **Documentation**: Add the DA model to [models_da.md](mdc:docs/modules/models_da.md), maintaining alphabetical order, and add an entry to [install_options.yml](mdc:docs/userguide/about/install_options.yml) when the model has an optional dependency extra or needs API install notes - DO NOT attempt to make a general base class with intent to reuse the wrapper - DO NOT over-populate the `load_model()` API - only expose essential parameters diff --git a/CHANGELOG.md b/CHANGELOG.md index aa78b513b..fa9bdf4a6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Added WeatherNext 2 Cyclones Mini prognostic model wrapper (`WeatherNext2CyclonesMini`) + ### Changed ### Deprecated diff --git a/docs/modules/models_px.md b/docs/modules/models_px.md index f504a5c70..b760025b2 100644 --- a/docs/modules/models_px.md +++ b/docs/modules/models_px.md @@ -57,6 +57,7 @@ earth2studio.models.px.StormScopeGOES earth2studio.models.px.StormScopeMeteosatEU earth2studio.models.px.StormScopeMRMS earth2studio.models.px.UCast +earth2studio.models.px.WeatherNext2CyclonesMini {% endautosummary %} diff --git a/docs/userguide/about/install_options.yml b/docs/userguide/about/install_options.yml index 24d9933b1..8ff8a4852 100644 --- a/docs/userguide/about/install_options.yml +++ b/docs/userguide/about/install_options.yml @@ -218,6 +218,20 @@ categories: - 'uv add "weathernext @ git+https://github.com/google-deepmind/weathernext.git@9c034db1ff412d5db6cbe6bb0c5c9afc5a267719"' notes: - GraphCast uses WeatherNext and requires Python 3.12 or newer. + - id: weathernext + label: WeatherNext 2 Cyclones Mini + extra: weathernext + summary: Google DeepMind WeatherNext 2 Cyclones Mini model dependencies. + api_refs: + - earth2studio.models.px.WeatherNext2CyclonesMini + preinstall: + pip: + - 'pip install "weathernext @ git+https://github.com/google-deepmind/weathernext.git@9c034db1ff412d5db6cbe6bb0c5c9afc5a267719"' + uv: + pypi: + - 'uv add "weathernext @ git+https://github.com/google-deepmind/weathernext.git@9c034db1ff412d5db6cbe6bb0c5c9afc5a267719"' + notes: + - WeatherNext 2 uses WeatherNext and requires Python 3.12 or newer. - id: interp-modafno label: InterpModAFNO extra: interp-modafno diff --git a/earth2studio/lexicon/wb2.py b/earth2studio/lexicon/wb2.py index bc921a0f8..c33b3cc0d 100644 --- a/earth2studio/lexicon/wb2.py +++ b/earth2studio/lexicon/wb2.py @@ -38,6 +38,8 @@ class WB2Lexicon(metaclass=LexiconType): VOCAB = { "u10m": "10m_u_component_of_wind::", "v10m": "10m_v_component_of_wind::", + "u100m": "100m_u_component_of_wind::", + "v100m": "100m_v_component_of_wind::", "t2m": "2m_temperature::", "sp": "surface_pressure::", "lsm": "land_sea_mask::", diff --git a/earth2studio/models/px/__init__.py b/earth2studio/models/px/__init__.py index 62aa5bcb8..64b28e27c 100644 --- a/earth2studio/models/px/__init__.py +++ b/earth2studio/models/px/__init__.py @@ -55,6 +55,7 @@ ) from earth2studio.models.px.stormscope_meteosat import StormScopeMeteosatEU from earth2studio.models.px.ucast import UCast +from earth2studio.models.px.weathernext2_cyclones_mini import WeatherNext2CyclonesMini # Silence warning spam from various models warnings.filterwarnings("ignore") diff --git a/earth2studio/models/px/weathernext2_cyclones_mini.py b/earth2studio/models/px/weathernext2_cyclones_mini.py new file mode 100644 index 000000000..763009d86 --- /dev/null +++ b/earth2studio/models/px/weathernext2_cyclones_mini.py @@ -0,0 +1,793 @@ +# 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 copy +import dataclasses +from collections import OrderedDict +from collections.abc import Callable, Generator, Iterator +from typing import Any + +import numpy as np +import torch +import xarray as xr +from loguru import logger + +from earth2studio.lexicon.wb2 import WB2Lexicon +from earth2studio.models.auto import AutoModelMixin, Package +from earth2studio.models.batch import batch_coords, batch_func +from earth2studio.models.px.base import PrognosticModel +from earth2studio.models.px.utils import PrognosticMixin +from earth2studio.utils.coords import map_coords +from earth2studio.utils.imports import ( + OptionalDependencyFailure, + check_optional_dependencies, +) +from earth2studio.utils.type import CoordSystem + +try: + import chex + import haiku as hk + import jax + import pandas as pd + from weathernext.cyclones import constants as cyclone_constants + from weathernext.cyclones import direct_tracker_6h_v1_config + from weathernext.utils import checkpoint, data_utils, fiddle_config_io, rollout + from weathernext.weathernext2 import fgn +except ImportError: + OptionalDependencyFailure("weathernext") + chex = None + checkpoint = None + data_utils = None + cyclone_constants = None + direct_tracker_6h_v1_config = None + fgn = None + fiddle_config_io = None + hk = None + jax = None + pd = None + rollout = None + + +SURFACE_INPUT_VARIABLES = [ + "t2m", + "msl", + "v10m", + "u10m", + "sst", +] +SURFACE_OUTPUT_VARIABLES = SURFACE_INPUT_VARIABLES + ["tp06"] +ATMOS_VARIABLES = ["t", "z", "u", "v", "w", "q"] +PRESSURE_LEVELS = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000] + +INPUT_VARIABLES = SURFACE_INPUT_VARIABLES + [ + f"{var}{level}" for var in ATMOS_VARIABLES for level in PRESSURE_LEVELS +] +OUTPUT_VARIABLES = SURFACE_OUTPUT_VARIABLES + [ + f"{var}{level}" for var in ATMOS_VARIABLES for level in PRESSURE_LEVELS +] + +WN2_TARGET_VARIABLES = tuple( + dict.fromkeys(WB2Lexicon.VOCAB[var].split("::")[0] for var in OUTPUT_VARIABLES) +) +INV_VOCAB = {v: k for k, v in WB2Lexicon.VOCAB.items()} +KNOTS_TO_METERS_PER_SECOND = 0.514444 + + +def _add_e2s_cyclone_columns(tracks: "pd.DataFrame") -> "pd.DataFrame": + """Add Earth2Studio TC observation aliases to WeatherNext track columns.""" + if tracks.empty: + return tracks + tracks = tracks.copy() + tracks["tcmsl"] = tracks["minimum_sea_level_pressure_hpa"] * 100.0 + tracks["tcw10m"] = ( + tracks["maximum_sustained_wind_speed_knots"] * KNOTS_TO_METERS_PER_SECOND + ) + return tracks + + +MODEL_NAME = "WeatherNextCyclones_Mini" +MODEL_SPLIT = "2024" +PARAMS_PATH = f"params/{MODEL_NAME}_<{MODEL_SPLIT}.npz" +SAMPLE_PATH = ( + "dataset/source-hres_forecast_init-2024-10-07 00:00:00_" + "res-1.0_levels-13_steps-01.nc" +) + + +@check_optional_dependencies() +class WeatherNext2CyclonesMini(torch.nn.Module, AutoModelMixin, PrognosticMixin): + """WeatherNext 2 Cyclones Mini medium-range forecast model. + + WeatherNext 2 is Google DeepMind's global medium-range weather forecasting + model family. This wrapper currently uses the public + ``WeatherNextCyclones_Mini`` checkpoint and 1 degree sample grid, which is the + mini model configuration that can be validated on the available single-GPU + test hardware. + + The model requires two input states, valid at ``-6h`` and ``0h`` lead time, + and predicts 6 hours forward per model call. By default this wrapper returns + only the gridded weather fields expected by Earth2Studio prognostic models. + Cyclone tracking can be enabled with ``track_cyclones=True`` to accumulate + WeatherNext's tropical cyclone track diagnostics in the ``cyclone_tracks`` + property without changing the model output type. + + + Note + ---- + For more information see the following references: + + - https://doi.org/10.1038/s41586-026-10953-2 + - https://github.com/google-deepmind/weathernext + + Warning + ------- + We encourage users to familiarize themselves with the license restrictions of this + model's checkpoints. + + Parameters + ---------- + ckpt : fgn.CheckPoint + Model checkpoint containing weights. + land_sea_mask : np.ndarray + Land-sea mask on the WeatherNext grid. + geopotential_at_surface : np.ndarray + Surface geopotential on the WeatherNext grid. + seed : int, optional + Initial random seed for the stochastic FGN noise generator, by default 0. + jit_compile : bool, optional + JIT-compile the model forward pass, by default True. + track_cyclones : bool, optional + Accumulate tropical cyclone tracks in the ``cyclone_tracks`` property, + by default False. + + Examples + -------- + Access tropical cyclone tracks after a model call: + + >>> model = WeatherNext2CyclonesMini.load_model( + ... WeatherNext2CyclonesMini.load_default_package(), + ... track_cyclones=True, + ... ) + >>> x, coords = model(x, coords) + >>> tracks = model.cyclone_tracks + >>> tracks[["track_id", "lead_time", "lat", "lon", "tcmsl", "tcw10m"]] + + The ``tcmsl`` and ``tcw10m`` columns provide Earth2Studio-compatible names + for the minimum sea-level pressure and surface wind speed diagnostics. + + The tracker filters short-lived cyclogenesis tracks, so short rollouts can + return an empty dataframe even when cyclone tracking is active. The active + duration threshold is set by + `model._cyclone_tracker.cyclogenesis_minimum_duration`. + + Badges + ------ + region:global class:medium-range product:wind product:precip product:temp product:atmos + product:ocean year:2026 gpu:40gb provider:google backend:jax + """ + + def __init__( + self, + ckpt: "fgn.CheckPoint", + land_sea_mask: np.ndarray, + geopotential_at_surface: np.ndarray, + seed: int = 0, + jit_compile: bool = True, + track_cyclones: bool = False, + ): + super().__init__() + + self.ckpt = ckpt + self.land_sea_mask = land_sea_mask + self.geopotential_at_surface = geopotential_at_surface + self.seed = seed + self.prng_key = jax.random.PRNGKey(seed) + self.track_cyclones = track_cyclones + self._cyclone_tracks = pd.DataFrame() + self._cyclone_prediction_history: list[xr.Dataset] = [] + self._cyclone_tracker = None + if self.track_cyclones: + tracker_config = direct_tracker_6h_v1_config.get_config() + self._cyclone_tracker = tracker_config.tracker_constructor( + **tracker_config.tracker_kwargs + ) + self.task_config = self._load_task_config() + self.run_forward = self._load_run_forward_from_checkpoint( + jit_compile=jit_compile + ) + + n_lat = land_sea_mask.shape[0] + n_lon = land_sea_mask.shape[1] + self._input_coords = OrderedDict( + { + "batch": np.empty(0), + "time": np.empty(0), + "lead_time": np.array( + [np.timedelta64(-6, "h"), np.timedelta64(0, "h")] + ), + "variable": np.array(INPUT_VARIABLES), + "lat": np.linspace(90, -90, n_lat, endpoint=True), + "lon": np.linspace(0, 360, n_lon, endpoint=False), + } + ) + self._output_coords = OrderedDict( + { + "batch": np.empty(0), + "time": np.empty(0), + "lead_time": np.array([np.timedelta64(6, "h")]), + "variable": np.array(OUTPUT_VARIABLES), + "lat": np.linspace(90, -90, n_lat, endpoint=True), + "lon": np.linspace(0, 360, n_lon, endpoint=False), + } + ) + + @property + def cyclone_tracks(self) -> "pd.DataFrame": + """Tropical cyclone tracks accumulated during the latest model run.""" + if not self.track_cyclones: + logger.warning("Cyclone tracking is currently not active on this model.") + return pd.DataFrame() + return self._cyclone_tracks.copy() + + def _reset_cyclone_tracks(self) -> None: + """Reset accumulated cyclone track diagnostics.""" + self._cyclone_tracks = pd.DataFrame() + self._cyclone_prediction_history = [] + + @staticmethod + def _empty_initial_storms() -> "pd.DataFrame": + """Create an empty initial storm table for pure cyclogenesis tracking.""" + return pd.DataFrame( + columns=[ + cyclone_constants.TRACK_ID, + cyclone_constants.LEAD_TIME, + cyclone_constants.VALID_TIME, + cyclone_constants.LAT, + cyclone_constants.LON, + ] + ) + + def _update_cyclone_tracks( + self, + predictions: xr.Dataset, + coords: CoordSystem, + accumulate_predictions: bool, + ) -> None: + """Update cyclone tracks from native WeatherNext prediction fields.""" + if not self.track_cyclones: + return + if self._cyclone_tracker is None: + logger.warning("Cyclone tracking is active, but no tracker is available.") + return + + init_times = np.asarray(coords["time"]).reshape(-1) + if len(init_times) != 1: + logger.warning( + "Cyclone tracking currently supports one init time per model run." + ) + return + + cyclone_vars = [ + var for var in predictions.data_vars if var.startswith("cyclone") + ] + if not cyclone_vars: + logger.warning( + "Cyclone tracking is active, but this prediction did not include " + "cyclone fields." + ) + return + + cyclone_predictions = predictions[cyclone_vars].copy() + if "batch" in cyclone_predictions.dims: + if cyclone_predictions.sizes["batch"] != 1: + logger.warning("Cyclone tracking currently supports batch size one.") + return + cyclone_predictions = cyclone_predictions.isel(batch=0, drop=True) + cyclone_predictions = cyclone_predictions.assign_coords( + time=np.asarray(coords["lead_time"]), init_time=init_times[0] + ) + for name in cyclone_predictions.data_vars: + array = cyclone_predictions[name] + cyclone_predictions[name] = xr.DataArray( + np.asarray(array.data), + dims=array.dims, + coords=array.coords, + attrs=array.attrs, + name=name, + ) + + if accumulate_predictions: + self._cyclone_prediction_history.append(cyclone_predictions) + tracker_input = xr.concat(self._cyclone_prediction_history, dim="time") + tracker_input = tracker_input.sortby("time") + self._cyclone_tracks = _add_e2s_cyclone_columns( + self._cyclone_tracker( + tracker_input, initial_storms_df=self._empty_initial_storms() + ) + ) + return + + tracks = _add_e2s_cyclone_columns( + self._cyclone_tracker( + cyclone_predictions, initial_storms_df=self._empty_initial_storms() + ) + ) + if self._cyclone_tracks.empty: + self._cyclone_tracks = tracks + elif not tracks.empty: + self._cyclone_tracks = pd.concat( + [self._cyclone_tracks, tracks], ignore_index=True + ) + + def input_coords(self) -> CoordSystem: + """Input coordinate system of the prognostic model. + + Returns + ------- + CoordSystem + Coordinate system dictionary. + """ + return self._input_coords.copy() + + @batch_coords() + def output_coords(self, input_coords: CoordSystem) -> CoordSystem: + """Output coordinate system of the prognostic model. + + Parameters + ---------- + input_coords : CoordSystem + Input coordinate system to transform into output_coords. + + Returns + ------- + CoordSystem + Coordinate system dictionary. + """ + output_coords = self._output_coords.copy() + output_coords["batch"] = input_coords["batch"] + output_coords["time"] = input_coords["time"] + output_coords["lead_time"] = ( + input_coords["lead_time"][-1] + output_coords["lead_time"] + ) + return output_coords + + @classmethod + def load_default_package(cls) -> Package: + """Load default pre-trained WeatherNext 2 package from Google Cloud. + + Returns + ------- + Package + Model package. + """ + return Package( + "gs://dm_graphcast/weathernext2", + cache_options={ + "cache_storage": Package.default_cache("weathernext2"), + "same_names": True, + }, + ) + + @classmethod + @check_optional_dependencies() + def load_model( + cls, + package: Package, + seed: int = 0, + jit_compile: bool = True, + track_cyclones: bool = False, + ) -> PrognosticModel: + """Load prognostic model from package. + + Parameters + ---------- + package : Package + Package to load model from. + seed : int, optional + Initial random seed for the stochastic FGN noise generator, by default 0. + jit_compile : bool, optional + JIT-compile the model forward pass, by default True. + track_cyclones : bool, optional + Accumulate tropical cyclone tracks in the ``cyclone_tracks`` property, + by default False. + + Returns + ------- + PrognosticModel + Prognostic model. + """ + params_path = package.resolve(PARAMS_PATH) + with open(params_path, "rb") as f: + ckpt = checkpoint.load(f, fgn.CheckPoint) + + sample_input = xr.load_dataset(package.resolve(SAMPLE_PATH)) + land_sea_mask = sample_input["land_sea_mask"].values + geopotential_at_surface = sample_input["geopotential_at_surface"].values + + return cls( + ckpt, + land_sea_mask, + geopotential_at_surface, + seed=seed, + jit_compile=jit_compile, + track_cyclones=track_cyclones, + ) + + def _load_task_config(self) -> Any: + config = fiddle_config_io.get_fiddle_config_by_name( + f"weathernext2/configs/{MODEL_NAME}" + ) + target_variables = ( + config.task.target_variables + if self.track_cyclones + else WN2_TARGET_VARIABLES + ) + return dataclasses.replace(config.task, target_variables=target_variables) + + def _load_run_forward_from_checkpoint(self, jit_compile: bool = True) -> Callable: + """Build WeatherNext 2 inference function from checkpoint.""" + config = copy.deepcopy( + fiddle_config_io.get_fiddle_config_by_name( + f"weathernext2/configs/{MODEL_NAME}" + ) + ) + task_config = self.task_config + noisy_function_kwargs = config.predictor_kwargs["noisy_function_kwargs"] + noisy_function_kwargs["per_var_activation_fns"] = { + key: value + for key, value in noisy_function_kwargs.get( + "per_var_activation_fns", {} + ).items() + if key in task_config.target_variables + } + transformer_kwargs = noisy_function_kwargs["mesh_model_ctor"].keywords[ + "transformer_kwargs" + ] + if jax.default_backend() == "gpu": + transformer_kwargs["attention_type"] = "triblockdiag_mha" + + config_inference = fgn.PredictorConfig( + task=task_config, + predictor_constructor=config.predictor_constructor, + predictor_kwargs=config.predictor_kwargs, + predictor_wrappers=config.predictor_wrappers[:-1], + ) + + @hk.transform + def run_forward( + inputs: xr.Dataset, targets_template: xr.Dataset, forcings: xr.Dataset + ) -> xr.Dataset: + predictor = fgn.construct_predictor(config_inference) + return predictor( + inputs, targets_template=targets_template, forcings=forcings + ) + + def apply( + rng: "chex.PRNGKey", + inputs: xr.Dataset, + targets_template: xr.Dataset, + forcings: xr.Dataset, + ) -> xr.Dataset: + return run_forward.apply( + self.ckpt.params, rng, inputs, targets_template, forcings + ) + + if jit_compile: + return jax.jit(apply) + return apply + + def _chunked_prediction_generator( + self, + predictor_fn: Callable, + rng: "chex.PRNGKey", + inputs: xr.Dataset, + targets_template: xr.Dataset, + batch: xr.Dataset, + forcings: xr.Dataset, + ) -> Generator[xr.Dataset, None, None]: + """Generate an open-ended WeatherNext 2 rollout one chunk at a time.""" + inputs = xr.Dataset(inputs) + targets_template = xr.Dataset(targets_template) + forcings = xr.Dataset(forcings) + targets_chunk_time = targets_template.time.isel(time=slice(0, 1)) + current_inputs = inputs + forcing_variables = list(self.task_config.forcing_variables) + index = 0 + + while True: + forcings = forcings.assign_coords(time=targets_chunk_time).compute() + rng, step_rng = jax.random.split(rng) + predictions = predictor_fn( + rng=step_rng, + inputs=current_inputs, + targets_template=targets_template, + forcings=forcings, + ) + next_frame = xr.merge([predictions, forcings]) + current_inputs = rollout._get_next_inputs(current_inputs, next_frame) + current_inputs = current_inputs.assign_coords(time=inputs.coords["time"]) + predictions = predictions.assign_coords( + time=targets_template.coords["time"] + index * np.timedelta64(6, "h") + ) + yield predictions + + batch = batch.assign_coords( + datetime=batch.coords["datetime"] + np.timedelta64(6, "h") + ) + batch = batch.drop_vars( + forcing_variables + ["year_progress", "day_progress"], errors="ignore" + ) + data_utils.add_derived_vars(batch) + data_utils.add_tisr_var(batch) + batch = batch.compute() + forcings = batch.isel(time=slice(-1, None))[forcing_variables] + forcings = forcings.reset_coords("datetime", drop=True).compute() + index += 1 + + def iterator_result_to_tensor(self, dataset: xr.Dataset) -> torch.Tensor: + """Convert an xarray Dataset prediction to an Earth2Studio tensor.""" + dataset = dataset[ + [var for var in dataset.data_vars if var in WN2_TARGET_VARIABLES] + ] + for var in list(dataset.data_vars): + if "level" in dataset[var].dims: + for level in dataset[var].level: + dataset[f"{var}::{level.values}"] = dataset[var].sel(level=level) + dataset = dataset.drop_vars(var) + else: + dataset = dataset.rename({var: f"{var}::"}) + + if "level" in dataset.dims: + dataset = dataset.drop_dims("level") + if len(dataset.time) > 1: + dataset = dataset.rename({"time": "lead_time"}) + dataset = dataset.expand_dims(dim="time") + else: + dataset = dataset.expand_dims(dim="lead_time") + if "sample" in dataset.dims: + dataset = dataset.isel(sample=0, drop=True) + + dataset = dataset.rename({key: INV_VOCAB[key] for key in dataset.data_vars}) + if "batch" in dataset.dims: + dataarray = ( + dataset[OUTPUT_VARIABLES] + .to_dataarray() + .T.transpose( + ..., "batch", "time", "lead_time", "variable", "lat", "lon" + ) + ) + else: + dataarray = ( + dataset[OUTPUT_VARIABLES] + .to_dataarray() + .T.transpose(..., "time", "lead_time", "variable", "lat", "lon") + ) + out = torch.from_numpy(dataarray.to_numpy().copy()) + return out.flip(-2) + + @staticmethod + def get_jax_device_from_tensor(x: torch.Tensor) -> "jax.Device": + """From a tensor, get device and corresponding JAX device.""" + device_id = x.get_device() + if device_id == -1: + return jax.devices("cpu")[0] + return jax.devices("gpu")[device_id] + + def from_dataarray_to_dataset( + self, data: xr.DataArray, lead_time: int = 6, hour_steps: int = 6 + ) -> tuple[xr.Dataset, list[str]]: + """Convert an Earth2Studio DataArray to a WeatherNext 2 Dataset.""" + if len(data.time.values) > 1: + raise TypeError("WeatherNext 2 only supports one init_time per JAX call.") + if "lead_time" in data.dims: + data["lead_time"] = [ + data.time.values[0] + level for level in data.lead_time.values + ] + data = data.isel(time=0).reset_coords("time", drop=True) + data = data.rename({"lead_time": "time"}) + + lead_times = range(hour_steps, lead_time + hour_steps, hour_steps) + target_lead_times = [f"{h}h" for h in lead_times] + time_deltas = np.concatenate( + ( + self._input_coords["lead_time"], + [np.timedelta64(h, "h") for h in lead_times], + ) + ) + start_date = data.time.values[-1] + all_datetimes = [start_date + time_delta for time_delta in time_deltas] + + data = data.to_dataset(dim="variable") + data = data.rename({key: WB2Lexicon.VOCAB[key] for key in data.data_vars}) + out_data = xr.Dataset( + coords={ + "time": all_datetimes[0:2], + "lat": data.lat, + "lon": data.lon, + "level": PRESSURE_LEVELS, + } + ) + + pressure_level_vars: dict[str, list[xr.DataArray]] = {} + for var in data.data_vars: + wb2_variable, level = var.split("::") + if level: + pressure_level_vars.setdefault(wb2_variable, []).append( + data[var].expand_dims(dim=dict(level=[int(level)])) + ) + else: + out_data[wb2_variable] = data[var] + for var in pressure_level_vars: + out_data[var] = xr.concat(pressure_level_vars[var], dim="level") + + out_data = out_data.assign_coords( + datetime=all_datetimes[: len(out_data.time.values)] + ) + out_data = out_data.assign_coords(time=time_deltas[: len(out_data.time.values)]) + out_data["datetime"] = out_data.datetime.expand_dims(dict(batch=1)) + for var in out_data.data_vars: + if "batch" not in out_data[var].dims: + out_data[var] = out_data[var].expand_dims(dict(batch=1)) + + out_data = out_data.pad(pad_width=dict(time=(0, len(lead_times)))) + out_data = out_data.assign_coords( + coords=dict(time=time_deltas, datetime=(("batch", "time"), [all_datetimes])) + ) + out_data = out_data.reindex(lat=sorted(out_data.lat.values)) + out_data = out_data.transpose("batch", "time", "level", "lat", "lon", ...) + out_data["land_sea_mask"] = xr.DataArray( + self.land_sea_mask, dims=("lat", "lon") + ) + out_data["geopotential_at_surface"] = xr.DataArray( + self.geopotential_at_surface, dims=("lat", "lon") + ) + out_data["total_precipitation_6hr"] = xr.full_like( + out_data["2m_temperature"], np.nan + ) + for var in self.task_config.target_variables: + if var.startswith("cyclone") and var not in out_data: + out_data[var] = xr.full_like(out_data["2m_temperature"], np.nan) + for var in out_data.data_vars: + out_data[var] = out_data[var].astype(np.float32) + return out_data, target_lead_times + + @batch_func() + def __call__( + self, x: torch.Tensor, coords: CoordSystem + ) -> tuple[torch.Tensor, CoordSystem]: + """Runs prognostic model one step. + + Parameters + ---------- + x : torch.Tensor + Input tensor. + coords : CoordSystem + Input coordinate system. + + Returns + ------- + tuple[torch.Tensor, CoordSystem] + Output tensor and coordinate system 6 hours in the future. + """ + self._reset_cyclone_tracks() + device = x.device + with jax.default_device(self.get_jax_device_from_tensor(x)): + x, coords = map_coords(x, coords, self.input_coords()) + time_dim = list(coords.keys()).index("time") + results = [] + for t in range(len(coords["time"])): + x_t = x.narrow(time_dim, t, 1) + coords_t = coords.copy() + coords_t["time"] = coords["time"][t : t + 1] + data, target_lead_times = self.from_dataarray_to_dataset( + xr.DataArray(x_t.cpu(), coords=coords_t), 6 + ) + inputs, targets, forcings = data_utils.extract_inputs_targets_forcings( + data, + target_lead_times=target_lead_times, + **dataclasses.asdict(self.task_config), + ) + self.prng_key, rng = jax.random.split(self.prng_key) + predictions = rollout.chunked_prediction( + self.run_forward, + rng=rng, + inputs=inputs, + targets_template=targets * np.nan, + forcings=forcings, + ) + self._update_cyclone_tracks( + predictions, + self.output_coords(coords_t), + accumulate_predictions=False, + ) + results.append(self.iterator_result_to_tensor(predictions)) + + out = torch.cat(results, dim=1) if len(results) > 1 else results[0] + return out.to(device), self.output_coords(coords) + + @batch_func() + def _default_generator( + self, x: torch.Tensor, coords: CoordSystem + ) -> Generator[tuple[torch.Tensor, CoordSystem]]: + coords = coords.copy() + self.output_coords(coords) + device = x.device + coords_out = coords.copy() + coords_out["lead_time"] = coords["lead_time"][1:] + yield x[:, :, 1:, ...], coords_out + + while True: + coords = self.output_coords(coords) + predictions = [next(it) for it in self.iterators] + if len(predictions) == 1: + self._update_cyclone_tracks( + predictions[0], coords, accumulate_predictions=True + ) + elif self.track_cyclones: + logger.warning( + "Cyclone tracking currently supports one init time per iterator." + ) + results = [self.iterator_result_to_tensor(pred) for pred in predictions] + x = torch.cat(results, dim=1) if len(results) > 1 else results[0] + x, coords = self.rear_hook(x, coords) + yield x.to(device), coords.copy() + + def create_iterator( + self, x: torch.Tensor, coords: CoordSystem + ) -> Iterator[tuple[torch.Tensor, CoordSystem]]: + """Create a time-integration iterator for the prognostic model. + + Parameters + ---------- + x : torch.Tensor + Input tensor. + coords : CoordSystem + Input coordinate system. + + Yields + ------ + Iterator[tuple[torch.Tensor, CoordSystem]] + Iterator that generates model time steps. + """ + self.output_coords(coords) + self._reset_cyclone_tracks() + with jax.default_device(self.get_jax_device_from_tensor(x)): + time_dim = list(coords.keys()).index("time") + self.iterators = [] + for t in range(len(coords["time"])): + x_t = x.narrow(time_dim, t, 1) + coords_t = coords.copy() + coords_t["time"] = coords["time"][t : t + 1] + data, target_lead_times = self.from_dataarray_to_dataset( + xr.DataArray(x_t.cpu(), coords=coords_t), 6 + ) + inputs, targets, forcings = data_utils.extract_inputs_targets_forcings( + data, + target_lead_times=target_lead_times, + **dataclasses.asdict(self.task_config), + ) + self.prng_key, rng = jax.random.split(self.prng_key) + self.iterators.append( + self._chunked_prediction_generator( + predictor_fn=self.run_forward, + rng=rng, + inputs=inputs, + targets_template=targets * np.nan, + batch=data, + forcings=forcings, + ) + ) + yield from self._default_generator(x, coords) diff --git a/mkdocs.yml b/mkdocs.yml index 910f39f2e..4ce44a609 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -425,6 +425,8 @@ plugins: output_open: true collect_telemetry: true backreferences: true + cache_output_directory: false + invalidate_on_lock_change: false - blog: blog_dir: blog post_dir: '{blog}/posts' diff --git a/pyproject.toml b/pyproject.toml index e90524f5b..4872af0d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -202,6 +202,13 @@ graphcast = [ "flax>=0.10.6", "dm-tree>=0.1.9", ] +weathernext = [ + "weathernext; python_version >= '3.12'", + "dm-haiku>=0.0.14", + "jax[cuda13]>=0.4.26", + "flax>=0.10.6", + "dm-tree>=0.1.9", +] interp-modafno = [ "nvidia-physicsnemo>=2.0", ] @@ -314,7 +321,7 @@ da-stormcast = [ # [tool.uv.sources] and the conflicts declared in [tool.uv]). all = [ "earth2studio[data,perturbation,statistics,utils,serve]", - "earth2studio[atlas,aurora,dlwp,dlesym,fcn,fcn3,fengwu,gencast,interp-modafno,pangu,stormcast,sfno,stormscope,graphcast,ucast]", + "earth2studio[atlas,aurora,dlwp,dlesym,fcn,fcn3,fengwu,gencast,interp-modafno,pangu,stormcast,sfno,stormscope,graphcast,ucast,weathernext]", "earth2studio[cbottle,climatenet,corrdiff,cosmo,orbit,precip-afno,cyclone,precip-afno-v2,solarradiation-afno,windgust-afno]", "earth2studio[da-interp,da-stormcast,da-healda,da-cosmo]", ] diff --git a/skills/earth2studio-create-diagnostic/SKILL.md b/skills/earth2studio-create-diagnostic/SKILL.md index e330cefff..49cd7ea74 100644 --- a/skills/earth2studio-create-diagnostic/SKILL.md +++ b/skills/earth2studio-create-diagnostic/SKILL.md @@ -261,8 +261,8 @@ file that should not be exported. For public models: -- Add to `docs/modules/models_dx.rst` alphabetically so API docs include the generated page. -- Add to `docs/userguide/about/install.md` if a model extra exists. Include model notes plus both `pip install earth2studio[model-name]` and `uv add earth2studio --extra model-name` instructions. +- Add to `docs/modules/models_dx.md` alphabetically so API docs include the generated page. +- Add an entry to `docs/userguide/about/install_options.yml` if a model extra exists. Include model notes, source-specific preinstall commands, and `api_refs` for any model classes that should show a View Install Notes button in the generated API docs. - Update `CHANGELOG.md` under `### Added`. Format and lint: diff --git a/skills/earth2studio-create-diagnostic/references/validation-guide.md b/skills/earth2studio-create-diagnostic/references/validation-guide.md index d4e141f12..3b55b4db3 100644 --- a/skills/earth2studio-create-diagnostic/references/validation-guide.md +++ b/skills/earth2studio-create-diagnostic/references/validation-guide.md @@ -83,8 +83,9 @@ placeholders so the PR author can upload plots manually in the browser. Before opening or updating the PR, verify that packaged diagnostics have a `pyproject.toml` optional dependency extra, that the `all` extra includes it, -that install docs include both pip and uv commands, and that the model is listed -in `docs/modules/models_dx.rst` and `CHANGELOG.md` when it is public. +that `docs/userguide/about/install_options.yml` has model notes, preinstall +commands, and `api_refs`, and that the model is listed in +`docs/modules/models_dx.md` and `CHANGELOG.md` when it is public. Stage only implementation, tests, docs, changelog, dependency metadata, and skill updates that belong in the branch. Exclude validation scripts and outputs. @@ -96,8 +97,8 @@ git add \ test/models/dx/test_.py \ pyproject.toml \ CHANGELOG.md \ - docs/modules/models_dx.rst \ - docs/userguide/about/install.md + docs/modules/models_dx.md \ + docs/userguide/about/install_options.yml ``` Create the PR with the body template and then post the validation comment from diff --git a/skills/earth2studio-create-prognostic/SKILL.md b/skills/earth2studio-create-prognostic/SKILL.md index 71a1717d1..204d43ea5 100644 --- a/skills/earth2studio-create-prognostic/SKILL.md +++ b/skills/earth2studio-create-prognostic/SKILL.md @@ -179,12 +179,12 @@ while still loading real weights and running a forward pass. ### Step 9 — Documentation -- Add to `docs/modules/models_px.rst` (alphabetical). This is required for +- Add to `docs/modules/models_px.md` (alphabetical). This is required for every new prognostic model so the API docs include the generated page. -- Add to `docs/userguide/about/install.md` (alphabetical tab) for the - model extra, even when the extra is empty. Include model-specific notes plus - both `pip install earth2studio[model-name]` and - `uv add earth2studio --extra model-name` instructions. +- Add an entry to `docs/userguide/about/install_options.yml` for the model + extra, even when the extra is empty. Include model-specific notes, + source-specific preinstall commands, and `api_refs` for any model classes + that should show a View Install Notes button in the generated API docs. - Update `CHANGELOG.md` under `### Added`. This is required for every new prognostic model. @@ -208,8 +208,9 @@ Follow `references/validation-guide.md` and use: - `references/pr-comment-template.md` Before creating the PR, verify `pyproject.toml` has the model extra, the -`all` extra includes it, install docs include both pip and uv commands, and -`docs/modules/models_px.rst` plus `CHANGELOG.md` are updated. +`all` extra includes it, `docs/userguide/about/install_options.yml` has the +install entry and `api_refs`, and `docs/modules/models_px.md` plus +`CHANGELOG.md` are updated. Do not include machine names, absolute paths, device inventory, or uploaded image links in PR text. Use plot placeholders instead. diff --git a/skills/earth2studio-create-prognostic/references/validation-guide.md b/skills/earth2studio-create-prognostic/references/validation-guide.md index 551c50fda..8e9f47a73 100644 --- a/skills/earth2studio-create-prognostic/references/validation-guide.md +++ b/skills/earth2studio-create-prognostic/references/validation-guide.md @@ -82,13 +82,14 @@ machine-identifying details. Before opening or updating the PR, verify that new prognostic models have a `pyproject.toml` optional dependency extra, even if empty, that the `all` extra -includes it, that install docs include model notes plus both pip and uv commands, -and that the model is listed in `docs/modules/models_px.rst` and `CHANGELOG.md`. +includes it, that `docs/userguide/about/install_options.yml` has model notes, +preinstall commands, and `api_refs`, and that the model is listed in +`docs/modules/models_px.md` and `CHANGELOG.md`. Stage only implementation, tests, docs, changelog, dependency metadata, and skill updates that belong in the branch. Exclude validation scripts and outputs. ```bash -git add earth2studio/models/px/.py earth2studio/models/px/__init__.py test/models/px/test_.py pyproject.toml CHANGELOG.md docs/modules/models_px.rst docs/userguide/about/install.md +git add earth2studio/models/px/.py earth2studio/models/px/__init__.py test/models/px/test_.py pyproject.toml CHANGELOG.md docs/modules/models_px.md docs/userguide/about/install_options.yml ``` Create the PR with the body template and then post the validation comment from diff --git a/test/models/px/test_weathernext2.py b/test/models/px/test_weathernext2.py new file mode 100644 index 000000000..f8613aa45 --- /dev/null +++ b/test/models/px/test_weathernext2.py @@ -0,0 +1,265 @@ +# 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 collections import OrderedDict +from unittest import mock + +import numpy as np +import pandas as pd +import pytest +import torch +import xarray as xr + +try: + from weathernext.weathernext2 import fgn +except ImportError: + pytest.importorskip("weathernext") + +from earth2studio.data import Random, fetch_data +from earth2studio.models.px.weathernext2_cyclones_mini import ( + OUTPUT_VARIABLES, + WeatherNext2CyclonesMini, +) +from earth2studio.utils import handshake_dim + +TEST_TIME = np.array([np.datetime64("2025-01-01T00:00")]) +DEVICES = ["cpu", "cuda:0"] + + +def mocked_chunked_prediction( + predictor_fn, + rng, + inputs, + targets_template, + forcings, + num_steps_per_chunk=None, + verbose=None, +): + return targets_template + + +def mocked_chunked_prediction_generator( + self, + predictor_fn, + rng, + inputs, + targets_template, + batch, + forcings, +): + while True: + yield targets_template.isel(time=[0]) + + +@pytest.fixture +def mock_weathernext2_model(): + ckpt = fgn.CheckPoint(params={}, description="mock", license="license") + with mock.patch.object( + WeatherNext2CyclonesMini, "_load_run_forward_from_checkpoint", return_value=None + ): + return WeatherNext2CyclonesMini( + ckpt, + land_sea_mask=np.ones((9, 12), dtype=np.float32), + geopotential_at_surface=np.ones((9, 12), dtype=np.float32), + jit_compile=False, + ) + + +def fetch_random_input(model, time=TEST_TIME, device="cpu"): + input_coords = model.input_coords() + random_coords = input_coords.copy() + for dim in ("batch", "time", "lead_time", "variable"): + del random_coords[dim] + return fetch_data( + Random(random_coords), + time, + input_coords["variable"], + input_coords["lead_time"], + device=device, + ) + + +def assert_output(model, out, out_coords, coords, time): + assert out.shape == torch.Size([len(time), 1, len(OUTPUT_VARIABLES), 9, 12]) + assert (out_coords["variable"] == model.output_coords(coords)["variable"]).all() + assert (out_coords["time"] == time).all() + for dim, index in ( + ("lon", 4), + ("lat", 3), + ("variable", 2), + ("lead_time", 1), + ("time", 0), + ): + handshake_dim(out_coords, dim, index) + + +@pytest.mark.parametrize( + "time", + [ + TEST_TIME, + np.array( + [np.datetime64("2025-01-01T00:00"), np.datetime64("2025-01-02T00:00")] + ), + ], +) +@pytest.mark.parametrize("device", DEVICES) +@mock.patch("weathernext.utils.rollout.chunked_prediction", mocked_chunked_prediction) +def test_weathernext2_call(time, device, mock_weathernext2_model): + model = mock_weathernext2_model.to(device) + x, coords = fetch_random_input(model, time, device) + out, out_coords = model(x, coords) + assert_output(model, out, out_coords, coords, time) + + +@pytest.mark.parametrize("device", DEVICES) +@mock.patch.object( + WeatherNext2CyclonesMini, + "_chunked_prediction_generator", + mocked_chunked_prediction_generator, +) +def test_weathernext2_iter(device, mock_weathernext2_model): + model = mock_weathernext2_model.to(device) + x, coords = fetch_random_input(model, device=device) + model_iter = model.create_iterator(x, coords) + + out, out_coords = next(model_iter) + assert out_coords["lead_time"] == np.timedelta64(0, "h") + assert out.shape == torch.Size([1, 1, len(model.input_coords()["variable"]), 9, 12]) + + for i in range(7): + out, out_coords = next(model_iter) + assert_output(model, out, out_coords, coords, TEST_TIME) + assert out_coords["lead_time"] == np.timedelta64(6 * (i + 1), "h") + + +@mock.patch("weathernext.utils.rollout.chunked_prediction") +def test_weathernext2_rng_advances(chunked_prediction, mock_weathernext2_model): + rngs = [] + + def mock_prediction(predictor_fn, rng, inputs, targets_template, forcings): + rngs.append(np.asarray(rng)) + return targets_template + + chunked_prediction.side_effect = mock_prediction + x, coords = fetch_random_input(mock_weathernext2_model) + mock_weathernext2_model(x, coords) + mock_weathernext2_model(x, coords) + + assert len(rngs) == 2 + assert not np.array_equal(rngs[0], rngs[1]) + + +def test_weathernext2_cyclone_tracks_inactive(mock_weathernext2_model): + with mock.patch( + "earth2studio.models.px.weathernext2_cyclones_mini.logger.warning" + ) as warning: + tracks = mock_weathernext2_model.cyclone_tracks + + assert tracks.empty + warning.assert_called_once_with( + "Cyclone tracking is currently not active on this model." + ) + + +def test_weathernext2_cyclone_tracks_have_e2s_observation_names( + mock_weathernext2_model, +): + mock_weathernext2_model.track_cyclones = True + mock_weathernext2_model._cyclone_tracker = mock.Mock( + return_value=pd.DataFrame( + { + "track_id": ["storm-0"], + "lead_time": [pd.Timedelta(hours=6)], + "valid_time": [pd.Timestamp("2025-01-01T06:00")], + "lat": [10.0], + "lon": [20.0], + "minimum_sea_level_pressure_hpa": [990.0], + "maximum_sustained_wind_speed_knots": [20.0], + } + ) + ) + + mock_weathernext2_model._update_cyclone_tracks( + xr.Dataset( + { + "cyclone_probability": xr.DataArray( + np.ones((1, 1, 1)), dims=("time", "lat", "lon") + ) + } + ), + OrderedDict( + { + "time": TEST_TIME, + "lead_time": np.array([np.timedelta64(6, "h")]), + } + ), + accumulate_predictions=False, + ) + + tracks = mock_weathernext2_model.cyclone_tracks + assert {"lat", "lon", "tcmsl", "tcw10m"}.issubset(tracks.columns) + np.testing.assert_allclose(tracks[["lat", "lon", "tcmsl"]], [[10.0, 20.0, 99000.0]]) + np.testing.assert_allclose(tracks["tcw10m"], [10.28888]) + + +@mock.patch("weathernext.utils.rollout.chunked_prediction", mocked_chunked_prediction) +def test_weathernext2_call_updates_cyclone_tracks(mock_weathernext2_model): + mock_weathernext2_model.track_cyclones = True + x, coords = fetch_random_input(mock_weathernext2_model) + + with mock.patch.object(mock_weathernext2_model, "_reset_cyclone_tracks") as reset: + with mock.patch.object( + mock_weathernext2_model, "_update_cyclone_tracks" + ) as update: + out, out_coords = mock_weathernext2_model(x, coords) + + assert_output(mock_weathernext2_model, out, out_coords, coords, TEST_TIME) + reset.assert_called_once_with() + update.assert_called_once() + + +@pytest.mark.parametrize( + "dc", + [ + OrderedDict({"lat": np.random.randn(9)}), + OrderedDict({"lat": np.random.randn(9), "phoo": np.random.randn(12)}), + ], +) +@pytest.mark.parametrize("device", DEVICES) +def test_weathernext2_exceptions(dc, device, mock_weathernext2_model): + model = mock_weathernext2_model.to(device) + x, coords = fetch_data( + Random(dc), + TEST_TIME, + model.input_coords()["variable"], + model.input_coords()["lead_time"], + device=device, + ) + with pytest.raises((KeyError, ValueError)): + model(x, coords) + + +@pytest.mark.package +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_weathernext2_package(device): + torch.cuda.empty_cache() + model = WeatherNext2CyclonesMini.load_model( + WeatherNext2CyclonesMini.load_default_package(), jit_compile=False + ).to(device) + + assert model.input_coords()["lat"].shape == (181,) + assert model.input_coords()["lon"].shape == (360,) + assert model.output_coords(model.input_coords())["variable"].shape == (84,) diff --git a/uv.lock b/uv.lock index 98fee8b9d..e6b1e6209 100644 --- a/uv.lock +++ b/uv.lock @@ -2318,6 +2318,13 @@ utils = [ { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-atlas' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-fcn3' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-perturbation' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-samudrace' and extra == 'extra-12-earth2studio-sfno')" }, { name = "scipy", version = "1.18.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-atlas' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-fcn3' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-perturbation' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-samudrace' and extra == 'extra-12-earth2studio-sfno')" }, ] +weathernext = [ + { name = "dm-haiku" }, + { name = "dm-tree" }, + { name = "flax" }, + { name = "jax", extra = ["cuda13"] }, + { name = "weathernext", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-atlas' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-fcn3' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-perturbation' and extra == 'extra-12-earth2studio-samudrace') or (extra == 'extra-12-earth2studio-samudrace' and extra == 'extra-12-earth2studio-sfno')" }, +] windgust-afno = [ { name = "nvidia-physicsnemo" }, ] @@ -2408,9 +2415,11 @@ requires-dist = [ { name = "dm-haiku", marker = "extra == 'all'", specifier = ">=0.0.14" }, { name = "dm-haiku", marker = "extra == 'gencast'", specifier = ">=0.0.14" }, { name = "dm-haiku", marker = "extra == 'graphcast'", specifier = ">=0.0.14" }, + { name = "dm-haiku", marker = "extra == 'weathernext'", specifier = ">=0.0.14" }, { name = "dm-tree", marker = "extra == 'all'", specifier = ">=0.1.9" }, { name = "dm-tree", marker = "extra == 'gencast'", specifier = ">=0.1.9" }, { name = "dm-tree", marker = "extra == 'graphcast'", specifier = ">=0.1.9" }, + { name = "dm-tree", marker = "extra == 'weathernext'", specifier = ">=0.1.9" }, { name = "earth2grid", marker = "extra == 'all'", git = "https://github.com/NVlabs/earth2grid.git?rev=11dcf1b0787a7eb6a8497a3a5a5e1fdcc31232d3" }, { name = "earth2grid", marker = "extra == 'cbottle'", git = "https://github.com/NVlabs/earth2grid.git?rev=11dcf1b0787a7eb6a8497a3a5a5e1fdcc31232d3" }, { name = "earth2grid", marker = "extra == 'da-healda'", git = "https://github.com/NVlabs/earth2grid.git?rev=11dcf1b0787a7eb6a8497a3a5a5e1fdcc31232d3" }, @@ -2450,6 +2459,7 @@ requires-dist = [ { name = "flax", marker = "extra == 'all'", specifier = ">=0.10.6" }, { name = "flax", marker = "extra == 'gencast'", specifier = ">=0.10.6" }, { name = "flax", marker = "extra == 'graphcast'", specifier = ">=0.10.6" }, + { name = "flax", marker = "extra == 'weathernext'", specifier = ">=0.10.6" }, { name = "fme", marker = "extra == 'ace2'", git = "https://github.com/ai2cm/ace.git?rev=e211cad3e1a5cff0fa84e8d2f0ee67042eff3d4d" }, { name = "fme", marker = "extra == 'samudrace'", specifier = "==2026.4.0" }, { name = "fsspec", specifier = ">=2024.2.0" }, @@ -2476,6 +2486,7 @@ requires-dist = [ { name = "jax", extras = ["cuda13"], marker = "extra == 'all'", specifier = ">=0.4.26" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'gencast'", specifier = ">=0.4.26" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'graphcast'", specifier = ">=0.4.26" }, + { name = "jax", extras = ["cuda13"], marker = "extra == 'weathernext'", specifier = ">=0.4.26" }, { name = "loguru" }, { name = "makani", marker = "extra == 'all'", git = "https://github.com/NVIDIA/makani.git?rev=b38fcb2799d7dbc146fa60459f3f9823394a8bf1" }, { name = "makani", marker = "extra == 'fcn3'", git = "https://github.com/NVIDIA/makani.git?rev=b38fcb2799d7dbc146fa60459f3f9823394a8bf1" }, @@ -2627,11 +2638,12 @@ requires-dist = [ { name = "weathernext", marker = "python_full_version >= '3.12' and extra == 'all'", git = "https://github.com/google-deepmind/weathernext.git?rev=9c034db1ff412d5db6cbe6bb0c5c9afc5a267719" }, { name = "weathernext", marker = "python_full_version >= '3.12' and extra == 'gencast'", git = "https://github.com/google-deepmind/weathernext.git?rev=9c034db1ff412d5db6cbe6bb0c5c9afc5a267719" }, { name = "weathernext", marker = "python_full_version >= '3.12' and extra == 'graphcast'", git = "https://github.com/google-deepmind/weathernext.git?rev=9c034db1ff412d5db6cbe6bb0c5c9afc5a267719" }, + { name = "weathernext", marker = "python_full_version >= '3.12' and extra == 'weathernext'", git = "https://github.com/google-deepmind/weathernext.git?rev=9c034db1ff412d5db6cbe6bb0c5c9afc5a267719" }, { name = "xarray", marker = "extra == 'samudrace'", specifier = ">=2025.1.0" }, { name = "xarray", extras = ["parallel"], specifier = ">=2023.1.0" }, { name = "zarr", specifier = ">=3.1.3" }, ] -provides-extras = ["ace2", "aifs", "aifs2", "aifs2ens", "aifsens", "all", "atlas", "aurora", "cbottle", "climatenet", "corrdiff", "cosmo", "cyclone", "da-cosmo", "da-healda", "da-interp", "da-stormcast", "data", "derived", "dlesym", "dlwp", "fcn", "fcn3", "fengwu", "fuxi", "gencast", "graphcast", "interp-modafno", "orbit", "pangu", "perturbation", "precip-afno", "precip-afno-v2", "samudrace", "serve", "sfno", "solarradiation-afno", "statistics", "stormcast", "stormcast-conus", "stormscope", "ucast", "utils", "windgust-afno"] +provides-extras = ["ace2", "aifs", "aifs2", "aifs2ens", "aifsens", "all", "atlas", "aurora", "cbottle", "climatenet", "corrdiff", "cosmo", "cyclone", "da-cosmo", "da-healda", "da-interp", "da-stormcast", "data", "derived", "dlesym", "dlwp", "fcn", "fcn3", "fengwu", "fuxi", "gencast", "graphcast", "interp-modafno", "orbit", "pangu", "perturbation", "precip-afno", "precip-afno-v2", "samudrace", "serve", "sfno", "solarradiation-afno", "statistics", "stormcast", "stormcast-conus", "stormscope", "ucast", "utils", "weathernext", "windgust-afno"] [package.metadata.requires-dev] build = [