convergence_guarantees_ENCO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "20351b9a",
   "metadata": {},
   "source": [
    "# Convergence guarantees - ENCO\n",
    "\n",
    "This notebook contains an example of how we can calculate the convergence conditions for ENCO on a simple graph. This notebook assumes the following graph structure: $X_1\\to X_2\\to X_3$. The conditional distributions are specified below, and can be alternated freely. The code is fixed for the specific graph structure, but it can be extended to other graph structures by changing the calculation of the corresponding conditional distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3291908f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cd93d38",
   "metadata": {},
   "source": [
    "## Conditional distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c23a864d",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_1 = np.array([0.3, 0.7])                  # p(X_1) - First value: p(X_1=0), Second value: p(X_1=1)\n",
    "p_2_1 = np.array([[0.6, 0.4], [0.4, 0.6]])  # p(X_2|X_1) - First axis: X_1, Second axis: X_2\n",
    "p_3_2 = np.array([[0.2, 0.8], [0.8, 0.2]])  # p(X_3|X_2) - First axis: X_2, Second axis: X_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a06b915e",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_interv = np.array([0.5, 0.5])  # Distribution to use for interventions instead of original conditionals"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ffac97",
   "metadata": {},
   "source": [
    "### Determining relevant conditionals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "465cd912",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_joint=array([[[0.036, 0.144],\n",
      "        [0.096, 0.024]],\n",
      "\n",
      "       [[0.056, 0.224],\n",
      "        [0.336, 0.084]]])\n",
      "p_joint.sum()=1.0\n"
     ]
    }
   ],
   "source": [
    "p_joint = p_1[:,None,None] * p_2_1[:,:,None] * p_3_2[None,:,:]  # p(X_1,X_2,X_3) - First axis: X_1, Second axis: X_2, Third axis: X_3\n",
    "print(f'{p_joint=}')\n",
    "print(f'{p_joint.sum()=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d169a209",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_2=array([0.46, 0.54])\n",
      "p_3=array([0.524, 0.476])\n"
     ]
    }
   ],
   "source": [
    "p_2 = (p_1[:,None] * p_2_1).sum(axis=0)  # p(X_2)\n",
    "p_3 = (p_2[:,None] * p_3_2).sum(axis=0)  # p(X_3)\n",
    "print(f'{p_2=}')\n",
    "print(f'{p_3=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b982ce4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_2_3=array([[0.17557252, 0.82442748],\n",
      "       [0.77310924, 0.22689076]])\n"
     ]
    }
   ],
   "source": [
    "p_2_3 = (p_3_2 * p_2[:,None] / p_3[None,:]).transpose()  # p(X_2|X_3) - First axis: X_3, Second axis: X_2\n",
    "print(f'{p_2_3=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6371b2c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_3_1=array([[0.44, 0.56],\n",
      "       [0.56, 0.44]])\n"
     ]
    }
   ],
   "source": [
    "p_3_1 = (p_2_1[:,:,None] * p_3_2[None,:,:]).sum(axis=1)  # p(X_3|X_1) - First axis: X_1, Second axis: X_3\n",
    "print(f'{p_3_1=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5db1ddc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_2_1_3=array([[[0.27272727, 0.72727273],\n",
      "        [0.85714286, 0.14285714]],\n",
      "\n",
      "       [[0.14285714, 0.85714286],\n",
      "        [0.72727273, 0.27272727]]])\n"
     ]
    }
   ],
   "source": [
    "p_2_1_3 = np.transpose(p_3_2[None,:,:] * p_2_1[:,:,None] / p_3_1[:,None,:], (0, 2, 1))  # p(X_2|X_1,X_3) - First axis: X_1, Second axis: X_3, Third axis: X_2\n",
    "print(f'{p_2_1_3=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0f8ce1d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expect_diff(joint_prob, prob1, prob2):  # Expected log-likelihood difference\n",
    "    return (joint_prob * (np.log(prob1) - np.log(prob2))).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "632c93ff",
   "metadata": {},
   "source": [
    "### Edge $X_1\\to X_2$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85e48311",
   "metadata": {},
   "source": [
    "_Condition 1 - Part 1_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d6c6d1f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_interv_1_2=array([[0.3, 0.2],\n",
      "       [0.2, 0.3]])\n"
     ]
    }
   ],
   "source": [
    "p_interv_1_2 = p_interv[:,None] * p_2_1  # p(do(X_1),X_2) - First axis X_1, Second axis X_2\n",
    "print(f'{p_interv_1_2=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a5eeba77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.023345797452150232"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_1_2, p_2_1, p_2[None,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "959940a4",
   "metadata": {},
   "source": [
    "_Condition 1 - Part 2_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "89196bc5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_interv_1_2_3=array([[[0.06, 0.24],\n",
      "        [0.16, 0.04]],\n",
      "\n",
      "       [[0.04, 0.16],\n",
      "        [0.24, 0.06]]])\n"
     ]
    }
   ],
   "source": [
    "p_interv_1_2_3 = p_interv[:,None,None] * p_2_1[:,:,None] * p_3_2[None,:,:]  # p(do(X_1),X_2,X_3)\n",
    "print(f'{p_interv_1_2_3=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bc619261",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.014975087998617024"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_1_2_3.transpose((0, 2, 1)), p_2_1_3, p_2_3[None,:,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2fd6ad7",
   "metadata": {},
   "source": [
    "_Condition 3_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b2738491",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_interv_3_1_2=array([[[0.09, 0.06],\n",
      "        [0.14, 0.21]],\n",
      "\n",
      "       [[0.09, 0.06],\n",
      "        [0.14, 0.21]]])\n"
     ]
    }
   ],
   "source": [
    "p_interv_3_1_2 = p_interv[:,None,None] * p_1[None,:,None] * p_2_1[None,:,:]  # p(do(X_3),X_1,X_2)\n",
    "print(f'{p_interv_3_1_2=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "75f3931c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.02334579745215022"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_1_2_3, p_2_1[:,:,None], p_2[None,:,None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5790fcf0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.016932091449143066"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_3_1_2, p_2_1[None,:,:], p_2[None,None,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cf0310a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.020138944450646644"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "0.5 * expect_diff(p_interv_1_2_3, p_2_1[:,:,None], p_2[None,:,None]) + 0.5 * expect_diff(p_interv_3_1_2, p_2_1[None,:,:], p_2[None,None,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b8a5759",
   "metadata": {},
   "source": [
    "### Edge $X_2\\to X_3$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd6f9a42",
   "metadata": {},
   "source": [
    "_Condition 1 - Part 1_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "76a2ab90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_interv_2_3=array([[0.1, 0.4],\n",
      "       [0.4, 0.1]])\n"
     ]
    }
   ],
   "source": [
    "p_interv_2_3 = p_interv[:,None] * p_3_2  # p(do(X_2),X_3)\n",
    "print(f'{p_interv_2_3=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "379f72d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.19389808616771806"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_2_3, p_3_2, p_3[None,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6e889e4",
   "metadata": {},
   "source": [
    "_Condition 1 - Part 2_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d10fd7ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_interv_2_3_1=array([[[0.03, 0.07],\n",
      "        [0.12, 0.28]],\n",
      "\n",
      "       [[0.12, 0.28],\n",
      "        [0.03, 0.07]]])\n"
     ]
    }
   ],
   "source": [
    "p_interv_2_3_1 = p_interv[:,None,None] * p_3_2[:,:,None] * p_1[None,None,:]  # p(do(X_2),X_3,X_1)\n",
    "print(f'{p_interv_2_3_1=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8035f517",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.19999710012319816"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_2_3_1.transpose((0,2,1)), p_3_2[:,None,:], p_3_1[None,:,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb6bf5f7",
   "metadata": {},
   "source": [
    "_Condition 3 - Part 1_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "469ee9c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.19389808616771811"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "int1 = expect_diff(p_interv_1_2_3, p_3_2[None,:,:], p_3[None,None,:])\n",
    "int1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "902da495",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.1938980861677181"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "int2 = expect_diff(p_interv_2_3_1, p_3_2[:,:,None], p_3[None,:,None])\n",
    "int2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "879dfee8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.1938980861677181"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "0.5 * int1 + 0.5 * int2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f43c01dd",
   "metadata": {},
   "source": [
    "_Condition 3 - Part 2_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "470d09ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.1855273767141849"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "int1 = expect_diff(p_interv_1_2_3, p_3_2[None,:,:], p_3_1[:,None,:])\n",
    "int1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "576dbd09",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.19999710012319816"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "int2 = expect_diff(p_interv_2_3_1.transpose((0,2,1)), p_3_2[:,None,:], p_3_1[None,:,:])\n",
    "int2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a1f946f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.19276223841869153"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "0.5 * int1 + 0.5 * int2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20a67efe",
   "metadata": {},
   "source": [
    "## Ancestor-Descendant pair $X_1,X_3$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd41c526",
   "metadata": {},
   "source": [
    "_Condition 1 - Part 1_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "91e1fab5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_interv_1_3=array([[0.22, 0.28],\n",
      "       [0.28, 0.22]])\n"
     ]
    }
   ],
   "source": [
    "p_interv_1_3 = p_interv_1_2_3.sum(axis=1)  # p(do(X_1),X_3)\n",
    "print(f'{p_interv_1_3=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "27b10aa0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.008370709453533205"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_1_3, p_3_1, p_3[None,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4d76598",
   "metadata": {},
   "source": [
    "_Condition 1 - Part 2_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8f4970aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect_diff(p_interv_1_2_3, p_3_2[None,:,:], p_3_2[None,:,:])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}