"""Shared pieces of the resample contract: config checks, bucket labels, paths.
A *resample* aggregates a ``(timestamp, symbol)`` panel onto a coarser
regular time grid, one aggregation method per variable. Datasets and factors
carry the request in their config (``resample_freq`` and ``resample_how``),
map each source timestamp to the bar it belongs to, and hand the grouping to
their storage backend. The helpers here are the parts of that contract both
layers share; the backend's ``resample`` does the grouping itself.
"""
from pathlib import Path
import numpy as np
import pandas as pd
from quantlab.enums.data import RESAMPLE_FREQUENCY_SECONDS, RESAMPLE_METHODS
[docs]
def validate_resample_config(
freq: str | None, how: dict[str, str] | str | None, owner: str
) -> None:
"""Check a config's ``resample_freq`` and ``resample_how`` pair.
Parameters
----------
freq : str or None
The ``resample_freq`` field.
how : dict[str, str] or str or None
The ``resample_how`` field.
owner : str
Class name quoted in error messages.
Raises
------
ValueError
If ``freq`` is not a ``ResampleFrequency`` token, if ``how`` is set
without ``freq`` or ``freq`` without ``how``, or if any method is
not a ``ResampleMethod`` token.
Examples
--------
>>> validate_resample_config("1d", {"close": "last"}, "StockDataset")
>>> validate_resample_config("1d", None, "StockDataset")
Traceback (most recent call last):
ValueError: StockDataset: resample_freq='1d' needs resample_how ...
"""
if freq is None and how is None:
return
if freq is None:
raise ValueError(
f"{owner}: resample_how is set but resample_freq is None. Set "
f"both to resample, or neither to keep the store's own bars."
)
if freq not in RESAMPLE_FREQUENCY_SECONDS:
raise ValueError(
f"{owner}: resample_freq {freq!r} is not one of "
f"{list(RESAMPLE_FREQUENCY_SECONDS)}."
)
if how is None:
raise ValueError(
f"{owner}: resample_freq={freq!r} needs resample_how: one "
f"method for every variable (a str) or a {{variable: method}} "
f"dict. Methods: {list(RESAMPLE_METHODS)}."
)
methods = [how] if isinstance(how, str) else list(how.values())
if not isinstance(how, (str, dict)):
raise ValueError(
f"{owner}: resample_how must be a str or a dict, got "
f"{type(how).__name__}."
)
bad = sorted({m for m in methods if m not in RESAMPLE_METHODS})
if bad:
raise ValueError(
f"{owner}: resample_how uses unknown method(s) {bad}; choose "
f"from {list(RESAMPLE_METHODS)}."
)
[docs]
def resolve_resample_how(
how: dict[str, str] | str, variables: list[str], owner: str
) -> dict[str, str]:
"""Return one method per variable of the panel being resampled.
A string applies to every variable. A dict must name every variable:
an unlisted variable is refused rather than silently dropped or given
a default, and a listed name the panel does not have is refused too.
Parameters
----------
how : dict[str, str] or str
The ``resample_how`` field.
variables : list[str]
The panel's data variables.
owner : str
Class name quoted in error messages.
Examples
--------
>>> resolve_resample_how("last", ["bid", "ask"], "NbboPanelDataset")
{'bid': 'last', 'ask': 'last'}
>>> resolve_resample_how({"bid": "last"}, ["bid", "ask"], "NbboPanelDataset")
Traceback (most recent call last):
ValueError: NbboPanelDataset: resample_how does not name ['ask'] ...
"""
if isinstance(how, str):
return {name: how for name in variables}
missing = [name for name in variables if name not in how]
if missing:
raise ValueError(
f"{owner}: resample_how does not name {missing}; every variable "
f"of the panel needs a method (or pass one method as a str)."
)
unknown = [name for name in how if name not in variables]
if unknown:
raise ValueError(
f"{owner}: resample_how names {unknown}, which the panel does "
f"not have. Variables: {list(variables)}."
)
return {name: how[name] for name in variables}
[docs]
def resample_seconds(freq: str) -> int:
"""Return the length of a ``ResampleFrequency`` token in seconds.
Examples
--------
>>> resample_seconds("1h")
3600
"""
return RESAMPLE_FREQUENCY_SECONDS[freq]
[docs]
def source_interval_seconds(timestamps: np.ndarray) -> float:
"""Return the most common spacing of ``timestamps``, in seconds.
The mode is used rather than the minimum so that gaps such as
weekends do not distort the answer.
Examples
--------
>>> source_interval_seconds(pd.date_range("2024-01-01", periods=3, freq="min"))
60.0
"""
index = pd.DatetimeIndex(timestamps)
if len(index) < 2:
raise ValueError(
"a panel needs at least two timestamps to know its bar size"
)
return float(pd.Series(np.diff(index.values)).mode().iloc[0] / np.timedelta64(1, "s"))
[docs]
def assert_coarser(timestamps: np.ndarray, freq: str, owner: str) -> None:
"""Refuse a resample onto a grid no coarser than the panel's own bars.
Raises
------
ValueError
If ``freq`` is not longer than the panel's bar size.
Examples
--------
>>> minutes = pd.date_range("2024-01-01", periods=10, freq="min")
>>> assert_coarser(minutes, "5m", "DemoDataset")
>>> assert_coarser(minutes, "1m", "DemoDataset")
Traceback (most recent call last):
ValueError: DemoDataset: resample_freq='1m' (60s) is not coarser ...
"""
source = source_interval_seconds(timestamps)
target = resample_seconds(freq)
if target <= source:
raise ValueError(
f"{owner}: resample_freq={freq!r} ({target}s) is not coarser "
f"than the panel's own bars ({source:g}s). A resample only "
f"aggregates onto a coarser grid."
)
[docs]
def clock_labels(timestamps: np.ndarray, freq: str) -> np.ndarray:
"""Return each timestamp's bar, floored on the UTC clock.
This is the default bucketing: a bar starts at a multiple of ``freq``
counted from midnight UTC and is labelled at its start, which matches
bars labelled at their open time and daily stores stamped at midnight.
Examples
--------
>>> ts = pd.to_datetime(["2024-01-01 09:31", "2024-01-01 09:34", "2024-01-01 09:35"])
>>> clock_labels(ts, "5m").astype("datetime64[m]").tolist() # doctest: +SKIP
[2024-01-01T09:30, 2024-01-01T09:30, 2024-01-01T09:35]
"""
index = pd.DatetimeIndex(timestamps)
return index.floor(f"{resample_seconds(freq)}s").values
[docs]
def session_labels(
timestamps: np.ndarray, freq: str, sessions: pd.DataFrame, owner: str
) -> np.ndarray:
"""Return each timestamp's bar inside its trading session.
Bars here are labelled at their *end*, as ``NbboResampler`` labels them,
so a source timestamp belongs to the session whose ``open < t <=
close``. ``"1d"`` labels the whole session with its date at midnight, so
a daily panel lines up with daily stores. Any other frequency cuts the
session into right-closed bars from its open, labelled at their end; a
partial last bar is labelled at the session close.
Parameters
----------
timestamps : np.ndarray
Source timestamps, naive UTC.
freq : str
A ``ResampleFrequency`` token.
sessions : pd.DataFrame
Columns ``date``, ``open`` and ``close`` (naive UTC), one row per
session, sorted by ``open``.
owner : str
Class name quoted in error messages.
Raises
------
ValueError
If a timestamp falls in no session.
Examples
--------
>>> sessions = pd.DataFrame({
... "date": [pd.Timestamp("2024-01-24").date()],
... "open": [pd.Timestamp("2024-01-24 14:30")],
... "close": [pd.Timestamp("2024-01-24 21:00")],
... })
>>> ts = pd.to_datetime(["2024-01-24 14:31", "2024-01-24 21:00"])
>>> session_labels(ts, "1d", sessions, "Demo").astype("datetime64[D]").tolist() # doctest: +SKIP
[2024-01-24, 2024-01-24]
"""
index = pd.DatetimeIndex(timestamps).values.astype("datetime64[ns]")
opens = pd.to_datetime(sessions["open"]).values.astype("datetime64[ns]")
closes = pd.to_datetime(sessions["close"]).values.astype("datetime64[ns]")
dates = pd.to_datetime(sessions["date"]).values.astype("datetime64[ns]")
# The session whose close is the first at or after t.
idx = np.searchsorted(closes, index, side="left")
outside = idx >= len(closes)
idx_safe = np.where(outside, 0, idx)
outside |= ~(opens[idx_safe] < index)
if outside.any():
first = pd.Timestamp(index[outside][0])
raise ValueError(
f"{owner}: timestamp {first} falls in no trading session; a "
f"session-based resample needs every bar inside a session."
)
if freq == "1d":
return dates[idx]
step = np.timedelta64(resample_seconds(freq), "s").astype("timedelta64[ns]")
offset = index - opens[idx]
bars = -(-offset // step) # ceil: a bar is labelled at its end
labels = opens[idx] + bars * step
return np.minimum(labels, closes[idx])
[docs]
def resample_store_path(path: str | None, freq: str | None) -> str | None:
"""Return the store a resampled object reads and writes.
The resampled store sits beside the source store with ``_resample_<freq>``
added to its name: ``klines.zarr`` becomes ``klines_resample_1d.zarr``.
With ``freq`` unset the source path is returned unchanged; a ``None``
path stays ``None``.
Examples
--------
>>> resample_store_path("data/klines.zarr", "1d")
'data/klines_resample_1d.zarr'
>>> resample_store_path("data/klines.zarr", None)
'data/klines.zarr'
"""
if path is None or freq is None:
return path
source = Path(path)
return str(source.with_name(f"{source.stem}_resample_{freq}{source.suffix}"))