solution.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "import warnings\n",
    "import importlib\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import seaborn as sns\n",
    "from tqdm.notebook import tqdm\n",
    "from IPython import display\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder, OneHotEncoder\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.gaussian_process import GaussianProcessRegressor\n",
    "from sklearn.gaussian_process.kernels import DotProduct, WhiteKernel, Matern\n",
    "from sklearn.gaussian_process.kernels import ConstantKernel as C\n",
    "from sklearn.experimental import enable_iterative_imputer\n",
    "from sklearn.impute import IterativeImputer\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from scipy.interpolate import UnivariateSpline\n",
    "from sklearn.metrics import mean_squared_error\n",
    "from sklearn.model_selection import KFold\n",
    "\n",
    "import lightgbm as lgb\n",
    "import scipy.sparse\n",
    "from cylp.cy import CyClpSimplex\n",
    "from cylp.py.modeling.CyLPModel import CyLPArray\n",
    "\n",
    "import kmeans\n",
    "\n",
    "sns.set_style('darkgrid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_preparation():\n",
    "    \"\"\"\n",
    "    Чтение данных\n",
    "    \"\"\"\n",
    "    train_path = 'data/train.csv'\n",
    "    test_path = 'data/test.csv'\n",
    "\n",
    "    train = pd.read_csv(train_path)\n",
    "    test = pd.read_csv(test_path)\n",
    "\n",
    "    train = train[train['galaxy'] != 'NGC 5253']\n",
    "    train = train.reset_index(drop=True)\n",
    "\n",
    "    test_sorted = test.sort_values('galactic year')\n",
    "    test_sorted_ = test_sorted[test_sorted.columns[2:-1]]\n",
    "    years = np.unique(np.array(test_sorted['galactic year']))\n",
    "    year_fin = years[-1]\n",
    "    year_prev = years[-2]\n",
    "\n",
    "    test['galactic year'] = test['galactic year'].replace(year_fin, year_prev)\n",
    "\n",
    "    train_ = train.copy(deep=True)\n",
    "    train_ = train_.drop('y', axis=1)\n",
    "    train_ = pd.concat((train_, test))\n",
    "\n",
    "    enc = LabelEncoder()\n",
    "    train_['galaxy'] = enc.fit_transform(train_['galaxy'])\n",
    "    train['galaxy'] = enc.transform(train['galaxy'])\n",
    "    test['galaxy'] = enc.transform(test['galaxy'])\n",
    "\n",
    "    gxs = np.array(train['galaxy'].drop_duplicates())\n",
    "\n",
    "    train_ = train.copy(deep=True)\n",
    "    train_ = train_.drop('y', axis=1)\n",
    "    train_ = pd.concat((train_, test))\n",
    "\n",
    "    gxs = np.array(train['galaxy'].drop_duplicates())\n",
    "\n",
    "    enc_h = OneHotEncoder()\n",
    "    train_mat = enc_h.fit_transform(train[['galaxy']]).toarray()\n",
    "    test_mat = enc_h.transform(test[['galaxy']]).toarray()\n",
    "\n",
    "    train_mat = pd.DataFrame(train_mat)\n",
    "    test_mat = pd.DataFrame(test_mat)\n",
    "    train = pd.concat((train, train_mat), axis=1)\n",
    "    test = pd.concat((test, test_mat), axis=1)\n",
    "    \n",
    "    return train, test, train_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Регрессия на кластерах + оптимизация"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "importlib.reload(kmeans)\n",
    "\n",
    "def clustering(train, N_CLUSTERS, random_state):\n",
    "    \"\"\"\n",
    "    Функция для кластеризации временных рядов\n",
    "    \"\"\"\n",
    "    df_to_clust = train.pivot(index='galactic year',\n",
    "                              columns='galaxy', values='y').transpose()\n",
    "    df_to_clust_val = df_to_clust.values\n",
    "    df_to_clust[df_to_clust.columns] = df_to_clust_val\n",
    "    df_clust = df_to_clust.fillna(-9999)\n",
    "    array_to_clust = df_clust.to_numpy()\n",
    "    mask = array_to_clust != -9999\n",
    "    mask = mask.astype(int)\n",
    "\n",
    "    selection_rule = mask.sum(axis=1) > 5\n",
    "    mask = mask[selection_rule]\n",
    "    array_to_clust = array_to_clust[selection_rule]\n",
    "\n",
    "    km = kmeans.KMeans(array_to_clust[:, 15:], mask[:, 15:],\n",
    "                       N_CLUSTERS, resolve_empty='singleton')\n",
    "    km.initialise(seed=random_state)\n",
    "    km.cluster()\n",
    "    res = km.clustering_results\n",
    "    res = np.where(res == 1)[1]\n",
    "    df_to_clust = df_to_clust[selection_rule]\n",
    "\n",
    "    galaxy_clusters = dict()\n",
    "    galaxies = list(df_to_clust.index)\n",
    "    for galaxy, cluster in zip(galaxies, res):\n",
    "        galaxy_clusters[galaxy] = cluster\n",
    "    return galaxy_clusters, df_to_clust, res, array_to_clust\n",
    "\n",
    "\n",
    "def smoothing(df_to_clust, res, plot=False, mean=False):\n",
    "    \"\"\"\n",
    "    Функция для сглаживания сплайнами\n",
    "    \"\"\"\n",
    "    \n",
    "    est_list = dict()\n",
    "    for i in np.unique(res):\n",
    "        output_plot = df_to_clust[res==i].index\n",
    "        output = df_to_clust[res==i].transpose().to_numpy()\n",
    "        pred = np.nanmean(output, axis=1)\n",
    "        pred_arr = np.nanmean(output[:5], axis=0)\n",
    "        if mean:\n",
    "            output /= pred_arr\n",
    "        df1 = pd.DataFrame(pred)\n",
    "        df2 = pd.DataFrame(pred)\n",
    "        df1.fillna(method='ffill', inplace=True)\n",
    "        df2.fillna(method='bfill', inplace=True)\n",
    "        arr1 = df1[0].to_numpy()\n",
    "        arr2 = df2[0].to_numpy()\n",
    "        pred = (arr1 + arr2) / 2\n",
    "        spl = UnivariateSpline(np.arange(len(pred)), pred, k=5, s=10)\n",
    "        spl.set_smoothing_factor(0.1)\n",
    "        est = spl(np.arange(len(pred)))\n",
    "        est_list[i] = est\n",
    "        if plot:\n",
    "            fig, ax = plt.subplots(1, 2, figsize=(20,7))\n",
    "            plt.title(f'cluster: {i}')\n",
    "            for y_arr, label in zip(output.T, list(output_plot)):\n",
    "                ax[0].plot(y_arr, label=label, marker='o')\n",
    "            ax[0].legend()\n",
    "            ax[1].plot(pred)\n",
    "            ax[1].plot(est)\n",
    "            plt.show()\n",
    "    return est_list\n",
    "\n",
    "\n",
    "def reg_lin(train_, test, array_to_clust, df_to_clust,\n",
    "            est_list, galaxy_clusters, plot=False):\n",
    "    \"\"\"\n",
    "    Регрессия внутри кластеров\n",
    "    \"\"\"\n",
    "    array_to_clust[array_to_clust < -100] = np.nan\n",
    "\n",
    "    years = train_['galactic year'].drop_duplicates().to_numpy()\n",
    "    preds_dict = dict()\n",
    "\n",
    "    for num, row in enumerate(array_to_clust):\n",
    "        galaxy = df_to_clust.iloc[num].name\n",
    "        cluster = galaxy_clusters[galaxy]\n",
    "        est = est_list[cluster]\n",
    "        row = row.reshape(-1, 1)\n",
    "        est = est.reshape(-1, 1)\n",
    "        df = np.hstack((est, row))\n",
    "        df = pd.DataFrame(df)\n",
    "        df.columns = ['x', 'y']\n",
    "        df_train = df.dropna()\n",
    "        train_X = df_train[['x']]\n",
    "        train_y = df_train['y']\n",
    "        reg = GaussianProcessRegressor(kernel=DotProduct() + WhiteKernel())\n",
    "        reg.fit(train_X, train_y)\n",
    "        preds = reg.predict(df[['x']])\n",
    "        preds = pd.DataFrame(preds).set_index(years)\n",
    "        preds_dict[galaxy] = dict(zip(preds.index, preds.values))\n",
    "        if plot:\n",
    "            plt.figure(figsize=(16,4))\n",
    "            plt.title(f'cluster: {cluster}')\n",
    "            plt.plot(preds.to_numpy())\n",
    "            plt.plot(df['y'].to_numpy())\n",
    "            plt.show()\n",
    "            plt.close()\n",
    "\n",
    "    preds_list = list()\n",
    "\n",
    "    for row_num in range(len(test)):\n",
    "        row = test.iloc[row_num]\n",
    "        galaxy = row['galaxy']\n",
    "        year = row['galactic year']\n",
    "        prediction = preds_dict[galaxy][year][0]\n",
    "        preds_list.append(prediction)\n",
    "    return preds_list, preds_dict\n",
    "\n",
    "\n",
    "def regression(train, regressor, preds_dict, info=False):\n",
    "    \"\"\"\n",
    "    Регрессия - один раунд кросс-валидации\n",
    "    train содержит y\n",
    "    \"\"\"\n",
    "    preds_list_train = list()\n",
    "\n",
    "    for row_num in range(len(train)):\n",
    "        row = train.iloc[row_num]\n",
    "        galaxy = row['galaxy']\n",
    "        year = row['galactic year']\n",
    "        prediction = preds_dict[galaxy][year][0]\n",
    "        preds_list_train.append(prediction)\n",
    "\n",
    "    train_extended = train.copy(deep=True)\n",
    "    train_extended['spline'] = preds_list_train\n",
    "\n",
    "    train_extended = train_extended.fillna(-9999)\n",
    "\n",
    "    X = train_extended.drop('y', axis=1)\n",
    "    y = train_extended['y']\n",
    "\n",
    "    reg = regressor\n",
    "    reg.fit(X, y)\n",
    "    preds = reg.predict(X)\n",
    "\n",
    "    if info:\n",
    "        plt.figure(figsize=(16,4))\n",
    "        plt.plot(preds_list_train)\n",
    "        plt.plot(train['y'].to_numpy())\n",
    "        print(mean_squared_error(preds_list_train, train['y'].to_numpy()) ** 0.5)\n",
    "        print(mean_squared_error(preds, train['y'].to_numpy()) ** 0.5)\n",
    "    return regressor\n",
    "\n",
    "\n",
    "def preds_lists(X_train, X_test, train_, N_CLUSTERS=20, random_state=7):\n",
    "    \"\"\"\n",
    "    Функция для получения предикшенов по данным кластеризации\n",
    "    \"\"\"\n",
    "    galaxy_clusters, df_to_clust, res, array_to_clust = clustering(X_train,\n",
    "                                                                   N_CLUSTERS, random_state)\n",
    "\n",
    "    est_list = smoothing(df_to_clust, res, plot=False, mean=False)\n",
    "    preds_list, preds_dict = reg_lin(train_, X_test, array_to_clust,\n",
    "                                     df_to_clust, est_list, galaxy_clusters, plot=False)\n",
    "    train_preds_list, train_preds_dict = reg_lin(train_, X_train, array_to_clust,\n",
    "                                     df_to_clust, est_list, galaxy_clusters, plot=False)\n",
    "    test_preds_list, train_preds_dict = reg_lin(train_, test, array_to_clust,\n",
    "                                     df_to_clust, est_list, galaxy_clusters, plot=False)\n",
    "    return preds_list, train_preds_list, test_preds_list\n",
    "\n",
    "\n",
    "def cluster_res(train, test, train_, random_state):\n",
    "    \"\"\"\n",
    "    Обобщение результатов кластеризации\n",
    "    \"\"\"\n",
    "    train_base = train[train['galactic year'] < test['galactic year'].min()]\n",
    "    train_to_split = train[train['galactic year'] >= test['galactic year'].min()]\n",
    "    X_train, X_test = train_test_split(train_to_split, test_size=0.00004, random_state=6, shuffle=True)\n",
    "    X_train = pd.concat((train_base, X_train))\n",
    "\n",
    "    preds_clust, train_preds_clust, test_preds_clust = [], [], []\n",
    "    preds_list, train_preds_list, test_preds_list = preds_lists(X_train, X_test, train_,\n",
    "                                                                N_CLUSTERS=20, random_state=random_state)\n",
    "    preds_clust.append(preds_list)\n",
    "    train_preds_clust.append(train_preds_list)\n",
    "    test_preds_clust.append(test_preds_list)\n",
    "    return preds_clust, train_preds_clust, test_preds_clust, X_train, X_test\n",
    "\n",
    "\n",
    "def prediction(preds_clust, train_preds_clust, test_preds_clust, X_train, X_test, test):\n",
    "\n",
    "    test_for_pred = test.copy(deep=True)\n",
    "    spline_test_pred = pd.DataFrame(test_preds_clust).transpose()\n",
    "    spline_test_pred.columns = [str(i) + '_spline' for i in spline_test_pred.columns]\n",
    "    test_for_pred = test_for_pred.join(spline_test_pred)\n",
    "    test_for_pred_ = test_for_pred.fillna(0)\n",
    "\n",
    "    y_test = X_test['y'].copy(deep=True)\n",
    "    test_X = X_test.reset_index(drop=True).copy(deep=True)\n",
    "    spline_test = pd.DataFrame(preds_clust).transpose()\n",
    "    spline_test.columns = [str(i) + '_spline' for i in spline_test.columns]\n",
    "    test_X = test_X.drop('y', axis=1)\n",
    "    test_X = test_X.fillna(-10).reset_index(drop=True)\n",
    "    test_X = test_X.join(spline_test)\n",
    "\n",
    "    train_X = X_train.copy(deep=True)\n",
    "    spline_train = pd.DataFrame(train_preds_clust).transpose()\n",
    "    spline_train.columns = [str(i) + '_spline' for i in spline_train.columns]\n",
    "    train_y = train_X['y']\n",
    "    train_X = train_X.drop('y', axis=1)\n",
    "    train_X = train_X.fillna(0).reset_index(drop=True)\n",
    "    train_X = train_X.join(spline_train)\n",
    "    test_y = y_test\n",
    "\n",
    "    sc = StandardScaler()\n",
    "    train_X = sc.fit_transform(train_X)\n",
    "    test_X = sc.transform(test_X)\n",
    "    test_for_pred = sc.transform(test_for_pred_)\n",
    "\n",
    "    regressor_boost = lgb.LGBMRegressor(num_leaves=20,\n",
    "                                        learning_rate=0.05,\n",
    "                                        n_estimators=200000,\n",
    "                                        verbose=0)\n",
    "    regressor_boost.fit(train_X, train_y, eval_set=[(test_X, test_y)],\n",
    "                        eval_metric='l2', early_stopping_rounds=10, verbose=0)\n",
    "\n",
    "    preds_boost = regressor_boost.predict(test_X)\n",
    "    preds_on_test_boost_mult = regressor_boost.predict(test_for_pred)\n",
    "    \n",
    "    return preds_on_test_boost_mult"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize(preds, test, upper_bound_=None, lower_bound_=None):\n",
    "    \"\"\"\n",
    "    Функция для выполнения оптимизации\n",
    "    \"\"\"\n",
    "\n",
    "    preds_list = preds\n",
    "\n",
    "    constr_matrix_1 = test['existence expectancy index'] < 0.7\n",
    "    constr_matrix_1 = constr_matrix_1.to_numpy().astype(int).reshape(1, -1) * (-1)\n",
    "    constr_matrix_1 = np.vstack((\n",
    "        constr_matrix_1, np.ones(constr_matrix_1.size).reshape(1, -1)))\n",
    "    constr_matrix_1 = scipy.sparse.csr_matrix(constr_matrix_1)\n",
    "    constr_vector_1 = np.array([-5000, 50000])\n",
    "    constr_matrix_2 = scipy.sparse.eye(test.shape[0])\n",
    "    constr_vector_2 = np.ones(test.shape[0]) * 100\n",
    "    lower_bound = np.zeros(test.shape[0])\n",
    "\n",
    "    output = preds_list\n",
    "\n",
    "    solver = CyClpSimplex()\n",
    "    x_var = solver.addVariable('x', len(preds_list))\n",
    "\n",
    "    vec1 = CyLPArray(constr_vector_1)\n",
    "    solver.addConstraint(constr_matrix_1 * x_var <= vec1)\n",
    "    vec2 = CyLPArray(constr_vector_2)\n",
    "    solver.addConstraint(constr_matrix_2 * x_var <= vec2)\n",
    "    if lower_bound_ is not None:\n",
    "        solver.variablesLower = lower_bound_\n",
    "    else:\n",
    "        solver.variablesLower = lower_bound\n",
    "    if upper_bound_ is not None:\n",
    "        solver.variablesUpper = upper_bound_\n",
    "    objective = CyLPArray(output - 1)\n",
    "    solver.objective = objective * x_var\n",
    "    solver.primal()\n",
    "    out = solver.primalVariableSolution['x']\n",
    "\n",
    "    std = np.ones(preds_list.size) * 0.0093\n",
    "\n",
    "    output_matrix = np.repeat([preds_list], repeats=10000, axis=0)\n",
    "    for i in range(output_matrix.shape[1]):\n",
    "        np.random.seed(i)\n",
    "        vec = np.random.normal(loc=0, scale=std[i], size=output_matrix.shape[0])\n",
    "        output_matrix[:, i] = output_matrix[:, i] + vec\n",
    "\n",
    "    optimized_matrix = np.zeros(output_matrix.shape)\n",
    "\n",
    "    for row in tqdm(range(len(output_matrix))):\n",
    "        objective = CyLPArray(output_matrix[row] - 1)\n",
    "        solver.objective = objective * x_var\n",
    "        solver.primal()\n",
    "        optimized_matrix[row] = solver.primalVariableSolution['x']\n",
    "    opt_pred = optimized_matrix.mean(axis=0)\n",
    "    \n",
    "    return opt_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d957cfbb85a94e3388534aa36e7a03fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=10000.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "train, test, train_ = data_preparation()\n",
    "\n",
    "preds_gauss = list()\n",
    "preds_boost = list()\n",
    "\n",
    "for random_state in [7, 14]:\n",
    "\n",
    "    preds_clust, train_preds_clust, test_preds_clust, X_train, X_test =\\\n",
    "        cluster_res(train, test, train_, random_state)\n",
    "\n",
    "    preds_on_test_boost_mult =\\\n",
    "        prediction(preds_clust, train_preds_clust, test_preds_clust, X_train, X_test, test)\n",
    "    \n",
    "    preds_boost.append(preds_on_test_boost_mult)\n",
    "\n",
    "preds_boost_first = np.array(preds_boost).mean(axis=0)\n",
    "\n",
    "display.clear_output()\n",
    "\n",
    "opt_pred = optimize(preds_boost_first, test)\n",
    "opt_pred[0] -= 0.0001"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Вторая регрессия"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reg2(train, test, train_):\n",
    "    train_base = train[train['galactic year'] < test['galactic year'].min()]\n",
    "    train_to_split = train[train['galactic year'] >= test['galactic year'].min()]\n",
    "    X_train, X_test = train_test_split(train_to_split, test_size=0.0004, random_state=6, shuffle=True)\n",
    "    X_train = pd.concat((train_base, X_train))\n",
    "\n",
    "    random_states = [7, 14]\n",
    "\n",
    "    preds_clust, train_preds_clust, test_preds_clust = [], [], []\n",
    "    for random_state in random_states:\n",
    "        preds_list, train_preds_list, test_preds_list = preds_lists(X_train, X_test, train_,\n",
    "                                                                    N_CLUSTERS=25, random_state=random_state)\n",
    "        preds_clust.append(preds_list)\n",
    "        train_preds_clust.append(train_preds_list)\n",
    "        test_preds_clust.append(test_preds_list)\n",
    "    return preds_clust, train_preds_clust, test_preds_clust, X_train, X_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_features(train, test, train_):\n",
    "\n",
    "    # Делаем новые фичи\n",
    "    cols_to_use = train_.columns[2:]\n",
    "    new_feature_matrix = train_.copy(deep=True)\n",
    "    new_feature_matrix = new_feature_matrix.set_index('galaxy', drop=True)\n",
    "\n",
    "    new_features = list()\n",
    "\n",
    "    for column in cols_to_use:\n",
    "        df = new_feature_matrix[[column\n",
    "            ]].groupby(new_feature_matrix.index).agg([np.nanmean, np.nanmax, np.nanmin, np.nanstd])\n",
    "        df.columns = [i + '_' + column for i in ['nanmean', 'nanmax', 'nanmin', 'nanstd']]\n",
    "        new_features.append(deepcopy(df))\n",
    "    new_features = pd.concat(new_features, axis=1)\n",
    "\n",
    "    # Делаем агрегацию для таргета и вычисляем разницу первых и последних значений\n",
    "    df_to_clust = train.pivot(index='galactic year',\n",
    "                             columns='galaxy', values='y')\n",
    "    df_to_clust = df_to_clust.to_numpy()\n",
    "    diff = np.nanmean(df_to_clust[:10], axis=0) - np.nanmean(df_to_clust[-10:], axis=0)\n",
    "\n",
    "    new_features['diff'] = diff\n",
    "    new_features_ = new_features.dropna().copy(deep=True)\n",
    "    transformations = dict()\n",
    "\n",
    "    counter = 0\n",
    "    for feature in new_features_.columns:\n",
    "        cc = np.corrcoef(new_features_[feature], new_features_['diff'])\n",
    "        cc_sqr = np.corrcoef(new_features_[feature] ** 2, new_features_['diff'])\n",
    "        cc_inv = np.corrcoef(1 / new_features_[feature], new_features_['diff'])\n",
    "        cc_ = np.max([np.abs(cc[0][1]), np.abs(cc_sqr[0][1]), np.abs(cc_inv[0][1])])\n",
    "        tp = np.argmax([np.abs(cc[0][1]), np.abs(cc_sqr[0][1]), np.abs(cc_inv[0][1])])\n",
    "        transformation = [lambda x: x, lambda x: np.square(x), lambda x: 1 / x, lambda x: 1 / np.square(x)][tp]\n",
    "        transf_type = ['None', 'Square', 'Inversion', 'Quad inversion'][tp]\n",
    "        if cc_ > 0.6 and feature != 'diff':\n",
    "            transformations[feature] = transformation\n",
    "\n",
    "    new_features = new_features[transformations.keys()]\n",
    "    for key, value in transformations.items():\n",
    "        new_features[key] = new_features[key].apply(value)\n",
    "\n",
    "\n",
    "    train = train.join(new_features, on='galaxy')\n",
    "    test = test.join(new_features, on='galaxy')\n",
    "    \n",
    "    return train, test\n",
    "\n",
    "\n",
    "def fillna_y_train(train):\n",
    "    \"\"\"\n",
    "    Новый признак для обучения\n",
    "    \"\"\"\n",
    "    y = pd.pivot_table(train, values=['y'], index='galaxy', columns='galactic year').transpose()\n",
    "    y_prev = y.to_numpy()\n",
    "    y_transformed = y_prev.copy()\n",
    "\n",
    "    kf = KFold(n_splits=7, shuffle=True, random_state=10)\n",
    "\n",
    "    nan_idxs_list = list()\n",
    "\n",
    "    idxs_list = np.indices(y_prev.shape)\n",
    "    idxs_list = idxs_list.reshape(-1, np.product(y_prev.shape)).T\n",
    "\n",
    "    splits = list(kf.split(idxs_list))\n",
    "\n",
    "    for split in tqdm(splits):\n",
    "        y_to_mask = deepcopy(y_prev)\n",
    "        nan_idxs = idxs_list[split[1]]\n",
    "        y_to_mask[nan_idxs[:, 0], nan_idxs[:, 1]] = np.nan\n",
    "        imp = IterativeImputer(missing_values=np.nan, max_iter=40)\n",
    "        y_out = imp.fit_transform(y_to_mask)\n",
    "        y_transformed[nan_idxs[:, 0], nan_idxs[:, 1]] = y_out[nan_idxs[:, 0], nan_idxs[:, 1]].copy()\n",
    "\n",
    "    y_out = pd.DataFrame(y_transformed, columns=y.columns.astype(str))\n",
    "    y_out['galactic year'] = y.index.get_level_values(1)\n",
    "    y_out = pd.melt(y_out, id_vars=['galactic year'], value_vars=list(y_out.columns[:-1]))\n",
    "    y_out['galaxy'] = y_out['galaxy'].astype(int)\n",
    "    y_out_masked = y_out.set_index(['galaxy','galactic year']).copy(deep=True)\n",
    "    \n",
    "    return y_out_masked\n",
    "\n",
    "\n",
    "def fillna_y_final(train):\n",
    "    \"\"\"\n",
    "    Новый признак для предсказания\n",
    "    \"\"\"\n",
    "    y = pd.pivot_table(train, values=['y'], index='galaxy', columns='galactic year').transpose()\n",
    "    y_prev = y.to_numpy()\n",
    "    estimator = GaussianProcessRegressor(kernel=DotProduct() + WhiteKernel())\n",
    "    imp = IterativeImputer(estimator=estimator, missing_values=np.nan, max_iter=40)\n",
    "    y_out = imp.fit_transform(y_prev)\n",
    "    y_out = pd.DataFrame(y_out, columns=y.columns.astype(str))\n",
    "    y_out['galactic year'] = y.index.get_level_values(1)\n",
    "    y_out = pd.melt(y_out, id_vars=['galactic year'], value_vars=list(y_out.columns[:-1]))\n",
    "    y_out['galaxy'] = y_out['galaxy'].astype(int)\n",
    "    y_out_final = y_out.set_index(['galaxy','galactic year']).copy(deep=True)\n",
    "    \n",
    "    return y_out_final\n",
    "\n",
    "\n",
    "def regression_second(train, test, y_out_masked, y_out_final):\n",
    "\n",
    "    X_train, X_test = train_test_split(train, test_size=0.00003, random_state=7, shuffle=True)\n",
    "\n",
    "    test_for_pred = test.copy(deep=True)\n",
    "    test_for_pred_ = test_for_pred.fillna(0)\n",
    "    test_for_pred_ = test_for_pred_.join(y_out_final, on=['galaxy','galactic year'], how='left')\n",
    "\n",
    "    preds_rnd = deepcopy(y_out_masked)\n",
    "    n_out = preds_rnd.shape[0]\n",
    "    np.random.seed(1)\n",
    "    rnd_vec = np.random.normal(loc=0, scale=0.000002, size=n_out)\n",
    "\n",
    "    y_test = X_test['y'].copy(deep=True)\n",
    "    test_X = X_test.reset_index(drop=True).copy(deep=True)\n",
    "    test_X = test_X.drop('y', axis=1)\n",
    "    test_X = test_X.fillna(0).reset_index(drop=True)\n",
    "    test_X = test_X.join(y_out_masked, on=['galaxy','galactic year'], how='left')\n",
    "\n",
    "    train_X = X_train.copy(deep=True)\n",
    "    train_y = train_X['y']\n",
    "    train_X = train_X.drop('y', axis=1)\n",
    "    train_X = train_X.fillna(0).reset_index(drop=True)\n",
    "    train_X = train_X.join(preds_rnd, on=['galaxy','galactic year'], how='left')\n",
    "    test_y = y_test\n",
    "\n",
    "    sc = StandardScaler()\n",
    "    train_X = sc.fit_transform(train_X)\n",
    "    test_X = sc.transform(test_X)\n",
    "    test_for_pred = sc.transform(test_for_pred_)\n",
    "\n",
    "    regressor_gp = GaussianProcessRegressor(kernel=Matern() * C() \n",
    "                                            + WhiteKernel()\n",
    "                                            + DotProduct() * C())\n",
    "    regressor_gp.fit(train_X, train_y)\n",
    "    preds_gp = regressor_gp.predict(test_X)\n",
    "\n",
    "    regressor_boost = lgb.LGBMRegressor(num_leaves=27,\n",
    "                                        learning_rate=0.05,\n",
    "                                        n_estimators=5000,\n",
    "                                        verbose=0)\n",
    "    regressor_boost.fit(train_X, train_y,\n",
    "                        eval_set=[(test_X, test_y)],\n",
    "                        eval_metric='l2',\n",
    "                        early_stopping_rounds=10,\n",
    "                        verbose=0)\n",
    "\n",
    "    preds_boost = regressor_boost.predict(test_for_pred)\n",
    "    preds_gp = regressor_gp.predict(test_for_pred)\n",
    "    \n",
    "    return preds_boost, preds_gp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, test, train_ = data_preparation()\n",
    "\n",
    "preds_clust, train_preds_clust, test_preds_clust, X_train, X_test = reg2(train, test, train_)\n",
    "preds_boost_first =\\\n",
    "        prediction(preds_clust, train_preds_clust, test_preds_clust, X_train, X_test, test)\n",
    "\n",
    "train, test, train_ = data_preparation()\n",
    "train, test = add_features(train, test, train_)\n",
    "y_out_masked = fillna_y_train(train)\n",
    "y_out_final = fillna_y_final(train)\n",
    "\n",
    "preds_boost_second, preds_gp_second = regression_second(train, test, y_out_masked, y_out_final)\n",
    "\n",
    "pred = 0.29 * preds_boost_first + 0.29 * preds_gp_second + 0.42 * preds_boost_second\n",
    "\n",
    "display.clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = pd.read_csv('data/sample_submit.csv')\n",
    "output['pred'] = pred\n",
    "output['opt_pred'] = opt_pred\n",
    "output = output.set_index('index')\n",
    "output.to_csv(f'output.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}