levels_ridge_regression_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import pyspark\n",
    "from pyspark import SparkContext\n",
    "from pyspark.sql import SparkSession\n",
    "from glow.levels.linear_model import RidgeReducer, RidgeRegression\n",
    "import pyspark.sql.functions as f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = SparkSession.builder.appName('levels').getOrCreate()\n",
    "spark.conf.set('spark.sql.execution.arrow.enabled', 'true')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_root = '/Users/leland.barnard/glow/glow-wgr/test-data/levels/ridge-regression' #path to glow levels test data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need three objects to get started:\n",
    "* A Spark DataFrame representing the block genotype matrix\n",
    "* A Spark DataFrame containing a mapping of sample block ID to corresponding list of sample ids\n",
    "* A Pandas DataFrame containing phenotypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "blockdf_lvl0 = spark.read.parquet(f'{test_data_root}/blockedGT.snappy.parquet') #block genotype matrix\n",
    "indexdf = spark.read.parquet(f'{test_data_root}/groupedIDs.snappy.parquet') #sample block ID to sample list mapping\n",
    "labeldf =  pd.read_csv(f'{test_data_root}/pts.csv').set_index('sample_id') #phenotype data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The block genotype matrix as a DataFrame\n",
    "If we imagine the block genotype matrix conceptually, we think of an *NxM* matrix *X* where each row *n* represents an individual sample, each column *m* represents a variant, and each cell *(n, m)* contains a genotype value for sample *n* at variant *m*.  We then imagine laying a coarse grid on top of this matrix such that matrix cells within the same coarse grid cell are all assigned to the same block *x*.  Each block *x* is indexed by a sample block ID (corresponding to a list of rows belonging to the block) and a header block ID (corresponding to a list of columns belonging to the block).  The sample block IDs are generally just integers 0 through the number of sample blocks.  The header block IDs are strings of the form 'chr_C_block_B', which refers to the Bth block on chromosome C.  The Spark DataFrame representing this block matrix can be thought of as the transpose of each block *xT* all stacked one atop another.  Each row represents the values from a particular column from *X*, for the samples corresponding to a particular sample block.  The fields in the DataFrame are:\n",
    "* header:  Corresponds to a column name in the conceptual matrix *X*.\n",
    "* size:  If the matrix is sparse (so that genotype values of 0 are implicit), the values for this header are represented as a sparse vector, and this column contains the size of that sparse vector.  Corresponds to the number of individuals in the sample block for the row.\n",
    "* indices:  (Optional, present of the matrix is sparse) Indices of the non-zero entries in the sparse vector for this header in this sample block.\n",
    "* values:  Genotype values for this header in this sample block.  If the matrix is sparse, contains only non-zero values.\n",
    "* header_block:  An ID assigned to the block *x* containing this header.\n",
    "* sample_block:  An ID assigned to the block *x* containing the group of samples represented on this row.\n",
    "* position:  An integer assigned to this header that specifies the correct sort order for the headers in this block.\n",
    "* mu:  The mean of the genotype calls for this header\n",
    "* sig:  The standard deviation of the genotype calls for this header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------+----+--------------------+--------------------+-------------+------------+---------+------------------+------------------+\n",
      "|         header|size|             indices|              values| header_block|sample_block| position|                mu|               sig|\n",
      "+---------------+----+--------------------+--------------------+-------------+------------+---------+------------------+------------------+\n",
      "|2:231414300:T:C|   9|     [0, 1, 6, 7, 8]|[1.0, 1.0, 1.0, 1...|chr_2_block_6|           7|231414300|0.8686868686868686|0.6730002176294544|\n",
      "|2:231414300:T:C|  10|[0, 1, 2, 3, 4, 5...|[2.0, 1.0, 1.0, 1...|chr_2_block_6|           1|231414300|0.8686868686868686|0.6730002176294544|\n",
      "|2:231414300:T:C|  12|[1, 3, 4, 5, 7, 8...|[2.0, 1.0, 1.0, 1...|chr_2_block_6|           8|231414300|0.8686868686868686|0.6730002176294544|\n",
      "|2:231414300:T:C|  13|[0, 1, 2, 3, 4, 5...|[2.0, 1.0, 2.0, 1...|chr_2_block_6|           9|231414300|0.8686868686868686|0.6730002176294544|\n",
      "+---------------+----+--------------------+--------------------+-------------+------------+---------+------------------+------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "blockdf_lvl0.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The sample block mapping\n",
    "This is a comparitively simple key-value store where each key is a sample block ID and each value is a list of sample IDs contained in that sample block.  As a Spark DataFrame, this is represented as a two column DataFrame with the following fields:\n",
    "* sample_block:  ID for a sample block\n",
    "* sample_ids: Array of sample IDs for the samples in this sample block.  The order of these IDs must match the order of the values arrays in the block genotype DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+--------------------+\n",
      "|sample_block|          sample_ids|\n",
      "+------------+--------------------+\n",
      "|           3|[1008962444, 1035...|\n",
      "|           9|[1083737921, 1041...|\n",
      "|           7|[1048623585, 1030...|\n",
      "|           1|[1073111137, 1082...|\n",
      "+------------+--------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "indexdf.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The phenotype data\n",
    "The phenotype data is represented as a Pandas DataFrame indexed by the sample ID.  Each column represents a single phenotype, and it is assumed that there are no missing phenotype values, and that the phenotypes mean centered at 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sim100</th>\n",
       "      <th>sim92</th>\n",
       "      <th>sim58</th>\n",
       "      <th>sim16</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample_id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1042204109</th>\n",
       "      <td>-0.905058</td>\n",
       "      <td>-1.171217</td>\n",
       "      <td>-1.437376</td>\n",
       "      <td>-1.703535</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1035505158</th>\n",
       "      <td>-0.616539</td>\n",
       "      <td>-0.411283</td>\n",
       "      <td>-0.206027</td>\n",
       "      <td>-0.000770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1008166305</th>\n",
       "      <td>-0.946014</td>\n",
       "      <td>-0.482639</td>\n",
       "      <td>-0.019263</td>\n",
       "      <td>0.444112</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1068805020</th>\n",
       "      <td>-1.155375</td>\n",
       "      <td>-0.660005</td>\n",
       "      <td>-0.164634</td>\n",
       "      <td>0.330736</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1095012035</th>\n",
       "      <td>-1.024889</td>\n",
       "      <td>-0.492179</td>\n",
       "      <td>0.040530</td>\n",
       "      <td>0.573240</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              sim100     sim92     sim58     sim16\n",
       "sample_id                                         \n",
       "1042204109 -0.905058 -1.171217 -1.437376 -1.703535\n",
       "1035505158 -0.616539 -0.411283 -0.206027 -0.000770\n",
       "1008166305 -0.946014 -0.482639 -0.019263  0.444112\n",
       "1068805020 -1.155375 -0.660005 -0.164634  0.330736\n",
       "1095012035 -1.024889 -0.492179  0.040530  0.573240"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labeldf.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reducer model fitting\n",
    "The first step in the fitting procedure is to apply a dimensionality reduction to the block matrix *X* using the `RidgeReducer`.  This is accomplished by fitting multiple ridge models within each block *x* and producing a new block matrix where each column represents the prediction of one ridge model applied within one block.  This approach to model building is generally referred to as **stacking**.  We will call the block genotype matrix we started with the **level 0** matrix in the stack *X0*, and the output of the ridge reduction step the **level 1** matrix *X1*.  The `RidgeReducer` class is used for this step, which is initiallized with a list of ridge regularization values (referred to here as alpha).  Since ridge models are indexed by these alpha values, the `RidgeReducer` will generate one ridge model per value of alpha provided, which in turn will produce one column per block in *X0*, so the final dimensions of matrix *X1* will be *Nx(LxK)*, where *L* is the number of header blocks in *X0* and *K* is the number of alpha values provided to the `RidgeReducer`.  In practice, we can estimate a span of alpha values in a reasonable order of magnitude based on guesses at the heritability of the phenotype we are fitting, but here we will just pick some values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "alphas_lvl0 = np.logspace(2, 5, 10)\n",
    "stack_lvl0 = RidgeReducer(alphas_lvl0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the `RidgeReducer` is initialized, it will assign names to the provided alphas and store them in a dict accessible as `RidgeReducer.alphas`.  This is mostly just to give an easily readable and sortable name to the models produced for each ridge value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'alpha_0': 100.0,\n",
       " 'alpha_1': 215.44346900318845,\n",
       " 'alpha_2': 464.15888336127773,\n",
       " 'alpha_3': 1000.0,\n",
       " 'alpha_4': 2154.4346900318824,\n",
       " 'alpha_5': 4641.588833612777,\n",
       " 'alpha_6': 10000.0,\n",
       " 'alpha_7': 21544.346900318822,\n",
       " 'alpha_8': 46415.888336127726,\n",
       " 'alpha_9': 100000.0}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stack_lvl0.alphas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `RidgeReducer.fit(blockdf, labeldf, indexdf)` method generates a Spark DataFrame representing the model that we can use to reduce *X0* to *X1*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "modeldf_lvl0 = stack_lvl0.fit(blockdf_lvl0, labeldf, indexdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In explicit terms, the reduction of a block *x0* from *X0* to the corresponding block *x1* from *X1* is accomplished by the matrix multiplication *x0 * B = x1*, where *B* is a coefficient matrix of size *mxK*, where *m* is the number of columns in block *x0* and *K* is the number of alpha values used in the reduction.  As an added wrinkle, if the ridge reduction is being performed against multiple phenotypes at once, each phenotype will have its own *B*, and for convenience we panel these next to each other in the output into a single matrix, so *B* in that case has dimensions *mx(K*P)* where *P* is the number of phenotypes.  Each matrix *B* is specific to a particular block in *X0*, so the Spark DataFrame produced by the `RidgeReducer` can be thought of all of as the matrices *B* from all of the blocks stacked one atop another.  The fields in the model DataFrame are:\n",
    "* header_block:  An ID assigned to the block *x0* corresponding to the coefficients in this row.\n",
    "* sample_block:  An ID assigned to the block *x0* corresponding to the coefficients in this row.\n",
    "* header:  The name of a column from the conceptual matrix *X0* that correspond with a particular row from the coefficient matrix *B*.\n",
    "* alphas:  List of alpha names corresponding to the columns of *B*.\n",
    "* labels:  List of label (i.e., phenotypes) corresponding to the columns of *B*. \n",
    "* coefficients:  List of the actual values from a row in *B*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------------+------------+---------------+---------+--------------------+--------------------+--------------------+\n",
      "| header_block|sample_block|         header| position|              alphas|              labels|        coefficients|\n",
      "+-------------+------------+---------------+---------+--------------------+--------------------+--------------------+\n",
      "|chr_3_block_8|           0|3:160741710:G:A|160741710|[alpha_0, alpha_1...|[sim100, sim100, ...|[0.07462677364336...|\n",
      "|chr_3_block_8|           0|3:175345110:C:T|175345110|[alpha_0, alpha_1...|[sim100, sim100, ...|[0.07834053929928...|\n",
      "|chr_3_block_8|           0|3:183469890:A:G|183469890|[alpha_0, alpha_1...|[sim100, sim100, ...|[0.02152237814164...|\n",
      "|chr_3_block_8|           0|3:195047160:C:T|195047160|[alpha_0, alpha_1...|[sim100, sim100, ...|[0.01153728383795...|\n",
      "+-------------+------------+---------------+---------+--------------------+--------------------+--------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "modeldf_lvl0.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reducer transformation\n",
    "After fitting, the `RidgeReducer.transform(blockdf, labeldf, modeldf)` method can be used to generate `X1` from `X0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "blockdf_lvl1 = stack_lvl0.transform(blockdf_lvl0, labeldf, modeldf_lvl0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the transformation is closely analogous to the block matrix DataFrame we started with.  The main difference is that, rather than representing a single block matrix, it really represents multiple block matrices, with one such matrix per label (phenotype).  Comparing the schema of this block matrix DataFrame (`blockdf_lvl1`) with the DataFrame we started with (`blockdf_lvl0`), the new columns are:\n",
    "* alpha:  This is the name of the alpha value used in fitting the model that produced the values in this row.\n",
    "* label:  This is the label corresponding to the values in this row.  Since the genotype block matrix *X0* is phenotype-agnostic, the rows in `blockdf_lvl0` were not restricted to any label/phenotype, but the level 1 block matrix *X1* represents ridge model predictions for the labels the reducer was fit with, so each row is associated with a specific label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+----+--------------------+------------+------------+--------+--------------------+-------------------+-------+------+\n",
      "|              header|size|              values|header_block|sample_block|position|                  mu|                sig|  alpha| label|\n",
      "+--------------------+----+--------------------+------------+------------+--------+--------------------+-------------------+-------+------+\n",
      "|chr_3_block_8_alp...|  13|[0.08337895454032...|       chr_3|           0|      80| 0.04148112816674154|0.19099426058493266|alpha_0|sim100|\n",
      "|chr_3_block_8_alp...|  13|[0.04796003873174...|       chr_3|           0|      81| 0.02402075708176127|0.11316256614620662|alpha_1|sim100|\n",
      "|chr_3_block_8_alp...|  13|[0.02504256254617...|       chr_3|           0|      82|0.012596289114544081|0.06030642726717367|alpha_2|sim100|\n",
      "|chr_3_block_8_alp...|  13|[0.01234023662311...|       chr_3|           0|      83|0.006221371128544...|0.03006776034645892|alpha_3|sim100|\n",
      "+--------------------+----+--------------------+------------+------------+--------+--------------------+-------------------+-------+------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "blockdf_lvl1.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The headers in the *X1* block matrix are derived from a combination of the source block in *X0*, the alpha value used in fitting the ridge model, and the label they were fit with.  These headers are assigned to header blocks that correspond to the chromosome of the source block in *X0*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------+------------+\n",
      "|header                            |header_block|\n",
      "+----------------------------------+------------+\n",
      "|chr_3_block_8_alpha_0_label_sim100|chr_3       |\n",
      "|chr_3_block_8_alpha_1_label_sim100|chr_3       |\n",
      "|chr_3_block_8_alpha_2_label_sim100|chr_3       |\n",
      "|chr_3_block_8_alpha_3_label_sim100|chr_3       |\n",
      "+----------------------------------+------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "blockdf_lvl1.select('header', 'header_block').show(4, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Regression fitting\n",
    "The block matrix *X1* can be used to fit a final predictive model that can generate phenotype predictions *y_hat* using the `RidgeRegression` class.  As with the `RidgeReducer` class, this class is initialized with a list of alpha values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "alphas_lvl1 = np.logspace(1, 4, 10)\n",
    "estimator_lvl1 = RidgeRegression(alphas_lvl1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "modeldf_lvl1_est, cvdf_lvl1 = estimator_lvl1.fit(blockdf_lvl1, labeldf, indexdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `RidgeRegression.fit(blockdf, labeldf, indexdf)` works in much the same way as the `RidgeReducer.fit(blockdf, labeldf, indexdf)` method, except that it returns two DataFrames:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A model DataFrame analogous to the model DataFrame provided by the `RidgeReducer`.  An important difference is that the header block ID for all rows will be 'all', indicating that all headers from all blocks have been used in a single fit, rather than fitting within blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+------------+--------------------+--------+--------------------+--------------------+--------------------+\n",
      "|header_block|sample_block|              header|position|              alphas|              labels|        coefficients|\n",
      "+------------+------------+--------------------+--------+--------------------+--------------------+--------------------+\n",
      "|         all|           1|chr_1_block_0_alp...|       0|[alpha_0, alpha_1...|[sim16, sim16, si...|[0.02787784249249...|\n",
      "|         all|           1|chr_2_block_0_alp...|       0|[alpha_0, alpha_1...|[sim16, sim16, si...|[-0.0164002560049...|\n",
      "|         all|           1|chr_3_block_0_alp...|       0|[alpha_0, alpha_1...|[sim16, sim16, si...|[-0.0234168451974...|\n",
      "|         all|           1|chr_1_block_0_alp...|       1|[alpha_0, alpha_1...|[sim16, sim16, si...|[0.00381390574280...|\n",
      "+------------+------------+--------------------+--------+--------------------+--------------------+--------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "modeldf_lvl1_est.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A cross validation (cv) report DataFrame, which reports the results of the hyperparameter (i.e., alpha) value optimization routine.\n",
    "* label:  This is the label corresponding to the cross cv results on the row.\n",
    "* alpha:  The name of the optimal alpha value\n",
    "* r2_mean:  The mean out of fold r2 score for the optimal alpha value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+-------+--------------------+\n",
      "| label|  alpha|             r2_mean|\n",
      "+------+-------+--------------------+\n",
      "| sim92|alpha_5| 0.18389799898047948|\n",
      "| sim16|alpha_8|-0.22499071350515992|\n",
      "| sim58|alpha_6|-0.02504464471643515|\n",
      "|sim100|alpha_5|  0.2566748993770534|\n",
      "+------+-------+--------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cvdf_lvl1.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Producing phenotype predictions *y_hat*\n",
    "After fitting the `RidgeRegression` model, the model DataFrame and cv DataFrame are used to apply the model to the block matrix DataFrame to produce predictions (*y_hat*) for each label in each sample block using the `RidgeRegression.transform(blockdf, labeldf, modeldf, cvdf)` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_hat_lvl1 = estimator_lvl1.transform(blockdf_lvl1, labeldf, modeldf_lvl1_est, cvdf_lvl1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting *y_hat* DataFrame has the following fields:\n",
    "* sample_block:  The sample block ID for the samples corresponding to the *y_hat* values on this row.\n",
    "* label:  The label corresponding to the *y_hat* values on this row\n",
    "* alpha:  The name of the alpha value used to fit the model that produced the *y_hat* values on this row.\n",
    "* values:  The array of *y_hat* values for the samples in the sample block for this row."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+-----+-------+--------------------+\n",
      "|sample_block|label|  alpha|              values|\n",
      "+------------+-----+-------+--------------------+\n",
      "|           1|sim16|alpha_8|[0.08461773658136...|\n",
      "|           4|sim16|alpha_8|[0.08343907935865...|\n",
      "|           7|sim16|alpha_8|[-0.0976335915514...|\n",
      "|           8|sim16|alpha_8|[-0.0461222342349...|\n",
      "+------------+-----+-------+--------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "y_hat_lvl1.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Fitting a second round of ridge reduction instead of ridge regression\n",
    "After fitting the first ridge reduction step and producing *X1* from *X0*, we can go directly into fitting the final ridge regression model, as we have just seen.  Alternatively, we can fit a second round of ridge reduction to squeeze *X1* into an even smaller feature matrix, which we will call the **level 2** matrix *X2*.  This has some advantages when it comes to generating the leave-one-chromosome-out versions of the *y_hat*s and does not come at much additional cost.  The procedure for fitting the second round of ridge reduction is identical to the first (we will reuse the same alphas we chose for the ridge regression fit above):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "stack_lvl1 = RidgeReducer(alphas_lvl1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "modeldf_lvl1 = stack_lvl1.fit(blockdf_lvl1, labeldf, indexdf)\n",
    "blockdf_lvl2 = stack_lvl1.transform(blockdf_lvl1, labeldf, modeldf_lvl1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **level 2** block matrix DataFrame produced here has an identical schema to the **level 1** block matrix.  A key difference is that the header block ID for all headers is now \"all\" for all headers, indicating that there are now no more blocks to collapse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+----+--------------------+------------+------------+--------+--------------------+-------------------+-------+-----+\n",
      "|              header|size|              values|header_block|sample_block|position|                  mu|                sig|  alpha|label|\n",
      "+--------------------+----+--------------------+------------+------------+--------+--------------------+-------------------+-------+-----+\n",
      "|all_block_1_alpha...|  13|[-0.0796642628265...|         all|           9|      10|-1.49453099468771...| 0.4138556129030118|alpha_0|sim16|\n",
      "|all_block_1_alpha...|  13|[-0.0957216101643...|         all|           9|      11|                 0.0|0.40982580962691184|alpha_1|sim16|\n",
      "|all_block_1_alpha...|  13|[-0.1070572367031...|         all|           9|      12|5.124106267500723...|0.40028561461234197|alpha_2|sim16|\n",
      "|all_block_1_alpha...|  13|[-0.1163419313886...|         all|           9|      13|1.708035422500241...| 0.3798792747100316|alpha_3|sim16|\n",
      "+--------------------+----+--------------------+------------+------------+--------+--------------------+-------------------+-------+-----+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "blockdf_lvl2.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The headers for each column now follow the name convention 'all_block_B_alpha_A_label_L', which refer to the ridge model prediction using alpha A and for label L fit using the features from header block B from block matrix *X1*.  Since the blocks in *X1* refer to chromosomes, the block number B here can be interpreted as a chromosome.  The 'all' token reflects the fact that we are not assigning the columns in *X2* to any new blocks (i.e, *X2* only has sample blocks, but there is only one header block which encompasses the entire matrix)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------------------------------+\n",
      "|header                         |\n",
      "+-------------------------------+\n",
      "|all_block_1_alpha_0_label_sim16|\n",
      "|all_block_1_alpha_1_label_sim16|\n",
      "|all_block_1_alpha_2_label_sim16|\n",
      "|all_block_1_alpha_3_label_sim16|\n",
      "+-------------------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "blockdf_lvl2.select('header').show(4, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now fit a ridge regression model as we did above, except that we will use the matrix *X2* instead of *X1*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "alphas_lvl2 = np.logspace(0, 3, 10)\n",
    "estimator_lvl2 = RidgeRegression(alphas_lvl2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "modeldf_lvl2_est, cvdf_lvl2 = estimator_lvl2.fit(blockdf_lvl2, labeldf, indexdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+------------+--------------------+--------+--------------------+--------------------+--------------------+\n",
      "|header_block|sample_block|              header|position|              alphas|              labels|        coefficients|\n",
      "+------------+------------+--------------------+--------+--------------------+--------------------+--------------------+\n",
      "|         all|           1|all_block_1_alpha...|      10|[alpha_0, alpha_1...|[sim16, sim16, si...|[-0.0939792871878...|\n",
      "|         all|           1|all_block_1_alpha...|      11|[alpha_0, alpha_1...|[sim16, sim16, si...|[-0.0788683759104...|\n",
      "|         all|           1|all_block_1_alpha...|      12|[alpha_0, alpha_1...|[sim16, sim16, si...|[-0.0693010949556...|\n",
      "|         all|           1|all_block_1_alpha...|      13|[alpha_0, alpha_1...|[sim16, sim16, si...|[-0.0446945065691...|\n",
      "+------------+------------+--------------------+--------+--------------------+--------------------+--------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "modeldf_lvl2_est.show(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+-------+--------------------+\n",
      "| label|  alpha|             r2_mean|\n",
      "+------+-------+--------------------+\n",
      "| sim92|alpha_7|   0.199251090828654|\n",
      "| sim16|alpha_9|-0.22903758326079596|\n",
      "| sim58|alpha_7|0.005461670993813417|\n",
      "|sim100|alpha_8|  0.2314559298409073|\n",
      "+------+-------+--------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cvdf_lvl2.show(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_hat_lvl2 = estimator_lvl2.transform(blockdf_lvl2, labeldf, modeldf_lvl2_est, cvdf_lvl2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+-----+-------+--------------------+\n",
      "|sample_block|label|  alpha|              values|\n",
      "+------------+-----+-------+--------------------+\n",
      "|           9|sim58|alpha_7|[-0.2126330471314...|\n",
      "|           6|sim58|alpha_7|[0.18042213283121...|\n",
      "|           5|sim58|alpha_7|[-0.0126226427178...|\n",
      "|           2|sim58|alpha_7|[0.00871975701462...|\n",
      "+------------+-----+-------+--------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "y_hat_lvl2.show(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For producing the LOCO versions of the *y_hat* vectors, it is only necessary to filter out rows from `blockdf_lvl2` corresponding to the chromosome we wish to drop before applying the transformation.  For example, if we wanted to produce *y_hat* with chromosome 1 left out (recall that the chromosomes constitute the source blocks for the headers in `blockdf_lvl2`, so headers from chromosome 1 will have headers like %block_1%):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_hat_lvl2_loco1 = estimator_lvl2.transform(blockdf_lvl2.filter(f'header NOT LIKE \"%block_1%\"'), labeldf, modeldf_lvl2_est, cvdf_lvl2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+-----+-------+--------------------+\n",
      "|sample_block|label|  alpha|              values|\n",
      "+------------+-----+-------+--------------------+\n",
      "|           9|sim58|alpha_7|[-0.1347024836295...|\n",
      "|           6|sim58|alpha_7|[0.20213653390706...|\n",
      "|           5|sim58|alpha_7|[-0.1602333580401...|\n",
      "|           2|sim58|alpha_7|[-0.1511874717623...|\n",
      "+------------+-----+-------+--------------------+\n",
      "only showing top 4 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "y_hat_lvl2_loco1.show(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "glow",
   "language": "python",
   "name": "glow"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}