main_results_for_public_repo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "from constants_and_util import *\n",
    "import image_processing\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.api as sm\n",
    "import pickle\n",
    "import train_models\n",
    "import seaborn as snsage\n",
    "import non_image_data_processing\n",
    "from image_processing import PytorchImagesDataset\n",
    "import analysis\n",
    "import pandas as pd\n",
    "from scipy.stats import spearmanr\n",
    "assert USE_HELD_OUT_TEST_SET"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load in model and data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "all_results = analysis.load_all_results(binary=False, \n",
    "                                        min_timestring='2019_06_20', \n",
    "                                        thing_to_filter_by={'experiment_to_run':'train_best_model_continuous', \n",
    "                                                             'crop_to_just_the_knee':False})\n",
    "\n",
    "pd.set_option('precision', 3)\n",
    "all_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble_results, ensemble_test_yhat = analysis.try_ensembling(\n",
    "    all_results, 5, binary_prediction=False)\n",
    "ensemble_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for i in range(1):\n",
    "    print(\"Model %i\" % i)\n",
    "    pytorch_model, dataloaders, datasets, dataset_sizes, yhats, dataset_kwargs = train_models.load_model_and_data_from_timestring(\n",
    "        all_results.iloc[i]['timestring'],\n",
    "        compute_yhats=True,\n",
    "        make_the_cam_plots=True, \n",
    "        make_the_prediction_change_plots=True)\n",
    "    \n",
    "    assert np.allclose(all_results.iloc[0]['test_r'], \n",
    "                   analysis.assess_performance(yhat=yhats['test_yhat'],\n",
    "                            y=datasets['test'].non_image_data['koos_pain_subscore'].values, \n",
    "                            binary_prediction=False)['r'])\n",
    "\n",
    "    assert np.allclose(all_results.iloc[0]['negative_test_rmse'], \n",
    "                   analysis.assess_performance(yhat=yhats['test_yhat'],\n",
    "                            y=datasets['test'].non_image_data['koos_pain_subscore'].values, \n",
    "                            binary_prediction=False)['negative_rmse'])\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Make ensemble prediction change plots. \n",
    "\n",
    "rng = random.Random(42)\n",
    "image_idxs_for_interpretability_plots = rng.sample(range(len(datasets['test'].non_image_data)), 8)\n",
    "print(\"Images making ensemble plots for\", image_idxs_for_interpretability_plots) # these should be same as for individual models. \n",
    "\n",
    "train_models.make_prediction_change_plots(list(all_results.iloc[:5]['timestring'].values), \n",
    "                             dataset_kwargs=dataset_kwargs, \n",
    "                             n_images_to_plot=len(image_idxs_for_interpretability_plots), \n",
    "                             figtitle='prediction_change_ensemble_top_5_models.png', \n",
    "                             img_idxs=image_idxs_for_interpretability_plots)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make descriptive statistics table. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "analysis.make_descriptive_stats_table(train_df=datasets['train'].non_image_data, \n",
    "                                      val_df=datasets['val'].non_image_data,\n",
    "                                      test_df=datasets['test'].non_image_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract relevant variables. \n",
    "\n",
    "dataset_name = 'test'\n",
    "yhat = ensemble_test_yhat\n",
    "y = datasets[dataset_name].non_image_data['koos_pain_subscore'].values\n",
    "klg = datasets[dataset_name].non_image_data['xrkl'].values\n",
    "all_ses_vars, income_at_least_50k, graduated_college, race_black = analysis.extract_all_ses_vars(datasets[dataset_name].non_image_data)\n",
    "discretized_yhat = analysis.discretize_yhat_like_kl_grade(yhat_arr=yhat,\n",
    "                                                          kl_grade_arr=klg,\n",
    "                                                          y_col='koos_pain_subscore')\n",
    "ids = datasets[dataset_name].non_image_data['id'].values\n",
    "decile_yhat = analysis.cut_into_deciles(yhat, y_col='koos_pain_subscore')\n",
    "binarized_y = binarize_koos(y)\n",
    "\n",
    "\n",
    "print('y mean median and std', np.mean(y), np.median(y), np.std(y, ddof=1))\n",
    "print('yhat mean median and std', np.mean(yhat), np.median(yhat), np.std(yhat, ddof=1))\n",
    "print('KLG mean median and std',  np.mean(klg), np.median(klg), np.std(klg, ddof=1))\n",
    "\n",
    "for ses_var in all_ses_vars:\n",
    "    ses_arr = all_ses_vars[ses_var]\n",
    "    print(\"disaggregate by %s\" % ses_var)\n",
    "    print('ses_var = 1: mean y %2.3f, mean KLG = %2.3f, mean yhat = %2.3f' % (y[ses_arr == 1].mean(), \n",
    "                                                                              klg[ses_arr == 1].mean(), \n",
    "                                                                              yhat[ses_arr == 1].mean()))\n",
    "    print('ses_var = 0: mean y %2.3f, mean KLG = %2.3f, mean yhat = %2.3f' % (y[ses_arr == 0].mean(), \n",
    "                                                                              klg[ses_arr == 0].mean(), \n",
    "                                                                              yhat[ses_arr == 0].mean()))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compare yhat's performance on predicting binarized high pain to KLG's. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\n",
    "yhat_binary_performance = analysis.assess_performance(yhat=-ensemble_test_yhat,\n",
    "                                                      y=binarized_y,\n",
    "                                                      binary_prediction=True,\n",
    "                                                      return_tpr_and_fpr=True)\n",
    "                                                        \n",
    "yhat_binary_performance['predictor'] = 'yhat'\n",
    "\n",
    "discretized_yhat_binary_performance = analysis.assess_performance(yhat=discretized_yhat,\n",
    "                                                                  y=binarized_y,\n",
    "                                                                  binary_prediction=True, \n",
    "                                                                  return_tpr_and_fpr=True)\n",
    "discretized_yhat_binary_performance['predictor'] = 'discretized_yhat'\n",
    "\n",
    "klg_binary_performance = analysis.assess_performance(yhat=klg,\n",
    "                                                     y=binarized_y,\n",
    "                                                     binary_prediction=True, \n",
    "                                                     return_tpr_and_fpr=True)\n",
    "klg_binary_performance['predictor'] = 'klg'\n",
    "\n",
    "print(pd.DataFrame([yhat_binary_performance, \n",
    "                    discretized_yhat_binary_performance,\n",
    "                    klg_binary_performance])[['predictor', 'auc', 'auprc']])\n",
    "\n",
    "# Make TPR/FPR plot. \n",
    "plt.figure(figsize=[4, 4])\n",
    "plt.plot(yhat_binary_performance['fpr'], yhat_binary_performance['tpr'], label='$\\hat y$')\n",
    "plt.plot(discretized_yhat_binary_performance['fpr'], \n",
    "         discretized_yhat_binary_performance['tpr'], label='Discretized $\\hat y$')\n",
    "plt.plot(klg_binary_performance['fpr'], klg_binary_performance['tpr'], label='KLG')\n",
    "plt.xlabel(\"FPR\")\n",
    "plt.ylabel(\"TPR\")\n",
    "plt.xlim([0, 1])\n",
    "plt.ylim([0, 1])\n",
    "plt.legend()\n",
    "plt.plot([0, 1], [0, 1], '--', color='black')\n",
    "plt.savefig('auc.png', dpi=300)\n",
    "\n",
    "print(\"Computing CIs on yhat binarized performance compared to KLG\")\n",
    "analysis.bootstrap_CIs_on_model_performance(y=binarized_y,\n",
    "                                            yhat=-yhat, \n",
    "                                            yhat_from_klg=klg, \n",
    "                                            yhat_from_clinical_image_features=None,\n",
    "                                            ids=ids, \n",
    "                                            binary_prediction=True,\n",
    "                                            n_bootstraps=N_BOOTSTRAPS)\n",
    "\n",
    "print(\"Computing CIs on discretized yhat binarized performance compared to KLG\")\n",
    "analysis.bootstrap_CIs_on_model_performance(y=binarized_y,\n",
    "                                            yhat=discretized_yhat, \n",
    "                                            yhat_from_klg=klg, \n",
    "                                            yhat_from_clinical_image_features=None,\n",
    "                                            ids=ids, \n",
    "                                            binary_prediction=True,\n",
    "                                            n_bootstraps=N_BOOTSTRAPS)\n",
    "                                            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis.make_violin_nonredundancy_plot(yhat=yhat, klg=klg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make simple histogram for presentation. Checked. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis.make_simple_histogram_of_pain(y=y, \n",
    "                              binary_vector_to_use=race_black, \n",
    "                              positive_class_label='Black patients', \n",
    "                              negative_class_label='Other patients', \n",
    "                              plot_filename='simple_pain_histogram_by_race.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Robustness check: confirm results look similar when using womac pain score rather than Koos pain score. Checked. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Correlation between Womac and Koos subscores on test dataset: %2.3f\" % \n",
    "      pearsonr(datasets['train'].non_image_data['koos_pain_subscore'].values,\n",
    "               datasets['train'].non_image_data['womac_pain_subscore'].values)[0])\n",
    "\n",
    "womac_yhat_from_klg = analysis.compare_to_clinical_performance(train_df=datasets['train'].non_image_data, \n",
    "                                val_df=datasets['val'].non_image_data, \n",
    "                                test_df=datasets['test'].non_image_data, \n",
    "                                y_col='womac_pain_subscore', \n",
    "                                features_to_use=['C(xrkl)'], \n",
    "                                binary_prediction=False, \n",
    "                                use_nonlinear_model=False, \n",
    "                                do_ols_sanity_check=True)\n",
    "\n",
    "womac_y = datasets['test'].non_image_data['womac_pain_subscore'].values\n",
    "\n",
    "print(\"Comparing performance (just in terms of R^2)\") # because yhat and womac_y aren't on same scale, so can't do RMSE. \n",
    "print('KLG R^2 for predicting Womac pain score: %2.3f' % pearsonr(womac_yhat_from_klg, womac_y)[0] ** 2)\n",
    "print('yhat R^2 for predicting Womac pain score: %2.3f' % pearsonr(yhat, womac_y)[0] ** 2)\n",
    "\n",
    "print('KLG SPEARMAN R^2 for predicting Womac pain score: %2.3f' % spearmanr(womac_yhat_from_klg, womac_y)[0] ** 2)\n",
    "print('yhat SPEARMAN R^2 for predicting Womac pain score: %2.3f' % spearmanr(yhat, womac_y)[0] ** 2)\n",
    "\n",
    "print(\"Comparing reductions in WOMAC pain gap\")\n",
    "analysis.quantify_pain_gap_reduction_vs_rival(yhat=yhat, \n",
    "                                              y=womac_y, \n",
    "                                              rival_severity_measure=womac_yhat_from_klg, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# BOOTSTRAPPED CIs for stratification plot. \n",
    "\n",
    "analysis.make_ses_stratification_plot(ses=graduated_college, \n",
    "                                     y=y, \n",
    "                                     dict_of_severity_scores={'KLG':klg, \n",
    "                                                              '$\\hat y$ decile':decile_yhat},\n",
    "                                    severity_score_order=['KLG', \n",
    "                                                          '$\\hat y$ decile'], \n",
    "                                      ses_var_one_label='College grad', \n",
    "                                      ses_var_zero_label='Not college grad', \n",
    "                                      fig_title='education_pain_gap.png', \n",
    "                                      n_bootstraps=N_BOOTSTRAPS, \n",
    "                                      ids=ids)\n",
    "\n",
    "analysis.make_ses_stratification_plot(ses=~(race_black == 1), \n",
    "                                     y=y, \n",
    "                                     dict_of_severity_scores={'KLG':klg, \n",
    "                                                              '$\\hat y$ decile':decile_yhat},\n",
    "                                    severity_score_order=['KLG', \n",
    "                                                          '$\\hat y$ decile'], \n",
    "                                      ses_var_one_label='Other patients', \n",
    "                                      ses_var_zero_label='Black patients', \n",
    "                                      fig_title='black_pain_gap.png', \n",
    "                                      n_bootstraps=N_BOOTSTRAPS, \n",
    "                                      ids=ids)\n",
    "\n",
    "analysis.make_ses_stratification_plot(ses=income_at_least_50k, \n",
    "                                     y=y, \n",
    "                                     dict_of_severity_scores={'KLG':klg, \n",
    "                                                              '$\\hat y$ decile':decile_yhat},\n",
    "                                    severity_score_order=['KLG', \n",
    "                                                          '$\\hat y$ decile'], \n",
    "                                      ses_var_one_label='Income >= 50k', \n",
    "                                      ses_var_zero_label='Income < 50k', \n",
    "                                      fig_title='income_pain_gap.png', \n",
    "                                      ids=ids, \n",
    "                                      n_bootstraps=N_BOOTSTRAPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare_two_severity_scores_on_one_plot. \n",
    "\n",
    "analysis.make_ses_stratification_plot(ses=~(race_black == 1), \n",
    "                                     y=y, \n",
    "                                     dict_of_severity_scores={'KLG':klg, \n",
    "                                                              'Discretized yhat':discretized_yhat},\n",
    "                                    severity_score_order=['KLG', \n",
    "                                                          'Discretized yhat'], \n",
    "                                      ses_var_one_label='Other patients', \n",
    "                                      ses_var_zero_label='Black patients', \n",
    "                                      compare_two_severity_scores_on_one_plot=True)\n",
    "\n",
    "analysis.make_ses_stratification_plot(ses=graduated_college, \n",
    "                                     y=y, \n",
    "                                     dict_of_severity_scores={'KLG':klg, \n",
    "                                                              'Discretized yhat':discretized_yhat},\n",
    "                                    severity_score_order=['KLG', \n",
    "                                                          'Discretized yhat'], \n",
    "                                      ses_var_one_label='College grad', \n",
    "                                      ses_var_zero_label='Not college grad', \n",
    "                                      compare_two_severity_scores_on_one_plot=True)\n",
    "\n",
    "analysis.make_ses_stratification_plot(ses=income_at_least_50k, \n",
    "                                     y=y, \n",
    "                                     dict_of_severity_scores={'KLG':klg, \n",
    "                                                              'Discretized yhat':discretized_yhat},\n",
    "                                    severity_score_order=['KLG', \n",
    "                                                          'Discretized yhat'], \n",
    "                                      ses_var_one_label='Income >= 50k', \n",
    "                                      ses_var_zero_label='Income < 50k', \n",
    "                                      compare_two_severity_scores_on_one_plot=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basically, before, we were fitting a predictive model to map KLG -> y, \n",
    "# then using that same model to map discretized_yhat -> y. \n",
    "# But this seems quite unfair to discretized_yhat because mean(y) might not be the same for the 5 groups. \n",
    "# So now instead we fit a linear model directly on discretized_yhat. \n",
    "# The only problem is, we only have that for the test set, which isn't ideal. \n",
    "# Ideally we would fit on the val set. I don't think this is a big deal, because overfitting should be minimal \n",
    "# (very few parameters to fit) and we outperform KLG regardless. \n",
    "\n",
    "discretized_yhat_df = pd.DataFrame({'koos_pain_subscore':datasets['test'].non_image_data['koos_pain_subscore'].values,\n",
    "                                    'xrkl':discretized_yhat})\n",
    "\n",
    "fit_discretized_yhat_model = sm.OLS.from_formula('koos_pain_subscore ~ C(discretized_yhat)', \n",
    "                                                data=discretized_yhat_df,\n",
    "                                                ).fit()\n",
    "\n",
    "pain_prediction_from_discretized_yhat = fit_discretized_yhat_model.predict(discretized_yhat_df).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict from max of joint space narrowing. \n",
    "\n",
    "for dataset in ['train', 'val', 'test']:\n",
    "    assert 'joint_space_narrowing_max' not in datasets[dataset].non_image_data.columns\n",
    "    datasets[dataset].non_image_data['joint_space_narrowing_max'] = np.maximum(datasets[dataset].non_image_data['xrjsm'].values,\n",
    "                                                                               datasets[dataset].non_image_data['xrjsl'].values\n",
    "                                                                              )\n",
    "    print(\"Proportion of JSN values in %s dataset\" % dataset)\n",
    "    print(100 * datasets[dataset].non_image_data['joint_space_narrowing_max'].value_counts()/len(datasets[dataset].non_image_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# rival predictors. \n",
    "print(\"\\n\\n****predicting y from KLG\")\n",
    "yhat_from_klg = analysis.compare_to_clinical_performance(train_df=datasets['train'].non_image_data, \n",
    "                                val_df=datasets['val'].non_image_data, \n",
    "                                test_df=datasets['test'].non_image_data, \n",
    "                                y_col='koos_pain_subscore', \n",
    "                                features_to_use=['C(xrkl)'], \n",
    "                                binary_prediction=False, \n",
    "                                use_nonlinear_model=False, \n",
    "                                do_ols_sanity_check=True)\n",
    "\n",
    "print(\"\\n\\n****predicting y from JSN max\")\n",
    "yhat_from_joint_space_narrowing_max = analysis.compare_to_clinical_performance(train_df=datasets['train'].non_image_data, \n",
    "                                val_df=datasets['val'].non_image_data, \n",
    "                                test_df=datasets['test'].non_image_data, \n",
    "                                y_col='koos_pain_subscore', \n",
    "                                features_to_use=['C(joint_space_narrowing_max)'], \n",
    "                                binary_prediction=False, \n",
    "                                use_nonlinear_model=False, \n",
    "                                do_ols_sanity_check=True)\n",
    "\n",
    "print(\"\\n\\n****predicting y from LINEAR MODEL using all clinical image features (including KLG)\")\n",
    "linear_yhat_from_clinical_image_features = analysis.compare_to_clinical_performance(train_df=datasets['train'].non_image_data, \n",
    "                                val_df=datasets['val'].non_image_data, \n",
    "                                test_df=datasets['test'].non_image_data, \n",
    "                                y_col='koos_pain_subscore', \n",
    "                                features_to_use=['C(%s)' % a for a in CLINICAL_CONTROL_COLUMNS], \n",
    "                                binary_prediction=False,\n",
    "                                use_nonlinear_model=False, \n",
    "                                do_ols_sanity_check=True)                                                   \n",
    "\n",
    "nonlinear_interactions = []\n",
    "for c1 in CLINICAL_CONTROL_COLUMNS:\n",
    "    for c2 in CLINICAL_CONTROL_COLUMNS:\n",
    "        if c1 > c2:\n",
    "            nonlinear_interactions.append('C(%s) * C(%s)' % (c1, c2))\n",
    "\n",
    "print(\"\\n\\n****predicting y from NONLINEAR MODEL using all clinical image features (including KLG)\")\n",
    "nonlinear_yhat_from_clinical_image_features = analysis.compare_to_clinical_performance(train_df=datasets['train'].non_image_data, \n",
    "                                val_df=datasets['val'].non_image_data, \n",
    "                                test_df=datasets['test'].non_image_data, \n",
    "                                y_col='koos_pain_subscore', \n",
    "                                features_to_use=nonlinear_interactions, \n",
    "                                binary_prediction=False,\n",
    "                                use_nonlinear_model=False, \n",
    "                                do_ols_sanity_check=False)\n",
    "\n",
    "print(\"\\n\\n****predicting y using RANDOM FOREST MODEL using all clinical image features (including KLG)\")\n",
    "random_forest_yhat_from_clinical_image_features = analysis.compare_to_clinical_performance(train_df=datasets['train'].non_image_data, \n",
    "                                val_df=datasets['val'].non_image_data, \n",
    "                                test_df=datasets['test'].non_image_data, \n",
    "                                y_col='koos_pain_subscore', \n",
    "                                features_to_use=['C(%s)' % a for a in CLINICAL_CONTROL_COLUMNS], \n",
    "                                binary_prediction=False,\n",
    "                                use_nonlinear_model=True, \n",
    "                                do_ols_sanity_check=False)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict from MRI features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_non_image_data_for_filtering_mri = non_image_data_processing.NonImageData('all',\n",
    "    i_promise_i_really_want_to_use_the_blinded_hold_out_set=True,                                                        \n",
    "    timepoints_to_filter_for=['12 month follow-up', \n",
    "                              '24 month follow-up', \n",
    "                              '36 month follow-up', \n",
    "                              '48 month follow-up', \n",
    "                              '00 month follow-up: Baseline'], \n",
    "                             filter_out_special_values_in_mri_data=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Do this analysis totally separately (for pain gap as well) \n",
    "# because it's on a different, narrower subset (since most people don't have MRIs)\n",
    "# As an additional robustness check, we repeat the analysis filtering out special values in MRI data just because \n",
    "# I'm not exactly sure the binarization procedure makes sense for these guys. \n",
    "\n",
    "df_for_filtering = original_non_image_data_for_filtering_mri.processed_dataframes['david_mri_data'][['id', 'side', 'visit']].copy()\n",
    "df_for_filtering['no_special_values'] = True\n",
    "\n",
    "mri_features = ['car11plusm', 'car11plusl', 'car11pluspf','bml2plusm', 'bml2plusl', 'bml2pluspf','mentearm', 'mentearl', 'menextm', 'menextl']\n",
    "\n",
    "for k in mri_features:\n",
    "    print(\"MRI feature %s\" % k)\n",
    "    print(\"Values are\")\n",
    "    print(datasets['train'].non_image_data[k].value_counts(dropna=False)/len(datasets['train']))\n",
    "\n",
    "for also_include_xray_features in [False, True]:\n",
    "    for use_random_forest in [False, True]:\n",
    "        for df_for_filtering_out_special_values in [None, df_for_filtering]:\n",
    "            analysis.compare_to_mri_features(datasets=datasets, \n",
    "                                 y=y, \n",
    "                                 yhat=yhat, \n",
    "                                 all_ses_vars=all_ses_vars, \n",
    "                                 ids=ids,\n",
    "                                 df_for_filtering_out_special_values=df_for_filtering_out_special_values, \n",
    "                                also_include_xray_features=also_include_xray_features, \n",
    "                                 use_random_forest=use_random_forest, \n",
    "                                mri_features=mri_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Robustness check -- how do clinical image features do on KLG >= 2. (We don't fill in missing data for KLG >= 2, so these features are the true values.) This is just a robustness check to make sure that it isn't our adding noise to the features which is causing our superior performance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for use_nonlinear_model in [False, True]:\n",
    "    print(\"\\n\\n\\n****Assessing image feature performance on KLG>=2 with nonlinear model=%s\" % use_nonlinear_model)\n",
    "    rival_predictor_just_on_klg_geq_2 = analysis.compare_to_clinical_performance(\n",
    "        train_df=datasets['train'].non_image_data.loc[datasets['train'].non_image_data['xrkl'] >= 2],\n",
    "        val_df=datasets['val'].non_image_data.loc[datasets['val'].non_image_data['xrkl'] >= 2],\n",
    "        test_df=datasets['test'].non_image_data.loc[datasets['test'].non_image_data['xrkl'] >= 2], \n",
    "                                    y_col='koos_pain_subscore', \n",
    "                                    features_to_use=['C(%s)' % a for a in CLINICAL_CONTROL_COLUMNS], \n",
    "                                    binary_prediction=False,\n",
    "                                    use_nonlinear_model=use_nonlinear_model, \n",
    "                                    do_ols_sanity_check=True) \n",
    "    \n",
    "    klg_geq_2_idxs = (datasets['test'].non_image_data['xrkl'] >= 2).values\n",
    "    print(\"\\nIn contrast, our performance is\")\n",
    "    our_performance_on_klg_geq_2 = analysis.assess_performance(y=y[klg_geq_2_idxs], \n",
    "                                yhat=yhat[klg_geq_2_idxs], \n",
    "                                binary_prediction=False)\n",
    "    for k in our_performance_on_klg_geq_2:\n",
    "        print(\"%s: %2.3f\" % (k, our_performance_on_klg_geq_2[k]))\n",
    "    \n",
    "    print(\"Do we still reduce the pain gap more?\")\n",
    "    all_ses_vars_just_on_klg_geq_2 = {}\n",
    "    for var in all_ses_vars:\n",
    "        all_ses_vars_just_on_klg_geq_2[var] = all_ses_vars[var][klg_geq_2_idxs]\n",
    "\n",
    "    print(analysis.quantify_pain_gap_reduction_vs_rival(yhat=yhat[klg_geq_2_idxs], \n",
    "                                              y=y[klg_geq_2_idxs], \n",
    "                                              rival_severity_measure=rival_predictor_just_on_klg_geq_2, \n",
    "                                              all_ses_vars=all_ses_vars_just_on_klg_geq_2, \n",
    "                                             ids=ids[klg_geq_2_idxs]))\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Small sanity check. ses gap looks roughly the same regardless of whether you do y ~ ses + C(klg) or whether you first predict y from klg and do y ~ ses + yhat_from_klg. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_sanity_check_df = pd.DataFrame({'klg':klg, 'y':y, 'yhat_from_klg':yhat_from_klg})\n",
    "\n",
    "# make sure r^2 is the same.\n",
    "just_compute_r2_david_method = sm.OLS.from_formula('y ~ C(klg)', data=small_sanity_check_df).fit().rsquared\n",
    "just_compute_r2_paper_method = sm.OLS.from_formula('y ~ yhat_from_klg', data=small_sanity_check_df).fit().rsquared\n",
    "rel_diff = 100 * np.abs((just_compute_r2_david_method - just_compute_r2_paper_method) / just_compute_r2_paper_method)\n",
    "print(\"R^2 for y ~ C(KLG) is %2.6f as compared to %2.6f in paper; rel diff %2.1f%%\" % (\n",
    "    just_compute_r2_david_method, just_compute_r2_paper_method, rel_diff))\n",
    "assert rel_diff < 1\n",
    "\n",
    "for ses_var in all_ses_vars:\n",
    "    small_sanity_check_df['ses'] = all_ses_vars[ses_var] * 1.\n",
    "    david_method = sm.OLS.from_formula('y ~ C(klg) + ses', data=small_sanity_check_df).fit()\n",
    "    paper_method = sm.OLS.from_formula('y ~ yhat_from_klg + ses', data=small_sanity_check_df).fit()\n",
    "    rel_diff = 100 * np.abs((david_method.params['ses'] - paper_method.params['ses']) / paper_method.params['ses'])\n",
    "    print('%s coefficient with y ~ C(klg) + ses: %2.6f; %2.6f with method in paper; diff %2.1f%%' % (\n",
    "        ses_var, david_method.params['ses'], paper_method.params['ses'], rel_diff))\n",
    "    assert rel_diff < 1\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KLG performance stratified by SES. Note that these results are just on the test set and as such are noisy for small groups of people. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_klg_performances_stratified_by_ses = []\n",
    "for ses_var in all_ses_vars:\n",
    "    for ses_var_val in [0, 1]:\n",
    "        ses_idxs = all_ses_vars[ses_var] == ses_var_val\n",
    "        klg_performance_ses_subgroup = analysis.assess_performance(y=y[ses_idxs], \n",
    "                                                        yhat=yhat_from_klg[ses_idxs], \n",
    "                                                        binary_prediction=False)\n",
    "        klg_performance_ses_subgroup['klg mean'] = klg[ses_idxs].mean()\n",
    "        klg_performance_ses_subgroup['klg std'] = klg[ses_idxs].std(ddof=1)\n",
    "        klg_performance_ses_subgroup['pain mean'] = y[ses_idxs].mean()\n",
    "        klg_performance_ses_subgroup['pain std'] = y[ses_idxs].std(ddof=1)\n",
    "        klg_performance_ses_subgroup['Subset'] = '%s=%i' % (ses_var, ses_var_val)\n",
    "        del klg_performance_ses_subgroup['negative_rmse']\n",
    "        all_klg_performances_stratified_by_ses.append(klg_performance_ses_subgroup)\n",
    "pd.DataFrame(all_klg_performances_stratified_by_ses)\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pain gaps controlling for all three variables at once. Note that this is run on the full dataset, so results will not quite line up with stuff run only on the test dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from statsmodels.iolib.summary2 import summary_col\n",
    "df_for_regression = pd.concat([datasets['train'].non_image_data, \n",
    "                              datasets['val'].non_image_data, \n",
    "                              datasets['test'].non_image_data])\n",
    "pain_gaps_controlling_for_all_three_at_once = []\n",
    "for specification in ['binarized_income_at_least_50k', \n",
    "                      'binarized_education_graduated_college', \n",
    "                      'race_black', \n",
    "                      'binarized_income_at_least_50k + binarized_education_graduated_college + race_black', \n",
    "                      'binarized_income_at_least_50k + C(xrkl)', \n",
    "                      'binarized_education_graduated_college + C(xrkl)',\n",
    "                      'race_black + C(xrkl)',\n",
    "                      'binarized_income_at_least_50k + binarized_education_graduated_college + race_black + C(xrkl)']:\n",
    "                      \n",
    "                    \n",
    "                      \n",
    "    all_ses_vars_model = sm.OLS.from_formula('koos_pain_subscore ~ %s' % specification, data=df_for_regression).fit()\n",
    "    all_ses_vars_model = all_ses_vars_model.get_robustcov_results(cov_type='cluster', \n",
    "                                                              groups=df_for_regression['id'].astype(int))\n",
    "    pain_gaps_controlling_for_all_three_at_once.append(all_ses_vars_model)\n",
    "summary_col(pain_gaps_controlling_for_all_three_at_once, \n",
    "            stars=True,\n",
    "            model_names=range(len(pain_gaps_controlling_for_all_three_at_once)),\n",
    "            regressor_order=['race_black', 'binarized_income_at_least_50k', 'binarized_education_graduated_college'])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show we get slightly better performance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "yhat_performance = analysis.assess_performance(y=y, yhat=yhat, binary_prediction=False)\n",
    "discretized_yhat_performance = analysis.assess_performance(y=y, yhat=pain_prediction_from_discretized_yhat, binary_prediction=False)\n",
    "klg_performance = analysis.assess_performance(y=y, yhat=yhat_from_klg, binary_prediction=False)\n",
    "joint_space_narrowing_max_performance = analysis.assess_performance(y=y, \n",
    "                                                                   yhat=yhat_from_joint_space_narrowing_max, \n",
    "                                                                   binary_prediction=False)\n",
    "all_clinical_performance_linear = analysis.assess_performance(y=y, \n",
    "                                                       yhat=linear_yhat_from_clinical_image_features, \n",
    "                                                       binary_prediction=False)\n",
    "all_clinical_performance_nonlinear = analysis.assess_performance(y=y, \n",
    "                                                       yhat=nonlinear_yhat_from_clinical_image_features, \n",
    "                                                       binary_prediction=False)\n",
    "\n",
    "all_clinical_performance_random_forest = analysis.assess_performance(y=y, \n",
    "                                                       yhat=random_forest_yhat_from_clinical_image_features, \n",
    "                                                       binary_prediction=False)\n",
    "\n",
    "\n",
    "yhat_performance['predictor'] = 'yhat'\n",
    "klg_performance['predictor'] = 'C(klg)'\n",
    "joint_space_narrowing_max_performance['predictor'] = 'C(joint_space_narrowing_max)'\n",
    "discretized_yhat_performance['predictor'] = 'discretized yhat'\n",
    "all_clinical_performance_linear['predictor'] = 'all clinical linear'\n",
    "all_clinical_performance_nonlinear['predictor'] = 'all clinical nonlinear'\n",
    "all_clinical_performance_random_forest['predictor'] = 'all clinical random forest'\n",
    "\n",
    "combined_results_df = pd.DataFrame([klg_performance, \n",
    "                    all_clinical_performance_linear, \n",
    "                    all_clinical_performance_nonlinear,\n",
    "                    all_clinical_performance_random_forest,\n",
    "                    joint_space_narrowing_max_performance,\n",
    "                    discretized_yhat_performance,\n",
    "                    yhat_performance])[['predictor', 'rmse', 'r', 'spearman_r', 'r^2', 'spearman_r^2']]\n",
    "\n",
    "combined_results_df\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# CIs. \n",
    "\n",
    "analysis.bootstrap_CIs_on_model_performance(y=y,\n",
    "                                            yhat=yhat, \n",
    "                                            yhat_from_klg=yhat_from_klg, \n",
    "                                            yhat_from_clinical_image_features=linear_yhat_from_clinical_image_features,\n",
    "                                            ids=ids, \n",
    "                                            n_bootstraps=N_BOOTSTRAPS)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write out DF for Ziad.\n",
    "\n",
    "df_for_ziad = pd.DataFrame({'patient_id':ids, \n",
    "                            'side':datasets['test'].non_image_data['side'].values, \n",
    "                            'timepoint':datasets['test'].non_image_data['visit'].values,\n",
    "                            'income_at_least_50k':income_at_least_50k, \n",
    "                            'graduated_college':graduated_college, \n",
    "                            'race_black':race_black,\n",
    "                            'koos_pain_subscore':y, \n",
    "                            'klg':klg, \n",
    "                            'klg_p':yhat_from_klg, \n",
    "                            'alg_p':yhat, \n",
    "                            'discretized_alg_p':discretized_yhat\n",
    "                           })\n",
    "\n",
    "df_for_ziad = df_for_ziad[['patient_id', 'side', 'timepoint', \n",
    "                           'income_at_least_50k', 'graduated_college', 'race_black',\n",
    "                           'klg', 'klg_p', \n",
    "                           'koos_pain_subscore','alg_p', 'discretized_alg_p']] \n",
    "assert pd.isnull(df_for_ziad).values.sum() == 0\n",
    "assert Counter(df_for_ziad['discretized_alg_p']) == Counter(df_for_ziad['klg'])\n",
    "assert len(df_for_ziad) == 11320\n",
    "assert len(df_for_ziad.drop_duplicates(['patient_id', 'timepoint', 'side'])) == 11320\n",
    "assert len(df_for_ziad.drop_duplicates('patient_id')) == 1295\n",
    "assert len(df_for_ziad.drop_duplicates(['patient_id', 'graduated_college', 'income_at_least_50k', 'race_black'])) == 1295\n",
    "assert len(df_for_ziad.drop_duplicates(['klg', 'klg_p'])) == 5\n",
    "df_for_ziad.to_csv('df_for_ziad.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show we reduce pain gap. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare to reduction from KLG. \n",
    "analysis.quantify_pain_gap_reduction_vs_rival(yhat=yhat, \n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=yhat_from_klg, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Omitted variable bias decomposition. \n",
    "\n",
    "### \"Short equals long plus the effect of omitted times the regression of omitted on included.\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for ses_var in all_ses_vars:\n",
    "    all_ovb_results = []\n",
    "    print(\"\\n\\nSES var: %s\" % ses_var)\n",
    "    for severity_score in ['yhat', 'klg']:\n",
    "        if severity_score == 'yhat':\n",
    "            severity_score_to_use = yhat\n",
    "        else:\n",
    "            severity_score_to_use = yhat_from_klg\n",
    "        severity_score_to_use = copy.deepcopy(severity_score_to_use)\n",
    "        zscored_severity_score = (severity_score_to_use - severity_score_to_use.mean()) / severity_score_to_use.std(ddof=1)\n",
    "        df_for_ovb_regression = pd.DataFrame({'koos_pain_subscore':y, \n",
    "                      'ses':all_ses_vars[ses_var]* 1., \n",
    "                      'severity_score':zscored_severity_score})\n",
    "\n",
    "        short_reg = sm.OLS.from_formula('koos_pain_subscore ~ ses', data=df_for_ovb_regression).fit()\n",
    "        long_reg = sm.OLS.from_formula('koos_pain_subscore ~ ses + severity_score', data=df_for_ovb_regression).fit()\n",
    "        omitted_on_included = sm.OLS.from_formula('severity_score ~ ses', data=df_for_ovb_regression).fit()\n",
    "        all_ovb_results.append({'reduction_in_pain_gap':short_reg.params['ses'] - long_reg.params['ses'], \n",
    "                                'effect_of_omitted':long_reg.params['severity_score'], \n",
    "                                'omitted_on_included':omitted_on_included.params['ses'], \n",
    "                                'severity_score':severity_score})\n",
    "        assert np.allclose(short_reg.params['ses'] - long_reg.params['ses'], \n",
    "                           long_reg.params['severity_score'] * omitted_on_included.params['ses'])\n",
    "    all_ovb_results = pd.DataFrame(all_ovb_results)\n",
    "    all_ovb_results.index = all_ovb_results['severity_score']\n",
    "    del all_ovb_results['severity_score']\n",
    "    \n",
    "    all_ovb_results.loc['RATIO'] = all_ovb_results.iloc[0] / all_ovb_results.iloc[1]\n",
    "    print(all_ovb_results)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Bootstrap CIs on pain gap. \n",
    "\n",
    "analysis.bootstrap_CIs_on_pain_gap_reduction(y=y, \n",
    "                                             yhat=yhat, \n",
    "                                             yhat_from_klg=yhat_from_klg, \n",
    "                                             ids=ids, \n",
    "                                             all_ses_vars=all_ses_vars, \n",
    "                                             n_bootstraps=N_BOOTSTRAPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare to reduction from using all clinical features (linear model)\n",
    "\n",
    "analysis.quantify_pain_gap_reduction_vs_rival(yhat=yhat, \n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=linear_yhat_from_clinical_image_features, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare to reduction from using all clinical features (nonlinear model)\n",
    "analysis.quantify_pain_gap_reduction_vs_rival(yhat=yhat, \n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=nonlinear_yhat_from_clinical_image_features, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare to reduction from using joint space narrowing max. \n",
    "analysis.quantify_pain_gap_reduction_vs_rival(yhat=yhat, \n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=yhat_from_joint_space_narrowing_max, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare to reduction from using all clincial features + random forest (nonlinear model)\n",
    "analysis.quantify_pain_gap_reduction_vs_rival(yhat=yhat, \n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=random_forest_yhat_from_clinical_image_features, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare discretized yhat to reduction from KLG. \n",
    "\n",
    "analysis.quantify_pain_gap_reduction_vs_rival(yhat=pain_prediction_from_discretized_yhat, \n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=yhat_from_klg, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check main results hold when controlling for things. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "From our email chain, covariates are:\n",
    "Recruitment site\n",
    "Age*sex\n",
    "Timepoint\n",
    "Side\n",
    "Current BMI\n",
    "Max BMI\n",
    "Past knee injury\n",
    "Past knee surgery\n",
    "Current / former smoker\n",
    "Drinking status\n",
    "Marital status\n",
    "\n",
    "Note: we sometimes get \"condition number is large\" warnings when we do regressions of y ~ yhat. \n",
    "This does not appear to indicate a convergence problem. Rather, it's just that y and yhat are on a large scale (0-100). \n",
    "If you just divide both variables by 100, the coefficients (of course) remain identical and the warnings go away. \n",
    "\"\"\"\n",
    "all_comparisons = []\n",
    "for rival_name in ['klg', 'all_clinical_features']:\n",
    "    for control_set in [['C(v00site)', \n",
    "                         'C(age_at_visit)*C(p02sex)',\n",
    "                         'C(visit)', \n",
    "                         'side', \n",
    "                          'C(max_bmi)', \n",
    "                         'C(current_bmi)',\n",
    "                        'C(knee_surgery)', \n",
    "                         'C(knee_injury)', \n",
    "                         'C(cigarette_smoker)', \n",
    "                         'C(drinks_per_week)', \n",
    "                         'C(v00maritst)'] ,\n",
    "                        ['C(age_at_visit)*C(p02sex)'], \n",
    "                        ['C(age_at_visit)*C(p02sex)', 'side', 'C(visit)', 'C(v00site)']]:\n",
    "        if rival_name == 'klg':\n",
    "            rival_predictor = yhat_from_klg\n",
    "        else:\n",
    "            rival_predictor = linear_yhat_from_clinical_image_features\n",
    "        print(\"\\n\\n\\n\\n*************** COMPARING %s to yhat controlling for\" % rival_name)\n",
    "        print(control_set)\n",
    "        comparison = analysis.check_main_results_hold_when_controlling_for_things(df=datasets['test'].non_image_data, \n",
    "                                                             yhat=yhat, \n",
    "                                                             rival_predictor=rival_predictor, \n",
    "                                                             rival_name=rival_name,\n",
    "                                                             all_controls=control_set)\n",
    "        all_comparisons.append(comparison)\n",
    "\n",
    "pd.set_option('max_colwidth', 500)\n",
    "pd.concat(all_comparisons)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Additional check: we aren't predicting y merely by predicting SES. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "print(\"Check yhat coefficient is still coefficient when controlling for xrkl and SES\")\n",
    "\n",
    "df_to_test = pd.DataFrame({'y':y, \n",
    "                               'yhat_from_klg':yhat_from_klg,\n",
    "                               'yhat':yhat, \n",
    "                               'klg':klg,\n",
    "                               'id':ids})\n",
    "for ses_var in all_ses_vars:\n",
    "    df_to_test[ses_var] = all_ses_vars[ses_var]\n",
    "    print(\"SES var: %s\" % ses_var)\n",
    "    \n",
    "    print('R^2 for y ~ yhat + SES: %2.3f; for y ~ KLG + ses: %2.3f' % \n",
    "          (\n",
    "              sm.OLS.from_formula('y ~ yhat + %s' % ses_var, data=df_to_test).fit().rsquared, \n",
    "              sm.OLS.from_formula('y ~ yhat_from_klg + %s' % ses_var, data=df_to_test).fit().rsquared\n",
    "          )\n",
    "         )\n",
    "\n",
    "    ols_model = sm.OLS.from_formula('y ~ yhat + %s + yhat_from_klg' % ses_var, data=df_to_test).fit(\n",
    "        cov_type='cluster', cov_kwds={'groups':df_to_test['id']})\n",
    "    print('y ~ yhat + SES + KLG: yhat coefficient %2.3f, p=%2.3e' % (ols_model.params['yhat'], ols_model.pvalues['yhat']))\n",
    "\n",
    "print(\"yhat Coefficient when not controlling for any SES vars\")\n",
    "control_for_no_ses_vars = sm.OLS.from_formula('y ~ yhat', data=df_to_test).fit(cov_type='cluster', cov_kwds={'groups':df_to_test['id']})\n",
    "print(control_for_no_ses_vars.summary())\n",
    "print(\"yhat Coefficient when not controlling for all SES vars\")\n",
    "control_for_all_three_ses_vars = sm.OLS.from_formula('y ~ yhat + race_black + income_less_than_50k + did_not_graduate_college', \n",
    "                                                     data=df_to_test).fit(cov_type='cluster', cov_kwds={'groups':df_to_test['id']})\n",
    "print(control_for_all_three_ses_vars.summary())\n",
    "\n",
    "# Relatedly what is coefficient of KLG when controlling for yhat (don't need this in for loop because not using SES var)\n",
    "print(\"**\\n\\nIs KLG significant when controlling for yhat\")\n",
    "\n",
    "ols_model = sm.OLS.from_formula('y ~ yhat + yhat_from_klg', data=df_to_test).fit(\n",
    "    cov_type='cluster', cov_kwds={'groups':df_to_test['id']})\n",
    "print(ols_model.summary())\n",
    "\n",
    "# Try doing an f-test with categorical dummies. \n",
    "# Note that an f-test is not valid here when we cluster; there's no way to account for the clustering. \n",
    "# So I don't think the f-test is the way to do it. \n",
    "restricted_model_unclustered = sm.OLS.from_formula('y ~ yhat', data=df_to_test).fit()\n",
    "full_model_unclustered = sm.OLS.from_formula('y ~ yhat + C(klg)', data=df_to_test).fit()\n",
    "\n",
    "print(full_model_unclustered.compare_f_test(restricted_model_unclustered))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# another check: check that we outperform across a whole bunch of subgroups. \n",
    "analysis.stratify_performances(df=datasets['test'].non_image_data, \n",
    "                               yhat=yhat, \n",
    "                               y=y, \n",
    "                               yhat_from_klg=yhat_from_klg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Robustness check -- do we generalize across sites?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_site_generalization_results = analysis.load_all_results(binary=False, \n",
    "                                        min_timestring='2019_06_20', \n",
    "                                        thing_to_filter_by={'experiment_to_run':'hold_out_one_imaging_site'})\n",
    "\n",
    "all_site_generalization_results['site_to_remove'] = all_site_generalization_results['hold_out_one_imaging_site_kwargs'].map(\n",
    "    lambda x:x['site_to_remove'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check: does computing the val loss on the held out site make a difference in how models are ranked\n",
    "# either in terms of choosing the best epoch or the best model? \n",
    "# doesn't seem to, no. \n",
    "\n",
    "best_results_by_model = []\n",
    "for i in range(len(all_site_generalization_results)):\n",
    "    # first compare within models. \n",
    "    val_results_for_model = all_site_generalization_results['val_results_stratified_by_site'].iloc[i]\n",
    "    site_excluded = all_site_generalization_results['site_to_remove'].iloc[i]\n",
    "    epochs = sorted(val_results_for_model.keys())\n",
    "    assert epochs == list(range(15))\n",
    "    site_excluded_results = [val_results_for_model[a]['every_site_but_%s' % site_excluded]['negative_rmse']\n",
    "                             for a in epochs]\n",
    "    full_val_set_results = [val_results_for_model[a]['val_negative_rmse']\n",
    "                             for a in epochs]\n",
    "    best_results_by_model.append({'site_excluded_results':max(site_excluded_results), \n",
    "                                  'full_val_set_results':max(full_val_set_results), \n",
    "                                  'site_excluded':site_excluded})\n",
    "    plt.figure()\n",
    "    plt.title(\"Each point is one EPOCH for a single model site %s\\nspearman r %2.3f\" % \n",
    "              (site_excluded, spearmanr(full_val_set_results, site_excluded_results)[0]))\n",
    "    plt.xlabel(\"Full val set RMSE\")\n",
    "    plt.ylabel(\"RMSE excluding site %s\" % site_excluded)\n",
    "    plt.scatter(full_val_set_results, site_excluded_results)\n",
    "\n",
    "# then compare across models. \n",
    "best_results_by_model = pd.DataFrame(best_results_by_model)\n",
    "for site_excluded in sorted(list(set(best_results_by_model['site_excluded']))):\n",
    "    plt.figure()\n",
    "    plt.title(\"Each point is one MODEL for site %s\" % site_excluded)\n",
    "    plt.scatter(best_results_by_model.loc[best_results_by_model['site_excluded'] == site_excluded, 'full_val_set_results'],\n",
    "                best_results_by_model.loc[best_results_by_model['site_excluded'] == site_excluded, 'site_excluded_results'])\n",
    "    plt.xlabel(\"Full val set RMSE\")\n",
    "    plt.ylabel(\"RMSE excluding site %s\" % site_excluded)\n",
    "    plt.show()\n",
    "    \n",
    "assert np.allclose(list(best_results_by_model['full_val_set_results'].values), \n",
    "            sorted(best_results_by_model['full_val_set_results'].values)[::-1])\n",
    "assert((all_site_generalization_results['best_val_negative_rmse'] \n",
    "       == best_results_by_model['full_val_set_results'].values).all())\n",
    "all_site_generalization_results['exclude_held_out_site_best_val_negative_rmse'] = best_results_by_model['site_excluded_results'].values\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# We win consistently on pain gap reduction across all sites. \n",
    "# r^2 and spearman r^2 are generally better. \n",
    "# But yhat is somewhat miscalibrated for the new site (r^2 is good, but RMSE is less so). \n",
    "# You can fix this by recalibrating both yhat and KLG for each site\n",
    "# Recalibrating does not affect pain gap reduction or r^2. \n",
    "# Recalibrating seems a bit sketchy, though, so we don't do it in the results we report. \n",
    "\n",
    "key_to_sort_by = 'best_val_negative_rmse'#'exclude_held_out_site_best_val_negative_rmse' # 'best_val_negative_rmse'\n",
    "all_site_generalization_results = all_site_generalization_results.sort_values(by=key_to_sort_by)[::-1]\n",
    "\n",
    "for recalibrate_to_new_set in [True, False]:\n",
    "    print(\"\\n\\n***********Recalibrating to new set: %s\" % recalibrate_to_new_set)\n",
    "    analysis.analyze_performance_on_held_out_sites(all_site_generalization_results, \n",
    "                                               yhat=yhat, \n",
    "                                               y=y, \n",
    "                                               yhat_from_klg=yhat_from_klg, \n",
    "                                               site_vector=datasets['test'].non_image_data['v00site'].values, \n",
    "                                               all_ses_vars=all_ses_vars, \n",
    "                                               ids=ids, \n",
    "                                               recalibrate_to_new_set=recalibrate_to_new_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Robustness check: dataset where there is no correlation between pain and race/SES. I ultimately favored this robustness check over the one where we just train one race/SES group because that hurts predictive performance for reasons unrelated to what we're trying to assess (ie, we halve the training set size). (This robustness check does not end up in the paper.) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "all_ses_decorrelated_results = analysis.load_all_results(binary=False, \n",
    "                                        min_timestring='2019_06_20', \n",
    "                                        thing_to_filter_by={'experiment_to_run':'remove_correlation_between_pain_and_ses'})\n",
    "\n",
    "pd.set_option('precision', 3)\n",
    "all_ses_decorrelated_results['ses_decorrelation_col'] = all_ses_decorrelated_results['remove_correlation_between_pain_and_ses_kwargs'].map(lambda x:x['ses_col'])\n",
    "all_ses_decorrelated_results[['ses_decorrelation_col', 'test_r', 'negative_test_rmse']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Confirm that we still narrow the pain gap more than KLG and we still have better predictive performance than KLG. \n",
    "all_performances = []\n",
    "for ses_col in sorted(list(set(all_ses_decorrelated_results['ses_decorrelation_col']))):\n",
    "    print('\\n\\n***Robustness check for dataset where pain and %s are decorrelated' % ses_col)\n",
    "    _, ses_decorrelated_yhat = analysis.try_ensembling(\n",
    "        all_ses_decorrelated_results.loc[all_ses_decorrelated_results['ses_decorrelation_col'] == ses_col],\n",
    "        5,\n",
    "        binary_prediction=False)\n",
    "    performance = analysis.assess_performance(y=y, \n",
    "                                              yhat=ses_decorrelated_yhat, \n",
    "                                              binary_prediction=False)\n",
    "    performance['ses_decorrelation_col'] = ses_col\n",
    "    all_performances.append(performance)\n",
    "    print(\"Pain gap reduction\")\n",
    "    print(analysis.quantify_pain_gap_reduction_vs_rival(yhat=ses_decorrelated_yhat, \n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=yhat_from_klg, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids))\n",
    "    \n",
    "print(\"All performance\")\n",
    "klg_performance = analysis.assess_performance(y=y, yhat=yhat_from_klg, binary_prediction=False)\n",
    "klg_performance['ses_decorrelation_col'] = 'KLG BASELINE'\n",
    "all_performances.append(klg_performance)\n",
    "pd.DataFrame(all_performances)[['ses_decorrelation_col', 'rmse', 'r', 'spearman_r', 'r^2', 'spearman_r^2']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Followup predictions: does yhat help predict future KLG or pain? We try multiple specifications of this question. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.set_option(\"display.precision\", 5)\n",
    "pd.set_option('max_colwidth', 500)\n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data, yhat=yhat)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare to KLG.\n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=yhat_from_klg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# effect of discretized_yhat=4 on continuous pain score + KLG. \n",
    "\n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=(discretized_yhat == 4) * 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# effect of discretized_yhat>=2 on continuous pain score + KLG. \n",
    "\n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=(discretized_yhat >= 2) * 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# effect of one std increase in yhat on odds of being high pain. \n",
    "\n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=-(yhat - yhat.mean()) / yhat.std(ddof=1), \n",
    "                                         use_binary_pain=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# effect of discretized_yhat=4 on odds of being high pain. \n",
    "\n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=(discretized_yhat >= 4) * 1., \n",
    "                                         use_binary_pain=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# effect of discretized_yhat>=2 on odds of being high pain. \n",
    "\n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=(discretized_yhat >= 2) * 1., \n",
    "                                         use_binary_pain=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# again, yhat doesn't really have better predictive power for pain than KLG at followup when you control for pain at baseline.\n",
    "# (although it does when you don't control for KLG). \n",
    "# This makes sense because yhat is more strongly correlated with pain than is KLG, so it's not clear it should add more than KLG should over pain.  \n",
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=(klg >= 2) * 1., \n",
    "                                         use_binary_pain=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show we can predict KLG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# comparison to previous results:\n",
    "# 1. Automatic Knee Osteoarthritis Diagnosis from Plain Radiographs: A Deep Learning-Based Approach\n",
    "# Sophisticated model: manual cropping, siamese network etc. \n",
    "# We trained our method using the data solely from the Multicenter Osteoarthritis Study \n",
    "# and validated it on randomly selected 3,000 subjects (5,960 knees) from Osteoarthritis Initiative dataset (just baseline data)\n",
    "# The classification MSE value achieved was 0.48, which is lower than previously published results...\n",
    "# we also evaluated a fine-tuned ResNet-34 network because it performed similarly on the validation set. \n",
    "# On the test set, the baseline also performed similarly to our approach in terms of MSE (value of 0.51)\n",
    "# More sophisticated model, probably a harder prediction task. \n",
    "\n",
    "# Quantifying Radiographic Knee Osteoarthritis Severity using Deep Convolutional Neural Networks\n",
    "# MSE: 0.504\n",
    "# The data used for the experiments are bilateral PA fixed flexion knee X-ray images, \n",
    "# taken from the baseline (image release version O.E.1) radiographs of the Osteoarthritis Initiative (OAI) dataset \n",
    "# containing an entire cohort of 4,476 participants\n",
    "\n",
    "all_klg_results = analysis.load_all_results(binary=False, \n",
    "                                        min_timestring='2019_06_20', \n",
    "                                        thing_to_filter_by={'experiment_to_run':'predict_klg', \n",
    "                                                             'crop_to_just_the_knee':False})\n",
    "\n",
    "ensemble_results_klg_prediction, ensemble_prediction_klg_hat = analysis.try_ensembling(all_klg_results, 5, binary_prediction=False)\n",
    "\n",
    "pd.set_option('precision', 3)\n",
    "ensemble_results_klg_prediction['MSE'] = ensemble_results_klg_prediction['negative_rmse'] ** 2\n",
    "ensemble_results_klg_prediction\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare to baseline visit to get more precise comparison to previous work (which used baseline visit). \n",
    "true_klg = all_klg_results.iloc[0]['test_y']\n",
    "baseline_idxs = (datasets['test'].non_image_data['visit'].map(lambda x:'Baseline' in x).values)\n",
    "klg_prediction_results = analysis.assess_performance(y=true_klg[baseline_idxs],\n",
    "                            yhat=ensemble_prediction_klg_hat[baseline_idxs], \n",
    "                            binary_prediction=False)\n",
    "\n",
    "print(\"MSE just on baseline visits: %2.3f; r %2.3f\" % (klg_prediction_results['negative_rmse'] ** 2, \n",
    "                                                       klg_prediction_results['r']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# how well does this predicted KLG predict pain?\n",
    "\n",
    "print(\"Assessing how KLG_hat does compared to KLG\")\n",
    "# Note this might be a little generous to KLG_hat, because we just assess fit on the test set \n",
    "# but I don't think the problem should be very big because the number of free parameters is very small. \n",
    "# (The reason we do this is that we lack a train set KLG_hat to fit on, due to overfitting.) \n",
    "\n",
    "discretized_klg_hat = analysis.discretize_yhat_like_kl_grade(\n",
    "    yhat_arr=-ensemble_prediction_klg_hat,\n",
    "    kl_grade_arr=klg,\n",
    "    y_col='koos_pain_subscore')\n",
    "\n",
    "print(pd.DataFrame({'klg':klg,'discretized_klg_hat':discretized_klg_hat})\n",
    "      .groupby(['klg', 'discretized_klg_hat'])\n",
    "      .size()/len(klg))\n",
    "\n",
    "klg_to_predict_pain_df = pd.DataFrame({'klg_hat':ensemble_prediction_klg_hat, \n",
    "                                      'y':y, \n",
    "                                       'discretized_klg_hat':discretized_klg_hat})\n",
    "\n",
    "pain_predicted_from_klg_hat_model = sm.OLS.from_formula('y ~ klg_hat', data=klg_to_predict_pain_df).fit()\n",
    "pain_predicted_from_discretized_klg_hat_model = sm.OLS.from_formula('y ~ C(discretized_klg_hat)', data=klg_to_predict_pain_df).fit()\n",
    "\n",
    "all_predictors_to_compare = {'klg':yhat_from_klg, \n",
    "                             'klg_hat':pain_predicted_from_klg_hat_model.predict(klg_to_predict_pain_df).values, \n",
    "                             'discretized_klg_hat':pain_predicted_from_discretized_klg_hat_model.predict(klg_to_predict_pain_df).values}\n",
    "all_klg_versions_performance = []\n",
    "for predictor_name in ['klg', 'klg_hat', 'discretized_klg_hat']:\n",
    "    predictor_performance = analysis.assess_performance(yhat=all_predictors_to_compare[predictor_name], \n",
    "                                                        y=y, \n",
    "                                                        binary_prediction=False)\n",
    "    predictor_performance['predictor'] = predictor_name\n",
    "    all_klg_versions_performance.append(predictor_performance)\n",
    "print(pd.DataFrame(all_klg_versions_performance)[['predictor', 'r^2', 'spearman_r^2', 'rmse']])\n",
    "    \n",
    "\n",
    "for predictor_name in ['klg', 'klg_hat', 'discretized_klg_hat']:\n",
    "    print('pain gap reduction due to %s' % predictor_name)\n",
    "    print(analysis.quantify_pain_gap_reduction_vs_rival(yhat=all_predictors_to_compare[predictor_name],\n",
    "                                              y=y, \n",
    "                                              rival_severity_measure=yhat_from_klg, \n",
    "                                              all_ses_vars=all_ses_vars, \n",
    "                                             ids=ids)[['SES var', 'yhat/rival red. ratio']])\n",
    "\n",
    "\n",
    "                            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict at future timepoints. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis.predict_kl_at_future_timepoints(non_image_data=datasets['test'].non_image_data,\n",
    "                                         yhat=yhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interpretability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "correlations = analysis.assess_what_image_features_y_and_yhat_correlate_with(\n",
    "                non_image_data=datasets['test'].non_image_data, \n",
    "                y=y, \n",
    "                yhat=yhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interventions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# New Ziad surgery analysis. \n",
    "baseline_idxs = datasets['test'].non_image_data['visit'].values == '00 month follow-up: Baseline'\n",
    "print(\"Showing that when we allocate surgery on the basis of pain + KLG, we give more to black patients\")\n",
    "analysis.do_surgery_analysis_ziad_style(yhat=yhat,\n",
    "                                        y=y,\n",
    "                                        klg=klg,\n",
    "                                        all_ses_vars=all_ses_vars, \n",
    "                                        baseline_idxs=baseline_idxs,\n",
    "                                        have_actually_had_surgery=datasets['test'].non_image_data['knee_surgery'].values, \n",
    "                                        df_to_use=datasets['test'].non_image_data, \n",
    "                                        ids=ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Show that disadvantaged groups have higher rates of surgery (both controlling and not controlling for KLG). \n",
    "concatenated_df_for_surgery_gap_analysis = pd.concat([datasets['train'].non_image_data, \n",
    "                                      datasets['val'].non_image_data, \n",
    "                                      datasets['test'].non_image_data])\n",
    "concatenated_df_for_surgery_gap_analysis.index = range(len(concatenated_df_for_surgery_gap_analysis))\n",
    "\n",
    "analysis.assess_treatment_gaps_controlling_for_klg(klg=concatenated_df_for_surgery_gap_analysis['xrkl'].values, \n",
    "            all_ses_vars = analysis.extract_all_ses_vars(concatenated_df_for_surgery_gap_analysis)[0], \n",
    "            baseline_idxs=concatenated_df_for_surgery_gap_analysis['visit'].values == '00 month follow-up: Baseline',\n",
    "            df_to_use=concatenated_df_for_surgery_gap_analysis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis.make_painkillers_and_surgery_frequency_bar_plot(\n",
    "    datasets['test'].non_image_data.loc[baseline_idxs])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# New Ziad surgery analysis. This is just arthroscopies and shows no improvement in pain; \n",
    "# if anything, pain gets worse. This makes sense because the variable of real interest is arthroplasty, knee replacement\n",
    "# so we bring that in by going back to the raw data. \n",
    "# This analysis does not end up in the final paper. \n",
    "original_non_image_data = non_image_data_processing.NonImageData('all',\n",
    "    i_promise_i_really_want_to_use_the_blinded_hold_out_set=True,                                                        \n",
    "    timepoints_to_filter_for=['12 month follow-up', \n",
    "                              '24 month follow-up', \n",
    "                              '36 month follow-up', \n",
    "                              '48 month follow-up', \n",
    "                              '00 month follow-up: Baseline'])\n",
    "\n",
    "df_for_surgery_progression_analysis = pd.concat([datasets['train'].non_image_data, \n",
    "                                  datasets['val'].non_image_data, \n",
    "                                   datasets['test'].non_image_data])\n",
    "\n",
    "old_len = len(df_for_surgery_progression_analysis)\n",
    "df_for_surgery_progression_analysis = pd.merge(df_for_surgery_progression_analysis, \n",
    "                                              original_non_image_data.processed_dataframes['knee_replacement'], \n",
    "                                              how='inner', \n",
    "                                              on=['id', 'side', 'visit'])\n",
    "assert len(df_for_surgery_progression_analysis) == old_len\n",
    "\n",
    "# note that people get filtered out of our data if they have knee repalcement, so we can't even analyze this using our processed data. \n",
    "# (they get filtered out b/c KLG and other image features aren't meaningful)\n",
    "print(\"Confirmation: the vast majority of points in our datapoint do not have knee replacement: proportion %2.6f do\" % \n",
    "      df_for_surgery_progression_analysis['knee_replacement'].mean())\n",
    "\n",
    "\n",
    "# Does pain decrease after general knee arthroscopy -- seems like answer is NO. \n",
    "print(\"Do people who get general surgery experience lower pain?\")\n",
    "analysis.study_effect_of_surgery(df_for_surgery_progression_analysis, \n",
    "                             surgery_col_to_analyze='knee_surgery')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Does pain decrease after arthroplasty. \n",
    "# Note this analysis cannot be done on our original dataset\n",
    "# because all the people lack post surgery data, so we go back to the raw data. \n",
    "# This analysis does not end up in the final paper. \n",
    "\n",
    "pain_df = original_non_image_data.processed_dataframes['all_knee_pain_scores'].copy()\n",
    "knee_replacement_df = original_non_image_data.processed_dataframes['knee_replacement'].copy()\n",
    "\n",
    "merged_df = pd.merge(knee_replacement_df, pain_df, how='inner', on=['id', 'visit', 'side'], validate='one_to_one')\n",
    "merged_df['high_pain'] = binarize_koos(merged_df['koos_pain_subscore'])\n",
    "merged_df = merged_df.merge(original_non_image_data.processed_dataframes['medications'], \n",
    "                           how='inner', \n",
    "                           on=['id', 'visit'], validate='many_to_one')\n",
    "\n",
    "test_prognostic_df = datasets['test'].non_image_data[['id', 'visit', 'side', 'xrkl']].copy()\n",
    "test_prognostic_df['discretized_yhat'] = discretized_yhat\n",
    "assert Counter(test_prognostic_df['discretized_yhat']) == Counter(test_prognostic_df['xrkl'])\n",
    "\n",
    "merged_df = pd.merge(merged_df, test_prognostic_df, how='left', on=['id', 'visit', 'side'], validate='one_to_one')\n",
    "analysis.study_effect_of_surgery(merged_df, \n",
    "                                 surgery_col_to_analyze='knee_replacement')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rate of surgery increases with KLG. Note we compute this on full dataset. \n",
    "\n",
    "interventions_df = pd.concat([datasets['train'].non_image_data, \n",
    "                             datasets['val'].non_image_data, \n",
    "                             datasets['test'].non_image_data])\n",
    "interventions_df = interventions_df.loc[interventions_df['visit'] == '00 month follow-up: Baseline']\n",
    "interventions_df = interventions_df.dropna(subset=['knee_surgery'])\n",
    "interventions_df.index = range(len(interventions_df))\n",
    "analysis.make_rate_of_surgery_figure(interventions_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis.make_painkillers_and_surgery_frequency_bar_plot(interventions_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Show assignment to higher risk categories and reassignment of surgery. \n",
    "\n",
    "baseline_idxs = datasets['test'].non_image_data['visit'].map(lambda x:x == '00 month follow-up: Baseline').values\n",
    "\n",
    "analysis.make_scatter_plot_showing_severity_reassignment_under_yhat(yhat=yhat, \n",
    "                                                           y=y, \n",
    "                                                           klg=klg, \n",
    "                                                           all_ses_vars=all_ses_vars, \n",
    "                                                           idxs_to_use=baseline_idxs, \n",
    "                                                           interventions_df=interventions_df)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What is the effect of diversity? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Load in diversity results. \n",
    "all_diversity_results = analysis.load_all_results(binary=False,  min_timestring='2019_06_20', \n",
    "                                    thing_to_filter_by={'experiment_to_run':'increase_diversity', \n",
    "                                                         'crop_to_just_the_knee':False})\n",
    "\n",
    "for k in ['exclude_minority_group', 'ses_col', 'majority_group_seed']:\n",
    "    all_diversity_results[k] = all_diversity_results['increase_diversity_kwargs'].map(lambda x:x[k])\n",
    "assert ((pd.isnull(all_diversity_results['majority_group_seed'])) == \n",
    "        (all_diversity_results['exclude_minority_group'])).all()\n",
    "all_diversity_results['majority_group_seed'] = all_diversity_results['majority_group_seed'].map(str)\n",
    "\n",
    "\n",
    "# analyze results. \n",
    "analysis.analyze_effect_of_diversity(all_diversity_results, \n",
    "                                     all_ses_vars=all_ses_vars,\n",
    "                                     y=y, \n",
    "                                     yhat_from_klg=yhat_from_klg, \n",
    "                                     ids=ids, \n",
    "                                     n_bootstraps=N_BOOTSTRAPS)\n",
    "                                     \n",
    "                                     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Not using code below here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_this_code = False\n",
    "if run_this_code:\n",
    "    raise Exception(\"Not using this code at present; in theory it computes the painkiller gap, but in practice it's pretty messy and you should check it if you use it.\")\n",
    "    ids_on_painkillers_at_baseline = set(medication_df.loc[medication_df['Narcotic_analgesic'] == 1, 'id'])\n",
    "\n",
    "    painkiller_gap_df = datasets['test'].non_image_data[[\n",
    "        'id',\n",
    "        'visit',\n",
    "        'binarized_education_graduated_college', \n",
    "        'binarized_income_at_least_50k', \n",
    "        'race_black', \n",
    "        'xrkl']].copy()\n",
    "    painkiller_gap_df['discretized_yhat'] = discretized_yhat\n",
    "    painkiller_gap_df['on_painkillers_at_baseline'] = 1.*painkiller_gap_df['id'].map(lambda x:x in ids_on_painkillers_at_baseline)\n",
    "    painkiller_gap_df = painkiller_gap_df.loc[painkiller_gap_df['visit'].map(lambda x:'Baseline' in x)]\n",
    "\n",
    "    for k in ['binarized_education_graduated_college', \n",
    "              'binarized_income_at_least_50k', \n",
    "              'race_black']:\n",
    "        painkiller_model = sm.OLS.from_formula('on_painkillers_at_baseline ~ %s' % k, \n",
    "                                           data=painkiller_gap_df).fit()\n",
    "        print(painkiller_model.params[k])\n",
    "        painkiller_model = sm.OLS.from_formula('on_painkillers_at_baseline ~ %s + C(xrkl)' % k, \n",
    "                                           data=painkiller_gap_df).fit()\n",
    "        print(painkiller_model.params[k])\n",
    "        painkiller_model = sm.OLS.from_formula('on_painkillers_at_baseline ~ %s + C(discretized_yhat)' % k, \n",
    "                                           data=painkiller_gap_df).fit()\n",
    "        print(painkiller_model.params[k])\n",
    "                                                     "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "knee",
   "language": "python",
   "name": "knee"
  },
  "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.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}