diff --git a/rocketpy/environment/atmosphere_cache.py b/rocketpy/environment/atmosphere_cache.py new file mode 100644 index 000000000..393ac6e5b --- /dev/null +++ b/rocketpy/environment/atmosphere_cache.py @@ -0,0 +1,620 @@ +"""Disk cache for downloaded atmospheric datasets (netCDF profiles and JSON). + +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, 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 + +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_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) or DEFAULT_CACHE_ROOT).expanduser() + + +def get_atmosphere_cache_dir() -> Path: + """Return the atmosphere subdirectory under the cache root.""" + 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 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) + 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 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) + + +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, + variant: str = "", +) -> str: + """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") + 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: + """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://")) + + +# --------------------------------------------------------------------------- +# Raw byte / JSON helpers +# --------------------------------------------------------------------------- + + +def atomic_write_bytes(path: Path, data: bytes) -> bool: + """Write ``data`` to ``path`` atomically. Returns False on failure.""" + if ensure_atmosphere_cache_dir() is None: + return False + temp_name = None + 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, + ) + _discard(temp_name) + return False + + +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: + envelope = 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 + + 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(envelope).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) + + +# --------------------------------------------------------------------------- +# 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, + *, + height, + pressure, + temperature, + wind_u, + wind_v, + elevation: float, + max_expected_height: float, + kind: str = "forecast", + metadata=None, +) -> bool: + """Write extracted atmospheric profiles to a compact local netCDF file.""" + if ensure_atmosphere_cache_dir() is None: + return False + + 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: + dataset, temp_path = _open_for_write(path) + try: + _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[:] = np.asarray(values, dtype=float) + _write_metadata(dataset, metadata) + finally: + dataset.close() + temp_path.replace(path) + return True + except Exception as exc: # pylint: disable=broad-except + warnings.warn( + f"Failed to write atmosphere profile cache '{path}': {exc}.", + UserWarning, + stacklevel=2, + ) + _discard(temp_path) + return False + + +def read_profile_netcdf(path: Path) -> dict | None: + """Read a profile netCDF written by :func:`write_profile_netcdf`.""" + dataset = _open_valid_cache(path) + if dataset is None: + return None + try: + profiles = { + "kind": getattr(dataset, "rocketpy_cache_kind", "forecast"), + "elevation": float(dataset.getncattr("elevation")), + "max_expected_height": float(dataset.getncattr("max_expected_height")), + } + 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.", + UserWarning, + stacklevel=2, + ) + return None + finally: + dataset.close() + + +# --------------------------------------------------------------------------- +# Ensemble profiles +# --------------------------------------------------------------------------- + + +def write_ensemble_profile_netcdf( + path: Path, + *, + levels, + height_ensemble, + temperature_ensemble, + wind_u_ensemble, + wind_v_ensemble, + elevation: float, + max_expected_height: float, + metadata=None, +) -> bool: + """Write ensemble member profiles to a compact local netCDF file.""" + if ensure_atmosphere_cache_dir() is None: + return False + + height_ensemble = np.asarray(height_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: + dataset, temp_path = _open_for_write(path) + try: + _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[:] = np.asarray(levels, dtype=float) + + for name, values, units in columns: + variable = dataset.createVariable(name, "f8", ("member", "level")) + variable.units = units + variable[:] = np.asarray(values, dtype=float) + _write_metadata(dataset, metadata) + finally: + dataset.close() + temp_path.replace(path) + return True + except Exception as exc: # pylint: disable=broad-except + warnings.warn( + f"Failed to write ensemble atmosphere cache '{path}': {exc}.", + UserWarning, + stacklevel=2, + ) + _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`.""" + dataset = _open_valid_cache(path, expected_kind="ensemble") + if dataset is None: + return None + try: + profiles = { + "elevation": float(dataset.getncattr("elevation")), + "max_expected_height": float(dataset.getncattr("max_expected_height")), + "levels": np.array(dataset.variables["level"][:], dtype=float), + } + 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.", + UserWarning, + stacklevel=2, + ) + return None + finally: + dataset.close() diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index 460f0bc89..8ac6f4033 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 @@ -13,6 +14,18 @@ import numpy as np import pytz +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, + 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, @@ -1234,6 +1247,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. @@ -1354,6 +1368,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 ------- @@ -1401,7 +1420,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}) @@ -1468,15 +1490,14 @@ 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 - - if type in ["forecast", "reanalysis"]: - self.process_forecast_reanalysis( - dataset, dictionary, conversion_factor=conversion_factor - ) - else: - self.process_ensemble(dataset, dictionary, conversion_factor) + 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, + ) ground_pressure = self.pressure(self.elevation) if not 30000 <= ground_pressure <= 120_000: @@ -1514,6 +1535,253 @@ 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, 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"), 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, + source_label, + 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 + ): + """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"] + self.__restore_model_metadata(profiles.get("metadata")) + return True + + @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``.""" + 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 + + write_profile_netcdf( + cache_path, + height=heights, + elevation=float(self.elevation), + 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): + """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"] + + 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 = 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 getattr(self, "height_ensemble", None) is None: + 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) + ), + metadata=self.__collect_model_metadata(), + ) + # Atmospheric model processing methods def process_standard_atmosphere(self): @@ -1662,7 +1930,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 @@ -1672,6 +1942,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 ------ @@ -1686,9 +1958,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, kind="windy") + 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/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/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 new file mode 100644 index 000000000..df67819bf --- /dev/null +++ b/tests/unit/environment/test_atmosphere_cache.py @@ -0,0 +1,487 @@ +"""Unit tests for atmosphere netCDF/JSON disk caching (#654).""" + +import time +import warnings +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 + + +# --------------------------------------------------------------------------- +# 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() 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