batch_processor.py

import os
import numpy as np
import pandas as pd
import databento as db
from scipy.optimize import minimize
from scipy.special import ndtr
from datetime import datetime, timedelta
from dotenv import load_dotenv

# --- Configuration ---
START_DATE = "2024-01-08" 
END_DATE = "2024-02-08"    # Start with a 1-month sample
OUTPUT_DIR = "dataset_v1"
INTERVAL_MINUTES = 30      # Sample frequency
STRICT_IV_CLIP = 0.8       # Clip IV at 80% (Removes artifacts)

load_dotenv()
DATABENTO_KEY = os.getenv("DATABENTO_API_KEY")

if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

# --- 1. Math Helpers ---
def vectorized_implied_volatility(prices, S, K, T, r, flags):
    sigma = np.full_like(prices, 0.5)
    for i in range(15):
        d1 = (np.log(S / K) + (r + 0.5 * sigma**2) * T) / (sigma * np.sqrt(T))
        d2 = d1 - sigma * np.sqrt(T)
        bs_price = flags * S * ndtr(flags * d1) - flags * K * np.exp(-r * T) * ndtr(flags * d2)
        vega = S * np.sqrt(T) * (np.exp(-0.5 * d1**2) / np.sqrt(2 * np.pi))
        diff = prices - bs_price
        vega = np.where(vega < 1e-8, 1e-8, vega)
        step = diff / vega
        sigma += step
        if np.max(np.abs(step)) < 1e-4: break
    return sigma

def raw_svi(k, params):
    a, b, rho, m, sigma = params
    return a + b * (rho * (k - m) + np.sqrt((k - m)**2 + sigma**2))

def fit_svi_slice(k, w_obs):
    def obj(params):
        a, b, rho, m, sigma = params
        # Penalties for Arbitrage
        if b < 0 or abs(rho) >= 1 or sigma <= 0: return 1e9
        if a + b * sigma * np.sqrt(1 - rho**2) < 0: return 1e9
        return np.sum((raw_svi(k, params) - w_obs)**2)
    
    # Run optimizer with multiple starting guesses to avoid local minima
    guesses = [
        [0.04, 0.1, -0.5, 0.0, 0.1],
        [0.09, 0.3, -0.2, 0.0, 0.2]
    ]
    best_res = None
    best_err = float('inf')
    
    for g in guesses:
        res = minimize(obj, g, method='Nelder-Mead', tol=1e-3)
        if res.fun < best_err:
            best_err = res.fun
            best_res = res
            
    return best_res.x

# --- 2. Processing Logic ---
def process_day(client, day_str):
    print(f"--> Processing {day_str}...")
    try:
        # 1. Fetch Data (OHLCV fallback included implicitly by using ohlcv schema)
        start = f"{day_str}T14:30:00"
        end = f"{day_str}T21:00:00"
        
        defs = client.timeseries.get_range(
            dataset="OPRA.PILLAR", schema="definition", stype_in="parent",
            symbols=["SPY.OPT"], start=f"{day_str}T00:00:00", end=end
        ).to_df()
        
        # Filter Options
        defs = defs[defs['instrument_class'].isin(['C', 'P'])]
        defs['put_call'] = defs['instrument_class']
        meta = defs.set_index('instrument_id')[['strike_price', 'expiration', 'put_call']]
        
        candles = client.timeseries.get_range(
            dataset="OPRA.PILLAR", schema="ohlcv-1m", stype_in="parent",
            symbols=["SPY.OPT"], start=start, end=end, limit=500000
        ).to_df()
        
        # 2. Merge
        df = candles.join(meta, on='instrument_id', how='inner')
        df['mid_price'] = df['close'] # Use close for batch stability
        df['expiration'] = pd.to_datetime(df['expiration'])
        df['T'] = (df['expiration'] - df.index) / pd.Timedelta(days=365.25)
        
        # FILTER: Strict Time and Price filter
        df = df[(df['T'] > 0.04) & (df['mid_price'] > 0.05)]
        
        if len(df) < 500: return [], []

        # 3. Calc IV
        flags = np.where(df['put_call'] == 'C', 1, -1)
        # Approximate spot from median strike of active options
        est_spot = df['strike_price'].median()
        
        df['iv'] = vectorized_implied_volatility(
            df['mid_price'].values, est_spot, df['strike_price'].values, 
            df['T'].values, 0.05, flags
        )
        # Filter implausible IVs before fitting
        df = df[(df['iv'] > 0.01) & (df['iv'] < 1.5)]
        df['k'] = np.log(df['strike_price'] / est_spot)

        # 4. Rasterize Loop
        frames = []
        timestamps = []
        times = df.index.floor(f'{INTERVAL_MINUTES}min').unique()
        
        grid_k = np.linspace(-0.5, 0.5, 64)
        grid_T = np.linspace(0.04, 1.0, 64) 
        
        for t_stamp in times:
            slice_df = df[df.index.floor(f'{INTERVAL_MINUTES}min') == t_stamp]
            if len(slice_df) < 50: continue
            
            fitted_params = {}
            for exp in slice_df['expiration'].unique():
                sub = slice_df[slice_df['expiration'] == exp]
                if len(sub) < 5: continue
                T_val = sub['T'].iloc[0]
                w_obs = (sub['iv']**2) * T_val
                fitted_params[T_val] = fit_svi_slice(sub['k'].values, w_obs.values)
            
            if len(fitted_params) < 3: continue
            
            # Interpolate
            tensor = np.zeros((64, 64))
            sorted_T = sorted(fitted_params.keys())
            
            for j, t_target in enumerate(grid_T):
                nearest_t = min(sorted_T, key=lambda x: abs(x - t_target))
                params = fitted_params[nearest_t]
                w_grid = raw_svi(grid_k, params)
                
                # Convert to IV and Clip
                iv_slice = np.sqrt(np.maximum(w_grid, 0) / t_target)
                iv_slice = np.clip(iv_slice, 0, STRICT_IV_CLIP) # <--- THE FIX
                tensor[:, j] = iv_slice
            
            frames.append(tensor)
            timestamps.append(t_stamp)
            
        return frames, timestamps

    except Exception as e:
        print(f"Skipping {day_str}: {e}")
        return [], []

# --- 3. Execution ---
if __name__ == "__main__":
    client = db.Historical(DATABENTO_KEY)
    dates = pd.date_range(START_DATE, END_DATE, freq='B')
    
    all_frames = []
    
    print(f"Starting Batch Process: {START_DATE} to {END_DATE}")
    for d in dates:
        frames, ts = process_day(client, d.strftime("%Y-%m-%d"))
        if frames:
            all_frames.extend(frames)
            print(f"   Saved {len(frames)} frames")
            
    if all_frames:
        dataset = np.array(all_frames, dtype=np.float32)
        save_path = f"{OUTPUT_DIR}/spy_iv_surface_dataset.npy"
        np.save(save_path, dataset)
        print(f"DONE. Dataset saved to: {save_path}")
        print(f"Shape: {dataset.shape}")
    else:
        print("No data generated. Check date range or API limits.")