Add adjust_ohlcv.py

adjust_ohlcv.py

@@ -1,7 +1,7 @@
 """Split/dividend adjustment for a user's raw 1-minute OHLCV bars.
 
-Reproduces the convention the model was trained under (see
-data_pull/data/adjustments.py and preprocessing/quality.py in the source repo):
+Reproduces the convention the model was trained under (ported from the authors' training
+pipeline):
 
 * factors are cumulative and backward-looking, anchored at 1.0 on the latest bar date;
 * a split with ex-date D multiplies price by denominator/numerator and volume by
@@ -9,7 +9,8 @@ data_pull/data/adjustments.py and preprocessing/quality.py in the source repo):
 * a cash dividend with ex-date D multiplies price by (close_prior - amount) / close_prior
   on every bar dated strictly before D, where close_prior is the regular-session close on
   the last bar date before D; volume is unchanged;
-* dividends with amount >= close_prior are skipped (price left under-adjusted).
+* dividends with amount >= close_prior are skipped (price left under-adjusted);
+* two actions on the same ex-date are combined: amounts summed, split ratios multiplied.
 
 Only pandas and numpy are required.
 """
@@ -29,32 +30,43 @@ def clean_bars(bars: pd.DataFrame) -> pd.DataFrame:
     """Apply the training pipeline's raw-bar quality checks.
 
     Drops duplicate timestamps (keeps first), bars with high < low, bars whose close lies
-    outside [low, high], and bars with a non-positive price; negative volume is set to 0.
+    outside [low, high], and bars with a non-positive or missing price; negative or
+    missing volume is set to 0.
     """
     df = bars.drop_duplicates(subset="timestamp", keep="first")
     ok = (df["high"] >= df["low"]) & (df["close"] <= df["high"]) & (df["close"] >= df["low"])
     ok &= (df[["open", "high", "low", "close"]] > 0).all(axis=1)
     df = df[ok.fillna(False)].copy()
-    df["volume"] = df["volume"].clip(lower=0)
+    df["volume"] = df["volume"].fillna(0).clip(lower=0)
     return df.sort_values("timestamp").reset_index(drop=True)
 
 
 def _eastern(ts: pd.Series) -> pd.Series:
-    ts = pd.to_datetime(ts)
+    """Parse if needed, then localize naive as Eastern or convert aware to Eastern."""
+    if not pd.api.types.is_datetime64_any_dtype(ts):
+        try:
+            ts = pd.to_datetime(ts)
+        except ValueError:  # strings with mixed UTC offsets, e.g. across a DST change
+            ts = pd.to_datetime(ts, utc=True)
     return ts.dt.tz_localize(EASTERN) if ts.dt.tz is None else ts.dt.tz_convert(EASTERN)
 
 
 def _split_events(splits: pd.DataFrame | None) -> dict[dt.date, tuple[float, float]]:
-    """{ex_date: (numerator, denominator)}; accepts `ratio` (4.0 for 4:1, 0.5 for 1:2)."""
+    """{ex_date: (numerator, denominator)}; accepts `ratio` (10.0 = 10-for-1, 0.5 = 1-for-2).
+
+    Same-date splits are multiplied together.
+    """
     if splits is None or len(splits) == 0:
         return {}
     s = splits.copy()
     if "numerator" not in s:
         s["numerator"], s["denominator"] = s["ratio"].astype(float), 1.0
-    return {
-        pd.Timestamp(d).date(): (float(n), float(m))
-        for d, n, m in zip(s["date"], s["numerator"], s["denominator"], strict=True)
-    }
+    events: dict[dt.date, tuple[float, float]] = {}
+    for d, n, m in zip(s["date"], s["numerator"], s["denominator"], strict=True):
+        ex_date = pd.Timestamp(d).date()
+        n0, m0 = events.get(ex_date, (1.0, 1.0))
+        events[ex_date] = (n0 * float(n), m0 * float(m))
+    return events
 
 
 def _close_prior(df: pd.DataFrame, ex_date: dt.date) -> float | None:
@@ -67,14 +79,20 @@ def _close_prior(df: pd.DataFrame, ex_date: dt.date) -> float | None:
 
 
 def _dividend_events(df: pd.DataFrame, dividends: pd.DataFrame | None) -> dict[dt.date, tuple[float, float]]:
-    """{ex_date: (amount, close_prior)}; close_prior is taken from the bars unless given."""
+    """{ex_date: (amount, close_prior)}; close_prior is taken from the bars unless given.
+
+    Same-date dividends are summed against one close_prior.
+    """
     if dividends is None or len(dividends) == 0:
         return {}
-    events = {}
+    events: dict[dt.date, tuple[float, float]] = {}
     given = dividends["close_prior"] if "close_prior" in dividends else [None] * len(dividends)
     for d, amount, cp in zip(dividends["date"], dividends["amount"], given, strict=True):
         ex_date = pd.Timestamp(d).date()
-        close_prior = cp if cp is not None and not pd.isna(cp) else _close_prior(df, ex_date)
+        if ex_date in events:
+            events[ex_date] = (events[ex_date][0] + float(amount), events[ex_date][1])
+            continue
+        close_prior = cp if not pd.isna(cp) else _close_prior(df, ex_date)
         if close_prior is not None and close_prior > 0:
             events[ex_date] = (float(amount), float(close_prior))
     return events
@@ -120,10 +138,12 @@ def adjust_ohlcv(
     """Return cleaned bars with adjusted prices ready for `make_features`.
 
     Args:
-        bars: one security; columns `timestamp` (tz-aware, or naive = Eastern), `open`,
-            `high`, `low`, `close`, `volume`, all unadjusted.
-        splits: columns `date` (ex-date) and either `ratio` or `numerator`+`denominator`.
-        dividends: columns `date` (ex-date), `amount`, optional `close_prior`.
+        bars: one security; columns `timestamp` (bar start; tz-aware, or naive taken as
+            Eastern), `open`, `high`, `low`, `close`, `volume`, all unadjusted.
+        splits: columns `date` (ex-date) and either `ratio` or `numerator`+`denominator`
+            (numerator = shares after the split).
+        dividends: columns `date` (ex-date), `amount` (cash per share as declared),
+            optional `close_prior`.
 
     Returns:
         DataFrame with `timestamp_eastern`, `open_adj`, `high_adj`, `low_adj`, `close_adj`,
@@ -150,4 +170,4 @@ def adjust_ohlcv(
         "adj_factor_price",
         "adj_factor_volume",
     ]
-    return df[keep].reset_index(drop=True)
+    return df[keep]