zipline_polygon_bundle 0.1.8__py3-none-any.whl → 0.2.0__py3-none-any.whl

zipline_polygon_bundle/compute_signals.py

@@ -0,0 +1,261 @@
+from .config import PolygonConfig
+from .trades import custom_aggs_schema, custom_aggs_partitioning
+
+import datetime
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pa_compute
+import pyarrow.dataset as pa_ds
+import pandas_ta as ta
+import pandas as pd
+
+
+def calculate_mfi(typical_price: pd.Series, money_flow: pd.Series, period: int):
+    mf_sign = np.where(typical_price > np.roll(typical_price, shift=1), 1, -1)
+    signed_mf = money_flow * mf_sign
+
+    # Calculate gain and loss using vectorized operations
+    positive_mf = np.maximum(signed_mf, 0)
+    negative_mf = np.maximum(-signed_mf, 0)
+
+    mf_avg_gain = (
+        np.convolve(positive_mf, np.ones(period), mode="full")[: len(positive_mf)]
+        / period
+    )
+    mf_avg_loss = (
+        np.convolve(negative_mf, np.ones(period), mode="full")[: len(negative_mf)]
+        / period
+    )
+
+    epsilon = 1e-10  # Small epsilon value to avoid division by zero
+    mfi = 100 - (100 / (1 + mf_avg_gain / (mf_avg_loss + epsilon)))
+    return mfi
+
+
+# https://github.com/twopirllc/pandas-ta/blob/main/pandas_ta/momentum/stoch.py
+# https://github.com/twopirllc/pandas-ta/blob/development/pandas_ta/momentum/stoch.py
+# `k` vs `fast_k` arg names.
+# https://github.com/twopirllc/pandas-ta/issues/726
+# Results affected by values outside range
+# https://github.com/twopirllc/pandas-ta/issues/535
+
+
+def calculate_stoch(
+    high: pd.Series,
+    low: pd.Series,
+    close: pd.Series,
+    k: int = 14,
+    d: int = 3,
+    smooth_k: int = 3,
+    mamode: str = "sma",
+):
+    """Indicator: Stochastic Oscillator (STOCH)"""
+    lowest_low = low.rolling(k).min()
+    highest_high = high.rolling(k).max()
+
+    stoch = 100 * (close - lowest_low)
+    stoch /= ta.utils.non_zero_range(highest_high, lowest_low)
+
+    stoch_k = ta.overlap.ma(
+        mamode, stoch.loc[stoch.first_valid_index() :,], length=smooth_k
+    )
+    stoch_d = (
+        ta.overlap.ma(mamode, stoch_k.loc[stoch_k.first_valid_index() :,], length=d)
+        if stoch_k is not None
+        else None
+    )
+    # Histogram
+    stoch_h = stoch_k - stoch_d if stoch_d is not None else None
+
+    return stoch_k, stoch_d, stoch_h
+
+
+def compute_per_ticker_signals(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
+    df = df.set_index("window_start").sort_index()
+    session_index = pd.date_range(
+        start=df.index[0], end=df.index[-1], freq=pd.Timedelta(seconds=60)
+    )
+    df = df.reindex(session_index)
+    df.index.rename("window_start", inplace=True)
+
+    # df["minute_of_day"] = (df.index.hour * 60) + df.index.minute
+    # df["day_of_week"] = df.index.day_of_week
+
+    df.transactions = df.transactions.fillna(0)
+    df.volume = df.volume.fillna(0)
+    df.total = df.total.fillna(0)
+    df.close = df.close.ffill()
+    close = df.close
+    df.vwap = df.vwap.fillna(close)
+    df.high = df.high.fillna(close)
+    df.low = df.low.fillna(close)
+    df.open = df.open.fillna(close)
+    price_open = df.open
+    high = df.high
+    low = df.low
+    vwap = df.vwap
+    # volume = df.volume
+    total = df.total
+    next_close = close.shift()
+
+    # TODO: Odometer rollover signal.  Relative difference to nearest power of 10.
+    # Something about log10 being a whole number?  When is $50 the rollover vs $100 or $10?
+
+    # "True (Typical?) Price" which I think is an approximation of VWAP.
+    # Trouble with both is that if there are no trades in a bar we get NaN.
+    # That then means we get NaN for averages for the next period-1 bars too.
+    # Question is whether to ffill the price for these calculations.
+    df["TP"] = (high + low + close) / 3
+
+    # Gain/loss in this bar.
+    df["ret1bar"] = close.div(price_open).sub(1)
+
+    for t in range(2, period):
+        df[f"ret{t}bar"] = close.div(price_open.shift(t - 1)).sub(1)
+
+    # Average True Range (ATR)
+    true_range = pd.concat(
+        [high.sub(low), high.sub(next_close).abs(), low.sub(next_close).abs()], axis=1
+    ).max(1)
+    # Normalized ATR (NATR) or Average of Normalized TR.
+    # Choice of NATR operations ordering discussion: https://www.macroption.com/normalized-atr/
+    # He doesn't talk about VWAP but I think that is a better normalizing price for a bar.
+    # atr = true_range.ewm(span=period).mean()
+    # df["natr_c"] = atr / close
+    # df["antr_c"] = (true_range / close).ewm(span=period).mean()
+    # df["natr_v"] = atr / vwap
+    # df["antr_v"] = (true_range / vwap).ewm(span=period).mean()
+    df["NATR"] = (true_range / vwap).ewm(span=period).mean()
+
+    # True Price as HLC average VS VWAP.
+    # VWAP is better I think but is quite different than standard CCI.
+    # Three ways to compute CCI, all give the same value using TP.
+    # tp = (high + low + close) / 3
+    # df['SMA'] = ta.sma(tp, length=period)
+    # df['sma_r'] = tp.rolling(period).mean()
+    # df['MAD'] = ta.mad(tp, length=period)
+    # # Series.mad deprecated. mad = (s - s.mean()).abs().mean()
+    # df['mad_r'] = tp.rolling(period).apply(lambda x: (pd.Series(x) - pd.Series(x).mean()).abs().mean())
+
+    # df['cci_r'] = (tp - df['sma_r']) / (0.015 * df['mad_r'])
+    # df['CCI'] = (tp - df['SMA']) / (0.015 * df['MAD'])
+    # df['cci_ta'] = ta.cci(high=high, low=low, close=close, length=period)
+
+    df["taCCI"] = ta.cci(high=high, low=low, close=close, length=period)
+
+    # https://gist.github.com/quantra-go-algo/1b37bfb74d69148f0dfbdb5a2c7bdb25
+    # https://medium.com/@huzaifazahoor654/how-to-calculate-cci-in-python-a-step-by-step-guide-9a3f61698be6
+    sma = pd.Series(ta.sma(vwap, length=period))
+    mad = pd.Series(ta.mad(vwap, length=period))
+    df["CCI"] = (vwap - sma) / (0.015 * mad)
+
+    # df['MFI'] = calculate_mfi(high=high, low=low, close=close, volume=volume, period=period)
+    df["MFI"] = calculate_mfi(typical_price=vwap, money_flow=total, period=period)
+
+    # We use Stochastic (rather than MACD because we need a ticker independent indicator.
+    # IOW a percentage price oscillator (PPO) rather than absolute price oscillator (APO).
+    # https://www.alpharithms.com/moving-average-convergence-divergence-macd-031217/
+    # We're using 14/3 currently rather than the usual 26/12 popular for MACD though.
+    stoch_k, stoch_d, stoch_h = calculate_stoch(high, low, close, k=period)
+    df["STOCHk"] = stoch_k
+    df["STOCHd"] = stoch_d
+    df["STOCHh"] = stoch_h
+
+    return df
+
+
+def iterate_all_aggs_tables(
+    config: PolygonConfig,
+    valid_tickers: pa.Array,
+):
+    schedule = config.calendar.trading_index(
+        start=config.start_timestamp, end=config.end_timestamp, period="1D"
+    )
+    for timestamp in schedule:
+        date = timestamp.to_pydatetime().date()
+        aggs_ds = pa_ds.dataset(
+            config.aggs_dir,
+            format="parquet",
+            schema=custom_aggs_schema(tz=config.calendar.tz.key),
+            partitioning=custom_aggs_partitioning(),
+        )
+        date_filter_expr = (
+            (pa_compute.field("year") == date.year)
+            & (pa_compute.field("month") == date.month)
+            & (pa_compute.field("date") == date)
+        )
+        # print(f"{date_filter_expr=}")
+        for fragment in aggs_ds.get_fragments(filter=date_filter_expr):
+            session_filter = (
+                (pa_compute.field("window_start") >= start_dt)
+                & (pa_compute.field("window_start") < end_dt)
+                & pa_compute.is_in(pa_compute.field("ticker"), valid_tickers)
+            )
+            # Sorting table doesn't seem to avoid needing to sort the df.  Maybe use_threads=False on to_pandas would help?
+            # table = fragment.to_table(filter=session_filter).sort_by([('ticker', 'ascending'), ('window_start', 'descending')])
+            table = fragment.to_table(filter=session_filter)
+            if table.num_rows > 0:
+                metadata = (
+                    dict(table.schema.metadata) if table.schema.metadata else dict()
+                )
+                metadata["date"] = date.isoformat()
+                table = table.replace_schema_metadata(metadata)
+                yield table
+
+
+# def iterate_all_aggs_with_signals(config: PolygonConfig):
+#     for table in iterate_all_aggs_tables(config):
+#         df = table.to_pandas()
+#         df = df.groupby("ticker").apply(
+#             compute_per_ticker_signals, include_groups=False
+#         )
+#         yield pa.Table.from_pandas(df)
+
+
+def file_visitor(written_file):
+    print(f"{written_file.path=}")
+
+
+def compute_signals_for_all_aggs(
+    from_config: PolygonConfig,
+    to_config: PolygonConfig,
+    valid_tickers: pa.Array,
+    overwrite: bool = False,
+) -> str:
+    if overwrite:
+        print("WARNING: overwrite not implemented/ignored.")
+
+    # Need a different aggs_dir for the signals because schema is different.
+    print(f"{to_config.aggs_dir=}")
+
+    for aggs_table in iterate_all_aggs_tables(from_config, valid_tickers):
+        metadata = aggs_table.schema.metadata
+        date = datetime.date.fromisoformat(metadata[b"date"].decode("utf-8"))
+        print(f"{date=}")
+        df = aggs_table.to_pandas()
+        df = df.groupby("ticker").apply(
+            compute_per_ticker_signals, include_groups=False
+        )
+        table = pa.Table.from_pandas(df)
+        if table.num_rows > 0:
+            table = table.replace_schema_metadata(metadata)
+            table = table.append_column("date", pa.array(np.full(len(table), date)))
+            table = table.append_column(
+                "year", pa.array(np.full(len(table), date.year), type=pa.uint16())
+            )
+            table = table.append_column(
+                "month", pa.array(np.full(len(table), date.month), type=pa.uint8())
+            )
+            table = table.sort_by(
+                [("ticker", "ascending"), ("window_start", "ascending")]
+            )
+            pa_ds.write_dataset(
+                table,
+                filesystem=to_config.filesystem,
+                base_dir=to_config.aggs_dir,
+                partitioning=custom_aggs_partitioning(),
+                format="parquet",
+                existing_data_behavior="overwrite_or_ignore",
+                file_visitor=file_visitor,
+            )
+    return to_config.aggs_dir

zipline_polygon_bundle/concat_all_aggs.py

@@ -1,44 +1,43 @@
-from .config import PolygonConfig
+from .config import PolygonConfig, PARTITION_COLUMN_NAME, to_partition_key
 
 import shutil
-from typing import Iterator, Tuple, List, Union
+from typing import Iterator, Tuple, Union
 
 import argparse
-import glob
 import os
+import datetime
 
 import pyarrow as pa
-from pyarrow import dataset as pa_ds
-from pyarrow import csv as pa_csv
-from pyarrow import compute as pa_compute
+import pyarrow.compute as pa_compute
+import pyarrow.csv as pa_csv
+import pyarrow.dataset as pa_ds
+import pyarrow.fs as pa_fs
 
 import pandas as pd
 
 
-PARTITION_COLUMN_NAME = "part"
-PARTITION_KEY_LENGTH = 2
-
-
-def to_partition_key(s: str) -> str:
-    """
-    Partition key is low cardinality and must be filesystem-safe.
-    The reason for partitioning is to keep the parquet files from getting too big.
-    10 years of minute aggs for US stocks is 83GB gzipped.  A single parquet would be 62GB on disk.
-    Currently the first two characters so files stay under 1GB.  Weird characters are replaced with "A".
-    """
-    k = (s + "A")[0:PARTITION_KEY_LENGTH].upper()
-    if k.isalpha():
-        return k
-    # Replace non-alpha characters with "A".
-    k = "".join([c if c.isalpha() else "A" for c in k])
-    return k
+# def get_by_ticker_dates(config: PolygonConfig, schema) -> set[datetime.date]:
+#     file_info = config.filesystem.get_file_info(config.by_ticker_dir)
+#     if file_info.type == pa_fs.FileType.NotFound:
+#         return set()
+#     partitioning = None
+#     if PARTITION_COLUMN_NAME in schema.names:
+#         partitioning = pa_ds.partitioning(
+#             pa.schema([(PARTITION_COLUMN_NAME, pa.string())]), flavor="hive"
+#         )
+#     by_ticker_aggs_ds = pa_ds.dataset(config.by_ticker_aggs_arrow_dir, schema=schema, partitioning=partitioning)
+#     return set(
+#         [
+#             pa_ds.get_partition_keys(fragment.partition_expression).get("date")
+#             for fragment in by_ticker_aggs_ds.get_fragments()
+#         ]
+#     )
 
 
 def generate_tables_from_csv_files(
-    paths: Iterator[Union[str, os.PathLike]],
+    config: PolygonConfig,
     schema: pa.Schema,
-    start_timestamp: pd.Timestamp,
-    limit_timestamp: pd.Timestamp,
+    overwrite: bool = False,
 ) -> Iterator[pa.Table]:
     empty_table = schema.empty_table()
     # TODO: Find which column(s) need to be cast to int64 from the schema.
@@ -49,16 +48,35 @@ def generate_tables_from_csv_files(
     )
     csv_schema = empty_table.schema
 
+    existing_by_ticker_dates = set()
+    if not overwrite:
+        # print("Getting existing by_ticker_dates")
+        # existing_by_ticker_dates = get_by_ticker_dates(config, schema)
+        print(f"{len(existing_by_ticker_dates)=}")
+
+    schedule = config.calendar.trading_index(
+        start=config.start_timestamp, end=config.end_timestamp, period="1D"
+    )
+    start_timestamp = config.start_timestamp.tz_localize(config.calendar.tz.key)
+    limit_timestamp = (config.end_timestamp + pd.Timedelta(days=1)).tz_localize(
+        config.calendar.tz.key)
+    # print(f"{start_timestamp=} {limit_timestamp=} {config.calendar.tz=} {schedule[:2]=} {schedule[-2:]=}")
+
     tables_read_count = 0
     skipped_table_count = 0
-    for path in paths:
+    for timestamp in schedule:
+        date: datetime.date = timestamp.tz_localize(config.calendar.tz.key).to_pydatetime().date()
+        # print(f"{date=} {timestamp=}")
+        if date in existing_by_ticker_dates:
+            continue
+        csv_path = config.date_to_csv_file_path(date)
         convert_options = pa_csv.ConvertOptions(
             column_types=csv_schema,
             strings_can_be_null=False,
             quoted_strings_can_be_null=False,
         )
 
-        table = pa_csv.read_csv(path, convert_options=convert_options)
+        table = pa_csv.read_csv(csv_path, convert_options=convert_options)
         tables_read_count += 1
         table = table.set_column(
             table.column_names.index("window_start"),
@@ -96,18 +114,86 @@ def generate_tables_from_csv_files(
             skipped_table_count += 1
             continue
 
+        if PARTITION_COLUMN_NAME in schema.names:
+            print(f"{date=}")
         yield table
     print(f"{tables_read_count=} {skipped_table_count=}")
 
 
+# def generate_tables_from_csv_files(
+#     paths: Iterator[Union[str, os.PathLike]],
+#     schema: pa.Schema,
+#     start_timestamp: pd.Timestamp,
+#     limit_timestamp: pd.Timestamp,
+# ) -> Iterator[pa.Table]:
+#     empty_table = schema.empty_table()
+#     # TODO: Find which column(s) need to be cast to int64 from the schema.
+#     empty_table = empty_table.set_column(
+#         empty_table.column_names.index("window_start"),
+#         "window_start",
+#         empty_table.column("window_start").cast(pa.int64()),
+#     )
+#     csv_schema = empty_table.schema
+
+#     tables_read_count = 0
+#     skipped_table_count = 0
+#     for path in paths:
+#         convert_options = pa_csv.ConvertOptions(
+#             column_types=csv_schema,
+#             strings_can_be_null=False,
+#             quoted_strings_can_be_null=False,
+#         )
+
+#         table = pa_csv.read_csv(path, convert_options=convert_options)
+#         tables_read_count += 1
+#         table = table.set_column(
+#             table.column_names.index("window_start"),
+#             "window_start",
+#             table.column("window_start").cast(schema.field("window_start").type),
+#         )
+#         if PARTITION_COLUMN_NAME in schema.names:
+#             table = table.append_column(
+#                 PARTITION_COLUMN_NAME,
+#                 pa.array(
+#                     [
+#                         to_partition_key(ticker)
+#                         for ticker in table.column("ticker").to_pylist()
+#                     ]
+#                 ),
+#             )
+#         expr = (
+#             pa_compute.field("window_start")
+#             >= pa.scalar(start_timestamp, type=schema.field("window_start").type)
+#         ) & (
+#             pa_compute.field("window_start")
+#             < pa.scalar(
+#                 limit_timestamp,
+#                 type=schema.field("window_start").type,
+#             )
+#         )
+#         table = table.filter(expr)
+
+#         # TODO: Also check that these rows are within range for this file's date (not just the whole session).
+#         # And if we're doing that (figuring date for each file), we can just skip reading the file.
+#         # Might able to do a single comparison using compute.days_between.
+#         # https://arrow.apache.org/docs/python/generated/pyarrow.compute.days_between.html
+
+#         if table.num_rows == 0:
+#             skipped_table_count += 1
+#             continue
+
+#         yield table
+#     print(f"{tables_read_count=} {skipped_table_count=}")
+
+
 def generate_csv_agg_tables(
-    config: PolygonConfig,
+    config: PolygonConfig, overwrite: bool = False
 ) -> Tuple[pa.Schema, Iterator[pa.Table]]:
     """zipline does bundle ingestion one ticker at a time."""
 
     # Polygon Aggregate flatfile timestamps are in nanoseconds (like trades), not milliseconds as the docs say.
     # I make the timestamp timezone-aware because that's how Unix timestamps work and it may help avoid mistakes.
-    timestamp_type = pa.timestamp("ns", tz="UTC")
+    timestamp_type = pa.timestamp("ns", tz='UTC')
 
     # But we can't use the timestamp type in the schema here because it's not supported by the CSV reader.
     # So we'll use int64 and cast it after reading the CSV file.
@@ -141,14 +227,12 @@ def generate_csv_agg_tables(
             pa.field(PARTITION_COLUMN_NAME, pa.string(), nullable=False)
         )
 
-    # TODO: Use generator like os.walk for paths.
     return (
         polygon_aggs_schema,
         generate_tables_from_csv_files(
-            paths=config.csv_paths(),
+            config,
             schema=polygon_aggs_schema,
-            start_timestamp=config.start_timestamp,
-            limit_timestamp=config.end_timestamp + pd.to_timedelta(1, unit="day"),
+            overwrite=overwrite,
         ),
     )
 
@@ -163,16 +247,17 @@ def concat_all_aggs_from_csv(
     config: PolygonConfig,
     overwrite: bool = False,
 ) -> str:
-    schema, tables = generate_csv_agg_tables(config)
+    schema, tables = generate_csv_agg_tables(config, overwrite=overwrite)
 
     by_ticker_aggs_arrow_dir = config.by_ticker_aggs_arrow_dir
-    if os.path.exists(by_ticker_aggs_arrow_dir):
-        if overwrite:
-            print(f"Removing {by_ticker_aggs_arrow_dir=}")
-            shutil.rmtree(by_ticker_aggs_arrow_dir)
-        else:
-            print(f"Found existing {by_ticker_aggs_arrow_dir=}")
-            return by_ticker_aggs_arrow_dir
+    # if os.path.exists(by_ticker_aggs_arrow_dir):
+    #     if overwrite:
+    #         print(f"Removing {by_ticker_aggs_arrow_dir=}")
+    #         shutil.rmtree(by_ticker_aggs_arrow_dir)
+    #     else:
+    #         # TODO: Validate the existing data.
+    #         print(f"Found existing {by_ticker_aggs_arrow_dir=}")
+    #         return by_ticker_aggs_arrow_dir
 
     partitioning = None
     if PARTITION_COLUMN_NAME in schema.names:
@@ -180,7 +265,7 @@ def concat_all_aggs_from_csv(
             pa.schema([(PARTITION_COLUMN_NAME, pa.string())]), flavor="hive"
         )
 
-    # scanner = pa_ds.Scanner.from_batches(source=generate_batches_from_tables(tables), schema=schema)
+    print(f"Scattering aggregates by ticker to {by_ticker_aggs_arrow_dir=}")
     pa_ds.write_dataset(
         generate_batches_from_tables(tables),
         schema=schema,
@@ -189,7 +274,7 @@ def concat_all_aggs_from_csv(
         format="parquet",
         existing_data_behavior="overwrite_or_ignore",
     )
-    print(f"Concatenated aggregates to {by_ticker_aggs_arrow_dir=}")
+    print(f"Scattered aggregates by ticker to {by_ticker_aggs_arrow_dir=}")
     return by_ticker_aggs_arrow_dir