diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 09dbbcd..351083c 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -36,7 +36,10 @@ jobs: submodules: true coverage: codecov toxdeps: tox-pypi-filter - posargs: -n auto --color=yes + # No `-n auto` here: the online env forces `-n 1` (serial) in tox.ini to + # avoid the data server truncating concurrent FITS downloads. Passing + # `-n auto` would override that. + posargs: --color=yes envs: | - linux: py313-online secrets: diff --git a/.gitignore b/.gitignore index a4cfc2a..db317d5 100644 --- a/.gitignore +++ b/.gitignore @@ -92,7 +92,7 @@ ipython_config.py # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: -# .python-version +.python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. @@ -264,3 +264,4 @@ package.json # stixpy stixpy/data/*.fits .vscode/* +docs/generated/* diff --git a/changelog/165.bugfix.rst b/changelog/165.bugfix.rst new file mode 100644 index 0000000..e3b7527 --- /dev/null +++ b/changelog/165.bugfix.rst @@ -0,0 +1 @@ +Fixed `~stixpy.coordinates.transforms.get_hpc_info` not hitting the ephemeris cache for scalar point queries (each call re-ran a ``Fido`` search even when the day's file was already cached) and not hitting the cache for arrays of times either (the strict equality check raised on broadcast). Cache reads now use a nearest-row lookup with a configurable maximum distance. diff --git a/changelog/165.feature.rst b/changelog/165.feature.rst new file mode 100644 index 0000000..a20b8d3 --- /dev/null +++ b/changelog/165.feature.rst @@ -0,0 +1 @@ +Add an internal ephemeris data cache that lets coordinate-transform and imaging code paths avoid redundant ``Fido`` searches and downloads. ANC ephemeris FITS files are daily, so touching any time inside a day caches the whole day. Users can also pre-warm specific files via `~stixpy.coordinates.transforms.load_ephemeris_fits_to_cache`. Coverage detection now uses actual inter-row time gaps rather than a row-count heuristic, and queries near UTC midnight pull the neighbouring daily file. diff --git a/changelog/165.trivial.rst b/changelog/165.trivial.rst new file mode 100644 index 0000000..e51bf97 --- /dev/null +++ b/changelog/165.trivial.rst @@ -0,0 +1,3 @@ +The internal ephemeris ``TableLRUCache`` now uses binary search (``numpy.searchsorted``) over the sorted ``time`` column for point and range lookups instead of scanning every cached row, speeding up the common pattern of many queries against a large cache. Internal change only; query semantics are unchanged. + +CI: online (remote-data) tests now run serially (``-n 1``), avoiding the remote data server truncating large FITS responses under the concurrency of ``-n auto``. CI-only change. diff --git a/examples/ephemeris_demo.py b/examples/ephemeris_demo.py new file mode 100644 index 0000000..fb9a8bb --- /dev/null +++ b/examples/ephemeris_demo.py @@ -0,0 +1,242 @@ +""" +============== +Ephemeris Demo +============== + +How to query STIX ephemeris (SOLO pointing, roll, and HEEQ position) via +`~stixpy.coordinates.transforms.get_hpc_info`, and how the internal cache +keeps subsequent calls cheap by avoiding redundant Fido searches. + +ANC ephemeris FITS files are organised as daily files. The first query that +touches a given day downloads that whole file via Fido and writes every row +of it into the global cache. Any later query inside the same day is served +straight from memory — no network, no file I/O. + +Imports +""" + +import logging + +import matplotlib.pyplot as plt +import numpy as np + +import astropy.units as u +from astropy.coordinates import SkyCoord +from astropy.time import Time + +from sunpy.coordinates import Helioprojective +from sunpy.net import Fido +from sunpy.net import attrs as a + +from stixpy.coordinates.frames import STIXImaging +from stixpy.coordinates.transforms import STIX_EPHEMERIS_CACHE, get_hpc_info, load_ephemeris_fits_to_cache + +logger = logging.getLogger(__name__) + +############################################################################### +# Start from an empty cache so the demo's behaviour is reproducible. + +STIX_EPHEMERIS_CACHE.clear() +print(f"Rows in cache: {len(STIX_EPHEMERIS_CACHE)}") + +############################################################################### +# Case 1 — single time point +# -------------------------- +# +# Pass a scalar `~astropy.time.Time`. ``get_hpc_info`` interpolates the +# ANC data onto that exact instant. The first call to any time in +# 2023-01-01 triggers a Fido search and caches the entire daily file. + +t = Time("2023-01-01T12:00:00") +roll, solo_heeq, stix_pointing = get_hpc_info(t) +print(f"At {t}: roll={roll}, pointing={stix_pointing}") +print(f"Rows in cache after first query: {len(STIX_EPHEMERIS_CACHE)}") + +############################################################################### +# Case 2 — time range +# ------------------- +# +# A scalar ``times`` plus an ``end_time`` averages the ephemeris across +# ``[start, end]``. This is the calling pattern used by imaging code +# (see `~stixpy.calibration.visibility.calibrate_visibility`) where one +# averaged pointing is needed for an integration window. +# +# This call lands entirely inside the day already cached above, so no +# Fido search runs — the lookup is in-memory. + +start = Time("2023-01-01T15:20:00") +end = Time("2023-01-01T15:23:00") +roll, solo_heeq, stix_pointing = get_hpc_info(start, end) +print(f"Averaged over [{start}, {end}]: roll={roll}, pointing={stix_pointing}") + +############################################################################### +# Case 3 — array of times +# ----------------------- +# +# Passing an array returns one interpolated value per input time. This is +# what `astropy` coordinate frame transforms use under the hood when +# transforming a `~stixpy.coordinates.frames.STIXImaging` `~astropy.coordinates.SkyCoord` +# with a vector ``obstime``. + +times = Time("2023-01-01") + np.arange(0, 24, 2) * u.h # 12 points across the day +roll_array, _, _ = get_hpc_info(times) +print(f"Roll across the day shape: {roll_array.shape}") + +############################################################################### +# Inspect the cache +# ----------------- +# +# After the three calls above, the cache holds the full content of the +# daily ANC file for 2023-01-01 — a single Fido download served all three +# queries. + +print(f"Rows in cache: {len(STIX_EPHEMERIS_CACHE)}") +print(f"Source files in cache: {set(STIX_EPHEMERIS_CACHE.cache['__source'])}") + +############################################################################### +# Plot the cached roll-angle across the day to visualise that the whole +# day was cached after the first query. + +day_start = Time("2023-01-01T00:00:00") +cached_times = STIX_EPHEMERIS_CACHE.cache["time"] +cached_roll = STIX_EPHEMERIS_CACHE.cache["roll_angle_rpy"][:, 0] +cached_hours = (cached_times - day_start).to_value(u.h) +query_hours = (times - day_start).to_value(u.h) + +fig, ax = plt.subplots(figsize=(10, 4), layout="constrained") +ax.plot(cached_hours, cached_roll.to_value(u.deg), ".", markersize=2, label="cached rows") +ax.plot(query_hours, roll_array.to_value(u.deg), "o", color="C1", label="get_hpc_info(times)") +ax.set_xlabel("Hours since 2023-01-01 00:00 UTC") +ax.set_ylabel("Roll [deg]") +ax.set_title("Full day cached after first get_hpc_info call") +ax.legend() + +############################################################################### +# Coordinate transforms use the cache implicitly +# ---------------------------------------------- +# +# Astropy frame transforms involving `~stixpy.coordinates.frames.STIXImaging` +# call ``get_hpc_info`` under the hood — see ``stixim_to_hpc`` and +# ``hpc_to_stixim`` in `~stixpy.coordinates.transforms`. Because the day +# is already cached above, every transform below is served from memory: +# the cache row count does not change. + +n_rows_before = len(STIX_EPHEMERIS_CACHE) + +# Sample STIX (0,0) every 30 minutes through 2023-01-01. +sample_times = Time("2023-01-01") + np.arange(0, 24, 0.5) * u.h +stix_zero = SkyCoord( + np.zeros(sample_times.size) * u.arcsec, + np.zeros(sample_times.size) * u.arcsec, + frame=STIXImaging(obstime=sample_times), +) +hpc_zero = stix_zero.transform_to(Helioprojective(obstime=sample_times)) + +n_rows_after = len(STIX_EPHEMERIS_CACHE) +print(f"Cache rows before transforms: {n_rows_before}") +print(f"Cache rows after transforms: {n_rows_after} (unchanged → cache served the transforms)") + +############################################################################### +# Plot where STIX (0,0) lands on the Helioprojective disk through the day. +# Pure SOLO orbit + roll evolution — every point uses cached ephemeris. + +fig, ax = plt.subplots(figsize=(7, 7), layout="constrained") +hours = (sample_times - sample_times[0]).to_value(u.h) +tx = hpc_zero.Tx.to_value(u.arcsec) +ty = hpc_zero.Ty.to_value(u.arcsec) +sc = ax.scatter(tx, ty, c=hours, cmap="viridis") +ax.set_aspect("equal") +ax.set_xlabel("HPC Tx [arcsec]") +ax.set_ylabel("HPC Ty [arcsec]") +ax.set_title("STIX (0,0) → Helioprojective across 2023-01-01\n(every transform served from cache)") + +ax.set_xlim(32, 34) +ax.set_ylim(54, 56) +fig.colorbar(sc, ax=ax, label="Hours since 00:00 UTC") + +############################################################################### +# Pre-warm the cache from local FITS files +# ---------------------------------------- +# +# If you already have ANC FITS files on disk (e.g. mounted archive, prior +# download, or an offline environment), pre-warm the cache directly with +# `~stixpy.coordinates.transforms.load_ephemeris_fits_to_cache`. After +# that, calls to ``get_hpc_info`` for times inside those days hit the +# cache without any Fido call. + +# Clear and demonstrate pre-warming using files we fetch ourselves. +STIX_EPHEMERIS_CACHE.clear() + +query = Fido.search( + a.Time("2023-01-02", "2023-01-03"), + a.Instrument.stix, + a.Level.anc, + a.stix.DataType.asp, + a.stix.DataProduct.asp_ephemeris, +) +query["stix"].filter_for_latest_version() +aux_files = Fido.fetch(query["stix"]) + +for path in aux_files: + load_ephemeris_fits_to_cache(path) + +print(f"Rows in cache after pre-warming 2 days: {len(STIX_EPHEMERIS_CACHE)}") + +# This query now hits the cache directly; no Fido search happens. +roll, _, _ = get_hpc_info(Time("2023-01-02T08:00:00")) +print(f"Pre-warmed lookup roll: {roll}") + +############################################################################### +# Point Fido at a local archive (no custom client needed) +# ------------------------------------------------------- +# +# The cache is the recommended way to bypass network downloads for known +# time ranges (pre-warm via +# `~stixpy.coordinates.transforms.load_ephemeris_fits_to_cache`). If you +# want a *system-wide* "always look in this local archive instead of the +# online server" override, override +# `~stixpy.net.client.STIXClient.baseurl` to a ``file://`` URL. The class +# attribute is what `~sunpy.net.Fido` reads when it instantiates the client +# with no ``source`` argument, so every `~sunpy.net.Fido.search` afterwards +# — including the one inside +# `~stixpy.coordinates._ephemeris_fetcher.fetch_ephemeris_for_range` — +# resolves files from your local mirror. + +if False: # disabled in the docs build; uncomment locally to redirect Fido + from stixpy.net.client import STIXClient + + STIXClient.baseurl = "file:///data/solo/fits" # ← your local mirror root + + # From here on, any get_hpc_info / coordinate-transform call that misses + # the cache will resolve files from the local directory rather than the + # public archive. No stixpy code change required. + roll, _, _ = get_hpc_info(Time("2023-01-04T08:00:00")) + +# A per-call alternative (no global state mutation) is to instantiate +# `~stixpy.net.client.STIXClient` directly and bypass Fido entirely — useful +# when you want to feed a specific set of files into the cache and stop +# there. See the pre-warming section above for the recommended pattern. + +############################################################################### +# Cache controls +# -------------- +# +# A few knobs are exposed on the global cache: +# +# * ``STIX_EPHEMERIS_CACHE.clear()`` — empty the cache. +# * ``STIX_EPHEMERIS_CACHE.bypass_cache_read = True`` — force the next read +# to miss so Fido re-runs (useful if you suspect cached data is stale or +# to compare cached vs fresh results in tests). +# * ``len(STIX_EPHEMERIS_CACHE)`` — current row count. +# +# Note: ``bypass_cache_read`` does the *opposite* of pre-warming. To skip +# Fido for a known range, pre-load the relevant files via +# ``load_ephemeris_fits_to_cache`` — the cache itself is the bypass +# mechanism. + +STIX_EPHEMERIS_CACHE.bypass_cache_read = True +roll_fresh, _, _ = get_hpc_info(Time("2023-01-02T08:00:00")) # re-downloads +STIX_EPHEMERIS_CACHE.bypass_cache_read = False +print(f"Forced-fresh roll: {roll_fresh}") + +plt.show() diff --git a/stixpy/config/data b/stixpy/config/data index a213fbe..7e6de14 160000 --- a/stixpy/config/data +++ b/stixpy/config/data @@ -1 +1 @@ -Subproject commit a213fbe50a84ea7d20a53c9e42aeb266d0532b4e +Subproject commit 7e6de14fff19f1bb6b1da47d56fcb376ac87c3ea diff --git a/stixpy/coordinates/_ephemeris_fetcher.py b/stixpy/coordinates/_ephemeris_fetcher.py new file mode 100644 index 0000000..b39dfbc --- /dev/null +++ b/stixpy/coordinates/_ephemeris_fetcher.py @@ -0,0 +1,127 @@ +""" +ANC ephemeris fetcher. + +Internal module that searches, downloads, and loads STIX ANC asp_ephemeris +FITS files into a single QTable. Owns Fido / version filtering / day-boundary +padding so the cache and orchestrator layers don't have to. +""" + +from pathlib import Path + +import astropy.units as u +from astropy.table import QTable, vstack +from astropy.time import Time + +from sunpy.net import Fido +from sunpy.net import attrs as a + +from stixpy.product.product_factory import Product +from stixpy.product.sources.anc import Ephemeris +from stixpy.utils.logging import get_logger +from stixpy.utils.table import drop_fits_checksums + +logger = get_logger(__name__) + +__all__ = ["fetch_ephemeris_for_range"] + + +def fetch_ephemeris_for_range( + start: Time, + end: Time, + *, + midnight_pad: u.Quantity = 90 * u.s, +) -> QTable: + """ + Search, version-filter, download, and load ANC asp_ephemeris files + covering ``[start, end]``. + + The Fido query window is extended by ``midnight_pad`` whenever ``start`` + or ``end`` falls within ``midnight_pad`` of UTC midnight, so a query + near a day boundary pulls the neighbour-day file too. Per the ephemeris + requirements document, the row closest to a query just after midnight + may live in the previous day's file (and vice versa). + + Parameters + ---------- + start, end + Scalar `~astropy.time.Time` defining the requested range. + midnight_pad + Extension applied to the Fido query window near day boundaries. + + Returns + ------- + QTable + Vstacked-and-sorted ``Ephemeris.data`` rows from all loaded files, + with a ``__source`` column (filename) per row. + + Raises + ------ + ValueError + If the Fido search returns no results or downloads fail. + """ + query_start, query_end = _apply_midnight_pad(start, end, midnight_pad) + + logger.info(f"Fido searching ANC asp_ephemeris for {query_start} – {query_end}") + query = Fido.search( + a.Time(query_start, query_end), + a.Instrument.stix, + a.Level.anc, + a.stix.DataType.asp, + a.stix.DataProduct.asp_ephemeris, + ) + + if len(query["stix"]) == 0: + raise ValueError(f"No STIX pointing data found for time range {query_start} to {query_end}.") + + query["stix"].filter_for_latest_version() + logger.debug(f"Downloading {len(query['stix'])} ANC files") + aux_files = Fido.fetch(query["stix"]) + if len(aux_files.errors) > 0: + raise ValueError("There were errors downloading the ANC data.") + + # NOTE: the data server truncates large FITS responses under concurrent load; + # CI sidesteps this by running the online tests serially (``-n 1``, see tox.ini). + # If end-user resilience is wanted, re-introduce a ``stixpy.utils.io.is_valid_fits`` + # check on ``aux_files`` here and re-fetch with + # ``Fido.fetch(query["stix"], overwrite=True)`` on corruption. + + tables = [] + for aux_file in aux_files: + aux_file = Path(aux_file) + anc = Product(aux_file, data_only=True) + if not isinstance(anc, Ephemeris): + logger.warning(f"Skipping non-Ephemeris ANC file: {aux_file.name}") + continue + anc.data["__source"] = aux_file.name + tables.append(anc.data) + + if not tables: + raise ValueError("No Ephemeris files could be loaded from the Fido response.") + + drop_fits_checksums(*tables) + ephemeris = vstack(tables) + ephemeris.sort(keys=["time"]) + return ephemeris + + +def _apply_midnight_pad(start: Time, end: Time, pad: u.Quantity) -> tuple[Time, Time]: + """ + Widen ``[start, end]`` toward the neighbour day(s) when either endpoint + is within ``pad`` of UTC midnight, so cross-day queries pick up the + neighbouring daily file. + """ + pad_s = pad.to_value(u.s) + seconds_per_day = 86400.0 + + start_sod = _seconds_of_day(start) + end_sod = _seconds_of_day(end) + + query_start = start - pad if start_sod < pad_s else start + query_end = end + pad if (seconds_per_day - end_sod) < pad_s else end + return query_start, query_end + + +def _seconds_of_day(t: Time) -> float: + """Return seconds elapsed since UTC midnight for ``t``.""" + ymdhms = t.utc.ymdhms + return float(ymdhms.hour * 3600 + ymdhms.minute * 60 + ymdhms.second) diff --git a/stixpy/coordinates/tests/test_transforms.py b/stixpy/coordinates/tests/test_transforms.py index eaa99ad..cb1a12a 100644 --- a/stixpy/coordinates/tests/test_transforms.py +++ b/stixpy/coordinates/tests/test_transforms.py @@ -5,13 +5,39 @@ import astropy.units as u from astropy.coordinates import SkyCoord +from astropy.table import QTable from astropy.tests.helper import assert_quantity_allclose from astropy.time import Time from sunpy.coordinates import HeliographicStonyhurst, Helioprojective from stixpy.coordinates.frames import STIXImaging -from stixpy.coordinates.transforms import _get_ephemeris_data, get_hpc_info +from stixpy.coordinates.transforms import STIX_EPHEMERIS_CACHE, _get_ephemeris_data, get_hpc_info + + +@pytest.fixture(autouse=True) +def _clear_ephemeris_cache(): + """Isolate every test from cache pollution carried over from earlier tests.""" + STIX_EPHEMERIS_CACHE.clear() + yield + STIX_EPHEMERIS_CACHE.clear() + + +def _fake_ephemeris(start, *, n_rows=2000, dt=64 * u.s): + """Synthesize an ANC-shaped ephemeris table for warm-cache tests.""" + times = start + np.arange(n_rows) * dt + return QTable( + { + "time": times, + "time_end": times + dt, + "timedel": np.full(n_rows, 64.0) * u.s, + "roll_angle_rpy": np.zeros((n_rows, 3)) * u.deg, + "solo_loc_heeq_zxy": np.tile([1.0e8, 0.0, 0.0], (n_rows, 1)) * u.km, + "y_srf": np.zeros(n_rows) * u.arcsec, + "z_srf": np.zeros(n_rows) * u.arcsec, + "sas_ok": np.ones(n_rows, dtype=int), + } + ) @pytest.mark.skip(reason="Test data maybe incorrect") @@ -86,19 +112,34 @@ def test_stx_to_hpc_obstime_end(): assert np.all(stix_coord.obstime_end.isclose(stix_coord_rt.obstime_end)) +@pytest.mark.remote_data +def test_stx_to_hpc_obstime_end_x2(): + # strange error that the call the second times crashes + test_stx_to_hpc_obstime_end() + test_stx_to_hpc_obstime_end() + + @pytest.mark.remote_data def test_get_aux_data(): with pytest.raises(ValueError, match="No STIX pointing data found for time range"): _get_ephemeris_data(Time("2015-06-06")) # Before the mission started - aux_data = _get_ephemeris_data(Time("2022-08-28T16:02:00")) - assert len(aux_data) == 1341 + t1 = Time("2022-08-28T16:02:00") + aux_data = _get_ephemeris_data(t1) + assert len(aux_data) > 0 + assert aux_data["time"].min() <= t1 <= aux_data["time_end"].max() - aux_data = _get_ephemeris_data(Time("2022-08-28T16:02:00"), end_time=Time("2022-08-28T16:04:00")) - assert len(aux_data) == 1341 + t2 = Time("2022-08-28T16:04:00") + aux_data = _get_ephemeris_data(t1, end_time=t2) + assert len(aux_data) > 0 + # Cache returns a padded slice so min may be slightly before t1. + assert aux_data["time"].min() <= t1 < t2 <= aux_data["time_end"].max() - aux_data = _get_ephemeris_data(Time("2022-08-28T23:58:00"), end_time=Time("2022-08-29T00:02:00")) - assert len(aux_data) == 2691 + t1 = Time("2022-08-28T23:58:00") + t2 = Time("2022-08-29T00:02:00") + aux_data = _get_ephemeris_data(t1, end_time=t2) + assert len(aux_data) > 0 + assert aux_data["time"].min() <= t1 < t2 <= aux_data["time_end"].max() @pytest.mark.remote_data @@ -129,6 +170,46 @@ def test_get_hpc_info(): assert_quantity_allclose(solo_heeq, [-9.7671984e07, 6.2774768e07, -5547166.0] * u.km) +def test_get_hpc_info_array_times_hits_warm_cache(): + """Regression: array `times` against a warm cache must not blow up. + + Before the rework, _get_ephemeris_data passed the multi-element Time + straight to TableLRUCache.get, broadcasting an N-row column against an + M-element array. The bug was latent because the array case usually + arrives with a cold cache and falls into the Fido path. + """ + STIX_EPHEMERIS_CACHE.clear() + base = Time("2023-06-15T00:00:00") + STIX_EPHEMERIS_CACHE.put(_fake_ephemeris(base), source="fake.fits") + + times = base + 6 * u.h + np.arange(10) * u.min + roll, solo_heeq, stix_pointing = get_hpc_info(times) + + assert roll.shape == (10,) + assert solo_heeq.shape == (10, 3) + assert stix_pointing.shape == (10, 2) + + +def test_get_hpc_info_single_time_hits_warm_cache(): + """Regression: scalar point query must hit the warm cache. + + Before the rework, the cache's strict time-equality check made point + queries almost always miss (64 s ANC cadence vs. arbitrary timestamps), + so every call re-ran Fido. With nearest-row lookup and range padding, + a scalar query inside the cached day must return without fetching. + """ + STIX_EPHEMERIS_CACHE.clear() + base = Time("2023-06-15T00:00:00") + STIX_EPHEMERIS_CACHE.put(_fake_ephemeris(base), source="fake.fits") + + # Pick an instant deliberately not aligned to the 64 s grid. + t = base + 6 * u.h + 17.3 * u.s + roll, solo_heeq, stix_pointing = get_hpc_info(t) + + assert np.isscalar(roll.value) or roll.size == 1 + assert solo_heeq.shape == (3,) + + @pytest.mark.remote_data def test_get_hpc_info_shapes(): t = Time("2022-08-28T16:00:00") + np.arange(10) * u.min @@ -137,9 +218,9 @@ def test_get_hpc_info_shapes(): roll3, solo_heeq3, stix_pointing3 = get_hpc_info(t[5]) assert_quantity_allclose(roll1[5], roll2) - assert_quantity_allclose(solo_heeq1[5, :], solo_heeq2[0, :]) - assert_quantity_allclose(stix_pointing1[5, :], stix_pointing2[0, :]) + assert_quantity_allclose(solo_heeq1[5, :], solo_heeq2) + assert_quantity_allclose(stix_pointing1[5, :], stix_pointing2) assert_quantity_allclose(roll3, roll2[0]) - assert_quantity_allclose(solo_heeq3, solo_heeq2[0, :]) - assert_quantity_allclose(stix_pointing3, stix_pointing2[0, :]) + assert_quantity_allclose(solo_heeq3, solo_heeq2) + assert_quantity_allclose(stix_pointing3, stix_pointing2) diff --git a/stixpy/coordinates/transforms.py b/stixpy/coordinates/transforms.py index 0fa5a96..725fb6d 100644 --- a/stixpy/coordinates/transforms.py +++ b/stixpy/coordinates/transforms.py @@ -1,5 +1,5 @@ import warnings -from functools import lru_cache +from pathlib import Path import numpy as np @@ -7,25 +7,45 @@ import astropy.units as u from astropy.coordinates import frame_transform_graph from astropy.coordinates.matrix_utilities import matrix_transpose, rotation_matrix -from astropy.io import fits -from astropy.table import QTable, vstack -from astropy.time import Time from sunpy.coordinates import HeliographicStonyhurst, Helioprojective -from sunpy.net import Fido -from sunpy.net import attrs as a +from stixpy.coordinates._ephemeris_fetcher import fetch_ephemeris_for_range from stixpy.coordinates.frames import STIXImaging +from stixpy.product.product_factory import Product +from stixpy.product.sources.anc import Ephemeris from stixpy.utils.logging import get_logger +from stixpy.utils.table_lru import TableLRUCache STIX_X_SHIFT = 26.1 * u.arcsec # fall back to this when non sas solution available STIX_Y_SHIFT = 58.2 * u.arcsec # fall back to this when non sas solution available STIX_X_OFFSET = 60.0 * u.arcsec # remaining offset after SAS solution STIX_Y_OFFSET = 8.0 * u.arcsec # remaining offset after SAS solution +# Bracket pad applied around cache range queries so np.interp doesn't clamp +# at endpoints. Intentionally not a clean 64 s multiple per the ephemeris doc. +EDGE_PAD = 80 * u.s + logger = get_logger(__name__) -__all__ = ["get_hpc_info", "stixim_to_hpc", "hpc_to_stixim"] +__all__ = ["get_hpc_info", "stixim_to_hpc", "hpc_to_stixim", "STIX_EPHEMERIS_CACHE", "load_ephemeris_fits_to_cache"] + +# Create a global cache for STIX ephemeris data +STIX_EPHEMERIS_CACHE = TableLRUCache("STIX_EPHEMERIS_CACHE", maxsize=300000, nominal_bin=64 * u.s) + + +def load_ephemeris_fits_to_cache(anc_file): + """ + Load ephemeris data from a fits file into the ephemeris cache. + """ + logger.info(f"Loading STIX ephemeris data from {anc_file}. into cache") + anc_file = Path(anc_file) + try: + anc = Product(anc_file, data_only=True) + if isinstance(anc, Ephemeris): + STIX_EPHEMERIS_CACHE.put(anc.data, source=anc_file.name) + except (OSError, ValueError) as e: + logger.error(f"Error loading STIX ephemeris data from {anc_file}: {e}") def _get_rotation_matrix_and_position(obstime, obstime_end=None): @@ -57,82 +77,62 @@ def _get_rotation_matrix_and_position(obstime, obstime_end=None): return rmatrix, solo_position_heeq -def get_hpc_info(times, end_time=None): +def get_hpc_info(times, end_time=None, *, return_source=False): r""" - Get STIX pointing and SO location from L2 aspect files. + Get STIX pointing and SO location from ANC aspect files. Parameters ---------- times : `astropy.time.Time` - Time/s to get hpc info for at. + Time, array of times, or start of a time interval. + end_time : `astropy.time.Time`, optional + End of a time interval; combined with a scalar ``times`` selects the + range-averaging mode. + return_source : list, optional + If a list is passed, it is extended with the set of ANC FITS + filenames that contributed to the result. Returns ------- - + roll, solo_heeq, stix_pointing + SOLO roll, SOLO HEEQ position, and the STIX pointing (best of the + spacecraft and SAS solutions). """ - if end_time is not None: - end_time = end_time.max() - aux = _get_ephemeris_data(times.min(), end_time if end_time is not None else times.max()) - - if end_time is not None: - indices = np.argwhere((aux["time"] >= times.min()) & (aux["time"] <= end_time.max())) - else: - indices = np.argwhere((aux["time"] >= times.min()) & (aux["time"] <= times.max())) - - indices = indices.flatten() - - if end_time is not None and times.size == 1 and indices.size >= 2: - # mean - aux = aux[indices] - - roll, pitch, yaw = np.mean(aux["roll_angle_rpy"], axis=0) - solo_heeq = np.mean(aux["solo_loc_heeq_zxy"], axis=0) - - good_solution = np.where(aux["sas_ok"] == 1) - good_sas = aux[good_solution] - - if len(good_sas) == 0: - warnings.warn(f"No good SAS solution found for time range: {times} to {end_time}.") - sas_x = 0 - sas_y = 0 + # Explicit dispatch: range mode is "scalar start + scalar end_time". + # An array of times (with or without end_time) always uses interpolation. + is_range = end_time is not None and times.isscalar + + aux = _get_ephemeris_data(times, end_time=end_time) + + if is_range: + in_window = (aux["time"] >= times) & (aux["time"] <= end_time) + window = aux[in_window] + + if len(window) >= 2: + # Average across [times, end_time] + roll, pitch, yaw = np.mean(window["roll_angle_rpy"], axis=0) + solo_heeq = np.mean(window["solo_loc_heeq_zxy"], axis=0) + + good_sas = window[window["sas_ok"] == 1] + if len(good_sas) == 0: + warnings.warn(f"No good SAS solution found for time range: {times} to {end_time}.") + sas_x = 0 + sas_y = 0 + else: + sas_x = np.mean(good_sas["y_srf"]) + sas_y = np.mean(good_sas["z_srf"]) + sigma_x = np.std(good_sas["y_srf"]) + sigma_y = np.std(good_sas["z_srf"]) + tolerance = 3 * u.arcsec + if sigma_x > tolerance or sigma_y > tolerance: + warnings.warn(f"Pointing unstable: StD(X) = {sigma_x}, StD(Y) = {sigma_y}.") else: - sas_x = np.mean(good_sas["y_srf"]) - sas_y = np.mean(good_sas["z_srf"]) - sigma_x = np.std(good_sas["y_srf"]) - sigma_y = np.std(good_sas["z_srf"]) - tolerance = 3 * u.arcsec - if sigma_x > tolerance or sigma_y > tolerance: - warnings.warn(f"Pointing unstable: StD(X) = {sigma_x}, StD(Y) = {sigma_y}.") + # Fewer than 2 rows in [start, end]: degrade to a midpoint interpolation + midpoint = times + (end_time - times) * 0.5 + roll, pitch, yaw, solo_heeq, sas_x, sas_y, good_sas = _interpolate_aux(midpoint, aux) else: - if end_time is not None and indices.size < 2: - times = times + (end_time - times) * 0.5 - - # Interpolate all times - x = (times - aux["time"][0]).to_value(u.s) - xp = (aux["time"] - aux["time"][0]).to_value(u.s) - - roll = np.interp(x, xp, aux["roll_angle_rpy"][:, 0].value) << aux["roll_angle_rpy"].unit - pitch = np.interp(x, xp, aux["roll_angle_rpy"][:, 1].value) << aux["roll_angle_rpy"].unit - yaw = np.interp(x, xp, aux["roll_angle_rpy"][:, 2].value) << aux["roll_angle_rpy"].unit - - solo_heeq = ( - np.vstack( - [ - np.interp(x, xp, aux["solo_loc_heeq_zxy"][:, 0].value), - np.interp(x, xp, aux["solo_loc_heeq_zxy"][:, 1].value), - np.interp(x, xp, aux["solo_loc_heeq_zxy"][:, 2].value), - ] - ).T - << aux["solo_loc_heeq_zxy"].unit - ) - - sas_x = np.interp(x, xp, aux["y_srf"]) - sas_y = np.interp(x, xp, aux["z_srf"]) - if x.size == 1: - good_sas = [True] if np.interp(x, xp, aux["sas_ok"]).astype(bool) else [] - else: - sas_ok = np.interp(x, xp, aux["sas_ok"]).astype(bool) - good_sas = sas_ok[sas_ok == True] # noqa E712 + # Point / array mode: interpolate each requested time + roll, pitch, yaw, solo_heeq, sas_x, sas_y, good_sas = _interpolate_aux(times, aux) # Convert the spacecraft pointing to STIX frame rotated_yaw = -yaw * np.cos(roll) + pitch * np.sin(roll) @@ -153,63 +153,127 @@ def get_hpc_info(times, end_time=None): else: warnings.warn(f"SAS solution not available using spacecraft pointing: {stix_pointing}.") - if end_time is not None or times.ndim == 0: + if is_range or times.size == 1: solo_heeq = solo_heeq.squeeze() stix_pointing = stix_pointing.squeeze() + if isinstance(return_source, list) and "__source" in aux.colnames: + return_source.extend(set(aux["__source"].value)) + return roll, solo_heeq, stix_pointing -@lru_cache -def _get_ephemeris_data(start_time, end_time=None): +def _interpolate_aux(at_times, aux): + """ + Interpolate roll/pitch/yaw, SOLO HEEQ position, and SAS pointing from + the cached ANC rows ``aux`` onto the requested ``at_times``. + + Returns ``(roll, pitch, yaw, solo_heeq, sas_x, sas_y, good_sas)``. + """ + x = (at_times - aux["time"][0]).to_value(u.s) + xp = (aux["time"] - aux["time"][0]).to_value(u.s) + + roll = np.interp(x, xp, aux["roll_angle_rpy"][:, 0].value) << aux["roll_angle_rpy"].unit + pitch = np.interp(x, xp, aux["roll_angle_rpy"][:, 1].value) << aux["roll_angle_rpy"].unit + yaw = np.interp(x, xp, aux["roll_angle_rpy"][:, 2].value) << aux["roll_angle_rpy"].unit + + solo_heeq = ( + np.vstack( + [ + np.interp(x, xp, aux["solo_loc_heeq_zxy"][:, 0].value), + np.interp(x, xp, aux["solo_loc_heeq_zxy"][:, 1].value), + np.interp(x, xp, aux["solo_loc_heeq_zxy"][:, 2].value), + ] + ).T + << aux["solo_loc_heeq_zxy"].unit + ) + + sas_x = np.interp(x, xp, aux["y_srf"].value) << aux["y_srf"].unit + sas_y = np.interp(x, xp, aux["z_srf"].value) << aux["z_srf"].unit + + if np.ndim(x) == 0 or x.size == 1: + good_sas = [True] if np.interp(x, xp, aux["sas_ok"]).astype(bool) else [] + else: + sas_ok = np.interp(x, xp, aux["sas_ok"]).astype(bool) + good_sas = sas_ok[sas_ok == True] # noqa E712 + + return roll, pitch, yaw, solo_heeq, sas_x, sas_y, good_sas + + +def _get_ephemeris_data(times, end_time=None, *, interpolate=True): r""" - Search, download and read L2 pointing data. + Return ANC ephemeris rows sufficient to interpolate or average over the + requested time(s). Used by :func:`get_hpc_info`. + + Three input shapes are accepted and normalised to a single scalar + range ``[qstart, qend]``: + + * scalar ``times`` (no ``end_time``) → point at ``times`` + * scalar ``times`` + ``end_time`` → range ``[times, end_time]`` + * array ``times`` (with or without ``end_time``) → range covering the array Parameters ---------- - start_time : `astropy.time.Time` - Time or start of a time interval. + times : `astropy.time.Time` + Scalar or array of times. + end_time : `astropy.time.Time`, optional + End of a time interval. + interpolate : bool + Legacy parameter, kept for back-compatibility; currently unused. Returns ------- + QTable + Cache slice padded by ``EDGE_PAD`` on each side. Includes the + ``__source`` column populated from the ANC filename(s). + + Raises + ------ + ValueError + If the requested range cannot be covered after fetching (real data + gap larger than ``max_gap``). + """ + # Normalise both inputs to scalar Time. `end_time` may itself be an + # array (frame transforms with vector obstime default obstime_end to + # the same array — see STIXImaging.__init__). + qstart = times if times.isscalar else times.min() + if end_time is not None: + qend = end_time if end_time.isscalar else end_time.max() + else: + qend = times if times.isscalar else times.max() + + logger.info(f"Resolving ANC ephemeris for {qstart} – {qend}") + table = _query_cache(qstart, qend) + if table is not None: + return table + + fetched = fetch_ephemeris_for_range(qstart, qend) + STIX_EPHEMERIS_CACHE.put(fetched, source=fetched["__source"].value) + + # _force=True so we still re-read the data we just put even if + # bypass_cache_read is set (which only blocks initial reads). + table = _query_cache(qstart, qend, _force=True) + if table is None: + raise ValueError( + f"Ephemeris data for {qstart}–{qend} has gaps larger than " + f"{STIX_EPHEMERIS_CACHE.max_gap} after fetch (real data gap)." + ) + return table + + +def _query_cache(qstart, qend, *, _force=False): """ - if end_time is None: - end_time = start_time - # Find, download, read aux file with pointing, sas and position information - logger.debug(f"Searching for AUX data: {start_time} - {end_time}") - query = Fido.search( - a.Time(start_time, end_time), - a.Instrument.stix, - a.Level.anc, - a.stix.DataType.asp, - a.stix.DataProduct.asp_ephemeris, - ) - if len(query["stix"]) == 0: - raise ValueError(f"No STIX pointing data found for time range {start_time} to {end_time}.") - - logger.debug(f"Downloading {len(query['stix'])} AUX files") - aux_files = Fido.fetch(query["stix"]) - if len(aux_files.errors) > 0: - raise ValueError("There were errors downloading the data.") - # Read and extract data - logger.debug("Loading and extracting AUX data") - - aux_data = [] - for aux_file in aux_files: - hdu = fits.getheader(aux_file, ext=0) - aux = QTable.read(aux_file, hdu=2) - date_beg = Time(hdu.get("DATE-BEG")) - aux["time"] = ( - date_beg + aux["time"] - 32 * u.s - ) # Shift AUX data by half a time bin (starting time vs. bin centre) - [aux.meta.pop(key) for key in ["CHECKSUM", "DATASUM"]] - aux_data.append(aux) - - aux = vstack(aux_data) - aux.sort(keys=["time"]) - - return aux + Translate a normalised (qstart, qend) into a single cache call. + + Point queries (qstart == qend) widen the cache search window by + ``EDGE_PAD`` so the cache can find bracketing rows for downstream + interpolation. Range queries keep coverage strict on ``[qstart, qend]`` + and pad only the returned slice. + """ + if qstart == qend: + return STIX_EPHEMERIS_CACHE.get_range(qstart - EDGE_PAD, qend + EDGE_PAD, pad=0 * u.s, _force=_force) + return STIX_EPHEMERIS_CACHE.get_range(qstart, qend, pad=EDGE_PAD, _force=_force) @frame_transform_graph.transform(coord.FunctionTransform, STIXImaging, Helioprojective) diff --git a/stixpy/net/tests/test_client.py b/stixpy/net/tests/test_client.py index 4896c17..4b5da34 100644 --- a/stixpy/net/tests/test_client.py +++ b/stixpy/net/tests/test_client.py @@ -253,22 +253,33 @@ def test_search_date_product_sci(): assert len(res) == 1 -@pytest.mark.parametrize( - "query, expected_len, is_total", - [ - ([a.Instrument.stix], 67, True), - ([a.Instrument.stix, a.stix.DataType.ql], 6, False), - ([a.Instrument.stix, a.stix.DataType.sci], 58, False), - ([a.Instrument.stix, a.stix.DataType.hk], 1, False), - ([a.Instrument.stix, a.stix.DataType.asp], 1, False), - ([a.Instrument.stix, a.stix.DataType.cal], 1, False), - ], -) @pytest.mark.remote_data -def test_fido(query, expected_len, is_total): - res = Fido.search(a.Time("2020-11-17T00:00", "2020-11-17T23:59"), *query) - actual_len = len(res["stix"]) - assert actual_len == expected_len +def test_fido(): + res = Fido.search(a.Time("2020-11-17T00:00", "2020-11-17T23:59"), a.Instrument.stix) + len_total = len(res["stix"]) + assert len_total == 67 + + res_ql = Fido.search(a.Time("2020-11-17T00:00", "2020-11-17T23:59"), a.Instrument.stix, a.stix.DataType.ql) + len_ql = len(res_ql["stix"]) + assert len_ql == 6 + + res_sci = Fido.search(a.Time("2020-11-17T00:00", "2020-11-17T23:59"), a.Instrument.stix, a.stix.DataType.sci) + len_sci = len(res_sci["stix"]) + assert len_sci == 58 + + res_hk = Fido.search(a.Time("2020-11-17T00:00", "2020-11-17T23:59"), a.Instrument.stix, a.stix.DataType.hk) + len_kh = len(res_hk["stix"]) + assert len_kh == 1 + + res_asp = Fido.search(a.Time("2020-11-17T00:00", "2020-11-17T23:59"), a.Instrument.stix, a.stix.DataType.asp) + len_asp = len(res_asp["stix"]) + assert len_asp == 1 + + res_cal = Fido.search(a.Time("2020-11-17T00:00", "2020-11-17T23:59"), a.Instrument.stix, a.stix.DataType.cal) + len_cal = len(res_cal["stix"]) + assert len_cal == 1 + + assert len_ql + len_sci + len_kh + len_asp + len_cal == len_total @pytest.mark.remote_data diff --git a/stixpy/product/product.py b/stixpy/product/product.py index ce9deeb..f22fe5f 100644 --- a/stixpy/product/product.py +++ b/stixpy/product/product.py @@ -23,7 +23,7 @@ def __init_subclass__(cls, **kwargs): class GenericProduct(BaseProduct): - def __init__(self, *, meta, control, data, idb_versions=None, energies=None): + def __init__(self, *, meta, control, data, idb_versions=None, energies=None, **kwargs): """ Generic product composed of meta, control, data and optionally idb, and energy information diff --git a/stixpy/product/product_factory.py b/stixpy/product/product_factory.py index da6e7ee..0cf5a41 100644 --- a/stixpy/product/product_factory.py +++ b/stixpy/product/product_factory.py @@ -50,26 +50,31 @@ def read_qtable(file, hdu, hdul=None): `astropy.table.QTable` The corrected QTable with correct data types """ - if hdul is None: + own_hdul = hdul is None + if own_hdul: hdul = fits.open(file) - qtable = QTable.read(file, hdu) + try: + qtable = QTable.read(file, hdu) - for col in hdul[hdu].data.columns: - if col.unit: - logger.debug(f"Unit present dtype correction needed for {col}") - dtype = col.dtype + for col in hdul[hdu].data.columns: + if col.unit: + logger.debug(f"Unit present dtype correction needed for {col}") + dtype = col.dtype - if col.bzero: - logger.debug(f"Unit present dtype and bzero correction needed for {col}") - bits = np.log2(col.bzero) - if bits.is_integer(): - dtype = BITS_TO_UINT[int(bits + 1)] + if col.bzero: + logger.debug(f"Unit present dtype and bzero correction needed for {col}") + bits = np.log2(col.bzero) + if bits.is_integer(): + dtype = BITS_TO_UINT[int(bits + 1)] - if hasattr(dtype, "subdtype"): - dtype = dtype.base + if hasattr(dtype, "subdtype"): + dtype = dtype.base - qtable[col.name] = qtable[col.name].astype(dtype) + qtable[col.name] = qtable[col.name].astype(dtype) + finally: + if own_hdul: + hdul.close() return qtable @@ -114,18 +119,30 @@ def _read_file(self, fname, **kwargs): msg = f"Failed to read {fname}." raise OSError(msg) from e - if hdul[0].header.get("INSTRUME", "") != "STIX": - raise FileError(f"File '{fname}' is not a STIX fits file.") + try: + if hdul[0].header.get("INSTRUME", "") != "STIX": + raise FileError(f"File '{fname}' is not a STIX fits file.") - data = {"meta": hdul[0].header} - for name in ["CONTROL", "DATA", "IDB_VERSIONS", "ENERGIES"]: - try: - data[name.lower()] = read_qtable(fname, hdu=name) - except KeyError as e: - if name in ("IDB_VERSIONS", "ENERGIES"): - logger.debug(f"Extension '{name}' not in file '{fname}'") - else: - raise e + data = {"meta": hdul[0].header.copy()} + + # determine extensions to load + # allow for just data extension if requested + if kwargs.get("data_only", False): + extensions = ["CONTROL", "DATA"] + # normally read all extensions + else: + extensions = ["CONTROL", "DATA", "IDB_VERSIONS", "ENERGIES"] + + for name in extensions: + try: + data[name.lower()] = read_qtable(fname, hdu=name, hdul=hdul) + except KeyError as e: + if name in ("IDB_VERSIONS", "ENERGIES"): + logger.debug(f"Extension '{name}' not in file '{fname}'") + else: + raise e + finally: + hdul.close() return [data] @@ -214,6 +231,11 @@ def _parse_map(self, arg, **kwargs): @_parse_arg.register(Request) def _parse_url(self, arg, **kwargs): url = arg.full_url + # NOTE: the data server truncates large FITS responses under concurrent + # load; CI sidesteps this by running the online tests serially (``-n 1``, + # see tox.ini). If end-user resilience is wanted, re-introduce a validate + + # re-download guard here using ``stixpy.utils.io.is_valid_fits`` and + # ``cache.download(url, redownload=True)``. path = str(cache.download(url).absolute()) pairs = self._read_file(path, **kwargs) return pairs diff --git a/stixpy/product/sources/__init__.py b/stixpy/product/sources/__init__.py index 6814efe..f8ed7e4 100644 --- a/stixpy/product/sources/__init__.py +++ b/stixpy/product/sources/__init__.py @@ -1,3 +1,4 @@ +from stixpy.product.sources.anc import * from stixpy.product.sources.housekeeping import * from stixpy.product.sources.quicklook import * from stixpy.product.sources.science import * diff --git a/stixpy/product/sources/anc.py b/stixpy/product/sources/anc.py new file mode 100644 index 0000000..2139b96 --- /dev/null +++ b/stixpy/product/sources/anc.py @@ -0,0 +1,56 @@ +import astropy.units as u +from astropy.time import Time +from astropy.units import Quantity + +from sunpy.time import TimeRange + +from stixpy.product.product import L1Product + +__all__ = ["ANCProduct", "Ephemeris"] + + +class ANCProduct(L1Product): + """ + Basic ANC + """ + + @property + def time(self) -> Time: + return self.data["time"] + + @property + def exposure_time(self) -> Quantity[u.s]: + return self.data["timedel"].to(u.s) + + @property + def time_range(self) -> TimeRange: + """ + A `sunpy.time.TimeRange` for the data. + """ + return TimeRange(self.time[0] - self.exposure_time[0] / 2, self.time[-1] + self.exposure_time[-1] / 2) + + +class Ephemeris(ANCProduct): + """ + Ephemeris data in daily files normal with 64s time resolution + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + # TODO remove this when we have a better solution + # Shift ANC data by half a time bin (starting time vs. bin centre) + self.data["time_end"] = self.data["time"] + 32 * u.s + self.data["time"] = self.data["time"] - 32 * u.s + self.data["timedel"] = 64 * u.s + self.data["time_utc"] = [Time(t, format="isot", scale="utc") for t in self.data["time_utc"]] + + @classmethod + def is_datasource_for(cls, *, meta, **kwargs): + """Determines if meta data meach Raw Pixel Data""" + service_subservice_ssid = tuple(meta[name] for name in ["STYPE", "SSTYPE", "SSID"]) + level = meta["level"] + if service_subservice_ssid == (0, 0, 1) and level == "ANC": + return True + + def __repr__(self): + return f"{self.__class__.__name__}\n {self.time_range}" diff --git a/stixpy/product/sources/science.py b/stixpy/product/sources/science.py index f3863a1..de193b0 100644 --- a/stixpy/product/sources/science.py +++ b/stixpy/product/sources/science.py @@ -17,6 +17,7 @@ from stixpy.io.readers import read_subc_params from stixpy.product.product import L1Product +from stixpy.utils.table import drop_fits_checksums __all__ = [ "ScienceData", @@ -660,13 +661,7 @@ def concatenate(self, others): other.control["index"] = other.control["index"] + self_control_ind_max other.data["control_index"] = other.data["control_index"] + self_control_ind_max - try: - [ - (table.meta.pop("DATASUM"), table.meta.pop("CHECKSUM")) - for table in [control, other.control, data, other.data] - ] - except KeyError: - pass + drop_fits_checksums(control, other.control, data, other.data) control = vstack([control, other.control]) data = vstack([data, other.data]) diff --git a/stixpy/utils/io.py b/stixpy/utils/io.py new file mode 100644 index 0000000..937c54f --- /dev/null +++ b/stixpy/utils/io.py @@ -0,0 +1,65 @@ +""" +I/O helpers shared across stixpy. + +.. note:: + :func:`is_valid_fits` is currently **not wired into any download path**. + The data server truncates large FITS responses under concurrent load; CI + sidesteps this by running the online tests serially (``-n 1``, see tox.ini). + The helper and its tests are kept here on purpose: if end-user (non-CI) + resilience is wanted, reintroduce a validate + re-download guard at the two + download boundaries that read from the server — + + * ``stixpy.product.product_factory.ProductFactory._parse_url`` + (``cache.download(url, redownload=True)`` on corruption), and + * ``stixpy.coordinates._ephemeris_fetcher.fetch_ephemeris_for_range`` + (``Fido.fetch(query["stix"], overwrite=True)`` on corruption). + + Prefer an *on-error* redownload (try the read, redownload only if it fails) over + pre-validating every load, to avoid reading each file twice on the happy path. +""" + +import warnings + +from astropy.io import fits + +from stixpy.utils.logging import get_logger + +logger = get_logger(__name__) + +__all__ = ["is_valid_fits"] + + +def is_valid_fits(path) -> bool: + """ + Return whether ``path`` is a readable, structurally complete FITS file. + + Opens the file and forces a read of every HDU's data so that *truncated* + downloads — which open fine but fail once the missing bytes are touched — + are detected. Remote data servers occasionally truncate responses under + concurrent load, and parfive caches the partial file as if it were + complete; this check lets callers notice that and re-download. + + Warnings are suppressed during the check: valid STIX FITS files emit benign + verification warnings (e.g. the non-standard ``BLANK`` keyword) which must + not be mistaken for corruption. Only hard read errors count. + + Parameters + ---------- + path + Path-like to the FITS file. + + Returns + ------- + bool + ``True`` if the file opens and every HDU's data can be read. + """ + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with fits.open(path) as hdul: + for hdu in hdul: + _ = hdu.data + except Exception as e: + logger.debug(f"FITS validation failed for {path}: {e}") + return False + return True diff --git a/stixpy/utils/table.py b/stixpy/utils/table.py new file mode 100644 index 0000000..6fce63f --- /dev/null +++ b/stixpy/utils/table.py @@ -0,0 +1,27 @@ +"""Generic helpers for `astropy.table` objects used across stixpy.""" + +__all__ = ["drop_fits_checksums"] + +_NOISY_FITS_META_KEYS = ("DATASUM", "CHECKSUM") + + +def drop_fits_checksums(*tables) -> None: + """ + Remove per-file FITS integrity checksums from each table's ``meta`` in place. + + ``DATASUM`` and ``CHECKSUM`` are HDU-level integrity hashes that correctly + differ between source FITS files. When stacking tables from multiple files + (e.g. daily ANC files), the merged table has no meaningful single + checksum — dropping the keys before `~astropy.table.vstack` avoids the + ``MergeConflictWarning`` while keeping any *other* metadata conflicts + visible. + + Parameters + ---------- + *tables + One or more `~astropy.table.Table` (or subclass) instances. Each + table's ``meta`` is mutated in place; missing keys are ignored. + """ + for table in tables: + for key in _NOISY_FITS_META_KEYS: + table.meta.pop(key, None) diff --git a/stixpy/utils/table_lru.py b/stixpy/utils/table_lru.py new file mode 100644 index 0000000..30e954a --- /dev/null +++ b/stixpy/utils/table_lru.py @@ -0,0 +1,304 @@ +import warnings +from datetime import datetime + +import numpy as np + +import astropy.units as u +from astropy.table import QTable, unique, vstack +from astropy.time import Time + +from stixpy.utils.logging import get_logger +from stixpy.utils.table import drop_fits_checksums + +logger = get_logger(__name__) + +__all__ = ["TableLRUCache"] + + +class TableLRUCache: + """ + LRU cache of time-indexed astropy.table rows. + + The cache stores rows of a QTable keyed by the ``time`` column. Two + query modes are supported: + + * Point query (``get_point``): returns the row whose ``time`` is closest + to a requested instant, provided the distance is within + ``max_point_distance``. + * Range query (``get_range``): returns rows in ``[start - pad, end + pad]`` + iff the cache covers ``[start, end]`` without holes larger than + ``max_gap``. + + The cache knows nothing about ANC files, Fido, or daily file boundaries. + Those concerns live in the fetcher / orchestrator layers. + + Indexing + -------- + The ``time`` column is kept sorted ascending (enforced by :meth:`put` and + :meth:`_prune`). A float mirror of it in unix seconds, ``_times_unix``, is + held alongside the table and refreshed by :meth:`_rebuild_index` whenever + the cache is rebuilt or re-sorted. Both query modes binary-search this + array with ``numpy.searchsorted`` (O(log n)) rather than scanning every + row (O(n)): ``get_point`` locates the insertion point and inspects only the + two straddling neighbours, while ``get_range`` resolves its window to a + contiguous index slice. Precise, leap-second-aware ``Time`` arithmetic is + then applied only to those few candidate rows. This makes the common usage + pattern — one bulk day-add followed by many queries against a large cache — + cheap; an astropy table index (``add_index``) is deliberately *not* used, + as it would be invalidated by the ``vstack``/``unique``/``sort`` in + :meth:`put` and does not serve nearest-neighbour or gap-aware range lookups. + """ + + def __init__( + self, + name: str, + *, + maxsize: int = 300000, + nominal_bin: u.Quantity = 64 * u.s, + max_point_distance: u.Quantity = 90 * u.s, + max_gap: u.Quantity = 150 * u.s, + default_bin_duration: u.Quantity | None = None, + ): + """ + Parameters + ---------- + name + Cache identifier used in log messages. + maxsize + Maximum number of rows retained. Default is ~6 months at 64 s. + nominal_bin + Nominal time resolution of the underlying data. Informational; + used for log messages. Real spacing varies (the doc explicitly + notes 64 s is not exact). + max_point_distance + Maximum |row.time - query_time| accepted for a point query. + Beyond this, ``get_point`` returns ``None``. + max_gap + Maximum consecutive inter-row time gap accepted inside a range + query; also the edge tolerance for considering a range covered. + default_bin_duration + Deprecated alias for ``nominal_bin``. + """ + if default_bin_duration is not None: + warnings.warn( + "'default_bin_duration' is deprecated; use 'nominal_bin' instead.", + DeprecationWarning, + stacklevel=2, + ) + nominal_bin = default_bin_duration + + self.name = name + self.maxsize = maxsize + self.nominal_bin = nominal_bin + self.max_point_distance = max_point_distance + self.max_gap = max_gap + self.cache = QTable() + # Float (unix seconds) mirror of the sorted ``time`` column, kept in + # sync by ``_rebuild_index``. Enables O(log n) ``np.searchsorted`` + # lookups instead of scanning every row on each query. + self._times_unix = np.empty(0, dtype=float) + self.bypass_cache_read = False + + @property + def default_bin_duration(self) -> u.Quantity: + """Deprecated alias for ``nominal_bin``.""" + return self.nominal_bin + + def clear(self) -> None: + """Empty the cache.""" + self.cache = QTable() + self._rebuild_index() + logger.info(f"{self.name}: Cleared LRU cache.") + + def __len__(self) -> int: + return len(self.cache) + + def _rebuild_index(self) -> None: + """ + Refresh the binary-search lookup array from the cache. + + Must be called whenever ``self.cache`` is rebuilt or re-sorted (i.e. + after :meth:`put` / :meth:`_prune` / :meth:`clear`). The ``time`` + column is assumed sorted ascending, matching the cache invariant. + """ + if len(self.cache) == 0: + self._times_unix = np.empty(0, dtype=float) + else: + self._times_unix = np.asarray(self.cache["time"].unix, dtype=float) + + def put(self, anc_data, *, source="unknown") -> None: + """ + Add rows to the cache. Existing rows with the same ``time`` are + replaced (``keep="last"``). Injects ``__lcu`` (LRU timestamp) and + ``__source`` columns; pass ``source`` as a string (broadcast to all + rows) or a per-row array. + """ + anc_data = anc_data.copy() + anc_data["__lcu"] = datetime.now().timestamp() + anc_data["__source"] = source + drop_fits_checksums(anc_data, self.cache) + logger.info(f"{self.name}: Adding {len(anc_data)} rows.") + self.cache = vstack([self.cache, anc_data]) + self.cache = unique(self.cache, keys=["time"], keep="last") + self.cache.sort(keys=["time"]) + self._prune() + self._rebuild_index() + + def _prune(self) -> None: + """Evict least-recently-used rows down to 80% of maxsize.""" + if len(self.cache) > self.maxsize: + logger.info(f"{self.name}: Pruning LRU cache.") + sorted_idx = self.cache.argsort(keys=["__lcu", "time"]) + self.cache = self.cache[sorted_idx[-int(self.maxsize * 0.8) :]] + self.cache.sort(keys=["time"]) + self._rebuild_index() + + def get_point(self, t: Time, *, _force: bool = False) -> QTable | None: + """ + Nearest-row lookup. + + Returns a 1-row deep copy of the row whose ``time`` is closest to ``t`` + if that distance is within ``max_point_distance``; otherwise ``None``. + + Parameters + ---------- + t + Scalar `~astropy.time.Time`. + _force + Internal-only escape hatch that ignores ``bypass_cache_read``. + Used by the orchestrator immediately after a fresh fetch to read + back what was just put. + + Raises + ------ + TypeError + If ``t`` is not scalar. + """ + if not isinstance(t, Time) or not t.isscalar: + raise TypeError("get_point requires a scalar astropy.time.Time") + + if len(self.cache) == 0 or (self.bypass_cache_read and not _force): + return None + + # Binary search for the insertion point; the nearest row is one of the + # two straddling neighbours. Precise (leap-second aware) Time + # arithmetic is then done on at most those two rows. + pos = int(np.searchsorted(self._times_unix, t.unix)) + cand = [j for j in (pos - 1, pos) if 0 <= j < len(self._times_unix)] + dt = np.abs((self.cache["time"][cand] - t).to_value(u.s)) + k = int(np.argmin(dt)) + i = cand[k] + if dt[k] > self.max_point_distance.to_value(u.s): + return None + + self.cache["__lcu"][i] = datetime.now().timestamp() + logger.info(f"{self.name}: point hit for {t}, nearest row at {self.cache['time'][i]} (Δ={dt[k]:.2f} s).") + ret = self.cache[[i]].copy(copy_data=True) + ret["__lcu"] = 1 + return ret + + def get_range(self, start: Time, end: Time, *, pad: u.Quantity = 0 * u.s, _force: bool = False) -> QTable | None: + """ + Range lookup with honest coverage detection. + + ``[start, end]`` is considered "covered" iff: + + * A cached row exists within ``max_gap`` of ``start``. + * A cached row exists within ``max_gap`` of ``end``. + * Between those two brackets, no consecutive inter-row gap exceeds + ``max_gap``. + + This bracketing semantic — rather than strict containment — is what + the ephemeris doc requires. It also makes short range queries + (shorter than the nominal bin cadence, e.g. a 10 s query against + 64 s data) work correctly: the brackets straddle the window. + + Returns rows in ``[start - pad, end + pad]``. If that slice is + empty (very short range with ``pad=0``), the bracketing rows from + the coverage check are returned so callers always have something + to interpolate from. Returns ``None`` on coverage failure. + + Parameters + ---------- + start, end + Scalar `~astropy.time.Time`. ``start <= end``. + pad + Extra time to widen the returned slice on each side. Useful to + bracket the requested range for downstream interpolation. + + Raises + ------ + TypeError + If ``start`` or ``end`` is not scalar. + """ + if not isinstance(start, Time) or not start.isscalar: + raise TypeError("get_range requires a scalar astropy.time.Time for 'start'") + if not isinstance(end, Time) or not end.isscalar: + raise TypeError("get_range requires a scalar astropy.time.Time for 'end'") + + if len(self.cache) == 0 or (self.bypass_cache_read and not _force): + return None + + times = self.cache["time"] + times_unix = self._times_unix + max_gap_s = self.max_gap.to_value(u.s) + + # Coverage uses an expanded window so rows just outside [start, end] + # can satisfy the bracket condition (essential for short queries). + # Binary search bounds the window to a contiguous index slice rather + # than scanning every row. + lo = int(np.searchsorted(times_unix, (start - self.max_gap).unix, side="left")) + hi = int(np.searchsorted(times_unix, (end + self.max_gap).unix, side="right")) + expanded_idx = np.arange(lo, hi) + if expanded_idx.size == 0: + logger.info(f"{self.name}: no rows near [{start}, {end}] — not covered.") + return None + + expanded_times = times[lo:hi] + + left_dt = float(np.min(np.abs((expanded_times - start).to_value(u.s)))) + right_dt = float(np.min(np.abs((expanded_times - end).to_value(u.s)))) + if left_dt > max_gap_s or right_dt > max_gap_s: + logger.info( + f"{self.name}: edge not covered (left={left_dt:.1f}s, right={right_dt:.1f}s, max_gap={max_gap_s}s)." + ) + return None + + if expanded_idx.size >= 2: + diffs = np.diff(times_unix[lo:hi]) + if diffs.max() > max_gap_s: + logger.info(f"{self.name}: interior gap {diffs.max():.1f}s exceeds max_gap={max_gap_s}s — not covered.") + return None + + pad_s = pad.to_value(u.s) + if pad_s > 0: + # Avoid `start - 0*u.s` arithmetic: it's not bit-equal to `start` + # in astropy.Time, which would exclude rows lying exactly on the + # boundary. + s_lo = int(np.searchsorted(times_unix, (start - pad).unix, side="left")) + s_hi = int(np.searchsorted(times_unix, (end + pad).unix, side="right")) + else: + s_lo = int(np.searchsorted(times_unix, start.unix, side="left")) + s_hi = int(np.searchsorted(times_unix, end.unix, side="right")) + slice_idx = np.arange(s_lo, s_hi) + if slice_idx.size == 0: + # Very short range with no pad: fall back to the bracketing rows + # so the caller still has something to interpolate from. + slice_idx = expanded_idx + + self.cache["__lcu"][slice_idx] = datetime.now().timestamp() + logger.info(f"{self.name}: range hit [{start}, {end}] — returning {slice_idx.size} rows (pad={pad_s}s).") + ret = self.cache[slice_idx].copy(copy_data=True) + ret["__lcu"] = 1 + return ret + + def get(self, start: Time, end: Time | None = None) -> QTable | None: + """ + Back-compat dispatcher. + + * ``end is None`` or ``end == start`` → :meth:`get_point`. + * Otherwise → :meth:`get_range` with ``pad=0``. + """ + if end is None or end == start: + return self.get_point(start) + return self.get_range(start, end, pad=0 * u.s) diff --git a/stixpy/utils/tests/test_io.py b/stixpy/utils/tests/test_io.py new file mode 100644 index 0000000..a783b85 --- /dev/null +++ b/stixpy/utils/tests/test_io.py @@ -0,0 +1,37 @@ +import numpy as np + +from astropy.io import fits + +from stixpy.utils.io import is_valid_fits + + +def _write_fits(path, shape=(100, 10)): + fits.HDUList([fits.PrimaryHDU(np.arange(np.prod(shape)).reshape(shape))]).writeto(path) + return path + + +def test_is_valid_fits_complete(tmp_path): + """A complete, readable FITS file validates.""" + path = _write_fits(tmp_path / "ok.fits") + assert is_valid_fits(path) is True + + +def test_is_valid_fits_truncated(tmp_path): + """A truncated download (data section cut off) is rejected.""" + path = _write_fits(tmp_path / "ok.fits") + raw = path.read_bytes() + truncated = tmp_path / "truncated.fits" + truncated.write_bytes(raw[: len(raw) // 3]) + assert is_valid_fits(truncated) is False + + +def test_is_valid_fits_empty(tmp_path): + """An empty file is rejected.""" + path = tmp_path / "empty.fits" + path.write_bytes(b"") + assert is_valid_fits(path) is False + + +def test_is_valid_fits_missing(tmp_path): + """A non-existent path is rejected rather than raising.""" + assert is_valid_fits(tmp_path / "does-not-exist.fits") is False diff --git a/stixpy/utils/tests/test_table.py b/stixpy/utils/tests/test_table.py new file mode 100644 index 0000000..682cb20 --- /dev/null +++ b/stixpy/utils/tests/test_table.py @@ -0,0 +1,33 @@ +from astropy.table import QTable + +from stixpy.utils.table import drop_fits_checksums + + +def test_drop_fits_checksums_removes_known_keys(): + t = QTable({"x": [1, 2, 3]}) + t.meta["DATASUM"] = "1868371976" + t.meta["CHECKSUM"] = "UGARV94QUGAQU93Q" + t.meta["INSTRUME"] = "STIX" + + drop_fits_checksums(t) + + assert "DATASUM" not in t.meta + assert "CHECKSUM" not in t.meta + assert t.meta["INSTRUME"] == "STIX" # unrelated keys preserved + + +def test_drop_fits_checksums_handles_missing_keys(): + t = QTable({"x": [1, 2]}) # no meta keys set + drop_fits_checksums(t) # must not raise + + +def test_drop_fits_checksums_accepts_multiple_tables(): + t1 = QTable({"x": [1]}) + t2 = QTable({"y": [2]}) + t1.meta["DATASUM"] = "a" + t2.meta["CHECKSUM"] = "b" + + drop_fits_checksums(t1, t2) + + assert "DATASUM" not in t1.meta + assert "CHECKSUM" not in t2.meta diff --git a/stixpy/utils/tests/test_table_lru_cache.py b/stixpy/utils/tests/test_table_lru_cache.py new file mode 100644 index 0000000..6adbc74 --- /dev/null +++ b/stixpy/utils/tests/test_table_lru_cache.py @@ -0,0 +1,395 @@ +from datetime import datetime, timedelta + +import numpy as np +import pytest + +import astropy.units as u +from astropy.table import QTable +from astropy.time import Time + +from sunpy.net import Fido +from sunpy.net import attrs as a + +from stixpy.coordinates.transforms import STIX_EPHEMERIS_CACHE, get_hpc_info, load_ephemeris_fits_to_cache +from stixpy.utils.table_lru import TableLRUCache + + +@pytest.fixture +def cache(): + """Fixture to create a fresh TableLRUCache instance.""" + return TableLRUCache( + "testcache", + maxsize=10, + nominal_bin=1 * u.s, + max_point_distance=2 * u.s, + max_gap=2.5 * u.s, + ) + + +@pytest.fixture +def mock_data_table(): + """Fixture to create mock ephemeris data.""" + now = Time(datetime.now()) + times = now + np.arange(10) * timedelta(seconds=1) + data = QTable({"time": times, "value": np.arange(10)}) + return data + + +def test_initialization_default(): + """Test default initialization of TableLRUCache.""" + cache = TableLRUCache("testcache") + assert cache.maxsize == 300000 + assert cache.nominal_bin == 64 * u.s + assert cache.default_bin_duration == 64 * u.s # deprecated alias still works + assert len(cache.cache) == 0 + + +def test_initialization_deprecated_alias_warns(): + """`default_bin_duration` is accepted but emits a DeprecationWarning.""" + with pytest.warns(DeprecationWarning, match="default_bin_duration"): + cache = TableLRUCache("testcache", default_bin_duration=42 * u.s) + assert cache.nominal_bin == 42 * u.s + + +def test_initialization_custom_maxsize(): + """Test initialization with a custom maxsize.""" + cache = TableLRUCache("testcache", maxsize=100) + assert cache.maxsize == 100 + + +def test_clear_cache(cache, mock_data_table): + """Test clearing the cache.""" + cache.put(mock_data_table) + assert len(cache.cache) == len(mock_data_table) + cache.clear() + assert len(cache.cache) == 0 + + +def test_put_data(cache, mock_data_table): + """Test adding data to the cache.""" + cache.put(mock_data_table) + assert len(cache.cache) == len(mock_data_table) + assert "__lcu" in cache.cache.colnames + + +def test_get_data(cache, mock_data_table): + """Test retrieving data from the cache.""" + cache.put(mock_data_table) + start_time = mock_data_table["time"][2] + end_time = mock_data_table["time"][5] + result = cache.get(start_time, end_time) + assert len(result) == 4 + assert all(result["time"] >= start_time) + assert all(result["time"] <= end_time) + assert all(result["__lcu"] == 1) # code for HIT + + +def test_get_data_no_match(cache, mock_data_table): + """Test retrieving data when no matching time range exists.""" + cache.put(mock_data_table) + start_time = Time(datetime.now() + timedelta(days=1)) + result = cache.get(start_time) + assert result is None + + +def test_get_data_single_time(cache, mock_data_table): + """Test retrieving data for a single time point.""" + cache.put(mock_data_table) + start_time = mock_data_table["time"][3] + result = cache.get(start_time) + assert len(result) == 1 + assert result["time"][0] == start_time + assert result["__lcu"][0] == 1 # code for HIT + + +def test_prune_logic(cache, mock_data_table): + """Test internal pruning logic. + + pruning should be cald if the cache is full and new data is added + """ + cache.maxsize = len(mock_data_table) - 3 + cache.put(mock_data_table) + assert len(cache.cache) <= cache.maxsize + + +def test_put_new_data_overrides_same_old_data(cache): + now = Time(datetime.now()) + times = now + np.arange(5) * timedelta(seconds=1) + data_first = QTable({"time": times, "value": np.full((5,), 1)}) + data_last = QTable({"time": times, "value": np.full((5,), 2)}) + + cache.put(data_first) + assert len(cache.cache) == len(data_first) + + found = cache.get(times[0]) + assert len(found) == 1 + assert found["value"][0] == 1 + + # Add new data with the same time but different value + # This should override the old data + cache.put(data_last) + + found = cache.get(times[0]) + assert len(found) == 1 + assert found["value"][0] == 2 + + +def test_get_range_in_block(cache): + """Test retrieving data for a range with time gaps. + This test checks if the cache can handle time gaps correctly. + should return all data in the range or None if any gaps in range + """ + now = Time(datetime.now()) + times = now + np.arange(10) * u.s + times[5:] += 2 * u.s # introduce a gap in the data + data = QTable({"time": times, "value": np.full((10,), 1)}) + cache.put(data) + start_time = data["time"][1] + end_time = data["time"][3] + result = cache.get(start_time, end_time) + assert len(result) == 3 + assert all(result["time"] >= start_time) + assert all(result["time"] <= end_time) + + # Test with a range that includes a gap + end_time = data["time"][6] + result = cache.get(start_time, end_time) + assert result is None + + +def test_get_point_returns_nearest_within_threshold(cache, mock_data_table): + """get_point returns the closest row by min |dt| inside max_point_distance.""" + cache.put(mock_data_table) + # Query halfway between rows 3 and 4 (1 s spacing → 0.5 s either side). + query = mock_data_table["time"][3] + 0.4 * u.s + result = cache.get_point(query) + assert result is not None + assert len(result) == 1 + assert result["time"][0] == mock_data_table["time"][3] + assert result["__lcu"][0] == 1 + + +def test_get_point_returns_none_outside_threshold(cache, mock_data_table): + """get_point returns None when no row is within max_point_distance.""" + cache.put(mock_data_table) + # cache fixture sets max_point_distance=2 s; query 1 day away → no hit. + query = mock_data_table["time"][0] + 1 * u.day + assert cache.get_point(query) is None + + +def test_get_point_picks_closer_of_two_neighbours(cache, mock_data_table): + """When two rows bracket the query, the closer one wins.""" + cache.put(mock_data_table) + # Query 0.3 s past row 3 → closer to row 3 (0.3 s) than row 4 (0.7 s). + query = mock_data_table["time"][3] + 0.3 * u.s + result = cache.get_point(query) + assert result["time"][0] == mock_data_table["time"][3] + + +def test_get_range_pad_returns_bracketing_rows(cache, mock_data_table): + """pad>0 widens the returned slice so callers can interpolate.""" + cache.put(mock_data_table) + start = mock_data_table["time"][3] + end = mock_data_table["time"][5] + strict = cache.get_range(start, end, pad=0 * u.s) + padded = cache.get_range(start, end, pad=1.5 * u.s) + assert len(strict) == 3 # rows 3,4,5 + assert len(padded) == 5 # rows 2..6 + + +def test_get_range_shorter_than_cadence_hits_cache(): + """Range queries narrower than the bin cadence must still hit the cache. + + Regression for a bug where the coverage check required at least one row + strictly inside [start, end]. With 64 s ANC cadence and a 10 s query + window (e.g. ``get_hpc_info("2023-01-01T00:00:00", "2023-01-01T00:00:10")``) + no row falls inside, so coverage was wrongly reported as missing — + forcing a Fido fetch even when the day was already cached. + """ + cache_64s = TableLRUCache("c64", maxsize=2000, nominal_bin=64 * u.s, max_gap=150 * u.s) + base = Time("2023-01-01T00:00:00") + times = base + 26 * u.s + np.arange(1350) * 64 * u.s # full day at 64s cadence + cache_64s.put(QTable({"time": times, "value": np.arange(1350)})) + + # 10 s window with no row strictly inside — must still be served from cache. + result = cache_64s.get_range(base, base + 10 * u.s, pad=80 * u.s) + assert result is not None + assert len(result) >= 1 + + +def test_get_range_real_gap_returns_none(cache): + """A consecutive gap larger than max_gap fails coverage.""" + # max_gap is 2.5 s on this fixture; inject a 3 s hole. + now = Time(datetime.now()) + times = now + np.array([0, 1, 2, 5, 6, 7]) * u.s + cache.put(QTable({"time": times, "value": np.arange(6)})) + assert cache.get_range(times[0], times[-1], pad=0 * u.s) is None + + +def test_get_range_jitter_within_max_gap_ok(cache): + """Sub-max_gap row-spacing jitter is accepted (no false 'gap').""" + # max_gap = 2.5 s; jitter rows at 0, 1, 2.2, 3.3, 4 (max diff 1.2 s). + now = Time(datetime.now()) + times = now + np.array([0.0, 1.0, 2.2, 3.3, 4.0]) * u.s + cache.put(QTable({"time": times, "value": np.arange(5)})) + result = cache.get_range(times[0], times[-1], pad=0 * u.s) + assert result is not None + assert len(result) == 5 + + +def test_get_point_requires_scalar_time(cache, mock_data_table): + """Array input to get_point is a programmer error.""" + cache.put(mock_data_table) + with pytest.raises(TypeError, match="scalar"): + cache.get_point(mock_data_table["time"]) + + +def test_get_range_requires_scalar_times(cache, mock_data_table): + """Array input to get_range is a programmer error.""" + cache.put(mock_data_table) + with pytest.raises(TypeError, match="scalar"): + cache.get_range(mock_data_table["time"], mock_data_table["time"][-1]) + with pytest.raises(TypeError, match="scalar"): + cache.get_range(mock_data_table["time"][0], mock_data_table["time"]) + + +def test_lru_data_stays(cache): + """Test that the LRU data stays in the cache.""" + now = Time(datetime.now()) + times = now + np.arange(10) * timedelta(seconds=1) + data = QTable({"time": times, "value": np.arange(10)}) + cache.put(data) + assert len(cache.cache) == len(data) + + # Simulate accessing some data to update the LRU + assert cache.get(times[2])["value"] == 2 + assert cache.get(times[5])["value"] == 5 + + # reduce the cache size to trigger pruning + cache.maxsize = 3 + cache._prune() + + # Check that the accessed data is still in the cache + assert len(cache.cache) > 0 + assert len(cache.cache) <= 3 + assert cache.get(times[2]) is not None + assert cache.get(times[5]) is not None + + +def test_gloabl_ephemeris_cache(): + """Test that the global ephemeris cache is initialized.""" + assert STIX_EPHEMERIS_CACHE is not None + assert isinstance(STIX_EPHEMERIS_CACHE, TableLRUCache) + + +@pytest.mark.remote_data +def test_get_hpc_info_fills_cache(): + STIX_EPHEMERIS_CACHE.clear() + assert len(STIX_EPHEMERIS_CACHE.cache) == 0 + res = get_hpc_info(Time("2023-01-01T00:00:00"), Time("2023-01-01T00:00:10")) + assert res is not None + assert len(STIX_EPHEMERIS_CACHE.cache) > 0 + + +@pytest.mark.remote_data +def test_get_hpc_info_cache_hit(): + """Test that the cache is used when available.""" + STIX_EPHEMERIS_CACHE.clear() + assert len(STIX_EPHEMERIS_CACHE.cache) == 0 + res1 = get_hpc_info(Time("2023-01-01T12:00:00"), Time("2023-01-01T12:00:10")) + assert res1 is not None + assert len(STIX_EPHEMERIS_CACHE.cache) > 0 + + # Call again with the same time range to check if cache is used + res2 = STIX_EPHEMERIS_CACHE.get(Time("2023-01-01T12:00:00"), Time("2023-01-01T12:30:00")) + assert res2 is not None + assert res2["__lcu"][0] == 1 # code for HIT + + +@pytest.mark.remote_data +def test_get_hpc_cache_same_as_first_call(): + """Test that the cache is used when available.""" + STIX_EPHEMERIS_CACHE.clear() + assert len(STIX_EPHEMERIS_CACHE.cache) == 0 + + # just on time point + res1 = get_hpc_info(Time("2023-01-01T12:00:00")) + assert res1 is not None + assert len(STIX_EPHEMERIS_CACHE.cache) > 0 + + res2 = get_hpc_info(Time("2023-01-01T12:00:00")) + assert res1[0] == res2[0] # should be the same data + + +@pytest.mark.remote_data +def test_get_hpc_bypass_cache(): + """Test that the cache is used when available.""" + STIX_EPHEMERIS_CACHE.clear() + assert len(STIX_EPHEMERIS_CACHE.cache) == 0 + + # just on time point + res1 = get_hpc_info(Time("2023-01-01T12:00:00")) + assert res1 is not None + assert len(STIX_EPHEMERIS_CACHE.cache) > 0 + + # corrupt the cache to test bypass + for row_idx in range(len(STIX_EPHEMERIS_CACHE.cache)): + for angle_idx in range(3): + STIX_EPHEMERIS_CACHE.cache["roll_angle_rpy"][row_idx][angle_idx] = -9999 * u.deg + res2 = get_hpc_info(Time("2023-01-01T12:00:00")) + + assert res2[0] == -9999 * u.deg # should be from cache + STIX_EPHEMERIS_CACHE.bypass_cache_read = True + res_fresh = get_hpc_info(Time("2023-01-01T12:00:00")) + assert res_fresh[0] == res1[0] # should be fresh data not from cache so the same as res1 + + STIX_EPHEMERIS_CACHE.bypass_cache_read = False + res_refreshed = get_hpc_info(Time("2023-01-01T12:00:00")) + assert res_refreshed[0] == res_fresh[0] + + +@pytest.mark.remote_data +def test_get_hpc_source_cache(): + """Test that the cache is used when available.""" + STIX_EPHEMERIS_CACHE.clear() + assert len(STIX_EPHEMERIS_CACHE.cache) == 0 + + hpc_source = list() + res1 = get_hpc_info(Time("2023-01-01T12:00:00"), Time("2023-01-03T13:00:00"), return_source=hpc_source) + assert res1 is not None + assert len(STIX_EPHEMERIS_CACHE.cache) > 0 + # multiple days so multiple source files + assert len(hpc_source[0]) > 1 + assert "solo_ANC_stix-asp-ephemeris" in hpc_source[0] + + +@pytest.mark.remote_data +def test_load_anc_file(): + """Test loading an ANC file.""" + start_time = Time("2023-01-01T12:00:00") + end_time = Time("2023-01-01T12:30:00") + query = Fido.search( + a.Time(start_time, end_time), + a.Instrument.stix, + a.Level.anc, + a.stix.DataType.asp, + a.stix.DataProduct.asp_ephemeris, + ) + if len(query["stix"]) == 0: + raise ValueError(f"No STIX pointing data found for time range {start_time} to {end_time}.") + else: + query["stix"].filter_for_latest_version() + aux_files = Fido.fetch(query["stix"]) + + assert len(aux_files) > 0 + + STIX_EPHEMERIS_CACHE.clear() + for file in aux_files: + load_ephemeris_fits_to_cache(file) + + assert len(STIX_EPHEMERIS_CACHE.cache) > 0 + + # Call again with the same time range to check if cache is used + res2 = STIX_EPHEMERIS_CACHE.get(Time("2023-01-01T12:00:00"), Time("2023-01-01T12:30:00")) + assert res2 is not None + assert res2["__lcu"][0] == 1 # code for HIT diff --git a/stixpy/visualisation/plotters.py b/stixpy/visualisation/plotters.py index ad06781..710f144 100644 --- a/stixpy/visualisation/plotters.py +++ b/stixpy/visualisation/plotters.py @@ -118,10 +118,7 @@ def _setup_plot_elements(self, cmap): self.quadrant_font = {"weight": "regular", "size": 15} if cmap is None: - self.clrmap = copy.copy(plt.colormaps["viridis"]) - self.clrmap.set_over("gray") - self.clrmap.set_under("white") - self.clrmap.set_bad("gray") + self.clrmap = plt.colormaps["viridis"].with_extremes(over="gray", under="white", bad="gray") elif isinstance(cmap, str): self.clrmap = copy.copy(plt.colormaps[cmap]) diff --git a/tox.ini b/tox.ini index aba9133..4ef222f 100644 --- a/tox.ini +++ b/tox.ini @@ -59,11 +59,15 @@ commands_pre = commands = # To amend the pytest command for different factors you can add a line # which starts with a factor like `online: --remote-data=any \` - # If you have no factors which require different commands this is all you need: + # Online (remote-data) tests run serially (-n 1): the data server's network + # edge truncates large FITS responses under the concurrency of `-n auto`, so + # we avoid parallel downloads for the online job. (CI must NOT pass `-n auto` + # in posargs for the online env, or it would override this.) pytest \ -vvv \ -r fEs \ online: --remote-data=any \ + online: -n 1 \ --pyargs stixpy \ --cov-report=xml \ --cov=stixpy \