From 091b812f4c7a6399098501049a84a13e18e36033 Mon Sep 17 00:00:00 2001 From: Timo Date: Sun, 5 Jul 2026 18:15:27 +0200 Subject: [PATCH 1/7] Fixes for easier weather forecasting datasets with expanded documentation --- .gitignore | 1 + graphbench/_evaluator.py | 68 +- graphbench/_loader/_loader.py | 19 +- .../_weatherforecasting_helpers/__init__.py | 6 +- .../_weatherforecasting_helpers/_losses.py | 274 ++----- .../_weatherforecasting_helpers/utils.py | 385 +++++++++ graphbench/datasets.csv | 2 +- graphbench/datasets/_weatherforecasting.py | 728 +++++++++++++++++- 8 files changed, 1239 insertions(+), 244 deletions(-) create mode 100644 graphbench/_weatherforecasting_helpers/utils.py diff --git a/.gitignore b/.gitignore index 19ab38d..5d663e0 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ train.sh /graphbench_lib.egg-info/ .DS_Store +/scripts/ \ No newline at end of file diff --git a/graphbench/_evaluator.py b/graphbench/_evaluator.py index a7f2a2e..3e1123d 100644 --- a/graphbench/_evaluator.py +++ b/graphbench/_evaluator.py @@ -1,4 +1,4 @@ -from typing import Callable, Optional, Union, TypeAlias +from typing import Callable, Dict, Optional, Union, TypeAlias import numpy as np import torch @@ -17,13 +17,11 @@ from graphbench._metadata import get_master_df from graphbench._helpers import VectorizedCircuitSimulator from graphbench._weatherforecasting_helpers import ( - compute_latitude_weights, - compute_pressure_level_weights, - get_default_pressure_levels, - get_variable_weights, - masked_loss, + compute_weather_metric_breakdown, ) +from graphbench.datasets._weatherforecasting import _load_static_files + # note that this is not the same having the Union inside the Callable, which would be too permissive _Metric: TypeAlias = Union[ @@ -172,7 +170,7 @@ def _get_metric_from_name(self, metric_name: str) -> _Metric: 'RMSE': self.get_rmse(), 'RSE': self.get_rse(), 'ChipDesignScore': self.get_chip_design_score(), - 'Weather_MSE': self.get_weather_mse(), + 'Weather_MSE': self.get_weather_metric_breakdown(), 'ClosedGap': self.get_closed_gap(), 'MisSize': self.get_mis_size(), 'MaxCutSize': self.get_max_cut_size(), @@ -200,6 +198,7 @@ def evaluate( y_pred: Union[Tensor, np.ndarray, list[Data]], y_true: Optional[Union[Tensor, np.ndarray, list[Data]]] = None, batch: Optional[Batch] = None, + root: Optional[str] = None, ) -> Union[float, list[float]]: """ Computes the selected metric(s) for the given predictions and true values. @@ -212,6 +211,7 @@ def evaluate( y_pred: Predicted values as a torch tensor or numpy array of shape (N,K) y_true: True values as a torch tensor or numpy array of shape (N,K) or (N,1), defaults to None batch: Optional batch information for unsupervised tasks, defaults to None + root: Root folder for data downloads, used for weatherforecasting data to obtain variable information. """ metric = self._get_metric() if batch is not None: @@ -220,8 +220,12 @@ def evaluate( return [met(y_pred, batch).item() for met in metric] return metric(y_pred, batch).item() - y_pred, y_true = self._check_input(y_pred, y_true) + y_pred, y_true = self._check_input(y_pred, y_true) + if root is not None: + if isinstance(metric, list): + return [met(y_pred, y_true, root) for met in metric] + return metric(y_pred, y_true, root) if isinstance(metric, list): return [met(y_pred, y_true).item() for met in metric] return metric(y_pred, y_true).item() @@ -286,13 +290,6 @@ def get_chip_design_score(self) -> Callable[[list[Data], list[Data]], Tensor]: """ return lambda y_pred, y_true: self._get_chip_design_score(y_pred, y_true) - def get_weather_mse(self) -> Callable[[Tensor, Tensor], Tensor]: - """ - Returns: - A callable that takes ``y_pred, y_true`` and computes the mean squared error (MSE) for weather - forecasting data. - """ - return lambda y_pred, y_true: self._get_weather_mse(y_pred, y_true) def get_mis_size(self) -> Callable[[Tensor, Batch], Tensor]: """ @@ -468,26 +465,29 @@ def _equivalence_score( # No match return 0.0 - # TODO I annotated y_true as Tensor, because masked_loss expects a Tensor. However, in the comment it looks like - # y_true needs to be a Data object. To me it looks like the implementation is simply incorrect. - def _get_weather_mse(self, y_pred: Tensor, y_true: Tensor) -> Tensor: - grid_variables = [ - '2m_temperature', 'mean_sea_level_pressure', '10m_v_component_of_wind', - '10m_u_component_of_wind', 'total_precipitation_6hr', 'temperature', - 'geopotential', 'u_component_of_wind', 'v_component_of_wind', - 'vertical_velocity', 'specific_humidity' - ] + def get_weather_metric_breakdown(self) -> Callable[[Tensor, Tensor, str], Dict[str, float]]: + """ + Returns: + A callable that takes ``y_pred, y_true`` and returns a dict of per-variable-group MSE + values, plus per-pressure-level MSE for multi-level variables (e.g. temperature at + each pressure level), for weather forecasting data. + """ + return lambda y_pred, y_true, root: self._get_weather_metric_breakdown(y_pred, y_true, root) + + def _get_weather_metric_breakdown(self, y_pred: Tensor, y_true: Tensor, root: str) -> Dict[str, float]: + #first get metadata information about slices and variables + + metadata, _ = _load_static_files(root) + + feature_group_slices = metadata['variable_channel_slices'] + feature_group_names = metadata['grid_variables'] - # assuming y_true is the data object and not only the prediction tensor - # TODO: change format of y_true in evaluator call if needed - return masked_loss( - predictions=y_pred, - targets=y_true, - variable_slices=None, - variable_weights=get_variable_weights(grid_variables), - variable_names=grid_variables, - latitude_weights=compute_latitude_weights(y_true.grid_lat), - pressure_level_weights=compute_pressure_level_weights(get_default_pressure_levels()), + #TODO get correct metadata and spatially weighting for evaluation + return compute_weather_metric_breakdown( + y_pred=y_pred, + y_true=y_true, + feature_group_slices=feature_group_slices, + feature_group_names=feature_group_names, ) def get_mse(self) -> Callable[[Tensor, Tensor], Tensor]: diff --git a/graphbench/_loader/_loader.py b/graphbench/_loader/_loader.py index 6d67543..91020ba 100644 --- a/graphbench/_loader/_loader.py +++ b/graphbench/_loader/_loader.py @@ -10,7 +10,7 @@ from graphbench._metadata import expand_dataset_names from ._dataset_registry import DatasetRegistry from ._split_strategies import AlgoReasSplitStrategy, FixedSplitStrategy, RatioSplitStrategy, TrainValTestSet - +from graphbench.datasets._weatherforecasting import _prepare_weather_cache_once, EfficientWeatherGraphDataset class Loader(): @@ -252,15 +252,22 @@ def _make_weather_dataset( split: str, name_override: Optional[str] = None, ) -> InMemoryDataset: - from graphbench.datasets import WeatherforecastingDataset + #from graphbench.datasets import WeatherforecastingDataset + + timeout_s = 3600 + skip_cache_build = False + #hier noch anpassen auf unsere konvention + use_prebuilt = True + + _prepare_weather_cache_once(self.root, timeout_s=timeout_s, skip_cache_build=skip_cache_build, task_name=name_override or dataset_name, use_prebuilt=use_prebuilt) + #dataset = EfficientWeatherGraphDataset(root=root, pre_transform=None, transform=None) - return WeatherforecastingDataset( + #hier noch den return anpassen + return EfficientWeatherGraphDataset( root=self.root, - name=name_override or dataset_name, pre_filter=self.pre_filter, pre_transform=self.pre_transform, - transform=self.transform, - split=split, + transform=self.transform ) def _make_co_dataset( diff --git a/graphbench/_weatherforecasting_helpers/__init__.py b/graphbench/_weatherforecasting_helpers/__init__.py index 3b4d6c5..0c089da 100644 --- a/graphbench/_weatherforecasting_helpers/__init__.py +++ b/graphbench/_weatherforecasting_helpers/__init__.py @@ -1,16 +1,14 @@ from ._losses import ( compute_latitude_weights, compute_pressure_level_weights, + compute_weather_metric_breakdown, get_default_pressure_levels, - get_variable_weights, - masked_loss, ) __all__ = [ "compute_latitude_weights", "compute_pressure_level_weights", + "compute_weather_metric_breakdown", "get_default_pressure_levels", - "get_variable_weights", - "masked_loss", ] diff --git a/graphbench/_weatherforecasting_helpers/_losses.py b/graphbench/_weatherforecasting_helpers/_losses.py index fd9344a..92e6703 100644 --- a/graphbench/_weatherforecasting_helpers/_losses.py +++ b/graphbench/_weatherforecasting_helpers/_losses.py @@ -1,9 +1,4 @@ -""" -Weighted loss function for the weather prediction model. - -This module provides GPU-aware loss functions that automatically detect and use -the same device as the input tensors, making them work seamlessly on both CPU and GPU. -""" +"""Weighted MSE loss and metric-breakdown helpers for weather forecasting.""" from typing import Dict, List, Optional, Tuple @@ -11,209 +6,98 @@ import torch +"""def _weather_loss_sum(y_pred, y_true, grid_lat): + device = y_pred.device + if y_pred.dim() == 2: + y_pred = y_pred.unsqueeze(0) + y_true = y_true.unsqueeze(0) + mask = torch.isfinite(y_true) + if not mask.any(): + return torch.zeros((), device=device, dtype=y_pred.dtype) + y_pred = torch.where(mask, y_pred, torch.zeros_like(y_pred)) + y_true = torch.where(mask, y_true, torch.zeros_like(y_true)) + lat_rad = torch.tensor(np.deg2rad(grid_lat), dtype=torch.float32, device=device) + lat_w = torch.cos(lat_rad) + lat_w = lat_w / lat_w.mean() + num_nodes = y_pred.shape[1] + num_grids = num_nodes // lat_w.shape[0] + lat_w = lat_w.repeat(num_grids) + sq_err = (y_pred - y_true) ** 2 + return (sq_err * lat_w.view(1, -1, 1)).mean()""" + def compute_latitude_weights(latitude_values: np.ndarray) -> torch.Tensor: - """Compute latitude-based area weights for grid cells.""" - # converting it to radians and computing the cos(lat) weights + """Compute cos(lat) area weights for grid cells, normalised to mean 1.""" lat_rad = np.deg2rad(latitude_values) weights = np.cos(lat_rad) - weights = weights / np.mean(weights) - return torch.tensor(weights, dtype=torch.float32) def compute_pressure_level_weights(pressure_levels: np.ndarray) -> torch.Tensor: - """Compute pressure-level weights using proportio to pressure level.""" - + """Compute pressure-level weights proportional to the pressure level, normalised to mean 1.""" weights = pressure_levels / np.mean(pressure_levels) - return torch.tensor(weights, dtype=torch.float32) -def _spatially_weighted_mse( - predictions: torch.Tensor, - targets: torch.Tensor, - latitude_weights: torch.Tensor, - variable_weights: Optional[Dict[str, float]] = None, - variable_slices: Optional[List[Tuple[int, int]]] = None, - variable_names: Optional[List[str]] = None, - pressure_level_weights: Optional[torch.Tensor] = None, - eval_mode: bool = False -) -> torch.Tensor: - """Compute spatially weighted MSE loss with pressure level weighting.""" - - - device = predictions.device - - predictions = predictions.to(device) - targets = targets.to(device) - latitude_weights = latitude_weights.to(device) - if pressure_level_weights is not None: - pressure_level_weights = pressure_level_weights.to(device) - #print(pressure_level_weights) - # handle batch dimension - if predictions.dim() == 3: - batch_size = predictions.shape[0] - num_nodes = predictions.shape[1] - num_features = predictions.shape[2] - - predictions_flat = predictions.view(batch_size, -1) - targets_flat = targets.view(batch_size, -1) - - lat_weights_expanded = latitude_weights.unsqueeze(1).expand(-1, num_features).contiguous().view(-1) - - # creating pressure level weights - if pressure_level_weights is not None and variable_slices is not None and variable_names is not None: - pressure_weights_expanded = torch.ones(num_nodes * num_features, device=device, dtype=torch.float32) - - for (start, end), var_name in zip(variable_slices, variable_names): - if var_name in variable_weights: # multi-level variable - var_channels = end - start - if var_channels > 1: # multi-level variable - # applying pressure level weights to this variable's channels - for level_idx in range(var_channels): - feature_idx = start + level_idx - # expanding pressure weight to all spatial locations for this feature - pressure_weights_expanded[feature_idx::num_features] = pressure_level_weights[level_idx] - else: - pressure_weights_expanded = torch.ones_like(lat_weights_expanded) - - sq_error = (predictions_flat - targets_flat) ** 2 - - # applying spatial and pressure level weighting - combined_weights = lat_weights_expanded.unsqueeze(0) * pressure_weights_expanded.unsqueeze(0) - weighted_sq_error = sq_error * combined_weights - - # applying variable-level weighting - if variable_weights is not None and variable_slices is not None and variable_names is not None: - var_weighted_error = torch.zeros_like(weighted_sq_error) - - for (start, end), var_name in zip(variable_slices, variable_names): - if var_name in variable_weights: - weight = variable_weights[var_name] - # applying the weight to all grid points for this variable - var_weighted_error[:, start:end] = weighted_sq_error[:, start:end] * weight - else: - var_weighted_error[:, start:end] = weighted_sq_error[:, start:end] - - weighted_sq_error = var_weighted_error - - loss = weighted_sq_error.mean() - - else: - num_nodes = predictions.shape[0] - num_features = predictions.shape[1] - #print(latitude_weights.shape) - # creating latitude weights for all the features - - lat_weights_expanded = latitude_weights.unsqueeze(1).repeat(1,num_features).repeat(64,1) - #lat_weights_expanded = latitude_weights.unsqueeze(1).expand(-1, num_features).contiguous().view(-1) - #rint(lat_weights_expanded.shape) - #print(pressure_level_weights) - - var_weights = torch.ones((11,1), device=device, dtype=torch.float32) - for i, var_val in enumerate(variable_weights.values()): - var_weights[i,:] = var_val - #print(var_weights) - pressure_weights = pressure_level_weights.repeat(6) - - pressure_weights_expanded = torch.ones((num_features,1), device=device, dtype=torch.float32) - - pressure_weights_expanded[:5,:] = var_weights[:5] - - pressure_weights_expanded[5:,:] = pressure_weights_expanded[5:,:] * pressure_weights.unsqueeze(1) - #print(pressure_weights_expanded.shape) - pressure_weights_expanded = pressure_weights_expanded.transpose(0,1).repeat(2048,1) - - - - - sq_error = (predictions - targets) ** 2 - #print("sq_error") - #print(sq_error.shape) - # applying spatial and pressure level weighting - if batch_size is None: - combined_weights = (lat_weights_expanded * pressure_weights_expanded).repeat(4,1) - else: - combined_weights = (lat_weights_expanded * pressure_weights_expanded).repeat(batch_size,1) - #print(combined_weights) - #print(combined_weights.shape) - if eval_mode: - weighted_sq_error = sq_error - else: - weighted_sq_error = sq_error * combined_weights - - - - loss = weighted_sq_error.mean(0).sum() - - return loss - - def get_default_pressure_levels() -> np.ndarray: """Get default pressure levels matching the GraphCast configuration.""" return np.array([50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000], dtype=np.float32) -def get_variable_weights(grid_variables: List[str]) -> Dict[str, float]: - """Get variable weights.""" - # we're using GraphCast's variable weights here - graphcast_weights = { - # single-level variables have lower weights - '10m_u_component_of_wind': 0.1, - '10m_v_component_of_wind': 0.1, - 'mean_sea_level_pressure': 0.1, - 'total_precipitation_6hr': 0.1, - - # 2m temperature is important though, thus gets a higher weighting - '2m_temperature': 1.0, - - # multi-level atmospheric variables have a default to 1.0 - 'temperature': 1.0, - 'geopotential': 1.0, - 'u_component_of_wind': 1.0, - 'v_component_of_wind': 1.0, - 'vertical_velocity': 1.0, - 'specific_humidity': 1.0, - } - - return {var: graphcast_weights.get(var, 1.0) for var in grid_variables} - - -def masked_loss( - predictions: torch.Tensor, - targets: torch.Tensor, - latitude_weights: torch.Tensor, - variable_weights: Optional[Dict[str, float]] = None, - variable_slices: Optional[List[Tuple[int, int]]] = None, - variable_names: Optional[List[str]] = None, - pressure_level_weights: Optional[torch.Tensor] = None, - device: Optional[str] = None, - eval_mode: bool = False, -) -> torch.Tensor: - """Computes spatially weighted MSE loss with NaN/Inf masking.""" - - if device is None: - device = predictions.device - - # handling nan/inf values by masking - mask = torch.isfinite(targets) - if not mask.any(): - return torch.zeros((), device=device, dtype=predictions.dtype) - - # replacing nan/inf with zeros for computation - pred_clean = torch.where(mask, predictions, torch.zeros_like(predictions)) - target_clean = torch.where(mask, targets, torch.zeros_like(targets)) - - loss = _spatially_weighted_mse( - predictions=pred_clean, - targets=target_clean, - latitude_weights=latitude_weights, - variable_weights=variable_weights, - variable_slices=variable_slices, - variable_names=variable_names, - pressure_level_weights=pressure_level_weights, - eval_mode=eval_mode - ) - - return loss +def compute_weather_metric_breakdown( + y_pred: torch.Tensor, + y_true: torch.Tensor, + feature_group_slices: List[Tuple[int, int]], + feature_group_names: List[str], + multi_level_variable_names: Optional[set] = None, + pressure_levels: Optional[List[float]] = None, +) -> Dict[str, float]: + """Compute per-variable-group MSE, and per-pressure-level MSE for multi-level variables. + + NaN/Inf targets are masked out of the mean rather than replaced with zero, so + padded or missing grid points don't bias the reported MSE. + + Args: + y_pred: Predictions [..., num_features]. + y_true: Targets, same shape as `y_pred`. + feature_group_slices: [start, end) channel slice per variable group. + feature_group_names: Variable name per slice in `feature_group_slices`. + multi_level_variable_names: Names of variables to additionally break down + by pressure level. Defaults to the standard GraphCast multi-level variables. + pressure_levels: Pressure level values corresponding to each channel within a + multi-level variable's slice. Defaults to `get_default_pressure_levels()`. + + Returns: + Dict mapping ``"{name}_mse"`` and ``"{name}_z{level}_mse"`` to scalar MSE values + (``nan`` for groups/levels with no finite targets). + """ + if multi_level_variable_names is None: + multi_level_variable_names = { + "specific_humidity", "vertical_velocity", "u_component_of_wind", + "v_component_of_wind", "geopotential", "temperature", + } + if pressure_levels is None: + pressure_levels = get_default_pressure_levels() + + mask = torch.isfinite(y_true) + sq_err = (y_pred - y_true) ** 2 + reduce_dims = tuple(range(sq_err.dim() - 1)) if sq_err.dim() > 1 else () + masked_sq_err = torch.where(mask, sq_err, torch.zeros_like(sq_err)) + feat_sum = masked_sq_err.sum(dim=reduce_dims).detach().cpu().to(torch.float64).numpy() + feat_count = mask.sum(dim=reduce_dims).detach().cpu().to(torch.float64).numpy() + + def _mse(sum_val: float, count_val: float) -> float: + return float(sum_val / count_val) if count_val > 0 else float("nan") + + metrics: Dict[str, float] = {} + for name, (start, end) in zip(feature_group_names, feature_group_slices): + start, end = int(start), int(end) + metrics[f"{name}_mse"] = _mse(feat_sum[start:end].sum(), feat_count[start:end].sum()) + + if name in multi_level_variable_names: + levels_count = end - start + for i in range(min(len(pressure_levels), levels_count)): + level = pressure_levels[i] + metrics[f"{name}_z{level}_mse"] = _mse(feat_sum[start + i], feat_count[start + i]) + + return metrics diff --git a/graphbench/_weatherforecasting_helpers/utils.py b/graphbench/_weatherforecasting_helpers/utils.py new file mode 100644 index 0000000..c502c87 --- /dev/null +++ b/graphbench/_weatherforecasting_helpers/utils.py @@ -0,0 +1,385 @@ +"""Spatial features and dataset utilities for grid-mesh weather forecasting graphs. + +Provides coordinate transformations, spatial feature computation, temporal dataset +splitting, and a PyG InMemoryDataset for GraphCast-style weather models. +""" + +import numpy as np +from typing import Tuple, Optional +import torch +from pathlib import Path +import pickle +from typing import List, Dict, Any, Optional, Tuple, Sequence +from torch.utils.data import Subset +import xarray as xr +try: + from tqdm.auto import tqdm +except Exception: + def tqdm(iterable, *args, **kwargs): + return iterable +from torch_geometric.data import Data, Dataset, InMemoryDataset +from dataclasses import dataclass + + +def get_weather_variables() -> List[str]: + return [ + '2m_temperature', 'mean_sea_level_pressure', '10m_v_component_of_wind', + '10m_u_component_of_wind', 'total_precipitation_6hr', 'temperature', + 'geopotential', 'u_component_of_wind', 'v_component_of_wind', + 'vertical_velocity', 'specific_humidity' + ] + +def get_pressure_levels() -> List[int]: + return [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000] + +def lat_lon_deg_to_spherical(lat: np.ndarray, lon: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Convert lat/lon in degrees to spherical angles (phi, theta).""" + phi = np.deg2rad(lon) + theta = np.deg2rad(90 - lat) + return phi, theta + + +def spherical_to_cartesian(phi: np.ndarray, theta: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Map spherical angles to unit-sphere Cartesian coordinates (x, y, z).""" + x = np.cos(phi) * np.sin(theta) + y = np.sin(phi) * np.sin(theta) + z = np.cos(theta) + return x, y, z + + +def get_relative_positions(senders_pos: np.ndarray, receivers_pos: np.ndarray) -> np.ndarray: + """Return receiver_pos - sender_pos for each edge.""" + return receivers_pos - senders_pos + + +def get_graph_spatial_features( + node_lat: np.ndarray, + node_lon: np.ndarray, + senders: np.ndarray, + receivers: np.ndarray, + add_node_positions: bool = False, + add_node_latitude: bool = True, + add_node_longitude: bool = True, + add_relative_positions: bool = True, + edge_normalization_factor: Optional[float] = None, + sine_cosine_encoding: bool = False, + encoding_num_frequencies: int = 10, + encoding_multiplier: float = 1.2, +) -> Tuple[np.ndarray, np.ndarray]: + """Compute node and edge spatial features for a homogeneous graph. + + Node features: optional Cartesian position, cos(lat), cos/sin(lon). + Edge features: L2-normalised relative position vector (and its norm) in 3-D space. + Optionally applies a multi-frequency sine/cosine encoding to both. + """ + num_nodes = node_lat.shape[0] + num_edges = senders.shape[0] + dtype = node_lat.dtype + node_phi, node_theta = lat_lon_deg_to_spherical(node_lat, node_lon) + + node_features = [] + if add_node_positions: + node_features.extend(spherical_to_cartesian(node_phi, node_theta)) + if add_node_latitude: + # cos(theta): 1 at north pole, -1 at south pole + node_features.append(np.cos(node_theta)) + if add_node_longitude: + node_features.append(np.cos(node_phi)) + node_features.append(np.sin(node_phi)) + + if not node_features: + node_features = np.zeros([num_nodes, 0], dtype=dtype) + else: + node_features = np.stack(node_features, axis=-1) + + edge_features = [] + if add_relative_positions: + relative_position = get_relative_position_in_receiver_local_coordinates( + node_phi=node_phi, + node_theta=node_theta, + senders=senders, + receivers=receivers, + latitude_local_coordinates=None, + longitude_local_coordinates=None, + ) + # L2 distance in 3-D space (not geodesic) + relative_edge_distances = np.linalg.norm(relative_position, axis=-1, keepdims=True) + if edge_normalization_factor is None: + edge_normalization_factor = relative_edge_distances.max() + edge_features.append(relative_edge_distances / edge_normalization_factor) + edge_features.append(relative_position / edge_normalization_factor) + + if not edge_features: + edge_features = np.zeros([num_edges, 0], dtype=dtype) + else: + edge_features = np.concatenate(edge_features, axis=-1) + + if sine_cosine_encoding: + def sine_cosine_transform(x: np.ndarray) -> np.ndarray: + freqs = encoding_multiplier**np.arange(encoding_num_frequencies) + phases = freqs * x[..., None] + x_cat = np.concatenate([np.sin(phases), np.cos(phases)], axis=-1) + return x_cat.reshape([x.shape[0], -1]) + + node_features = sine_cosine_transform(node_features) + edge_features = sine_cosine_transform(edge_features) + + return node_features, edge_features + +def get_bipartite_graph_spatial_features( + senders_node_lat: np.ndarray, + senders_node_lon: np.ndarray, + receivers_node_lat: np.ndarray, + receivers_node_lon: np.ndarray, + senders: np.ndarray, + receivers: np.ndarray, + add_node_positions: bool = False, + add_node_latitude: bool = True, + add_node_longitude: bool = True, + add_relative_positions: bool = True, + edge_normalization_factor: Optional[float] = None, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Compute spatial features for a bipartite graph (e.g. grid-to-mesh). + + Returns separate feature arrays for sender nodes, receiver nodes, and edges. + Edge features are L2-normalised relative positions in 3-D Cartesian space. + """ + num_senders = senders_node_lat.shape[0] + num_receivers = receivers_node_lat.shape[0] + num_edges = senders.shape[0] + dtype = senders_node_lat.dtype + + senders_phi, senders_theta = lat_lon_deg_to_spherical(senders_node_lat, senders_node_lon) + receivers_phi, receivers_theta = lat_lon_deg_to_spherical(receivers_node_lat, receivers_node_lon) + + # --- sender node features --- + senders_features = [] + if add_node_positions: + senders_features.extend(spherical_to_cartesian(senders_phi, senders_theta)) + if add_node_latitude: + senders_features.append(np.cos(senders_theta)) + if add_node_longitude: + senders_features.append(np.cos(senders_phi)) + senders_features.append(np.sin(senders_phi)) + senders_features = ( + np.zeros([num_senders, 0], dtype=dtype) if not senders_features + else np.stack(senders_features, axis=-1) + ) + + # --- receiver node features --- + receivers_features = [] + if add_node_positions: + receivers_features.extend(spherical_to_cartesian(receivers_phi, receivers_theta)) + if add_node_latitude: + receivers_features.append(np.cos(receivers_theta)) + if add_node_longitude: + receivers_features.append(np.cos(receivers_phi)) + receivers_features.append(np.sin(receivers_phi)) + receivers_features = ( + np.zeros([num_receivers, 0], dtype=dtype) if not receivers_features + else np.stack(receivers_features, axis=-1) + ) + + # --- edge features --- + edge_features = [] + if add_relative_positions: + #TODO: fix this + relative_positions = get_bipartite_relative_position_in_receiver_local_coordinates( + senders_node_phi=senders_phi, + senders_node_theta=senders_theta, + senders=senders, + receivers_node_phi=receivers_phi, + receivers_node_theta=receivers_theta, + receivers=receivers, + latitute_local_coordinates=None, + longitude_local_coordinates=None, + ) + distances = np.linalg.norm(relative_positions, axis=-1, keepdims=True) + if edge_normalization_factor is None: + edge_normalization_factor = distances.max() + edge_features.append(distances / edge_normalization_factor) + edge_features.append(relative_positions / edge_normalization_factor) + + edge_features = ( + np.zeros([num_edges, 0], dtype=dtype) if not edge_features + else np.concatenate(edge_features, axis=-1) + ) + + return senders_features, receivers_features, edge_features + +def get_bipartite_relative_position_in_receiver_local_coordinates( + senders_node_phi, + senders_node_theta, + senders, + receivers_node_phi, + receivers_node_theta, + receivers, + latitute_local_coordinates, + longitude_local_coordinates, +): + """Return sender_pos - receiver_pos in 3-D Cartesian space for each bipartite edge. + + Local-coordinate rotation is not yet implemented; pass None for both coordinate args. + """ + s_x, s_y, s_z = spherical_to_cartesian(senders_node_phi, senders_node_theta) + r_x, r_y, r_z = spherical_to_cartesian(receivers_node_phi, receivers_node_theta) + senders_node_pos = np.stack([s_x, s_y, s_z], axis=-1) + receivers_node_pos = np.stack([r_x, r_y, r_z], axis=-1) + if not (latitute_local_coordinates or longitude_local_coordinates): + return senders_node_pos[senders.cpu()] - receivers_node_pos[receivers.cpu()] + +def get_relative_position_in_receiver_local_coordinates( + node_phi, + node_theta, + senders, + receivers, + latitude_local_coordinates, + longitude_local_coordinates, +): + """Return sender_pos - receiver_pos in 3-D Cartesian space for each homogeneous edge. + + Local-coordinate rotation is not yet implemented; pass None for both coordinate args. + """ + x, y, z = spherical_to_cartesian(node_phi, node_theta) + node_pos = np.stack([x, y, z], axis=-1) + if not (latitude_local_coordinates or longitude_local_coordinates): + return node_pos[senders.cpu()] - node_pos[receivers.cpu()] + +@dataclass(frozen=True) +class TemporalSplits: + """Immutable container for train/val/test index arrays.""" + train_idx: np.ndarray + val_idx: np.ndarray + test_idx: np.ndarray + + def as_slices(self) -> Tuple[slice, slice, slice]: + """Convert index arrays to contiguous slices (assumes sorted, contiguous indices).""" + def to_slice(indices: np.ndarray) -> slice: + if indices.size == 0: + return slice(0, 0) + return slice(int(indices[0]), int(indices[-1]) + 1) + + return to_slice(self.train_idx), to_slice(self.val_idx), to_slice(self.test_idx) + + +def _find_first_index_at_or_after(datetimes: np.ndarray, threshold: np.datetime64) -> int: + # datetimes must be sorted ascending + return int(np.searchsorted(datetimes, threshold, side="left")) + + +def compute_temporal_splits( + datetimes: Sequence[np.datetime64], +) -> TemporalSplits: + """Split datetimes into train (1979-2015), val (2016-2017), test (2018-2021). + + Falls back to 80/10/10 proportional split when any year-based split is empty. + """ + return compute_fixed_year_splits( + datetimes, + train_years=(1979, 2015), + val_years=(2016, 2017), + test_years=(2018, 2021), + ) + + +def slice_sequence_by_indices(sequence: Sequence, indices: np.ndarray) -> List: + """Select elements from a sequence by integer index array.""" + return [sequence[int(i)] for i in indices] + + +def compute_fixed_year_splits( + datetimes: Sequence[np.datetime64], + *, + train_years: tuple[int, int] = (1979, 2015), + val_years: tuple[int, int] = (2016, 2017), + test_years: tuple[int, int] = (2018, 2021), +) -> TemporalSplits: + """Split a datetime sequence by calendar-year ranges. + + Sorts unsorted input internally and maps indices back to the original order. + Falls back to 80/10/10 proportional split when any year-based split is empty. + """ + dts = np.asarray(datetimes) + if dts.ndim != 1 or dts.size == 0: + return TemporalSplits( + train_idx=np.arange(dts.size, dtype=np.int64), + val_idx=np.array([], dtype=np.int64), + test_idx=np.array([], dtype=np.int64), + ) + + if not np.all(dts[:-1] <= dts[1:]): + order = np.argsort(dts, kind="stable") + dts_sorted = dts[order] + else: + order = None + dts_sorted = dts + + years = dts_sorted.astype('datetime64[Y]').astype(int) + 1970 + + def _range_mask(yrs: np.ndarray, yr_range: tuple[int, int]) -> np.ndarray: + start, end = int(yr_range[0]), int(yr_range[1]) + return (yrs >= start) & (yrs <= end) + + train_mask = _range_mask(years, train_years) + val_mask = _range_mask(years, val_years) + test_mask = _range_mask(years, test_years) + + train_idx_sorted = np.nonzero(train_mask)[0].astype(np.int64) + val_idx_sorted = np.nonzero(val_mask)[0].astype(np.int64) + test_idx_sorted = np.nonzero(test_mask)[0].astype(np.int64) + + # Proportional fallback when any split is empty + if ( + train_idx_sorted.size < 1 + or val_idx_sorted.size < 1 + or test_idx_sorted.size < 1 + ): + n = dts.size + n_train = max(int(0.8 * n), 1) + n_val = max(int(0.1 * n), 1) + n_test = max(n - n_train - n_val, 1) + if n_train + n_val + n_test > n: + n_test = n - n_train - n_val + train_idx_sorted = np.arange(0, n_train, dtype=np.int64) + val_idx_sorted = np.arange(n_train, n_train + n_val, dtype=np.int64) + test_idx_sorted = np.arange(n_train + n_val, n, dtype=np.int64) + + if order is not None: + def backmap(sorted_idx: np.ndarray) -> np.ndarray: + return order[sorted_idx] + + train_idx = backmap(train_idx_sorted) + val_idx = backmap(val_idx_sorted) + test_idx = backmap(test_idx_sorted) + train_idx.sort(); val_idx.sort(); test_idx.sort() + else: + train_idx, val_idx, test_idx = train_idx_sorted, val_idx_sorted, test_idx_sorted + + return TemporalSplits(train_idx=train_idx, val_idx=val_idx, test_idx=test_idx) + +def load_statistics(path: str, pressure_levels, grid_variables) -> Dict[str, Any]: + """Load per-variable normalisation statistics from xarray NetCDF files. + + Returns (mean, std, diff_std) tensors shaped [num_features], where features + are ordered as: surface variables first, then pressure-level variables × levels. + """ + mean = xr.open_dataset(f"{path}/stats-mean_by_level.nc").load() + std = xr.open_dataset(f"{path}/stats-stddev_by_level.nc").load() + diff_std = xr.open_dataset(f"{path}/stats-diffs_stddev_by_level.nc").load() + + mean_list, std_list, diff_std_list = [], [], [] + + # First 5 variables are surface-level (no pressure dimension) + for grid_var in grid_variables[:5]: + mean_list.append(torch.from_numpy(mean[grid_var].values).reshape(1)) + std_list.append(torch.from_numpy(std[grid_var].values).reshape(1)) + diff_std_list.append(torch.from_numpy(diff_std[grid_var].values).reshape(1)) + + # Remaining variables have one entry per pressure level + for grid_var in grid_variables[5:]: + for level in pressure_levels: + mean_list.append(torch.from_numpy(mean[grid_var].sel(level=level).values).reshape(1)) + std_list.append(torch.from_numpy(std[grid_var].sel(level=level).values).reshape(1)) + diff_std_list.append(torch.from_numpy(diff_std[grid_var].sel(level=level).values).reshape(1)) + + return torch.cat(mean_list), torch.cat(std_list), torch.cat(diff_std_list) \ No newline at end of file diff --git a/graphbench/datasets.csv b/graphbench/datasets.csv index 227f29c..68c10d0 100644 --- a/graphbench/datasets.csv +++ b/graphbench/datasets.csv @@ -2,7 +2,7 @@ dataset_name, datasets algorithmic_reasoning_easy,"algoreas_bipartitematching_easy;algoreas_bridges_easy;algoreas_mst_easy;algoreas_steinertree_easy;algoreas_maxclique_easy;algoreas_flow_easy;algoreas_topologicalorder_easy" algorithmic_reasoning_medium,"algoreas_bipartitematching_medium;algoreas_bridges_medium;algoreas_mst_medium;algoreas_steinertree_medium;algoreas_maxclique_medium;algoreas_flow_medium;algoreas_topologicalorder_medium" algorithmic_reasoning_hard,"algoreas_bipartitematching_hard;algoreas_bridges_hard;algoreas_mst_hard;algoreas_steinertree_hard;algoreas_maxclique_hard;algoreas_flow_hard;algoreas_topologicalorder_hard" -weather,"weather_64" +weather,"weather" co,"co_ba_large;co_ba_small;co_er_large;co_er_small;co_rb_large;co_rb_small" sat,"sat_lcg_as;sat_lcg_epm;sat_vcg_as;sat_vcg_epm;sat_vg_as;sat_vg_epm" socialnetwork,"bluesky_quotes;bluesky_replies;bluesky_retweets" diff --git a/graphbench/datasets/_weatherforecasting.py b/graphbench/datasets/_weatherforecasting.py index 0fcb98a..5bc0cf5 100644 --- a/graphbench/datasets/_weatherforecasting.py +++ b/graphbench/datasets/_weatherforecasting.py @@ -14,11 +14,20 @@ from typing import Callable, Dict, List, Literal, Optional, Union from torch_geometric.data import Data - +from loguru import logger from graphbench._helpers import download_and_unpack, SourceSpec, get_logger from ._base import GraphDataset - - +from torch_geometric.data import Data, Dataset, InMemoryDataset +from dataclasses import dataclass +import torch +import pickle +import xarray as xr +try: + from tqdm.auto import tqdm +except Exception: + def tqdm(iterable, *args, **kwargs): + return iterable +from torch.utils.data import Subset # (i) helper functions # -----------------------------------------------------------------------------# @@ -28,8 +37,716 @@ _logger = get_logger(__name__) + +# HuggingFace download URLs for prebuilt processed files. +_PREBUILT_PT_URLS = { + "weather": "https://huggingface.co/datasets/log-rwth-aachen/Graphbench_Weather/resolve/main/weather_64_combined.pt", + "weather_subset": "https://huggingface.co/datasets/log-rwth-aachen/GraphBench_Weather_Subset/resolve/main/weather_64_combined.pt", +} + +# HuggingFace repo IDs for raw component files. +_HF_WEATHER_REPOS = { + "weather": "log-rwth-aachen/Graphbench_Weather", + "weather_subset": "log-rwth-aachen/GraphBench_Weather_Subset", +} + +# Static component files that are shared across tasks and always sourced from +# the original weather repo (log-rwth-aachen/Graphbench_Weather). +_HF_WEATHER_STATIC_FILES = { + "metadata.pkl": "metadata.pkl", + "static_components.pkl": "static_components.pkl", + "stats-mean_by_level.nc": "stats-mean_by_level.nc", + "stats-stddev_by_level.nc": "stats-stddev_by_level.nc", + "stats-diffs_stddev_by_level.nc": "stats-diffs_stddev_by_level.nc", +} + +# Data files specific to each task (sourced from the task-specific repo). +# weather_64.pt is the raw timestep data; saved as weather_dataset.pt which is +# what EfficientWeatherGraphDataset.process() expects. +_HF_WEATHER_DATA_FILES = { + "weather_64.pt": "weather_dataset.pt", +} + +# Combined view for the full weather task (backward-compat). +_HF_WEATHER_RAW_FILES = {**_HF_WEATHER_STATIC_FILES, **_HF_WEATHER_DATA_FILES} + + + +def _load_static_files(root: str): + root = Path(root) + metadata_path = root /"weather" /"metadata.pkl" + if metadata_path.exists(): + with open(metadata_path, 'rb') as f: + metadata = pickle.load(f) + else: + + raise FileNotFoundError(f"Missing metadata.pkl in {metadata_path}. Please ensure the weather dataset is downloaded and processed correctly.") + + static_components_path = root /"weather" /"static_components.pkl" + if static_components_path.exists(): + with open(static_components_path, 'rb') as f: + static_components = pickle.load(f) + else: + raise FileNotFoundError(f"Missing static_components.pkl in {static_components_path}. Please ensure the weather dataset is downloaded and processed correctly.") + + return metadata, static_components + +def _dist_rank_world(): + """Return the current process's (rank, world_size), or (0, 1) outside of DDP.""" + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank(), torch.distributed.get_world_size() + return 0, 1 + + +def _wait_for_weather_cache(ready_path: Path, timeout_s: int, poll_s: float = 5.0): + """Block until rank 0 writes the ready marker, then returns. Used for DDP processes to coordinate loading of files.""" + start = time.time() + while not ready_path.exists(): + elapsed = time.time() - start + if elapsed > timeout_s: + raise TimeoutError( + f"Timed out after {timeout_s}s waiting for weather cache marker: {ready_path}" + ) + time.sleep(poll_s) + +def _prepare_weather_cache_once(root: str, timeout_s: int, skip_cache_build: bool, task_name: str = "weather", use_prebuilt: bool = False): + """In DDP runs, rank 0 builds the processed cache; all other ranks wait for it. + + If use_prebuilt is True, rank 0 downloads the prebuilt processed file from + HuggingFace instead of building from raw data (no-op if already present). + Raw component files (static_components.pkl, metadata.pkl, etc.) are + downloaded from HuggingFace automatically when missing. + """ + #from wf_utils import EfficientWeatherGraphDataset + + root_path = Path(root) + processed_path = root_path /"weather" / "processed" / "weather_graph_data_processed.pt" + ready_path = root_path /"weather" / "processed" / "weather_graph_data_processed.ready" + rank, world_size = _dist_rank_world() + + if world_size == 1: + if use_prebuilt: + _ensure_prebuilt_pt(task_name, processed_path) + elif not processed_path.exists(): + _ensure_raw_weather_files(root, task_name) + _download_static_weather_files(root) + return + + if rank == 0: + if use_prebuilt: + _ensure_prebuilt_pt(task_name, processed_path) + elif skip_cache_build: + logger.warning("Rank 0 skipping weather cache build due to --skip_weather_cache_build") + elif not processed_path.exists(): + _ensure_raw_weather_files(root, task_name) + logger.info(f"Rank 0 building weather processed cache at {processed_path}") + EfficientWeatherGraphDataset(root=root, pre_transform=None, transform=None) + else: + logger.info(f"Rank 0 reusing existing weather processed cache at {processed_path}") + + if processed_path.exists(): + ready_path.parent.mkdir(parents=True, exist_ok=True) + ready_path.write_text(f"ready\nrank=0\nts={int(time.time())}\n", encoding="ascii") + _download_static_weather_files(root) + return + + # Non-zero ranks skip the wait if the cache already exists without a ready marker + if not ready_path.exists() and processed_path.exists(): + logger.info( + f"Rank {rank} found existing weather processed cache without ready marker; proceeding." + ) + return + + logger.info(f"Rank {rank} waiting for weather cache marker at {ready_path}") + _wait_for_weather_cache(ready_path, timeout_s=timeout_s) + logger.info(f"Rank {rank} detected weather cache ready marker") + + +def _download_prebuilt_pt(url: str, dst: Path) -> None: + """Download a prebuilt .pt file from a url to a destination path dst with a progress indicator.""" + import urllib.request + + dst.parent.mkdir(parents=True, exist_ok=True) + tmp = dst.with_suffix(".tmp") + logger.info(f"Downloading prebuilt processed file from {url} ...") + + def _reporthook(count, block_size, total_size): + if total_size > 0 and count % 500 == 0: + mb_done = count * block_size / (1024 ** 2) + mb_total = total_size / (1024 ** 2) + logger.info(f" {mb_done:.1f} / {mb_total:.1f} MB") + + try: + urllib.request.urlretrieve(url, str(tmp), reporthook=_reporthook) + tmp.rename(dst) + logger.info(f"Saved prebuilt processed file to {dst}") + except Exception: + tmp.unlink(missing_ok=True) + raise + + +def _ensure_prebuilt_pt(task_name: str, processed_path: Path) -> None: + """Download the prebuilt .pt for task_name if it is not already at processed_path.""" + if processed_path.exists(): + logger.info(f"Prebuilt processed file already present at {processed_path}; skipping download") + return + url = _PREBUILT_PT_URLS[task_name] + _download_prebuilt_pt(url, processed_path) + + +def _download_static_weather_files(root: str) -> None: + """Download the named static component files from the original weather + repo (log-rwth-aachen/Graphbench_Weather) into `root/weather`. + + These files (metadata.pkl, static_components.pkl, stats-*.nc) are shared + across all weather tasks and are always sourced from the same repository + as the prebuilt datasets, regardless of which task variant is used. + """ + dest_dir = Path(root) / "weather" + origin_repo_id = _HF_WEATHER_REPOS["weather"] + + missing = [ + hf_name for hf_name, local_name in _HF_WEATHER_STATIC_FILES.items() + if not (dest_dir / local_name).exists() + ] + if not missing: + logger.info(f"All static weather component files already present in {dest_dir}") + return + + logger.info(f"Missing static weather files: {missing}; downloading from {origin_repo_id}") + + try: + from huggingface_hub import hf_hub_download + except ImportError: + logger.warning( + "huggingface_hub is not installed; cannot auto-download static weather files. " + "Install it with: pip install huggingface_hub" + ) + return + + dest_dir.mkdir(parents=True, exist_ok=True) + for hf_name in missing: + local_name = _HF_WEATHER_STATIC_FILES[hf_name] + dst = dest_dir / local_name + if dst.exists(): + continue + logger.info(f" Downloading {hf_name} from {origin_repo_id} -> {dst}") + try: + hf_hub_download( + repo_id=origin_repo_id, + filename=hf_name, + repo_type="dataset", + local_dir=str(dest_dir), + local_dir_use_symlinks=False, + ) + # hf_hub_download saves to local_dir/; rename if needed + downloaded = dest_dir / hf_name + if downloaded.exists() and downloaded != dst: + downloaded.rename(dst) + logger.info(f" Renamed {hf_name} -> {local_name}") + except Exception as exc: + logger.warning(f" Failed to download {hf_name}: {exc}") + + +def _ensure_raw_weather_files(root: str, task_name: str) -> None: + """Download raw component files from HuggingFace if they are missing. + + For weather_subset, static component files (metadata, static_components, + stats NetCDFs) are sourced from the original Graphbench_Weather repo since + the subset shares the same graph structure and normalisation statistics. + Only the timestep data file (weather_64.pt) comes from the subset repo. + """ + root_path = Path(root) + task_repo_id = _HF_WEATHER_REPOS.get(task_name) + origin_repo_id = _HF_WEATHER_REPOS["weather"] + + if task_repo_id is None: + logger.warning(f"No HF repo configured for task {task_name}; skipping raw file download") + return + + all_files = _HF_WEATHER_RAW_FILES + missing = [ + hf_name for hf_name, local_name in all_files.items() + if not (root_path / local_name).exists() + ] + if not missing: + logger.info(f"All raw weather component files already present in {root_path}") + return + + logger.info(f"Missing raw weather files for '{task_name}': {missing}; downloading") + + try: + from huggingface_hub import hf_hub_download + except ImportError: + logger.warning( + "huggingface_hub is not installed; cannot auto-download raw weather files. " + "Install it with: pip install huggingface_hub" + ) + return + + root_path.mkdir(parents=True, exist_ok=True) + for hf_name in missing: + local_name = all_files[hf_name] + dst = root_path / local_name + if dst.exists(): + continue + # Static components are always sourced from the original weather repo. + if hf_name in _HF_WEATHER_STATIC_FILES: + repo_id = origin_repo_id + else: + repo_id = task_repo_id + logger.info(f" Downloading {hf_name} from {repo_id} -> {dst}") + try: + hf_hub_download( + repo_id=repo_id, + filename=hf_name, + repo_type="dataset", + local_dir=str(root_path), + local_dir_use_symlinks=False, + ) + # hf_hub_download saves to local_dir/; rename if needed + downloaded = root_path / hf_name + if downloaded.exists() and downloaded != dst: + downloaded.rename(dst) + logger.info(f" Renamed {hf_name} -> {local_name}") + except Exception as exc: + logger.warning(f" Failed to download {hf_name}: {exc}") + + + +class EfficientWeatherGraphDataset(InMemoryDataset): + """ + Weather forecasting dataset. + + Note: + This class **should not be used directly**, please use :class:`graphbench.Loader` instead to access the provided + datasets. + The purpose of this page is merely to provide details on the dataset. + + + Overview: + We provide a graph-based medium-range weather forecasting dataset derived from the + `ERA5 `_ reanalysis dataset. + We use a down-sampled version of ERA5 with a ``64 x 32`` equiangular grid and a temporal resolution of six + hours. + + The task is to model medium-range weather evolution by predicting the residual change in the atmospheric + state over a fixed 12-hour horizon. + Given an initial snapshot of the current atmospheric state, the model forecasts the 12-hour future change + in meteorological variables at each grid location. + + Since the dataset contains node features of various scales and units, we provide pre-computed mean and standard deviation values from the GraphCast team for each variable across all pressure levels. + We refer to the `GraphCast paper `_ for details on the computation and usage of these statistics. These statistics are downloaded alongside the weather dataset as "stats-mean_by_level.nc", + "stats-stddev_by_level.nc" and "stats-diffs_stddev_by_level.nc". + + Additional metadata used during dataset preparation and model training such as latitude and longitude are provided in metadata.pkl and static_components.pkl. + These metadata and static component values can also be accessed from each split of the dataset via the ``metadata`` and ``static_components`` attributes of the returned PyG Data objects. + Example: given the training split via ``train_dataset = Loader("data", "weather").load()[0]["train"]``, the metadata can be accessed via ``train_dataset.metadata`` and the static components via ``train_dataset.static_components``. + + Graph Attributes: + .. list-table:: + :header-rows: 1 + + * - Attribute + - Size + - Description + * - ``grid_x`` + - ``[num_nodes, 166]`` + - Node features: weather variables across all pressure levels for each grid coordinate of the current ([0:83]) and previous time step ([83:]). + * - timestep_idx + - ``[1]`` + - Given timestep index to attribute a specific time to each graph. + * - edge_index + - ``[2,]`` + - Contains the complete edge index for both grid and mesh nodes. + * - mesh_edge_index + - ``[2,]`` + - The edge index for the mesh graph as used in GraphCast. Contains only the edges of the mesh nodes. + * - node_type + - ``[4610]`` + - Denotes whether a node belongs to the grid or mesh graph. + * - x + - ``[2562, 1]`` + - Placeholder weather variable for the mesh nodes computed during the model forward pass. + + Please refer to the `GraphBench paper `_ for a detailed list of the weather + variables included in the dataset. + + Metadata Attributes: + .. list-table:: + :header-rows: 1 + + * - Attribute + - Size + - Description + * - ``grid_variables`` + - [11] + - The variable names for all grid variables. Includes both pressure dependent and independent variables. + * - ``mesh_splits`` + - [1] + - The number of mesh splits. + * - ``connectivity_radius`` + - [1] + - The connectivity radius used during mesh generation. + * - ``grid_lat`` + - [64] + - Grid latitute values used for loss computation and distance computation. + * - ``grid_lon`` + - [64] + - Grid longitude values used for loss computation and distance computation. + * - ``num_timesteps`` + - [1] + - Number of total timesteps in the dataset. + * - ``datetimes`` + - [93544] + - Datetime objects for each timestep in the dataset. Can be used to compute temporal splits for training, validation and testing. + * - ``variable_channel_counts`` + - [11] + - Number of appearances of variables for each variable type. + * - ``variable_channel_slices`` + - [11] + - Channel slices for each variable channel in grid_x features to assign correct loss values. + + Static components Attributes: + .. list-table:: + :header-rows: 1 + + * - ``mesh_vertices`` + - [2562,3] + - 3D coordinates of each mesh vertex used in the mesh graph. Not used in current computations. + * - ``variable_channel_counts`` + - [11] + - Number of appearances of variables for each variable type. + * - ``variable_channel_slices`` + - [11] + - Channel slices for each variable channel in grid_x features to assign correct loss values. + * - ``mesh_faces`` + - [2562,3] + - Node IDs of the mesh faces used to compute the mesh edge index. + * - ``mesh_lat`` + - [2562] + - Mesh latitude values used for loss computation and distance computation. + * - ``mesh_lon`` + - [2562] + - Mesh longitude values used for loss computation and distance computation. + * - ``grid2mesh_data`` + - NaN + - Data object holding the information for the grid to mesh data. Used in construction from single time steps. Not used for the already combined dataset (default). + * - ``grid_lat`` + - [64] + - Grid latitude values used for loss computation and distance computation. + * - ``grid_lon`` + - [64] + - Grid longitude values used for loss computation and distance computation. + * - ``num_grid_nodes`` + - [1] + - The number of grid nodes. + * - ``num_mesh_nodes`` + - [1] + - The number of mesh nodes. + + Note: + Currently the weather dataset is only evaluated on a twelve hour forecast time window. The target values can be obtained by using the current time step values for each node of timestep t+2 with timestep t as an input. + Therefore no value is set for the graph attribute y. + + List of Available Datasets: + We currently provide a single dataset, called ``weather``. + + It can be loaded like this: + + .. code:: python + + from graphbench import Loader + dataset = Loader("data", "weather").load() + """ + + def __init__(self, root: str, transform=None, pre_transform=None, pre_filter=None): + """ + Args: + root: Root directory containing the ``weather`` folder with raw and processed files. + transform: Optional PyG transform applied to data objects before every access. + pre_transform: Optional PyG transform applied before saving data objects to disk. + pre_filter: A function that indicates whether a data object should be included in the final dataset. + """ + self.root_path = Path(root) + + self.metadata_path = self.root_path /"weather" /"metadata.pkl" + if self.metadata_path.exists(): + with open(self.metadata_path, 'rb') as f: + self.metadata = pickle.load(f) + else: + self.metadata = {} + + self.static_components_path = self.root_path /"weather" /"static_components.pkl" + if self.static_components_path.exists(): + with open(self.static_components_path, 'rb') as f: + self.static_components = pickle.load(f) + else: + self.static_components = {} + + # Cached references to static tensors shared across all samples + self._cached_splits: Optional[Dict[str, np.ndarray]] = None + self._static_edge_index: Optional[torch.Tensor] = None + self._static_mesh_edge_index: Optional[torch.Tensor] = None + self._static_node_type: Optional[torch.Tensor] = None + self._static_mesh_x: Optional[torch.Tensor] = None + self._refresh_static_cache() + + super().__init__(str(self.root_path / "weather"), transform, pre_transform) + self.load(self.processed_paths[0]) + self._refresh_static_cache() + + def load(self, path: str, data_cls=Data) -> None: + """Load processed data, trying mmap first to reduce RSS spikes.""" + load_attempts = [ + {"map_location": "cpu", "weights_only": False, "mmap": True}, + {"map_location": "cpu", "mmap": True}, + {"map_location": "cpu", "weights_only": False}, + {"map_location": "cpu"}, + ] + + out = None + for kwargs in load_attempts: + try: + out = torch.load(path, **kwargs) + break + except (TypeError, RuntimeError, ValueError): + continue + + if out is None: + out = torch.load(path) + + assert isinstance(out, tuple) + assert len(out) == 2 or len(out) == 3 + if len(out) == 2: + data, self.slices = out + else: + data, self.slices, data_cls = out + + self.data = data if not isinstance(data, dict) else data_cls.from_dict(data) + + @property + def raw_file_names(self): + """Returns list of files that must be present to trigger process().""" + files = ["static_components.pkl", "metadata.pkl"] + if (self.root_path /"weather"/"weather_dataset.pt").exists(): + files.append("weather_dataset.pt") + return files + + @property + def processed_file_names(self): + """Returns name of the cached file produced by process().""" + return ['weather_graph_data_processed.pt'] + + def _load_torch_payload(self, file_path: Path) -> Any: + """Load a torch file with mmap-first fallback to reduce RSS spikes.""" + load_attempts = [ + {"map_location": "cpu", "weights_only": False, "mmap": True}, + {"map_location": "cpu", "mmap": True}, + {"map_location": "cpu", "weights_only": False}, + {"map_location": "cpu"}, + ] + for kwargs in load_attempts: + try: + return torch.load(file_path, **kwargs) + except (TypeError, RuntimeError, ValueError): + continue + return torch.load(file_path) + + def _refresh_static_cache(self) -> None: + """Pre-compute and cache static graph tensors shared across all timesteps.""" + if not self.static_components: + return + + grid2mesh_data = self.static_components.get('grid2mesh_data') + mesh_faces = self.static_components.get('mesh_faces') + num_mesh_nodes = int(self.static_components.get('num_mesh_nodes', 0)) + + if grid2mesh_data is None or mesh_faces is None or num_mesh_nodes <= 0: + return + + self._static_edge_index = grid2mesh_data.edge_index + self._static_node_type = grid2mesh_data.node_type + + # Build bidirectional mesh edge index from triangular faces + faces = torch.as_tensor(mesh_faces, dtype=torch.long) + tri_edges = torch.cat([ + faces[:, [0, 1]], + faces[:, [1, 2]], + faces[:, [2, 0]], + ], dim=0).t().contiguous() + self._static_mesh_edge_index = torch.cat([tri_edges, tri_edges.flip(0)], dim=1) + self._static_mesh_x = torch.zeros(num_mesh_nodes, 1, dtype=torch.float32) + + def _build_dynamic_graph( + self, + grid_features: torch.Tensor, + grid_features_prev: torch.Tensor, + timestep_idx: int, + ) -> Data: + """Create a minimal per-timestep Data object with only dynamic tensors.""" + return Data( + grid_x=torch.cat([grid_features, grid_features_prev], dim=-1), + timestep_idx=timestep_idx, + ) + + def _attach_static_fields(self, data: Data) -> Data: + """Attach cached static tensors to a dynamic Data object at read time.""" + if self._static_edge_index is None or self._static_mesh_edge_index is None or self._static_node_type is None: + return data + + data.edge_index = self._static_edge_index + data.mesh_edge_index = self._static_mesh_edge_index + data.node_type = self._static_node_type + + if self._static_mesh_x is not None: + grid_x = getattr(data, 'grid_x', None) + if isinstance(grid_x, torch.Tensor) and grid_x.dtype != self._static_mesh_x.dtype: + data.x = self._static_mesh_x.to(dtype=grid_x.dtype) + else: + data.x = self._static_mesh_x + + return data + + def process(self): + """Reads time step data, provides additional static and dynamic features if not prebuilt and collates data into a single processed .pt file.""" + print("Processing dataset... this might take a while.") + + with open(self.root_path / "weather" / "static_components.pkl", 'rb') as f: + static_components = pickle.load(f) + self.static_components = static_components + self._refresh_static_cache() + + # When no pre_filter/pre_transform is set, store dynamic fields only to save memory + use_dynamic_only_storage = self.pre_filter is None and self.pre_transform is None + + combined_file_path = self.root_path /"weather" / "weather_dataset.pt" + data_list = [] + + if combined_file_path.exists(): + # Load all timesteps from a single combined file [T, N, F] + print(f"Loading raw data from {combined_file_path}...") + loaded_raw = self._load_torch_payload(combined_file_path) + all_grid_features = loaded_raw['grid_features'] + num_timesteps = all_grid_features.shape[0] + prev_grid_features: Optional[torch.Tensor] = None + + for t in tqdm(range(num_timesteps), desc="Loading timesteps", unit="step"): + grid_features = all_grid_features[t] + if prev_grid_features is None: + prev_grid_features = torch.zeros_like(grid_features) + + if use_dynamic_only_storage: + graph = self._build_dynamic_graph(grid_features, prev_grid_features, t) + else: + timestep_data = {'grid_features': grid_features} + timestep_data_prev = {'grid_features': prev_grid_features} + graph = self._combine_static_and_dynamic(static_components, timestep_data, t, timestep_data_prev) + + if self.pre_filter is not None and not self.pre_filter(graph): + prev_grid_features = grid_features + continue + if self.pre_transform is not None: + graph = self.pre_transform(graph) + + data_list.append(graph) + prev_grid_features = grid_features + + else: + # Load timesteps from individual timestep_*.pt files + timestep_files = sorted(list(self.root_path.glob("timestep_*.pt"))) + num_timesteps = len(timestep_files) + print(f"Loading raw data from {num_timesteps} individual files...") + prev_grid_features: Optional[torch.Tensor] = None + + for t, file_path in enumerate(tqdm(timestep_files, desc="Loading timesteps", unit="file")): + timestep_data = self._load_torch_payload(file_path) + grid_features = timestep_data['grid_features'] + if prev_grid_features is None: + prev_grid_features = torch.zeros_like(grid_features) + + if use_dynamic_only_storage: + graph = self._build_dynamic_graph(grid_features, prev_grid_features, t) + else: + timestep_data_prev = {'grid_features': prev_grid_features} + graph = self._combine_static_and_dynamic(static_components, timestep_data, t, timestep_data_prev) + + if self.pre_filter is not None and not self.pre_filter(graph): + prev_grid_features = grid_features + continue + if self.pre_transform is not None: + graph = self.pre_transform(graph) + + data_list.append(graph) + prev_grid_features = grid_features + + print("Collating and saving...") + self.save(data_list, self.processed_paths[0]) + print("Done!") + + def _combine_static_and_dynamic(self, static: Dict, timestep_data: Dict[str, Any], timestep_idx: int, timestep_data_prev: Dict[str, Any]) -> Data: + """Builds the completed weather data object by combining static and dynamic features.""" + grid_features = timestep_data['grid_features'] + grid_features_prev = timestep_data_prev['grid_features'] + graph_data = self._build_dynamic_graph(grid_features, grid_features_prev, timestep_idx) + return self._attach_static_fields(graph_data) + + def _get_split_indices(self) -> Dict[str, np.ndarray]: + """Return cached train/val/test index arrays, computing them on first call.""" + if self._cached_splits is not None: + return self._cached_splits + + num_timesteps = len(self) + datetimes = self.metadata.get('datetimes', None) + split_dict = {} + computed = False + + if datetimes is not None and len(datetimes) == num_timesteps: + try: + splits = compute_temporal_splits(datetimes, test_duration_days=365, val_duration_days=180) + split_dict = { + 'train': splits.train_idx, + 'val': splits.val_idx, + 'test': splits.test_idx, + } + computed = True + except Exception: + pass + + if not computed: + # Proportional 80/10/10 fallback + n_train = max(int(0.8 * num_timesteps), 1) + n_val = max(int(0.1 * num_timesteps), 1) + idx = np.arange(num_timesteps, dtype=np.int64) + split_dict = { + 'train': idx[:n_train], + 'val': idx[n_train:n_train + n_val], + 'test': idx[n_train + n_val:], + } + + self._cached_splits = split_dict + return split_dict + + def __getitem__(self, index): + """Support string split keys ('train', 'val'/'valid', 'test') alongside integer indexing.""" + if isinstance(index, str): + key = index.lower() + if key == 'valid': + key = 'val' + if key in ('train', 'val', 'test'): + splits = self._get_split_indices() + return self.index_select(splits[key]) + + # Re-attach static fields for dynamic-only storage mode + data_or_subset = super().__getitem__(index) + if isinstance(data_or_subset, Data): + return self._attach_static_fields(data_or_subset) + return data_or_subset + + + class WeatherforecastingDataset(GraphDataset): - r""" + """ Weather forecasting dataset. Note: @@ -140,6 +857,7 @@ def _prepare(self) -> None: ) def _load_graphs(self) -> List[Data]: + """Load the weather graphs for this dataset's split.""" return self._load_weather_graphs() def _load_weather_graphs(self) -> List[Data]: @@ -164,8 +882,10 @@ def _find_matching_files(self,directory, task, size, split): @property def raw_file_names(self) -> list[str]: + """Unused, we drive our own cache.""" return [] @property def processed_file_names(self) -> list[str]: + """Unused, we drive our own cache.""" return [f"{self.name}_{self.split}.pt"] From ec7804035dfa6f89fdb396248a3a868213f5ca2b Mon Sep 17 00:00:00 2001 From: Timo Date: Sun, 5 Jul 2026 18:17:47 +0200 Subject: [PATCH 2/7] removed old weatherforecasting class and documentation --- graphbench/datasets/__init__.py | 4 +- graphbench/datasets/_weatherforecasting.py | 145 --------------------- 2 files changed, 2 insertions(+), 147 deletions(-) diff --git a/graphbench/datasets/__init__.py b/graphbench/datasets/__init__.py index 43259ac..ae33d01 100644 --- a/graphbench/datasets/__init__.py +++ b/graphbench/datasets/__init__.py @@ -5,7 +5,7 @@ from ._combinatorial_optimization import CODataset from ._electroniccircuits import ECDataset from ._sat import SATDataset -from ._weatherforecasting import WeatherforecastingDataset +from ._weatherforecasting import EfficientWeatherforecastingDataset __all__ = [ @@ -16,5 +16,5 @@ "CODataset", "ECDataset", "SATDataset", - "WeatherforecastingDataset", + "EfficientWeatherforecastingDataset", ] diff --git a/graphbench/datasets/_weatherforecasting.py b/graphbench/datasets/_weatherforecasting.py index 5bc0cf5..33768dd 100644 --- a/graphbench/datasets/_weatherforecasting.py +++ b/graphbench/datasets/_weatherforecasting.py @@ -744,148 +744,3 @@ def __getitem__(self, index): return data_or_subset - -class WeatherforecastingDataset(GraphDataset): - """ - Weather forecasting dataset. - - Note: - This class **should not be used directly**, please use :class:`graphbench.Loader` instead to access the provided - datasets. - The purpose of this page is merely to provide details on the dataset. - - - Overview: - We provide a graph-based medium-range weather forecasting dataset derived from the - `ERA5 `_ reanalysis dataset. - We use a down-sampled version of ERA5 with a ``64 x 32`` equiangular grid and a temporal resolution of six - hours. - - The task is to model medium-range weather evolution by predicting the residual change in the atmospheric - state over a fixed 12-hour horizon. - Given an initial snapshot of the current atmospheric state, the model forecasts the 12-hour future change - in meteorological variables at each grid location. - - - Graph Attributes: - .. list-table:: - :header-rows: 1 - - * - Attribute - - Size - - Description - * - ``x`` - - ``[num_nodes, 83]`` - - Node features: weather variables across all pressure levels for each grid coordinate. - * - . - - ``[num_nodes, 83]`` - - Target: the atmospheric state at a future timestep 12 hours later. - - Please refer to the `GraphBench paper `_ for a detailed list of the weather - variables included in the dataset. - - - List of Available Datasets: - We currently provide a single dataset, called ``weather``. - - It can be loaded like this: - - .. code:: python - - from graphbench import Loader - dataset = Loader("data", "weather").load() - """ - - def __init__( - self, - name: str, - split: Literal["train", "val", "test"], - root: Union[str, Path], - transform: Optional[Callable[[Data], Data]] = None, - pre_transform: Optional[Callable[[Data], Data]] = None, - pre_filter: Optional[Callable[[Data], bool]] = None, - size : Optional[int] = 64, - ): - """ - Args: - name: Dataset identifier. Must be ``weather_64`` in order to load the provided dataset. - Note that this differs from the name provided to :class:`graphbench.Loader`. - split: Whether to load the train, validation, or test split of the dataset. - root: Root directory where the dataset folder will be created. - transform: Optional PyG transform applied to data objects before every access. - pre_transform: Optional PyG transform applied before saving data objects to disk. - pre_filter: A function that indicates whether a data object should be included in the final dataset. - size: The grid size of the dataset. Currently, only ``64`` is supported for loading the provided dataset. - """ - #currently downloads everything at once for a single dataset. Up to the user to manually unpack it so far - self.SOURCES: Dict[str, SourceSpec] = { - "weather_64": SourceSpec( - url="https://huggingface.co/datasets/log-rwth-aachen/Graphbench_Weather/resolve/main/weather_64.pt", - raw_folder="weather_64", - ), - } - self.name = name.lower() - if self.name not in self.SOURCES: - raise ValueError(f"Unsupported dataset name: {self.name}") - assert split in ["train", "val", "test"], "Only 'train', 'val', 'test' splits are supported." - self.split = split - self.source = self.SOURCES[self.name] - self._logger = _logger - self.size = size - - self.weather_dir = Path(root) / "weatherforecasting" - self._raw_dir = (self.weather_dir / self.SOURCES[self.name].raw_folder) / "raw" - self.processed_path = self.weather_dir / self.SOURCES[self.name].raw_folder / "processed" / f"{self.name}.pt" - super().__init__(str(self.weather_dir), transform, pre_transform, pre_filter) - - self._load_cached_or_prepare( - processed_path=self.processed_path, - cleanup_raw=False, - logger=_logger, - ) - - def _prepare(self) -> None: - """ - Download and unpack the weather data if it is not already cached. - """ - - download_and_unpack( - source=self.source, - raw_dir=self._raw_dir, - processed_dir=self.processed_path, - logger=_logger, - ) - - def _load_graphs(self) -> List[Data]: - """Load the weather graphs for this dataset's split.""" - return self._load_weather_graphs() - - def _load_weather_graphs(self) -> List[Data]: - """ - Load weather graph data files matching the dataset split and size. - """ - filepaths = self._find_matching_files(task=self.name, split=self.split, directory=self._raw_dir, size=self.size) - self.load(filepaths[0]) - return [self.get(i) for i in range(len(self))] - - def _find_matching_files(self,directory, task, size, split): - """ - Find and return filenames matching the expected pattern for this dataset split and size. - """ - pattern = f"weather_{size}.pt" - print(directory) - return [os.path.join(directory, fname) - for fname in os.listdir(directory) - if fname == pattern] - - # --- InMemoryDataset API (not used directly but kept for PyG hygiene) ----- - - @property - def raw_file_names(self) -> list[str]: - """Unused, we drive our own cache.""" - return [] - - @property - def processed_file_names(self) -> list[str]: - """Unused, we drive our own cache.""" - return [f"{self.name}_{self.split}.pt"] From f71d3c49f2f79f799c9139ea0828323800c0e6d5 Mon Sep 17 00:00:00 2001 From: Timo Date: Sun, 5 Jul 2026 18:24:44 +0200 Subject: [PATCH 3/7] added missing method for timestep computation --- graphbench/datasets/_weatherforecasting.py | 35 +++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/graphbench/datasets/_weatherforecasting.py b/graphbench/datasets/_weatherforecasting.py index 33768dd..1b5591f 100644 --- a/graphbench/datasets/_weatherforecasting.py +++ b/graphbench/datasets/_weatherforecasting.py @@ -10,8 +10,9 @@ from __future__ import annotations import os +import time from pathlib import Path -from typing import Callable, Dict, List, Literal, Optional, Union +from typing import Callable, Dict, List, Literal, Optional, Union, Any from torch_geometric.data import Data from loguru import logger @@ -22,6 +23,8 @@ import torch import pickle import xarray as xr +import numpy as np + try: from tqdm.auto import tqdm except Exception: @@ -312,6 +315,36 @@ def _ensure_raw_weather_files(root: str, task_name: str) -> None: except Exception as exc: logger.warning(f" Failed to download {hf_name}: {exc}") +@dataclass(frozen=True) +class TemporalSplits: + """Immutable container for train/val/test index arrays.""" + train_idx: np.ndarray + val_idx: np.ndarray + test_idx: np.ndarray + + def as_slices(self) -> Tuple[slice, slice, slice]: + """Convert index arrays to contiguous slices (assumes sorted, contiguous indices).""" + def to_slice(indices: np.ndarray) -> slice: + if indices.size == 0: + return slice(0, 0) + return slice(int(indices[0]), int(indices[-1]) + 1) + + return to_slice(self.train_idx), to_slice(self.val_idx), to_slice(self.test_idx) + + +def compute_temporal_splits( + datetimes: Sequence[np.datetime64], +) -> TemporalSplits: + """Split datetimes into train (1979-2015), val (2016-2017), test (2018-2021). + + Falls back to 80/10/10 proportional split when any year-based split is empty. + """ + return compute_fixed_year_splits( + datetimes, + train_years=(1979, 2015), + val_years=(2016, 2017), + test_years=(2018, 2021), + ) class EfficientWeatherGraphDataset(InMemoryDataset): From 63eeda79541639b4d4de963e23729cec077b5826 Mon Sep 17 00:00:00 2001 From: Timo Date: Sun, 5 Jul 2026 18:27:13 +0200 Subject: [PATCH 4/7] small fix for weatherforecasting helpers --- graphbench/datasets/_weatherforecasting.py | 71 +++++++++++++++++++++- 1 file changed, 70 insertions(+), 1 deletion(-) diff --git a/graphbench/datasets/_weatherforecasting.py b/graphbench/datasets/_weatherforecasting.py index 1b5591f..a85a0ca 100644 --- a/graphbench/datasets/_weatherforecasting.py +++ b/graphbench/datasets/_weatherforecasting.py @@ -12,7 +12,7 @@ import os import time from pathlib import Path -from typing import Callable, Dict, List, Literal, Optional, Union, Any +from typing import Callable, Dict, List, Literal, Optional, Union, Any, Sequence, Tuple from torch_geometric.data import Data from loguru import logger @@ -346,6 +346,75 @@ def compute_temporal_splits( test_years=(2018, 2021), ) +def compute_fixed_year_splits( + datetimes: Sequence[np.datetime64], + *, + train_years: tuple[int, int] = (1979, 2015), + val_years: tuple[int, int] = (2016, 2017), + test_years: tuple[int, int] = (2018, 2021), +) -> TemporalSplits: + """Split a datetime sequence by calendar-year ranges. + + Sorts unsorted input internally and maps indices back to the original order. + Falls back to 80/10/10 proportional split when any year-based split is empty. + """ + dts = np.asarray(datetimes) + if dts.ndim != 1 or dts.size == 0: + return TemporalSplits( + train_idx=np.arange(dts.size, dtype=np.int64), + val_idx=np.array([], dtype=np.int64), + test_idx=np.array([], dtype=np.int64), + ) + + if not np.all(dts[:-1] <= dts[1:]): + order = np.argsort(dts, kind="stable") + dts_sorted = dts[order] + else: + order = None + dts_sorted = dts + + years = dts_sorted.astype('datetime64[Y]').astype(int) + 1970 + + def _range_mask(yrs: np.ndarray, yr_range: tuple[int, int]) -> np.ndarray: + start, end = int(yr_range[0]), int(yr_range[1]) + return (yrs >= start) & (yrs <= end) + + train_mask = _range_mask(years, train_years) + val_mask = _range_mask(years, val_years) + test_mask = _range_mask(years, test_years) + + train_idx_sorted = np.nonzero(train_mask)[0].astype(np.int64) + val_idx_sorted = np.nonzero(val_mask)[0].astype(np.int64) + test_idx_sorted = np.nonzero(test_mask)[0].astype(np.int64) + + # Proportional fallback when any split is empty + if ( + train_idx_sorted.size < 1 + or val_idx_sorted.size < 1 + or test_idx_sorted.size < 1 + ): + n = dts.size + n_train = max(int(0.8 * n), 1) + n_val = max(int(0.1 * n), 1) + n_test = max(n - n_train - n_val, 1) + if n_train + n_val + n_test > n: + n_test = n - n_train - n_val + train_idx_sorted = np.arange(0, n_train, dtype=np.int64) + val_idx_sorted = np.arange(n_train, n_train + n_val, dtype=np.int64) + test_idx_sorted = np.arange(n_train + n_val, n, dtype=np.int64) + + if order is not None: + def backmap(sorted_idx: np.ndarray) -> np.ndarray: + return order[sorted_idx] + + train_idx = backmap(train_idx_sorted) + val_idx = backmap(val_idx_sorted) + test_idx = backmap(test_idx_sorted) + train_idx.sort(); val_idx.sort(); test_idx.sort() + else: + train_idx, val_idx, test_idx = train_idx_sorted, val_idx_sorted, test_idx_sorted + + return TemporalSplits(train_idx=train_idx, val_idx=val_idx, test_idx=test_idx) class EfficientWeatherGraphDataset(InMemoryDataset): """ From 4ff8900b4f2e3448efe0d135ccd44eb921664b9b Mon Sep 17 00:00:00 2001 From: Timo Date: Sun, 5 Jul 2026 18:33:33 +0200 Subject: [PATCH 5/7] fixed linting errors --- graphbench/_weatherforecasting_helpers/utils.py | 11 ++++------- graphbench/datasets/_weatherforecasting.py | 16 ++++++---------- 2 files changed, 10 insertions(+), 17 deletions(-) diff --git a/graphbench/_weatherforecasting_helpers/utils.py b/graphbench/_weatherforecasting_helpers/utils.py index c502c87..08aad34 100644 --- a/graphbench/_weatherforecasting_helpers/utils.py +++ b/graphbench/_weatherforecasting_helpers/utils.py @@ -5,19 +5,14 @@ """ import numpy as np -from typing import Tuple, Optional import torch -from pathlib import Path -import pickle -from typing import List, Dict, Any, Optional, Tuple, Sequence -from torch.utils.data import Subset +from typing import List, Dict, Any, Sequence, Tuple, Optional import xarray as xr try: from tqdm.auto import tqdm except Exception: def tqdm(iterable, *args, **kwargs): return iterable -from torch_geometric.data import Data, Dataset, InMemoryDataset from dataclasses import dataclass @@ -351,7 +346,9 @@ def backmap(sorted_idx: np.ndarray) -> np.ndarray: train_idx = backmap(train_idx_sorted) val_idx = backmap(val_idx_sorted) test_idx = backmap(test_idx_sorted) - train_idx.sort(); val_idx.sort(); test_idx.sort() + train_idx.sort() + val_idx.sort() + test_idx.sort() else: train_idx, val_idx, test_idx = train_idx_sorted, val_idx_sorted, test_idx_sorted diff --git a/graphbench/datasets/_weatherforecasting.py b/graphbench/datasets/_weatherforecasting.py index a85a0ca..c65f8d4 100644 --- a/graphbench/datasets/_weatherforecasting.py +++ b/graphbench/datasets/_weatherforecasting.py @@ -9,20 +9,15 @@ from __future__ import annotations -import os import time from pathlib import Path -from typing import Callable, Dict, List, Literal, Optional, Union, Any, Sequence, Tuple - -from torch_geometric.data import Data +from typing import Dict, Optional, Any, Sequence, Tuple from loguru import logger -from graphbench._helpers import download_and_unpack, SourceSpec, get_logger -from ._base import GraphDataset -from torch_geometric.data import Data, Dataset, InMemoryDataset +from graphbench._helpers import get_logger +from torch_geometric.data import Data, InMemoryDataset from dataclasses import dataclass import torch import pickle -import xarray as xr import numpy as np try: @@ -30,7 +25,6 @@ except Exception: def tqdm(iterable, *args, **kwargs): return iterable -from torch.utils.data import Subset # (i) helper functions # -----------------------------------------------------------------------------# @@ -410,7 +404,9 @@ def backmap(sorted_idx: np.ndarray) -> np.ndarray: train_idx = backmap(train_idx_sorted) val_idx = backmap(val_idx_sorted) test_idx = backmap(test_idx_sorted) - train_idx.sort(); val_idx.sort(); test_idx.sort() + train_idx.sort() + val_idx.sort() + test_idx.sort() else: train_idx, val_idx, test_idx = train_idx_sorted, val_idx_sorted, test_idx_sorted From e2864c306f8112033fe8052643f34f575bdc2a5a Mon Sep 17 00:00:00 2001 From: Timo Date: Mon, 6 Jul 2026 10:40:26 +0200 Subject: [PATCH 6/7] added version 0.2 information --- docs/conf.py | 2 +- graphbench/datasets/__init__.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index af8e409..d08f3b6 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -24,7 +24,7 @@ project = "GraphBench" author = "Timo Stoll, Chendi Qian, Ben Finkelshtein, Ali Parviz, Darius Weber, Fabrizio Frasca, Hadar Shavit, Antoine Siraudin, Arman Mielke, Marie Anastacio, Erik Müller, Maya Bechler-Speicher, Michael Bronstein, Mikhail Galkin, Holger Hoos, Mathias Niepert, Bryan Perozzi, Jan Tönshoff, Christopher Morris" copyright = "2026, " + author -version = "0.1.3" +version = "0.2.0" diff --git a/graphbench/datasets/__init__.py b/graphbench/datasets/__init__.py index ae33d01..6f3e1b8 100644 --- a/graphbench/datasets/__init__.py +++ b/graphbench/datasets/__init__.py @@ -5,7 +5,7 @@ from ._combinatorial_optimization import CODataset from ._electroniccircuits import ECDataset from ._sat import SATDataset -from ._weatherforecasting import EfficientWeatherforecastingDataset +from ._weatherforecasting import EfficientWeatherGraphDataset __all__ = [ @@ -16,5 +16,5 @@ "CODataset", "ECDataset", "SATDataset", - "EfficientWeatherforecastingDataset", + "EfficientWeatherGraphDataset", ] From 67666a78c8082bf5dc9810dc6f13452011c064dc Mon Sep 17 00:00:00 2001 From: Timo Date: Mon, 6 Jul 2026 10:45:30 +0200 Subject: [PATCH 7/7] fixed dependencies --- pyproject.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3150515..ea4c867 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "graphbench-lib" -version = "0.1.3" # IMPORTANT: when bumping the version, also update the version in docs/conf.py +version = "0.2.0" # IMPORTANT: when bumping the version, also update the version in docs/conf.py requires-python = ">=3.10" description = "Next-generation graph learning benchmarking." readme = "README.md" @@ -57,6 +57,8 @@ dependencies = [ "networkx", "tqdm", "scikit-learn", + "xarray", + "loguru" ] [project.optional-dependencies]