time-split 0.1.0__py3-none-any.whl

time_split/_support.py

@@ -0,0 +1,14 @@
+from typing import Protocol, TypeVar, runtime_checkable
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+@runtime_checkable
+class _Computable(Protocol[T_co]):
+    def compute(self) -> T_co:
+        """Compute results of a delayed Dask object."""
+
+
+def handle_dask(arg: T_co | _Computable[T_co]) -> T_co:
+    """Run Dask compute function on a delayed object."""
+    return arg.compute() if isinstance(arg, _Computable) else arg

time_split/integration/__init__.py

@@ -0,0 +1 @@
+"""Convenience functions and classes for common libraries."""

time_split/integration/_log_progress.py

@@ -0,0 +1,25 @@
+import logging
+from collections.abc import Iterable
+from typing import Any, Union
+
+from .. import log_split_progress
+from .._docstrings import add_docstrings
+from ..types import DatetimeSplitBounds, DatetimeSplits
+
+LogProgressArg = Union[str, bool, dict[str, Any], logging.Logger, logging.LoggerAdapter]  # type: ignore[type-arg]
+
+add_docstrings(log_progress="Controls logging of fold progress. See :func:`~.log_split_progress` for details.")
+
+
+def handle_log_progress_arg(
+    log_progress: LogProgressArg, *, splits: DatetimeSplits
+) -> Iterable[DatetimeSplitBounds] | None:
+    """Wrapper function for integrations."""
+    if log_progress is True:
+        return log_split_progress(splits)
+    elif isinstance(log_progress, (str, logging.Logger, logging.LoggerAdapter)):
+        return log_split_progress(splits, logger=log_progress)
+    elif isinstance(log_progress, dict):
+        return log_split_progress(splits, **log_progress)
+
+    return None

time_split/integration/pandas/__init__.py

@@ -0,0 +1,42 @@
+r"""Integration with the Pandas library.
+
+Examples:
+    Splitting a ``pandas.Series`` with :func:`split_pandas`.
+
+    >>> import pandas as pd
+    >>> index = pd.date_range("2022", "2022-1-10", freq="h")
+    >>> series = pd.Series(range(len(index)), index=index)
+    >>> series.sample(3, random_state=1999)
+    2022-01-04 00:00:00    72
+    2022-01-03 21:00:00    69
+    2022-01-04 14:00:00    86
+    dtype: int64
+
+    Series may only be split on the index. The `schedule` keyword argument is required, but `log_progress` is not.
+
+    >>> for fold in split_pandas(
+    ...     series, schedule="1d", log_progress="progress"
+    ... ):  # doctest: +SKIP
+    ...     print(
+    ...         f"Summary of fold {tuple(map(pd.Timestamp.isoformat, fold.bounds))}:"
+    ...         f"\n  {fold.data.mean()=}"
+    ...         f"\n  {fold.future_data.mean()=}",
+    ...     )
+    INFO:progress:Begin fold 1/2: ('2022-01-01' <= [schedule: '2022-01-08' (Saturday)] < '2022-01-09').
+    Summary of fold ('2022-01-01T00:00:00', '2022-01-08T00:00:00', '2022-01-09T00:00:00'):
+      fold.data.mean()=83.5
+      fold.future_data.mean()=179.5
+    INFO:progress:Finished fold 1/2 [schedule: '2022-01-08' (Saturday)] after 1ms.  # doctest: +SKIP
+    INFO:progress:Begin fold 2/2: ('2022-01-02' <= [schedule: '2022-01-09' (Sunday)] < '2022-01-10').
+    Summary of fold ('2022-01-02T00:00:00', '2022-01-09T00:00:00', '2022-01-10T00:00:00'):
+      fold.data.mean()=107.5
+      fold.future_data.mean()=203.5
+    INFO:progress:Finished fold 2/2 [schedule: '2022-01-09' (Sunday)] after 873μs.
+
+    When splitting dataframes, you may optionally pass a `time_column` argument as well. By default, both frames and
+    series are split along the index.
+"""
+
+from ._impl import PandasT, split_pandas
+
+__all__ = ["PandasT", "split_pandas"]

time_split/integration/pandas/_impl.py

@@ -0,0 +1,81 @@
+from collections.abc import Hashable, Iterable
+from datetime import date, datetime
+from typing import Generic, TypeVar, Unpack
+
+from pandas import DataFrame, DatetimeIndex, Series, Timestamp
+from rics.misc import tname
+
+from ..._docstrings import docs
+from ...types import DatetimeIndexSplitterKwargs
+from .._log_progress import LogProgressArg
+from ..split_data import DatetimeSplit, split_data
+
+PandasT = TypeVar("PandasT", Series, DataFrame)
+"""A splittable pandas type."""
+
+
+@docs
+def split_pandas(
+    data: PandasT,
+    time_column: Hashable = None,
+    *,
+    log_progress: LogProgressArg = False,
+    **kwargs: Unpack[DatetimeIndexSplitterKwargs],
+) -> Iterable[DatetimeSplit[PandasT]]:
+    """Split a pandas type.
+
+    This function splits indexed data (i.e. ``Series`` and ``DataFrame``, not the index itself. Use
+    :func:`time_split.split <time_split.split>` for pandas ``Index`` types, setting ``available=data.index``.
+
+    Args:
+        data: A pandas data container type to split; either ``Series`` or a ``DataFrame``.
+        time_column: A column in `data` to split on. Use ``data`.index`` if ``None``.
+        log_progress: {log_progress}
+        **kwargs: {DatetimeIndexSplitterKwargs}
+
+    {USER_GUIDE}
+
+    Yields:
+        Tuples ``(data, future_data, bounds)``.
+
+    Raises:
+        TypeError: If `time_column` does not denote a datetime index-like field.
+
+    """
+    indexer = _Indexer(time_column)
+
+    yield from split_data(
+        data,
+        log_progress=log_progress,
+        as_available=indexer.as_available,
+        select=indexer.select,
+        **kwargs,
+    )
+
+
+class _Indexer(Generic[PandasT]):
+    def __init__(self, time_column: Hashable | None) -> None:
+        self.time_column = time_column
+
+    def as_available(self, data: PandasT) -> Series | DatetimeIndex:
+        time = self._get_time(data)
+
+        first = time.iloc[0] if hasattr(time, "iloc") else time[0]
+        if isinstance(first, date):
+            return time
+
+        type_str = "data.index" if self.time_column is None else f"data[{self.time_column!r}]"
+        msg = f"Elements of {type_str} element are {tname(first)}, expected datetime-like."
+        raise TypeError(msg)
+
+    def _get_time(self, data: DataFrame) -> Series:
+        return data.index if self.time_column is None else data[self.time_column]
+
+    def select(self, data: PandasT, left: datetime, right: datetime) -> PandasT:
+        """Select data based on the given bounds."""
+        time = self._get_time(data)
+        if isinstance(time, Series):
+            return data[time.between(left, right, inclusive="left")]
+        else:
+            # Index slicing is a lot faster than boolean masks (empirically, seems to be a factor ~10).
+            return data[left : right - Timestamp.resolution]  # type: ignore[misc]

time_split/integration/polars/__init__.py

@@ -0,0 +1,57 @@
+r"""Integration with the Polars library.
+
+Examples:
+    Splitting a ``polars.DataFrame`` with :func:`split_polars`.
+
+    >>> import polars as pl
+    >>> from datetime import date
+    >>> ts = pl.datetime_range(
+    ...     date.fromisoformat("2022-01-01"),
+    ...     date.fromisoformat("2022-01-10"),
+    ...     interval="1h",
+    ...     eager=True,
+    ... )
+    >>> df = pl.DataFrame({"timestamp": ts, "ints": range(len(ts))})
+    >>> df.sample(5, seed=1999)
+    shape: (5, 2)
+    ┌─────────────────────┬──────┐
+    │ timestamp           ┆ ints │
+    │ ---                 ┆ ---  │
+    │ datetime[μs]        ┆ i64  │
+    ╞═════════════════════╪══════╡
+    │ 2022-01-06 16:00:00 ┆ 136  │
+    │ 2022-01-08 17:00:00 ┆ 185  │
+    │ 2022-01-01 06:00:00 ┆ 6    │
+    │ 2022-01-01 05:00:00 ┆ 5    │
+    │ 2022-01-03 08:00:00 ┆ 56   │
+    └─────────────────────┴──────┘
+
+    To split the frame, pass it to ``split_polars`` along with the column to split on. The `schedule` keyword argument
+    is required, but `log_progress` is not.
+
+    >>> for fold in split_polars(
+    ...     df, schedule="1d", log_progress="progress", time_column="timestamp"
+    ... ):  # doctest: +SKIP
+    ...     print(
+    ...         f"Summary of fold {tuple(map(pd.Timestamp.isoformat, fold.bounds))}:"
+    ...         f"\n  {fold.data['ints'].mean()=}"
+    ...         f"\n  {fold.future_data['ints'].mean()=}",
+    ...     )
+    INFO:progress:Begin fold 1/2: '2022-01-01' <= [schedule: '2022-01-08' (Saturday)] < '2022-01-09'.
+    Summary of fold ('2022-01-01T00:00:00', '2022-01-08T00:00:00', '2022-01-09T00:00:00'):
+      fold.data['ints'].mean()=83.5
+      fold.future_data['ints'].mean()=179.5
+    INFO:progress:Finished fold 1/2: [schedule: '2022-01-08' (Saturday)] after 542μs.
+    INFO:progress:Begin fold 2/2: '2022-01-02' <= [schedule: '2022-01-09' (Sunday)] < '2022-01-10'.
+    Summary of fold ('2022-01-02T00:00:00', '2022-01-09T00:00:00', '2022-01-10T00:00:00'):
+      fold.data['ints'].mean()=107.5
+      fold.future_data['ints'].mean()=203.5
+    INFO:progress:Finished fold 2/2: [schedule: '2022-01-09' (Sunday)] after 433μs.
+
+"""
+
+from ._impl import split_polars
+
+__all__ = [
+    "split_polars",
+]

time_split/integration/polars/_impl.py

@@ -0,0 +1,69 @@
+from collections.abc import Iterable
+from datetime import date, datetime
+from typing import Unpack
+
+from polars import DataFrame, Series
+from rics.misc import tname
+
+from ..._docstrings import docs
+from ...types import DatetimeIndexSplitterKwargs
+from .._log_progress import LogProgressArg
+from ..split_data import DatetimeSplit, split_data
+
+
+@docs
+def split_polars(
+    data: DataFrame,
+    time_column: str,
+    *,
+    log_progress: LogProgressArg = False,
+    **kwargs: Unpack[DatetimeIndexSplitterKwargs],
+) -> Iterable[DatetimeSplit[DataFrame]]:
+    """Split a polars frame.
+
+    Args:
+        data: A ``polars.DataFrame``.
+        time_column: A column to split on.
+        log_progress: {log_progress}
+        **kwargs: {DatetimeIndexSplitterKwargs}
+
+    {USER_GUIDE}
+
+    Yields:
+        Tuples ``(data, future_data, bounds)``.
+
+    Raises:
+        TypeError: If `time_column` does not denote a datetime index-like field.
+
+    """
+    indexer = _Indexer(time_column)
+
+    yield from split_data(
+        data,
+        log_progress=log_progress,
+        as_available=indexer.as_available,
+        select=indexer.select,
+        **kwargs,
+    )
+
+
+class _Indexer:
+    def __init__(self, time_column: str) -> None:
+        self.time_column = time_column
+
+    def as_available(self, data: DataFrame) -> Series:
+        time = self._get_time(data)
+
+        if isinstance(time[0], date):
+            return time
+
+        msg = f"Elements of data[{self.time_column!r}] element are {tname(time[0])}, expected datetime-like."
+        raise TypeError(msg)
+
+    def _get_time(self, data: DataFrame) -> Series:
+        return data[self.time_column]
+
+    def select(self, data: DataFrame, left: datetime, right: datetime) -> DataFrame:
+        """Select data based on the given bounds."""
+        time = self._get_time(data)
+        return data.filter(time.is_between(left, right, closed="left"))

time_split/integration/sklearn/__init__.py

@@ -0,0 +1,7 @@
+"""Integration with the scikit-learn library."""
+
+from ._impl import ScikitLearnSplitter
+
+__all__ = [
+    "ScikitLearnSplitter",
+]

time_split/integration/sklearn/_impl.py

@@ -0,0 +1,170 @@
+from collections.abc import Iterable, Sequence
+from typing import Any, cast, get_args
+
+from numpy import array, datetime64, logical_and, ndarray, nonzero
+from numpy.typing import NDArray
+
+from ..._backend import DatetimeIndexSplitter
+from ..._docstrings import docs
+from ...types import (
+    DatetimeIterable,
+    DatetimeSplits,
+    DatetimeTypes,
+    Flex,
+    Schedule,
+    Span,
+)
+from .._log_progress import LogProgressArg, handle_log_progress_arg
+
+try:
+    from sklearn.model_selection import BaseCrossValidator  # type: ignore[import-untyped]
+
+except ModuleNotFoundError:
+    BaseCrossValidator = object
+
+IndexTuple = tuple[Sequence[int], Sequence[int]]
+
+
+@docs
+class ScikitLearnSplitter(BaseCrossValidator):  # type: ignore[misc]
+    """A scikit-learn compatible datetime splitter.
+
+    This class may be used to create temporal folds from heterogeneous/unaggregated data, typically used for training
+    models (e.g. on raw transaction data). If your data is a well-formed time series, consider using the
+    `TimeSeriesSplit`_ class from scikit-learn instead.
+
+    .. _TimeSeriesSplit: https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.TimeSeriesSplit.html
+
+    If a ``pandas`` type is passed to the :meth:`ScikitLearnSplitter.split`-method, the index will be used.
+
+    Args:
+        schedule: {schedule}
+        before: {before}
+        after: {after}
+        step: {step}
+        n_splits: {n_splits}
+        flex: {flex}
+        log_progress: {log_progress}
+        verify_xy: If ``True``, split X and y independently and verify that they are equal.
+
+    {USER_GUIDE}
+
+    """
+
+    def __init__(
+        self,
+        schedule: Schedule,
+        *,
+        before: Span = "7d",
+        after: Span = 1,
+        n_splits: int | None = None,
+        flex: Flex = "auto",
+        step: int = 1,
+        log_progress: LogProgressArg = False,
+        verify_xy: bool = True,
+    ) -> None:
+        super().__init__()
+        self._splitter = DatetimeIndexSplitter(
+            schedule,
+            before=before,
+            after=after,
+            step=step,
+            n_splits=n_splits,
+            flex=flex,
+        )
+        self.log_progress = log_progress
+        self.verify_xy = verify_xy
+
+    def get_n_splits(
+        self, X: DatetimeIterable | None = None, y: DatetimeIterable | None = None, groups: Any = None
+    ) -> int:
+        """Returns the number of splitting iterations in the cross-validator.
+
+        Equivalent to ``len(list(split(X, y, groups))``.
+
+        Args:
+            X: Training data (features).
+            y: Target variable.
+            groups: Always ignored, exists for compatibility.
+
+        Returns:
+            Number of splits with given arguments.
+
+        Raises:
+            ValueError: If both `X` and `y` are ``None``.
+            ValueError: If splits of `X` and `y` are not equal when ``verify_xy=True``.
+
+        """
+        _, splits = self._get_splits(X, y)
+        return len(splits)
+
+    def split(
+        self,
+        X: DatetimeIterable | None = None,
+        y: DatetimeIterable | None = None,
+        groups: Any = None,
+    ) -> Iterable[IndexTuple]:
+        """Generate indices to split data into training and test set.
+
+        Args:
+            X: Training data (features).
+            y: Target variable.
+            groups: Always ignored, exists for compatibility.
+
+        Yields:
+            The training/test set indices for that split.
+
+        Raises:
+            ValueError: If both `X` and `y` are ``None``.
+            ValueError: If splits of `X` and `y` are not equal when ``verify_xy=True``.
+            TypeError: If `X` or `y` have an ``index``-attribute, but index elements are not datetime-like.
+
+        """
+        index, splits = self._get_splits(self._handle_pandas(X, "X"), self._handle_pandas(y, "y"))
+        for fold in handle_log_progress_arg(self.log_progress, splits=splits) or splits:
+            yield cast(
+                IndexTuple,
+                (
+                    nonzero(logical_and(fold.start <= index, index < fold.mid))[0],
+                    nonzero(logical_and(fold.mid <= index, index < fold.end))[0],
+                ),
+            )
+
+    @staticmethod
+    def _handle_pandas(arg: DatetimeIterable | None, name: str) -> DatetimeIterable | None:
+        if arg is None or not hasattr(arg, "index"):
+            return arg
+
+        index = arg.index
+
+        if type(arg.index[0]) in get_args(DatetimeTypes):
+            return cast(DatetimeIterable, index)
+
+        raise TypeError(f"{name}.index does not appear to be a datetime-like iterable.")
+
+    def _get_splits(
+        self,
+        X: DatetimeIterable | None = None,
+        y: DatetimeIterable | None = None,
+    ) -> tuple[NDArray[datetime64], DatetimeSplits]:
+        splits: DatetimeSplits | None = None
+        y_splits: DatetimeSplits | None = None
+        timestamps: Any | None = None
+
+        if X is not None:
+            splits = self._splitter.get_splits(X)
+            timestamps = X
+        if y is not None:
+            y_splits = self._splitter.get_splits(y)
+            timestamps = y
+        if self.verify_xy and splits and y_splits and splits != y_splits:
+            raise ValueError("Splits of X and y are not equal.")
+
+        # Cast should not be needed, but MyPy things this is nullable.
+        if splits is None and y_splits is None:
+            raise ValueError("At least one of (X, y) must be given.")
+
+        timestamps = timestamps if isinstance(timestamps, ndarray) else array(timestamps)
+        if len(timestamps.shape) > 1:
+            raise NotImplementedError(f"shape {timestamps.shape} not supported")
+        return timestamps, cast(DatetimeSplits, splits or y_splits)

time_split/integration/split_data.py

@@ -0,0 +1,72 @@
+"""Base implementations for splitting generic data types.
+
+Users may implement splitting of any data type by implementing suitable ``as_available`` and ``select`` functions.
+"""
+
+import typing as _t
+from datetime import datetime as _datetime
+
+from .. import _frontend, types
+from .._docstrings import docs as _docs
+from . import _log_progress
+
+DataT = _t.TypeVar("DataT")
+"""Type of data to split."""
+
+DataAsAvailableFn = _t.Callable[[DataT], types.DatetimeIterable]
+"""A callable ``(data: DataT) -> DatetimeIterable``."""
+DataSelectFn = _t.Callable[[DataT, _datetime, _datetime], DataT]
+"""A callable ``(data: DataT, left_inclusive: datetime, end_exclusive: datetime) -> DataT)``."""
+
+
+class DatetimeSplit(_t.NamedTuple, _t.Generic[DataT]):
+    """Time-based split of a generic data type."""
+
+    data: DataT
+    """Data before ``bounds.mid``."""
+    future_data: DataT
+    """Data after ``bounds.mid``."""
+    bounds: types.DatetimeSplitBounds
+    """The underlying bounds that produced this split."""
+
+
+@_docs
+def split_data(
+    data: DataT,
+    *,
+    log_progress: _log_progress.LogProgressArg = False,
+    as_available: DataAsAvailableFn[DataT],
+    select: DataSelectFn[DataT],
+    **kwargs: _t.Unpack[types.DatetimeIndexSplitterKwargs],
+) -> _t.Iterable[DatetimeSplit[DataT]]:
+    """Base implementation for splitting integrated `data` types.
+
+    The required ``as_available`` and ``select`` callables provided perform the actual integration.
+
+    Args:
+        data: The data to split.
+        log_progress: {log_progress}
+        as_available: A callable ``(data: DataT) -> DatetimeIterable``.
+        select: A callable ``(data: DataT, left_inclusive: datetime, end_exclusive: datetime) -> DataT)``.
+        **kwargs: Keyword arguments for :func:`.split`-function.
+
+    Yields:
+        Tuples ``(data, future_data, bounds)``.
+
+    See Also:
+        To get started with your own integration, copy :func:`~time_split.integration.pandas.split_pandas` or
+        :func:`~time_split.integration.polars.split_polars` and use it as the baseline (click ``[source]``) on
+        the linked function.
+
+    """
+    available = as_available(data)
+    splits = _frontend.split(**kwargs, available=available)
+
+    tracked_splits = _log_progress.handle_log_progress_arg(log_progress, splits=splits)
+
+    for bounds in splits if tracked_splits is None else tracked_splits:
+        yield DatetimeSplit(
+            select(data, bounds.start, bounds.mid),
+            future_data=select(data, bounds.mid, bounds.end),
+            bounds=bounds,
+        )