VA_Final.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RK92jpzh1CjC"
      },
      "source": [
        "# Manual inputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UVFYc9r8tbGa"
      },
      "source": [
        "This script should be executed at the start of the project. It requires access to Google Drive for data retrieval, so ensure that your Google account is connected. Additionally, the script reinstalls a necessary package, which requires manual input during execution. Please follow the prompts carefully to complete the installation process before proceeding with further code execution."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c87b9hWW1Fft"
      },
      "source": [
        "## Prelims"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The results presented in this paper rely on the following package versions:\n",
        "\n",
        "JAX: 0.4.31; TensorFlow: 2.17.0; Neural Tangents: 0.6.6."
      ],
      "metadata": {
        "id": "DmVN-GFk23-u"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "61nHgfr1E7kp"
      },
      "outputs": [],
      "source": [
        "!pip install -q --upgrade pip\n",
        "!pip install --upgrade jax==0.4.31 jaxlib==0.4.31\n",
        "!pip install tensorflow==2.17.0\n",
        "!pip install -q git+https://www.github.com/google/neural-tangents"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import tensorflow as tf\n",
        "import neural_tangents as nt\n",
        "\n",
        "print(\"JAX version:\", jax.__version__) # 0.4.31\n",
        "print(\"TensorFlow version:\", tf.__version__) # 2.17.0\n",
        "print(\"Neural Tangents version:\", nt.__version__) #0.6.6"
      ],
      "metadata": {
        "id": "6wu9_aixy0ve"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Connect to Google Drive to access the required files. If you plan to save the file locally, modifications to the code will be necessary."
      ],
      "metadata": {
        "id": "pyAlgr_Z3zgW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MMJ_m4CO2SB9",
        "outputId": "36ac4554-2159-4211-b59b-1d005bd7200d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Prep"
      ],
      "metadata": {
        "id": "iFh8p6KV6jJV"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eJWzXizHtl8k"
      },
      "source": [
        "This code loads the data and performs minor pre-processing. Before running the code, ensure you have the following two files: inforce.csv; fmv_seriatim.csv. These can be downloaded at https://www2.math.uconn.edu/~gan/software.html, under *Aggregate results*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-GTb0wxh4bWG"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab Notebooks/VA/\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "\n",
        "datpolicies = pd.read_csv('inforce.csv')\n",
        "datvals = pd.read_csv('fmv_seriatim.csv')\n",
        "datcombined = datpolicies.join(datvals.set_index('recordID'), on = 'recordID')\n",
        "#use matDate, birthDate, issueTime and currentDate to define\n",
        "datcombined['matTime'] = (datcombined.matDate - datcombined.currentDate)/365.25 #time to maturity\n",
        "datcombined['Age'] =   (datcombined.currentDate - datcombined.birthDate)/365.25 #age\n",
        "datgreeks = pd.read_csv('Greek.csv')\n",
        "#datcombined['issueTime'] =(datcombined.currentDate - datcombined.issueDate)/365.25 #time since issue\n",
        "mydat = datcombined #190000 policies\n",
        "mydat_y = mydat[\"base:base\"].astype('float') #contract valuation: base (dependent)\n",
        "mydat_y = np.array(mydat_y).flatten()\n",
        "\n",
        "mydat = mydat.iloc[:, np.r_[2, 3, 11:13, 15:25, 72,73]] #gender, productType, gbAmt, GMWBBalance, FundValue1-10, 2 derived variables\n",
        "#all products: include withdrawal variables\n",
        "mydat['gender'] = np.where(mydat['gender']=='M', 1, 0)\n",
        "mave = np.absolute(mydat_y).mean()\n",
        "datcombined2 = datcombined.join(datgreeks.set_index('recordID'), on = 'recordID')\n",
        "\n",
        "#for Greeks\n",
        "datgreeks2 = datcombined2.iloc[:, -13:]\n",
        "list1 = ['base'] + [str(x)+'_'+str(y) for x in range(1, 6) for y in ['D', 'U']]+['base']*16\n",
        "list2 = ['base']*11 + [str(x)+'y_'+str(y) for x in [1, 2, 3, 4, 5, 7, 10, 30] for y in ['D', 'U']]\n",
        "colnames = [str(x)+':'+str(y) for (x, y) in zip(list2, list1)]\n",
        "datvals2 = datcombined2[colnames]\n",
        "\n",
        "#for numeric variables, perform minmax scaling\n",
        "for i in range(2, len(mydat.columns)):\n",
        "  currentcolumn = mydat.iloc[:, i]\n",
        "  mydat.iloc[:, i] = (currentcolumn - currentcolumn.min())/(currentcolumn.max()-currentcolumn.min())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dpj5Pn6IVKnD"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from random import seed, sample\n",
        "from keras import backend as K\n",
        "from keras.callbacks import EarlyStopping\n",
        "from keras.models import Model as KerasModel\n",
        "from keras.models import Sequential\n",
        "from keras.layers import Input, Dense, Activation, Reshape\n",
        "from keras.layers import Concatenate, Softmax\n",
        "#from keras.layers.embeddings import Embedding\n",
        "from keras.utils import to_categorical, set_random_seed\n",
        "from keras.optimizers import Adam, SGD\n",
        "from keras.constraints import NonNeg\n",
        "from keras import initializers\n",
        "\n",
        "#create one hot encoding of categorical\n",
        "from keras.utils import to_categorical\n",
        "from copy import deepcopy\n",
        "\n",
        "\n",
        "productTypes = mydat['productType'].unique()\n",
        "for i in range(len(productTypes)):\n",
        "  mydat.loc[mydat['productType'] == productTypes[i], 'productType'] = i\n",
        "mydat_emb = deepcopy(mydat)\n",
        "#mydat_emb will be used for embedding encoding for NNs\n",
        "productType_one_hot = to_categorical(mydat['productType'])\n",
        "productType_one_hot = pd.DataFrame(productType_one_hot, columns = productTypes)\n",
        "mydat = mydat.loc[:, mydat.columns != 'productType']\n",
        "mydat = mydat.join(productType_one_hot)\n",
        "#mydat will be used for one-hot encoding for NNs\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XQ9gTfzn3M3f"
      },
      "source": [
        "# Sampling methods"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nIwmNJFy3P4w"
      },
      "source": [
        "At the beginning of each run, select only one of the following methods: stratified random sampling, hierarchical k-means, or conditional cLHS. Ensure that the random state is set based on the run number for reproducibility."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "mydat_array = mydat.values #certain methods require data to be stored in an array rather than a data frame\n",
        "random_state = 10\n",
        "n = 680\n",
        "ndat = len(mydat_y) #190,000\n",
        "\n",
        "random.seed(1)"
      ],
      "metadata": {
        "id": "rzmJfB99z2El"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bwIKXq2GuJNK"
      },
      "source": [
        "## Stratified random sampling"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The results presented in the paper are the average of replicates generated using the following five seeds: 10, 20, 30, 40, and 50. For single replicates, the results are based on seed 10."
      ],
      "metadata": {
        "id": "W7ERM2i36oHe"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qa_NdxR-qeCj"
      },
      "outputs": [],
      "source": [
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "\n",
        "mydat_train, mydat_val, mydat_train_emb, mydat_val_emb, mydat_train_y, mydat_val_y  = train_test_split(mydat_array, mydat_emb, mydat_y, test_size=(n/2)/ndat,train_size=n/ndat,\n",
        "                                                                                                       stratify = mydat_emb['productType'], random_state =random_state)\n",
        "def makenpfloat(vec):\n",
        "  return np.array(vec, dtype = np.float32)\n",
        "\n",
        "mydat_train = makenpfloat(mydat_train)\n",
        "mydat_val = makenpfloat(mydat_val)\n",
        "mydat_train_y = makenpfloat(mydat_train_y).flatten()\n",
        "mydat_val_y = makenpfloat(mydat_val_y).flatten()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following code is intended exclusively for Greeks."
      ],
      "metadata": {
        "id": "nYwhqTuu7Ajp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3twpC9cjHSOu"
      },
      "outputs": [],
      "source": [
        "mydat_greeks_train = datgreeks2.loc[mydat_train_emb.index]\n",
        "mydat_greeks_val = datgreeks2.loc[mydat_val_emb.index]\n",
        "mydat_vals_train = datvals2.loc[mydat_train_emb.index]\n",
        "mydat_vals_val = datvals2.loc[mydat_val_emb.index]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8JSaYLTPKgq0"
      },
      "source": [
        "## Hierarchical k-means\n",
        "\n",
        "This code implements hierarchical k-means as described in the Gan and Valdez book. In broader machine learning contexts, this method is commonly referred to as bisecting k-means. The results presented in the paper are the average of replicates generated using the following five seeds: 10, 20, 30, 40, and 50."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1zYktCwKL_0o"
      },
      "outputs": [],
      "source": [
        "from sklearn.cluster import BisectingKMeans\n",
        "from scipy.cluster.vq import vq\n",
        "mydat_array =mydat.values\n",
        "\n",
        "\n",
        "mykmeans = BisectingKMeans(n_clusters = n, random_state = random_state, init = 'k-means++', n_init = 1, bisecting_strategy = 'largest_cluster').fit(mydat_array)\n",
        "\n",
        "\n",
        "def getclosestinds(dat, kmeans):\n",
        "  all_data = [ i for i in range(dat.shape[0]) ]\n",
        "  m_clusters = kmeans.labels_.tolist()\n",
        "  centers = np.array(kmeans.cluster_centers_)\n",
        "  n_clusters = centers.shape[0]\n",
        "\n",
        "  closest_data = []\n",
        "  for i in range(n_clusters):\n",
        "      center_vec = centers[i].reshape(1,-1)\n",
        "      data_idx_within_i_cluster = [ idx for idx, clu_num in enumerate(m_clusters) if clu_num == i ]\n",
        "\n",
        "      one_cluster_dat = np.zeros( (  len(data_idx_within_i_cluster) , centers.shape[1] ) )\n",
        "      for row_num, data_idx in enumerate(data_idx_within_i_cluster):\n",
        "          one_row = dat[data_idx]\n",
        "          one_cluster_dat[row_num] = one_row\n",
        "\n",
        "      closest, _ = vq(center_vec, one_cluster_dat)\n",
        "      closest_idx_in_one_cluster_dat = closest[0]\n",
        "      closest_data_row_num = data_idx_within_i_cluster[closest_idx_in_one_cluster_dat]\n",
        "      data_id = all_data[closest_data_row_num]\n",
        "\n",
        "      closest_data.append(data_id)\n",
        "\n",
        "  closest_data = list(set(closest_data))\n",
        "\n",
        "  assert len(closest_data) == n_clusters\n",
        "\n",
        "  return closest_data\n",
        "\n",
        "mydat_train_inds = getclosestinds(mydat_array, mykmeans)\n",
        "\n",
        "\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "mydat_train = mydat_array[mydat_train_inds]\n",
        "mydat_train_y = mydat_y[mydat_train_inds]\n",
        "mydat_train_emb = mydat_emb.iloc[mydat_train_inds]\n",
        "\n",
        "\n",
        "mydat_y_notrain = np.delete(mydat_y, mydat_train_inds)\n",
        "mydat_array_notrain = np.delete(mydat_array, mydat_train_inds, axis=0)\n",
        "mydat_emb_notrain = mydat_emb[~mydat_emb.index.isin(mydat_train_inds)]\n",
        "productType_notrain = mydat_emb_notrain['productType']\n",
        "mydat_val, _, mydat_val_emb, _, mydat_val_y, _ = train_test_split(mydat_array_notrain, mydat_emb_notrain, mydat_y_notrain, train_size=(n/2)/(ndat-n), stratify = productType_notrain, random_state =random_state)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0_DZtKBgA2cd"
      },
      "source": [
        "## Conditioned latin hypercube sampling (cLHS)\n",
        "\n",
        "This code performs cLHS. The results presented in the paper are the average of replicates generated using the following five seeds: 10, 20, 30, 40, and 50."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wg2_4xTwA0gS"
      },
      "outputs": [],
      "source": [
        "!pip install clhs\n",
        "import clhs as cl\n",
        "\n",
        "\n",
        "sampled=cl.clhs(mydat_emb, n, max_iterations=1000, random_state = random_state)\n",
        "# then, validation indices are sampled from outside of training indices, using\n",
        "# stratified random sampling\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "mydat_train_inds = sampled['sample_indices']\n",
        "mydat_train = mydat_array[mydat_train_inds]\n",
        "mydat_train_y = mydat_y[mydat_train_inds]\n",
        "mydat_train_emb = mydat_emb.iloc[mydat_train_inds]\n",
        "\n",
        "\n",
        "mydat_y_notrain = np.delete(mydat_y, mydat_train_inds)\n",
        "mydat_array_notrain = np.delete(mydat_array, mydat_train_inds, axis=0)\n",
        "mydat_emb_notrain = mydat_emb[~mydat_emb.index.isin(mydat_train_inds)]\n",
        "productType_notrain = mydat_emb_notrain['productType']\n",
        "mydat_val, _, mydat_val_emb, _, mydat_val_y, _ = train_test_split(mydat_array_notrain, mydat_emb_notrain, mydat_y_notrain, train_size=(n/2)/(190000-n), stratify = productType_notrain, random_state =random_state)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j1Aq104Q4XRZ"
      },
      "source": [
        "#Trees"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C90_xzTA5KdY"
      },
      "source": [
        "The paper explores three methods: bagging, bias-corrected bagging, and boosting."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A6u80n5FwkCc"
      },
      "source": [
        "## Bagging"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H7YzBOHR4ZtV"
      },
      "outputs": [],
      "source": [
        "from sklearn.ensemble import BaggingRegressor\n",
        "\n",
        "reg = BaggingRegressor(n_estimators=300, random_state=0, oob_score=True).fit(mydat_train, mydat_train_y)\n",
        "\n",
        "mypreds_full = reg.predict(mydat_array)\n",
        "mypreds_val = reg.predict(mydat_val)\n",
        "mypreds = reg.predict(mydat_train)\n",
        "\n",
        "errs = mypreds_full.flatten() - mydat_y.flatten()\n",
        "\n",
        "print(np.mean(np.square(errs))/1000000)\n",
        "print(np.mean(np.absolute(errs))/1000)\n",
        "print(np.sqrt(np.mean(np.square(errs)))/mave*100)\n",
        "print(np.mean(np.absolute(errs))/mave*100)\n",
        "print(np.mean(errs)/mydat_y.mean())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pdxi8NALwoxj"
      },
      "source": [
        "##Bias-corrected bagging"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BoshUrNDtHaH"
      },
      "outputs": [],
      "source": [
        "oob_pred_train = reg.oob_prediction_\n",
        "bias_train = oob_pred_train - mydat_train_y\n",
        "regbias = BaggingRegressor(n_estimators=300, random_state=0).fit(mydat_train, bias_train)\n",
        "\n",
        "mypreds_full = reg.predict(mydat_array)-regbias.predict(mydat_array)\n",
        "mypreds_val = reg.predict(mydat_val)-regbias.predict(mydat_val)\n",
        "mypreds = reg.predict(mydat_train)-regbias.predict(mydat_train)\n",
        "\n",
        "errs = mypreds_full.flatten() - mydat_y.flatten()\n",
        "\n",
        "print(np.mean(np.square(errs))/1000000)\n",
        "print(np.mean(np.absolute(errs))/1000)\n",
        "print(np.sqrt(np.mean(np.square(errs)))/mave*100)\n",
        "print(np.mean(np.absolute(errs))/mave*100)\n",
        "print(np.mean(errs)/mydat_y.mean())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bz2xT-i9wtn8"
      },
      "source": [
        "##Boosting"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E-0aCCvSK8C5"
      },
      "outputs": [],
      "source": [
        "from sklearn.ensemble import GradientBoostingRegressor\n",
        "nest = [500, 1000, 2000]\n",
        "maxdepth = [1, 3, 5]\n",
        "minsamplesleaf = [1, 5, 10]\n",
        "learningrate = [0.01, 0.1, 0.2]\n",
        "boost_mat = np.zeros((81, 5))\n",
        "\n",
        "\n",
        "ncount = 0\n",
        "for i in range(3):\n",
        "  currentnest = nest[i]\n",
        "  for j in range(3):\n",
        "    currentmaxdepth = maxdepth[j]\n",
        "    for k in range(3):\n",
        "      currentminsamplesleaf = minsamplesleaf[k]\n",
        "      for l in range(3):\n",
        "        print(ncount)\n",
        "        currentlearningrate = learningrate[l]\n",
        "        booster = GradientBoostingRegressor(n_estimators=currentnest,\n",
        "                                            learning_rate = currentlearningrate,\n",
        "                                            max_depth = currentmaxdepth,\n",
        "                                            min_samples_leaf=currentminsamplesleaf,\n",
        "                                            random_state=0).fit(mydat_train, mydat_train_y)\n",
        "        val_preds = booster.predict(mydat_val)\n",
        "        val_error = np.mean(np.square(val_preds.flatten() - mydat_val_y.flatten()))\n",
        "        boost_mat[ncount] = [currentnest, currentmaxdepth, currentminsamplesleaf, currentlearningrate, val_error]\n",
        "\n",
        "\n",
        "\n",
        "        ncount = ncount+1\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DzY0W_U3De8r"
      },
      "outputs": [],
      "source": [
        "from sklearn.ensemble import GradientBoostingRegressor\n",
        "minind = np.argmin(boost_mat[:, 4])\n",
        "mynest = int(boost_mat[minind, 0])\n",
        "mymaxdepth = int(boost_mat[minind, 1])\n",
        "myminsamples = int(boost_mat[minind, 2])\n",
        "mylearning = boost_mat[minind, 3]\n",
        "booster = GradientBoostingRegressor(n_estimators=mynest, learning_rate = mylearning, max_depth = mymaxdepth, min_samples_leaf=myminsamples, random_state=0).fit(mydat_train, mydat_train_y)\n",
        "\n",
        "mypreds_full = booster.predict(mydat_array)\n",
        "mypreds_val = booster.predict(mydat_val)\n",
        "mypreds = booster.predict(mydat_train)\n",
        "\n",
        "errs = mypreds_full.flatten() - mydat_y.flatten()\n",
        "\n",
        "print(np.mean(np.square(errs))/1000000)\n",
        "print(np.mean(np.absolute(errs))/1000)\n",
        "print(np.sqrt(np.mean(np.square(errs)))/mave*100)\n",
        "print(np.mean(np.absolute(errs))/mave*100)\n",
        "print(np.mean(errs)/mydat_y.mean())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZlR_01EQymlm"
      },
      "source": [
        "# Finite NN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iftaU3lV_wR4"
      },
      "source": [
        "## One-hot encoding"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "K.clear_session()\n",
        "set_random_seed(1)\n",
        "\n",
        "\n",
        "k = 128\n",
        "inputdata = Input((34, ))\n",
        "mydense = Dense(k, activation = 'relu')(inputdata)\n",
        "mydense = Dense(int(k/2), activation = 'relu')(mydense)\n",
        "myval = Dense(1)(mydense)\n",
        "\n",
        "inputs = inputdata\n",
        "outputs = myval\n",
        "\n",
        "m = KerasModel(inputs = inputs, outputs=outputs)\n",
        "m.compile(optimizer=Adam(), loss='mse')\n",
        "#m.summary()\n",
        "callback = EarlyStopping(monitor='val_loss', patience=50)\n",
        "m.fit(x=mydat_train, y=mydat_train_y, epochs = 20000, validation_data = (mydat_val, mydat_val_y), callbacks = [callback], verbose = 0, batch_size = n)"
      ],
      "metadata": {
        "id": "3n1LHyrr01yg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mypreds_full = m.predict(x=mydat_array)\n",
        "mypreds_val = m.predict(x=mydat_val)\n",
        "mypreds = m.predict(x=mydat_train)\n",
        "\n",
        "errs = mypreds_full.flatten() - mydat_y.flatten()\n",
        "\n",
        "print(np.mean(np.square(errs))/1000000)\n",
        "print(np.mean(np.absolute(errs))/1000)\n",
        "print(np.sqrt(np.mean(np.square(errs)))/mave*100)\n",
        "print(np.mean(np.absolute(errs))/mave*100)\n",
        "print(np.mean(errs)/mydat_y.mean())\n"
      ],
      "metadata": {
        "id": "lu_iIrFf090M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(19):\n",
        "  print(\"Product type:\" + str(i+1))\n",
        "  currenterrs = errs[10000*i:(10000*i+10000)]\n",
        "  currenty = mydat_y[10000*i:(10000*i+10000)]\n",
        "  currentmave = np.mean(np.abs(currenty))\n",
        "  currentmse = np.mean(np.square(currenterrs))\n",
        "  print(currentmse/1000000)\n",
        "  print(np.sqrt(currentmse)/currentmave)\n",
        "  print(np.mean(currenterrs)/np.mean(currenty))"
      ],
      "metadata": {
        "id": "TlevJtqo1ODr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m63Xjz9RNYN-"
      },
      "source": [
        "## With embeddings"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from keras.layers import Embedding\n",
        "K.clear_session()\n",
        "set_random_seed(1)\n",
        "\n",
        "\n",
        "\n",
        "input_product_type = Input(shape=(1,))\n",
        "output_product_type = Embedding(19, 5, name='product_type_embedding')(input_product_type)\n",
        "output_product_type = Reshape(target_shape=(5,))(output_product_type)\n",
        "input_mat = Input(shape = (15, ))\n",
        "output_mat = input_mat\n",
        "input_list = [input_mat, input_product_type]\n",
        "outputs =[output_mat, output_product_type]\n",
        "outputs = Concatenate()(outputs)\n",
        "k = 128\n",
        "\n",
        "outputs = Dense(k, activation = 'relu')(outputs)\n",
        "outputs = Dense(int(k/2), activation = 'relu')(outputs)\n",
        "outputs = Dense(1)(outputs)\n",
        "\n",
        "m2 = KerasModel(inputs = input_list, outputs=outputs)\n",
        "m2.compile(optimizer=Adam(), loss='mse')\n",
        "#m2.summary()\n",
        "callback = EarlyStopping(monitor='val_loss', patience=50)\n",
        "\n",
        "\n",
        "def split_features(dat_emb):\n",
        "  product_type = np.array(dat_emb['productType']).astype('float')\n",
        "  dat_emb = np.array(dat_emb.loc[:, dat_emb.columns != 'productType'].values)\n",
        "  return [dat_emb, product_type]\n",
        "\n",
        "m2.fit(x=split_features(mydat_train_emb), y=mydat_train_y, epochs = 20000, validation_data = (split_features(mydat_val_emb), mydat_val_y), callbacks = [callback], verbose = 2, batch_size = n)\n"
      ],
      "metadata": {
        "id": "N_fAYaCy11R0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "m2preds_full = m2.predict(x=split_features(mydat_emb))\n",
        "m2preds_val = m2.predict(x=split_features(mydat_val_emb))\n",
        "m2preds = m2.predict(x=split_features(mydat_train_emb))\n",
        "\n",
        "errs2 = m2preds_full.flatten() - mydat_y.flatten()\n",
        "\n",
        "print(np.mean(np.square(errs2))/1000000)\n",
        "print(np.mean(np.absolute(errs2))/1000)\n",
        "print(np.sqrt(np.mean(np.square(errs2)))/mave*100)\n",
        "print(np.mean(np.absolute(errs2))/mave*100)\n",
        "print(np.mean(errs2)/mydat_y.mean())"
      ],
      "metadata": {
        "id": "CBvMMxrl18Zg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d06823fd-e485-4a7d-84db-3177dd88d600"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m5938/5938\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m11s\u001b[0m 2ms/step\n",
            "\u001b[1m11/11\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 2ms/step \n",
            "\u001b[1m22/22\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 2ms/step \n",
            "699.4219014133438\n",
            "14.932977704436835\n",
            "26.143891565186976\n",
            "14.76206239025223\n",
            "0.010636112906442091\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(19):\n",
        "  print(\"Product type:\" + str(i+1))\n",
        "  currenterrs = errs[10000*i:(10000*i+10000)]\n",
        "  currenty = mydat_y[10000*i:(10000*i+10000)]\n",
        "  currentmave = np.mean(np.abs(currenty))\n",
        "  currentmse = np.mean(np.square(currenterrs))\n",
        "  print(currentmse/1000000)\n",
        "  print(np.sqrt(currentmse)/currentmave)\n",
        "  print(np.mean(currenterrs)/np.mean(currenty))"
      ],
      "metadata": {
        "id": "kCpUnHzj2Gph"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bxAfkOqdki1q"
      },
      "source": [
        "## Hejazi methodology"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hWbqR-f1kmkR"
      },
      "outputs": [],
      "source": [
        "from sklearn.model_selection import train_test_split\n",
        "from tensorflow.keras import backend as K\n",
        "\n",
        "mydat_trainr, mydat_rep, mydat_trainr_y, mydat_rep_y = train_test_split(mydat_train_emb, mydat_train_y, test_size = 0.5, train_size = 0.5, random_state = random_state)\n",
        "#print(len(mydat_rep), len(mydat_trainr), len(mydat_val))\n",
        "\n",
        "def getcatdist(catvec2, catvec1):\n",
        "  return 1*np.array(catvec2 != catvec1)\n",
        "\n",
        "def getmaxdist(numvec2, numvec1):\n",
        "  numvec2 = np.array(numvec2)\n",
        "  numvec1 = np.array(numvec1)\n",
        "  return 1*np.maximum(numvec2-numvec1, 0)\n",
        "\n",
        "def getmindist(numvec2, numvec1):\n",
        "  numvec2 = np.array(numvec2)\n",
        "  numvec1 = np.array(numvec1)\n",
        "  return 1*np.maximum(numvec1-numvec2, 0)\n",
        "\n",
        "def getdist (list_rep, list_pred):\n",
        "  datcat_rep = list_rep[0]\n",
        "  datnum_rep = list_rep[1]\n",
        "  datcat_pred = list_pred[0]\n",
        "  datnum_pred = list_pred[1]\n",
        "  distlist = []\n",
        "  for i in range(len(list_rep[0])):\n",
        "    currentcat = getcatdist(datcat_rep.iloc[i], datcat_pred)\n",
        "    currentmax = getmaxdist(datnum_rep.iloc[i], datnum_pred)\n",
        "    currentmin = getmindist(datnum_rep.iloc[i], datnum_pred)\n",
        "    currentdist = np.concatenate([currentcat, currentmax, currentmin], axis=1)\n",
        "    distlist.append(currentdist)\n",
        "  return distlist\n",
        "\n",
        "mydat_trainr_cat = mydat_trainr[['gender', 'productType']]\n",
        "mydat_trainr_num = mydat_trainr.loc[:, ~mydat_trainr.columns.isin(['gender', 'productType'])]\n",
        "mydat_trainr_list = [mydat_trainr_cat, mydat_trainr_num]\n",
        "mydat_rep_cat = mydat_rep[['gender', 'productType']]\n",
        "mydat_rep_num = mydat_rep.loc[:, ~mydat_rep.columns.isin(['gender', 'productType'])]\n",
        "mydat_rep_list = [mydat_rep_cat, mydat_rep_num]\n",
        "mydat_val_cat = mydat_val_emb[['gender', 'productType']]\n",
        "mydat_val_num = mydat_val_emb.loc[:, ~mydat_val_emb.columns.isin(['gender', 'productType'])]\n",
        "mydat_val_list = [mydat_val_cat, mydat_val_num]\n",
        "mydat_cat = mydat_emb[['gender', 'productType']]\n",
        "mydat_num = mydat_emb.loc[:, ~mydat_emb.columns.isin(['gender', 'productType'])]\n",
        "mydat_list = [mydat_cat, mydat_num]\n",
        "\n",
        "\n",
        "K.clear_session()\n",
        "set_random_seed(1)\n",
        "inputlist = []\n",
        "denselist = []\n",
        "nrep = len(mydat_rep)\n",
        "for i in range(nrep):\n",
        "  inputlist.append(Input((30, ), name = 'input'+str(i)))\n",
        "  denselist.append(Dense(1, name = 'dense'+str(i))(inputlist[i]))\n",
        "\n",
        "weights = Concatenate()(denselist)\n",
        "weights = Softmax()(weights)\n",
        "output = Dense(1, use_bias = False, name = 'weightedval', trainable=False)(weights)\n",
        "m = KerasModel(inputs = inputlist, outputs = output)\n",
        "m.get_layer('weightedval').set_weights([np.reshape(np.array(mydat_rep_y), (len(mydat_rep_y), 1))])\n",
        "#this forces the weights to each multiply by the representative portfolio values\n",
        "m.compile(optimizer=Adam(), loss='mse')\n",
        "#m.summary()\n",
        "callback = EarlyStopping(monitor='val_loss', patience=50)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K6_tkufpidvP"
      },
      "outputs": [],
      "source": [
        "trainr_dist = getdist(mydat_rep_list, mydat_trainr_list)\n",
        "val_dist = getdist(mydat_rep_list, mydat_val_list)\n",
        "m.fit(x=trainr_dist, y=mydat_trainr_y, epochs = 20000, validation_data = (val_dist, mydat_val_y), callbacks = [callback], verbose = 1, batch_size = 16)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SngA3vulY9Wo"
      },
      "outputs": [],
      "source": [
        "neval = 2048\n",
        "nfull = len(mydat_y)\n",
        "niters = nfull//neval +1\n",
        "predslist = []\n",
        "for i in range(niters):\n",
        "  start = int(neval*i)\n",
        "  stop = int(np.min([neval*(i+1), nfull]))\n",
        "  currentlist = [mydat_cat[start:stop], mydat_num[start:stop]]\n",
        "  currentdist = getdist(mydat_rep_list, currentlist)\n",
        "  predslist.append(m.predict(x=currentdist))\n",
        "\n",
        "mypreds_full = np.concatenate(predslist)\n",
        "mypreds_val = m.predict(x=val_dist)\n",
        "mypreds_trainr = m.predict(x=trainr_dist)\n",
        "\n",
        "print(np.mean((np.array(mydat_val_y)-np.reshape(mypreds_val, newshape = (len(mypreds_val))))**2)/1e+6)\n",
        "print(np.mean((np.array(mydat_y)-np.reshape(mypreds_full, newshape = (len(mypreds_full))))**2)/1e+6)\n",
        "print(np.mean(np.abs(np.array(mydat_y)-np.reshape(mypreds_full, newshape = (len(mypreds_full)))))/1e+3)\n",
        "print(np.sqrt(np.mean((np.array(mydat_y)-np.reshape(mypreds_full, newshape = (len(mypreds_full))))**2))/np.abs(mydat_y).mean())\n",
        "print(np.mean(np.abs(np.array(mydat_y)-np.reshape(mypreds_full, newshape = (len(mypreds_full)))))/np.abs(mydat_y).mean())\n",
        "print(np.mean(np.array(mydat_y)-np.reshape(mypreds_full, newshape = (len(mypreds_full))))/mydat_y.mean())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xXbQwBeMz89I"
      },
      "source": [
        "# NTK"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W-vzD482_8Dg"
      },
      "source": [
        "## Plain NTK"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ihTWRLhOLJcR"
      },
      "source": [
        "This code is used for both plain NTK and regularized NTK. To obtain the regularized NTK, set diag_reg = 1e-4. Additionally, you will need to manually change the activation functions (e.g., stax.relu, stax.erf, or stax.gelu) and adjust the architecture from a 1-layer to a 2-layer network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G7Y8LL0-GUfj"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import predict, stax\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "from datetime import datetime\n",
        "\n",
        "def unflatten(x):\n",
        "  return np.reshape(x, (len(x), 1))\n",
        "\n",
        "mydat_train_y_ntk = jnp.array(unflatten(mydat_train_y))\n",
        "mydat_train_ntk = jnp.array(mydat_train)\n",
        "mydat_val_y_ntk = jnp.array(unflatten(mydat_val_y))\n",
        "mydat_val_ntk = jnp.array(mydat_val)\n",
        "mydat_y_ntk = jnp.array(unflatten(mydat_y))\n",
        "mydat_array_ntk = jnp.array(mydat_array)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_m5cwOhFJGbt"
      },
      "outputs": [],
      "source": [
        "K = 150\n",
        "mae_val = np.zeros(K)\n",
        "mse_val = np.zeros(K)\n",
        "mae_test = np.zeros(K)\n",
        "mse_test = np.zeros(K)\n",
        "mse_train = np.zeros(K)\n",
        "perr = np.zeros(K)\n",
        "b_std = 0.1\n",
        "diag_reg = 0\n",
        "myactivation = stax.Relu #stax.Erf, stax.Gelu\n",
        "\n",
        "for i in range(K):\n",
        "  random.seed(10)\n",
        "\n",
        "  W_std = (i+1)/10\n",
        "  print(W_std)\n",
        "\n",
        "  init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "    stax.Dense(128, W_std=W_std, b_std=b_std), myactivation(),\n",
        "    #stax.Dense(128, W_std=W_std, b_std=b_std), myactivation(),\n",
        "    stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "  )\n",
        "\n",
        "\n",
        "\n",
        "  kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "  get = 'ntk'\n",
        "\n",
        "  #Ktrain_train = np.array(kernel_fn_jit(mydat_train, mydat_train, get))\n",
        "  #Ktrain_train_reg = Ktrain_train + 0*np.identity(len(mydat_train_y))\n",
        "  #Ktrain_full = np.array(kernel_fn_jit(mydat_array, mydat_train, get))\n",
        "  #Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "  #alpha = np.linalg.solve(Ktrain_train_reg, mydat_train_y)\n",
        "  #mypreds_train_ntk = np.matmul(Ktrain_train, alpha)\n",
        "  #mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "  #mypreds_full_ntk = np.matmul(Ktrain_full, alpha)\n",
        "  predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn_jit, mydat_train_ntk, mydat_train_y_ntk, diag_reg = diag_reg)\n",
        "  mypreds_train_ntk = predict_fn(x_test=mydat_train_ntk, get=get, compute_cov=False).flatten()\n",
        "  mypreds_val_ntk = predict_fn(x_test=mydat_val_ntk, get=get, compute_cov=False).flatten()\n",
        "  mypreds_full_ntk = predict_fn(x_test=mydat_array_ntk, get=get, compute_cov=False).flatten()\n",
        "  errs_ntk_train = mypreds_train_ntk - mydat_train_y\n",
        "  errs_ntk_val =mypreds_val_ntk-mydat_val_y\n",
        "  errs_ntk_test = mypreds_full_ntk-mydat_y\n",
        "  mae_val[i] = np.mean(np.absolute(errs_ntk_val))\n",
        "  mse_val[i] = np.mean(np.square(errs_ntk_val))\n",
        "  mae_test[i] = np.mean(np.absolute(errs_ntk_test))\n",
        "  mse_test[i] = np.mean(np.square(errs_ntk_test))\n",
        "  perr[i] = np.mean(errs_ntk_test)/mydat_y.mean()\n",
        "  mse_train[i] = np.mean(np.square(errs_ntk_train))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "apGEHpFBiN69"
      },
      "outputs": [],
      "source": [
        "random.seed(10)\n",
        "#W_std = (sum(np.isnan(mse_test)) + np.argmin(mse_test[~np.isnan(mse_test)]))/10+0.1\n",
        "W_std = (sum(np.isnan(mse_val)) + np.argmin(mse_val[~np.isnan(mse_val)]))/10+0.1\n",
        "\n",
        "print(W_std)\n",
        "b_std = 0.1\n",
        "myactivation = stax.Relu #stax.Erf, stax.Gelu\n",
        "\n",
        "init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "    stax.Dense(128, W_std=W_std, b_std=b_std), myactivation(),\n",
        "    stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        ")\n",
        "\n",
        "diag_reg = 0\n",
        "start = time.time()\n",
        "kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "Ktrain_train = np.array(kernel_fn_jit(mydat_train, mydat_train, 'ntk'))\n",
        "Ktrain_train_reg = Ktrain_train + diag_reg*np.identity(len(mydat_train_y))\n",
        "Ktrain_full = np.array(kernel_fn_jit(mydat_array, mydat_train, 'ntk'))\n",
        "Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, 'ntk'))\n",
        "alpha = np.linalg.solve(Ktrain_train_reg, mydat_train_y)\n",
        "mypreds_train_ntk = np.matmul(Ktrain_train, alpha).flatten()\n",
        "mypreds_val_ntk = np.matmul(Ktrain_val, alpha).flatten()\n",
        "mypreds_full_ntk = np.matmul(Ktrain_full, alpha).flatten()\n",
        "#get = 'ntk'\n",
        "#predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn_jit, mydat_train_ntk, mydat_train_y_ntk, diag_reg = diag_reg)\n",
        "#mypreds_train_ntk = predict_fn(x_test=mydat_train_ntk, get=get, compute_cov=False).flatten()\n",
        "#mypreds_val_ntk = predict_fn(x_test=mydat_val_ntk, get=get, compute_cov=False).flatten()\n",
        "#mypreds_full_ntk = predict_fn(x_test=mydat_array_ntk, get=get, compute_cov=False).flatten()\n",
        "end = time.time()\n",
        "\n",
        "print(np.mean(np.square(mypreds_val_ntk-mydat_val_y)))\n",
        "errs_ntk = mydat_y - mypreds_full_ntk\n",
        "print(np.mean(np.square(errs_ntk)))\n",
        "print(np.mean(np.absolute(errs_ntk)))\n",
        "print(np.sqrt(np.mean(np.square(errs_ntk)))/mave*100)\n",
        "print(np.mean(np.absolute(errs_ntk))/mave*100)\n",
        "print(np.mean(errs_ntk)/mydat_y.mean()*100)\n",
        "print(end-start)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D9lgBtOyK8SM"
      },
      "source": [
        "## Further tuning for bias\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IWlag5XaeDHm"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import predict, stax\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "from datetime import datetime\n",
        "\n",
        "def unflatten(x):\n",
        "  return np.reshape(x, (len(x), 1))\n",
        "\n",
        "mydat_train_y_ntk = jnp.array(unflatten(mydat_train_y))\n",
        "mydat_train_ntk = jnp.array(mydat_train)\n",
        "mydat_val_y_ntk = jnp.array(unflatten(mydat_val_y))\n",
        "mydat_val_ntk = jnp.array(mydat_val)\n",
        "mydat_y_ntk = jnp.array(unflatten(mydat_y))\n",
        "mydat_array_ntk = jnp.array(mydat_array)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bjJ4v5IGdf3_"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab Notebooks/VA/\n",
        "ntk_mat = np.loadtxt('ntk_mat_gelu_srs_reg.csv', delimiter = ',')\n",
        "mse_val = ntk_mat[1]\n",
        "W_std = (sum(np.isnan(mse_val)) + np.argmin(mse_val[~np.isnan(mse_val)]))/10+0.1\n",
        "print(W_std)\n",
        "\n",
        "myactivation = stax.Gelu\n",
        "L = 21\n",
        "mae_val = np.zeros(L)\n",
        "mse_val = np.zeros(L)\n",
        "mae_test = np.zeros(L)\n",
        "mse_test = np.zeros(L)\n",
        "mse_train = np.zeros(L)\n",
        "perr = np.zeros(L)\n",
        "for i in range(L):\n",
        "  random.seed(10)\n",
        "  b_std = i/10\n",
        "  print(b_std)\n",
        "\n",
        "  init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "    stax.Dense(128, W_std=W_std, b_std=b_std), myactivation(),\n",
        "    #stax.Dense(128, W_std=W_std, b_std=b_std), myactivation(),\n",
        "    stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "  )\n",
        "\n",
        "  kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "  get = 'ntk'\n",
        "\n",
        "  #Ktrain_train = np.array(kernel_fn(mydat_train, mydat_train, 'ntk'))\n",
        "  #Ktrain_train_reg = Ktrain_train + 0.001*np.identity(len(mydat_train_y))\n",
        "  #Ktrain_full = np.array(kernel_fn(mydat_array, mydat_train, 'ntk'))\n",
        "  #Ktrain_val = np.array(kernel_fn(mydat_val, mydat_train, 'ntk'))\n",
        "  #alpha = np.linalg.solve(Ktrain_train_reg, mydat_train_y)\n",
        "  #mypreds_train_ntk = np.matmul(Ktrain_train, alpha)\n",
        "  #mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "  #mypreds_full_ntk = np.matmul(Ktrain_full, alpha)\n",
        "  predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn_jit, mydat_train_ntk, mydat_train_y_ntk, diag_reg = 1e-4)\n",
        "  mypreds_train_ntk = predict_fn(x_test=mydat_train_ntk, get=get, compute_cov=False).flatten()\n",
        "  mypreds_val_ntk = predict_fn(x_test=mydat_val_ntk, get=get, compute_cov=False).flatten()\n",
        "  mypreds_full_ntk = predict_fn(x_test=mydat_array_ntk, get=get, compute_cov=False).flatten()\n",
        "  errs_ntk_train = mypreds_train_ntk - mydat_train_y\n",
        "  errs_ntk_val =mypreds_val_ntk-mydat_val_y\n",
        "  errs_ntk_test = mypreds_full_ntk-mydat_y\n",
        "  mae_val[i] = np.mean(np.absolute(errs_ntk_val))\n",
        "  mse_val[i] = np.mean(np.square(errs_ntk_val))\n",
        "  mae_test[i] = np.mean(np.absolute(errs_ntk_test))\n",
        "  mse_test[i] = np.mean(np.square(errs_ntk_test))\n",
        "  perr[i] = np.mean(errs_ntk_test)/mydat_y.mean()\n",
        "  mse_train[i] = np.mean(np.square(errs_ntk_train))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IyPexY1ef7-K"
      },
      "outputs": [],
      "source": [
        "L=21\n",
        "plt.figure(figsize = (3, 2))\n",
        "xvals = np.array(range(L))/(L-1)*2\n",
        "W_val = xvals[(sum(np.isnan(mse_val)) + np.argmin(mse_val[~np.isnan(mse_val)]))]\n",
        "W_test = xvals[(sum(np.isnan(mse_test)) + np.argmin(mse_test[~np.isnan(mse_test)]))]\n",
        "\n",
        "plt.plot(xvals,np.log10(mse_val))\n",
        "plt.plot(xvals, np.log10(mse_test))\n",
        "plt.ylim([8.55, 9.05])\n",
        "#plt.plot([1.5, 1.5], [8.5, 9])\n",
        "plt.ylabel('log(MSE)')\n",
        "plt.xlabel('sigma_b')\n",
        "#plt.legend(['val', 'test'])\n",
        "print(np.log10(mse_test))\n",
        "print(W_val)\n",
        "print(W_test)\n",
        "mse_test_min = np.min(mse_test[~np.isnan(mse_test)])\n",
        "mse_test_val = mse_test[~np.isnan(mse_test)][np.argmin(mse_val[~np.isnan(mse_val)])]\n",
        "print(mse_test_min)\n",
        "print(mse_test_val)\n",
        "print(mse_test[0])\n",
        "print(100*(mse_test_min/mse_test[0]-1))\n",
        "print(100*(mse_test_val/mse_test[0]-1))\n",
        "plt.plot(W_val, np.log10(np.min(mse_val[~np.isnan(mse_val)])), marker = \"*\", ls = None, color = \"#1f77b4\")\n",
        "plt.plot(W_test, np.log10(np.min(mse_test[~np.isnan(mse_test)])), marker = \"*\", ls = None, color = \"#ff7f0e\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H7I3K4eT6b86"
      },
      "source": [
        "## Coordinate descent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ovVrX7Qx60In"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import time\n",
        "from sklearn.model_selection import train_test_split\n",
        "mydat_array = mydat.values\n",
        "def makenpfloat(vec):\n",
        "  return np.array(vec, dtype = np.float32)\n",
        "\n",
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import predict, stax\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "from datetime import datetime\n",
        "\n",
        "def unflatten(x):\n",
        "  return np.reshape(x, (len(x), 1))\n",
        "\n",
        "activations = [stax.Relu]\n",
        "K = 150\n",
        "B = 21\n",
        "R = 81\n",
        "W_vec = np.array(range(K))/10+0.1\n",
        "B_vec = np.array(range(B))/10\n",
        "R_vec = np.concatenate(([0], 10**((np.array(range(R))-(R-1))/10)))\n",
        "seeds = [10]\n",
        "nreps = len(seeds)\n",
        "ntk_mat_cd = np.zeros((nreps*len(activations), 10))\n",
        "for L in range(len(activations)):\n",
        "  kernel_fn_list = [[] for _ in range(150)]\n",
        "  activation = activations[L]\n",
        "  for i in range(K):\n",
        "    W_std = (i+1)/10\n",
        "    for j in range(B):\n",
        "      b_std = j/10\n",
        "      init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "        stax.Dense(128, W_std=W_std, b_std=b_std), activation(),\n",
        "        stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "      )\n",
        "      kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "      kernel_fn_list[i].append(kernel_fn_jit)\n",
        "  for M in range(nreps):\n",
        "\n",
        "    current_seed = seeds[M]\n",
        "\n",
        "    mydat_train, mydat_val, mydat_train_emb, mydat_val_emb, mydat_train_y, mydat_val_y  = train_test_split(\n",
        "        mydat_array, mydat_emb, mydat_y, test_size=340/190000,train_size=680/190000,\n",
        "        stratify = mydat_emb['productType'], random_state =current_seed)\n",
        "\n",
        "    mydat_train = makenpfloat(mydat_train)\n",
        "    mydat_val = makenpfloat(mydat_val)\n",
        "    mydat_train_y = makenpfloat(mydat_train_y).flatten()\n",
        "    mydat_val_y = makenpfloat(mydat_val_y).flatten()\n",
        "\n",
        "    mse_val_w = np.zeros(K)\n",
        "    mse_val_b = np.zeros(B)\n",
        "    mse_val_r = np.zeros(R+1)\n",
        "    b_std = 0.1\n",
        "    b_ind = int(10*b_std)\n",
        "    W_std = 3\n",
        "    diag_reg = 1e-4\n",
        "    tol = 1e-4\n",
        "    rel_improvement = 1\n",
        "    error_old = 1e+10\n",
        "    while rel_improvement > tol:\n",
        "      print(\"W_std\")\n",
        "      for i in range(K):\n",
        "        print(i)\n",
        "        random.seed(10)\n",
        "        currentW_std = (i+1)/10\n",
        "        W_ind = int(10*(currentW_std -0.1))\n",
        "        kernel_fn_jit = kernel_fn_list[i][b_ind]\n",
        "        get = 'ntk'\n",
        "        Ktrain_train = np.array(kernel_fn_jit(mydat_train, None, get))\n",
        "        mytrace = np.mean(np.trace(Ktrain_train))\n",
        "        Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(mydat_train_y))\n",
        "        Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "        alpha = np.linalg.solve(Ktrain_train_reg, mydat_train_y)\n",
        "        mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "        errs_val_ntk =mypreds_val_ntk-mydat_val_y\n",
        "        mse_val_w[i] = np.mean(np.square(errs_val_ntk))\n",
        "      W_ind = np.argmin(mse_val_w)\n",
        "      W_std = W_ind/10+0.1\n",
        "      if error_old == 1e+10:\n",
        "        error_old = np.min(mse_val_w)\n",
        "      kernel_fn_jit = kernel_fn_list[W_ind][b_ind]\n",
        "      Ktrain_train = np.array(kernel_fn_jit(mydat_train, None, get))\n",
        "      Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "      mytrace = np.mean(np.trace(Ktrain_train))\n",
        "\n",
        "\n",
        "      print(\"Diag\")\n",
        "      for i in range(R+1):\n",
        "        diag_reg = R_vec[i]\n",
        "        print(i)\n",
        "        Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(mydat_train_y))\n",
        "        alpha = np.linalg.solve(Ktrain_train_reg, mydat_train_y)\n",
        "        mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "        errs_val_ntk = mypreds_val_ntk - mydat_val_y\n",
        "        mse_val_r[i] = np.mean(np.square(errs_val_ntk))\n",
        "      diag_reg = R_vec[np.argmin(mse_val_r)]\n",
        "      print(\"b_std\")\n",
        "      for i in range(B):\n",
        "        print(i)\n",
        "        kernel_fn_jit = kernel_fn_list[W_ind][i]\n",
        "        Ktrain_train = np.array(kernel_fn_jit(mydat_train, None, get))\n",
        "        mytrace = np.mean(np.trace(Ktrain_train))\n",
        "        Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(mydat_train_y))\n",
        "        Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "        alpha = np.linalg.solve(Ktrain_train_reg, mydat_train_y)\n",
        "        mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "        errs_val_ntk = mypreds_val_ntk - mydat_val_y\n",
        "        mse_val_b[i] = np.mean(np.square(errs_val_ntk))\n",
        "      error_new = np.min(mse_val_b)\n",
        "      b_ind = np.argmin(mse_val_b)\n",
        "      rel_improvement = np.log(error_old)-np.log(error_new)\n",
        "      error_old = error_new\n",
        "    W_std = W_ind/10 + 0.1\n",
        "    b_std = b_ind/10\n",
        "    kernel_fn_jit = kernel_fn_list[W_ind][b_ind]\n",
        "    start = time.time()\n",
        "    Ktrain_train = np.array(kernel_fn_jit(mydat_train, mydat_train, get))\n",
        "    mytrace = np.mean(np.trace(Ktrain_train))\n",
        "    Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(mydat_train_y))\n",
        "    Ktrain_full = np.array(kernel_fn_jit(mydat_array, mydat_train, 'ntk'))\n",
        "    Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "    alpha = np.linalg.solve(Ktrain_train_reg, mydat_train_y)\n",
        "    mypreds_full_ntk = np.matmul(Ktrain_full, alpha)\n",
        "    mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "    end = time.time()\n",
        "    time_eval = start-end\n",
        "    errs_full_ntk =mypreds_full_ntk-mydat_y\n",
        "    errs_val_ntk = mypreds_val_ntk - mydat_val_y\n",
        "    currentind = 5*L+M\n",
        "    ntk_mat_cd[currentind, 0] = np.mean(np.square(errs_full_ntk))/1e+6\n",
        "    ntk_mat_cd[currentind, 1] = np.mean(np.absolute(errs_full_ntk))/1e+3\n",
        "    ntk_mat_cd[currentind, 2] = np.sqrt(ntk_mat_cd[currentind, 0])/(mave/1e+3)\n",
        "    ntk_mat_cd[currentind, 3] = ntk_mat_cd[currentind, 1]/(mave/1e+3)\n",
        "    ntk_mat_cd[currentind, 4] = np.mean(errs_full_ntk)/mave\n",
        "    ntk_mat_cd[currentind, 5] = time_eval\n",
        "    ntk_mat_cd[currentind, 6] = W_std\n",
        "    ntk_mat_cd[currentind, 7] = diag_reg\n",
        "    ntk_mat_cd[currentind, 8] = b_std\n",
        "    ntk_mat_cd[currentind, 9] = np.mean(np.square(errs_val_ntk))/1e+6"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jITVb88lkHiI"
      },
      "source": [
        "## Separate networks NTK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZenBAoJWjsQG"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import predict, stax\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "#import kernel\n",
        "#import eigenpro\n",
        "import jax\n",
        "from jax.lib import xla_bridge\n",
        "import jax.profiler\n",
        "from datetime import datetime\n",
        "\n",
        "torch.cuda.is_available()\n",
        "#print(jax.devices())\n",
        "#print(xla_bridge.get_backend().platform)\n",
        "#!nvidia-smi\n",
        "K=150\n",
        "b_std = 0.1\n",
        "diag_reg = 0\n",
        "activation = stax.Gelu\n",
        "kernelfn_vec = []\n",
        "for j in range(K):\n",
        "  W_std = (j+1)/10\n",
        "  init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "     stax.Dense(128, W_std=W_std, b_std=b_std), activation(),\n",
        "     stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "  )\n",
        "  kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "  kernelfn_vec.append(kernel_fn_jit)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "accRxLFqXGyc"
      },
      "outputs": [],
      "source": [
        "def rmprod(dat):\n",
        "  return np.array(dat.loc[:, dat.columns != 'productType'].values)\n",
        "\n",
        "mse_val_ntk_mat = np.zeros((19, K))\n",
        "mse_train_ntk_mat = np.zeros((19, K))\n",
        "mse_remainder_ntk_mat = np.zeros((19, K))\n",
        "mae_val_ntk_mat = np.zeros((19, K))\n",
        "mae_remainder_ntk_mat = np.zeros((19, K))\n",
        "percenterr_ntk_mat = np.zeros((19, K))\n",
        "dat_train_list = []\n",
        "dat_val_list = []\n",
        "dat_full_list = []\n",
        "dat_train_y_list = []\n",
        "dat_val_y_list = []\n",
        "dat_full_y_list = []\n",
        "\n",
        "for i in range(19):\n",
        "  traininds = (mydat_train_emb['productType'] == i)\n",
        "  valinds = (mydat_val_emb['productType'] == i)\n",
        "  shift = 10000*i\n",
        "  dat_train_list.append(rmprod(mydat_train_emb[traininds]))\n",
        "  dat_val_list.append(rmprod(mydat_val_emb[valinds]))\n",
        "  dat_full_list.append(rmprod(mydat_emb[shift+0:shift+10000]))\n",
        "  dat_train_y_list.append(mydat_train_y[traininds].flatten())\n",
        "  dat_val_y_list.append(mydat_val_y[valinds].flatten())\n",
        "  dat_full_y_list.append(mydat_y[shift+0:shift+10000].flatten())\n",
        "\n",
        "for j in range(K):\n",
        "  print(j)\n",
        "  kernel_fn_jit = kernelfn_vec[j]\n",
        "\n",
        "  for i in range(19):\n",
        "    currentdat_train = dat_train_list[i]\n",
        "    currentdat_val = dat_val_list[i]\n",
        "    currentdat_full = dat_full_list[i]\n",
        "    currentdat_train_y = dat_train_y_list[i]\n",
        "    currentdat_val_y = dat_val_y_list[i]\n",
        "    currentdat_y = dat_full_y_list[i]\n",
        "    Ktrain_train = np.array(kernel_fn_jit(currentdat_train, None, 'ntk'))\n",
        "    #Ktrain_train_reg = Ktrain_train + diag_reg* np.mean(np.trace(Ktrain_train))*np.identity(len(currentdat_train_y))\n",
        "    Ktrain_train_reg = Ktrain_train + diag_reg* np.mean(np.trace(Ktrain_train))*np.identity(len(currentdat_train_y))\n",
        "    Ktrain_full = np.array(kernel_fn_jit(currentdat_full, currentdat_train, 'ntk'))\n",
        "    Ktrain_val = np.array(kernel_fn_jit(currentdat_val, currentdat_train, 'ntk'))\n",
        "    alpha = np.linalg.solve(Ktrain_train_reg, currentdat_train_y)\n",
        "\n",
        "    currentpreds_train_ntk = np.matmul(Ktrain_train, alpha)\n",
        "    currentpreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "    currentpreds_full_ntk = np.matmul(Ktrain_full, alpha)\n",
        "\n",
        "    currenterrs_train_ntk = currentpreds_train_ntk - currentdat_train_y\n",
        "    currenterrs_full_ntk = currentpreds_full_ntk - currentdat_y\n",
        "    currenterrs_val_ntk = currentpreds_val_ntk - currentdat_val_y\n",
        "    mse_val_ntk_mat[i, j] = np.mean(np.square(currenterrs_val_ntk))\n",
        "    mse_train_ntk_mat[i, j] = np.mean(np.square(currenterrs_train_ntk))\n",
        "    mae_val_ntk_mat[i, j] = np.mean(np.absolute(currenterrs_val_ntk))\n",
        "    mse_remainder_ntk_mat[i, j] = np.mean(np.square(currenterrs_full_ntk))\n",
        "    mae_remainder_ntk_mat[i, j] = np.mean(np.absolute(currenterrs_full_ntk))\n",
        "    percenterr_ntk_mat[i, j] = np.mean(currenterrs_full_ntk)/currentdat_y.mean()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MT03xq6tMyyz"
      },
      "source": [
        "### Coordinate descent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BE9EMLVKM7vZ"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import time\n",
        "from sklearn.model_selection import train_test_split\n",
        "mydat_array = mydat.values\n",
        "def makenpfloat(vec):\n",
        "  return np.array(vec, dtype = np.float32)\n",
        "def rmprod(dat):\n",
        "  return np.array(dat.loc[:, dat.columns != 'productType'].values)\n",
        "\n",
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import predict, stax\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "from datetime import datetime\n",
        "\n",
        "def unflatten(x):\n",
        "  return np.reshape(x, (len(x), 1))\n",
        "\n",
        "activations = [stax.Relu, stax.Erf, stax.Gelu]\n",
        "K = 150\n",
        "B = 21\n",
        "R = 81\n",
        "W_vec = np.array(range(K))/10+0.1\n",
        "B_vec = np.array(range(B))/10\n",
        "R_vec = np.concatenate(([0], 10**((np.array(range(R))-(R-1))/10)))\n",
        "seeds = [10, 20, 30, 40, 50]\n",
        "nreps = 5\n",
        "ntk_mat_cd_prod = np.zeros((nreps*len(activations), 7))\n",
        "ntk_mat_cd_prod_sep = np.zeros((nreps*len(activations), 4, 19))\n",
        "\n",
        "\n",
        "for L in range(len(activations)):\n",
        "#for L in range(1):\n",
        "#  if ntk_mat_b_prod[5*L+4, 0] != 0:\n",
        "#    continue\n",
        "  kernel_fn_list = [[] for _ in range(150)]\n",
        "  activation = activations[L]\n",
        "  for i in range(K):\n",
        "    W_std = (i+1)/10\n",
        "    for j in range(B):\n",
        "      b_std = j/10\n",
        "      _, _, kernel_fn = stax.serial(\n",
        "        stax.Dense(128, W_std=W_std, b_std=b_std), activation(),\n",
        "        stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "      )\n",
        "      kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "      kernel_fn_list[i].append(kernel_fn_jit)\n",
        "  for M in range(nreps):\n",
        "  # for M in range(1):\n",
        "    currentind = 5*L+M\n",
        "    print(currentind*100)\n",
        "#    if ntk_mat_b_prod[currentind, 0] !=0:\n",
        "#      continue\n",
        "\n",
        "    current_seed = seeds[M]\n",
        "\n",
        "    mydat_train, mydat_val, mydat_train_emb, mydat_val_emb, mydat_train_y, mydat_val_y  = train_test_split(\n",
        "        mydat_array, mydat_emb, mydat_y, test_size=340/190000,train_size=680/190000,\n",
        "        stratify = mydat_emb['productType'], random_state =current_seed)\n",
        "\n",
        "    mydat_train = makenpfloat(mydat_train)\n",
        "    mydat_val = makenpfloat(mydat_val)\n",
        "    mydat_train_y = makenpfloat(mydat_train_y).flatten()\n",
        "    mydat_val_y = makenpfloat(mydat_val_y).flatten()\n",
        "    mypreds_full_ntk = np.zeros(len(mydat_y))\n",
        "    mypreds_val_ntk = np.zeros(len(mydat_val_y))\n",
        "\n",
        "\n",
        "\n",
        "    mse_val_w = np.zeros(K)\n",
        "    mse_val_b = np.zeros(B)\n",
        "    mse_val_r = np.zeros(R+1)\n",
        "    b_vec = np.array([0.1]*19)\n",
        "    b_vec_ind = [1]*19\n",
        "    b1_vec = np.array([0.1]*19)\n",
        "    b1_vec_ind = np.array([1]*19)\n",
        "    w_vec = np.array([3]*19)\n",
        "    w_vec_ind = [31]*19\n",
        "    w1_vec = [0]*19\n",
        "    w1_vec_ind = [0]*19\n",
        "    r_vec = np.array([0]*19)\n",
        "    r1_vec = np.array([1e-4]*19)\n",
        "    r_vec_ind = [0]*19\n",
        "    tol = 1e-4\n",
        "    dat_train_list = []\n",
        "    dat_val_list = []\n",
        "    dat_full_list = []\n",
        "    dat_train_y_list = []\n",
        "    dat_val_y_list = []\n",
        "    dat_full_y_list = []\n",
        "    valinds_list = []\n",
        "\n",
        "\n",
        "    for k in range(19):\n",
        "      print(10*k)\n",
        "      traininds = (mydat_train_emb['productType'] == k)\n",
        "      valinds = (mydat_val_emb['productType'] == k)\n",
        "      valinds_list.append(valinds)\n",
        "      shift = 10000*k\n",
        "      dat_train_list.append(rmprod(mydat_train_emb[traininds]))\n",
        "      dat_val_list.append(rmprod(mydat_val_emb[valinds]))\n",
        "      dat_full_list.append(rmprod(mydat_emb[shift+0:shift+10000]))\n",
        "      dat_train_y_list.append(mydat_train_y[traininds].flatten())\n",
        "      dat_val_y_list.append(mydat_val_y[valinds].flatten())\n",
        "      dat_full_y_list.append(mydat_y[shift+0:shift+10000].flatten())\n",
        "      diag_reg = 0\n",
        "      b_ind = 1\n",
        "      b_std = 0.1\n",
        "      rel_improvement = 1\n",
        "      niter = 0\n",
        "\n",
        "      error_old = 1e+10\n",
        "\n",
        "      currentdat_train = dat_train_list[k]\n",
        "      currentdat_val = dat_val_list[k]\n",
        "      currentdat_train_y = dat_train_y_list[k]\n",
        "      currentdat_val_y = dat_val_y_list[k]\n",
        "      while rel_improvement > tol:\n",
        "        print(niter)\n",
        "        print(\"W_std\")\n",
        "        print(diag_reg)\n",
        "        print(b_std)\n",
        "        for i in range(K):\n",
        "          print(i)\n",
        "          random.seed(10)\n",
        "          kernel_fn_jit = kernel_fn_list[i][b_ind]\n",
        "          get = 'ntk'\n",
        "          Ktrain_train = np.array(kernel_fn_jit(currentdat_train, None, get))\n",
        "          mytrace = np.mean(np.trace(Ktrain_train))\n",
        "          Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(currentdat_train_y))\n",
        "          Ktrain_val = np.array(kernel_fn_jit(currentdat_val, currentdat_train, get))\n",
        "          alpha = np.linalg.solve(Ktrain_train_reg, currentdat_train_y)\n",
        "          currentpreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "          errs_val_ntk =currentpreds_val_ntk-currentdat_val_y\n",
        "          mse_val_w[i] = np.mean(np.square(errs_val_ntk))\n",
        "        W_ind = np.argmin(mse_val_w)\n",
        "        w_vec_ind[k] = W_ind\n",
        "        W_std = W_ind/10+0.1\n",
        "        w_vec[k] = W_std\n",
        "        if rel_improvement == 1:\n",
        "          w1_vec[k] = W_std\n",
        "          w1_vec_ind[k] = W_ind\n",
        "\n",
        "\n",
        "        if error_old == 1e+10:\n",
        "          error_old = np.min(mse_val_w)\n",
        "\n",
        "        kernel_fn_jit = kernel_fn_list[W_ind][b_ind]\n",
        "        Ktrain_train = np.array(kernel_fn_jit(currentdat_train, None, get))\n",
        "        Ktrain_val = np.array(kernel_fn_jit(currentdat_val, currentdat_train, get))\n",
        "        mytrace = np.mean(np.trace(Ktrain_train))\n",
        "        print(\"Diag\")\n",
        "        print(W_std)\n",
        "        print(0.1*b_ind)\n",
        "        for i in range(R+1):\n",
        "          diag_reg = R_vec[i]\n",
        "          print(i)\n",
        "          Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(currentdat_train_y))\n",
        "          alpha = np.linalg.solve(Ktrain_train_reg, currentdat_train_y)\n",
        "          currentpreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "          errs_val_ntk = currentpreds_val_ntk - currentdat_val_y\n",
        "          mse_val_r[i] = np.mean(np.square(errs_val_ntk))\n",
        "        diag_reg = R_vec[np.argmin(mse_val_r)]\n",
        "        r_vec[k] = diag_reg\n",
        "        r_vec_ind[k] =np.argmin(mse_val_r)\n",
        "        if rel_improvement == 1:\n",
        "          r1_vec[k] = diag_reg\n",
        "\n",
        "        print(\"b_std\")\n",
        "        print(W_std)\n",
        "        print(diag_reg)\n",
        "        for i in range(B):\n",
        "          print(i)\n",
        "          kernel_fn_jit = kernel_fn_list[W_ind][i]\n",
        "          Ktrain_train = np.array(kernel_fn_jit(currentdat_train, None, get))\n",
        "          mytrace = np.mean(np.trace(Ktrain_train))\n",
        "          Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(currentdat_train_y))\n",
        "          Ktrain_val = np.array(kernel_fn_jit(currentdat_val, currentdat_train, get))\n",
        "          alpha = np.linalg.solve(Ktrain_train_reg, currentdat_train_y)\n",
        "          currentpreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "          errs_val_ntk = currentpreds_val_ntk - currentdat_val_y\n",
        "          mse_val_b[i] = np.mean(np.square(errs_val_ntk))\n",
        "        b_ind = np.argmin(mse_val_b)\n",
        "        b_vec_ind[k] = b_ind\n",
        "        b_std = 0.1*b_ind\n",
        "        b_vec[k] = b_std\n",
        "\n",
        "        if rel_improvement == 1:\n",
        "          b1_vec[k] = b_ind/10\n",
        "          b1_vec_ind[k] = b_ind\n",
        "\n",
        "        error_new = np.min(mse_val_b)\n",
        "        rel_improvement = np.log(error_old)-np.log(error_new)\n",
        "        error_old = error_new\n",
        "        niter +=1\n",
        "\n",
        "\n",
        "\n",
        "    print(\"Coordinate descent\")\n",
        "    print(w_vec_ind)\n",
        "    print(b_vec_ind)\n",
        "    print(R_vec[r_vec_ind])\n",
        "    currentind = 5*L+M\n",
        "    start = time.time()\n",
        "    for k in range(19):\n",
        "      W_ind = w_vec_ind[k]\n",
        "      W_std = W_ind/10 + 0.1\n",
        "      b_ind = b_vec_ind[k]\n",
        "      b_std = b_ind/10\n",
        "      diag_reg = R_vec[r_vec_ind[k]]\n",
        "      init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "        stax.Dense(128, W_std=W_std, b_std=b_std), activation(),\n",
        "        stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "      )\n",
        "      currentdat_train = dat_train_list[k]\n",
        "      currentdat_full = dat_full_list[k]\n",
        "      currentdat_train_y = dat_train_y_list[k]\n",
        "      currentdat_val_y = dat_val_y_list[k]\n",
        "      currentdat_val = dat_val_list[k]\n",
        "      currentdat_y = dat_full_y_list[k]\n",
        "      kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "\n",
        "      Ktrain_train = np.array(kernel_fn_jit(currentdat_train, None, get))\n",
        "      mytrace = np.mean(np.trace(Ktrain_train))\n",
        "      Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(currentdat_train_y))\n",
        "      Ktrain_full = np.array(kernel_fn_jit(currentdat_full, currentdat_train, 'ntk'))\n",
        "      Ktrain_val = np.array(kernel_fn_jit(currentdat_val, currentdat_train, get))\n",
        "      alpha = np.linalg.solve(Ktrain_train_reg, currentdat_train_y)\n",
        "      shift = k*10000\n",
        "      currentpreds_ntk = np.matmul(Ktrain_full, alpha)\n",
        "      mypreds_full_ntk[shift+0:shift+10000] = currentpreds_ntk\n",
        "      currentpreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "      currentvalinds = valinds_list[k]\n",
        "      mypreds_val_ntk[currentvalinds] = currentpreds_val_ntk\n",
        "    end = time.time()\n",
        "    time_eval = end-start\n",
        "    errs_full_ntk =mypreds_full_ntk-mydat_y\n",
        "    errs_val_ntk =mypreds_val_ntk-mydat_val_y\n",
        "    currentind = 5*L+M\n",
        "\n",
        "    ntk_mat_cd_prod[currentind, 0] = np.mean(np.square(errs_val_ntk))/1e+6\n",
        "    ntk_mat_cd_prod[currentind, 1] = np.mean(np.square(errs_full_ntk))/1e+6\n",
        "    ntk_mat_cd_prod[currentind, 2] = np.mean(np.absolute(errs_full_ntk))/1e+3\n",
        "    ntk_mat_cd_prod[currentind, 3] = np.sqrt(ntk_mat_cd_prod[currentind, 0])/(mave/1e+3)\n",
        "    ntk_mat_cd_prod[currentind, 4] = ntk_mat_cd_prod[currentind, 1]/(mave/1e+3)\n",
        "    ntk_mat_cd_prod[currentind, 5] = np.mean(errs_full_ntk)/mave\n",
        "    ntk_mat_cd_prod[currentind, 6] = time_eval\n",
        "\n",
        "\n",
        "    #np.savetxt('NTK_mat2_CD_prod.csv', ntk_mat_cd_prod)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W9QH4pPoxYO_"
      },
      "source": [
        "#Kernel regression"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vZORORskxajn"
      },
      "outputs": [],
      "source": [
        "from statsmodels.nonparametric.kernel_regression import KernelReg\n",
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "from jax.lax import fori_loop\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "from jax.numpy.linalg import solve\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "\n",
        "def unflatten(x):\n",
        "  return np.reshape(x, (len(x), 1))\n",
        "\n",
        "mydat_train_y_ntk = jnp.array(unflatten(mydat_train_y))\n",
        "mydat_train_ntk = jnp.array(mydat_train)\n",
        "mydat_val_y_ntk = jnp.array(unflatten(mydat_val_y))\n",
        "mydat_val_ntk = jnp.array(mydat_val)\n",
        "mydat_y_ntk = jnp.array(unflatten(mydat_y))\n",
        "mydat_array_ntk = jnp.array(mydat_array)\n",
        "\n",
        "\n",
        "def gauss_const(h):\n",
        "    \"\"\"\n",
        "    Returns the normalization constant for a gaussian\n",
        "    \"\"\"\n",
        "    return 1/(h*np.sqrt(np.pi*2))\n",
        "\n",
        "def makegauss_exp(ker_x, xi, h):\n",
        "    const = gauss_const(h)\n",
        "    den = h*h\n",
        "    diffs = jnp.sum(jnp.square((xi - ker_x)), axis=xi.ndim-1)\n",
        "    return const*jnp.exp(-0.5*diffs/den)\n",
        "\n",
        "makegauss_exp_jit = jit(makegauss_exp, static_argnames = 'h')\n",
        "\n",
        "def makegauss_exp_mat(ker_x, xi, h):\n",
        "    gauss_mat = np.zeros((len(ker_x), len(xi)))\n",
        "    for i in range(len(ker_x)):\n",
        "        gauss_mat[i] = (makegauss_exp_jit(jnp.array(ker_x[i]), xi, h))\n",
        "    return jnp.array(gauss_mat)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "txvH78H56c3c"
      },
      "outputs": [],
      "source": [
        "K = 50\n",
        "mae_val = np.zeros(K)\n",
        "mse_val = np.zeros(K)\n",
        "mae_test = np.zeros(K)\n",
        "mse_test = np.zeros(K)\n",
        "mse_train = np.zeros(K)\n",
        "perr = np.zeros(K)\n",
        "diag_reg = 1e-4\n",
        "\n",
        "\n",
        "for i in range(K):\n",
        "  random.seed(10)\n",
        "\n",
        "  h = (i+1)/10\n",
        "  print(h)\n",
        "\n",
        "\n",
        "  Ktrain_train = makegauss_exp_mat(mydat_train_ntk, mydat_train_ntk, h)\n",
        "  Ktrain_train_reg = Ktrain_train + jnp.trace(Ktrain_train)*diag_reg*jnp.identity(len(mydat_train_y))\n",
        "  Ktrain_full = jnp.transpose(jnp.array(makegauss_exp_mat(mydat_train_ntk, mydat_array_ntk, h)))\n",
        "  Ktrain_val = jnp.array(makegauss_exp_mat(mydat_val_ntk, mydat_train_ntk, h))\n",
        "  alpha = solve(Ktrain_train_reg, mydat_train_y_ntk)\n",
        "  mypreds_train_ntk = jnp.matmul(Ktrain_train, alpha)\n",
        "  mypreds_val_ntk = jnp.matmul(Ktrain_val, alpha)\n",
        "  mypreds_full_ntk = jnp.matmul(Ktrain_full, alpha)\n",
        "  errs_ntk_train = mypreds_train_ntk - mydat_train_y_ntk\n",
        "  errs_ntk_val =mypreds_val_ntk-mydat_val_y_ntk\n",
        "  errs_ntk_test = mypreds_full_ntk-mydat_y_ntk\n",
        "  mae_val[i] = np.mean(np.absolute(errs_ntk_val))\n",
        "  mse_val[i] = np.mean(np.square(errs_ntk_val))\n",
        "  mae_test[i] = np.mean(np.absolute(errs_ntk_test))\n",
        "  mse_test[i] = np.mean(np.square(errs_ntk_test))\n",
        "  perr[i] = np.mean(errs_ntk_test)/mydat_y.mean()\n",
        "  mse_train[i] = np.mean(np.square(errs_ntk_train))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HVgC9IYWaU4U"
      },
      "source": [
        "# Greeks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D_5XTT-9chfK"
      },
      "outputs": [],
      "source": [
        "def getdelta1(mydat_greeks, mydat_vals):\n",
        "    greeks = mydat_greeks[['Delta1']]\n",
        "    greeks = greeks.join(mydat_vals[['base:1_U', 'base:1_D']])\n",
        "    return greeks\n",
        "\n",
        "greeks_train = getdelta1(mydat_greeks_train, mydat_vals_train)\n",
        "greeks_val = getdelta1(mydat_greeks_val, mydat_vals_val)\n",
        "greeks = getdelta1(datgreeks2, datvals2)\n",
        "nmodels = len(greeks_train.columns)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epRARr42eJZ3"
      },
      "source": [
        "##Boosting"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 401
        },
        "collapsed": true,
        "id": "_apI9d-PeJZ4",
        "outputId": "cf43bcba-fb9b-4b75-867f-faa488f61765"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0\n",
            "1\n",
            "2\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-67-43ad0517bc71>\u001b[0m in \u001b[0;36m<cell line: 9>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     24\u001b[0m                                               \u001b[0mmax_depth\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcurrentmaxdepth\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m                                               \u001b[0mmin_samples_leaf\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcurrentminsamplesleaf\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m                                               random_state=0).fit(mydat_train, current_train_y)\n\u001b[0m\u001b[1;32m     27\u001b[0m           \u001b[0mval_preds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbooster\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpredict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmydat_val\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m           \u001b[0mval_error\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msquare\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mval_preds\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mcurrent_val_y\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/sklearn/base.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(estimator, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1150\u001b[0m                 )\n\u001b[1;32m   1151\u001b[0m             ):\n\u001b[0;32m-> 1152\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfit_method\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mestimator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1153\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1154\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/sklearn/ensemble/_gb.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, X, y, sample_weight, monitor)\u001b[0m\n\u001b[1;32m    530\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    531\u001b[0m         \u001b[0;31m# fit the boosting stages\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 532\u001b[0;31m         n_stages = self._fit_stages(\n\u001b[0m\u001b[1;32m    533\u001b[0m             \u001b[0mX_train\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    534\u001b[0m             \u001b[0my_train\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/sklearn/ensemble/_gb.py\u001b[0m in \u001b[0;36m_fit_stages\u001b[0;34m(self, X, y, raw_predictions, sample_weight, random_state, X_val, y_val, sample_weight_val, begin_at_stage, monitor)\u001b[0m\n\u001b[1;32m    608\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    609\u001b[0m             \u001b[0;31m# fit next stage of trees\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 610\u001b[0;31m             raw_predictions = self._fit_stage(\n\u001b[0m\u001b[1;32m    611\u001b[0m                 \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    612\u001b[0m                 \u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/sklearn/ensemble/_gb.py\u001b[0m in \u001b[0;36m_fit_stage\u001b[0;34m(self, i, X, y, raw_predictions, sample_weight, sample_mask, random_state, X_csc, X_csr)\u001b[0m\n\u001b[1;32m    243\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    244\u001b[0m             \u001b[0mX\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mX_csr\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mX_csr\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mX\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 245\u001b[0;31m             \u001b[0mtree\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresidual\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_weight\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msample_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcheck_input\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    246\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    247\u001b[0m             \u001b[0;31m# update tree leaves\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/sklearn/base.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(estimator, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1150\u001b[0m                 )\n\u001b[1;32m   1151\u001b[0m             ):\n\u001b[0;32m-> 1152\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfit_method\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mestimator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1153\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1154\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/sklearn/tree/_classes.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, X, y, sample_weight, check_input)\u001b[0m\n\u001b[1;32m   1318\u001b[0m         \"\"\"\n\u001b[1;32m   1319\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1320\u001b[0;31m         super()._fit(\n\u001b[0m\u001b[1;32m   1321\u001b[0m             \u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1322\u001b[0m             \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/sklearn/tree/_classes.py\u001b[0m in \u001b[0;36m_fit\u001b[0;34m(self, X, y, sample_weight, check_input, missing_values_in_feature_mask)\u001b[0m\n\u001b[1;32m    441\u001b[0m             )\n\u001b[1;32m    442\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 443\u001b[0;31m         \u001b[0mbuilder\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuild\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtree_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmissing_values_in_feature_mask\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    444\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    445\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mn_outputs_\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mis_classifier\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ],
      "source": [
        "from sklearn.ensemble import GradientBoostingRegressor\n",
        "nest = [500, 1000, 2000]\n",
        "maxdepth = [1, 3, 5]\n",
        "minsamplesleaf = [1, 5, 10]\n",
        "learningrate = [0.01, 0.1, 0.2]\n",
        "boost_mat = np.zeros((nmodels, 81, 5))\n",
        "\n",
        "\n",
        "for m in range(nmodels):\n",
        "  current_train_y =np.array(greeks_train.iloc[:, m])\n",
        "  current_val_y =np.array(greeks_val.iloc[:, m])\n",
        "  ncount = 0\n",
        "  for i in range(3):\n",
        "    currentnest = nest[i]\n",
        "    for j in range(3):\n",
        "      currentmaxdepth = maxdepth[j]\n",
        "      for k in range(3):\n",
        "        currentminsamplesleaf = minsamplesleaf[k]\n",
        "        for l in range(3):\n",
        "          print(ncount)\n",
        "          currentlearningrate = learningrate[l]\n",
        "          booster = GradientBoostingRegressor(n_estimators=currentnest,\n",
        "                                              learning_rate = currentlearningrate,\n",
        "                                              max_depth = currentmaxdepth,\n",
        "                                              min_samples_leaf=currentminsamplesleaf,\n",
        "                                              random_state=0).fit(mydat_train, current_train_y)\n",
        "          val_preds = booster.predict(mydat_val)\n",
        "          val_error = np.mean(np.square(val_preds.flatten() - current_val_y.flatten()))\n",
        "          boost_mat[m, ncount] = [currentnest, currentmaxdepth, currentminsamplesleaf, currentlearningrate, val_error]\n",
        "\n",
        "\n",
        "\n",
        "          ncount = ncount+1\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7HkGdWzleJZ5"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "mypreds_full_mat = np.zeros((nmodels, len(mydat_y)))\n",
        "mypreds_val_mat = np.zeros((nmodels, len(mydat_val_y)))\n",
        "\n",
        "\n",
        "for m in range(nmodels):\n",
        "  current_train_y =np.array(greeks_train.iloc[:, m])\n",
        "  current_val_y =np.array(greeks_val.iloc[:, m])\n",
        "  current_y = np.array(greeks.iloc[:, m])\n",
        "  mave = np.mean(np.abs(current_y))\n",
        "  minind = np.argmin(boost_mat[m, :, 4])\n",
        "  mynest = int(boost_mat[m, minind, 0])\n",
        "  mymaxdepth = int(boost_mat[m, minind, 1])\n",
        "  myminsamples = int(boost_mat[m, minind, 2])\n",
        "  mylearning = boost_mat[m, minind, 3]\n",
        "  booster = GradientBoostingRegressor(n_estimators=mynest, learning_rate = mylearning, max_depth = mymaxdepth, min_samples_leaf=myminsamples, random_state=0).fit(mydat_train, current_train_y)\n",
        "\n",
        "  mypreds_full = booster.predict(mydat_array)\n",
        "  mypreds_full_mat[m] = mypreds_full\n",
        "  mypreds_val = booster.predict(mydat_val)\n",
        "  mypreds_val_mat[m] = mypreds_val\n",
        "  mypreds = booster.predict(mydat_train)\n",
        "  errs = mypreds_full.flatten() - current_y.flatten()\n",
        "  errs_val = mypreds_val.flatten()-current_val_y.flatten()\n",
        "\n",
        "  print(np.mean(np.square(errs_val))/1000000)\n",
        "  print(np.mean(np.square(errs))/1000000)\n",
        "  print(np.mean(np.absolute(errs))/1000)\n",
        "  print(np.sqrt(np.mean(np.square(errs)))/mave*100)\n",
        "  print(np.mean(np.absolute(errs))/mave*100)\n",
        "  print(np.mean(errs)/mave)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XZ8wbjve3JxQ"
      },
      "outputs": [],
      "source": [
        "#boost_mat.tofile('boost_mat_delta1.csv', sep = ',')\n",
        "delta1_preds_val = (mypreds_val_mat[1] - mypreds_val_mat[2])/0.02\n",
        "delta1_preds_full = (mypreds_full_mat[1] - mypreds_full_mat[2])/0.02\n",
        "mave = np.mean(np.abs(datgreeks2['Delta1']))\n",
        "errs_preds_val = delta1_preds_val - greeks_val['Delta1']\n",
        "errs_preds_full = delta1_preds_full - datgreeks2['Delta1']\n",
        "print(np.mean(np.square(errs_preds_val))/1000000)\n",
        "print(np.mean(np.square(errs_preds_full))/1000000)\n",
        "print(np.mean(np.abs(errs_preds_full))/1000)\n",
        "print(np.sqrt(np.mean(np.square(errs_preds_full)))/mave)\n",
        "print(np.mean(np.abs(errs_preds_full))/mave)\n",
        "print(np.mean(errs_preds_full)/np.mean(datgreeks2['Delta1']))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OT1r9x-7b1fp"
      },
      "source": [
        "## Finite NN (Best NN)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s_CTkEYtIg_C"
      },
      "outputs": [],
      "source": [
        "from keras.layers import Embedding\n",
        "from tensorflow.keras import backend as K\n",
        "nepochs = 20000\n",
        "mypreds_full_mat = np.zeros((nmodels, len(mydat_y)))\n",
        "mypreds_val_mat = np.zeros((nmodels, len(mydat_val_y)))\n",
        "for m in range(nmodels):\n",
        "  K.clear_session()\n",
        "  set_random_seed(1)\n",
        "\n",
        "\n",
        "  current_train_y =np.array(greeks_train.iloc[:, m])\n",
        "  current_val_y =np.array(greeks_val.iloc[:, m])\n",
        "  current_y = np.array(greeks.iloc[:, m])\n",
        "  k = 128\n",
        "  inputdata = Input((34, ))\n",
        "  mydense = Dense(k, activation = 'relu')(inputdata)\n",
        "  mydense = Dense(int(k/2), activation = 'relu')(mydense)\n",
        "  myval = Dense(1)(mydense)\n",
        "\n",
        "  inputs = inputdata\n",
        "  outputs = myval\n",
        "\n",
        "  m1 = KerasModel(inputs = inputs, outputs=outputs)\n",
        "  m1.compile(optimizer=Adam(), loss='mse')\n",
        "  #m1.summary()\n",
        "  callback = EarlyStopping(monitor='val_loss', patience=50)\n",
        "  m1.fit(x=mydat_train, y=current_train_y, epochs = nepochs, validation_data = (mydat_val, current_val_y), callbacks = [callback], verbose = 0, batch_size = n)\n",
        "\n",
        "\n",
        "  mave = np.mean(np.abs(current_y))\n",
        "  mypreds_train = m1.predict(x=mydat_train).flatten()\n",
        "  errs_train = mypreds_train - current_train_y\n",
        "  mypreds_val  = m1.predict(x=mydat_val).flatten()\n",
        "  errs_val = mypreds_val - current_val_y\n",
        "  mse_val_1 = np.mean(np.square(errs_val)/1000000)\n",
        "  mypreds_full  = m1.predict(x=mydat_array).flatten()\n",
        "  errs = mypreds_full - current_y\n",
        "\n",
        "  K.clear_session()\n",
        "  set_random_seed(1)\n",
        "\n",
        "\n",
        "\n",
        "  input_product_type = Input(shape=(1,))\n",
        "  output_product_type = Embedding(19, 5, name='product_type_embedding')(input_product_type)\n",
        "  output_product_type = Reshape(target_shape=(5,))(output_product_type)\n",
        "  input_mat = Input(shape = (15, ))\n",
        "  output_mat = input_mat\n",
        "  input_list = [input_mat, input_product_type]\n",
        "  outputs =[output_mat, output_product_type]\n",
        "  outputs = Concatenate()(outputs)\n",
        "  k = 128\n",
        "\n",
        "  outputs = Dense(k, activation = 'relu')(outputs)\n",
        "  outputs = Dense(int(k/2), activation = 'relu')(outputs)\n",
        "  outputs = Dense(1)(outputs)\n",
        "\n",
        "  m2 = KerasModel(inputs = input_list, outputs=outputs)\n",
        "  m2.compile(optimizer=Adam(), loss='mse')\n",
        "  #m2.summary()\n",
        "  callback = EarlyStopping(monitor='val_loss', patience=50)\n",
        "\n",
        "\n",
        "  def split_features(dat_emb):\n",
        "    product_type = np.array(dat_emb['productType']).astype('float')\n",
        "    dat_emb = np.array(dat_emb.loc[:, dat_emb.columns != 'productType'].values)\n",
        "    return [dat_emb, product_type]\n",
        "\n",
        "\n",
        "  m2.fit(x=split_features(mydat_train_emb), y=current_train_y, epochs = nepochs, validation_data = (split_features(mydat_val_emb), current_val_y), callbacks = [callback], verbose = 0, batch_size = n)\n",
        "\n",
        "  mypreds2_val = m2.predict(x=split_features(mydat_val_emb)).flatten()\n",
        "  mypreds2_train = m2.predict(x=split_features(mydat_train_emb)).flatten()\n",
        "  errs2_train = mypreds2_train - current_train_y\n",
        "  errs2_val = mypreds2_val - current_val_y\n",
        "  mse_val_2 = np.mean(np.square(errs2_val)/1000000)\n",
        "  mypreds2_full  = m2.predict(x=split_features(mydat_emb)).flatten()\n",
        "  errs2 = mypreds2_full - current_y\n",
        "\n",
        "  if mse_val_1 > mse_val_2:\n",
        "    errs_best = errs2\n",
        "    mypreds_full_mat[m] = mypreds2_full\n",
        "  else:\n",
        "    errs_best = errs\n",
        "    mypreds_full_mat[m] = mypreds_full\n",
        "\n",
        "  print(np.min([mse_val_1, mse_val_2]))\n",
        "  print(np.mean(np.square(errs_best))/1000000)\n",
        "  print(np.mean(np.abs(errs_best))/1000)\n",
        "  print(np.sqrt(np.mean(np.square(errs_best)))/mave)\n",
        "  print(np.mean(np.abs(errs_best))/mave)\n",
        "  print(np.mean(errs_best)/np.mean(current_y))\n",
        "\n",
        "delta1_preds_val = (mypreds_val_mat[1] - mypreds_val_mat[2])/0.02\n",
        "delta1_preds_full = (mypreds_full_mat[1] - mypreds_full_mat[2])/0.02\n",
        "mave = np.mean(np.abs(datgreeks2['Delta1']))\n",
        "errs_preds_val = delta1_preds_val - greeks_val['Delta1']\n",
        "errs_preds_full = delta1_preds_full - datgreeks2['Delta1']\n",
        "print(np.mean(np.square(errs_preds_val))/1000000)\n",
        "print(np.mean(np.square(errs_preds_full))/1000000)\n",
        "print(np.mean(np.abs(errs_preds_full))/1000)\n",
        "print(np.sqrt(np.mean(np.square(errs_preds_full)))/mave)\n",
        "print(np.mean(np.abs(errs_preds_full))/mave)\n",
        "print(np.mean(errs_preds_full)/np.mean(datgreeks2['Delta1']))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GTVOu-T_bxwF"
      },
      "source": [
        "## NTK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kS9lAILwphMx"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import predict, stax\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "from datetime import datetime\n",
        "\n",
        "def unflatten(x):\n",
        "  return np.reshape(x, (len(x), 1))\n",
        "\n",
        "#greeks_train_ntk = jnp.array(greeks_train)\n",
        "mydat_train_ntk = jnp.array(mydat_train)\n",
        "#greeks_val_ntk = jnp.array(greeks_val)\n",
        "mydat_val_ntk = jnp.array(mydat_val)\n",
        "#greeks_ntk = jnp.array(greeks)\n",
        "mydat_array_ntk = jnp.array(mydat_array)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dpTSPkuNqBTE"
      },
      "outputs": [],
      "source": [
        "K = 150\n",
        "mae_val = np.zeros((nmodels, K))\n",
        "mse_val = np.zeros((nmodels, K))\n",
        "mae_test = np.zeros((nmodels, K))\n",
        "mse_test = np.zeros((nmodels, K))\n",
        "#pmae_test = np.zeros((nmodels, K))\n",
        "#prmse_test = np.zeros((nmodels, K))\n",
        "mse_train = np.zeros((nmodels, K))\n",
        "perr = np.zeros((nmodels, K))\n",
        "mypreds_full_mat = np.zeros((nmodels, len(mydat_y)))\n",
        "mypreds_val_mat = np.zeros((nmodels, len(mydat_val_y)))\n",
        "b_std = 0.1\n",
        "diag_reg = 1e-4\n",
        "myactivation = stax.Gelu\n",
        "\n",
        "for m in [0]:\n",
        "  current_train_y = jnp.array(unflatten(greeks_train.iloc[:, m]))\n",
        "  current_val_y = jnp.array(unflatten(greeks_val.iloc[:, m]))\n",
        "  current_y = jnp.array(unflatten(greeks.iloc[:, m]))\n",
        "  mave = np.mean(np.abs(current_y))\n",
        "  for i in range(K):\n",
        "    random.seed(10)\n",
        "\n",
        "    W_std = (i+1)/10\n",
        "    print(W_std)\n",
        "\n",
        "    init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "      stax.Dense(128, W_std=W_std, b_std=b_std), myactivation(),\n",
        "      stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "    )\n",
        "\n",
        "\n",
        "\n",
        "    kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "    get = 'ntk'\n",
        "\n",
        "    predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn_jit, mydat_train_ntk, current_train_y, diag_reg = diag_reg)\n",
        "    mypreds_train_ntk = predict_fn(x_test=mydat_train_ntk, get=get, compute_cov=False).flatten()\n",
        "    mypreds_val_ntk = predict_fn(x_test=mydat_val_ntk, get=get, compute_cov=False).flatten()\n",
        "    mypreds_full_ntk = predict_fn(x_test=mydat_array_ntk, get=get, compute_cov=False).flatten()\n",
        "    errs_ntk_train = mypreds_train_ntk - current_train_y\n",
        "    errs_ntk_val =mypreds_val_ntk-current_val_y.flatten()\n",
        "    errs_ntk_test = mypreds_full_ntk-current_y.flatten()\n",
        "\n",
        "    mae_val[m, i] = np.mean(np.absolute(errs_ntk_val))\n",
        "    mse_val[m, i] = np.mean(np.square(errs_ntk_val))\n",
        "    mae_test[m, i] = np.mean(np.absolute(errs_ntk_test))\n",
        "    mse_test[m, i] = np.mean(np.square(errs_ntk_test))\n",
        "    #pmae_test[m, i] = np.mean(np.absolute(errs_ntk_test))/mave\n",
        "    #prmse_test[m, i] = np.sqrt(np.mean(np.square(errs_ntk_test)))/mave\n",
        "    perr[m, i] = np.mean(errs_ntk_test)/current_y.mean()\n",
        "    mse_train[m, i] = np.mean(np.square(errs_ntk_train))\n",
        "\n",
        "  random.seed(10)\n",
        "  current_mse_val = mse_val[m]\n",
        "  W_std = (sum(np.isnan(current_mse_val)) + np.argmin(current_mse_val[~np.isnan(current_mse_val)]))/10+0.1\n",
        "  b_std = 0.1\n",
        "\n",
        "  init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "      stax.Dense(128, W_std=W_std, b_std=b_std), myactivation(),\n",
        "      stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "  )\n",
        "\n",
        "\n",
        "\n",
        "  start = time.time()\n",
        "  kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "  get = 'ntk'\n",
        "  predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn_jit, mydat_train_ntk, current_train_y, diag_reg = diag_reg)\n",
        "  mypreds_train_ntk = predict_fn(x_test=mydat_train_ntk, get=get, compute_cov=False).flatten()\n",
        "  mypreds_val_ntk = predict_fn(x_test=mydat_val_ntk, get=get, compute_cov=False).flatten()\n",
        "  mypreds_full_ntk = predict_fn(x_test=mydat_array_ntk, get=get, compute_cov=False).flatten()\n",
        "  end = time.time()\n",
        "  mypreds_full_mat[:,m] = mypreds_full_ntk[m]\n",
        "  mypreds_val_mat[:,m] = mypreds_val_ntk[m]\n",
        "\n",
        "  print(np.mean(np.square(mypreds_val_ntk-current_val_y.flatten())))\n",
        "  errs_ntk = current_y.flatten() - mypreds_full_ntk\n",
        "  print(np.mean(np.square(errs_ntk)))\n",
        "  print(np.mean(np.absolute(errs_ntk)))\n",
        "  print(np.sqrt(np.mean(np.square(errs_ntk)))/mave*100)\n",
        "  print(np.mean(np.absolute(errs_ntk))/mave*100)\n",
        "  print(np.mean(errs_ntk)/current_y.mean()*100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q2I4J6D2cFvs"
      },
      "source": [
        "## NTK CD"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KGQYJl1EaTyA"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import time\n",
        "from sklearn.model_selection import train_test_split\n",
        "mydat_array = mydat.values\n",
        "def makenpfloat(vec):\n",
        "  return np.array(vec, dtype = np.float32)\n",
        "\n",
        "\n",
        "import jax.numpy as jnp\n",
        "from jax import random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax import jit, grad, vmap\n",
        "import jax.numpy as jnp\n",
        "import functools\n",
        "import neural_tangents as nt\n",
        "from neural_tangents import predict, stax\n",
        "import numpy as np\n",
        "from random import seed, sample\n",
        "import os\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import warnings\n",
        "from scipy import stats\n",
        "import time\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import random\n",
        "from datetime import datetime\n",
        "\n",
        "def unflatten(x):\n",
        "  return np.reshape(x, (len(x), 1))\n",
        "\n",
        "activations = [stax.Relu, stax.Erf, stax.Gelu]\n",
        "K = 150\n",
        "B = 21\n",
        "R = 81\n",
        "W_vec = np.array(range(K))/10+0.1\n",
        "B_vec = np.array(range(B))/10\n",
        "R_vec = np.concatenate(([0], 10**((np.array(range(R))-(R-1))/10)))\n",
        "seeds = [10, 20, 30, 40, 50]\n",
        "nreps = len(seeds)\n",
        "ntk_mat_cd = np.zeros((nmodels, nreps*len(activations), 10))\n",
        "mypreds_full_mat = np.zeros((nmodels, nreps*len(activations), len(mydat_array)))\n",
        "mypreds_val_mat = np.zeros((nmodels, nreps*len(activations), len(mydat_val)))\n",
        "for L in range(len(activations)):\n",
        "  kernel_fn_list = [[] for _ in range(150)]\n",
        "  activation = activations[L]\n",
        "  for i in range(K):\n",
        "    W_std = (i+1)/10\n",
        "    for j in range(B):\n",
        "      b_std = j/10\n",
        "      init_fn, apply_fn, kernel_fn = stax.serial(\n",
        "        stax.Dense(128, W_std=W_std, b_std=b_std), activation(),\n",
        "        stax.Dense(1, W_std=W_std, b_std=b_std)\n",
        "      )\n",
        "      kernel_fn_jit = jit(kernel_fn, static_argnames='get')\n",
        "      kernel_fn_list[i].append(kernel_fn_jit)\n",
        "  for M in range(nreps):\n",
        "\n",
        "    current_seed = seeds[M]\n",
        "\n",
        "    mydat_train, mydat_val, mydat_train_emb, mydat_val_emb, mydat_train_y, mydat_val_y  = train_test_split(\n",
        "        mydat_array, mydat_emb, mydat_y, test_size=340/190000,train_size=680/190000,\n",
        "        stratify = mydat_emb['productType'], random_state =current_seed)\n",
        "\n",
        "    mydat_train = makenpfloat(mydat_train)\n",
        "    mydat_val = makenpfloat(mydat_val)\n",
        "    mydat_greeks_train = datgreeks2.loc[mydat_train_emb.index]\n",
        "    mydat_greeks_val = datgreeks2.loc[mydat_val_emb.index]\n",
        "    mydat_vals_train = datvals2.loc[mydat_train_emb.index]\n",
        "    mydat_vals_val = datvals2.loc[mydat_val_emb.index]\n",
        "    for m in [0, 1, 2]:\n",
        "      current_train_y = greeks_train.iloc[:, m]\n",
        "      current_val_y = greeks_val.iloc[:, m]\n",
        "      current_y = greeks.iloc[:, m]\n",
        "      mave = np.mean(np.abs(current_y))\n",
        "\n",
        "      mse_val_w = np.zeros(K)\n",
        "      mse_val_b = np.zeros(B)\n",
        "      mse_val_r = np.zeros(R+1)\n",
        "      b_std = 0.1\n",
        "      b_ind = int(10*b_std)\n",
        "      W_std = 3\n",
        "      diag_reg = 1e-4\n",
        "      tol = 1e-4\n",
        "      rel_improvement = 1\n",
        "      error_old = 1e+10\n",
        "      while rel_improvement > tol:\n",
        "        print(\"W_std\")\n",
        "        for i in range(K):\n",
        "          print(i)\n",
        "          random.seed(10)\n",
        "          currentW_std = (i+1)/10\n",
        "          W_ind = int(10*(currentW_std -0.1))\n",
        "          kernel_fn_jit = kernel_fn_list[i][b_ind]\n",
        "          get = 'ntk'\n",
        "          Ktrain_train = np.array(kernel_fn_jit(mydat_train, None, get))\n",
        "          mytrace = np.mean(np.trace(Ktrain_train))\n",
        "          Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(current_train_y))\n",
        "          Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "          alpha = np.linalg.solve(Ktrain_train_reg, current_train_y)\n",
        "          mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "          errs_val_ntk =mypreds_val_ntk-current_val_y\n",
        "          mse_val_w[i] = np.mean(np.square(errs_val_ntk))\n",
        "        W_ind = np.argmin(mse_val_w)\n",
        "        W_std = W_ind/10+0.1\n",
        "        if error_old == 1e+10:\n",
        "          error_old = np.min(mse_val_w)\n",
        "        kernel_fn_jit = kernel_fn_list[W_ind][b_ind]\n",
        "        Ktrain_train = np.array(kernel_fn_jit(mydat_train, None, get))\n",
        "        Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "        mytrace = np.mean(np.trace(Ktrain_train))\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "        print(\"Diag\")\n",
        "        for i in range(R+1):\n",
        "          diag_reg = R_vec[i]\n",
        "          print(i)\n",
        "          Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(current_train_y))\n",
        "          alpha = np.linalg.solve(Ktrain_train_reg, current_train_y)\n",
        "          mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "          errs_val_ntk = mypreds_val_ntk - current_val_y\n",
        "          mse_val_r[i] = np.mean(np.square(errs_val_ntk))\n",
        "        diag_reg = R_vec[np.argmin(mse_val_r)]\n",
        "        print(\"b_std\")\n",
        "        for i in range(B):\n",
        "          print(i)\n",
        "          kernel_fn_jit = kernel_fn_list[W_ind][i]\n",
        "          Ktrain_train = np.array(kernel_fn_jit(mydat_train, None, get))\n",
        "          mytrace = np.mean(np.trace(Ktrain_train))\n",
        "          Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(current_train_y))\n",
        "          Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, get))\n",
        "          alpha = np.linalg.solve(Ktrain_train_reg, current_train_y)\n",
        "          mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "          errs_val_ntk = mypreds_val_ntk - current_val_y\n",
        "          mse_val_b[i] = np.mean(np.square(errs_val_ntk))\n",
        "        error_new = np.min(mse_val_b)\n",
        "        b_ind = np.argmin(mse_val_b)\n",
        "        rel_improvement = np.log(error_old)-np.log(error_new)\n",
        "        error_old = error_new\n",
        "    W_std = W_ind/10 + 0.1\n",
        "    b_std = b_ind/10\n",
        "    kernel_fn_jit = kernel_fn_list[W_ind][b_ind]\n",
        "    start = time.time()\n",
        "    Ktrain_train = np.array(kernel_fn_jit(mydat_train, mydat_train, get))\n",
        "    mytrace = np.mean(np.trace(Ktrain_train))\n",
        "    Ktrain_train_reg = Ktrain_train + diag_reg*mytrace*np.identity(len(current_train_y))\n",
        "    Ktrain_full = np.array(kernel_fn_jit(mydat_array, mydat_train, 'ntk'))\n",
        "    Ktrain_val = np.array(kernel_fn_jit(mydat_val, mydat_train, 'ntk'))\n",
        "    alpha = np.linalg.solve(Ktrain_train_reg, current_train_y)\n",
        "    mypreds_full_ntk = np.matmul(Ktrain_full, alpha)\n",
        "    mypreds_val_ntk = np.matmul(Ktrain_val, alpha)\n",
        "    end = time.time()\n",
        "    time_eval = start-end\n",
        "    errs_full_ntk =mypreds_full_ntk-current_y\n",
        "    errs_val_ntk = mypreds_val_ntk - current_val_y\n",
        "    currentind = nreps*L + M\n",
        "    mse_val_ntk = np.mean(np.square(errs_val_ntk))\n",
        "    mse_full_ntk = np.mean(np.square(errs_full_ntk))\n",
        "    mae_full_ntk = np.mean(np.abs(errs_full_ntk))\n",
        "\n",
        "    ntk_mat_cd[m, currentind, 0] = mse_full_ntk/1e+6\n",
        "    ntk_mat_cd[m, currentind, 1] = mae_full_ntk/1e+3\n",
        "    ntk_mat_cd[m, currentind, 2] = np.sqrt(mse_full_ntk)/(mave/1e+3)\n",
        "    ntk_mat_cd[m, currentind, 3] = mae_full_ntk/(mave/1e+3)\n",
        "    ntk_mat_cd[m, currentind, 4] = np.mean(errs_full_ntk)/current_y.mean()\n",
        "    ntk_mat_cd[m, currentind, 5] = time_eval\n",
        "    ntk_mat_cd[m, currentind, 6] = W_std\n",
        "    ntk_mat_cd[m, currentind, 7] = diag_reg\n",
        "    ntk_mat_cd[m, currentind, 8] = b_std\n",
        "    ntk_mat_cd[m, currentind, 9] = mse_val_ntk/1e+6\n",
        "    mypreds_full_mat[m, currentind, :] = mypreds_full_ntk\n",
        "    mypreds_val_mat[m, currentind, :] = mypreds_val_ntk\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zxlom6dD7qkr"
      },
      "outputs": [],
      "source": [
        "delta1_preds_val = (mypreds_val_mat[1, :, :] - mypreds_val_mat[2, :, :])/0.02\n",
        "delta1_preds_full = (mypreds_full_mat[1, :, :] - mypreds_full_mat[2, :, :])/0.02\n",
        "mave = np.mean(np.abs(datgreeks2['Delta1']))\n",
        "errs_preds_val = np.zeros(delta1_preds_val.shape)\n",
        "errs_preds_full = np.zeros(delta1_preds_full.shape)\n",
        "for i in range(len(delta1_preds_val)):\n",
        "  errs_preds_val[i] = delta1_preds_val[i] - greeks_val['Delta1']\n",
        "  errs_preds_full[i] = delta1_preds_full[i] - datgreeks2['Delta1']\n",
        "  print(np.mean(np.square(errs_preds_val[i]))/1000000)\n",
        "  print(np.mean(np.square(errs_preds_full[i]))/1000000)\n",
        "  print(np.mean(np.abs(errs_preds_full[i]))/1000)\n",
        "  print(np.sqrt(np.mean(np.square(errs_preds_full[i])))/mave)\n",
        "  print(np.mean(np.abs(errs_preds_full[i]))/mave)\n",
        "  print(np.mean(errs_preds_full[i])/np.mean(datgreeks2['Delta1']))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LtgdtkYH_qqC"
      },
      "source": [
        "## Hejazi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wnBxLK_g_w6U",
        "outputId": "368f6ec9-0196-4482-af9c-0bf44d66346d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "340 340 340\n"
          ]
        }
      ],
      "source": [
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "mydat_trainr, mydat_rep, mydat_trainr_y, mydat_rep_y = train_test_split(mydat_train_emb, mydat_train_y, test_size = 0.5, train_size = 0.5, random_state = random_state)\n",
        "print(len(mydat_rep), len(mydat_trainr), len(mydat_val))\n",
        "\n",
        "mydat_greeks_trainr = datgreeks2.loc[mydat_trainr.index]\n",
        "mydat_greeks_rep = datgreeks2.loc[mydat_rep.index]\n",
        "delta_trainr = np.array(mydat_greeks_trainr['Delta1'])\n",
        "delta_rep = np.array(mydat_greeks_rep['Delta1'])\n",
        "delta_val = np.array(mydat_greeks_val['Delta1'])\n",
        "delta_full = np.array(greeks['Delta1'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9U0hyiD9BJlZ"
      },
      "outputs": [],
      "source": [
        "#first, we need a function which evaluates the \"distance vector\"\n",
        "#there are three components:\n",
        "\n",
        "#first, distance on each categorical variable is a Hamming distance\n",
        "def getcatdist(catvec2, catvec1):\n",
        "  return 1*np.array(catvec2 != catvec1)\n",
        "\n",
        "#second, positive distance on numerics: max(y2-y1, 0)\n",
        "#third, (absolute) negative distance on numerics: max(y1-y2, 0)\n",
        "\n",
        "def getmaxdist(numvec2, numvec1):\n",
        "  numvec2 = np.array(numvec2)\n",
        "  numvec1 = np.array(numvec1)\n",
        "  return 1*np.maximum(numvec2-numvec1, 0)\n",
        "\n",
        "def getmindist(numvec2, numvec1):\n",
        "  numvec2 = np.array(numvec2)\n",
        "  numvec1 = np.array(numvec1)\n",
        "  return 1*np.maximum(numvec1-numvec2, 0)\n",
        "\n",
        "#if we concatenate catdist, maxdist and mindist, we get a vector of length 30\n",
        "\n",
        "#this creates a list of n_rep matrices of shape (n_pred x 30)\n",
        "#this allows defining in Tensorflow the list of inputs [dat1, dat2, ..., datn_rep] where each dat is a dataset with n_pred rows\n",
        "def getdist (list_rep, list_pred):\n",
        "  datcat_rep = list_rep[0]\n",
        "  datnum_rep = list_rep[1]\n",
        "  datcat_pred = list_pred[0]\n",
        "  datnum_pred = list_pred[1]\n",
        "  distlist = []\n",
        "  for i in range(len(list_rep[0])):\n",
        "    currentcat = getcatdist(datcat_rep.iloc[i], datcat_pred)\n",
        "    currentmax = getmaxdist(datnum_rep.iloc[i], datnum_pred)\n",
        "    currentmin = getmindist(datnum_rep.iloc[i], datnum_pred)\n",
        "    currentdist = np.concatenate([currentcat, currentmax, currentmin], axis=1)\n",
        "    distlist.append(currentdist)\n",
        "  return distlist\n",
        "\n",
        "mydat_trainr_cat = mydat_trainr[['gender', 'productType']]\n",
        "mydat_trainr_num = mydat_trainr.loc[:, ~mydat_trainr.columns.isin(['gender', 'productType'])]\n",
        "mydat_trainr_list = [mydat_trainr_cat, mydat_trainr_num]\n",
        "mydat_rep_cat = mydat_rep[['gender', 'productType']]\n",
        "mydat_rep_num = mydat_rep.loc[:, ~mydat_rep.columns.isin(['gender', 'productType'])]\n",
        "mydat_rep_list = [mydat_rep_cat, mydat_rep_num]\n",
        "mydat_val_cat = mydat_val_emb[['gender', 'productType']]\n",
        "mydat_val_num = mydat_val_emb.loc[:, ~mydat_val_emb.columns.isin(['gender', 'productType'])]\n",
        "mydat_val_list = [mydat_val_cat, mydat_val_num]\n",
        "mydat_cat = mydat_emb[['gender', 'productType']]\n",
        "mydat_num = mydat_emb.loc[:, ~mydat_emb.columns.isin(['gender', 'productType'])]\n",
        "mydat_list = [mydat_cat, mydat_num]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qmbM0_7GBWmA"
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras import backend as K\n",
        "K.clear_session()\n",
        "set_random_seed(1)\n",
        "inputlist = []\n",
        "denselist = []\n",
        "nrep = len(mydat_rep)\n",
        "for i in range(nrep):\n",
        "  inputlist.append(Input((30, ), name = 'input'+str(i)))\n",
        "  #for each policy in the representative portfolio, calculate distance vector of target policy to the rep policy\n",
        "  denselist.append(Dense(1, name = 'dense'+str(i))(inputlist[i]))\n",
        "  #map each distance vector via an affine function to a single neuron. affine functions are distinct\n",
        "weights = Concatenate()(denselist)\n",
        "weights = Softmax()(weights)\n",
        "#concatenate the n_rep neurons, and then apply softmax to obtain weights\n",
        "output = Dense(1, use_bias = False, name = 'weightedval', trainable=False)(weights)\n",
        "#by introducing a dense layer on these weights, these weights will each be multiplied by a number and summed\n",
        "#we force the dense layer to be nontrainable, and set the number which the weights will be multiplied by\n",
        "#to be the representative portfolio values, in order to achieve this effect\n",
        "m = KerasModel(inputs = inputlist, outputs = output)\n",
        "m.get_layer('weightedval').set_weights([np.reshape(np.array(delta_rep), (len(delta_rep), 1))])\n",
        "#this forces the weights to each multiply by the representative portfolio values\n",
        "m.compile(optimizer=Adam(), loss='mse')\n",
        "#m.summary()\n",
        "callback = EarlyStopping(monitor='val_loss', patience=50)\n",
        "trainr_dist = getdist(mydat_rep_list, mydat_trainr_list)\n",
        "val_dist = getdist(mydat_rep_list, mydat_val_list)\n",
        "m.fit(x=trainr_dist, y=delta_trainr, epochs = 20000, validation_data = (val_dist, delta_val), callbacks = [callback], verbose = 1, batch_size = 16)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZBt7eIA5CUAx"
      },
      "outputs": [],
      "source": [
        "neval = 2048\n",
        "nfull = len(mydat_y)\n",
        "niters = nfull//neval +1\n",
        "predslist = []\n",
        "for i in range(niters):\n",
        "  start = int(neval*i)\n",
        "  stop = int(np.min([neval*(i+1), nfull]))\n",
        "  currentlist = [mydat_cat[start:stop], mydat_num[start:stop]]\n",
        "  currentdist = getdist(mydat_rep_list, currentlist)\n",
        "  predslist.append(m.predict(x=currentdist))\n",
        "mypreds_full = np.concatenate(predslist)\n",
        "mypreds_val = m.predict(x=val_dist)\n",
        "mypreds_trainr = m.predict(x=trainr_dist)\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "#plt.scatter(mydat_remainder_y, mypreds_remainder)\n",
        "plt.scatter(delta_full, mypreds_full)\n",
        "plt.scatter(delta_val, mypreds_val)\n",
        "plt.scatter(delta_trainr, mypreds_trainr)\n",
        "errs = mypreds_full.flatten() - delta_full\n",
        "errs_val = mypreds_val.flatten() - delta_val\n",
        "errs_trainr = mypreds_trainr.flatten() - delta_trainr\n",
        "mave_delta = np.mean(np.abs(delta_full))\n",
        "\n",
        "errs_zero = errs[delta_full != 0]\n",
        "errs_val_zero = errs_val[delta_val != 0]\n",
        "mave_delta_zero = np.mean(np.abs(delta_full[delta_full!=0]))\n",
        "\n",
        "mse_full = np.mean(np.square(errs))\n",
        "mae_full = np.mean(np.abs(errs))\n",
        "mse_val = np.mean(np.abs(errs_val))\n",
        "mse_trainr = np.mean(np.square(errs_trainr))\n",
        "\n",
        "mse_full_zero = np.mean(np.square(errs_zero))\n",
        "mae_full_zero = np.mean(np.abs(errs_zero))\n",
        "mse_val_zero = np.mean(np.square(errs_val_zero))\n",
        "\n",
        "print(mse_val/1e+6)\n",
        "print(mse_full/1e+6)\n",
        "print(mae_full/1e+6)\n",
        "print(np.sqrt(mse_full)/mave_delta)\n",
        "print(mae_full/mave_delta)\n",
        "\n",
        "print(mse_val_zero/1e+6)\n",
        "print(mse_full_zero/1e+6)\n",
        "print(mae_full_zero/1e+6)\n",
        "print(np.sqrt(mse_full_zero)/mave_delta_zero)\n",
        "print(mae_full_zero/mave_delta_zero)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "errs = mypreds_full.flatten() - delta_full\n",
        "errs_val = mypreds_val.flatten() - delta_val\n",
        "errs_trainr = mypreds_trainr.flatten() - delta_trainr\n",
        "mave_delta = np.mean(np.abs(delta_full))\n",
        "\n",
        "errs_zero = errs[delta_full != 0]\n",
        "errs_val_zero = errs_val[delta_val != 0]\n",
        "mave_delta_zero = np.mean(np.abs(delta_full[delta_full!=0]))\n",
        "\n",
        "mse_full = np.mean(np.square(errs))\n",
        "mae_full = np.mean(np.abs(errs))\n",
        "mse_val = np.mean(np.abs(errs_val))\n",
        "mse_trainr = np.mean(np.square(errs_trainr))\n",
        "\n",
        "mse_full_zero = np.mean(np.square(errs_zero))\n",
        "mae_full_zero = np.mean(np.abs(errs_zero))\n",
        "mse_val_zero = np.mean(np.square(errs_val_zero))\n",
        "\n",
        "\n",
        "print(mse_val/1e+6)\n",
        "print(mse_full/1e+6)\n",
        "print(mae_full/1e+6)\n",
        "print(np.sqrt(mse_full)/mave_delta)\n",
        "print(mae_full/mave_delta)\n",
        "\n",
        "print(mse_val_zero/1e+6)\n",
        "print(mse_full_zero/1e+6)\n",
        "print(mae_full_zero/1e+6)\n",
        "print(np.sqrt(mse_full_zero)/mave_delta_zero)\n",
        "print(mae_full_zero/mave_delta_zero)\n",
        "print(np.min(delta_rep))\n",
        "print(np.max(delta_rep))\n",
        "print(np.min(delta_full))\n",
        "print(np.max(delta_full))"
      ],
      "metadata": {
        "id": "KayNbPwjGD9m"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}