codoc_experiments.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UQB_vSBKPbjK"
      },
      "source": [
        "# Code for Replicating CoDoC Experiment Results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lKEnGwQfOUL9"
      },
      "source": [
        "This file is designed for the replication of the results of the experiments presented in the paper: _Enhancing the reliability and accuracy of AI-enabled diagnosis via complementarity-driven deferral to clinicians (CoDoC)_, by Dvijotham et al., 2023, published in _Nature Medicine_.\n",
        "\n",
        "A detailed set of instructions for using this file to replicate the experiments is provided in the supplementary material for the manuscript. In addition, this file also includes some instructive commentary, as well as guidance with respect to the specifics of the implementation. If you encounter any difficulty with the code and instructions in this file, feel free to contact the authors of the paper."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "drnAgH2Q5f3S"
      },
      "source": [
        "## Obtaining and Utilizing Data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rwOG3eL15f3T"
      },
      "source": [
        "Please see the `README` file in this repository for information regarding accessing the datasets featured in the paper.\n",
        "\n",
        "This notebook describes how to utilize our code once you obtain the data, or using your own data. See below for instructions regarding using this repository with your own dataset.\n",
        "\n",
        "The dataset to be used must consist of three files `tune.csv`, `val.csv`, and `test.csv`, containing the tune, validation, and test splits. All three files must include the columns for ground truth (`y_true`), ML model predictions (`y_model`), and clinical workflow opinions (`reader_score`).\n",
        "\n",
        "Based on the experiment you are attempting to replicate, the three files are placed under the following subfolders:\n",
        "\n",
        "- `data/uk_mammo/single_reader/` for experiments with UK Mammography Dataset and single reader.\n",
        "- `data/uk_mammo/arbitration/` for experiments with UK Mammography Dataset and arbitration.\n",
        "- `data/us_mammo/` for experiments with US Mammography Dataset.\n",
        "\n",
        "For any custom data, prepare the splits as described above and place them under the directory `data/[experiment_name]`, with an `experiment_name` of your own choosing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9f6_mv3DJEYC"
      },
      "source": [
        "## Training CoDoC with A Single Hyperparameter Setting"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VNJV207mghQ0"
      },
      "source": [
        "### Setting experiment name and target metric, loading data and hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gmzk6_ve5f3T"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'\n",
        "os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'\n",
        "\n",
        "import jax\n",
        "import codoc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YifKw63TmPQT"
      },
      "source": [
        "This section can be used to reproduce specific experiment results from the main paper. With the provided `experiment_name`, the data files and the hyperparameters will be automatically loaded based on the experiment setting and target metric you specified. For replicating results, please make sure you select `experiment_name` from among the provided options below in the long comment at the start of the cell. In cases where you are running this notebook with your own data, make sure you have stored the data under `data/[experiment_name]` as described above, according to the `experiment_name` you chose to use.\n",
        "\n",
        "In the cell below, `load_data` loads tune, validation, and test splits for your data as three separate `pandas.DataFrame` objects. Based on the `experiment_name` you provided, the code accesses the folder `data/[experiment_name]` and loads the files `tune.csv`, `val.csv`, `test.csv` as the respective data splits.\n",
        "\n",
        "Following this, `load_hyperparameters` function accesses the file `hyperparameters.json` and loads it as a `dict` object. This dictionary includes specific model hyperparameters to replicate the results in the main paper.\n",
        "\n",
        "Specify `target_metric` based on which model result you are trying to obtain: Models in the main paper have been chosen according to either to maximize sensitivity (`sens`) or specificity (`spec`).\n",
        "\n",
        "The last lines of the cell assign to the variable `hp` the desired model's hyperparameters for easy access. If you get a `KeyError` exception, that means you are attempting to access an experiment that is not present in the main paper or is not available through this repository. If you are using a custom dataset, it is normal to see this exception, and you are free to provide hyperparameters as desired (or select them using our model selection methodology, see next section)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LGpT7epSRuOg"
      },
      "outputs": [],
      "source": [
        "\"\"\"Experiment Name Choices\n",
        "\n",
        "Featured datasets:\n",
        "- \"uk_mammo_single\", \"uk_mammo_arbitration\" for UK Mammography Dataset with\n",
        "      either single reader or arbitration respectively.\n",
        "- \"us_mammo\" for US Mammography Dataset.\n",
        "\n",
        "For own datasets, pass the experiment name you have assigned once you store the\n",
        "data under `data/[experiment_name]`.\n",
        "\"\"\"\n",
        "experiment_name = \"uk_mammo_single\"\n",
        "df_tune, df_val, df_test = codoc.load_data(experiment_name)\n",
        "all_hp = codoc.load_hyperparameters()\n",
        "\n",
        "# Target metric: \"sens\" for sensitivity, \"spec\" for specificity.\n",
        "target_metric = \"spec\"\n",
        "\n",
        "try:\n",
        "  hp = all_hp[experiment_name + \"_\" + target_metric]\n",
        "except KeyError:\n",
        "  hp = {}\n",
        "  print(\n",
        "      \"The experiment + target_metric combination\"\n",
        "      f\" {experiment_name + '_' + target_metric} you requested cannot be found!\"\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HBWKCSkIgmCP"
      },
      "source": [
        "### Running the experiment with loaded hyperparameters and plotting the results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LL8K3COGmPQU"
      },
      "source": [
        "In the following cell, the function `estimate_model` takes the tune and validation splits we have loaded above, as well as model hyperparameters to estimate a CoDoC model. If you are using a custom dataset, feel free to provide custom hyperparameters by changing the inputs to `estimate_model`. This function returns a `dict` object that includes the estimated model's parameters, namely thresholds within which the decisions will be deferred to a clinician. This dictionary will also include sensitivity and specificity estimates in the tune and validation splits, as well as auxiliary statistics such as the percent of cases deferred to a clinician in tune and validation splits.\n",
        "\n",
        "Following the estimation of the model, `evaluate_codoc_model` uses the model parameters and the test data to estimate CoDoC model's sensitivity and specificity on test set. Similarly, `evaluate_baseline_reader` and `evaluate_baseline_model` computes these estimates for the clinician and AI model respectively.\n",
        "\n",
        "To visualize the CoDoC decision procedure, `plot_advantage_z` takes the model's output and plots $Advantage(z)$ for $z \\in [0,1]$, the operating point $\\theta$, and colors this plot according to $Defer(z)$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-zkW2OQOSk02"
      },
      "outputs": [],
      "source": [
        "# Obtaining results for a single hyperparameter combination\n",
        "model = codoc.estimate_model(\n",
        "    df_tune=df_tune,  # The tune dataset split.\n",
        "    df_val=df_val,  # The validation dataset split.\n",
        "    num_bins=hp[\"num_bins\"],  # $T$ from the main paper, number of bins.\n",
        "    tau=hp[\n",
        "        \"tau\"\n",
        "    ],  # Index for operating point for predictive AI: $\\tau = \\theta * T$.\n",
        "    lam=hp[\n",
        "        \"lam\"\n",
        "    ],  # $\\lambda$ from the main paper, sens-spec trade off hyperparameter.\n",
        "    pseudocounts=hp[\"pseudocounts\"],  # $\\kappa$ from the main paper.\n",
        "    smoothing_bandwidth=hp[\n",
        "        \"smoothing_bandwidth\"\n",
        "    ],  # $\\sigma$ from the main paper.\n",
        ")\n",
        "\n",
        "test_sens, test_spec = codoc.evaluate_codoc_model(\n",
        "    df_test, model[\"operating_point\"], model[\"thresholds\"]\n",
        ")\n",
        "print(\"CoDoC sensitivity on test set:\", test_sens)\n",
        "print(\"CoDoC specificity on test set:\", test_spec)\n",
        "codoc.plot_advantage_z(model[\"phis\"], model[\"params\"][\"tau\"], model[\"a_z\"])\n",
        "\n",
        "print(\"Number of cases in test set is {}\".format(len(df_test)))\n",
        "print(\"Number of positive cases in test set is {}\".format(df_test.y_true.sum()))\n",
        "\n",
        "test_reader_sens, test_reader_spec = codoc.evaluate_baseline_reader(df_test)\n",
        "print(\"Clinical workflow sensitivity on test set: {}\".format(test_reader_sens))\n",
        "print(\"Clinical workflow specificity on test set: {}\".format(test_reader_spec))\n",
        "\n",
        "test_ai_sens, test_ai_spec = codoc.evaluate_baseline_model(df_test)\n",
        "print(\"Predictive AI sensitivity on test set: {}\".format(test_ai_sens))\n",
        "print(\"Predictive AI specificity on test set: {}\".format(test_ai_spec))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MGz_fnoYvzbf"
      },
      "source": [
        "## Conducting Hyperparameter Sweep and Model Selection with CoDoC"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O0bHmkibYtJA"
      },
      "source": [
        "You can use this section of the code to conduct the hyperparameter sweep for an experiment, and do model selection to produce a single model for the said experiment. As described in the main paper, these hyperparameters include: $T, \\kappa, \\sigma, \\theta$, $\\lambda$; corresponding to number of bins, pseuodcounts, smoothing bandwidth, operating point, and metric trade-off parameter. Please make sure you have imported relevant functions and loaded the tune, validation, and test splits of your data before you proceed, see above for more details.\n",
        "\n",
        "A note on numerical precision: The parameter search can produce models that have equal performance on $D_{\\text{val}}$, and the ordering between equivalent models might change based on the floating point arithmetic implementation of the kernel used. This can at times lead to slightly different test set performances than in the main paper; however the results should be very close to the original results in all cases."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wxx_9MjfuwTR"
      },
      "source": [
        "### Running the Parameter Sweep"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "esqLnzUFmPQW"
      },
      "source": [
        "In the cell below, the function `parameter_sweep` will conduct a grid search over the hyperparameter ranges provided. The ranges are provided through the `sweep_params` parameter, and must include lists of values under the corresponding keys for the relevant hyperparameters `no_bins_range`, `pseudocounts_range`, `smoothing_bandwith_range`, `lam_range`; corresponding to ranges for $T, \\kappa, \\sigma, \\lambda$ from the main paper respectively. All possible operating points $\\theta$ will be automatically included in the sweep.\n",
        "\n",
        "To replicate the main paper model selection results, use the hyperparameter ranges provided in the `hyperparameters.json`, which will have been loaded through the `load_hyperparameters` function above. You can also use these default ranges for experiments on your custom dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yR8vbop1iWnO"
      },
      "outputs": [],
      "source": [
        "results = codoc.parameter_sweep(\n",
        "    df_tune, df_val, sweep_params=all_hp[\"sweep_hp_ranges\"]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ov3ssxSsuy3l"
      },
      "source": [
        "### Running Model Selection"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FIGm_pFvmPQW"
      },
      "source": [
        "After the parameter sweep has been conducted and candidate models have been stored in `results` in the cell above, the `select_model` function conducts model selection among the viable models. Leave the cell below as is for the experiment with the target metric you are attempting to replicate from the main paper. The model selection parameters are provided for sparse vs. rich data regimes and the two target metrics through `hyperparameter.json` as loaded above.\n",
        "\n",
        "For any custom dataset and `experiment_name`, sparse data regime parameters will be used by default. However feel free to consult the docstring of `select_model` to create a custom model selection scheme.\n",
        "\n",
        "The following code will extract the selected model's hyperparameters to the variable `selected_model_hp`, which can then be used to conduct a single model estimation and obtain the model for the selected hyperparameters, as shown furthre below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-Sm2AaggrZPO"
      },
      "outputs": [],
      "source": [
        "results = codoc.select_model(\n",
        "    results,\n",
        "    experiment_name=experiment_name,\n",
        "    **all_hp[\"model_selection_params\"][codoc.data_regime(experiment_name)][\n",
        "        target_metric\n",
        "    ],\n",
        ")\n",
        "selected_model_hp = results[\"params\"][results[\"val_idx\"]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cXyAYiHnwIuQ"
      },
      "source": [
        "### Examining the Selected Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-qaaY_4tmPQW"
      },
      "source": [
        "The following code repeats the original model estimation procedure described above, but this time with hyperparameters that are determined through model selection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d36SEQC7s0vR"
      },
      "outputs": [],
      "source": [
        "model = codoc.estimate_model(\n",
        "    df_tune=df_tune,  # The tune dataset split.\n",
        "    df_val=df_val,  # The validation dataset split.\n",
        "    num_bins=selected_model_hp[\n",
        "        \"num_bins\"\n",
        "    ],  # $T$ from the main paper, number of bins.\n",
        "    tau=selected_model_hp[\n",
        "        \"tau\"\n",
        "    ],  # Index for operating point for predictive AI: $\\tau = \\theta * T$.\n",
        "    lam=selected_model_hp[\n",
        "        \"lam\"\n",
        "    ],  # $\\lambda$ from the main paper, sens-spec trade off hyperparameter.\n",
        "    pseudocounts=selected_model_hp[\n",
        "        \"pseudocounts\"\n",
        "    ],  # $\\kappa$ from the main paper.\n",
        "    smoothing_bandwidth=selected_model_hp[\n",
        "        \"smoothing_bandwidth\"\n",
        "    ],  # $\\sigma$ from the main paper.\n",
        ")\n",
        "\n",
        "test_sens, test_spec = codoc.evaluate_codoc_model(\n",
        "    df_test, model[\"operating_point\"], model[\"thresholds\"]\n",
        ")\n",
        "print(\"CoDoC sensitivity on test set:\", test_sens)\n",
        "print(\"CoDoC specificity on test set:\", test_spec)\n",
        "codoc.plot_advantage_z(model[\"phis\"], model[\"params\"][\"tau\"], model[\"a_z\"])\n",
        "\n",
        "print(\"Number of cases in test set is {}\".format(len(df_test)))\n",
        "print(\"Number of positive cases in test set is {}\".format(df_test.y_true.sum()))\n",
        "\n",
        "test_reader_sens, test_reader_spec = codoc.evaluate_baseline_reader(df_test)\n",
        "print(\"Clinical workflow sensitivity on test set: {}\".format(test_reader_sens))\n",
        "print(\"Clinical workflow specificity on test set: {}\".format(test_reader_spec))\n",
        "\n",
        "test_ai_sens, test_ai_spec = codoc.evaluate_baseline_model(df_test)\n",
        "print(\"Predictive AI sensitivity on test set: {}\".format(test_ai_sens))\n",
        "print(\"Predictive AI specificity on test set: {}\".format(test_ai_spec))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}