{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":146992,"sourceType":"modelInstanceVersion","modelInstanceId":124705,"modelId":147730}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Import necessary libraries\nimport numpy as np\nimport polars as pl\nimport pandas as pd\nimport lightgbm as lgb\nimport xgboost as xgb\nimport os\nimport joblib\nfrom sklearn.model_selection import GridSearchCV\nfrom sklearn.metrics import mean_squared_error\n\n# Set up constants\nTARGET = 'responder_6'\nFEAT_COLS = [f\"feature_{i:02d}\" for i in range(79)]\n\n# Function to load data with optional filtering\ndef load_data(date_id_range=None, time_id_range=None, columns=None, return_type='pl'):\n    data_dir = '../input/jane-street-real-time-market-data-forecasting'\n    data = pl.scan_parquet(f\"{data_dir}/train.parquet\")\n\n    if date_id_range is not None:\n        start_date, end_date = date_id_range\n        data = data.filter((pl.col(\"date_id\") >= start_date) & (pl.col(\"date_id\") <= end_date))\n\n    if time_id_range is not None:\n        start_time, end_time = time_id_range\n        data = data.filter((pl.col(\"time_id\") >= start_time) & (pl.col(\"time_id\") <= end_time))\n\n    if columns is not None:\n        data = data.select(columns)\n\n    if return_type == 'pd':\n        return data.collect().to_pandas()\n    else:\n        return data.collect()\n\n# Function to add rolling features\ndef add_rolling_features(data):\n    for feature in FEAT_COLS:\n        data[f'rolling_mean_{feature}'] = data[feature].rolling(window=5).mean()\n        data[f'rolling_std_{feature}'] = data[feature].rolling(window=5).std()\n    return data\n\n# Function to calculate R² score\ndef calculate_r2(y_true, y_pred, weights):\n    numerator = np.sum(weights * (y_true - y_pred) ** 2)\n    denominator = np.sum(weights * (y_true ** 2))\n    r2_score = 1 - (numerator / denominator)\n    return r2_score\n\n# Function to evaluate the model\ndef evaluate_model(model, test_data):\n    y_pred = model.predict(test_data[FEAT_COLS])\n    y_true = test_data[TARGET].to_numpy() \n    weights = test_data['weight'].to_numpy()  \n    r2_score = calculate_r2(y_true, y_pred, weights)\n    print(f\"Sample weighted zero-mean R-squared score (R2) on test data: {r2_score}\")\n\n# Class to manage a group of models\nclass ModelGroup:\n    def __init__(self):\n        self.models = []\n\n    def add_model(self, model):\n        self.models.append(model)\n\n    def predict(self, test_data):\n        preds = []\n        for model in self.models:\n            if isinstance(model, lgb.Booster):\n                pred = model.predict(test_data[FEAT_COLS])\n            elif isinstance(model, xgb.Booster):\n                pred = model.predict(xgb.DMatrix(test_data[FEAT_COLS]))\n            elif hasattr(model, 'predict'):\n                pred = model.predict(test_data[FEAT_COLS])\n            else:\n                raise ValueError(\"Unsupported model type\")\n            preds.append(pred)\n\n        avg_pred = np.mean(preds, axis=0)\n        return avg_pred\n    \n    @classmethod\n    def load(cls, file_path):\n        model_group = joblib.load(file_path)\n        return model_group\n\n# Function to train XGBoost with K-Folds\ndef train_xgb_kfold(total_days=1498, n_splits=5, save_models=False):\n    model_group = ModelGroup()\n    fold_size = total_days // n_splits\n    folds = [(i * fold_size, min((i + 1) * fold_size - 1, total_days - 1)) for i in range(n_splits)]\n\n    for fold_idx in range(n_splits):\n        valid_range = folds[fold_idx]\n        train_ranges = [folds[i] for i in range(n_splits) if i != fold_idx]\n\n        print(f\"Fold {fold_idx}: validation range {valid_range}, train parts: {train_ranges}\")\n\n        valid_data = load_data(date_id_range=valid_range, columns=[\"date_id\", \"weight\"] + FEAT_COLS + [TARGET], return_type='pl')\n        valid_data = add_rolling_features(valid_data)\n        valid_weight = valid_data['weight'].to_pandas()\n\n        train_data = None\n        for train_range in train_ranges:\n            partial_train_data = load_data(date_id_range=train_range, columns=[\"date_id\", \"weight\"] + FEAT_COLS + [TARGET], return_type='pl')\n            partial_train_data = add_rolling_features(partial_train_data)\n            if train_data is None:\n                train_data = partial_train_data\n            else:\n                train_data = train_data.vstack(partial_train_data)\n\n        train_weight = train_data['weight'].to_pandas()\n\n        dtrain = xgb.DMatrix(train_data.select(FEAT_COLS).to_pandas(), label=train_data[TARGET].to_pandas(), weight=train_weight)\n        dvalid = xgb.DMatrix(valid_data.select(FEAT_COLS).to_pandas(), label=valid_data[TARGET].to_pandas(), weight=valid_weight)\n\n        # Set up hyperparameter tuning\n        xgb_model = xgb.XGBRegressor(\n            eval_metric='rmse',\n            learning_rate=0.1,\n            max_depth=6,\n            min_child_weight=1,\n            subsample=0.8,\n            colsample_bytree=0.8,\n            random_state=42,\n            tree_method='gpu_hist',\n        )\n\n        # Use GridSearchCV for hyperparameter tuning\n        param_grid = {\n            'n_estimators': [100, 200, 500],\n            'learning_rate': [0.01, 0.1, 0.2],\n            'max_depth': [4, 6, 8],\n            'min_child_weight': [1, 2],\n            'subsample': [0.7, 0.8],\n            'colsample_bytree': [0.7, 0.8]\n        }\n\n        grid_search = GridSearchCV(xgb_model, param_grid, scoring='neg_mean_squared_error', cv=3)\n        grid_search.fit(train_data[FEAT_COLS].to_pandas(), train_data[TARGET].to_pandas())\n\n        best_model = grid_search.best_estimator_\n        print(\"Best parameters found: \", grid_search.best_params_)\n\n        y_valid_pred = best_model.predict(valid_data[FEAT_COLS].to_pandas())\n        r2_score = calculate_r2(valid_data[TARGET].to_pandas(), y_valid_pred, valid_weight)\n        print(f\"Fold {fold_idx} validation R2 score: {r2_score}\")\n\n        model_group.add_model(best_model)\n\n    if save_models:\n        joblib.dump(model_group, \"xgb_model_group.pkl\")\n        print(\"Saved the model group to xgb_model_group.pkl\")\n    \n    return model_group\n\n# Uncomment to train a new model\n# total_days = 1699\n# xgb_models = train_xgb_kfold(total_days=total_days, n_splits=5, save_models=True)\n\n# Load pre-trained model group\nxgb_models = ModelGroup.load(\"/kaggle/input/model_an/other/default/1/xgb_model_group.pkl\")\n\n# Prediction function for the inference server\ndef predict(test: pl.DataFrame, lags: pl.DataFrame | None) -> pl.DataFrame | pd.DataFrame:\n    feat = test[FEAT_COLS].to_pandas()\n    pred = xgb_models.predict(feat)\n\n    predictions = test.select('row_id').with_columns(pl.Series('responder_6', pred.ravel()))\n\n    assert isinstance(predictions, (pl.DataFrame, pd.DataFrame))\n    assert list(predictions.columns) == ['row_id', 'responder_6']\n    assert len(predictions) == len(test)\n\n    return predictions\n\n# Set up the inference server\nimport kaggle_evaluation.jane_street_inference_server\ninference_server = kaggle_evaluation.jane_street_inference_server.JSInferenceServer(predict)\n\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    inference_server.serve()\nelse:\n    inference_server.run_local_gateway(\n        (\n            '../input/jane-street-real-time-market-data-forecasting/test.parquet',\n            '../input/jane-street-real-time-market-data-forecasting/lags.parquet',\n        )\n    )\n\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-10-26T13:04:35.403377Z","iopub.execute_input":"2024-10-26T13:04:35.403722Z","iopub.status.idle":"2024-10-26T13:04:42.513342Z","shell.execute_reply.started":"2024-10-26T13:04:35.403685Z","shell.execute_reply":"2024-10-26T13:04:42.512271Z"},"trusted":true},"execution_count":null,"outputs":[]}]}