mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-08 10:05:23 +08:00
Scott Sanderson
62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
"""
|
|
Helpers for downsampling code.
|
|
"""
|
|
from operator import attrgetter
|
|
|
|
from zipline.utils.input_validation import expect_element
|
|
from zipline.utils.numpy_utils import changed_locations
|
|
from zipline.utils.sharedoc import (
|
|
templated_docstring,
|
|
PIPELINE_DOWNSAMPLING_FREQUENCY_DOC,
|
|
)
|
|
|
|
_dt_to_period = {
|
|
'year_start': attrgetter('year'),
|
|
'quarter_start': attrgetter('quarter'),
|
|
'month_start': attrgetter('month'),
|
|
'week_start': attrgetter('week'),
|
|
}
|
|
|
|
SUPPORTED_DOWNSAMPLE_FREQUENCIES = frozenset(_dt_to_period)
|
|
|
|
|
|
expect_downsample_frequency = expect_element(
|
|
frequency=SUPPORTED_DOWNSAMPLE_FREQUENCIES,
|
|
)
|
|
|
|
|
|
@expect_downsample_frequency
|
|
@templated_docstring(frequency=PIPELINE_DOWNSAMPLING_FREQUENCY_DOC)
|
|
def select_sampling_indices(dates, frequency):
|
|
"""
|
|
Choose entries from ``dates`` to use for downsampling at ``frequency``.
|
|
|
|
Parameters
|
|
----------
|
|
dates : pd.DatetimeIndex
|
|
Dates from which to select sample choices.
|
|
{frequency}
|
|
|
|
Returns
|
|
-------
|
|
indices : np.array[int64]
|
|
An array condtaining indices of dates on which samples should be taken.
|
|
|
|
The resulting index will always include 0 as a sample index, and it
|
|
will include the first date of each subsequent year/quarter/month/week,
|
|
as determined by ``frequency``.
|
|
|
|
Notes
|
|
-----
|
|
This function assumes that ``dates`` does not have large gaps.
|
|
|
|
In particular, it assumes that the maximum distance between any two entries
|
|
in ``dates`` is never greater than a year, which we rely on because we use
|
|
``np.diff(dates.<frequency>)`` to find dates where the sampling
|
|
period has changed.
|
|
"""
|
|
return changed_locations(
|
|
_dt_to_period[frequency](dates),
|
|
include_first=True
|
|
)
|