"""Top-level orchestration: fetch -> process -> Dataset."""
from datetime import datetime
import numpy as np
from .dataset import build_dataset
from .dates import coerce_datetime, date_value, iso_timestamp, yesterday
from .processing import process_channels
from .sources import bcwind_samples, omni_samples, omni_samples_hapi
DEFAULT_OMNI_START = "1982-01-01"
[docs]
def generate_imf(
start=None,
end=None,
source="omni",
window=10,
cache_dir=None,
bcwind_path=None,
download=True,
omni_access="hapi",
verbose=False,
model="tiegcm",
):
"""Generate an IMF ``xarray.Dataset``.
Parameters
----------
start, end : date-like or None
Inclusive bounds (``YYYY-MM-DD``, ``YYYYDDD``, ISO, or ``datetime``). For
``omni`` they default to ``1982-01-01`` and yesterday. For ``bcwind`` they
optionally filter the file (default: the file's full span).
source : {"omni", "bcwind"}
OMNI 1-minute data (default) or a BCWIND HDF5 file.
window : int
Trailing-average window in minutes for the OMNI pipeline (default 10).
Ignored for ``bcwind`` (raw pass-through).
cache_dir : str or None
Directory holding / receiving ``omni_min<year>.asc`` files
(``omni_access="asc"`` only).
bcwind_path : str or None
Path to the BCWIND HDF5 file (required when ``source="bcwind"``).
download : bool
Fetch missing OMNI year files over FTP (``omni_access="asc"`` only;
default True).
omni_access : {"hapi", "asc"}
How to obtain OMNI data. ``"hapi"`` (default) fetches only the requested
window from CDAWeb's HAPI server (no whole-year download) -- best for
short ranges. ``"asc"`` downloads/parses the SPDF ``omni_min<year>.asc``
files and reproduces the legacy per-year output exactly. Ignored for
``bcwind``.
verbose : bool
Print progress.
model : {"tiegcm", "waccmx"}
Target model input format (default ``"tiegcm"``). ``"waccmx"`` emits the
WACCM-X format (unlimited ``time``; ``date`` as ``YYYYMMDD`` plus
``datefrac``/``datesec``; a ``WACCMX`` filename tag).
Returns
-------
xarray.Dataset
Use :func:`imfgen.save_imf` to write it to NetCDF.
"""
if source == "omni":
start_dt = coerce_datetime(start if start is not None else DEFAULT_OMNI_START)
end_dt = (coerce_datetime(end, end_of_day=True)
if end is not None else yesterday())
if verbose:
print(f"Generating OMNI IMF {start_dt.date()} -> {end_dt.date()} "
f"(window={window}, access={omni_access})")
if omni_access == "hapi":
samples = omni_samples_hapi(start_dt, end_dt, window=window,
verbose=verbose)
elif omni_access == "asc":
samples = omni_samples(start_dt, end_dt, window=window,
cache_dir=cache_dir, download=download,
verbose=verbose)
else:
raise ValueError(
f"Unknown omni_access {omni_access!r}; expected 'hapi' or 'asc'."
)
source_path = None
elif source == "bcwind":
if not bcwind_path:
raise ValueError("source='bcwind' requires bcwind_path=<file.h5>.")
start_dt = coerce_datetime(start) if start is not None else None
end_dt = (coerce_datetime(end, end_of_day=True)
if end is not None else None)
if verbose:
print(f"Converting BCWIND file {bcwind_path}")
samples = bcwind_samples(bcwind_path, start_dt=start_dt, end_dt=end_dt)
source_path = bcwind_path
else:
raise ValueError(f"Unknown source {source!r}; expected 'omni' or 'bcwind'.")
timestamps = samples["timestamps"]
n_out = len(timestamps)
processed = process_channels(
samples["channels"], samples["window"], n_out, samples["interpolate"]
)
dates = np.array([date_value(t) for t in timestamps])
iso = np.array([iso_timestamp(t) for t in timestamps])
ds = build_dataset(processed, dates, iso, source=source,
source_path=source_path, model=model, datetimes=timestamps)
if verbose:
print(f"Built IMF dataset: {ds.attrs['yearday_beg']} -> "
f"{ds.attrs['yearday_end']} ({n_out} minutes)")
return ds
[docs]
def generate_imf_years(start=None, end=None, window=10, cache_dir=None,
download=True, omni_access="hapi", verbose=False,
model="tiegcm"):
"""Yield one OMNI ``Dataset`` per calendar year in ``[start, end]``.
Each year is generated **independently** (its own within-year interpolation),
so the per-year files reproduce the legacy ``imf_OMNI_YYYY001-YYYYddd.nc``
files. This is what ``imfgen --split-years`` writes. (BCWIND files are a
single span and are not split.) For bit-for-bit reproduction of the legacy
files use ``omni_access="asc"``.
"""
start_dt = coerce_datetime(start if start is not None else DEFAULT_OMNI_START)
end_dt = (coerce_datetime(end, end_of_day=True)
if end is not None else yesterday())
for year in range(start_dt.year, end_dt.year + 1):
y_start = max(start_dt, datetime(year, 1, 1))
y_end = min(end_dt, datetime(year, 12, 31, 23, 59))
yield generate_imf(
start=y_start, end=y_end, source="omni", window=window,
cache_dir=cache_dir, download=download, omni_access=omni_access,
verbose=verbose, model=model,
)