From f036ba58aa8316ae13fcc93ade230f28d13c6ea0 Mon Sep 17 00:00:00 2001 From: Simon Meierhans Date: Mon, 14 Sep 2026 05:23:27 -0700 Subject: [PATCH] Add jena climate benchmark. PiperOrigin-RevId: 981091703 --- dgf/src/api/io.py | 1 + dgf/src/io/BUILD | 1 + dgf/src/io/dataset_loader.py | 381 ++++++++++++++++++++++++++++++ dgf/src/io/dataset_loader_test.py | 219 +++++++++++++++++ 4 files changed, 602 insertions(+) diff --git a/dgf/src/api/io.py b/dgf/src/api/io.py index 375322e..edcfc30 100644 --- a/dgf/src/api/io.py +++ b/dgf/src/api/io.py @@ -18,6 +18,7 @@ from dgf.src.io.dataset_loader import fetch_ogb_graph from dgf.src.io.dataset_loader import fetch_graphland_graph +from dgf.src.io.dataset_loader import fetch_jena_climate_graph from dgf.src.io.hgraph_in_memory import read_graphai_hgraph diff --git a/dgf/src/io/BUILD b/dgf/src/io/BUILD index bf21fe7..871ac2a 100644 --- a/dgf/src/io/BUILD +++ b/dgf/src/io/BUILD @@ -587,6 +587,7 @@ py_test( "//dgf/src/util:test_util", "//dgf/src/validate:in_memory_graph", # numpy dep, + # pandas dep, ], ) diff --git a/dgf/src/io/dataset_loader.py b/dgf/src/io/dataset_loader.py index 64cd900..01468c3 100644 --- a/dgf/src/io/dataset_loader.py +++ b/dgf/src/io/dataset_loader.py @@ -64,6 +64,7 @@ class Repo(str, enum.Enum): OGB = "OGB" CNS = "CNS" ZENODO = "ZENODO" + WEB = "WEB" def download_ogb_graph(name: str) -> Tuple[Any, Any, Any]: @@ -728,3 +729,383 @@ def load_graph(): return load_graph() else: return cache_lib.cache(cache_graph_path, load_graph) + + +JENA_CLIMATE_URL = ( + "https://storage.googleapis.com/tensorflow/tf-keras-datasets/jena_climate_2009_2016.csv.zip" +) + +JENA_CLIMATE_COLUMN_RENAME_MAP = { + "p (mbar)": "p_mbar", + "T (degC)": "t_degc", + "Tpot (K)": "tpot_k", + "Tdew (degC)": "tdew_degc", + "rh (%)": "rh_percent", + "VPmax (mbar)": "vpmax_mbar", + "VPact (mbar)": "vpact_mbar", + "VPdef (mbar)": "vpdef_mbar", + "sh (g/kg)": "sh_g_per_kg", + "H2OC (mmol/mol)": "h2oc_mmol_per_mol", + "rho (g/m**3)": "rho_g_per_cubic_m", + "wv (m/s)": "wv_m_per_s", + "max. wv (m/s)": "max_wv_m_per_s", + "wd (deg)": "wd_deg", +} + +JENA_WEATHER_FEATURE_NAMES = list(JENA_CLIMATE_COLUMN_RENAME_MAP.values()) + + +def download_jena_climate_csv(source: Optional[str] = None) -> pd.DataFrame: + """Downloads and parses the Jena Climate CSV. + + Cleans sentinel values (-9999.0 in wind velocities) and parses timestamps. + + Args: + source: Optional URL or file path. If None, downloads from the official + TensorFlow datasets public GCS archive. + + Returns: + A cleaned pandas DataFrame with timestamps and renamed columns. + """ + url_or_path = source or JENA_CLIMATE_URL + if url_or_path.startswith("http://") or url_or_path.startswith("https://"): + log.info("Downloading Jena Climate dataset from %s", url_or_path) + request = urllib.request.Request( + url_or_path, headers={"User-Agent": "Mozilla/5.0"} + ) + with urllib.request.urlopen(request) as response: + content = response.read() + with zipfile.ZipFile(io.BytesIO(content)) as zip_file: + csv_names = [ + name for name in zip_file.namelist() if name.endswith(".csv") + ] + if not csv_names: + raise ValueError(f"No CSV file found in archive from {url_or_path}") + with zip_file.open(csv_names[0]) as csv_file: + climate_df = pd.read_csv(csv_file) + else: + if url_or_path.endswith(".zip"): + with zipfile.ZipFile(url_or_path) as zip_file: + csv_names = [ + name for name in zip_file.namelist() if name.endswith(".csv") + ] + if not csv_names: + raise ValueError(f"No CSV file found in {url_or_path}") + with zip_file.open(csv_names[0]) as csv_file: + climate_df = pd.read_csv(csv_file) + else: + climate_df = pd.read_csv(url_or_path) + + missing_columns = [ + column + for column in ["Date Time", *JENA_CLIMATE_COLUMN_RENAME_MAP] + if column not in climate_df.columns + ] + if missing_columns: + raise ValueError( + f"Jena Climate data from {url_or_path} is missing the columns" + f" {missing_columns}." + ) + + # Clean sentinel values: wv (m/s) and max. wv (m/s) have -9999.0 for missing + for column in ["wv (m/s)", "max. wv (m/s)"]: + climate_df[column] = climate_df[column].replace(-9999.0, 0.0) + + # Parse Date Time to unix timestamp (seconds) + date_times = pd.to_datetime( + climate_df["Date Time"], format="%d.%m.%Y %H:%M:%S" + ) + climate_df["timestamp"] = (date_times.astype("int64") // 10**9).astype( + np.int64 + ) + + return climate_df.rename(columns=JENA_CLIMATE_COLUMN_RENAME_MAP) + + +def build_jena_climate_graph( + climate_df: pd.DataFrame, + forecast_horizon_seconds: int = 3600, + query_step: int = 6, + subsample_station_step: int = 1, +) -> Tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema]: + """Constructs a DGF InMemoryGraph and GraphSchema from Jena Climate data. + + Graph structure: + - Single station node containing 14 time series weather metrics. + - Query nodes at chronological intervals, each with `creation_time` + and regression label `temperature` (at `creation_time + horizon`). + - Bidirectional edges between query nodes and the single station node. + + Args: + climate_df: Jena Climate pandas DataFrame, as returned by + `download_jena_climate_csv`. + forecast_horizon_seconds: Horizon into the future to predict (default: 3600s + = 1 hour). + query_step: Stride for sampling query nodes (default: 6, i.e. 1 query/hour + for 10-minute data). + subsample_station_step: Stride for station time series observations + (default: 1). + + Returns: + A tuple of (InMemoryGraph, GraphSchema). + """ + missing_columns = [ + column + for column in ["timestamp", *JENA_WEATHER_FEATURE_NAMES] + if column not in climate_df.columns + ] + if missing_columns: + raise ValueError( + f"climate_df is missing the columns {missing_columns}. Expected the" + " output of download_jena_climate_csv." + ) + + all_timestamps = climate_df["timestamp"].to_numpy(dtype=np.int64) + all_temperatures = climate_df["t_degc"].to_numpy(dtype=np.float32) + + # Subsample station observations if requested + station_timestamps = all_timestamps[::subsample_station_step] + + # Station features schema and data + station_node_features: dict[str, np.ndarray] = { + "#id": np.array([b"station_0"], dtype=np.bytes_), + } + station_schema_features: dict[str, schema_lib.FeatureSchema] = { + "#id": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.BYTES, + semantic=schema_lib.FeatureSemantic.PRIMARY_ID, + ), + } + + station_times = np.empty(1, dtype=object) + station_times[0] = station_timestamps + station_node_features["time"] = station_times + station_schema_features["time"] = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_timeseries=True, + is_creation_time=True, + group="weather_ts", + shape=(None,), + ) + + for column_name in JENA_WEATHER_FEATURE_NAMES: + station_values = np.empty(1, dtype=object) + station_values[0] = ( + climate_df[column_name] + .iloc[::subsample_station_step] + .to_numpy(dtype=np.float32) + ) + station_node_features[column_name] = station_values + station_schema_features[column_name] = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + is_timeseries=True, + group="weather_ts", + shape=(None,), + ) + + station_node_set = in_memory_graph_lib.InMemoryNodeSet( + num_nodes=1, + features=station_node_features, + ) + + # Query nodes construction + candidate_query_indices = np.arange(0, len(climate_df), query_step) + candidate_query_times = all_timestamps[candidate_query_indices] + target_times = candidate_query_times + forecast_horizon_seconds + + target_indices = np.searchsorted(all_timestamps, target_times) + valid_mask = (target_indices < len(all_timestamps)) & ( + all_timestamps[np.minimum(target_indices, len(all_timestamps) - 1)] + == target_times + ) + + valid_query_times = candidate_query_times[valid_mask] + valid_target_temperatures = all_temperatures[target_indices[valid_mask]] + num_queries = len(valid_query_times) + + if num_queries == 0: + raise ValueError( + f"No valid query points found with horizon {forecast_horizon_seconds}s." + ) + + # Chronological split: 70% train, 20% valid, 10% test + num_train = int(num_queries * 0.7) + num_valid = int(num_queries * 0.2) + split_labels = np.full(num_queries, "n/a", dtype="S5") + split_labels[:num_train] = b"train" + split_labels[num_train : num_train + num_valid] = b"valid" + split_labels[num_train + num_valid :] = b"test" + + query_ids = np.array( + [f"query_{i}".encode("utf-8") for i in range(num_queries)], + dtype=np.bytes_, + ) + query_node_features: dict[str, np.ndarray] = { + "#id": query_ids, + "creation_time": valid_query_times, + "temperature": valid_target_temperatures, + "#split": split_labels, + } + + query_schema_features: dict[str, schema_lib.FeatureSchema] = { + "#id": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.BYTES, + semantic=schema_lib.FeatureSemantic.PRIMARY_ID, + ), + "creation_time": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_creation_time=True, + ), + "temperature": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + ), + "#split": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.BYTES, + semantic=schema_lib.FeatureSemantic.CATEGORICAL, + ), + } + + query_node_set = in_memory_graph_lib.InMemoryNodeSet( + num_nodes=num_queries, + features=query_node_features, + ) + + # Edge sets: bidirectional between queries and station node (0) + query_indices = np.arange(num_queries, dtype=np.int64) + zero_indices = np.zeros(num_queries, dtype=np.int64) + + query_to_station_edges = in_memory_graph_lib.InMemoryEdgeSet( + adjacency=np.stack([query_indices, zero_indices], axis=0), + features={}, + ) + station_to_query_edges = in_memory_graph_lib.InMemoryEdgeSet( + adjacency=np.stack([zero_indices, query_indices], axis=0), + features={}, + ) + + graph = in_memory_graph_lib.InMemoryGraph( + node_sets={ + "queries": query_node_set, + "station": station_node_set, + }, + edge_sets={ + "query_to_station": query_to_station_edges, + "station_to_query": station_to_query_edges, + }, + ) + + schema = schema_lib.GraphSchema( + node_sets={ + "queries": schema_lib.NodeSchema(features=query_schema_features), + "station": schema_lib.NodeSchema(features=station_schema_features), + }, + edge_sets={ + "query_to_station": schema_lib.EdgeSchema( + source="queries", + target="station", + features={}, + ), + "station_to_query": schema_lib.EdgeSchema( + source="station", + target="queries", + features={}, + ), + }, + ) + + return graph, schema + + +def fetch_jena_climate_graph( + name: str = "jena_climate_1h", + cache_dir: Optional[str] = "AUTO", + verbose: bool = True, + forecast_horizon_seconds: int = 3600, + query_step: int = 6, + subsample_station_step: int = 1, + repo: Union[Repo, str] = Repo.AUTO, + source: Optional[str] = None, +) -> Tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema]: + """Downloads and loads the Jena Climate time series benchmark into memory. + + This function loads the Jena Climate dataset + (https://www.bgc-jena.mpg.de/wetter/) and represents it as an in-memory graph + node prediction regression task with a single weather station node containing + 14 meteorological time-series features and query nodes representing prediction + time points. + + Usage example: + + ```python + graph, schema = dgf.io.fetch_jena_climate_graph() + dgf.analyse.print_schema(schema) + ``` + + Args: + name: The name of the dataset under CNS fetch_repo (e.g. 'jena_climate_1h' + or 'jena_climate_24h'). + cache_dir: Optional. Directory to cache the graph in order to avoid + re-downloading/re-parsing it each time. If "AUTO", uses OS default + temporary directory. If None, does not cache the graph. + verbose: Optional. Whether to print cache and download progress. + forecast_horizon_seconds: Forecasting horizon in seconds (default: 3600 for + 1 hour ahead). + query_step: Step size to subsample query nodes (default: 6, i.e. 1 query per + hour for 10-minute data). + subsample_station_step: Optional step to subsample the station time series + (default: 1, full 10-minute resolution). + repo: Define the source of the data (Repo.AUTO, Repo.CNS, Repo.WEB). + source: Optional URL or file path to the Jena Climate zip/CSV. + + Returns: + An InMemoryGraph instance and its GraphSchema. + """ + if cache_dir == "AUTO": + cache_dir = os.path.join(tempfile.gettempdir(), "gf_fetch_jena_climate") + + if cache_dir is not None: + fs.makedirs(cache_dir) + cache_key = ( + f"{name}_h{forecast_horizon_seconds}_qs{query_step}_" + f"ss{subsample_station_step}.cache" + ) + cache_graph_path = os.path.join(cache_dir, cache_key) + if verbose: + log.info("Caching Jena Climate graph at %s", cache_graph_path) + else: + cache_graph_path = None + + if isinstance(repo, str): + repo = Repo(repo) + + # Select the right repo. + if repo == Repo.AUTO: + repo = Repo.WEB + + if repo == Repo.CNS: + loader = functools.partial(load_from_cns, name=name) + elif repo == Repo.WEB: + + def loader(): + if verbose: + log.info( + "Loading Jena Climate data from %s", source or JENA_CLIMATE_URL + ) + climate_df = download_jena_climate_csv(source=source) + return build_jena_climate_graph( + climate_df=climate_df, + forecast_horizon_seconds=forecast_horizon_seconds, + query_step=query_step, + subsample_station_step=subsample_station_step, + ) + else: + raise ValueError(f"Unsupported repo for Jena Climate: {repo}") + + if cache_graph_path is None: + return loader() + else: + return cache_lib.cache(cache_graph_path, loader) diff --git a/dgf/src/io/dataset_loader_test.py b/dgf/src/io/dataset_loader_test.py index 07fd101..9d3bbbd 100644 --- a/dgf/src/io/dataset_loader_test.py +++ b/dgf/src/io/dataset_loader_test.py @@ -27,10 +27,23 @@ from dgf.src.util import test_util from dgf.src.validate import in_memory_graph as in_memory_graph_validate_lib import numpy as np +import pandas as pd test_util.disable_diff_truncation() +def jena_climate_dataframe(num_rows: int) -> pd.DataFrame: + """Returns a dummy frame in the format of `download_jena_climate_csv`.""" + date_range = pd.date_range("2009-01-01", periods=num_rows, freq="10min") + columns = { + name: np.linspace(1.0, 2.0, num_rows, dtype=np.float32) + for name in dataset_loader.JENA_WEATHER_FEATURE_NAMES + } + columns["t_degc"] = np.linspace(-5.0, 25.0, num_rows, dtype=np.float32) + columns["timestamp"] = (date_range.astype("int64") // 10**9).astype(np.int64) + return pd.DataFrame(columns) + + class LoadDatasetTest(parameterized.TestCase): def test_build_split_idx(self): @@ -213,6 +226,212 @@ def test_real_graphland(self, graph_name): ) in_memory_graph_validate_lib.validate_graph(graph, schema) + def test_download_jena_climate_csv(self): + raw_data = { + "Date Time": ["01.01.2009 00:10:00", "01.01.2009 00:20:00"], + "p (mbar)": [996.5, 996.6], + "T (degC)": [-8.0, -8.1], + "Tpot (K)": [265.4, 265.3], + "Tdew (degC)": [-8.9, -9.0], + "rh (%)": [93.3, 93.4], + "VPmax (mbar)": [3.3, 3.2], + "VPact (mbar)": [3.1, 3.0], + "VPdef (mbar)": [0.2, 0.2], + "sh (g/kg)": [1.9, 1.9], + "H2OC (mmol/mol)": [3.1, 3.1], + "rho (g/m**3)": [1307.8, 1308.0], + "wv (m/s)": [-9999.0, 0.7], + "max. wv (m/s)": [1.8, -9999.0], + "wd (deg)": [152.3, 136.1], + } + path = os.path.join(self.create_tempdir().full_path, "jena.csv") + pd.DataFrame(raw_data).to_csv(path, index=False) + + climate_df = dataset_loader.download_jena_climate_csv(source=path) + + # Columns are renamed to their normalized names. + for name in dataset_loader.JENA_WEATHER_FEATURE_NAMES: + self.assertIn(name, climate_df.columns) + # Date Time is parsed into unix seconds. + np.testing.assert_array_equal( + climate_df["timestamp"].to_numpy(), np.array([1230768600, 1230769200]) + ) + # The -9999.0 sentinels are replaced by zeros. + np.testing.assert_allclose(climate_df["wv_m_per_s"].to_numpy(), [0.0, 0.7]) + np.testing.assert_allclose( + climate_df["max_wv_m_per_s"].to_numpy(), [1.8, 0.0] + ) + + def test_download_jena_climate_csv_missing_column_raises(self): + path = os.path.join(self.create_tempdir().full_path, "jena.csv") + pd.DataFrame({"Date Time": ["01.01.2009 00:10:00"]}).to_csv( + path, index=False + ) + with self.assertRaisesRegex(ValueError, "missing the columns"): + dataset_loader.download_jena_climate_csv(source=path) + + @mock.patch.object(dataset_loader, "download_jena_climate_csv", autospec=True) + def test_fetch_jena_climate_graph_mocked(self, mock_download): + num_rows = 20 + climate_df = jena_climate_dataframe(num_rows) + mock_download.return_value = climate_df + + tmpdir = self.create_tempdir().full_path + graph, schema = dataset_loader.fetch_jena_climate_graph( + cache_dir=tmpdir, + forecast_horizon_seconds=1200, # 2 steps + query_step=1, + repo="WEB", + ) + + in_memory_graph_validate_lib.validate_graph(graph, schema) + + # Station node assertions + self.assertEqual(graph.node_sets["station"].num_nodes, 1) + station_features = graph.node_sets["station"].features + station_schema = schema.node_sets["station"].features + self.assertTrue(station_schema["time"].is_timeseries) + self.assertTrue(station_schema["time"].is_creation_time) + self.assertEqual(station_schema["time"].group, "weather_ts") + self.assertLen(station_features["time"][0], num_rows) + self.assertTrue(station_schema["t_degc"].is_timeseries) + self.assertEqual(station_schema["t_degc"].group, "weather_ts") + self.assertLen(station_features["t_degc"][0], num_rows) + for name in dataset_loader.JENA_WEATHER_FEATURE_NAMES: + self.assertIn(name, station_features) + + # Query node assertions + expected_num_queries = num_rows - 2 + self.assertEqual( + graph.node_sets["queries"].num_nodes, expected_num_queries + ) + query_features = graph.node_sets["queries"].features + query_schema = schema.node_sets["queries"].features + self.assertTrue(query_schema["creation_time"].is_creation_time) + self.assertFalse(query_schema["creation_time"].is_timeseries) + self.assertIn("temperature", query_features) + self.assertLen(query_features["temperature"], expected_num_queries) + self.assertIn("#split", query_features) + + # Check that temperature target is shifted by 2 steps + np.testing.assert_allclose( + query_features["temperature"], + climate_df["t_degc"].to_numpy()[2:], + rtol=1e-5, + ) + + # Edge set assertions + self.assertEqual( + graph.edge_sets["query_to_station"].adjacency.shape, + (2, expected_num_queries), + ) + self.assertEqual( + graph.edge_sets["station_to_query"].adjacency.shape, + (2, expected_num_queries), + ) + + # Verify loading from cache on subsequent call + graph_cached, _ = dataset_loader.fetch_jena_climate_graph( + cache_dir=tmpdir, + forecast_horizon_seconds=1200, + query_step=1, + repo="WEB", + ) + self.assertEqual(mock_download.call_count, 1) + self.assertEqual( + graph_cached.node_sets["queries"].num_nodes, expected_num_queries + ) + + @parameterized.parameters((3600, 6), (86400, 144)) + @mock.patch.object(dataset_loader, "download_jena_climate_csv", autospec=True) + def test_jena_climate_temporal_lookahead_and_clipping( + self, horizon_seconds, expected_step_offset, mock_download + ): + num_rows = 200 + climate_df = jena_climate_dataframe(num_rows) + mock_download.return_value = climate_df + + tmpdir = self.create_tempdir().full_path + graph, _ = dataset_loader.fetch_jena_climate_graph( + cache_dir=tmpdir, + forecast_horizon_seconds=horizon_seconds, + query_step=1, + repo="WEB", + ) + + query_creation_times = graph.node_sets["queries"].features["creation_time"] + query_target_temperatures = graph.node_sets["queries"].features[ + "temperature" + ] + station_times = graph.node_sets["station"].features["time"][0] + station_temperatures = graph.node_sets["station"].features["t_degc"][0] + + num_queries = len(query_creation_times) + self.assertEqual(num_queries, num_rows - expected_step_offset) + + # 1. Target temperature matches measurement at target time: + expected_targets = climate_df["t_degc"].to_numpy()[expected_step_offset:] + np.testing.assert_allclose( + query_target_temperatures, expected_targets, rtol=1e-5 + ) + + # 2. Query creation_time is exactly target_time - horizon_seconds: + for i in range(num_queries): + query_time = query_creation_times[i] + target_time = query_time + horizon_seconds + target_index = np.searchsorted(station_times, target_time) + self.assertEqual(station_times[target_index], target_time) + self.assertEqual( + query_target_temperatures[i], station_temperatures[target_index] + ) + + # 3. Validation of causal clipping: + # In temporal sampling, station observations are clipped at + # target_timestamp <= query_time: + visible_mask = station_times <= query_time + self.assertTrue(np.all(station_times[visible_mask] <= query_time)) + self.assertEqual(station_times[visible_mask][-1], query_time) + + # Ensure no future observations (lookahead) are visible: + hidden_mask = station_times > query_time + self.assertTrue(np.all(station_times[hidden_mask] > query_time)) + self.assertIn(target_time, station_times[hidden_mask]) + # Specifically, time difference between target measurement and latest + # visible observation is exactly horizon_seconds: + self.assertEqual( + target_time - station_times[visible_mask][-1], horizon_seconds + ) + + @unittest.skipIf( + os.environ.get("TEST_STRATEGY") != "local", + "Manual test that requires internet access and only runs on a" + " workstation with --test_strategy=local", + ) + def test_real_jena_climate(self): + r"""Download Jena Climate and check it. + + Usage example: + + ```shell + blaze test -c opt --test_strategy=local --test_output=streamed \ + --test_arg=--alsologtostderr \ + --test_filter=LoadDatasetTest.test_real_jena_climate \ + //third_party/py/dgf/src/io:dataset_loader_test + ``` + """ + tmpdir = self.create_tempdir().full_path + graph, schema = dataset_loader.fetch_jena_climate_graph( + cache_dir=tmpdir, + query_step=36, + repo="WEB", + ) + in_memory_graph_validate_lib.validate_graph(graph, schema) + self.assertEqual(graph.node_sets["station"].num_nodes, 1) + num_queries = graph.node_sets["queries"].num_nodes + self.assertIsNotNone(num_queries) + assert num_queries is not None + self.assertGreater(num_queries, 1000) + if __name__ == "__main__": absltest.main()