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