Results_on_semi_synthetic_linear_Gaussian_graphs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gciw4GKDUS0_"
   },
   "source": [
    "Code for generating the results for semi-synthetic linear Gaussian graphs from [*bnlearn*](https://www.bnlearn.com/) (Figs. 11,17)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "8Xqc_6q-WNW1"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import networkx as nx\n",
    "import numpy\n",
    "import pandas\n",
    "import networkx\n",
    "from itertools import combinations, permutations\n",
    "import logging\n",
    "from causallearn.utils.cit import CIT\n",
    "import collections\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle as pkl\n",
    "from datetime import datetime\n",
    "import copy\n",
    "\n",
    "from causal_discovery.utils import *\n",
    "from causal_discovery.pc_alg import PCAlgorithm\n",
    "from causal_discovery.mb_by_mb import MBbyMBAlgorithm\n",
    "from causal_discovery.sd_alg import SequentialDiscoveryAlgorithm\n",
    "from causal_discovery.ldecc import LDECCAlgorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "pnj3tN45cbZq"
   },
   "outputs": [],
   "source": [
    "def get_bnlearn_graph_params(name):\n",
    "  import pickle\n",
    "  return pickle.load(open(\"data/bnlearn-%s.pkl\" % name, \"rb\"))\n",
    "\n",
    "def get_cov_matrix(G_weighted, G_std):\n",
    "  A = G_weighted\n",
    "  B = np.linalg.inv(np.eye(A.shape[0]) - A)\n",
    "  noise_cov = G_std**2\n",
    "  cov_matrix = B @ noise_cov @ B.T\n",
    "  return cov_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "sYvrSkYZpoL_"
   },
   "outputs": [],
   "source": [
    "graph_to_tmt_outcome = {\n",
    "    \"magic-niab\": {\n",
    "        \"treatment\": \"G266\",\n",
    "        \"outcome\": \"HT\",\n",
    "    },\n",
    "    \"magic-irri\": {\n",
    "        \"treatment\": \"G6003\",\n",
    "        \"outcome\": \"BROWN\",\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "yw0LzlRmpced"
   },
   "outputs": [],
   "source": [
    "# one of \"magic-niab\" or \"magic-irri\".\n",
    "graph_name = \"magic-niab\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "id": "0qSr_-irdHwp"
   },
   "outputs": [],
   "source": [
    "G_weighted, G_std, node_names = get_bnlearn_graph_params(graph_name)\n",
    "graph_true = adj_matrix_to_nx_graph(G_weighted, node_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "id": "MDmFAlz6qFGC"
   },
   "outputs": [],
   "source": [
    "graph_true = nx.relabel_nodes(graph_true, {\n",
    "    graph_to_tmt_outcome[graph_name][\"treatment\"]: \"X\",\n",
    "    graph_to_tmt_outcome[graph_name][\"outcome\"]: \"Y\"\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 255
    },
    "id": "xoFVii0MdmFC",
    "outputId": "7b62af84-c574-4881-a928-a90ed60503a5"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>YR.GLASS</th>\n",
       "      <th>Y</th>\n",
       "      <th>YR.FIELD</th>\n",
       "      <th>MIL</th>\n",
       "      <th>FT</th>\n",
       "      <th>G418</th>\n",
       "      <th>G311</th>\n",
       "      <th>G1217</th>\n",
       "      <th>G800</th>\n",
       "      <th>G866</th>\n",
       "      <th>...</th>\n",
       "      <th>G2318</th>\n",
       "      <th>G1294</th>\n",
       "      <th>G1800</th>\n",
       "      <th>YLD</th>\n",
       "      <th>FUS</th>\n",
       "      <th>G1750</th>\n",
       "      <th>G524</th>\n",
       "      <th>G775</th>\n",
       "      <th>G2835</th>\n",
       "      <th>G43</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.420137</td>\n",
       "      <td>-9.412403</td>\n",
       "      <td>-0.641280</td>\n",
       "      <td>-0.548300</td>\n",
       "      <td>-4.104396</td>\n",
       "      <td>-1.234455</td>\n",
       "      <td>-0.115014</td>\n",
       "      <td>0.269181</td>\n",
       "      <td>1.177147</td>\n",
       "      <td>0.493421</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.867749</td>\n",
       "      <td>-0.775289</td>\n",
       "      <td>-0.026453</td>\n",
       "      <td>-0.076053</td>\n",
       "      <td>0.333144</td>\n",
       "      <td>0.091130</td>\n",
       "      <td>0.756651</td>\n",
       "      <td>-0.503202</td>\n",
       "      <td>0.314705</td>\n",
       "      <td>0.617827</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.002750</td>\n",
       "      <td>-4.332171</td>\n",
       "      <td>0.020122</td>\n",
       "      <td>-0.989963</td>\n",
       "      <td>0.954850</td>\n",
       "      <td>-0.858694</td>\n",
       "      <td>-1.436846</td>\n",
       "      <td>0.102548</td>\n",
       "      <td>1.271658</td>\n",
       "      <td>-0.096132</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.243383</td>\n",
       "      <td>-1.825503</td>\n",
       "      <td>0.110410</td>\n",
       "      <td>0.632123</td>\n",
       "      <td>-1.325769</td>\n",
       "      <td>0.258811</td>\n",
       "      <td>-0.944720</td>\n",
       "      <td>-1.031226</td>\n",
       "      <td>0.257486</td>\n",
       "      <td>0.620786</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-1.117361</td>\n",
       "      <td>3.640469</td>\n",
       "      <td>0.100454</td>\n",
       "      <td>-0.244524</td>\n",
       "      <td>-1.813873</td>\n",
       "      <td>-1.017628</td>\n",
       "      <td>0.392061</td>\n",
       "      <td>-1.361735</td>\n",
       "      <td>-0.468206</td>\n",
       "      <td>0.275902</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.883446</td>\n",
       "      <td>-0.426219</td>\n",
       "      <td>-0.255114</td>\n",
       "      <td>0.077915</td>\n",
       "      <td>-0.001940</td>\n",
       "      <td>0.847443</td>\n",
       "      <td>-2.619560</td>\n",
       "      <td>0.229620</td>\n",
       "      <td>-0.641768</td>\n",
       "      <td>-0.194753</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.216346</td>\n",
       "      <td>3.526878</td>\n",
       "      <td>-0.247618</td>\n",
       "      <td>0.678260</td>\n",
       "      <td>1.121236</td>\n",
       "      <td>0.071027</td>\n",
       "      <td>1.054274</td>\n",
       "      <td>1.330939</td>\n",
       "      <td>0.096987</td>\n",
       "      <td>0.863503</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.164306</td>\n",
       "      <td>0.591188</td>\n",
       "      <td>-0.123854</td>\n",
       "      <td>-0.655673</td>\n",
       "      <td>-1.221081</td>\n",
       "      <td>1.406668</td>\n",
       "      <td>-0.180540</td>\n",
       "      <td>0.215197</td>\n",
       "      <td>-0.161733</td>\n",
       "      <td>-0.432399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-0.565661</td>\n",
       "      <td>-5.829983</td>\n",
       "      <td>-0.334910</td>\n",
       "      <td>-0.453876</td>\n",
       "      <td>7.003238</td>\n",
       "      <td>0.478129</td>\n",
       "      <td>-0.347818</td>\n",
       "      <td>0.436353</td>\n",
       "      <td>-0.373046</td>\n",
       "      <td>1.388168</td>\n",
       "      <td>...</td>\n",
       "      <td>-1.933598</td>\n",
       "      <td>0.456514</td>\n",
       "      <td>-0.064923</td>\n",
       "      <td>-0.237212</td>\n",
       "      <td>-0.112231</td>\n",
       "      <td>-0.715062</td>\n",
       "      <td>0.641959</td>\n",
       "      <td>0.660831</td>\n",
       "      <td>-0.151781</td>\n",
       "      <td>-0.023642</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 44 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   YR.GLASS         Y  YR.FIELD       MIL        FT      G418      G311  \\\n",
       "0  0.420137 -9.412403 -0.641280 -0.548300 -4.104396 -1.234455 -0.115014   \n",
       "1 -0.002750 -4.332171  0.020122 -0.989963  0.954850 -0.858694 -1.436846   \n",
       "2 -1.117361  3.640469  0.100454 -0.244524 -1.813873 -1.017628  0.392061   \n",
       "3  0.216346  3.526878 -0.247618  0.678260  1.121236  0.071027  1.054274   \n",
       "4 -0.565661 -5.829983 -0.334910 -0.453876  7.003238  0.478129 -0.347818   \n",
       "\n",
       "      G1217      G800      G866  ...     G2318     G1294     G1800       YLD  \\\n",
       "0  0.269181  1.177147  0.493421  ... -0.867749 -0.775289 -0.026453 -0.076053   \n",
       "1  0.102548  1.271658 -0.096132  ... -0.243383 -1.825503  0.110410  0.632123   \n",
       "2 -1.361735 -0.468206  0.275902  ... -0.883446 -0.426219 -0.255114  0.077915   \n",
       "3  1.330939  0.096987  0.863503  ... -0.164306  0.591188 -0.123854 -0.655673   \n",
       "4  0.436353 -0.373046  1.388168  ... -1.933598  0.456514 -0.064923 -0.237212   \n",
       "\n",
       "        FUS     G1750      G524      G775     G2835       G43  \n",
       "0  0.333144  0.091130  0.756651 -0.503202  0.314705  0.617827  \n",
       "1 -1.325769  0.258811 -0.944720 -1.031226  0.257486  0.620786  \n",
       "2 -0.001940  0.847443 -2.619560  0.229620 -0.641768 -0.194753  \n",
       "3 -1.221081  1.406668 -0.180540  0.215197 -0.161733 -0.432399  \n",
       "4 -0.112231 -0.715062  0.641959  0.660831 -0.151781 -0.023642  \n",
       "\n",
       "[5 rows x 44 columns]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = get_dataset(G_weighted, np.diag(G_std)**2, 1000, list(graph_true.nodes()))\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_REfaO6ID9N5"
   },
   "outputs": [],
   "source": [
    "sample_cov = data.cov().values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ocEqP4jDejM1"
   },
   "outputs": [],
   "source": [
    "mb_by_mb_alg = MBbyMBAlgorithm(use_ci_oracle=False)\n",
    "result_mb_by_mb = mb_by_mb_alg.run(data)\n",
    "print(\"Total CI tests done: %d\" % mb_by_mb_alg.ci_test_calls[\"total\"])\n",
    "print(\"Estimated ATE set: %s\" % str(\n",
    "    set(get_ATE_using_nodes_and_cov(list(data.columns),\n",
    "                                    list(result_mb_by_mb[\"tmt_parents\"]) + par,\n",
    "                                    sample_cov)\n",
    "    for par in get_all_combinations(result_mb_by_mb[\"unoriented\"], \n",
    "                                    result_mb_by_mb[\"non_colliders\"]))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AhlAd4GUjgA5"
   },
   "outputs": [],
   "source": [
    "sd_alg = SequentialDiscoveryAlgorithm(use_ci_oracle=False)\n",
    "result_sd = sd_alg.run(data)\n",
    "print(\"Total CI tests done: %d\" % sd_alg.ci_test_calls[\"total\"])\n",
    "print(\"Estimated ATE set: %s\" % str(\n",
    "    set(get_ATE_using_nodes_and_cov(list(data.columns),\n",
    "                                    list(result_sd[\"tmt_parents\"]) + par,\n",
    "                                    sample_cov)\n",
    "    for par in get_all_combinations(result_sd[\"unoriented\"], \n",
    "                                    result_sd[\"non_colliders\"]))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TV5TXUenuTbq"
   },
   "outputs": [],
   "source": [
    "ldecc_alg = LDECCAlgorithm(use_ci_oracle=False, ldecc_do_checks=True)\n",
    "result_ldecc = ldecc_alg.run(data)\n",
    "print(\"Total CI tests done: %d\" % ldecc_alg.ci_test_calls[\"total\"])\n",
    "print(\"Estimated ATE set: %s\" % str(\n",
    "    set(get_ATE_using_nodes_and_cov(list(data.columns),\n",
    "                                    list(result_ldecc[\"tmt_parents\"]) + par,\n",
    "                                    sample_cov)\n",
    "    for par in get_all_combinations(result_ldecc[\"unoriented\"], \n",
    "                                    result_ldecc[\"non_colliders\"]))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R74t7mjIfRDj"
   },
   "outputs": [],
   "source": [
    "import ipyparallel as ipp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "dsMrRbQdhDnY"
   },
   "outputs": [],
   "source": [
    "# Verify that ipcluster is running and import the necessary Python packages.\n",
    "\n",
    "parallel_client = ipp.Client(debug=False)\n",
    "dview = parallel_client[:]\n",
    "# Execute an identity map in parallel.\n",
    "ar = dview.map(lambda x: x, (i for i in range(0, 20000, 2)))\n",
    "assert ar.get()[0] == 0\n",
    "\n",
    "# Import the required Python packages.\n",
    "with dview.sync_imports():\n",
    "  import numpy\n",
    "  import os\n",
    "  import pandas\n",
    "  import networkx\n",
    "  from itertools import combinations, permutations\n",
    "  import logging\n",
    "  from causallearn.utils.cit import CIT\n",
    "  import collections\n",
    "  import copy\n",
    "  import ipyparallel\n",
    "\n",
    "  from causal_discovery.utils import orient_colliders, get_dataset, get_connected_component_with_node, remove_directed_only_edges, apply_meek_rules, adj_matrix_to_nx_graph, dag_to_cpdag, get_neighbor_types, get_all_combinations, get_ATE_using_nodes_and_cov, get_ATE_using_cov, generate_local_dags_from, to_dict    \n",
    "  from causal_discovery.pc_alg import PCAlgorithm\n",
    "  from causal_discovery.mb_by_mb import MBbyMBAlgorithm\n",
    "  from causal_discovery.sd_alg import SequentialDiscoveryAlgorithm\n",
    "  from causal_discovery.ldecc import LDECCAlgorithm\n",
    "  \n",
    "  try:\n",
    "    from cPickle import dumps, loads, HIGHEST_PROTOCOL as PICKLE_PROTOCOL\n",
    "  except ImportError:\n",
    "    from pickle import dumps, loads, HIGHEST_PROTOCOL as PICKLE_PROTOCOL\n",
    "\n",
    "# Make sure ipyparallel is still able to execute functions.\n",
    "dview = parallel_client[:]\n",
    "ar = dview.map(lambda x: x, (i for i in range(0, 20000, 2)))\n",
    "assert ar.get()[0] == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mMjn9PkChGOp"
   },
   "outputs": [],
   "source": [
    "def execute_iteration(params, num_samples, seed, progress_log_filepath=None):\n",
    "  np = numpy\n",
    "  nx = networkx\n",
    "    \n",
    "  np.random.seed(seed)\n",
    "  \n",
    "  # linearly scale down alpha.\n",
    "  alpha = (32000 / num_samples) * 0.01\n",
    "\n",
    "  MAX_TESTS = 7000\n",
    "  \n",
    "  result = {\n",
    "      \"graph_true\": None,\n",
    "      \"graph_true_sample_cov_ate\": None,\n",
    "      \"oracle\": {\n",
    "          \"ate\": None,\n",
    "          \"neighbors\": None,\n",
    "          \"parent_sets\": None,\n",
    "      },\n",
    "      \"oracle_sample_cov\": {\n",
    "          \"ate\": None,\n",
    "      },\n",
    "      \"ldecc\": {\n",
    "          \"ci_tests\": None,\n",
    "          \"ate\": None,\n",
    "          \"neighbors\": None,\n",
    "          \"parent_sets\": None,\n",
    "      },\n",
    "      \"ldecc-checks\": {\n",
    "          \"ci_tests\": None,\n",
    "          \"ate\": None,\n",
    "          \"neighbors\": None,\n",
    "          \"parent_sets\": None,\n",
    "      },\n",
    "      \"mb-by-mb\": {\n",
    "          \"ci_tests\": None,\n",
    "          \"ate\": None,\n",
    "          \"neighbors\": None,\n",
    "          \"parent_sets\": None,\n",
    "      },\n",
    "      \"sd-alg\": {\n",
    "          \"ci_tests\": None,\n",
    "          \"ate\": None,\n",
    "          \"neighbors\": None,\n",
    "          \"parent_sets\": None,\n",
    "      },\n",
    "  }\n",
    "\n",
    "  G_weighted, G_std, node_names, graph_name, cov_matrix = params\n",
    "  graph_true = adj_matrix_to_nx_graph(G_weighted, node_names)\n",
    "  graph_true = nx.relabel_nodes(graph_true, {\n",
    "      graph_to_tmt_outcome[graph_name][\"treatment\"]: \"X\",\n",
    "      graph_to_tmt_outcome[graph_name][\"outcome\"]: \"Y\"\n",
    "  })\n",
    "  data = get_dataset(G_weighted, np.diag(G_std)**2, num_samples, \n",
    "                     list(graph_true.nodes()))\n",
    "  sample_cov = data.cov().values\n",
    "\n",
    "  result[\"graph_name\"] = graph_name\n",
    "  result[\"graph_true\"] = graph_true\n",
    "  result[\"graph_true_sample_cov_ate\"] = get_ATE_using_cov(graph_true, sample_cov)\n",
    "\n",
    "  cpdag_oracle = dag_to_cpdag(graph_true)\n",
    "  result[\"oracle\"][\"ate\"] = [get_ATE_using_cov(g, cov_matrix) \n",
    "                             for g in generate_local_dags_from(cpdag_oracle)]\n",
    "  result[\"oracle\"][\"neighbors\"] = get_neighbor_types(cpdag_oracle, \"X\")\n",
    "  result[\"oracle\"][\"parent_sets\"] = [list(g.predecessors(\"X\")) for g in generate_local_dags_from(cpdag_oracle)]                           \n",
    "  \n",
    "  result[\"oracle_sample_cov\"][\"ate\"] = [get_ATE_using_cov(g, sample_cov)\n",
    "                                        for g in generate_local_dags_from(cpdag_oracle)]\n",
    "  \n",
    "  ldecc_alg = LDECCAlgorithm(use_ci_oracle=False, alpha=alpha,\n",
    "                             ldecc_do_checks=False, max_tests=MAX_TESTS)\n",
    "  result_ldecc = ldecc_alg.run(data)\n",
    "  result[\"ldecc\"][\"ci_tests\"] = to_dict(ldecc_alg.ci_test_calls)\n",
    "  result[\"ldecc\"][\"neighbors\"] = {\n",
    "      \"parents\": result_ldecc[\"tmt_parents\"],\n",
    "      \"children\": result_ldecc[\"tmt_children\"],\n",
    "      \"unoriented\": result_ldecc[\"unoriented\"],\n",
    "      \"non_colliders\": result_ldecc[\"non_colliders\"], \n",
    "  }\n",
    "  tmt_par = list(result_ldecc[\"tmt_parents\"])\n",
    "  result[\"ldecc\"][\"parent_sets\"] = list(list(p) + tmt_par for p in get_all_combinations(result_ldecc[\"unoriented\"], result_ldecc[\"non_colliders\"]))\n",
    "  ate_ldecc = [get_ATE_using_nodes_and_cov(list(data.columns), tmt_par + par, sample_cov)\n",
    "               for par in get_all_combinations(result_ldecc[\"unoriented\"], result_ldecc[\"non_colliders\"])]\n",
    "  result[\"ldecc\"][\"ate\"] = ate_ldecc\n",
    "\n",
    "  ldecc_alg = LDECCAlgorithm(use_ci_oracle=False, alpha=alpha,\n",
    "                             ldecc_do_checks=True, max_tests=MAX_TESTS)\n",
    "  result_ldecc = ldecc_alg.run(data)\n",
    "  result[\"ldecc-checks\"][\"ci_tests\"] = to_dict(ldecc_alg.ci_test_calls)\n",
    "  result[\"ldecc-checks\"][\"neighbors\"] = {\n",
    "      \"parents\": result_ldecc[\"tmt_parents\"],\n",
    "      \"children\": result_ldecc[\"tmt_children\"],\n",
    "      \"unoriented\": result_ldecc[\"unoriented\"],\n",
    "      \"non_colliders\": result_ldecc[\"non_colliders\"], \n",
    "  }\n",
    "  tmt_par = list(result_ldecc[\"tmt_parents\"])\n",
    "  ate_ldecc = [get_ATE_using_nodes_and_cov(list(data.columns), tmt_par + par, sample_cov)\n",
    "               for par in get_all_combinations(result_ldecc[\"unoriented\"], result_ldecc[\"non_colliders\"])]\n",
    "  result[\"ldecc-checks\"][\"parent_sets\"] = list(list(p) + tmt_par for p in get_all_combinations(result_ldecc[\"unoriented\"], result_ldecc[\"non_colliders\"]))\n",
    "  result[\"ldecc-checks\"][\"ate\"] = ate_ldecc\n",
    "\n",
    "  mb_by_mb_alg = MBbyMBAlgorithm(use_ci_oracle=False, alpha=alpha, \n",
    "                                 max_tests=MAX_TESTS)\n",
    "  result_mb = mb_by_mb_alg.run(data)\n",
    "  result[\"mb-by-mb\"][\"ci_tests\"] = to_dict(mb_by_mb_alg.ci_test_calls)\n",
    "  result[\"mb-by-mb\"][\"neighbors\"] = {\n",
    "      \"parents\": result_mb[\"tmt_parents\"],\n",
    "      \"children\": result_mb[\"tmt_children\"],\n",
    "      \"unoriented\": result_mb[\"unoriented\"],\n",
    "      \"non_colliders\": result_mb[\"non_colliders\"], \n",
    "  }\n",
    "  tmt_par = list(result_mb[\"tmt_parents\"])\n",
    "  ate_mb = [get_ATE_using_nodes_and_cov(list(data.columns), tmt_par + par, sample_cov)\n",
    "            for par in get_all_combinations(result_mb[\"unoriented\"], result_mb[\"non_colliders\"])]\n",
    "  result[\"mb-by-mb\"][\"parent_sets\"] = list(list(p) + tmt_par for p in get_all_combinations(result_mb[\"unoriented\"], result_mb[\"non_colliders\"]))\n",
    "  result[\"mb-by-mb\"][\"ate\"] = ate_mb\n",
    "\n",
    "  sd_alg = SequentialDiscoveryAlgorithm(use_ci_oracle=False, alpha=alpha,\n",
    "                                        max_tests=MAX_TESTS)\n",
    "  result_sd = sd_alg.run(data)\n",
    "  result[\"sd-alg\"][\"ci_tests\"] = to_dict(sd_alg.ci_test_calls)\n",
    "  result[\"sd-alg\"][\"neighbors\"] = {\n",
    "      \"parents\": result_sd[\"tmt_parents\"],\n",
    "      \"children\": result_sd[\"tmt_children\"],\n",
    "      \"unoriented\": result_sd[\"unoriented\"],\n",
    "      \"non_colliders\": result_sd[\"non_colliders\"], \n",
    "  }\n",
    "  tmt_par = list(result_sd[\"tmt_parents\"])\n",
    "  result[\"sd-alg\"][\"parent_sets\"] = list(list(p) + tmt_par for p in get_all_combinations(result_sd[\"unoriented\"], result_sd[\"non_colliders\"]))\n",
    "  ate_sd = [get_ATE_using_nodes_and_cov(list(data.columns), tmt_par + par, sample_cov)\n",
    "            for par in get_all_combinations(result_sd[\"unoriented\"], result_sd[\"non_colliders\"])]\n",
    "  result[\"sd-alg\"][\"ate\"] = ate_sd\n",
    "  \n",
    "  if progress_log_filepath is not None:\n",
    "    os.system(\"echo '%d' >> %s\" % (seed, progress_log_filepath))\n",
    "\n",
    "  np.random.seed(None)\n",
    "  return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7ii8I2ZPlWYh"
   },
   "outputs": [],
   "source": [
    "def async_results_to_results(async_results, results):\n",
    "  for N, ar in zip(results[\"sample_sizes\"], async_results):\n",
    "    if N not in results:\n",
    "      results[N] = []\n",
    "    for r in ar.get():\n",
    "      results[N].append(r)\n",
    "  return results\n",
    "\n",
    "def execute_estimation_in_parallel(results, start_graph, end_graph, iters_per_graph=1):\n",
    "\n",
    "  num_threads = len(parallel_client.ids)\n",
    "  dview = parallel_client[:]\n",
    "\n",
    "  num_graphs = end_graph - start_graph\n",
    "  print(\"Executing %d iterations across %d engines\" % (\n",
    "      num_graphs*iters_per_graph, num_threads))\n",
    "\n",
    "  dview[\"get_dataset\"] = get_dataset\n",
    "  dview[\"graph_to_tmt_outcome\"] = graph_to_tmt_outcome\n",
    "  dview[\"get_noise_matrix\"] = get_noise_matrix\n",
    "  dview[\"execute_iteration\"] = execute_iteration\n",
    "\n",
    "  graph_name = \"magic-niab\"\n",
    "  \n",
    "  G_weighted, G_std, node_names = get_bnlearn_graph_params(graph_name)\n",
    "  cov_matrix = get_cov_matrix(G_weighted, G_std)\n",
    "  params = (G_weighted, G_std, node_names, graph_name, cov_matrix)\n",
    "  graph_true = adj_matrix_to_nx_graph(G_weighted, node_names)\n",
    "  graph_true = nx.relabel_nodes(graph_true, {\n",
    "      graph_to_tmt_outcome[graph_name][\"treatment\"]: \"X\",\n",
    "      graph_to_tmt_outcome[graph_name][\"outcome\"]: \"Y\"\n",
    "  })\n",
    "\n",
    "  RANDOM_SEED = 716697\n",
    "  np.random.seed(RANDOM_SEED)\n",
    "  new_param_list = []\n",
    "  for _ in range(start_graph, end_graph):\n",
    "\n",
    "    for graph_iter in range(iters_per_graph):\n",
    "      new_param_list.append(params)\n",
    "      results[\"params\"].append(params)\n",
    "      results[\"ground_truth_ate\"].append(get_ATE_using_cov(graph_true, cov_matrix))\n",
    "      \n",
    "  sample_sizes = [8000, 16000, 24000, 32000]\n",
    "  results[\"sample_sizes\"] = sample_sizes\n",
    "  async_results = []\n",
    "  for N in sample_sizes:\n",
    "    print(\"%s: Starting iterations %d for N=%d\" % (\n",
    "        str(datetime.now()), len(new_param_list), N))\n",
    "    \n",
    "    def execute_estimation_iteration(i):\n",
    "      return execute_iteration(new_param_list[i], N, (RANDOM_SEED + N + i))\n",
    "    \n",
    "    async_result = dview.map(execute_estimation_iteration, range(len(new_param_list)))\n",
    "    async_results.append(async_result)\n",
    "\n",
    "  np.random.seed(None)\n",
    "  return results, async_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rvMPJwZfodBA"
   },
   "outputs": [],
   "source": [
    "results = {}\n",
    "results.update({\n",
    "    \"params\": [],\n",
    "    \"ground_truth_ate\": [],\n",
    "})\n",
    "\n",
    "results, async_results = execute_estimation_in_parallel(results, \n",
    "                                                        start_graph=0,\n",
    "                                                        end_graph=100,\n",
    "                                                        iters_per_graph=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Vfbe49lusOGd"
   },
   "outputs": [],
   "source": [
    "# Monitor progress of the tasks.\n",
    "for ar in async_results:\n",
    "  ar.wait_interactive(timeout=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Fg8iX-g7sYBv"
   },
   "outputs": [],
   "source": [
    "results = async_results_to_results(async_results, results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_tCCeKkT4vis"
   },
   "source": [
    "The code below plots the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lDvbGe-Y1cOP"
   },
   "outputs": [],
   "source": [
    "# For the color map:\n",
    "# https://gist.github.com/AndiH/c957b4d769e628f506bd\n",
    "\n",
    "# Tableau 20 Colors\n",
    "tableau20 = [(31, 119, 180), (174, 199, 232), (255, 127, 14), (255, 187, 120),  \n",
    "             (44, 160, 44), (152, 223, 138), (214, 39, 40), (255, 152, 150),  \n",
    "             (148, 103, 189), (197, 176, 213), (140, 86, 75), (196, 156, 148),  \n",
    "             (227, 119, 194), (247, 182, 210), (127, 127, 127), (199, 199, 199),  \n",
    "             (188, 189, 34), (219, 219, 141), (23, 190, 207), (158, 218, 229)]\n",
    "             \n",
    "# Tableau Color Blind 10\n",
    "tableau20blind = [(0, 107, 164), (255, 128, 14), (171, 171, 171), (89, 89, 89),\n",
    "             (95, 158, 209), (200, 82, 0), (137, 137, 137), (163, 200, 236),\n",
    "             (255, 188, 121), (207, 207, 207)]\n",
    "  \n",
    "# Rescale to values between 0 and 1 \n",
    "for i in range(len(tableau20)):  \n",
    "    r, g, b = tableau20[i]  \n",
    "    tableau20[i] = (r / 255., g / 255., b / 255.)\n",
    "for i in range(len(tableau20blind)):  \n",
    "    r, g, b = tableau20blind[i]  \n",
    "    tableau20blind[i] = (r / 255., g / 255., b / 255.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U9mLTMRc1kkK"
   },
   "outputs": [],
   "source": [
    "SMALL_SIZE = 8\n",
    "MEDIUM_SIZE = 8\n",
    "BIGGER_SIZE = 8\n",
    "\n",
    "plt.rc('font', size=SMALL_SIZE)          # controls default text sizes\n",
    "plt.rc('axes', titlesize=SMALL_SIZE+6)     # fontsize of the axes title\n",
    "plt.rc('axes', labelsize=MEDIUM_SIZE+6)    # fontsize of the x and y labels\n",
    "plt.rc('xtick', labelsize=SMALL_SIZE + 4)    # fontsize of the tick labels\n",
    "plt.rc('ytick', labelsize=SMALL_SIZE+4)    # fontsize of the tick labels\n",
    "plt.rc('legend', fontsize=SMALL_SIZE+2)    # legend fontsize\n",
    "plt.rc('figure', titlesize=BIGGER_SIZE+20)  # fontsize of the figure title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2MAKzoHd1mJj"
   },
   "outputs": [],
   "source": [
    "plot_prop = {\n",
    "    \"pc-alg\": [\"dashed\", tableau20blind[0], \"o\"],\n",
    "    \"ldecc\": [\"dashdot\", tableau20blind[1], \"^\"],\n",
    "    \"ldecc-checks\": [\"solid\", tableau20blind[5], \"s\"],\n",
    "    \"mb-by-mb\": [\"dashdot\", tableau20blind[7], \"D\"],\n",
    "    \"sd-alg\": [\"dotted\", tableau20blind[3], \"v\"],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NmlANMMf2HlC"
   },
   "outputs": [],
   "source": [
    "def get_hausdorff_distance(arr_1, arr_2):\n",
    "\n",
    "  arr_1 = np.array(arr_1)\n",
    "  arr_2 = np.array(arr_2)\n",
    "\n",
    "  def get_distance_point_from_array(point, arr):\n",
    "    return np.min(np.abs(arr - point))\n",
    "\n",
    "  dist_1_to_2 = np.max([get_distance_point_from_array(p, arr_2) for p in arr_1])\n",
    "  dist_2_to_1 = np.max([get_distance_point_from_array(p, arr_1) for p in arr_2])\n",
    "\n",
    "  return max(dist_1_to_2, dist_2_to_1)\n",
    "\n",
    "def plot_mse(results):\n",
    "  plt.figure(figsize=(6, 4))\n",
    "  samples = results[\"sample_sizes\"]\n",
    "\n",
    "  dict_to_name = {\n",
    "      'ldecc': \"LDECC\", \n",
    "      'ldecc-checks': \"LDECC-checks\",\n",
    "      'sd-alg': \"SD\",\n",
    "      'mb-by-mb': \"MB-by-MB\",\n",
    "  }\n",
    "\n",
    "  errors = collections.defaultdict(list)\n",
    "\n",
    "  for N in samples:\n",
    "    for k in dict_to_name.keys():\n",
    "      errors[k].append([])\n",
    "\n",
    "    for i, res in enumerate(results[N]):\n",
    "      for k in dict_to_name.keys():\n",
    "        errors[k][-1].append(get_hausdorff_distance(\n",
    "            res[\"oracle\"][\"ate\"], res[k][\"ate\"]\n",
    "        )**2)\n",
    "\n",
    "      \n",
    "  samples = np.array(samples)\n",
    "  mses = {}\n",
    "  stds = {}\n",
    "  for k in dict_to_name.keys():\n",
    "    errors[k] = np.array(errors[k])\n",
    "    mses[k] = np.median(errors[k], axis=-1)\n",
    "    stds[k] = np.std(errors[k], axis=-1)\n",
    "\n",
    "    plt.plot(samples, mses[k], label=dict_to_name[k],\n",
    "             color=plot_prop[k][1], linestyle=plot_prop[k][0],\n",
    "             marker=plot_prop[k][2])\n",
    "  \n",
    "  plt.title(\"Median SE vs Sample size\")\n",
    "  plt.xlabel(\"Sample size\")\n",
    "  plt.ylabel(\"Median SE (Hausdorff distance)\")\n",
    "  plt.legend()\n",
    "  plt.show()\n",
    "  \n",
    "plot_mse(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7q0qKVkeAZd8"
   },
   "outputs": [],
   "source": [
    "def plot_accuracy(results):\n",
    "\n",
    "  dict_to_name = {\n",
    "      'ldecc': \"LDECC\", \n",
    "      'ldecc-checks': \"LDECC-checks\",\n",
    "      'sd-alg': \"SD\",\n",
    "      'mb-by-mb': \"MB-by-MB\",\n",
    "  }\n",
    "\n",
    "  def is_correct_adj_set(r, key):\n",
    "    return set(frozenset(p) for p in r[\"oracle\"][\"parent_sets\"]) == set(frozenset(p) for p in r[key][\"parent_sets\"])\n",
    "\n",
    "  N = len(results[results[\"sample_sizes\"][0]])\n",
    "  print(\"Total = %d\" % N)\n",
    "\n",
    "  same_count = collections.defaultdict(lambda: 0)\n",
    "  correct_count = collections.defaultdict(lambda: collections.defaultdict(lambda: 0))\n",
    "  \n",
    "  for s in results[\"sample_sizes\"]:\n",
    "    for i in range(N):\n",
    "      parents_or, children_or, uo_or = get_detected_neighbors(results[s][i][\"oracle\"])\n",
    "      \n",
    "      for k in dict_to_name.keys():\n",
    "        if is_correct_adj_set(r=results[s][i], key=k):\n",
    "          correct_count[k][s] += 1\n",
    "        \n",
    "  plt.figure(figsize=(6, 4))\n",
    "\n",
    "  accuracy = {}\n",
    "  samples = results[\"sample_sizes\"]\n",
    "  for k in dict_to_name.keys():\n",
    "    accuracy[k] = np.array([correct_count[k][s] * 100 / N for s in samples])\n",
    "\n",
    "    plt.plot(samples, accuracy[k], label=dict_to_name[k],\n",
    "             color=plot_prop[k][1], linestyle=plot_prop[k][0],\n",
    "             marker=plot_prop[k][2])\n",
    "  \n",
    "  plt.title(\"Accuracy vs Sample size\")\n",
    "  plt.xlabel(\"Sample size\")\n",
    "  plt.ylabel(\"Accuracy (%)\")\n",
    "  plt.legend()\n",
    "  plt.show()\n",
    "\n",
    "plot_accuracy(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FOmTtLHaBBX2"
   },
   "outputs": [],
   "source": [
    "def plot_recall(results):\n",
    "\n",
    "  dict_to_name = {\n",
    "      'ldecc': \"LDECC\", \n",
    "      'ldecc-checks': \"LDECC-checks\",\n",
    "      'sd-alg': \"SD\",\n",
    "      'mb-by-mb': \"MB-by-MB\",\n",
    "  }\n",
    "\n",
    "  N = len(results[results[\"sample_sizes\"][0]])\n",
    "  \n",
    "  correct_count = collections.defaultdict(lambda: collections.defaultdict(lambda: 0))\n",
    "  \n",
    "  for s in results[\"sample_sizes\"]:\n",
    "    for i in range(N):\n",
    "      for k in dict_to_name.keys():\n",
    "        if frozenset(results[s][i][\"graph_true\"].predecessors(\"X\")) in set(frozenset(p) for p in results[s][i][k][\"parent_sets\"]):\n",
    "          correct_count[k][s] += 1\n",
    "        \n",
    "  plt.figure(figsize=(6, 4))\n",
    "  \n",
    "  coverage = {}\n",
    "  samples = results[\"sample_sizes\"]\n",
    "  for k in dict_to_name.keys():\n",
    "    coverage[k] = np.array([correct_count[k][s] * 100 / N for s in samples])\n",
    "\n",
    "    plt.plot(samples, coverage[k], label=dict_to_name[k],\n",
    "             color=plot_prop[k][1], linestyle=plot_prop[k][0],\n",
    "             marker=plot_prop[k][2])\n",
    "  \n",
    "  plt.title(\"Recall vs Sample size\")\n",
    "  plt.xlabel(\"Sample size\")\n",
    "  plt.ylabel(\"Recall (%)\")\n",
    "  plt.legend()\n",
    "  plt.show()\n",
    "\n",
    "plot_recall(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aDhNjpg8g_sG"
   },
   "outputs": [],
   "source": [
    "def plot_number_of_tests(results):\n",
    "\n",
    "  plt.figure(figsize=(6, 4))\n",
    "\n",
    "  dict_to_name = {\n",
    "      'ldecc': \"LDECC\", \n",
    "      'ldecc-checks': \"LDECC-checks\",\n",
    "      'sd-alg': \"SD\",\n",
    "      'mb-by-mb': \"MB-by-MB\",\n",
    "  }\n",
    "\n",
    "  samples = results[\"sample_sizes\"]\n",
    "  samples = np.array(samples)\n",
    "  \n",
    "  tests_count = collections.defaultdict(lambda: collections.defaultdict(lambda: 0))\n",
    "  for k in dict_to_name.keys():\n",
    "    for s in samples:\n",
    "      tests_count[k][s] = np.mean([t[k][\"ci_tests\"][\"total\"] for t in results[s]])\n",
    "\n",
    "  for k in dict_to_name.keys():\n",
    "    plt.plot(samples, [tests_count[k][s] for s in samples], label=dict_to_name[k],\n",
    "             color=plot_prop[k][1], linestyle=plot_prop[k][0],\n",
    "             marker=plot_prop[k][2])\n",
    "\n",
    "  \n",
    "  plt.title(\"Number of CI tests\")\n",
    "  plt.xlabel(\"Sample size\")\n",
    "  plt.ylabel(\"Number of CI tests\")\n",
    "  plt.legend()\n",
    "  plt.show()\n",
    "\n",
    "plot_number_of_tests(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3vokLVf55E5O"
   },
   "source": [
    "The code below plots the distribution of conditional independence (CI) tests performed with a CI oracle by SD, LDECC, and MB-by-MB by repeatedly setting each node as the treatment (for computational reasons, we cap the maximum number of tests for each node to 20000)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "OHmW5je2CVeF"
   },
   "outputs": [],
   "source": [
    "def run_local_discovery_for_each_node(graph_true):\n",
    "\n",
    "  MAX_TESTS = 20000\n",
    "\n",
    "  node_to_tests = {}\n",
    "\n",
    "  for node in graph_true.nodes():\n",
    "    empty_df = pd.DataFrame(columns=list(graph_true.nodes()))\n",
    "\n",
    "    mb_by_mb_alg = MBbyMBAlgorithm(use_ci_oracle=True, graph_true=graph_true,\n",
    "                                   treatment_node=node, max_tests=MAX_TESTS)\n",
    "    result_mb_by_mb = mb_by_mb_alg.run(empty_df)\n",
    "\n",
    "    sd_alg = SequentialDiscoveryAlgorithm(use_ci_oracle=True, graph_true=graph_true,\n",
    "                                          treatment_node=node, max_tests=MAX_TESTS)\n",
    "    result_sd = sd_alg.run(empty_df)\n",
    "\n",
    "    ldecc_alg = LDECCAlgorithm(use_ci_oracle=True, graph_true=graph_true, \n",
    "                              treatment_node=node, outcome_node=node,\n",
    "                               max_tests=MAX_TESTS)\n",
    "    result_ldecc = ldecc_alg.run(empty_df)\n",
    "    \n",
    "    node_to_tests[node] = {\n",
    "        \"mb-by-mb\": mb_by_mb_alg.ci_test_calls[\"total\"],\n",
    "        \"ldecc\": ldecc_alg.ci_test_calls[\"total\"],\n",
    "        \"sd-alg\": sd_alg.ci_test_calls[\"total\"],\n",
    "    }\n",
    "\n",
    "    print(\"Node done: %s, Tests: %s\" % (node, node_to_tests[node]))\n",
    "  \n",
    "  return node_to_tests\n",
    "  \n",
    "node_to_tests = run_local_discovery_for_each_node(graph_true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MMTWlvMAd6_P"
   },
   "outputs": [],
   "source": [
    "def plot_test_statistics(node_to_tests):\n",
    "  plt.figure(figsize=(6, 4))\n",
    "\n",
    "  tests_mb_by_mb = [node_to_tests[k][\"mb-by-mb\"] for k in node_to_tests.keys()]\n",
    "  tests_sd = [node_to_tests[k][\"sd-alg\"] for k in node_to_tests.keys()]\n",
    "  tests_ldecc = [node_to_tests[k][\"ldecc\"] for k in node_to_tests.keys()]\n",
    "\n",
    "  plt.hist([tests_mb_by_mb, tests_sd, tests_ldecc],\n",
    "           label=[\"MB-by-MB\", \"SD\", \"LDECC\"], \n",
    "           color=[plot_prop[\"mb-by-mb\"][1], \n",
    "                  plot_prop[\"sd-alg\"][1], plot_prop[\"ldecc\"][1]])\n",
    "  plt.legend()\n",
    "  plt.title(\"Distribution of CI tests\")\n",
    "  plt.xlabel(\"Number of CI tests\")\n",
    "  plt.ylabel(\"Number of nodes\")\n",
    "  plt.show()\n",
    "\n",
    "\n",
    "plot_test_statistics(node_to_tests)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}