main.py

import optuna
from sklearn.metrics import precision_score


from data import StockData,NIFTY50_TICKERS
from features import add_features
from model import backtest,create_target,backtest_ensemble
import pandas as pd
sd=StockData()
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from xgboost import XGBClassifier
from lstm_model import predict_lstm,train_lstm
import evaluate

def run(ticker):
    stocks=sd.load_data()
    df=stocks[ticker]
    df=create_target(df)
    df=add_features(df)
    df.dropna(inplace=True)
    df = df.replace([float('inf'), float('-inf')], pd.NA)
    df = df.dropna()
    df = df.astype(float, errors='ignore')
    predictors=df.drop(['Target','Tomorrow','Close','Open','Volume','Low','High','Stock Splits','Dividends'],axis=1).columns.tolist()
    study=optuna.create_study(direction="maximize")
    def objective_xgboost(trial):
        n_estimators=trial.suggest_int("n_estimators",100,500)
        max_depth=trial.suggest_int("max_depth",3,20)
        subsample=trial.suggest_float("subsample",0.5,1)
        colsample_bytree=trial.suggest_float("colsample_bytree",0.5,1)
        learning_rate=trial.suggest_float("learning_rate",0.01,0.3)
        model=XGBClassifier(n_estimators=n_estimators,max_depth=max_depth,learning_rate=learning_rate,colsample_bytree=colsample_bytree,subsample=subsample)
        predictions=backtest(df,model,predictors)
        score=precision_score(predictions['Target'],predictions['Predictions'])
        return score
    study.optimize(objective_xgboost,n_trials=5)

    study_1 = optuna.create_study(direction="maximize")

    def objective_LSTM(trial):
        hidden_size=trial.suggest_int("hidden_size",32,128)
        num_layers=trial.suggest_int("num_layers",1,3)
        lr=trial.suggest_float("lr",0.001,0.1)
        epoch=trial.suggest_int("epoch",20,100)
        lstm=train_lstm(df,predictors,hidden_size=hidden_size,num_layers=num_layers,epoch=epoch,lr=lr)

        test=df.iloc[int(len(df)*0.8):]
        probs=[]
        start = int(len(df) * 0.8)
        for i in range(len(test)):
            prob=predict_lstm(lstm,df.iloc[:start+i+60],predictors)
            probs.append(prob)
        preds=(np.array(probs)>=0.5).astype(int)
        score=precision_score(test["Target"].values,preds)
        return score
    study_1.optimize(objective_LSTM, n_trials=3)


    study1=optuna.create_study(direction="maximize")
    def objective_random_forest(trial):

        n_estimators=trial.suggest_int("n_estimators",100,500)
        min_sample_split=trial.suggest_int("min_samples_split",10,100)
        max_depth=trial.suggest_int("max_depth",3,20)

        model=RandomForestClassifier(n_estimators=n_estimators,random_state=1,min_samples_split=min_sample_split,max_depth=max_depth)
        predictions=backtest(df,model,predictors)
        score=precision_score(predictions["Target"],predictions['Predictions'])
        return score
    study1.optimize(objective_random_forest,n_trials=5)



    model=XGBClassifier(**study.best_params,random_state=1)
    model1=RandomForestClassifier(**study1.best_params,random_state=1)
    model_1=train_lstm(df,predictors,**study_1.best_params)
    preds=backtest_ensemble(df,model,model1,model_1,predictors)
    from StockDB import StockDB
    db = StockDB()
    results = evaluate.evaluate(preds, df)
    db.save_results(ticker, results["precision"], results["sharpe"], results["max_dd"])
    db.save_predictions(ticker,preds)

    return preds,df,predictors,model,model1,model_1

if __name__=="__main__":
    for ticker in NIFTY50_TICKERS:
        print(f"Training {ticker}...")
        try:
            run(ticker)
            print(f"Done")
        except Exception as e:
            print(f"Failed {ticker}:{e}")
            continue