# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import pandas as pd
import numpy as np
import xgboost as xgb

# +
df_data = pd.read_csv('dix.csv', names=['Date','SP500','DIX','GEX'], index_col='Date', parse_dates=True, header=0)
df_vix = pd.read_csv('vixcurrent.csv', skiprows=1, index_col='Date', parse_dates=True)
df_vix = df_vix['VIX Close'].rename('VIX')

df_data = df_data.join(df_vix, how='left')

#df_data = df_data.tail(500)

def get_signal(x):
    ret = np.exp(x) - 1
    if ret > 0.008:
        return 1
    elif ret < -0.008:
        return -1
    else:
        return 0

def get_flip(row):
    a = row['GEX dir']
    b = row['GEX last dir']
    
    if a > 0 and b < 0:
        return 1
    elif a < 0 and b > 0:
        return -1
    else:
        return 0

def process_data(df):
    
    df['SP500 r'] = 100 * df['SP500'].pct_change()
    for col in ['DIX', 'GEX', 'VIX']:
        df[col + ' diff'] = df[col].diff()
    df['SP500 r next'] = df['SP500 r'].shift(periods=-1)
    #df['signal'] = df['SP500 lr next'].map(get_signal)
    df['GEX dir'] = df['GEX'].map(lambda x : -1 if x < 0 else 1)
    df['GEX last dir'] = df['GEX dir'].shift(periods=1)
    df['GEX flip'] = df[['GEX dir', 'GEX last dir']].apply(get_flip, axis=1)
    
    df.drop('SP500', axis=1, inplace=True)
    
    return df.iloc[1:]

df_data = process_data(df_data)
#df_data['signal'].value_counts()
df_data.head()

# +
features = [
    'DIX',
    'GEX',
    'GEX flip',
    'VIX',
    'DIX diff',
    'SP500 r'
]

target = 'SP500 r next'

prediction_size = 10
labelled_size = len(df_data) - prediction_size

X_to_predict = df_data.tail(prediction_size)[features]

df_labelled = df_data.head(labelled_size)

# +
model = xgb.XGBRegressor(objective='reg:squarederror', learning_rate=0.05, n_estimators=300)

X_train = df_labelled[features]
y_train = df_labelled[target]

model.fit(X_train, y_train)


# -

#df_data['Returns'] = np.exp(df_data['SP500 lr next']) - 1
df_data.tail(prediction_size)

pred = model.predict(X_to_predict)
growth = np.exp(pred) - 1
#X_to_predict['Projected lr'] = pred
X_to_predict['Projected Returns'] = pred
X_to_predict

StarcoderdataJupyterScriptsDedupFiltered