mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-01 03:40:51 +08:00
Eddie Hebert
e33f6dcdcd
Put the logic for reading and writing the equity price and adjustment data into a module located in data, making it distinct from the pipeline loader usage of the formats. This prepares for both incoming changes of how adjustments are written, (which includes using the bcolz daily reader as an input), as well as eventually providing the readers to a DataPortal object.
197 lines
6.4 KiB
Cython
197 lines
6.4 KiB
Cython
#
|
|
# Copyright 2015 Quantopian, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import bcolz
|
|
cimport cython
|
|
|
|
from numpy import (
|
|
array,
|
|
float64,
|
|
intp,
|
|
uint32,
|
|
zeros,
|
|
)
|
|
from numpy cimport (
|
|
float64_t,
|
|
intp_t,
|
|
ndarray,
|
|
uint32_t,
|
|
uint8_t,
|
|
)
|
|
from numpy.math cimport NAN
|
|
|
|
ctypedef object ctable_t
|
|
ctypedef object Timestamp_t
|
|
ctypedef object DatetimeIndex_t
|
|
ctypedef object Int64Index_t
|
|
|
|
|
|
@cython.boundscheck(False)
|
|
@cython.wraparound(False)
|
|
cpdef _compute_row_slices(dict asset_starts_absolute,
|
|
dict asset_ends_absolute,
|
|
dict asset_starts_calendar,
|
|
intp_t query_start,
|
|
intp_t query_end,
|
|
Int64Index_t requested_assets):
|
|
"""
|
|
Core indexing functionality for loading raw data from bcolz.
|
|
|
|
Parameters
|
|
----------
|
|
asset_starts_absolute : dict
|
|
Dictionary containing the index of the first row of each asset in the
|
|
bcolz file from which we will query.
|
|
|
|
asset_ends_absolute : dict
|
|
Dictionary containing the index of the last row of each asset in the
|
|
bcolz file from which we will query.
|
|
|
|
asset_starts_calendar : dict
|
|
Dictionary containing the index of in our calendar corresponding to the
|
|
start date of each asset
|
|
|
|
query_start : intp
|
|
query_end : intp
|
|
Start and end indices in our calendar of the dates for which we're
|
|
querying.
|
|
|
|
requested_assets : pandas.Int64Index
|
|
The assets for which we want to load data.
|
|
|
|
For each asset in requested assets, computes three values:
|
|
1.) The index in the raw bcolz data of first row to load.
|
|
2.) The index in the raw bcolz data of the last row to load.
|
|
3.) The index in the dates of our query corresponding to the first row for
|
|
each asset. This is non-zero iff the asset's lifetime begins partway
|
|
through the requested query dates.
|
|
|
|
Returns
|
|
-------
|
|
first_rows, last_rows, offsets : 3-tuple of ndarrays
|
|
"""
|
|
cdef:
|
|
intp_t nassets = len(requested_assets)
|
|
|
|
# For each sid, we need to compute the following:
|
|
ndarray[dtype=intp_t, ndim=1] first_row_a = zeros(nassets, dtype=intp)
|
|
ndarray[dtype=intp_t, ndim=1] last_row_a = zeros(nassets, dtype=intp)
|
|
ndarray[dtype=intp_t, ndim=1] offset_a = zeros(nassets, dtype=intp)
|
|
|
|
# Loop variables.
|
|
intp_t i
|
|
intp_t asset
|
|
intp_t asset_start_data
|
|
intp_t asset_end_data
|
|
intp_t asset_start_calendar
|
|
intp_t asset_end_calendar
|
|
|
|
for i, asset in enumerate(requested_assets):
|
|
asset_start_data = asset_starts_absolute[asset]
|
|
asset_end_data = asset_ends_absolute[asset]
|
|
asset_start_calendar = asset_starts_calendar[asset]
|
|
asset_end_calendar = (
|
|
asset_start_calendar + (asset_end_data - asset_start_data)
|
|
)
|
|
|
|
# If the asset started during the query, then start with the asset's
|
|
# first row.
|
|
# Otherwise start with the asset's first row + the number of rows
|
|
# before the query on which the asset existed.
|
|
first_row_a[i] = (
|
|
asset_start_data + max(0, (query_start - asset_start_calendar))
|
|
)
|
|
# If the asset ended during the query, the end with the asset's last
|
|
# row.
|
|
# Otherwise, end with the asset's last row minus the number of rows
|
|
# after the query for which the asset
|
|
last_row_a[i] = (
|
|
asset_end_data - max(0, asset_end_calendar - query_end)
|
|
)
|
|
# If the asset existed on or before the query, no offset.
|
|
# Otherwise, offset by the number of rows in the query in which the
|
|
# asset did not yet exist.
|
|
offset_a[i] = max(0, asset_start_calendar - query_start)
|
|
|
|
return first_row_a, last_row_a, offset_a
|
|
|
|
|
|
@cython.boundscheck(False)
|
|
@cython.wraparound(False)
|
|
cpdef _read_bcolz_data(ctable_t table,
|
|
tuple shape,
|
|
list columns,
|
|
intp_t[:] first_rows,
|
|
intp_t[:] last_rows,
|
|
intp_t[:] offsets):
|
|
"""
|
|
Load raw bcolz data for the given columns and indices.
|
|
|
|
Parameters
|
|
----------
|
|
table : bcolz.ctable
|
|
The table from which to read.
|
|
shape : tuple (length 2)
|
|
The shape of the expected output arrays.
|
|
columns : list[str]
|
|
List of column names to read.
|
|
|
|
first_rows : ndarray[intp]
|
|
last_rows : ndarray[intp]
|
|
offsets : ndarray[intp
|
|
Arrays in the format returned by _compute_row_slices.
|
|
|
|
Returns
|
|
-------
|
|
results : list of ndarray
|
|
A 2D array of shape `shape` for each column in `columns`.
|
|
"""
|
|
cdef:
|
|
int nassets
|
|
str column_name
|
|
ndarray[dtype=uint32_t, ndim=1] raw_data
|
|
ndarray[dtype=uint32_t, ndim=2] outbuf
|
|
ndarray[dtype=uint8_t, ndim=2, cast=True] where_nan
|
|
ndarray[dtype=float64_t, ndim=2] outbuf_as_float
|
|
intp_t asset
|
|
intp_t out_idx
|
|
intp_t raw_idx
|
|
intp_t first_row
|
|
intp_t last_row
|
|
intp_t offset
|
|
list results = []
|
|
|
|
nassets = shape[1]
|
|
if not nassets== len(first_rows) == len(last_rows) == len(offsets):
|
|
raise ValueError("Incompatible index arrays.")
|
|
|
|
for column_name in columns:
|
|
raw_data = table[column_name][:]
|
|
outbuf = zeros(shape=shape, dtype=uint32)
|
|
for asset in range(nassets):
|
|
first_row = first_rows[asset]
|
|
last_row = last_rows[asset]
|
|
offset = offsets[asset]
|
|
for out_idx, raw_idx in enumerate(range(first_row, last_row + 1)):
|
|
outbuf[out_idx + offset, asset] = raw_data[raw_idx]
|
|
|
|
if column_name in {'open', 'high', 'low', 'close'}:
|
|
where_nan = (outbuf == 0)
|
|
outbuf_as_float = outbuf.astype(float64) * .001
|
|
outbuf_as_float[where_nan] = NAN
|
|
results.append(outbuf_as_float)
|
|
else:
|
|
results.append(outbuf)
|
|
return results
|