{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "99d0d245",
   "metadata": {},
   "source": [
    "# Validation of a Surrogate Endpoint: Prentice Criteria\n",
    "\n",
    "<br>\n",
    "Author: Martin Horvat, March 2026\n",
    "\n",
    "## Introduction\n",
    "\n",
    "This notebook outlines the **Prentice Criteria** (1989) for validating a surrogate endpoint ($S$) for a true clinical endpoint ($T$) in a randomized trial with treatment assignment ($Z$).\n",
    "\n",
    "**Variable Definitions**\n",
    "*   $T \\in \\{0, 1\\}$: **True Clinical Endpoint** (Binary)\n",
    "*   $S \\in \\mathbb{R}$: **Surrogate Endpoint** (Continuous)\n",
    "*   $Z \\in \\{0, 1\\}$: **Treatment Assignment** (Binary)\n",
    "\n",
    "Original paper was discussing binary variables $T, S, Z \\in \\{0,1\\}$\n",
    "\n",
    "## The Four Criteria\n",
    "\n",
    "To validate $S$ as a surrogate for $T$, the following probabiliy conditions must be satisfied:\n",
    "\n",
    "1. **Treatment Effect on True Endpoint**: The treatment must have a statistically significant impact on the true endpoint.\n",
    "\n",
    "$$\n",
    "P(T \\mid Z) \\neq P(T)\n",
    "$$\n",
    "\n",
    "The opposite is $P(T \\mid Z) = P(T)$ implying $P(T, Z) = P(T) P(Z)$, which means that $T$ and $Z$ are independent.   \n",
    "\n",
    "2. **Treatment Effect on Surrogate**: The treatment must have a statistically significant impact on the surrogate endpoint.\n",
    "\n",
    "$$\n",
    "P(S \\mid Z) \\neq P(S)\n",
    "$$\n",
    "\n",
    "Note: this notation only makes literal sense if $S$ is discrete (especially binary/categorical).\n",
    "\n",
    "3. **Surrogate Association with True Endpoint**: The surrogate must be significantly associated with the true endpoint.\n",
    "\n",
    "$$\n",
    "P(T \\mid S) \\neq P(T)\n",
    "$$\n",
    "\n",
    "4. **The Specific Prentice Criterion (Full Mediation)**: The surrogate must fully capture the treatment effect on the true endpoint. \n",
    "Mathematically, $T$ must be **conditionally independent** of $Z$ given $S$:\n",
    "\n",
    "$$\n",
    "P(T \\mid S, Z) = P(T \\mid S)\n",
    "$$\n",
    "\n",
    "Ref:\n",
    "- Wikipedia contributors. *Surrogate endpoint*. Wikipedia.  https://en.wikipedia.org/wiki/Surrogate_endpoint\n",
    "- Prentice RL. (1989). *Surrogate endpoints in clinical trials: definition and operational criteria*. *Statistics in Medicine*. DOI: 10.1002/sim.4780080407  \n",
    "- Freedman LS, Graubard BI, Schatzkin A. (1992). *Statistical validation of intermediate endpoints for chronic diseases*. PMID: 1557551  \n",
    "- Buyse M, Molenberghs G. (1998). *Criteria for the validation of surrogate endpoints in randomized experiments*. PMID: 9840970\n",
    "- Google/DeepMind AI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4261f8d",
   "metadata": {},
   "source": [
    "## Common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eba6ef4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "import src.utils as utils\n",
    "\n",
    "from rich.pretty import pprint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c151183",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1338b541",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Data interpretation:\n",
    "        'Lung_SUV_95th_Percentile' ~ S or log(S)\n",
    "        'Lung_AE_Flag' ~ T\n",
    "        'Treatment' ~ Z\n",
    "\"\"\"\n",
    "\n",
    "df_data, df_work = utils.read_data(\n",
    "    \"data/prenticedataa.xlsx\",\n",
    "    ['Lung_SUV_95th_Percentile', 'Lung_AE_Flag', 'Treatment'],\n",
    "    [\"S\", \"T\", \"Z\"],\n",
    "    logscale = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1d533fe3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Patient_ID",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "Lung_SUV_95th_Percentile",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Bowel_SUV_95th_Percentile",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Thyroid_SUV_75th_Percentile",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Lung_AE_Flag",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Bowel_AE_Flag",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Thyroid_AE_Flag",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Treatment",
         "rawType": "object",
         "type": "string"
        }
       ],
       "ref": "fa6f2da9-e7ce-4d2b-a902-6ed191c7e0b1",
       "rows": [
        [
         "0",
         "112_14",
         "2.105301856994629",
         "2.269335746765137",
         "2.06760048866272",
         "1",
         "0",
         "0",
         "PEMBRO"
        ],
        [
         "1",
         "1127_06",
         "1.536382555961609",
         "5.6944580078125",
         "3.189261913299561",
         "0",
         "0",
         "1",
         "PEMBRO"
        ],
        [
         "2",
         "1414_15",
         "1.274404644966125",
         "2.373678207397461",
         "1.930612087249756",
         "0",
         "0",
         "0",
         "PEMBRO"
        ],
        [
         "3",
         "1705_16",
         "1.18793511390686",
         "2.522944688796997",
         "7.049728393554688",
         "0",
         "0",
         "0",
         "PEMBRO"
        ],
        [
         "4",
         "2278_07",
         "1.674635887145996",
         "3.006704330444336",
         "1.58749532699585",
         "1",
         "0",
         "0",
         "PEMBRO"
        ]
       ],
       "shape": {
        "columns": 8,
        "rows": 5
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Patient_ID</th>\n",
       "      <th>Lung_SUV_95th_Percentile</th>\n",
       "      <th>Bowel_SUV_95th_Percentile</th>\n",
       "      <th>Thyroid_SUV_75th_Percentile</th>\n",
       "      <th>Lung_AE_Flag</th>\n",
       "      <th>Bowel_AE_Flag</th>\n",
       "      <th>Thyroid_AE_Flag</th>\n",
       "      <th>Treatment</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>112_14</td>\n",
       "      <td>2.105302</td>\n",
       "      <td>2.269336</td>\n",
       "      <td>2.067600</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>PEMBRO</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1127_06</td>\n",
       "      <td>1.536383</td>\n",
       "      <td>5.694458</td>\n",
       "      <td>3.189262</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>PEMBRO</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1414_15</td>\n",
       "      <td>1.274405</td>\n",
       "      <td>2.373678</td>\n",
       "      <td>1.930612</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>PEMBRO</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1705_16</td>\n",
       "      <td>1.187935</td>\n",
       "      <td>2.522945</td>\n",
       "      <td>7.049728</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>PEMBRO</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2278_07</td>\n",
       "      <td>1.674636</td>\n",
       "      <td>3.006704</td>\n",
       "      <td>1.587495</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>PEMBRO</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Patient_ID  Lung_SUV_95th_Percentile  Bowel_SUV_95th_Percentile  \\\n",
       "0     112_14                  2.105302                   2.269336   \n",
       "1    1127_06                  1.536383                   5.694458   \n",
       "2    1414_15                  1.274405                   2.373678   \n",
       "3    1705_16                  1.187935                   2.522945   \n",
       "4    2278_07                  1.674636                   3.006704   \n",
       "\n",
       "   Thyroid_SUV_75th_Percentile  Lung_AE_Flag  Bowel_AE_Flag  Thyroid_AE_Flag  \\\n",
       "0                     2.067600             1              0                0   \n",
       "1                     3.189262             0              0                1   \n",
       "2                     1.930612             0              0                0   \n",
       "3                     7.049728             0              0                0   \n",
       "4                     1.587495             1              0                0   \n",
       "\n",
       "  Treatment  \n",
       "0    PEMBRO  \n",
       "1    PEMBRO  \n",
       "2    PEMBRO  \n",
       "3    PEMBRO  \n",
       "4    PEMBRO  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e130984c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "S",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "T",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Z",
         "rawType": "int64",
         "type": "integer"
        }
       ],
       "ref": "65f5bd4f-e917-4132-af97-1aa54ad404c6",
       "rows": [
        [
         "0",
         "0.7444588568509963",
         "1",
         "0"
        ],
        [
         "1",
         "0.4294306635860589",
         "0",
         "0"
        ],
        [
         "2",
         "0.24247912443295785",
         "0",
         "0"
        ],
        [
         "3",
         "0.17221660152532148",
         "0",
         "0"
        ],
        [
         "4",
         "0.5155957608380367",
         "1",
         "0"
        ]
       ],
       "shape": {
        "columns": 3,
        "rows": 5
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>S</th>\n",
       "      <th>T</th>\n",
       "      <th>Z</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.744459</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.429431</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.242479</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.172217</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.515596</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          S  T  Z\n",
       "0  0.744459  1  0\n",
       "1  0.429431  0  0\n",
       "2  0.242479  0  0\n",
       "3  0.172217  0  0\n",
       "4  0.515596  1  0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_work.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "61ac62b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33728 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(ncols=2, figsize =(8,4), layout = \"constrained\", sharey=True)\n",
    "\n",
    "for ax, lab, name in zip(axs, [\"T\",\"Z\"], [\"True endpoint\", \"Treatment\"]):\n",
    "    \n",
    "    for val in [0,1]:\n",
    "        mask = df_work[lab] == val\n",
    "        ax.scatter(df_work[\"S\"][mask], df_work[lab][mask], label=f\"{val}\", alpha=0.5)\n",
    "    \n",
    "    ax.grid()\n",
    "    ax.set_xlabel(\"Surrogate - S\")\n",
    "    #ax.set_ylabel(f\"{name} - {lab}\")\n",
    "    ax.set_title(f\"{name} - {lab}\")\n",
    "\n",
    "# build one shared legend, positioned to the right of the right subplot\n",
    "handles, labels = axs[0].get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc=\"center left\", bbox_to_anchor=(1.02, 0.5), title=\"Value\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c1f4f8a",
   "metadata": {},
   "source": [
    "## Discussing first Prentice criteria"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8513d6c4-8047-447a-9163-351e269d5950",
   "metadata": {},
   "source": [
    "### T-Z statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0adcb6fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "rTZ = utils.binary_independence_analysis(df_work, \"T\", \"Z\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f88fd897",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counts F =\n"
     ]
    },
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "0",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "1",
         "rawType": "int64",
         "type": "integer"
        }
       ],
       "ref": "6a990d52-0086-4c59-8aa5-d7acc2d83825",
       "rows": [
        [
         "0",
         "47",
         "6"
        ],
        [
         "1",
         "4",
         "1"
        ]
       ],
       "shape": {
        "columns": 2,
        "rows": 2
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>47</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    0  1\n",
       "0  47  6\n",
       "1   4  1"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Probabilities P(T,Z) =\n"
     ]
    },
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "0",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "1",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "431bd473-d76e-451e-ab03-9ca002c3a0c4",
       "rows": [
        [
         "0",
         "0.8103448275862069",
         "0.10344827586206896"
        ],
        [
         "1",
         "0.06896551724137931",
         "0.017241379310344827"
        ]
       ],
       "shape": {
        "columns": 2,
        "rows": 2
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.810345</td>\n",
       "      <td>0.103448</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.068966</td>\n",
       "      <td>0.017241</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          0         1\n",
       "0  0.810345  0.103448\n",
       "1  0.068966  0.017241"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 10008 characters]",
      "text/plain": [
       "<Figure size 600x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Contingency table \n",
    "print(\"Counts F =\")\n",
    "display(pd.DataFrame(rTZ[\"counts\"][\"F\"]))\n",
    "\n",
    "# Probabilities\n",
    "print(\"Probabilities P(T,Z) =\")\n",
    "prob = rTZ[\"probabilities\"]\n",
    "display(pd.DataFrame(prob[\"P_TZ\"]))\n",
    "\n",
    "# Create the heatmap for 2x2 Contingency Table\n",
    "plt.figure(figsize=(6, 5))\n",
    "sns.heatmap(rTZ[\"counts\"][\"F\"], \n",
    "            annot=True,           # Show the numbers in the cells\n",
    "            fmt='d',              # Use integer formatting\n",
    "            cmap='YlGnBu',        # Color scheme (Yellow-Green-Blue)\n",
    "            xticklabels= [\"0\", \"1\"], \n",
    "            yticklabels= [\"0\", \"1\"],)\n",
    "\n",
    "plt.xlabel('T')\n",
    "plt.ylabel('Z')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d9d2aa8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P(T) = [0.9137931 0.0862069]\n",
      "P(Z) = [0.87931034 0.12068966]\n"
     ]
    }
   ],
   "source": [
    "# Marginal probabilities \n",
    "print(f\"P(T) = {prob['P_T']}\")\n",
    "print(f\"P(Z) = {prob['P_Z']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "42f7d99e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conditional probabilities P(T|Z) =\n"
     ]
    },
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "0",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "1",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "b07a4aad-7af4-4c42-9300-d0d0733cf936",
       "rows": [
        [
         "0",
         "0.9215686274509803",
         "0.8571428571428572"
        ],
        [
         "1",
         "0.0784313725490196",
         "0.14285714285714288"
        ]
       ],
       "shape": {
        "columns": 2,
        "rows": 2
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.921569</td>\n",
       "      <td>0.857143</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.078431</td>\n",
       "      <td>0.142857</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          0         1\n",
       "0  0.921569  0.857143\n",
       "1  0.078431  0.142857"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Conditional probabilities P(T|Z) =\")\n",
    "pd.DataFrame(prob[\"P_T_given_Z\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5c1118e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Differences: P(T|Z) - P(T) =\n"
     ]
    },
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "0",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "1",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "80760edf-29c5-49fb-90a7-e16f1026f646",
       "rows": [
        [
         "0",
         "0.007775524002704537",
         "-0.056650246305418595"
        ],
        [
         "1",
         "-0.007775524002704537",
         "0.05665024630541873"
        ]
       ],
       "shape": {
        "columns": 2,
        "rows": 2
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.007776</td>\n",
       "      <td>-0.05665</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.007776</td>\n",
       "      <td>0.05665</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          0        1\n",
       "0  0.007776 -0.05665\n",
       "1 -0.007776  0.05665"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Differences: P(T|Z)/P(T) - 1 1^T =\n"
     ]
    },
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "0",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "1",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "eb14d2c3-8c28-47e2-b4b8-979e7b202dde",
       "rows": [
        [
         "0",
         "0.008509064002959699",
         "-0.061994609164420345"
        ],
        [
         "1",
         "-0.09019607843137267",
         "0.6571428571428573"
        ]
       ],
       "shape": {
        "columns": 2,
        "rows": 2
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.008509</td>\n",
       "      <td>-0.061995</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.090196</td>\n",
       "      <td>0.657143</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          0         1\n",
       "0  0.008509 -0.061995\n",
       "1 -0.090196  0.657143"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Difference between Prob_T_at_Z and Prob_T\n",
    "prob = rTZ[\"probabilities\"]\n",
    "print(\"Differences: P(T|Z) - P(T) =\")\n",
    "display(pd.DataFrame(prob[\"P_T_given_Z\"] - prob[\"P_T\"][:, None]))\n",
    "\n",
    "print(\"Differences: P(T|Z)/P(T) - 1 1^T =\")\n",
    "display(pd.DataFrame(prob[\"P_T_given_Z\"]/prob[\"P_T\"][:,None] - np.ones(shape = (2,2))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c2af5bc-b65c-4ff5-9991-714dd4ee9172",
   "metadata": {},
   "source": [
    "### Tests results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "60014129",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'stat'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3243168965699486</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'dof'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'pvalue'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.569024824759762</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'expected'</span>: <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">array</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">46.60344828</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.39655172</span><span style=\"font-weight: bold\">]</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│      </span><span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.39655172</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.60344828</span><span style=\"font-weight: bold\">]])</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'reject_alpha'</span>: <span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>\n",
       "<span style=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m{\u001b[0m\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'stat'\u001b[0m: \u001b[1;36m0.3243168965699486\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'dof'\u001b[0m: \u001b[1;36m1\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'pvalue'\u001b[0m: \u001b[1;36m0.569024824759762\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'expected'\u001b[0m: \u001b[1;35marray\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m46.60344828\u001b[0m,  \u001b[1;36m6.39655172\u001b[0m\u001b[1m]\u001b[0m,\n",
       "\u001b[2;32m│      \u001b[0m\u001b[1m[\u001b[0m \u001b[1;36m4.39655172\u001b[0m,  \u001b[1;36m0.60344828\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'reject_alpha'\u001b[0m: \u001b[3;91mFalse\u001b[0m\n",
       "\u001b[1m}\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Pearson chi-square test\n",
    "pprint(rTZ[\"global_tests\"][\"chi2\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "eca15f98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'odds_ratio'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.9583333333333333</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'pvalue'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.48734165612568525</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'reject_alpha'</span>: <span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>\n",
       "<span style=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m{\u001b[0m\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'odds_ratio'\u001b[0m: \u001b[1;36m1.9583333333333333\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'pvalue'\u001b[0m: \u001b[1;36m0.48734165612568525\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'reject_alpha'\u001b[0m: \u001b[3;91mFalse\u001b[0m\n",
       "\u001b[1m}\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Fishers exact test\n",
    "pprint(rTZ[\"global_tests\"][\"fisher\"], expand_all=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "089bf384",
   "metadata": {},
   "source": [
    "### Conclusion\n",
    "\n",
    "To assess whether the binary variables $T$ and $Z$ are independent, we applied both **Pearson’s chi-square test** and **Fisher’s exact test** to the $2 \\times 2$ contingency table.\n",
    "\n",
    "The obtained p-values were:\n",
    "\n",
    "- Pearson’s chi-square test: $p = 0.569$\n",
    "- Fisher’s exact test: $p = 0.487$\n",
    "\n",
    "At significance level $\\alpha = 0.05$, both p-values are greater than $\\alpha$, so we **fail to reject** the null hypothesis\n",
    "\n",
    "$$\n",
    "H_0:\\; T \\perp Z.\n",
    "$$\n",
    "\n",
    "Thus, the data do **not** provide statistically significant evidence of an association between $T$ and $Z$. In other words, based on this sample, we do not have evidence that the distribution of $T$ depends on $Z$.\n",
    "\n",
    "For a binary variable $T$, the null hypothesis of independence can equivalently be written as\n",
    "\n",
    "$$\n",
    "P(T=1 \\mid Z=0) = P(T=1 \\mid Z=1).\n",
    "$$\n",
    "\n",
    "One expected cell count is very small (approximately $0.603$), which means the chi-square approximation may be unreliable. Therefore, **Fisher’s exact test is the more appropriate test in this case**. Since Fisher’s exact test is also not significant, the final conclusion remains:\n",
    "\n",
    "**There is insufficient evidence to reject independence between $T$ and $Z$ at the 5\\% significance level.**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b46ab90",
   "metadata": {},
   "source": [
    "## Prentice Criteria (PC) evaluation framework\n",
    "\n",
    "It was by Google AI. In practice, these are tested using a nested series of regression models:\n",
    "\n",
    "1. **Model for PC 1 (Effect of $Z$):**\n",
    "\n",
    "$$\n",
    "    \\operatorname{logit}(P(T=1)) = \\mu_1 + \\beta Z\n",
    "$$  \n",
    "\n",
    "*Requirement:* $\\beta \\neq 0$\n",
    "\n",
    "3. **Model for PC 2 (Effect of $Z$ on $S$):**\n",
    "\n",
    "$$\n",
    "    E[S] = \\mu_2 + \\alpha Z\n",
    "$$  \n",
    "\n",
    "*Requirement:* $\\alpha \\neq 0$\n",
    "\n",
    "5. **Model for PC 3 (The Full Model):**\n",
    "\n",
    "$$\n",
    "    \\operatorname{logit}(P(T=1)) = \\mu_3 + \\gamma S + \\beta_S Z\n",
    "$$\n",
    "\n",
    "*Requirement (Criterion 4):* $\\gamma \\neq 0$ and $\\beta_S \\to 0$\n",
    "\n",
    "**Note:** PC 4 is notoriously difficult to satisfy perfectly. Researchers often calculate the **Proportion Explained (PE)**:\n",
    "$$\n",
    "    PE = \\frac{\\beta - \\beta_S}{\\beta}\n",
    "$$\n",
    "\n",
    "Ideally $PE \\approx 1$ as $\\beta_S \\approx 0$. \n",
    "\n",
    "Note: $\\beta_S$ and $\\beta$ is sometimes called $\\beta_{\\rm un-adjusted}$, $\\beta_{\\rm adjusted}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "30841dae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Criterion",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "Method",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "Model",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "Estimate",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "P-value",
         "rawType": "object",
         "type": "unknown"
        },
        {
         "name": "Pass",
         "rawType": "boolean",
         "type": "boolean"
        }
       ],
       "ref": "e26dba4c-dc0c-4b6f-855a-e61accb0f5a1",
       "rows": [
        [
         "0",
         "1. Treatment -> True",
         "Pearson's Chi-squared test",
         "Z x T (2x2 table)",
         "0.6720937713621131",
         "0.569024824759762",
         "False"
        ],
        [
         "1",
         "2. Treatment -> Surrogate",
         "Kolmogorov–Smirnov test",
         "S by Z groups",
         "0.28851540616246496",
         "0.5879664894834571",
         "False"
        ],
        [
         "2",
         "3. Surrogate -> True",
         "Logistic",
         "T ~ S",
         "10.73815933915853",
         "0.004114055500597262",
         "True"
        ],
        [
         "3",
         "4. Full Mediation",
         "Logistic",
         "T ~ S + Z",
         "1.0343270483046114",
         "0.6441919482170702",
         "True"
        ],
        [
         "4",
         "Proportion Explained (PE)",
         "Logistic coefficients",
         "Compare Z in (T ~ Z) vs (T ~ S + Z).",
         "-0.5389624072967832",
         null,
         "False"
        ]
       ],
       "shape": {
        "columns": 6,
        "rows": 5
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Criterion</th>\n",
       "      <th>Method</th>\n",
       "      <th>Model</th>\n",
       "      <th>Estimate</th>\n",
       "      <th>P-value</th>\n",
       "      <th>Pass</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1. Treatment -&gt; True</td>\n",
       "      <td>Pearson's Chi-squared test</td>\n",
       "      <td>Z x T (2x2 table)</td>\n",
       "      <td>0.672094</td>\n",
       "      <td>0.569025</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2. Treatment -&gt; Surrogate</td>\n",
       "      <td>Kolmogorov–Smirnov test</td>\n",
       "      <td>S by Z groups</td>\n",
       "      <td>0.288515</td>\n",
       "      <td>0.587966</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3. Surrogate -&gt; True</td>\n",
       "      <td>Logistic</td>\n",
       "      <td>T ~ S</td>\n",
       "      <td>10.738159</td>\n",
       "      <td>0.004114</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4. Full Mediation</td>\n",
       "      <td>Logistic</td>\n",
       "      <td>T ~ S + Z</td>\n",
       "      <td>1.034327</td>\n",
       "      <td>0.644192</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Proportion Explained (PE)</td>\n",
       "      <td>Logistic coefficients</td>\n",
       "      <td>Compare Z in (T ~ Z) vs (T ~ S + Z).</td>\n",
       "      <td>-0.538962</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   Criterion                      Method  \\\n",
       "0       1. Treatment -> True  Pearson's Chi-squared test   \n",
       "1  2. Treatment -> Surrogate     Kolmogorov–Smirnov test   \n",
       "2       3. Surrogate -> True                    Logistic   \n",
       "3          4. Full Mediation                    Logistic   \n",
       "4  Proportion Explained (PE)       Logistic coefficients   \n",
       "\n",
       "                                  Model   Estimate   P-value   Pass  \n",
       "0                     Z x T (2x2 table)   0.672094  0.569025  False  \n",
       "1                         S by Z groups   0.288515  0.587966  False  \n",
       "2                                 T ~ S  10.738159  0.004114   True  \n",
       "3                             T ~ S + Z   1.034327  0.644192   True  \n",
       "4  Compare Z in (T ~ Z) vs (T ~ S + Z).  -0.538962      <NA>  False  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_prentice = utils.prentice_criteria(\n",
    "    df_work,\n",
    "    endpoint=\"T\",\n",
    "    surrogate=\"S\",\n",
    "    treatment=\"Z\",\n",
    "    alpha=0.05,\n",
    "    return_models=False,\n",
    "    criterion1_test=\"chi2\",\n",
    "    criterion2_test=\"ks\",\n",
    ")\n",
    "\n",
    "df_prentice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469afd91",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
