From 19889c9a08b0a02f310dcf903f99ac97adf71c11 Mon Sep 17 00:00:00 2001 From: Joshua Terranova Date: Mon, 10 Aug 2026 19:02:26 -0700 Subject: [PATCH 1/4] ENH: cache downloaded atmosphere netCDF datasets (#654) --- rocketpy/environment/atmosphere_cache.py | 381 ++++++++++++++++++ rocketpy/environment/environment.py | 217 +++++++++- .../unit/environment/test_atmosphere_cache.py | 187 +++++++++ 3 files changed, 773 insertions(+), 12 deletions(-) create mode 100644 rocketpy/environment/atmosphere_cache.py create mode 100644 tests/unit/environment/test_atmosphere_cache.py diff --git a/rocketpy/environment/atmosphere_cache.py b/rocketpy/environment/atmosphere_cache.py new file mode 100644 index 000000000..75cfa3c49 --- /dev/null +++ b/rocketpy/environment/atmosphere_cache.py @@ -0,0 +1,381 @@ +"""Disk cache for downloaded atmospheric datasets (netCDF profiles and JSON). + +Cache root defaults to ``~/.rocketpy_cache/atmosphere``. Override with the +``ROCKETPY_CACHE`` environment variable (the ``atmosphere`` subfolder is +created under that root). + +OPeNDAP "Best" aggregations are virtual catalogs, not downloadable files. For +Forecast/Ensemble shortcuts this module therefore stores the **location-and-time +profiles** RocketPy extracts after the first successful fetch, as a compact +``.nc`` file. Subsequent identical requests load those profiles from disk. +Windy responses are stored as ``.json``. +""" + +from __future__ import annotations + +import json +import os +import re +import tempfile +import warnings +from pathlib import Path + +import netCDF4 +import numpy as np + +CACHE_ENV_VAR = "ROCKETPY_CACHE" +DEFAULT_CACHE_ROOT = Path.home() / ".rocketpy_cache" +PROFILE_FORMAT_ATTR = "rocketpy_atmosphere_profiles_v1" + + +def get_cache_root() -> Path: + """Return the root cache directory (honors ``ROCKETPY_CACHE``).""" + return Path(os.environ.get(CACHE_ENV_VAR, DEFAULT_CACHE_ROOT)).expanduser() + + +def get_atmosphere_cache_dir() -> Path: + """Return the atmosphere subdirectory under the cache root.""" + return get_cache_root() / "atmosphere" + + +def ensure_atmosphere_cache_dir() -> Path | None: + """Create the atmosphere cache directory. + + Returns + ------- + pathlib.Path or None + The directory path, or ``None`` if creation failed (caching disabled). + """ + cache_dir = get_atmosphere_cache_dir() + try: + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir + except OSError as exc: + warnings.warn( + f"Could not create atmosphere cache directory '{cache_dir}': {exc}. " + "Caching disabled for this request.", + UserWarning, + stacklevel=2, + ) + return None + + +def sanitize_cache_key(key: str) -> str: + """Replace characters that are unsafe in filenames.""" + return re.sub(r"[^A-Za-z0-9_.-]", "_", key) + + +def cache_path_for(key: str, suffix: str) -> Path: + """Build a cache file path for ``key`` with the given suffix (e.g. ``.nc``).""" + if not suffix.startswith("."): + suffix = f".{suffix}" + return get_atmosphere_cache_dir() / f"{sanitize_cache_key(key)}{suffix}" + + +def build_atmosphere_cache_key( + kind: str, + source: str, + latitude: float, + longitude: float, + datetime_date, +) -> str: + """Build a stable cache key for a Forecast/Ensemble/Windy request.""" + if datetime_date is None: + date_part = "nodate" + else: + date_part = datetime_date.strftime("%Y%m%d%H") + return sanitize_cache_key( + f"{kind}_{source}_{latitude:.4f}_{longitude:.4f}_{date_part}" + ) + + +def is_remote_url(path_or_url) -> bool: + """Return True if ``path_or_url`` looks like an HTTP(S)/OPeNDAP URL.""" + if not isinstance(path_or_url, str): + return False + lowered = path_or_url.lower() + return lowered.startswith(("http://", "https://", "dods://")) + + +def atomic_write_bytes(path: Path, data: bytes) -> bool: + """Write ``data`` to ``path`` atomically. Returns False on failure.""" + cache_dir = ensure_atmosphere_cache_dir() + if cache_dir is None: + return False + try: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + dir=path.parent, delete=False, suffix=".tmp" + ) as handle: + handle.write(data) + temp_name = handle.name + Path(temp_name).replace(path) + return True + except OSError as exc: + warnings.warn( + f"Failed to write atmosphere cache file '{path}': {exc}.", + UserWarning, + stacklevel=2, + ) + try: + Path(temp_name).unlink(missing_ok=True) + except (OSError, NameError): + pass + return False + + +def load_json_cache(path: Path) -> dict | None: + """Load a JSON cache file, or ``None`` if missing/unreadable.""" + if not path.is_file(): + return None + try: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: + warnings.warn( + f"Failed to read cached atmosphere JSON '{path}': {exc}. " + "Fetching fresh data.", + UserWarning, + stacklevel=2, + ) + return None + + +def save_json_cache(path: Path, payload: dict) -> bool: + """Serialize ``payload`` as JSON to ``path``. Returns False on failure.""" + try: + data = json.dumps(payload).encode("utf-8") + except (TypeError, ValueError) as exc: + warnings.warn( + f"Failed to serialize atmosphere JSON for cache '{path}': {exc}.", + UserWarning, + stacklevel=2, + ) + return False + return atomic_write_bytes(path, data) + + +def write_profile_netcdf( + path: Path, + *, + height, + pressure, + temperature, + wind_u, + wind_v, + elevation: float, + max_expected_height: float, + kind: str = "forecast", +) -> bool: + """Write extracted atmospheric profiles to a compact local netCDF file.""" + cache_dir = ensure_atmosphere_cache_dir() + if cache_dir is None: + return False + + height = np.asarray(height, dtype=float) + pressure = np.asarray(pressure, dtype=float) + temperature = np.asarray(temperature, dtype=float) + wind_u = np.asarray(wind_u, dtype=float) + wind_v = np.asarray(wind_v, dtype=float) + + try: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + dir=path.parent, delete=False, suffix=".nc.tmp" + ) as handle: + temp_path = Path(handle.name) + + dataset = netCDF4.Dataset(temp_path, mode="w", format="NETCDF4") + try: + dataset.setncattr("rocketpy_cache_format", PROFILE_FORMAT_ATTR) + dataset.setncattr("rocketpy_cache_kind", kind) + dataset.setncattr("elevation", float(elevation)) + dataset.setncattr("max_expected_height", float(max_expected_height)) + + dataset.createDimension("level", height.size) + for name, values, units in ( + ("height", height, "m"), + ("pressure", pressure, "Pa"), + ("temperature", temperature, "K"), + ("wind_u", wind_u, "m s-1"), + ("wind_v", wind_v, "m s-1"), + ): + variable = dataset.createVariable(name, "f8", ("level",)) + variable.units = units + variable[:] = values + finally: + dataset.close() + + temp_path.replace(path) + return True + except OSError as exc: + warnings.warn( + f"Failed to write atmosphere profile cache '{path}': {exc}.", + UserWarning, + stacklevel=2, + ) + try: + temp_path.unlink(missing_ok=True) + except (OSError, NameError): + pass + return False + + +def read_profile_netcdf(path: Path) -> dict | None: + """Read a profile netCDF written by :func:`write_profile_netcdf`.""" + if not path.is_file(): + return None + try: + dataset = netCDF4.Dataset(path, mode="r") + except OSError as exc: + warnings.warn( + f"Failed to open atmosphere profile cache '{path}': {exc}. " + "Fetching fresh data.", + UserWarning, + stacklevel=2, + ) + return None + + try: + fmt = getattr(dataset, "rocketpy_cache_format", None) + if fmt != PROFILE_FORMAT_ATTR: + warnings.warn( + f"Ignoring atmosphere cache '{path}' with unknown format '{fmt}'.", + UserWarning, + stacklevel=2, + ) + return None + return { + "kind": getattr(dataset, "rocketpy_cache_kind", "forecast"), + "elevation": float(dataset.getncattr("elevation")), + "max_expected_height": float(dataset.getncattr("max_expected_height")), + "height": np.array(dataset.variables["height"][:], dtype=float), + "pressure": np.array(dataset.variables["pressure"][:], dtype=float), + "temperature": np.array(dataset.variables["temperature"][:], dtype=float), + "wind_u": np.array(dataset.variables["wind_u"][:], dtype=float), + "wind_v": np.array(dataset.variables["wind_v"][:], dtype=float), + } + except (AttributeError, KeyError, ValueError, OSError) as exc: + warnings.warn( + f"Failed to read atmosphere profile cache '{path}': {exc}. " + "Fetching fresh data.", + UserWarning, + stacklevel=2, + ) + return None + finally: + dataset.close() + + +def write_ensemble_profile_netcdf( + path: Path, + *, + levels, + height_ensemble, + temperature_ensemble, + wind_u_ensemble, + wind_v_ensemble, + elevation: float, + max_expected_height: float, +) -> bool: + """Write ensemble member profiles to a compact local netCDF file.""" + cache_dir = ensure_atmosphere_cache_dir() + if cache_dir is None: + return False + + levels = np.asarray(levels, dtype=float) + height_ensemble = np.asarray(height_ensemble, dtype=float) + temperature_ensemble = np.asarray(temperature_ensemble, dtype=float) + wind_u_ensemble = np.asarray(wind_u_ensemble, dtype=float) + wind_v_ensemble = np.asarray(wind_v_ensemble, dtype=float) + num_members, num_levels = height_ensemble.shape + + try: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + dir=path.parent, delete=False, suffix=".nc.tmp" + ) as handle: + temp_path = Path(handle.name) + + dataset = netCDF4.Dataset(temp_path, mode="w", format="NETCDF4") + try: + dataset.setncattr("rocketpy_cache_format", PROFILE_FORMAT_ATTR) + dataset.setncattr("rocketpy_cache_kind", "ensemble") + dataset.setncattr("elevation", float(elevation)) + dataset.setncattr("max_expected_height", float(max_expected_height)) + + dataset.createDimension("member", num_members) + dataset.createDimension("level", num_levels) + + level_var = dataset.createVariable("level", "f8", ("level",)) + level_var.units = "Pa" + level_var[:] = levels + + for name, values, units in ( + ("height", height_ensemble, "m"), + ("temperature", temperature_ensemble, "K"), + ("wind_u", wind_u_ensemble, "m s-1"), + ("wind_v", wind_v_ensemble, "m s-1"), + ): + variable = dataset.createVariable(name, "f8", ("member", "level")) + variable.units = units + variable[:] = values + finally: + dataset.close() + + temp_path.replace(path) + return True + except OSError as exc: + warnings.warn( + f"Failed to write ensemble atmosphere cache '{path}': {exc}.", + UserWarning, + stacklevel=2, + ) + try: + temp_path.unlink(missing_ok=True) + except (OSError, NameError): + pass + return False + + +def read_ensemble_profile_netcdf(path: Path) -> dict | None: + """Read an ensemble profile netCDF written by :func:`write_ensemble_profile_netcdf`.""" + if not path.is_file(): + return None + try: + dataset = netCDF4.Dataset(path, mode="r") + except OSError as exc: + warnings.warn( + f"Failed to open ensemble atmosphere cache '{path}': {exc}. " + "Fetching fresh data.", + UserWarning, + stacklevel=2, + ) + return None + + try: + fmt = getattr(dataset, "rocketpy_cache_format", None) + kind = getattr(dataset, "rocketpy_cache_kind", None) + if fmt != PROFILE_FORMAT_ATTR or kind != "ensemble": + return None + return { + "elevation": float(dataset.getncattr("elevation")), + "max_expected_height": float(dataset.getncattr("max_expected_height")), + "levels": np.array(dataset.variables["level"][:], dtype=float), + "height_ensemble": np.array(dataset.variables["height"][:], dtype=float), + "temperature_ensemble": np.array( + dataset.variables["temperature"][:], dtype=float + ), + "wind_u_ensemble": np.array(dataset.variables["wind_u"][:], dtype=float), + "wind_v_ensemble": np.array(dataset.variables["wind_v"][:], dtype=float), + } + except (AttributeError, KeyError, ValueError, OSError) as exc: + warnings.warn( + f"Failed to read ensemble atmosphere cache '{path}': {exc}. " + "Fetching fresh data.", + UserWarning, + stacklevel=2, + ) + return None + finally: + dataset.close() diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index edf3a342c..5272ddee2 100644 --- a/rocketpy/environment/environment.py +++ b/rocketpy/environment/environment.py @@ -1,5 +1,6 @@ # pylint: disable=too-many-public-methods, too-many-instance-attributes, too-many-lines import bisect +import hashlib import json import logging import os @@ -12,6 +13,17 @@ import numpy as np import pytz +from rocketpy.environment.atmosphere_cache import ( + build_atmosphere_cache_key, + cache_path_for, + is_remote_url, + load_json_cache, + read_ensemble_profile_netcdf, + read_profile_netcdf, + save_json_cache, + write_ensemble_profile_netcdf, + write_profile_netcdf, +) from rocketpy.environment.fetchers import ( fetch_aigfs_file_return_dataset, fetch_atmospheric_data_from_meteomatics, @@ -1233,6 +1245,7 @@ def set_atmospheric_model( # pylint: disable=too-many-statements pressure_conversion_factor=None, username=None, password=None, + no_cache=False, ): """Define the atmospheric model for this Environment. @@ -1353,6 +1366,11 @@ def set_atmospheric_model( # pylint: disable=too-many-statements Meteomatics account password. Only used when ``type`` is ``"meteomatics"``. If None (the default), the value is read from the ``METEOMATICS_PASSWORD`` environment variable. + no_cache : bool, optional + If True, bypass the on-disk atmosphere cache and force a fresh + download for remote Forecast/Ensemble/Windy sources. Cached files + live under ``~/.rocketpy_cache/atmosphere`` (or ``ROCKETPY_CACHE``). + Default is False. Returns ------- @@ -1400,7 +1418,10 @@ def set_atmospheric_model( # pylint: disable=too-many-statements case "custom_atmosphere": self.process_custom_atmosphere(pressure, temperature, wind_u, wind_v) case "windy": - self.process_windy_atmosphere(file) + self.process_windy_atmosphere( + **({} if file is None else {"model": file}), + no_cache=no_cache, + ) case "open_meteo": self.process_open_meteo_atmosphere( **({} if file is None else {"model": file}) @@ -1467,15 +1488,38 @@ def set_atmospheric_model( # pylint: disable=too-many-statements except KeyError: fetch_function = None - # Fetches the dataset using OpenDAP protocol or uses the file path - dataset = fetch_function() if fetch_function is not None else file + cache_path = self.__atmosphere_cache_path_for_request( + type, file, fetch_function + ) + loaded_from_cache = False + if cache_path is not None and not no_cache: + if type == "ensemble": + loaded_from_cache = self.__apply_cached_ensemble_profiles( + cache_path + ) + else: + loaded_from_cache = self.__apply_cached_forecast_profiles( + cache_path + ) - if type in ["forecast", "reanalysis"]: - self.process_forecast_reanalysis( - dataset, dictionary, conversion_factor=conversion_factor - ) - else: - self.process_ensemble(dataset, dictionary, conversion_factor) + if not loaded_from_cache: + # Fetches the dataset using OpenDAP protocol or uses the file path + dataset = fetch_function() if fetch_function is not None else file + + if type in ["forecast", "reanalysis"]: + self.process_forecast_reanalysis( + dataset, dictionary, conversion_factor=conversion_factor + ) + else: + self.process_ensemble( + dataset, dictionary, conversion_factor=conversion_factor + ) + + if cache_path is not None: + if type == "ensemble": + self.__save_ensemble_profiles_to_cache(cache_path) + else: + self.__save_forecast_profiles_to_cache(cache_path) ground_pressure = self.pressure(self.elevation) if not 30000 <= ground_pressure <= 120_000: @@ -1513,6 +1557,138 @@ def set_atmospheric_model( # pylint: disable=too-many-statements self.atmospheric_model_file = file self.atmospheric_model_dict = dictionary + def __atmosphere_cache_path_for_request(self, atm_type, file, fetch_function): + """Return a cache path for remote Forecast/Ensemble sources, else None.""" + if fetch_function is not None and isinstance(file, str): + source_label = file + elif is_remote_url(file): + source_label = ( + "url_" + hashlib.md5(file.encode("utf-8")).hexdigest()[:12] + ) + else: + return None + + return cache_path_for( + build_atmosphere_cache_key( + atm_type, + source_label, + self.latitude, + self.longitude, + self.datetime_date, + ), + ".nc", + ) + + def __apply_profiles_from_arrays( + self, height, pressure, temperature, wind_u, wind_v + ): + """Install forecast-style profile Functions from 1-D arrays.""" + wind_speed = calculate_wind_speed(wind_u, wind_v) + wind_heading = calculate_wind_heading(wind_u, wind_v) + wind_direction = convert_wind_heading_to_direction(wind_heading) + data_array = mask_and_clean_dataset( + pressure, + height, + temperature, + wind_u, + wind_v, + wind_heading, + wind_direction, + wind_speed, + ) + self.__set_pressure_function(data_array[:, (1, 0)]) + self.__set_barometric_height_function(data_array[:, (0, 1)]) + self.__set_temperature_function(data_array[:, (1, 2)]) + self.__set_wind_velocity_x_function(data_array[:, (1, 3)]) + self.__set_wind_velocity_y_function(data_array[:, (1, 4)]) + self.__set_wind_heading_function(data_array[:, (1, 5)]) + self.__set_wind_direction_function(data_array[:, (1, 6)]) + self.__set_wind_speed_function(data_array[:, (1, 7)]) + return data_array + + def __apply_cached_forecast_profiles(self, cache_path): + """Load Forecast/Reanalysis profiles from disk. Return True on success.""" + profiles = read_profile_netcdf(cache_path) + if profiles is None: + return False + self.__apply_profiles_from_arrays( + profiles["height"], + profiles["pressure"], + profiles["temperature"], + profiles["wind_u"], + profiles["wind_v"], + ) + self.elevation = profiles["elevation"] + self._max_expected_height = profiles["max_expected_height"] + return True + + def __save_forecast_profiles_to_cache(self, cache_path): + """Persist the active Forecast/Reanalysis profiles to ``cache_path``.""" + if not self.pressure.is_array_source(): + return + pressure_source = np.asarray(self.pressure.source, dtype=float) + temperature_source = np.asarray(self.temperature.source, dtype=float) + wind_u_source = np.asarray(self.wind_velocity_x.source, dtype=float) + wind_v_source = np.asarray(self.wind_velocity_y.source, dtype=float) + write_profile_netcdf( + cache_path, + height=pressure_source[:, 0], + pressure=pressure_source[:, 1], + temperature=temperature_source[:, 1], + wind_u=wind_u_source[:, 1], + wind_v=wind_v_source[:, 1], + elevation=float(self.elevation), + max_expected_height=float(self._max_expected_height), + kind="forecast", + ) + + def __apply_cached_ensemble_profiles(self, cache_path): + """Load Ensemble member profiles from disk. Return True on success.""" + profiles = read_ensemble_profile_netcdf(cache_path) + if profiles is None: + return False + + height = profiles["height_ensemble"] + temper = profiles["temperature_ensemble"] + wind_u = profiles["wind_u_ensemble"] + wind_v = profiles["wind_v_ensemble"] + levels = profiles["levels"] + + wind_speed = calculate_wind_speed(wind_u, wind_v) + wind_heading = calculate_wind_heading(wind_u, wind_v) + wind_direction = convert_wind_heading_to_direction(wind_heading) + + self.level_ensemble = levels + self.height_ensemble = height + self.temperature_ensemble = temper + self.wind_u_ensemble = wind_u + self.wind_v_ensemble = wind_v + self.wind_heading_ensemble = wind_heading + self.wind_direction_ensemble = wind_direction + self.wind_speed_ensemble = wind_speed + self.num_ensemble_members = height.shape[0] + self.elevation = profiles["elevation"] + self._max_expected_height = profiles["max_expected_height"] + self.select_ensemble_member() + return True + + def __save_ensemble_profiles_to_cache(self, cache_path): + """Persist Ensemble member profiles to ``cache_path``.""" + if not hasattr(self, "height_ensemble"): + return + write_ensemble_profile_netcdf( + cache_path, + levels=self.level_ensemble, + height_ensemble=self.height_ensemble, + temperature_ensemble=self.temperature_ensemble, + wind_u_ensemble=self.wind_u_ensemble, + wind_v_ensemble=self.wind_v_ensemble, + elevation=float(self.elevation), + max_expected_height=float( + getattr(self, "_max_expected_height", self.max_expected_height) + ), + ) + # Atmospheric model processing methods def process_standard_atmosphere(self): @@ -1661,7 +1837,9 @@ def wind_heading_func(h): # TODO: create another custom reset for heading self._max_expected_height = max_expected_height - def process_windy_atmosphere(self, model="ECMWF"): # pylint: disable=too-many-statements + def process_windy_atmosphere( # pylint: disable=too-many-statements + self, model="ECMWF", no_cache=False + ): """Process data from Windy.com to retrieve atmospheric forecast data. Parameters @@ -1671,6 +1849,8 @@ def process_windy_atmosphere(self, model="ECMWF"): # pylint: disable=too-many-s ``ECMWF`` for the `ECMWF-HRES` model, ``GFS`` for the `GFS` model, ``ICON`` for the `ICON-Global` model or ``ICONEU`` for the `ICON-EU` model. + no_cache : bool, optional + If True, force a fresh download even when a JSON cache entry exists. Raises ------ @@ -1685,9 +1865,22 @@ def process_windy_atmosphere(self, model="ECMWF"): # pylint: disable=too-many-s "Valid options are 'ECMWF', 'GFS', 'ICON' or 'ICONEU'." ) - response = fetch_atmospheric_data_from_windy( - self.latitude, self.longitude, model + cache_path = cache_path_for( + build_atmosphere_cache_key( + "windy", + model, + self.latitude, + self.longitude, + self.datetime_date, + ), + ".json", ) + response = None if no_cache else load_json_cache(cache_path) + if response is None: + response = fetch_atmospheric_data_from_windy( + self.latitude, self.longitude, model + ) + save_json_cache(cache_path, response) # Determine time index from model time_array = np.array(response["data"]["hours"]) diff --git a/tests/unit/environment/test_atmosphere_cache.py b/tests/unit/environment/test_atmosphere_cache.py new file mode 100644 index 000000000..d880b9a2a --- /dev/null +++ b/tests/unit/environment/test_atmosphere_cache.py @@ -0,0 +1,187 @@ +"""Unit tests for atmosphere netCDF/JSON disk caching (#654).""" + +from datetime import datetime +from unittest.mock import MagicMock + +import netCDF4 +import numpy as np +import pytest + +from rocketpy import Environment +from rocketpy.environment import atmosphere_cache + + +def _write_minimal_profile_nc(path, elevation=1400.0): + """Create a tiny valid profile cache file for apply tests.""" + height = np.array([1400.0, 5000.0, 10000.0]) + pressure = np.array([85000.0, 54000.0, 26500.0]) + temperature = np.array([288.0, 255.0, 223.0]) + wind_u = np.array([1.0, 2.0, 3.0]) + wind_v = np.array([-1.0, 0.0, 1.0]) + assert atmosphere_cache.write_profile_netcdf( + path, + height=height, + pressure=pressure, + temperature=temperature, + wind_u=wind_u, + wind_v=wind_v, + elevation=elevation, + max_expected_height=10000.0, + kind="forecast", + ) + + +def test_cache_root_honors_rocketpy_cache_env(monkeypatch, tmp_path): + """``ROCKETPY_CACHE`` redirects the atmosphere cache root.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + assert atmosphere_cache.get_cache_root() == tmp_path + assert atmosphere_cache.get_atmosphere_cache_dir() == tmp_path / "atmosphere" + + +def test_profile_netcdf_roundtrip(monkeypatch, tmp_path): + """Write and read forecast profile netCDF through the cache helpers.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + path = atmosphere_cache.cache_path_for("forecast_test_key", ".nc") + _write_minimal_profile_nc(path) + loaded = atmosphere_cache.read_profile_netcdf(path) + assert loaded is not None + assert loaded["elevation"] == pytest.approx(1400.0) + np.testing.assert_allclose(loaded["height"], [1400.0, 5000.0, 10000.0]) + np.testing.assert_allclose(loaded["pressure"], [85000.0, 54000.0, 26500.0]) + + +def test_json_cache_roundtrip(monkeypatch, tmp_path): + """Windy-style JSON cache round-trips through disk.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + path = atmosphere_cache.cache_path_for("windy_test_key", ".json") + payload = {"data": {"hours": [1, 2, 3], "temp-surface": [288]}} + assert atmosphere_cache.save_json_cache(path, payload) + assert atmosphere_cache.load_json_cache(path) == payload + + +def test_forecast_shortcut_reuses_disk_cache(monkeypatch, tmp_path): + """Second Forecast shortcut call loads profiles from disk (no re-fetch).""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + fixture = "data/weather/SpaceportAmerica_2018_ERA-5.nc" + fetch_calls = [] + + def fake_fetch(): + fetch_calls.append(1) + return netCDF4.Dataset(fixture) + + env = Environment( + latitude=32.990254, + longitude=-106.974998, + elevation=1400, + datum="WGS84", + ) + env.set_date((2018, 10, 15, 12)) + env._Environment__atm_type_file_to_function_map["forecast"]["GFS"] = fake_fetch + + env.set_atmospheric_model( + type="Forecast", + file="GFS", + dictionary="ECMWF", + pressure_conversion_factor="hPa", + ) + assert len(fetch_calls) == 1 + pressure_first = env.pressure(env.elevation) + cached_files = list((tmp_path / "atmosphere").glob("*.nc")) + assert cached_files, "Expected a profile .nc cache file after first fetch" + + env2 = Environment( + latitude=32.990254, + longitude=-106.974998, + elevation=1400, + datum="WGS84", + ) + env2.set_date((2018, 10, 15, 12)) + env2._Environment__atm_type_file_to_function_map["forecast"]["GFS"] = fake_fetch + env2.set_atmospheric_model( + type="Forecast", + file="GFS", + dictionary="ECMWF", + pressure_conversion_factor="hPa", + ) + assert len(fetch_calls) == 1, "Second call should reuse disk cache" + assert env2.pressure(env2.elevation) == pytest.approx(pressure_first, rel=1e-6) + + +def test_forecast_no_cache_bypasses_disk(monkeypatch, tmp_path): + """``no_cache=True`` forces a re-fetch even when a cache file exists.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + fixture = "data/weather/SpaceportAmerica_2018_ERA-5.nc" + fetch_calls = [] + + def fake_fetch(): + fetch_calls.append(1) + return netCDF4.Dataset(fixture) + + env = Environment( + latitude=32.990254, + longitude=-106.974998, + elevation=1400, + datum="WGS84", + ) + env.set_date((2018, 10, 15, 12)) + env._Environment__atm_type_file_to_function_map["forecast"]["GFS"] = fake_fetch + + env.set_atmospheric_model( + type="Forecast", + file="GFS", + dictionary="ECMWF", + pressure_conversion_factor="hPa", + ) + env.set_atmospheric_model( + type="Forecast", + file="GFS", + dictionary="ECMWF", + pressure_conversion_factor="hPa", + no_cache=True, + ) + assert len(fetch_calls) == 2 + + +def test_windy_json_cache_hit(monkeypatch, tmp_path): + """Windy response is cached as JSON; second call skips the network.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + + # Minimal Windy payload matching __parse_windy_file expectations. + levels = [1000, 950, 925, 900, 850, 800, 700, 600, 500, 400, 300, 250, 200, 150] + payload = { + "header": {"elevation": 1234.0}, + "data": { + "hours": [1_540_000_000_000, 1_540_003_600_000], + }, + } + for level in levels: + # Geopotential heights increasing with altitude (decreasing pressure). + payload["data"][f"gh-{level}h"] = [ + float(2000 + (1000 - level) * 10), + float(2000 + (1000 - level) * 10), + ] + payload["data"][f"temp-{level}h"] = [280.0, 281.0] + payload["data"][f"wind_u-{level}h"] = [1.0, 1.5] + payload["data"][f"wind_v-{level}h"] = [-1.0, -0.5] + + fetch_mock = MagicMock(return_value=payload) + monkeypatch.setattr( + "rocketpy.environment.environment.fetch_atmospheric_data_from_windy", + fetch_mock, + ) + + env = Environment(latitude=45.0, longitude=10.0, elevation=100) + env.set_date(datetime(2018, 10, 15, 12)) + env.set_atmospheric_model(type="Windy", file="ECMWF") + assert fetch_mock.call_count == 1 + assert list((tmp_path / "atmosphere").glob("*.json")) + + env2 = Environment(latitude=45.0, longitude=10.0, elevation=100) + env2.set_date(datetime(2018, 10, 15, 12)) + env2.set_atmospheric_model(type="Windy", file="ECMWF") + assert fetch_mock.call_count == 1 + + env3 = Environment(latitude=45.0, longitude=10.0, elevation=100) + env3.set_date(datetime(2018, 10, 15, 12)) + env3.set_atmospheric_model(type="Windy", file="ECMWF", no_cache=True) + assert fetch_mock.call_count == 2 From b3f6a8ffdf9ae313d4672f5623d89049264391d2 Mon Sep 17 00:00:00 2001 From: Gui-FernandesBR Date: Fri, 14 Aug 2026 23:57:30 -0300 Subject: [PATCH 2/4] MNT: satisfy ruff format on the atmosphere cache changes Collapse the cache-key assignment the formatter wants on one line. No behaviour change. Co-Authored-By: Claude Opus 5 (1M context) --- rocketpy/environment/environment.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index 5272ddee2..40edba604 100644 --- a/rocketpy/environment/environment.py +++ b/rocketpy/environment/environment.py @@ -1562,9 +1562,7 @@ def __atmosphere_cache_path_for_request(self, atm_type, file, fetch_function): if fetch_function is not None and isinstance(file, str): source_label = file elif is_remote_url(file): - source_label = ( - "url_" + hashlib.md5(file.encode("utf-8")).hexdigest()[:12] - ) + source_label = "url_" + hashlib.md5(file.encode("utf-8")).hexdigest()[:12] else: return None From d949f0dd04c30db3d161be352a95d0440828abc3 Mon Sep 17 00:00:00 2001 From: Gui-FernandesBR Date: Sat, 12 Sep 2026 03:32:59 -0300 Subject: [PATCH 3/4] BUG: make an atmosphere cache hit reproduce a fresh fetch The cache restored only the eight profile Functions, elevation and _max_expected_height, so every attribute Environment derives from the dataset was missing on the second run of an otherwise identical script. atmospheric_model_type was still set to "Forecast", so info() and all_info() walked into the Forecast branch of the prints and raised AttributeError on atmospheric_model_init_date, and to_dict() serialized those fields as None. Persist the full metadata (date range, grid bounds, and the raw interpolation inputs) alongside the profiles and reinstate it on a hit, bumping the cache format to v2. Saving also indexed temperature and both wind profiles as 2-D arrays while guarding only pressure, which raised IndexError for a constant wind and broke three existing tests. Guarding pressure alone is not sufficient either: np.asarray(None, dtype=float) yields nan, so a missing column would be written out as an entry full of NaN winds. Check every column against the pressure grid before writing. Also: - Expire forecast entries after ROCKETPY_CACHE_TTL seconds (default 6h, the GFS cycle). A launch date days out would otherwise pin the first forecast ever downloaded for it and silently reuse it forever. Reanalysis is immutable and never expires; a TTL of 0 disables expiry. - Let ROCKETPY_CACHE=0/off/false/no/none/disabled turn the cache off. It previously only relocated the directory, so there was no way to opt out globally, and an empty value cached into the working directory. - Include the variable dictionary and pressure conversion factor in the cache key, so the same source decoded two ways no longer collides. - Add clear_atmosphere_cache() to remove the entries. - Treat a damaged cache file as a miss rather than letting netCDF4's RuntimeError escape and take down the simulation with it. - Pass usedforsecurity=False to hashlib.md5 so the key still builds under a FIPS-enabled Python. - Isolate ROCKETPY_CACHE per test via an autouse fixture. The suite was writing to the real ~/.rocketpy_cache, which let one test read profiles cached by an earlier one. Co-Authored-By: Claude Opus 5 (1M context) --- rocketpy/environment/atmosphere_cache.py | 521 +++++++++++++----- rocketpy/environment/environment.py | 208 +++++-- tests/conftest.py | 12 + .../unit/environment/test_atmosphere_cache.py | 300 ++++++++++ 4 files changed, 843 insertions(+), 198 deletions(-) diff --git a/rocketpy/environment/atmosphere_cache.py b/rocketpy/environment/atmosphere_cache.py index 75cfa3c49..393ac6e5b 100644 --- a/rocketpy/environment/atmosphere_cache.py +++ b/rocketpy/environment/atmosphere_cache.py @@ -1,14 +1,22 @@ """Disk cache for downloaded atmospheric datasets (netCDF profiles and JSON). -Cache root defaults to ``~/.rocketpy_cache/atmosphere``. Override with the -``ROCKETPY_CACHE`` environment variable (the ``atmosphere`` subfolder is -created under that root). +Cache root defaults to ``~/.rocketpy_cache/atmosphere``. Override the root with +the ``ROCKETPY_CACHE`` environment variable (the ``atmosphere`` subfolder is +created under it), or disable caching entirely by setting it to one of ``0``, +``off``, ``false``, ``no``, ``none`` or ``disabled``. OPeNDAP "Best" aggregations are virtual catalogs, not downloadable files. For Forecast/Ensemble shortcuts this module therefore stores the **location-and-time profiles** RocketPy extracts after the first successful fetch, as a compact -``.nc`` file. Subsequent identical requests load those profiles from disk. +``.nc`` file, together with every derived attribute ``Environment`` publishes +for that model (date range, grid bounds and the raw interpolation inputs) so a +cache hit reproduces the same object a fresh download would have produced. Windy responses are stored as ``.json``. + +Forecasts are re-issued by their providers on a fixed cycle, so cache entries +expire after ``ROCKETPY_CACHE_TTL`` seconds (default: 6 hours, matching the GFS +cycle). Reanalysis data is immutable and never expires. Set the TTL to ``0`` to +disable expiry. """ from __future__ import annotations @@ -16,21 +24,62 @@ import json import os import re +import shutil import tempfile +import time import warnings +from datetime import datetime from pathlib import Path import netCDF4 import numpy as np CACHE_ENV_VAR = "ROCKETPY_CACHE" +CACHE_TTL_ENV_VAR = "ROCKETPY_CACHE_TTL" DEFAULT_CACHE_ROOT = Path.home() / ".rocketpy_cache" -PROFILE_FORMAT_ATTR = "rocketpy_atmosphere_profiles_v1" +PROFILE_FORMAT_ATTR = "rocketpy_atmosphere_profiles_v2" +JSON_FORMAT_KEY = "rocketpy_cache_format" +CREATED_AT_ATTR = "rocketpy_cache_created_at" +DEFAULT_CACHE_TTL_SECONDS = 6 * 3600 + +#: Model kinds whose data never changes once published, so they never expire. +IMMUTABLE_KINDS = frozenset({"reanalysis"}) + +_DISABLED_VALUES = frozenset({"", "0", "off", "false", "no", "none", "disabled"}) +_DATE_FORMAT = "%Y-%m-%dT%H:%M:%S" + +# Scalar metadata persisted as netCDF attributes, with the caster used on read. +_SCALAR_METADATA = ( + ("atmospheric_model_interval", float), + ("atmospheric_model_init_lat", float), + ("atmospheric_model_end_lat", float), + ("atmospheric_model_init_lon", float), + ("atmospheric_model_end_lon", float), + ("lat_index", int), + ("lon_index", int), +) +_DATE_METADATA = ("atmospheric_model_init_date", "atmospheric_model_end_date") +_PAIR_METADATA = ("lat_array", "lon_array", "time_array") +#: Raw interpolation inputs, shaped ``(raw_level, lat_pair, lon_pair)``. +_CORNER_METADATA = ("geopotentials", "wind_us", "wind_vs", "temperatures") + + +# --------------------------------------------------------------------------- +# Cache location and policy +# --------------------------------------------------------------------------- + + +def is_cache_enabled() -> bool: + """Return False when ``ROCKETPY_CACHE`` opts out of disk caching.""" + raw = os.environ.get(CACHE_ENV_VAR) + if raw is None: + return True + return raw.strip().lower() not in _DISABLED_VALUES def get_cache_root() -> Path: """Return the root cache directory (honors ``ROCKETPY_CACHE``).""" - return Path(os.environ.get(CACHE_ENV_VAR, DEFAULT_CACHE_ROOT)).expanduser() + return Path(os.environ.get(CACHE_ENV_VAR) or DEFAULT_CACHE_ROOT).expanduser() def get_atmosphere_cache_dir() -> Path: @@ -38,14 +87,52 @@ def get_atmosphere_cache_dir() -> Path: return get_cache_root() / "atmosphere" +def get_cache_ttl() -> float: + """Return the cache lifetime in seconds (``0`` disables expiry).""" + raw = os.environ.get(CACHE_TTL_ENV_VAR) + if raw is None: + return float(DEFAULT_CACHE_TTL_SECONDS) + try: + return max(float(raw), 0.0) + except (TypeError, ValueError): + warnings.warn( + f"Invalid {CACHE_TTL_ENV_VAR}='{raw}'. " + f"Using the default of {DEFAULT_CACHE_TTL_SECONDS} seconds.", + UserWarning, + stacklevel=2, + ) + return float(DEFAULT_CACHE_TTL_SECONDS) + + +def is_entry_expired(created_at, kind) -> bool: + """Return True when a cache entry written at ``created_at`` is too old. + + Entries whose ``kind`` is in :data:`IMMUTABLE_KINDS` never expire, and a + TTL of ``0`` disables expiry for every kind. + """ + if kind in IMMUTABLE_KINDS: + return False + ttl = get_cache_ttl() + if ttl <= 0: + return False + try: + age = time.time() - float(created_at) + except (TypeError, ValueError): + return True # unreadable timestamp: treat as stale and re-fetch + return age > ttl + + def ensure_atmosphere_cache_dir() -> Path | None: """Create the atmosphere cache directory. Returns ------- pathlib.Path or None - The directory path, or ``None`` if creation failed (caching disabled). + The directory path, or ``None`` if caching is disabled or the + directory could not be created. """ + if not is_cache_enabled(): + return None cache_dir = get_atmosphere_cache_dir() try: cache_dir.mkdir(parents=True, exist_ok=True) @@ -60,6 +147,23 @@ def ensure_atmosphere_cache_dir() -> Path | None: return None +def clear_atmosphere_cache() -> bool: + """Delete every cached atmosphere file. Returns False on failure.""" + cache_dir = get_atmosphere_cache_dir() + if not cache_dir.is_dir(): + return True + try: + shutil.rmtree(cache_dir) + return True + except OSError as exc: + warnings.warn( + f"Could not clear atmosphere cache '{cache_dir}': {exc}.", + UserWarning, + stacklevel=2, + ) + return False + + def sanitize_cache_key(key: str) -> str: """Replace characters that are unsafe in filenames.""" return re.sub(r"[^A-Za-z0-9_.-]", "_", key) @@ -78,15 +182,22 @@ def build_atmosphere_cache_key( latitude: float, longitude: float, datetime_date, + variant: str = "", ) -> str: - """Build a stable cache key for a Forecast/Ensemble/Windy request.""" + """Build a stable cache key for a Forecast/Ensemble/Windy request. + + ``variant`` distinguishes requests that hit the same source and location + but decode it differently (a different variable dictionary or pressure + conversion factor), which would otherwise collide on one file. + """ if datetime_date is None: date_part = "nodate" else: date_part = datetime_date.strftime("%Y%m%d%H") - return sanitize_cache_key( - f"{kind}_{source}_{latitude:.4f}_{longitude:.4f}_{date_part}" - ) + key = f"{kind}_{source}_{latitude:.4f}_{longitude:.4f}_{date_part}" + if variant: + key = f"{key}_{variant}" + return sanitize_cache_key(key) def is_remote_url(path_or_url) -> bool: @@ -97,11 +208,16 @@ def is_remote_url(path_or_url) -> bool: return lowered.startswith(("http://", "https://", "dods://")) +# --------------------------------------------------------------------------- +# Raw byte / JSON helpers +# --------------------------------------------------------------------------- + + def atomic_write_bytes(path: Path, data: bytes) -> bool: """Write ``data`` to ``path`` atomically. Returns False on failure.""" - cache_dir = ensure_atmosphere_cache_dir() - if cache_dir is None: + if ensure_atmosphere_cache_dir() is None: return False + temp_name = None try: path.parent.mkdir(parents=True, exist_ok=True) with tempfile.NamedTemporaryFile( @@ -117,20 +233,27 @@ def atomic_write_bytes(path: Path, data: bytes) -> bool: UserWarning, stacklevel=2, ) - try: - Path(temp_name).unlink(missing_ok=True) - except (OSError, NameError): - pass + _discard(temp_name) return False -def load_json_cache(path: Path) -> dict | None: - """Load a JSON cache file, or ``None`` if missing/unreadable.""" - if not path.is_file(): +def _discard(path) -> None: + """Best-effort removal of a leftover temporary file.""" + if path is None: + return + try: + Path(path).unlink(missing_ok=True) + except OSError: + pass + + +def load_json_cache(path: Path, kind: str = "windy") -> dict | None: + """Load a JSON cache file, or ``None`` if missing/stale/unreadable.""" + if not is_cache_enabled() or not path.is_file(): return None try: with path.open("r", encoding="utf-8") as handle: - return json.load(handle) + envelope = json.load(handle) except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: warnings.warn( f"Failed to read cached atmosphere JSON '{path}': {exc}. " @@ -140,11 +263,28 @@ def load_json_cache(path: Path) -> dict | None: ) return None + if ( + not isinstance(envelope, dict) + or envelope.get(JSON_FORMAT_KEY) != PROFILE_FORMAT_ATTR + ): + return None + if is_entry_expired(envelope.get(CREATED_AT_ATTR), kind): + return None + payload = envelope.get("payload") + return payload if isinstance(payload, dict) else None + def save_json_cache(path: Path, payload: dict) -> bool: """Serialize ``payload`` as JSON to ``path``. Returns False on failure.""" + if not is_cache_enabled(): + return False + envelope = { + JSON_FORMAT_KEY: PROFILE_FORMAT_ATTR, + CREATED_AT_ATTR: time.time(), + "payload": payload, + } try: - data = json.dumps(payload).encode("utf-8") + data = json.dumps(envelope).encode("utf-8") except (TypeError, ValueError) as exc: warnings.warn( f"Failed to serialize atmosphere JSON for cache '{path}': {exc}.", @@ -155,6 +295,160 @@ def save_json_cache(path: Path, payload: dict) -> bool: return atomic_write_bytes(path, data) +# --------------------------------------------------------------------------- +# Model metadata persistence +# --------------------------------------------------------------------------- + + +def _write_metadata(dataset, metadata) -> None: + """Store the ``Environment`` model metadata on an open netCDF dataset.""" + metadata = metadata or {} + _write_metadata_attributes(dataset, metadata) + _write_metadata_arrays(dataset, metadata) + + +def _write_metadata_attributes(dataset, metadata) -> None: + """Store the scalar, date and coordinate-pair metadata as attributes.""" + for name, _ in _SCALAR_METADATA: + value = metadata.get(name) + if value is not None: + dataset.setncattr(name, float(value)) + + for name in _DATE_METADATA: + value = metadata.get(name) + if isinstance(value, datetime): + dataset.setncattr(name, value.strftime(_DATE_FORMAT)) + + for name in _PAIR_METADATA: + value = metadata.get(name) + if value is not None: + dataset.setncattr(name, [float(item) for item in value]) + + +def _write_metadata_arrays(dataset, metadata) -> None: + """Store the raw interpolation inputs as netCDF variables.""" + raw_levels = metadata.get("levels") + if raw_levels is None: + return + + raw_levels = np.asarray(raw_levels) + dataset.setncattr( + "levels_integer", int(np.issubdtype(raw_levels.dtype, np.integer)) + ) + dataset.createDimension("raw_level", raw_levels.size) + dataset.createDimension("lat_pair", 2) + dataset.createDimension("lon_pair", 2) + + variable = dataset.createVariable("raw_levels", "f8", ("raw_level",)) + variable[:] = np.asarray(raw_levels, dtype=float) + + raw_height = metadata.get("height") + if raw_height is not None: + variable = dataset.createVariable("raw_height", "f8", ("raw_level",)) + variable[:] = _filled(raw_height) + + for name in _CORNER_METADATA: + values = metadata.get(name) + if values is None: + continue + variable = dataset.createVariable( + name, "f8", ("raw_level", "lat_pair", "lon_pair") + ) + variable[:] = _filled(values) + + +def _read_metadata(dataset) -> dict: + """Rebuild the ``Environment`` model metadata from an open netCDF dataset.""" + metadata = {} + + for name, caster in _SCALAR_METADATA: + if hasattr(dataset, name): + metadata[name] = caster(dataset.getncattr(name)) + + for name in _DATE_METADATA: + if hasattr(dataset, name): + metadata[name] = datetime.strptime(dataset.getncattr(name), _DATE_FORMAT) + + for name in _PAIR_METADATA: + if hasattr(dataset, name): + metadata[name] = [ + float(item) for item in np.atleast_1d(dataset.getncattr(name)) + ] + + if "raw_levels" in dataset.variables: + levels = np.array(dataset.variables["raw_levels"][:], dtype=float) + if int(getattr(dataset, "levels_integer", 0)): + levels = levels.astype(np.int64) + metadata["levels"] = levels + + if "raw_height" in dataset.variables: + metadata["height"] = np.array(dataset.variables["raw_height"][:], dtype=float) + + for name in _CORNER_METADATA: + if name in dataset.variables: + metadata[name] = np.array(dataset.variables[name][:], dtype=float) + + return metadata + + +def _filled(values): + """Return a plain float array, replacing any masked entries with NaN.""" + return np.ma.filled(np.ma.asarray(values).astype(float), np.nan) + + +def _open_for_write(path: Path): + """Create a temporary netCDF file next to ``path``. Returns (dataset, temp).""" + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + dir=path.parent, delete=False, suffix=".nc.tmp" + ) as handle: + temp_path = Path(handle.name) + return netCDF4.Dataset(temp_path, mode="w", format="NETCDF4"), temp_path + + +def _set_common_attributes(dataset, kind, elevation, max_expected_height) -> None: + """Write the attributes shared by every cache file format.""" + dataset.setncattr("rocketpy_cache_format", PROFILE_FORMAT_ATTR) + dataset.setncattr("rocketpy_cache_kind", kind) + dataset.setncattr(CREATED_AT_ATTR, float(time.time())) + dataset.setncattr("elevation", float(elevation)) + dataset.setncattr("max_expected_height", float(max_expected_height)) + + +def _open_valid_cache(path: Path, expected_kind=None): + """Open a cache file, returning ``None`` if absent, foreign or expired.""" + if not is_cache_enabled() or not path.is_file(): + return None + try: + dataset = netCDF4.Dataset(path, mode="r") + except Exception as exc: # pylint: disable=broad-except + # netCDF4 raises OSError, RuntimeError or others for damaged files. A + # cache miss must never be louder than the download it replaces. + warnings.warn( + f"Failed to open atmosphere cache '{path}': {exc}. Fetching fresh data.", + UserWarning, + stacklevel=2, + ) + return None + + fmt = getattr(dataset, "rocketpy_cache_format", None) + kind = getattr(dataset, "rocketpy_cache_kind", None) + if fmt != PROFILE_FORMAT_ATTR or ( + expected_kind is not None and kind != expected_kind + ): + dataset.close() + return None + if is_entry_expired(getattr(dataset, CREATED_AT_ATTR, None), kind): + dataset.close() + return None + return dataset + + +# --------------------------------------------------------------------------- +# Forecast / Reanalysis profiles +# --------------------------------------------------------------------------- + + def write_profile_netcdf( path: Path, *, @@ -166,96 +460,60 @@ def write_profile_netcdf( elevation: float, max_expected_height: float, kind: str = "forecast", + metadata=None, ) -> bool: """Write extracted atmospheric profiles to a compact local netCDF file.""" - cache_dir = ensure_atmosphere_cache_dir() - if cache_dir is None: + if ensure_atmosphere_cache_dir() is None: return False - height = np.asarray(height, dtype=float) - pressure = np.asarray(pressure, dtype=float) - temperature = np.asarray(temperature, dtype=float) - wind_u = np.asarray(wind_u, dtype=float) - wind_v = np.asarray(wind_v, dtype=float) - + columns = ( + ("height", height, "m"), + ("pressure", pressure, "Pa"), + ("temperature", temperature, "K"), + ("wind_u", wind_u, "m s-1"), + ("wind_v", wind_v, "m s-1"), + ) + temp_path = None try: - path.parent.mkdir(parents=True, exist_ok=True) - with tempfile.NamedTemporaryFile( - dir=path.parent, delete=False, suffix=".nc.tmp" - ) as handle: - temp_path = Path(handle.name) - - dataset = netCDF4.Dataset(temp_path, mode="w", format="NETCDF4") + dataset, temp_path = _open_for_write(path) try: - dataset.setncattr("rocketpy_cache_format", PROFILE_FORMAT_ATTR) - dataset.setncattr("rocketpy_cache_kind", kind) - dataset.setncattr("elevation", float(elevation)) - dataset.setncattr("max_expected_height", float(max_expected_height)) - - dataset.createDimension("level", height.size) - for name, values, units in ( - ("height", height, "m"), - ("pressure", pressure, "Pa"), - ("temperature", temperature, "K"), - ("wind_u", wind_u, "m s-1"), - ("wind_v", wind_v, "m s-1"), - ): + _set_common_attributes(dataset, kind, elevation, max_expected_height) + dataset.createDimension("level", np.asarray(height, dtype=float).size) + for name, values, units in columns: variable = dataset.createVariable(name, "f8", ("level",)) variable.units = units - variable[:] = values + variable[:] = np.asarray(values, dtype=float) + _write_metadata(dataset, metadata) finally: dataset.close() - temp_path.replace(path) return True - except OSError as exc: + except Exception as exc: # pylint: disable=broad-except warnings.warn( f"Failed to write atmosphere profile cache '{path}': {exc}.", UserWarning, stacklevel=2, ) - try: - temp_path.unlink(missing_ok=True) - except (OSError, NameError): - pass + _discard(temp_path) return False def read_profile_netcdf(path: Path) -> dict | None: """Read a profile netCDF written by :func:`write_profile_netcdf`.""" - if not path.is_file(): - return None - try: - dataset = netCDF4.Dataset(path, mode="r") - except OSError as exc: - warnings.warn( - f"Failed to open atmosphere profile cache '{path}': {exc}. " - "Fetching fresh data.", - UserWarning, - stacklevel=2, - ) + dataset = _open_valid_cache(path) + if dataset is None: return None - try: - fmt = getattr(dataset, "rocketpy_cache_format", None) - if fmt != PROFILE_FORMAT_ATTR: - warnings.warn( - f"Ignoring atmosphere cache '{path}' with unknown format '{fmt}'.", - UserWarning, - stacklevel=2, - ) - return None - return { + profiles = { "kind": getattr(dataset, "rocketpy_cache_kind", "forecast"), "elevation": float(dataset.getncattr("elevation")), "max_expected_height": float(dataset.getncattr("max_expected_height")), - "height": np.array(dataset.variables["height"][:], dtype=float), - "pressure": np.array(dataset.variables["pressure"][:], dtype=float), - "temperature": np.array(dataset.variables["temperature"][:], dtype=float), - "wind_u": np.array(dataset.variables["wind_u"][:], dtype=float), - "wind_v": np.array(dataset.variables["wind_v"][:], dtype=float), } - except (AttributeError, KeyError, ValueError, OSError) as exc: + for name in ("height", "pressure", "temperature", "wind_u", "wind_v"): + profiles[name] = np.array(dataset.variables[name][:], dtype=float) + profiles["metadata"] = _read_metadata(dataset) + return profiles + except Exception as exc: # pylint: disable=broad-except warnings.warn( f"Failed to read atmosphere profile cache '{path}': {exc}. " "Fetching fresh data.", @@ -267,6 +525,11 @@ def read_profile_netcdf(path: Path) -> dict | None: dataset.close() +# --------------------------------------------------------------------------- +# Ensemble profiles +# --------------------------------------------------------------------------- + + def write_ensemble_profile_netcdf( path: Path, *, @@ -277,99 +540,75 @@ def write_ensemble_profile_netcdf( wind_v_ensemble, elevation: float, max_expected_height: float, + metadata=None, ) -> bool: """Write ensemble member profiles to a compact local netCDF file.""" - cache_dir = ensure_atmosphere_cache_dir() - if cache_dir is None: + if ensure_atmosphere_cache_dir() is None: return False - levels = np.asarray(levels, dtype=float) height_ensemble = np.asarray(height_ensemble, dtype=float) - temperature_ensemble = np.asarray(temperature_ensemble, dtype=float) - wind_u_ensemble = np.asarray(wind_u_ensemble, dtype=float) - wind_v_ensemble = np.asarray(wind_v_ensemble, dtype=float) + if height_ensemble.ndim != 2: + return False num_members, num_levels = height_ensemble.shape + columns = ( + ("height", height_ensemble, "m"), + ("temperature", temperature_ensemble, "K"), + ("wind_u", wind_u_ensemble, "m s-1"), + ("wind_v", wind_v_ensemble, "m s-1"), + ) + temp_path = None try: - path.parent.mkdir(parents=True, exist_ok=True) - with tempfile.NamedTemporaryFile( - dir=path.parent, delete=False, suffix=".nc.tmp" - ) as handle: - temp_path = Path(handle.name) - - dataset = netCDF4.Dataset(temp_path, mode="w", format="NETCDF4") + dataset, temp_path = _open_for_write(path) try: - dataset.setncattr("rocketpy_cache_format", PROFILE_FORMAT_ATTR) - dataset.setncattr("rocketpy_cache_kind", "ensemble") - dataset.setncattr("elevation", float(elevation)) - dataset.setncattr("max_expected_height", float(max_expected_height)) - + _set_common_attributes(dataset, "ensemble", elevation, max_expected_height) dataset.createDimension("member", num_members) dataset.createDimension("level", num_levels) level_var = dataset.createVariable("level", "f8", ("level",)) level_var.units = "Pa" - level_var[:] = levels - - for name, values, units in ( - ("height", height_ensemble, "m"), - ("temperature", temperature_ensemble, "K"), - ("wind_u", wind_u_ensemble, "m s-1"), - ("wind_v", wind_v_ensemble, "m s-1"), - ): + level_var[:] = np.asarray(levels, dtype=float) + + for name, values, units in columns: variable = dataset.createVariable(name, "f8", ("member", "level")) variable.units = units - variable[:] = values + variable[:] = np.asarray(values, dtype=float) + _write_metadata(dataset, metadata) finally: dataset.close() - temp_path.replace(path) return True - except OSError as exc: + except Exception as exc: # pylint: disable=broad-except warnings.warn( f"Failed to write ensemble atmosphere cache '{path}': {exc}.", UserWarning, stacklevel=2, ) - try: - temp_path.unlink(missing_ok=True) - except (OSError, NameError): - pass + _discard(temp_path) return False def read_ensemble_profile_netcdf(path: Path) -> dict | None: """Read an ensemble profile netCDF written by :func:`write_ensemble_profile_netcdf`.""" - if not path.is_file(): - return None - try: - dataset = netCDF4.Dataset(path, mode="r") - except OSError as exc: - warnings.warn( - f"Failed to open ensemble atmosphere cache '{path}': {exc}. " - "Fetching fresh data.", - UserWarning, - stacklevel=2, - ) + dataset = _open_valid_cache(path, expected_kind="ensemble") + if dataset is None: return None - try: - fmt = getattr(dataset, "rocketpy_cache_format", None) - kind = getattr(dataset, "rocketpy_cache_kind", None) - if fmt != PROFILE_FORMAT_ATTR or kind != "ensemble": - return None - return { + profiles = { "elevation": float(dataset.getncattr("elevation")), "max_expected_height": float(dataset.getncattr("max_expected_height")), "levels": np.array(dataset.variables["level"][:], dtype=float), - "height_ensemble": np.array(dataset.variables["height"][:], dtype=float), - "temperature_ensemble": np.array( - dataset.variables["temperature"][:], dtype=float - ), - "wind_u_ensemble": np.array(dataset.variables["wind_u"][:], dtype=float), - "wind_v_ensemble": np.array(dataset.variables["wind_v"][:], dtype=float), } - except (AttributeError, KeyError, ValueError, OSError) as exc: + for key, name in ( + ("height_ensemble", "height"), + ("temperature_ensemble", "temperature"), + ("wind_u_ensemble", "wind_u"), + ("wind_v_ensemble", "wind_v"), + ): + profiles[key] = np.array(dataset.variables[name][:], dtype=float) + profiles["metadata"] = _read_metadata(dataset) + return profiles + except Exception as exc: # pylint: disable=broad-except warnings.warn( f"Failed to read ensemble atmosphere cache '{path}': {exc}. " "Fetching fresh data.", diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index 7badffba9..8ac6f4033 100644 --- a/rocketpy/environment/environment.py +++ b/rocketpy/environment/environment.py @@ -17,6 +17,7 @@ from rocketpy.environment.atmosphere_cache import ( build_atmosphere_cache_key, cache_path_for, + is_cache_enabled, is_remote_url, load_json_cache, read_ensemble_profile_netcdf, @@ -1489,38 +1490,14 @@ def set_atmospheric_model( # pylint: disable=too-many-statements except KeyError: fetch_function = None - cache_path = self.__atmosphere_cache_path_for_request( - type, file, fetch_function + self.__load_or_fetch_atmospheric_model( + atm_type=type, + file=file, + dictionary=dictionary, + conversion_factor=conversion_factor, + fetch_function=fetch_function, + no_cache=no_cache, ) - loaded_from_cache = False - if cache_path is not None and not no_cache: - if type == "ensemble": - loaded_from_cache = self.__apply_cached_ensemble_profiles( - cache_path - ) - else: - loaded_from_cache = self.__apply_cached_forecast_profiles( - cache_path - ) - - if not loaded_from_cache: - # Fetches the dataset using OpenDAP protocol or uses the file path - dataset = fetch_function() if fetch_function is not None else file - - if type in ["forecast", "reanalysis"]: - self.process_forecast_reanalysis( - dataset, dictionary, conversion_factor=conversion_factor - ) - else: - self.process_ensemble( - dataset, dictionary, conversion_factor=conversion_factor - ) - - if cache_path is not None: - if type == "ensemble": - self.__save_ensemble_profiles_to_cache(cache_path) - else: - self.__save_forecast_profiles_to_cache(cache_path) ground_pressure = self.pressure(self.elevation) if not 30000 <= ground_pressure <= 120_000: @@ -1558,15 +1535,32 @@ def set_atmospheric_model( # pylint: disable=too-many-statements self.atmospheric_model_file = file self.atmospheric_model_dict = dictionary - def __atmosphere_cache_path_for_request(self, atm_type, file, fetch_function): + def __atmosphere_cache_path_for_request( + self, atm_type, file, fetch_function, dictionary, conversion_factor + ): """Return a cache path for remote Forecast/Ensemble sources, else None.""" + if not is_cache_enabled(): + return None + if fetch_function is not None and isinstance(file, str): source_label = file elif is_remote_url(file): - source_label = "url_" + hashlib.md5(file.encode("utf-8")).hexdigest()[:12] + source_label = ( + "url_" + + hashlib.md5(file.encode("utf-8"), usedforsecurity=False).hexdigest()[ + :12 + ] + ) else: return None + # The same source decoded with a different variable dictionary or + # pressure unit yields different profiles, so it needs its own entry. + variant = hashlib.md5( + repr((sorted(dictionary.items()), conversion_factor)).encode("utf-8"), + usedforsecurity=False, + ).hexdigest()[:8] + return cache_path_for( build_atmosphere_cache_key( atm_type, @@ -1574,10 +1568,91 @@ def __atmosphere_cache_path_for_request(self, atm_type, file, fetch_function): self.latitude, self.longitude, self.datetime_date, + variant=variant, ), ".nc", ) + def __load_or_fetch_atmospheric_model( + self, + *, + atm_type, + file, + dictionary, + conversion_factor, + fetch_function, + no_cache, + ): + """Apply a cached model when available, otherwise fetch and cache it.""" + cache_path = self.__atmosphere_cache_path_for_request( + atm_type, file, fetch_function, dictionary, conversion_factor + ) + is_ensemble = atm_type == "ensemble" + + if cache_path is not None and not no_cache: + apply_cached = ( + self.__apply_cached_ensemble_profiles + if is_ensemble + else self.__apply_cached_forecast_profiles + ) + if apply_cached(cache_path): + return + + # Fetches the dataset using OpenDAP protocol or uses the file path + dataset = fetch_function() if fetch_function is not None else file + if is_ensemble: + self.process_ensemble( + dataset, dictionary, conversion_factor=conversion_factor + ) + else: + self.process_forecast_reanalysis( + dataset, dictionary, conversion_factor=conversion_factor + ) + + if cache_path is not None: + if is_ensemble: + self.__save_ensemble_profiles_to_cache(cache_path) + else: + self.__save_forecast_profiles_to_cache(cache_path, atm_type) + + #: Attributes ``Environment`` derives from a Forecast/Ensemble dataset that + #: are not recoverable from the extracted profiles alone, so they travel + #: with the cache entry to keep a cache hit indistinguishable from a fetch. + __CACHED_MODEL_METADATA = ( + "atmospheric_model_init_date", + "atmospheric_model_end_date", + "atmospheric_model_interval", + "atmospheric_model_init_lat", + "atmospheric_model_end_lat", + "atmospheric_model_init_lon", + "atmospheric_model_end_lon", + "lat_array", + "lon_array", + "lat_index", + "lon_index", + "geopotentials", + "wind_us", + "wind_vs", + "levels", + "temperatures", + "time_array", + "height", + ) + + def __collect_model_metadata(self): + """Snapshot the model metadata that must survive a cache round-trip.""" + return { + name: getattr(self, name) + for name in self.__CACHED_MODEL_METADATA + if getattr(self, name, None) is not None + } + + def __restore_model_metadata(self, metadata): + """Reinstate the model metadata recovered from a cache entry.""" + for name, value in (metadata or {}).items(): + if name in self.__CACHED_MODEL_METADATA: + setattr(self, name, value) + def __apply_profiles_from_arrays( self, height, pressure, temperature, wind_u, wind_v ): @@ -1619,26 +1694,46 @@ def __apply_cached_forecast_profiles(self, cache_path): ) self.elevation = profiles["elevation"] self._max_expected_height = profiles["max_expected_height"] + self.__restore_model_metadata(profiles.get("metadata")) return True - def __save_forecast_profiles_to_cache(self, cache_path): + @staticmethod + def __profile_column(function, column=1): + """Return one column of an array-backed Function, or None.""" + if not isinstance(function, Function) or not function.is_array_source(): + return None + source = np.asarray(function.source, dtype=float) + if source.ndim != 2 or source.shape[1] <= column: + return None + return source[:, column] + + def __save_forecast_profiles_to_cache(self, cache_path, kind="forecast"): """Persist the active Forecast/Reanalysis profiles to ``cache_path``.""" - if not self.pressure.is_array_source(): + heights = self.__profile_column(self.pressure, column=0) + columns = { + "pressure": self.__profile_column(self.pressure), + "temperature": self.__profile_column(self.temperature), + "wind_u": self.__profile_column(self.wind_velocity_x), + "wind_v": self.__profile_column(self.wind_velocity_y), + } + # Every profile must be array-backed and share the pressure grid; + # constant (scalar) profiles carry nothing worth caching. + if heights is None or any( + column is None or column.shape != heights.shape + for column in columns.values() + ): return - pressure_source = np.asarray(self.pressure.source, dtype=float) - temperature_source = np.asarray(self.temperature.source, dtype=float) - wind_u_source = np.asarray(self.wind_velocity_x.source, dtype=float) - wind_v_source = np.asarray(self.wind_velocity_y.source, dtype=float) + write_profile_netcdf( cache_path, - height=pressure_source[:, 0], - pressure=pressure_source[:, 1], - temperature=temperature_source[:, 1], - wind_u=wind_u_source[:, 1], - wind_v=wind_v_source[:, 1], + height=heights, elevation=float(self.elevation), - max_expected_height=float(self._max_expected_height), - kind="forecast", + max_expected_height=float( + getattr(self, "_max_expected_height", self.max_expected_height) + ), + kind=kind, + metadata=self.__collect_model_metadata(), + **columns, ) def __apply_cached_ensemble_profiles(self, cache_path): @@ -1651,29 +1746,27 @@ def __apply_cached_ensemble_profiles(self, cache_path): temper = profiles["temperature_ensemble"] wind_u = profiles["wind_u_ensemble"] wind_v = profiles["wind_v_ensemble"] - levels = profiles["levels"] - - wind_speed = calculate_wind_speed(wind_u, wind_v) - wind_heading = calculate_wind_heading(wind_u, wind_v) - wind_direction = convert_wind_heading_to_direction(wind_heading) - self.level_ensemble = levels + self.level_ensemble = profiles["levels"] self.height_ensemble = height self.temperature_ensemble = temper self.wind_u_ensemble = wind_u self.wind_v_ensemble = wind_v - self.wind_heading_ensemble = wind_heading - self.wind_direction_ensemble = wind_direction - self.wind_speed_ensemble = wind_speed + self.wind_heading_ensemble = calculate_wind_heading(wind_u, wind_v) + self.wind_direction_ensemble = convert_wind_heading_to_direction( + self.wind_heading_ensemble + ) + self.wind_speed_ensemble = calculate_wind_speed(wind_u, wind_v) self.num_ensemble_members = height.shape[0] self.elevation = profiles["elevation"] self._max_expected_height = profiles["max_expected_height"] + self.__restore_model_metadata(profiles.get("metadata")) self.select_ensemble_member() return True def __save_ensemble_profiles_to_cache(self, cache_path): """Persist Ensemble member profiles to ``cache_path``.""" - if not hasattr(self, "height_ensemble"): + if getattr(self, "height_ensemble", None) is None: return write_ensemble_profile_netcdf( cache_path, @@ -1686,6 +1779,7 @@ def __save_ensemble_profiles_to_cache(self, cache_path): max_expected_height=float( getattr(self, "_max_expected_height", self.max_expected_height) ), + metadata=self.__collect_model_metadata(), ) # Atmospheric model processing methods @@ -1874,7 +1968,7 @@ def process_windy_atmosphere( # pylint: disable=too-many-statements ), ".json", ) - response = None if no_cache else load_json_cache(cache_path) + response = None if no_cache else load_json_cache(cache_path, kind="windy") if response is None: response = fetch_atmospheric_data_from_windy( self.latitude, self.longitude, model diff --git a/tests/conftest.py b/tests/conftest.py index a12c683e2..621ddc2a5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,18 @@ # Configure matplotlib to use non-interactive backend for tests matplotlib.use("Agg") + +@pytest.fixture(autouse=True) +def isolate_atmosphere_cache(monkeypatch, tmp_path): + """Keep the atmosphere disk cache out of the developer's home directory. + + Without this, ``set_atmospheric_model`` would write to + ``~/.rocketpy_cache`` during the test run and later tests could silently + read profiles cached by an earlier one. + """ + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path / "rocketpy_cache")) + + # Pytest configuration pytest_plugins = [ "tests.fixtures.environment.environment_fixtures", diff --git a/tests/unit/environment/test_atmosphere_cache.py b/tests/unit/environment/test_atmosphere_cache.py index d880b9a2a..df67819bf 100644 --- a/tests/unit/environment/test_atmosphere_cache.py +++ b/tests/unit/environment/test_atmosphere_cache.py @@ -1,5 +1,7 @@ """Unit tests for atmosphere netCDF/JSON disk caching (#654).""" +import time +import warnings from datetime import datetime from unittest.mock import MagicMock @@ -185,3 +187,301 @@ def test_windy_json_cache_hit(monkeypatch, tmp_path): env3.set_date(datetime(2018, 10, 15, 12)) env3.set_atmospheric_model(type="Windy", file="ECMWF", no_cache=True) assert fetch_mock.call_count == 2 + + +# --------------------------------------------------------------------------- +# Regression tests for the cache-hit / fresh-fetch parity contract +# --------------------------------------------------------------------------- + + +#: Attributes ``Environment`` derives from the dataset. A cache hit must +#: reproduce every one of them, otherwise ``info()`` and ``to_dict()`` break on +#: the second run of an otherwise identical script. +DERIVED_MODEL_ATTRIBUTES = [ + "atmospheric_model_init_date", + "atmospheric_model_end_date", + "atmospheric_model_interval", + "atmospheric_model_init_lat", + "atmospheric_model_end_lat", + "atmospheric_model_init_lon", + "atmospheric_model_end_lon", + "lat_array", + "lon_array", + "lat_index", + "lon_index", + "geopotentials", + "wind_us", + "wind_vs", + "levels", + "temperatures", + "time_array", + "height", +] + + +def _forecast_env(fetch): + """Build an Environment wired to ``fetch`` and load the Forecast model.""" + env = Environment( + latitude=32.990254, + longitude=-106.974998, + elevation=1400, + datum="WGS84", + ) + env.set_date((2018, 10, 15, 12)) + env._Environment__atm_type_file_to_function_map["forecast"]["GFS"] = fetch + env.set_atmospheric_model( + type="Forecast", + file="GFS", + dictionary="ECMWF", + pressure_conversion_factor="hPa", + ) + return env + + +def _counting_fetch(calls, fixture="data/weather/SpaceportAmerica_2018_ERA-5.nc"): + def fetch(): + calls.append(1) + return netCDF4.Dataset(fixture) + + return fetch + + +def test_cache_hit_restores_every_derived_attribute(monkeypatch, tmp_path): + """A cache hit must rebuild the same Environment a fresh fetch produces.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + calls = [] + fetch = _counting_fetch(calls) + + fresh = _forecast_env(fetch) + cached = _forecast_env(fetch) + assert len(calls) == 1, "second call should have been served from disk" + + for name in DERIVED_MODEL_ATTRIBUTES: + expected = getattr(fresh, name) + actual = getattr(cached, name, None) + assert actual is not None, f"'{name}' was lost on the cached path" + if isinstance(expected, datetime): + assert actual == expected, name + else: + np.testing.assert_allclose( + np.ma.filled(np.ma.asarray(actual, dtype=float), np.nan), + np.ma.filled(np.ma.asarray(expected, dtype=float), np.nan), + rtol=1e-10, + err_msg=name, + ) + + +def test_cache_hit_environment_can_print_info(monkeypatch, tmp_path): + """``info()`` used to raise AttributeError on the second (cached) run.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + calls = [] + fetch = _counting_fetch(calls) + + _forecast_env(fetch) + cached = _forecast_env(fetch) + + assert len(calls) == 1 + cached.info() # must not raise + + +def test_cache_hit_matches_fresh_profiles(monkeypatch, tmp_path): + """Profiles served from disk are numerically identical to a fresh fetch.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + calls = [] + fetch = _counting_fetch(calls) + + fresh = _forecast_env(fetch) + cached = _forecast_env(fetch) + + heights = np.linspace(fresh.elevation, fresh.max_expected_height, 25) + for name in ( + "pressure", + "temperature", + "wind_velocity_x", + "wind_velocity_y", + "wind_speed", + "wind_heading", + "wind_direction", + ): + np.testing.assert_allclose( + [getattr(cached, name)(h) for h in heights], + [getattr(fresh, name)(h) for h in heights], + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + + +def test_constant_wind_profile_does_not_break_caching(monkeypatch, tmp_path): + """Saving must tolerate scalar profiles instead of raising IndexError. + + ``set_atmospheric_model`` used to index every profile as a 2-D array while + only checking ``pressure``, so a constant wind blew up the cache write. + """ + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + env = Environment(latitude=0, longitude=0, elevation=0) + env._Environment__atm_type_file_to_function_map = { + "forecast": {"GFS": lambda: "fake-dataset"}, + "ensemble": {}, + } + env.process_forecast_reanalysis = lambda dataset, dictionary, conversion_factor: ( + None + ) + + env.set_atmospheric_model(type="Forecast", file="gfs") # must not raise + + assert not list((tmp_path / "atmosphere").glob("*.nc")), ( + "nothing worth caching should have been written" + ) + + +def test_cache_disabled_by_environment_variable(monkeypatch, tmp_path): + """``ROCKETPY_CACHE=0`` turns the disk cache off entirely.""" + monkeypatch.setenv("ROCKETPY_CACHE", "0") + assert atmosphere_cache.is_cache_enabled() is False + + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + assert atmosphere_cache.is_cache_enabled() is True + + +def test_disabled_cache_refetches_every_time(monkeypatch, tmp_path): + """With caching off, a repeated request hits the network again.""" + monkeypatch.setenv("ROCKETPY_CACHE", "off") + calls = [] + fetch = _counting_fetch(calls) + + _forecast_env(fetch) + _forecast_env(fetch) + + assert len(calls) == 2 + assert not list(tmp_path.rglob("*.nc")) + + +def test_expired_forecast_entry_is_refetched(monkeypatch, tmp_path): + """Forecast entries older than the TTL are discarded, not served.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + monkeypatch.setenv("ROCKETPY_CACHE_TTL", "3600") + calls = [] + fetch = _counting_fetch(calls) + + _forecast_env(fetch) + assert len(calls) == 1 + + # Pretend the entry was written two hours ago. The real clock has to be + # sampled before patching, or the stub would call itself. + two_hours_from_now = time.time() + 7200 + monkeypatch.setattr(atmosphere_cache.time, "time", lambda: two_hours_from_now) + _forecast_env(fetch) + assert len(calls) == 2, "a stale forecast must not be reused" + + +def test_zero_ttl_disables_expiry(monkeypatch): + """``ROCKETPY_CACHE_TTL=0`` keeps entries forever.""" + monkeypatch.setenv("ROCKETPY_CACHE_TTL", "0") + assert atmosphere_cache.is_entry_expired(0.0, "forecast") is False + + +def test_reanalysis_entries_never_expire(monkeypatch): + """Reanalysis data is immutable, so the TTL does not apply to it.""" + monkeypatch.setenv("ROCKETPY_CACHE_TTL", "1") + assert atmosphere_cache.is_entry_expired(0.0, "reanalysis") is False + assert atmosphere_cache.is_entry_expired(0.0, "forecast") is True + + +def test_different_dictionary_uses_a_separate_entry(monkeypatch, tmp_path): + """The same source decoded with another dictionary must not collide.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + calls = [] + fetch = _counting_fetch(calls) + + _forecast_env(fetch) + + env = Environment( + latitude=32.990254, longitude=-106.974998, elevation=1400, datum="WGS84" + ) + env.set_date((2018, 10, 15, 12)) + env._Environment__atm_type_file_to_function_map["forecast"]["GFS"] = fetch + env.set_atmospheric_model( + type="Forecast", + file="GFS", + dictionary="ECMWF_v0", + pressure_conversion_factor="hPa", + ) + + assert len(calls) == 2, "a different dictionary must miss the cache" + assert len(list((tmp_path / "atmosphere").glob("*.nc"))) == 2 + + +def test_corrupt_cache_file_falls_back_to_fetch(monkeypatch, tmp_path): + """A damaged cache file degrades to a re-fetch instead of raising.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + calls = [] + fetch = _counting_fetch(calls) + + _forecast_env(fetch) + cached_file = next((tmp_path / "atmosphere").glob("*.nc")) + cached_file.write_bytes(b"this is not a netCDF file") + + with pytest.warns(UserWarning): + _forecast_env(fetch) + + assert len(calls) == 2 + + +def test_clear_atmosphere_cache_removes_entries(monkeypatch, tmp_path): + """``clear_atmosphere_cache`` empties the cache directory.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + path = atmosphere_cache.cache_path_for("forecast_clear_me", ".nc") + _write_minimal_profile_nc(path) + assert path.is_file() + + assert atmosphere_cache.clear_atmosphere_cache() is True + assert not path.is_file() + # Clearing an already-absent cache is not an error. + assert atmosphere_cache.clear_atmosphere_cache() is True + + +def test_json_cache_respects_ttl(monkeypatch, tmp_path): + """Windy JSON entries expire like the netCDF ones.""" + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + monkeypatch.setenv("ROCKETPY_CACHE_TTL", "3600") + path = atmosphere_cache.cache_path_for("windy_ttl", ".json") + payload = {"data": {"hours": [1, 2, 3]}} + + assert atmosphere_cache.save_json_cache(path, payload) + assert atmosphere_cache.load_json_cache(path) == payload + + two_hours_from_now = time.time() + 7200 + monkeypatch.setattr(atmosphere_cache.time, "time", lambda: two_hours_from_now) + assert atmosphere_cache.load_json_cache(path) is None + + +def test_mismatched_profiles_are_skipped_without_warning(monkeypatch, tmp_path): + """A scalar wind profile is skipped cleanly, not written and not warned about. + + The save path used to check only ``pressure`` before slicing all four + profiles as 2-D arrays, so a constant wind raised ``IndexError``. Guarding + only ``pressure`` is not enough either: ``np.asarray(None, dtype=float)`` + silently yields ``nan``, so a missing column would be persisted as a cache + entry full of NaN winds. Checking every column keeps the write from being + attempted at all, which is why this asserts on the absence of a warning and + not just of a file. + """ + monkeypatch.setenv("ROCKETPY_CACHE", str(tmp_path)) + env = Environment(latitude=0, longitude=0, elevation=0) + env.set_atmospheric_model( + type="custom_atmosphere", + pressure=[[0.0, 101325.0], [1000.0, 89875.0]], + temperature=[[0.0, 288.0], [1000.0, 281.0]], + wind_u=5, + wind_v=-3, + ) + assert env.pressure.is_array_source() + assert not env.wind_velocity_x.is_array_source() + + path = atmosphere_cache.cache_path_for("forecast_mismatched", ".nc") + with warnings.catch_warnings(): + warnings.simplefilter("error") + env._Environment__save_forecast_profiles_to_cache(path) + + assert not path.exists() From 021ff137415419a71cd8df58ff932fdfd3c6142b Mon Sep 17 00:00:00 2001 From: Gui-FernandesBR Date: Sat, 12 Sep 2026 03:37:04 -0300 Subject: [PATCH 4/4] MNT: drop four pylint suppressions that no longer suppress anything The lint job installs pylint unpinned and has not run green on develop since 2026-07-19. Pylint 4.x reports these four disables as useless-suppression, which fails the job with exit code 8 on every open PR regardless of what the PR changes -- two of them already appear in this PR's first CI run from 2026-08-15. These files are untouched by the caching work; this commit only removes the stale comments so the lint job can go green. It is separate so it can be dropped and landed on its own if maintainers prefer. Co-Authored-By: Claude Opus 5 (1M context) --- rocketpy/motors/motor.py | 1 - rocketpy/sensors/sensor.py | 2 +- rocketpy/simulation/monte_carlo.py | 2 +- tests/unit/test_plots.py | 2 +- 4 files changed, 3 insertions(+), 4 deletions(-) diff --git a/rocketpy/motors/motor.py b/rocketpy/motors/motor.py index 9ecb001dd..f27664ec6 100644 --- a/rocketpy/motors/motor.py +++ b/rocketpy/motors/motor.py @@ -1386,7 +1386,6 @@ class GenericMotor(Motor): therefore for more accurate results, use the ``SolidMotor``, ``HybridMotor`` or ``LiquidMotor`` classes.""" - # pylint: disable=too-many-arguments def __init__( self, thrust_source, diff --git a/rocketpy/sensors/sensor.py b/rocketpy/sensors/sensor.py index de3461d2e..4e83ff9d8 100644 --- a/rocketpy/sensors/sensor.py +++ b/rocketpy/sensors/sensor.py @@ -446,7 +446,7 @@ class InertialSensor(Sensor): temperature drift. """ - def __init__( # pylint: disable=too-many-arguments + def __init__( self, sampling_rate, orientation=(0, 0, 0), diff --git a/rocketpy/simulation/monte_carlo.py b/rocketpy/simulation/monte_carlo.py index 09198bded..94d4702f6 100644 --- a/rocketpy/simulation/monte_carlo.py +++ b/rocketpy/simulation/monte_carlo.py @@ -564,7 +564,7 @@ def __validate_number_of_workers(self, n_workers): raise ValueError("Number of workers must be at least 2 for parallel mode.") return n_workers - def __sim_producer(self, seed, sim_monitor, mutex, error_event): # pylint: disable=too-many-statements + def __sim_producer(self, seed, sim_monitor, mutex, error_event): """Simulation producer to be used in parallel by multiprocessing. Parameters diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py index d6a529e8b..3e34d2a97 100644 --- a/tests/unit/test_plots.py +++ b/tests/unit/test_plots.py @@ -450,7 +450,7 @@ def test_animation_options_validation_errors(kwargs, error): @patch("matplotlib.pyplot.show") @pytest.mark.parametrize("filename", [None, "test_cp_evolution.png"]) -def test_flight_center_of_pressure_plot(mock_show, filename, flight_calisto): # pylint: disable=unused-argument +def test_flight_center_of_pressure_plot(mock_show, filename, flight_calisto): """Center-of-pressure evolution plot runs for a fixture flight. Parameters