Demonstration_seq.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LOGAN for testing sequential mediation effects\n",
    "\n",
    "### In the following, we conduct a simulation study to demonstrate the usage of LOGAN for sequential pathway inference. The data generating process is the same as the first simulation scenario in Li, Shi, Guo and Jagust (2021)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x115a13e80>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAWsAAAEzCAYAAAD3myLiAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAXrElEQVR4nO3dfawdd33n8fcn3iCqEClh7SQmD5hF1qoUFYOuDCirVVgS1lh0E1BTJa0gy7IyrcgKJP4oC1KhK60U7Ra6IBCpgYhESqGRIIkFLsFEVCFSAdtRnoyh8UYpMbZiGzYkEe0ik+/+ccbs5eY+nad7Zua+X9LRnacz851zbz75+Xd+M5OqQpLUbmfNugBJ0soMa0nqAMNakjrAsJakDjCsJakDDGtJ6gDDWpKGkOTSJN9OcjjJoSTvX2SbJPlUkiNJHk7yunGP+y/G3YEkrTOngQ9W1QNJzgUOJtlXVT+Yt81bga3N6/XAZ5ufI7NlLUlDqKrjVfVAM/0scBi4eMFmVwO31cB3gfOSbB7nuLasJfXWjqRODfmeg3AI+Od5i3ZX1e7Ftk2yBXgt8L0Fqy4Gnpw3f7RZdnzIcn7NsJbUW6eAA0O+J/DPVTW34nbJS4CvAB+oqmdeuJsXGOveHoa1pH47a8je3uefX3GTJGczCOrbq+qri2xyFLh03vwlwLHhCvlN9llL6rezzhrutYIkAb4AHK6qTyyx2R7gXc2okDcAP6+qkbtAwJa1pD5Lhm9Zr+xy4J3AI0kebJZ9GLgMoKpuBvYCO4EjwC+Ad4970DUN640bN9aWLVvW8pCSOurgwYOnqmrT2DuacFhX1f0s3ic9f5sC3jfJ444V1kl2AJ8ENgCfr6qbltt+y5YtHNi/f5xDSlonctZZ/zj+TqbSsp6JkcM6yQbgM8BVDDrT9yfZs2BguCTN1noPa2A7cKSqHgdI8mUGA8ENa0nt0ZOwHucslhr0/RuS7EpyIMmBkydPjnE4SRrSmW6QCY4GmZVxKlvVoO+q2l1Vc1U1t2nT+N8VSNJ6NE43yMQHfUvSxLW4tTyMccJ6P7A1ySuAnwDXAX84kaokaRIcDQJVdTrJjcA9DIbu3VJVhyZWmSRNwnoPa4Cq2svgSh1JaifDWpJazm4QSeoIw1qSWs6WtSR1hGEtSR1gWEtSy9kNIkkdYVhLUsv1qGXdj7OQpJ6zZS2p33rSsjasJfWbYS1JLdejPmvDWlK/GdaS1HK2rCWpIwxrSeoAw1qSWq5H3SD9OAtJWspZZw33WoUktyQ5keTRJdZfkeTnSR5sXn827mnYspbUX9NrWX8R+DRw2zLbfKeq3japAxrWkvptCmFdVfcl2TLxHS/DbhBJ/TZ8N8jGJAfmvXaNeOQ3Jnkoyd8m+Z1xT8OWtaT+Gq0b5FRVzY155AeAl1fVc0l2AncBW8fZoS1rSZqwqnqmqp5rpvcCZyfZOM4+bVlL6rcZDN1LchHwVFVVku0MGsY/HWefhrWk/prSaJAkXwKuYNC/fRT4KHA2QFXdDPw+8CdJTgP/BFxXVTXOMQ1rSf02ndEg16+w/tMMhvZNjGEtqd96cgWjYS2pv3p0ublhLanfDGtI8gTwLPAr4PQExiZK0uTYsv4Nb6qqUxPYjyRNnmEtSS3Xo5b1uGdRwDeTHFzq+vkku85cY3/y5MkxDydJQ5rCLVJnYdyW9eVVdSzJBcC+JD+sqvvmb1BVu4HdAHNzc2MNCpekobU4gIcx1llU1bHm5wngTmD7JIqSpIk40w3Sg5b1yJUlOSfJuWemgbcAiz41QZI0nnG6QS4E7kxyZj9/XVXfmEhVkjQpLW4tD2PksK6qx4HXTLAWSZqsHo0GceiepH4zrCWpAwxrSWo5u0EkqSMMa0lqOVvWktQRhrUkdYBhLUktZzeIJHWEYS1JLdejlnU/zkKSes6WtaR+60nL2rCW1G+GtSS1nH3WktQRU3hSTJJbkpxIsugDVzLwqSRHkjyc5HVjn8a4O5Ck1preY72+COxYZv1bga3Naxfw2bHOA8NaUt9NIaybB4P/bJlNrgZuq4HvAucl2TzOadhnLanfhu+z3pjkwLz53VW1e8h9XAw8OW/+aLPs+LDFnGFYS+qv0b5gPFVVc+MeeZFlNc4ODWtJ/Tab0SBHgUvnzV8CHBtnh/ZZS+qv6X3BuJI9wLuaUSFvAH5eVSN3gYAta0l9N4WWdZIvAVcw6N8+CnwUOBugqm4G9gI7gSPAL4B3j3tMw1pSv00hrKvq+hXWF/C+SR7TbhBJ6gBb1pL6q0eXmxvWkvrNsJaklrNlLUkdYVhLUgf0JKxXPIvFbgWY5KVJ9iV5rPl5/nTLlKQRzO6imIlbTWVf5IW3AvwQcG9VbQXubeYlqX3WS1gvcSvAq4Fbm+lbgWsmXJckja9HLetR+6wvPHOde1UdT3LBUhsm2cXg5ttcdtllIx5OkkbU4gAextTPoqp2V9VcVc1t2rRp2oeTpN+0zlvWTyXZ3LSqNwMnJlmUJE1Ej8ZZj3oWe4AbmukbgLsnU44kTVhPWtarGbr3JeDvgX+d5GiS9wA3AVcleQy4qpmXJE3Jit0gy9wK8M0TrkWSJqtH3SBewSip3wxrSWo5W9aS1BGGtSR1gGEtSS1nN4gkdYRhPYKnn4a77lp6/TXeD6oz/D2qC2xZS1JHGNaS1AGGtSS1XI+6QfpxFpK0lCncyCnJjiQ/SnIkyQuelJXkiiQ/T/Jg8/qzcU/DlrUkDSHJBuAzDG5idxTYn2RPVf1gwabfqaq3Teq4hrWk/ppON8h24EhVPT44RL7M4FGHC8N6ouwGkdRvw3eDbExyYN5r14I9Xgw8OW/+aLNsoTcmeSjJ3yb5nXFPY21b1ued5xjcvvD3qK4YvmV9qqrmllmfRZbVgvkHgJdX1XNJdgJ3AVuHLWQ+W9aS+ms6Tzc/Clw6b/4S4Nj8Darqmap6rpneC5ydZOM4p2JYS+q3yYf1fmBrklckeRFwHYNHHf5akouSpJneziBrfzrOafgFo6T+msIXjFV1OsmNwD3ABuCWqjqU5I+b9TcDvw/8SZLTwD8B11XVwq6SoRjWkvptChfFNF0bexcsu3ne9KeBT0/ymIa1pH7ryRWMhrWk/urR5eaGtaR+M6wlqeVsWUtSR/QkrPtxFpLUc7asJfVbT1rWhrWk/rLPWpI6wrCWpJazZS1JHdGTsF7xLJLckuREkkfnLftYkp/Me77YzumWKUkjmsIzGGdhNZV9EdixyPK/rKptzWvvIuslabamcz/rmVixG6Sq7kuyZfqlSNIUtDiAhzHOWdyY5OGmm+T8pTZKsuvMs8xOnjw5xuEkaUg9almPWtlngVcC24DjwMeX2rCqdlfVXFXNbdq0acTDSdKIehLWI40Gqaqnzkwn+RzwtYlVJEmT1OIAHsZIYZ1kc1Udb2bfDjy63PaSNBPraZx1ki8BVwAbkxwFPgpckWQbg8evPwG8d4o1StK6t5rRINcvsvgLU6hFkiZvvbSsJamz1lM3iCR1mmEtSS1ny1qSOsKwlqQOMKwlqeV61A3Sj7OQpKVM4XLzJDuS/CjJkSQfWmR9knyqWf9wkteNexq2rCX11xRa1kk2AJ8BrgKOAvuT7KmqH8zb7K3A1ub1egb3U3r9OMe1ZS2p3ybfst4OHKmqx6vql8CXgasXbHM1cFsNfBc4L8nmsU5jnDdLUtsVGerF4NYaB+a9di3Y5cXAk/PmjzbLht1mKHaDSOq1558f+i2nqmpumfVZZFmNsM1QbFlL0nCOApfOm78EODbCNkMxrCX1VtWgZT3MaxX2A1uTvCLJi4DrgD0LttkDvKsZFfIG4Ofzbis9ErtBJPXaCN0gy6qq00luBO4BNgC3VNWhJH/crL8Z2AvsBI4AvwDePe5xDWtJvXWmZT35/dZeBoE8f9nN86YLeN8kj2lYS+q1aYT1LBjWknrNsJaklptWN8gsGNaSes2wlqSWs2UtSR1hWEtSy9mylqSOMKwlqQP6EtbeG0SSOsCWtaTess9akjrCsJaklrNlLUkd0ZewXvELxiSXJvl2ksNJDiV5f7P8pUn2JXms+Xn+9MuVpOFM4eEDM7Ga0SCngQ9W1W8DbwDel+RVwIeAe6tqK3BvMy9JrTGlJ8XMxIphXVXHq+qBZvpZ4DCDp/ReDdzabHYrcM20ipSkUfUlrIfqs06yBXgt8D3gwjPPFKuq40kuWOI9u4BdAJdddtk4tUrSUPr0BeOqL4pJ8hLgK8AHquqZ1b6vqnZX1VxVzW3atGmUGiVpZOuqZZ3kbAZBfXtVfbVZ/FSSzU2rejNwYlpFStKo2hzAw1jNaJAAXwAOV9Un5q3aA9zQTN8A3D358iRpdH36gnE1LevLgXcCjyR5sFn2YeAm4I4k7wF+DFw7nRIlaXRtDuBhrBjWVXU/kCVWv3my5UiSFuMVjJJ6q0+jQQxrSb1mWEtSB6x1WCd5KfA3wBbgCeAPqur/LLLdE8CzwK+A01U1t9x+ffiApN6a0WiQYW7F8aaq2rZSUINhLannZhDWU7kVx9p2gzz9NNx119Lrr/H2Ip3h71GrcfvtS6/7oz+a+uFn9AXjqm7FARTwzSQF/FVV7V5up/ZZS+q1EcJ6Y5ID8+Z3LwzSJN8CLlrkvR8Z4jiXV9WxJsz3JflhVd231MaGtaReGyGsT63Uh1xVVy61LsmqbsVRVceanyeS3AlsB5YMa/usJfXWjL5gXPFWHEnOSXLumWngLcCjy+3UsJbUazMI65uAq5I8BlzVzJPkZUn2NttcCNyf5CHg+8DXq+oby+3UbhBJvTWLLxir6qcsciuOpttjZzP9OPCaYfZrWEvqtb5cwWg3iCR1wNq2rM87zzG4feHvUauxBmOpV9KXlrXdIJJ6y7vuSVJHGNaS1HK2rCWpIwxrSeoAw1qSWs5uEEnqCMNaklrOlrUkdYRhLUktZ8takjqiL2HtjZwkqQNsWUvqtb60rA1rSb1ln7UkdURfwnrFPusklyb5dpLDSQ4leX+z/GNJfpLkwea1c/rlStLqzeiBuVOxmpb1aeCDVfVA8zTeg0n2Nev+sqr+YnrlSdJ42hzAw1gxrKvqOHC8mX42yWHg4mkXJkmT0JewHmroXpItwGuB7zWLbkzycJJbkpy/xHt2JTmQ5MDJkyfHKlaShtGnbpBVh3WSlwBfAT5QVc8AnwVeCWxj0PL++GLvq6rdVTVXVXObNm2aQMmStHp9CetVjQZJcjaDoL69qr4KUFVPzVv/OeBrU6lQkka0robuJQnwBeBwVX1i3vLNTX82wNuBR6dToiSNbt2ENXA58E7gkSQPNss+DFyfZBtQwBPAe6dSoSSNYd2EdVXdD2SRVXsnX44kTc4sukGSXAt8DPhtYHtVHVhiux3AJ4ENwOer6qbl9uuNnCRpsh4F3gHct9QGSTYAnwHeCryKQU/Fq5bbqZebS+q1tW5ZV9VhgMHXfUvaDhypqsebbb8MXA38YKk3GNaSemvEbpCNSeZ3Xeyuqt2TqwoYXFj45Lz5o8Drl3uDYS2p10YI61NVNbfcBkm+BVy0yKqPVNXdqzjGYs3uWu4NhrWkXptGN0hVXTnmLo4Cl86bvwQ4ttwbDGtJvdXii2L2A1uTvAL4CXAd8IfLvcHRIJJ6ba0vN0/y9iRHgTcCX09yT7P8ZUn2AlTVaeBG4B7gMHBHVR1abr+2rCX11ixa1lV1J3DnIsuPATvnze9liOtVDGtJvdbSbpChGdaSes2wlqSWa/EXjEMzrCX1mmEtSS3Xp5a1Q/ckqQNsWUvqtb60rA1rSb1mWEtSy/Wpz9qwltRrhrUktZwta0nqCMNakjrAsJaklrMbRJI6wrCWpJazZS1JHWFYS1IH9CWsvZGTJHWALWtJvWWftSR1RF/CesVukCQvTvL9JA8lOZTkz5vlL02yL8ljzc/zp1+uJK3emZb1MK+2Wk2f9f8F/l1VvQbYBuxI8gbgQ8C9VbUVuLeZl6RWWTdhXQPPNbNnN68CrgZubZbfClwzlQolaUTrrWVNkg1JHgROAPuq6nvAhVV1HKD5ecES792V5ECSAydPnpxU3ZK0KusqrKvqV1W1DbgE2J7k1as9QFXtrqq5qprbtGnTqHVK0kjWVVifUVVPA38H7ACeSrIZoPl5YuLVSdIYZtENkuTaZjDG80nmltnuiSSPJHkwyYGV9rua0SCbkpzXTP8WcCXwQ2APcEOz2Q3A3as5EUlaSzNoWT8KvAO4bxXbvqmqtlXVkqF+xmrGWW8Gbk2ygUG431FVX0vy98AdSd4D/Bi4dhX7kqQ1M4uLYqrqMECSie53xbCuqoeB1y6y/KfAmydajSRN2AhhvXFBt8Tuqto9uYp+rYBvJingr1Y6hlcwSuq1EcL61ErdEkm+BVy0yKqPVNVqu4Qvr6pjSS4A9iX5YVUt2XViWEvSkKrqygns41jz80SSO4HtLNPP7V33JPVWWy+KSXJOknPPTANvYfDF5JIMa0m9NoOhe29PchR4I/D1JPc0y1+WZG+z2YXA/UkeAr4PfL2qvrHcfu0GkdRbMxoNcidw5yLLjwE7m+nHgdcMs1/DWlKvtfmqxGEY1pJ6zbCWpJbzSTGS1BGGtSS1nC1rSeoIw1qSOsCwlqSWsxtEkjqiL2Ht5eaS1AG2rCX1lt0gktQRhrUkdYBhLUktZzeIJHWEYS1JLWfLWpI6wrCWpA4wrCWp5ewGkaSOMKwlqeVsWUtSRxjWktQBfQlr77onSR2wpi3rgwcPnspZZ/3jvEUbgVNrWcMqWdfw2lqbdQ2nTXW9fNwd2Gc9oqraNH8+yYGqmlvLGlbDuobX1tqsazhtrWscax3WSf4n8HvAL4H/Dby7qp5eZLsdwCeBDcDnq+qm5fZrN4ik3jrTsh7mNQH7gFdX1e8C/wD814UbJNkAfAZ4K/Aq4Pokr1pup4a1pF5b67Cuqm9W1elm9rvAJYtsth04UlWPV9UvgS8DVy+331mPBtk94+MvxbqG19barGs4ba1rRAfvef75bBzyTS9OcmDe/O6qGvVz+U/A3yyy/GLgyXnzR4HXL7ejVNWINUjS+pTkW8BFi6z6SFXd3WzzEWAOeEctCNok1wL/vqr+czP/TmB7Vf2XpY4565a1JHVOVV253PokNwBvA968MKgbR4FL581fAhxbbp/2WUvSBDWjPP4U+A9V9YslNtsPbE3yiiQvAq4D9iy335mEdZIdSX6U5EiSD82ihqUkeSLJI0keXNBvtdZ13JLkRJJH5y17aZJ9SR5rfp7fkro+luQnzWf2YJKdM6jr0iTfTnI4yaEk72+Wz/QzW6auNnxmL07y/SQPNbX9ebN85n9nHfdp4FxgX/O7vRkgycuS7AVovoC8EbgHOAzcUVWHltvpmvdZN0NW/gG4isE/BfYD11fVD9a0kCUkeQKYq6qZXhiQ5N8CzwG3VdWrm2X/A/hZVd3U/E/u/Kr60xbU9THguar6i7WsZUFdm4HNVfVAknOBg8A1wH9khp/ZMnX9AbP/zAKcU1XPJTkbuB94P/AOZvx3pheaRct66CEr61FV3Qf8bMHiq4Fbm+lbGfxHv6aWqGvmqup4VT3QTD/LoLVyMTP+zJapa+Zq4Llm9uzmVbTg70wvNIuwXmzISiv+eBsFfDPJwSS7Zl3MAhdW1XEYhABwwYzrme/GJA833SQz/Wdzki3Aa4Hv0aLPbEFd0ILPLMmGJA8CJ4B9VdWqz0z/3yzCOossa9P4wcur6nUMrix6X/PPfi3vs8ArgW3AceDjsyokyUuArwAfqKpnZlXHQovU1YrPrKp+VVXbGIxG2J7k1bOoQyubRVgPPWRlLVXVsebnCeBOBt02bfFU0wd6pi/0xIzrAaCqnmr+o38e+Bwz+syaftevALdX1VebxTP/zBarqy2f2RnNvSv+DthBCz4zvdAswnroIStrJck5zZdAJDkHeAvw6PLvWlN7gBua6RuAu2dYy6+d+Q+78XZm8Jk1X5Z9AThcVZ+Yt2qmn9lSdbXkM9uU5Lxm+reAK4Ef0tK/s/VuJlcwNsOU/heDu03dUlX/fc2LWESSf8WgNQ2DC4b+ela1JfkScAWDW1Y+BXwUuAu4A7gM+DFwbVWt6Zd9S9R1BYN/zhfwBPDeM32ea1jXvwG+AzwCnLnDw4cZ9A/P7DNbpq7rmf1n9rsMvkDcwKDhdkdV/bck/5IZ/53phbzcXJI6wCsYJakDDGtJ6gDDWpI6wLCWpA4wrCWpAwxrSeoAw1qSOuD/AT8eVA1tBe9FAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scipy\n",
    "import networkx\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import notears.utils as ut\n",
    "import matplotlib.cm as cm\n",
    "ut.set_random_seed(123456)\n",
    "d1 = 35\n",
    "d2 = d1\n",
    "p = 70\n",
    "W, _ = ut.generate_dag_exam1(p, prob=[0.05, 0.1], edge_coefficient_range=[0.6, 1.4])\n",
    "\n",
    "W_star = abs(W).copy()\n",
    "W_tem = W_star.copy()\n",
    "for j in range(p+1):\n",
    "    for i in range(len(W)):\n",
    "        W_star[i, :] = np.amax(np.minimum(np.outer(W_star[i, :], np.ones(len(W))), abs(W)), axis=0)\n",
    "    W_tem = np.maximum(W_tem, W_star)\n",
    "W_star = W_tem.copy()\n",
    "strength = np.array([[min(W_star[i,0],W_star[j,i],W_star[p+1,j]) for i in np.arange(1, d1+1)] for j in np.arange(d1+1, p+1)])\n",
    "cmap = plt.get_cmap('bwr')\n",
    "cmap.set_bad(color='white')\n",
    "fig, ax = plt.subplots(figsize=(6,6))\n",
    "img = ax.imshow(strength, cmap=cmap, vmin=-2, vmax=2)\n",
    "np.nonzero(strength)\n",
    "plt.colorbar(img, shrink=0.85)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We next run L simulations to compute the p-values of the proposed test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "1\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "Training is over.\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "import pycasso\n",
    "import numpy.random as nr\n",
    "import notears.notears as nt\n",
    "from numpy import linalg as LA\n",
    "import scipy.linalg as slin\n",
    "# define penalized regression with mcp penalty\n",
    "def _mcp_reg(y, x, n_sample, p): # p is the number of columns in x\n",
    "    if p==1:\n",
    "        x = x.reshape((n_sample,1))\n",
    "    lambda_list = np.exp(np.arange(-5,3,0.1))\n",
    "    for j in range(p):\n",
    "        x[:,j] = x[:,j] - np.mean(x[:,j])\n",
    "    std = np.sqrt(np.sum(x * x, axis=0))/np.sqrt(n_sample)\n",
    "    x = x/std\n",
    "    mcp = pycasso.Solver(x, y-np.mean(y), penalty=\"mcp\", gamma=1.25, prec=1e-4, lambdas=lambda_list)\n",
    "    mcp.train()\n",
    "    BIC = np.zeros(len(lambda_list))\n",
    "    for k in range(len(lambda_list)):\n",
    "        BIC[k] = np.sum(np.square(y - np.mean(y) - x @ mcp.coef()['beta'][k])) +                         sum(mcp.coef()['beta'][k]!=0)*np.log(n_sample)\n",
    "    return mcp.coef()['beta'][np.argmin(BIC)]/std\n",
    "\n",
    "# reestimate the coefficients based on mcp penalized regression\n",
    "def _refit1(trt, med1, n_sample, d1, W11): \n",
    "    W11 = W11.transpose()\n",
    "    W1_refit = np.zeros((d1+1,d1+1))\n",
    "    for j in np.arange(1,d1+1):\n",
    "        Indices = W11[j-1,:]!=0\n",
    "        W1_refit[j,np.append(True, Indices)] = _mcp_reg(med1[:,j-1], np.c_[trt, med1[:,Indices]], n_sample, sum(Indices)+1)\n",
    "    return W1_refit\n",
    "\n",
    "def _refit2(trt, med1, med2, n_sample, d1, d2, W22):\n",
    "    res2 = med2 - med2@W22\n",
    "    W22 = W22.transpose()\n",
    "    W2_refit = np.zeros((d2,d1+d2+1))\n",
    "    for j in np.arange(0,d2):\n",
    "        Indices2 = W22[j,:]!=0\n",
    "        gamma = _mcp_reg(res2[:,j], np.c_[trt, med1], n_sample, d1+1)\n",
    "        Indices1 = gamma!=0\n",
    "        Indices1[0] = True\n",
    "        Indices = np.append(Indices1, Indices2)\n",
    "        W2_refit[j,Indices] = _mcp_reg(med2[:,j], np.c_[trt, med1, med2][:,Indices], n_sample, sum(Indices))\n",
    "    return W2_refit\n",
    "\n",
    "def _refit(trt, med1, med2, res, n_sample, d1, d2, W11, W22):\n",
    "    W1_refit = _refit1(trt, med1, n_sample, d1, W11)\n",
    "    W2_refit = _refit2(trt, med1, med2, n_sample, d1, d2, W22)\n",
    "    W2_refit = np.row_stack((W2_refit, np.zeros([1,d1+d2+1])))\n",
    "    W2_refit[d2,:] = _mcp_reg(res, np.c_[trt, med1, med2], n_sample, d1+d2+1)\n",
    "    return W1_refit, W2_refit\n",
    "\n",
    "n_sample = 200\n",
    "L = 1\n",
    "for k in np.arange(1,1+L):\n",
    "    ut.set_random_seed(12345*k)\n",
    "    data = ut.simulate_linear_sem(1, W.transpose(), n_sample)\n",
    "    trt = data[:, 0].reshape(-1)\n",
    "    med1 = np.copy(data[:,1:(d1+1)])\n",
    "    med2 = np.copy(data[:,(d1+1):(p+1)])\n",
    "    res = data[:, p+1].reshape(-1)\n",
    "    \n",
    "    ## Step 1: sample splitting \n",
    "    import numpy.random as nr\n",
    "    import notears.notears as nt\n",
    "    from numpy import linalg as LA\n",
    "    import scipy.linalg as slin\n",
    "    ## set random seed\n",
    "    ut.set_random_seed(12345)\n",
    "    sample1 = nr.choice(n_sample, int(n_sample/2), replace=False)\n",
    "    sample2 = [x for x in range(n_sample) if x not in set(sample1)]\n",
    "    \n",
    "    ## Step 2: Estimation of W_{1,1}\n",
    "    # compute residuals\n",
    "    from sklearn.linear_model import LinearRegression\n",
    "    med1_re1 = med1[sample1,:].copy()\n",
    "    med1_re2 = med1[sample2,:].copy()\n",
    "    for j in range(np.shape(med1_re1)[1]):\n",
    "        reg = LinearRegression().fit(trt[sample1].reshape(np.shape(sample1)[0],1), med1_re1[:,j])\n",
    "        med1_re1[:,j] = med1_re1[:,j] - reg.predict(trt[sample1].reshape(np.shape(sample1)[0],1))\n",
    "    for j in range(np.shape(med1_re2)[1]):\n",
    "        reg = LinearRegression().fit(trt[sample2].reshape(np.shape(sample2)[0],1), med1_re2[:,j])\n",
    "        med1_re2[:,j] = med1_re2[:,j] - reg.predict(trt[sample2].reshape(np.shape(sample2)[0],1))   \n",
    "    # estimate W11\n",
    "    w=1e-3\n",
    "    # tuning parameter\n",
    "    lambda1 = 0.1*np.sqrt(np.log(d1)/(n_sample/2))\n",
    "    # apply notears\n",
    "    W111 = nt.notears_linear_l1(med1_re1, lambda1=lambda1, max_iter=2000, loss_type='l2', w_threshold=w) \n",
    "    W112 = nt.notears_linear_l1(med1_re2, lambda1=lambda1, max_iter=2000, loss_type='l2', w_threshold=w)\n",
    "    while True:\n",
    "        M = np.eye(d1) + W111 * W111 / d1\n",
    "        E = np.linalg.matrix_power(M, d1-1)\n",
    "        h = (E.T * M).sum() - d1\n",
    "        if h<=0:\n",
    "            break\n",
    "        else:\n",
    "            w = 2*w \n",
    "            W111[np.abs(W111)<=w] = 0\n",
    "    w=1e-3\n",
    "    while True:\n",
    "        M = np.eye(d1) + W112 * W112 / d1\n",
    "        E = np.linalg.matrix_power(M, d1-1)\n",
    "        h = (E.T * M).sum() - d1\n",
    "        if h<=0:\n",
    "            break\n",
    "        else:\n",
    "            w = 2*w \n",
    "            W112[np.abs(W112)<=w] = 0\n",
    "    \n",
    "    ## Estimation of W_{2,2}\n",
    "    # compute residuals\n",
    "    med2_re1 = med2[sample1,:].copy()\n",
    "    med2_re2 = med2[sample2,:].copy()\n",
    "    for j in range(np.shape(med2_re1)[1]):\n",
    "        gamma = _mcp_reg(med2_re1[:,j], np.c_[trt[sample1], med1_re1], np.shape(sample1)[0], d1+1)\n",
    "        med2_re1[:,j] = med2_re1[:,j] - np.c_[trt[sample1], med1_re1] @ gamma\n",
    "        med2_re1[:,j] = med2_re1[:,j] - np.mean(med2_re1[:,j])\n",
    "    for j in range(np.shape(med2_re2)[1]):\n",
    "        gamma = _mcp_reg(med2_re2[:,j], np.c_[trt[sample2], med1_re2], np.shape(sample2)[0], d1+1)\n",
    "        med2_re2[:,j] = med2_re2[:,j] - np.c_[trt[sample2], med1_re2] @ gamma\n",
    "        med2_re2[:,j] = med2_re2[:,j] - np.mean(med2_re2[:,j])\n",
    "    \n",
    "    ## estimate W12\n",
    "    w=1e-3\n",
    "    ## tuning parameter\n",
    "    lambda2 = 0.1*np.sqrt(np.log(d2)/(n_sample/2))\n",
    "    ## apply notears\n",
    "    W221 = nt.notears_linear_l1(med2_re1, lambda1=lambda2, max_iter=2000, loss_type='l2', w_threshold=w) \n",
    "    W222 = nt.notears_linear_l1(med2_re2, lambda1=lambda2, max_iter=2000, loss_type='l2', w_threshold=w)\n",
    "    while True:\n",
    "        M = np.eye(d2) + W221 * W221 / d2\n",
    "        E = np.linalg.matrix_power(M, d2-1)\n",
    "        h = (E.T * M).sum() - d2\n",
    "        if h<=0:\n",
    "            break\n",
    "        else:\n",
    "            w = 2*w \n",
    "            W221[np.abs(W221)<=w] = 0\n",
    "    w=1e-3\n",
    "    while True:\n",
    "        M = np.eye(d2) + W222 * W222 / d2\n",
    "        E = np.linalg.matrix_power(M, d2-1)\n",
    "        h = (E.T * M).sum() - d2\n",
    "        if h<=0:\n",
    "            break\n",
    "        else:\n",
    "            w = 2*w \n",
    "            W222[np.abs(W222)<=w] = 0\n",
    "            \n",
    "    ## Step 4: Refitting the weighted adjancency matrix\n",
    "    W1_refit1, W2_refit1 = _refit(trt[sample1], med1[sample1,:], med2[sample1,:], res[sample1], np.shape(sample1)[0], d1, d2, W111, W221)\n",
    "    W1_refit2, W2_refit2 = _refit(trt[sample2], med1[sample2,:], med2[sample2,:], res[sample2], np.shape(sample2)[0], d1, d2, W112, W222)\n",
    "    W_refit1 = np.row_stack((np.c_[W1_refit1, np.zeros((d1+1,d2))], W2_refit1))\n",
    "    W_refit1 = np.c_[W_refit1, np.zeros((d1+d2+2,))]\n",
    "    W_refit2 = np.row_stack((np.c_[W1_refit2, np.zeros((d1+1,d2))], W2_refit2))\n",
    "    W_refit2 = np.c_[W_refit2, np.zeros((d1+d2+2,))]\n",
    "    \n",
    "    np.savetxt('W_refit1'+str(k)+'.csv', W_refit1, delimiter=',')\n",
    "    np.savetxt('W_refit2'+str(k)+'.csv', W_refit2, delimiter=',')\n",
    "    print(k)\n",
    "    \n",
    "import numpy.random as nr\n",
    "import notears.notears as nt\n",
    "from numpy import linalg as LA\n",
    "import scipy.linalg as slin\n",
    "\n",
    "## determine the set of ancestors for each node\n",
    "def _ancestor(W_refit): \n",
    "    p = len(W_refit)\n",
    "    B0 = np.abs(W_refit)>0\n",
    "    B = B0.copy()\n",
    "    B_tem = B0.copy()\n",
    "    for j in range(p-1):\n",
    "        B = np.dot(B, B0)\n",
    "        B_tem = np.maximum(B_tem, B)\n",
    "    B_tem[-1,0:-1] = True\n",
    "    B_tem[0:-1,0] = True\n",
    "    return B_tem\n",
    "\n",
    "## calculate the decorrelated score statistic\n",
    "def _decor_score(x, W_refit, B, L=1000):\n",
    "    # B denote the set of ancestors\n",
    "    n_sample, p = np.shape(x)\n",
    "    W_ds = W_refit.copy()\n",
    "    Boot_W = np.zeros((p, p, L))\n",
    "    for j in range(p):\n",
    "        x[:,j] = x[:,j] - np.mean(x[:,j])\n",
    "    for i in np.arange(1,p):\n",
    "        for j in np.arange(0,p-1):\n",
    "            if W_refit[i,j]!=0:\n",
    "                Indices = B[i,:].copy()\n",
    "                Indices[0] = True\n",
    "                Indices[i] = False\n",
    "                Indices[j] = False\n",
    "                Indices[p-1] = False\n",
    "                gamma = np.zeros(p)\n",
    "                if (sum(Indices)>0):\n",
    "                    gamma[Indices] = _mcp_reg(x[:,j], x[:,Indices], n_sample, sum(Indices))\n",
    "                tem_vec1 = x[:,j] - x @ gamma\n",
    "                tem_vec2 = x[:,i] - x @ W_refit[i,:] + x[:,j]*W_refit[i,j]\n",
    "                W_ds[i,j] = np.dot(tem_vec1, tem_vec2)/np.dot(x[:,j], tem_vec1)\n",
    "                for l in range(L):\n",
    "                    Boot_W[i,j,l] = np.dot(tem_vec1, nr.normal(size=n_sample))/np.dot(x[:,j], tem_vec1)\n",
    "    return W_ds, Boot_W\n",
    "\n",
    "## calculate W_star\n",
    "def _W_star(W, p):\n",
    "    W_star = abs(W).copy()\n",
    "    W_tem = W_star.copy()\n",
    "    for j in range(p+1):\n",
    "        for i in range(len(W)):\n",
    "            W_star[i, :] = np.amax(np.minimum(np.outer(W_star[i, :], np.ones(len(W))), abs(W)), axis=0)\n",
    "        W_tem = np.maximum(W_tem, W_star)\n",
    "    return W_tem.copy()\n",
    "\n",
    "import pycasso\n",
    "import numpy.random as nr\n",
    "import notears.notears as nt\n",
    "from numpy import linalg as LA\n",
    "import scipy.linalg as slin\n",
    "# define penalized regression with mcp penalty\n",
    "def _mcp_reg(y, x, n_sample, p): # p is the number of columns in x\n",
    "    if p==1:\n",
    "        x = x.reshape((n_sample,1))\n",
    "    lambda_list = np.exp(np.arange(-5,3,0.1))\n",
    "    for j in range(p):\n",
    "        x[:,j] = x[:,j] - np.mean(x[:,j])\n",
    "    std = np.sqrt(np.sum(x * x, axis=0))/np.sqrt(n_sample)\n",
    "    x = x/std\n",
    "    mcp = pycasso.Solver(x, y-np.mean(y), penalty=\"mcp\", gamma=1.25, prec=1e-4, lambdas=lambda_list)\n",
    "    mcp.train()\n",
    "    BIC = np.zeros(len(lambda_list))\n",
    "    for k in range(len(lambda_list)):\n",
    "        BIC[k] = np.sum(np.square(y - np.mean(y) - x @ mcp.coef()['beta'][k])) +sum(mcp.coef()['beta'][k]!=0)*np.log(n_sample)\n",
    "    return mcp.coef()['beta'][np.argmin(BIC)]/std\n",
    "\n",
    "\n",
    "# In[ ]:\n",
    "\n",
    "\n",
    "n_sample = 200\n",
    "test11 = np.zeros((d1, 2, L))\n",
    "test21 = np.zeros((d1, d2, 2, L))\n",
    "test31 = np.zeros((d2, 2, L))\n",
    "test41 = np.zeros((d1, 2, L))\n",
    "for k in np.arange(1,1+L):\n",
    "    ut.set_random_seed(12345*k)\n",
    "    data = ut.simulate_linear_sem(1, W.transpose(), n_sample)\n",
    "    trt = data[:, 0].reshape(-1)\n",
    "    med1 = np.copy(data[:,1:(d1+1)])\n",
    "    med2 = np.copy(data[:,(d1+1):(p+1)])\n",
    "    res = data[:, p+1].reshape(-1)\n",
    "    \n",
    "    ## sample splitting \n",
    "    import numpy.random as nr\n",
    "    import notears.notears as nt\n",
    "    from numpy import linalg as LA\n",
    "    import scipy.linalg as slin\n",
    "    \n",
    "    ## set random seed\n",
    "    ut.set_random_seed(12345)\n",
    "    sample1 = nr.choice(n_sample, int(n_sample/2), replace=False)\n",
    "    sample2 = [x for x in range(n_sample) if x not in set(sample1)]\n",
    "    \n",
    "    ## load matrices\n",
    "    W_refit1 = np.loadtxt('W_refit1'+str(k)+'.csv', delimiter=',')\n",
    "    W_refit2 = np.loadtxt('W_refit2'+str(k)+'.csv', delimiter=',')\n",
    "    \n",
    "    critic1 = np.zeros((1000,))\n",
    "    critic2 = np.zeros((1000,))\n",
    "    data1 = np.c_[trt, med1, med2, res][sample1,:]\n",
    "    data2 = np.c_[trt, med1, med2, res][sample2,:]\n",
    "    sigma1 = np.mean(np.square(data2 - data2 @ W_refit1.transpose()))\n",
    "    sigma2 = np.mean(np.square(data1 - data1 @ W_refit2.transpose()))\n",
    "    sigma = np.sqrt((sigma1+sigma2)/2)\n",
    "    \n",
    "    ## compute individual p-values\n",
    "    ## compute p-values using one half of the data\n",
    "    B1 = _ancestor(W_refit1)\n",
    "    W_ds2, Boot_W2 = _decor_score(data2, W_refit1, B1)\n",
    "    W_est_star = _W_star(W_ds2.copy(), d1+d2+2)\n",
    "\n",
    "    for q1 in np.arange(1, d1+1):\n",
    "        Ind1 = np.full(d1+d2+2, True)\n",
    "        Ind2 = np.full(d1+d2+2, True)\n",
    "        Ind1[B1[q1,:]==False] = False\n",
    "        Ind1[q1] = True\n",
    "        Ind1[B1[:,0]==False] = False\n",
    "        Ind2[B1[:,0]==False] = False\n",
    "        Ind2[0] = True\n",
    "        Ind2[B1[q1,:]==False] = False\n",
    "        for l in range(1000):\n",
    "            Boot_W = np.copy(Boot_W2[:,:,l])\n",
    "            Boot_W[W_refit1==0] = 0\n",
    "            if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                critic1[l] = 0\n",
    "            else:\n",
    "                critic1[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma  \n",
    "        test11[q1-1,0,k-1] = np.mean(critic1>=W_est_star[q1,0])\n",
    "\n",
    "    for q2 in np.arange(d1+1, d1+d2+1):\n",
    "        Ind1 = np.full(d1+d2+2, True)\n",
    "        Ind2 = np.full(d1+d2+2, True)\n",
    "        Ind1[B1[:,q2]==False] = False\n",
    "        Ind1[0:-1][B1[d1+d2+1,-1]==False] = False\n",
    "        Ind2[B1[:,q2]==False] = False\n",
    "        Ind2[q2] = True\n",
    "        Ind2[B1[d1+d2+1,:]==False] = False\n",
    "        for l in range(1000):\n",
    "            Boot_W = np.copy(Boot_W2[:,:,l])\n",
    "            Boot_W[W_refit1==0] = 0\n",
    "            if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                critic1[l] = 0\n",
    "            else:\n",
    "                critic1[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma\n",
    "        test31[q2-d1-1,0,k-1] = np.mean(critic1>=W_est_star[d1+d2+1,q2])\n",
    "\n",
    "    for q1 in np.arange(1, d1+1):\n",
    "        for q2 in np.arange(d1+1, d1+d2+1):\n",
    "            Ind1 = np.full(d1+d2+2, True)\n",
    "            Ind2 = np.full(d1+d2+2, True)\n",
    "            Ind1[B1[q2,:]==False] = False\n",
    "            Ind1[q2] = True\n",
    "            Ind1[B1[:,q1]==False] = False\n",
    "            Ind2[B1[:,q1]==False] = False\n",
    "            Ind2[q1] = True\n",
    "            Ind2[B1[q2,:]==False] = False\n",
    "            for l in range(1000):\n",
    "                Boot_W = np.copy(Boot_W2[:,:,l])\n",
    "                Boot_W[W_refit1==0] = 0\n",
    "                if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                    critic1[l] = 0\n",
    "                else:\n",
    "                    critic1[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma\n",
    "            test21[q1-1,q2-d1-1,0,k-1] = np.mean(critic1>=W_est_star[q2,q1])\n",
    "\n",
    "    for q1 in np.arange(1, d1+1):\n",
    "        Ind1 = np.full(d1+d2+2, True)\n",
    "        Ind2 = np.full(d1+d2+2, True)\n",
    "        Ind1[B1[:,q1]==False] = False\n",
    "        Ind1[0:-1][B1[d1+d2+1,-1]==False] = False\n",
    "        Ind2[B1[:,q1]==False] = False\n",
    "        Ind2[q1] = True\n",
    "        Ind2[B1[d1+d2+1,:]==False] = False\n",
    "        for l in range(1000):\n",
    "            Boot_W = np.copy(Boot_W2[:,:,l])\n",
    "            Boot_W[W_refit1==0] = 0\n",
    "            if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                critic1[l] = 0\n",
    "            else:\n",
    "                critic1[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma\n",
    "        test41[q1-1,0,k-1] = np.mean(critic1>=W_est_star[d1+d2+1,q1])\n",
    "    \n",
    "    ## compute p-values on another set of data\n",
    "    B2 = _ancestor(W_refit2)\n",
    "    W_ds1, Boot_W1 = _decor_score(data1, W_refit2, B2)\n",
    "    W_est_star = _W_star(W_ds1.copy(), d1+d2+2)\n",
    "\n",
    "    for q1 in np.arange(1, d1+1):\n",
    "        Ind1 = np.full(d1+d2+2, True)\n",
    "        Ind2 = np.full(d1+d2+2, True)\n",
    "        Ind1[B2[q1,:]==False] = False\n",
    "        Ind1[q1] = True\n",
    "        Ind1[B2[:,0]==False] = False\n",
    "        Ind2[B2[:,0]==False] = False\n",
    "        Ind2[0] = True\n",
    "        Ind2[B2[q1,:]==False] = False\n",
    "        for l in range(1000):\n",
    "            Boot_W = np.copy(Boot_W1[:,:,l])\n",
    "            Boot_W[W_refit2==0] = 0\n",
    "            if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                critic2[l] = 0\n",
    "            else:\n",
    "                critic2[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma  \n",
    "        test11[q1-1,1,k-1] = np.mean(critic2>=W_est_star[q1,0])\n",
    "\n",
    "    for q2 in np.arange(d1+1, d1+d2+1):\n",
    "        Ind1 = np.full(d1+d2+2, True)\n",
    "        Ind2 = np.full(d1+d2+2, True)\n",
    "        Ind1[B2[:,q2]==False] = False\n",
    "        Ind1[0:-1][B2[d1+d2+1,-1]==False] = False\n",
    "        Ind2[B2[:,q2]==False] = False\n",
    "        Ind2[q2] = True\n",
    "        Ind2[B2[d1+d2+1,:]==False] = False\n",
    "        for l in range(1000):\n",
    "            Boot_W = np.copy(Boot_W1[:,:,l])\n",
    "            Boot_W[W_refit2==0] = 0\n",
    "            if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                critic2[l] = 0\n",
    "            else:\n",
    "                critic2[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma\n",
    "        test31[q2-d1-1,1,k-1] = np.mean(critic2>=W_est_star[d1+d2+1,q2])\n",
    "\n",
    "    for q1 in np.arange(1, d1+1):\n",
    "        for q2 in np.arange(d1+1, d1+d2+1):\n",
    "            Ind1 = np.full(d1+d2+2, True)\n",
    "            Ind2 = np.full(d1+d2+2, True)\n",
    "            Ind1[B2[q2,:]==False] = False\n",
    "            Ind1[q2] = True\n",
    "            Ind1[B2[:,q1]==False] = False\n",
    "            Ind2[B2[:,q1]==False] = False\n",
    "            Ind2[q1] = True\n",
    "            Ind2[B2[q2,:]==False] = False\n",
    "            for l in range(1000):\n",
    "                Boot_W = np.copy(Boot_W1[:,:,l])\n",
    "                Boot_W[W_refit2==0] = 0\n",
    "                if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                    critic2[l] = 0\n",
    "                else:\n",
    "                    critic2[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma\n",
    "            test21[q1-1,q2-d1-1,1,k-1] = np.mean(critic2>=W_est_star[q2,q1])\n",
    "\n",
    "    for q1 in np.arange(1, d1+1):\n",
    "        Ind1 = np.full(d1+d2+2, True)\n",
    "        Ind2 = np.full(d1+d2+2, True)\n",
    "        Ind1[B2[:,q1]==False] = False\n",
    "        Ind1[0:-1][B2[d1+d2+1,-1]==False] = False\n",
    "        Ind2[B2[:,q1]==False] = False\n",
    "        Ind2[q1] = True\n",
    "        Ind2[B2[d1+d2+1,:]==False] = False\n",
    "        for l in range(1000):\n",
    "            Boot_W = np.copy(Boot_W1[:,:,l])\n",
    "            Boot_W[W_refit2==0] = 0\n",
    "            if (sum(Ind1)==0)|(sum(Ind2)==0):\n",
    "                critic2[l] = 0\n",
    "            else:\n",
    "                critic2[l] = np.amax(abs(Boot_W[np.ix_(Ind1,Ind2)]))*sigma\n",
    "        test41[q1-1,1,k-1] = np.mean(critic2>=W_est_star[d1+d2+1,q1])\n",
    "            \n",
    "    print(k)\n",
    "    \n",
    "test10 = 2*np.minimum(test11[:,0,:], test11[:,1,:])\n",
    "test20 = 2*np.minimum(test21[:,:,0,:], test21[:,:,1,:])\n",
    "test30 = 2*np.minimum(test31[:,0,:], test31[:,1,:])\n",
    "pvalue = np.zeros((d1,d2,L))\n",
    "for k in range(L):\n",
    "    pvalue[:,:,k] = np.array([[max(test10[i,k],test20[i,j,k],test30[j,k]) for i in np.arange(0, d1)] for j in np.arange(0, d2)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualise the p-values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x115ec34c0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAWMAAAEzCAYAAADkTGIWAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAWYUlEQVR4nO3df+xldX3n8eeLEaOLZMHOQMcBCnVnu0tMwfYrmNjdpVq6A/+gm+0usFE0NiNZ2WjSPyT9Y8HdbMJuq203SyWjTqSNSkjEMmtmSwnRpaZqmTGIDCMyoRSGmTBQaJWalgy+9497vs31u98f936/937Pj3k+kpPvPeee7+e8vX553/e8z+eck6pCktSu09oOQJJkMpakTjAZS1IHmIwlqQNMxpLUASZjSeoAk7EkTSnJ3iQnkjy6wvtJ8j+THEnySJJfWGtMk7EkTe9zwK5V3r8K2Nksu4FPrTWgyViSplRVDwIvrrLLNcAf1Mg3gbOSbF9tzNfMMkBJ6pJ/kjPqR7w61e8c5+8PAX83tmlPVe2Z8tA7gGfG1o82246v9AsmY0mD9SNe5UP8zFS/cyvf/7uqWtjgobPMtlXvPWEyljRYobVe7FHg/LH184Bjq/2CPWNJg3balMuM7APe18yqeDvwN1W1YosCrIwlDdi8KuMkXwSuALYmOQrcApwOUFV3APuBq4EjwI+AD6w15qYm461bz6oLL3zTZh5SUk8dPHj4harattFx5pGMq+q6Nd4v4MPTjLmhZJxkF/B7wBbgM1V122r7X3jhmzhw4A83ckhJp4hk4S9nMU5ferHrjjPJFuB2RpObLwauS3LxrAKTpI3KOpa2bORL4zLgSFU9WVWvAHcxmugsSZ3R0gm8qW3k2CtNav4JSXYnOZDkwPPPv7SBw0nS9E6FZDzRpOaq2lNVC1W1sG3b2Rs4nCQN10ZO4E09qVmSNlOLF31MbSNxPgTsTHJRktcC1zKa6CxJndGXNsW6K+OqOpnkJuA+RlPb9lbVoZlFJkkb1KfKeEPzjKtqP6MrTSSpk06JZCxJXWcylqSWnTJtCknqOpOxJLXMyliSOsJkLEkd0ObNf6ZhMpY0WLYpJKkj+pKM+xKnJA2albGkwbJNIUkdYTKWpA4wGUtSy2xTSFJHmIwlqWVWxpLUESZjSeoAk7EktSx4bwpJ6gQrY0nqAJOxJLXM2RSS1BGZtmlccwljTX350pCkQbMyljRop2XKUrelythkLGmwwjraFC0xGUsatJ7kYpOxpGHLtG2KlpiMJQ1XbFNIUiecEsk4yVPAD4FXgZNVtTCLoCRpFsI6ZlO0ZBaV8S9X1QszGEeSZq4nhbFtCknD1pc2xUavwCvgT5IcTLJ7uR2S7E5yIMmB559/aYOHk6TpJNMtbdloZfyOqjqW5Bzg/iTfq6oHx3eoqj3AHoCFhYv70byRNAhJ9WZq24Yq46o61vw8AXwZuGwWQUnSrJyW6ZbW4lzvLyY5I8mZi6+BXwUenVVgknQq2Uib4lzgyxk1WV4DfKGq/ngmUUnSjPTlBN66k3FVPQlcMsNYJGnm0tZt2Kbk/YwlDdbiXdtmPZsiya4kjyc5kuTmZd7/x0n+d5LvJDmU5ANrjWkyljRos07GSbYAtwNXARcD1yW5eMluHwYeq6pLgCuATyR57WrjetGHpOHKXC6Hvgw40rRqSXIXcA3w2Ng+BZyZ0Um1NwAvAidXG9RkLGnQ1nECb2uSA2Pre5rrJRbtAJ4ZWz8KXL5kjP8F7AOOAWcC/76qfrzaQU3GkgYrrOveFC+scdOz5YZcWn7/a+Bh4J3AmxldFPenVfWDlQa1Zyxp0Bavwpt0mcBR4Pyx9fMYVcDjPgDcUyNHgL8A/tlqg5qMJQ3aHGZTPATsTHJRc1LuWkYtiXFPA+8aHT/nAj8HPLnaoLYpJA3arC9xrqqTSW4C7gO2AHur6lCSG5v37wD+K/C5JN9l1Nb42Fq3GjYZSxqsUbU7+4s+qmo/sH/JtjvGXh9jdIuIiZmMJQ1aT66GtmcsSV1gZSxp0AZ/oyBJ6r7+3FzeZCxpsEZPh247ismYjCUNmm0KSWpbyw8ZnYbJWNKg9eXm8iZjSYNmZSxJLQuQnpzBMxlLGq5AenJpm8lY0qDZppCk1qU3E41NxpKGq0dtip6EKUnDZmUsadDSk6axyVjSYI2mtrUdxWRMxpKGzcpYklrWoxN4JmNJg9aXK/DW/M5IsjfJiSSPjm17Y5L7kzzR/Dx7vmFK0vok0y1tmaSA/xywa8m2m4EHqmon8ECzLknd0rQpplnasuahq+pB4MUlm68B7mxe3wm8e8ZxSdJsnJbplpast2d8blUdB6iq40nOWWnHJLuB3QAXXPDT6zycJE0v9GYyxfyvwKuqPVW1UFUL27bZWpa0iRJy2nRLW9ZbGT+XZHtTFW8HTswyKEmalb5MbVtvmPuAG5rXNwD3ziYcSZqtJFMtbZlkatsXgW8AP5fkaJIPArcBVyZ5AriyWZckrdOabYqqum6Ft94141gkabZCb+5N6RV4kgarT7MpTMaSBq0vl0ObjCUNlzcKkqSO6EmfwmQsadCsjCWpbbFnvKxjB5/k47l+xfdvqS9sYjTaCP9/VB84m0KSOqHd+01Mw2QsabjSLD1gMpY0aJ7Ak6QO6EuboiffGZI0bFbGkoar5YeMTsNkLGmwQn/aFJuajN/0iz/LLQf+cDMPqTlxLrF6oyfNWCtjScMVWn3i8zRMxpKGrSeVcU/ClKR1WKyMp1kmGTbZleTxJEeS3LzCPlckeTjJoST/d60xrYwlDduMS84kW4DbGT3/8yjwUJJ9VfXY2D5nAb8P7Kqqp5Ocs8lhSlKHZMqqeLLK+DLgSFU9WVWvAHcB1yzZ53rgnqp6GqCqTqw1qMlY0rBNn4y3JjkwtuxeMuIO4Jmx9aPNtnH/FDg7ydeSHEzyvrXCtE0hadimLzlfqKqFVd5frnyuJeuvAX4ReBfweuAbSb5ZVd9faVCTsaThms/UtqPA+WPr5wHHltnnhar6W+BvkzwIXAKsmIxtU0gattOmXNb2ELAzyUVJXgtcC+xbss+9wL9I8pok/wi4HDi82qBWxpI0hao6meQm4D5gC7C3qg4lubF5/46qOpzkj4FHgB8Dn6mqR1cb12QsabjmdAVeVe0H9i/ZdseS9d8CfmvSMU3GkoatH1dDm4wlDZz3ppCklmXiCzlaZzKWNGw9mTO2ZphJ9iY5keTRsW23Jnm2uQnGw0munm+YkrROc7hR0FzCnGCfzwG7ltn+O1V1abPsX+Z9SWpXRk+HnmZpy5ptiqp6MMmF8w9FkuagJz3jjXwP3JTkkaaNcfZKOyXZvXjDjeeff2kDh5OkKYV5XIE3F+s99KeANwOXAseBT6y0Y1XtqaqFqlrYtm3FnC1J89GTnvG6ZlNU1XOLr5N8GvjKzCKSpFkacpsiyfax1fcAq15zLUmt6FGbYs3KOMkXgSsY3XD5KHALcEWSSxndw/Mp4ENzjFGSBm+S2RTXLbP5s3OIRZJmzCvwJKl9i22KHjAZSxo2K2NJ6gArY0lq2ZxuLj8PJmNJw2ZlLEktszKWpI4wGUtSB9imkKSW+dglSeoIK2NJ6oD0ozLuyXeGJA2blbGk4Uqz9IDJWNKw9aRNYTKWNGz9yMUmY0kDZ2UsSS3zfsaS1BFWxpLUAf3IxSZjSUMWK2NJ6oR+5GKTsaQBC1bGktQJzqaQpA7oSWXck+8MSRo2K2NJw9aPwthkLGnAPIEnSd3Qk1xsMpY0cD3JxmuewEtyfpKvJjmc5FCSjzTb35jk/iRPND/Pnn+4kjSlTLm0ZJLZFCeB36iqfw68HfhwkouBm4EHqmon8ECzLkndsfh06GmWlqyZjKvqeFV9u3n9Q+AwsAO4Briz2e1O4N3zClKS1q0nlfFUPeMkFwJvBb4FnFtVx2GUsJOcs8Lv7AZ2A1xwwU9vJFZJmt5QesaLkrwB+BLw0ar6waS/V1V7qmqhqha2bbOtLGmT9aQynigZJzmdUSL+fFXd02x+Lsn25v3twIn5hChJ67Q4z3iaZZJhk11JHk9yJMmK58uSvC3Jq0n+7VpjTjKbIsBngcNV9cmxt/YBNzSvbwDuXWssSdp0M66Mk2wBbgeuAi4GrmsmNSy3338H7pskzEkq43cA7wXemeThZrkauA24MskTwJXNuiR1y+xnU1wGHKmqJ6vqFeAuRhMalvpPjDoKE3UN1jyBV1VfZ+Xvi3dNchBJ6pGtSQ6Mre+pqj1j6zuAZ8bWjwKXjw+QZAfwHuCdwNsmOahX4EkatulPyr1QVQtTjlhL1n8X+FhVvZoJ+9AmY0nDNZ8bBR0Fzh9bPw84tmSfBeCuJhFvBa5OcrKq/milQU3GkoZt9tPVHgJ2JrkIeBa4Frh+fIequugfDp98DvjKaokYTMaSBm32T4euqpNJbmI0S2ILsLeqDiW5sXn/jvWMazKWNGxzuJCjqvYD+5dsWzYJV9X7JxnTZKx1+XiuX/G9W+oLmxiJuqz1v5PQ6s1/pmEyljRsPbk3hclY0rCZjCWpA0zGktS2QCa+OWWrTMaShssTeJLUET1pU/SjfpekgbMy1ro4l1iTaP/vxJ6xJHVDT9oUJmNJw+UJPEnqAtsUktQNtikkqQNMxpLUsmCbQpLaN/ETn1tnMpY0bLYpJKkDbFNIUsvm83TouTAZSxqw/vSM+1G/S9LAWRlLGjZ7xpLUAfaMJallPTqBt2b9nuT8JF9NcjjJoSQfabbfmuTZJA83y9XzD1eSppFRMp5mackklfFJ4Deq6ttJzgQOJrm/ee93quq35xeeJG3QaQPpGVfVceB48/qHSQ4DO+YdmCTNxFDaFOOSXAi8FfhWs+mmJI8k2Zvk7BV+Z3eSA0kOPP/8SxsKVpKmstgz7kGbYuJknOQNwJeAj1bVD4BPAW8GLmVUOX9iud+rqj1VtVBVC9u2LZuvJWlOmpvLT7O0ZKLZFElOZ5SIP19V9wBU1XNj738a+MpcIpSkjRjKFXhJAnwWOFxVnxzbvn1st/cAj84+PEnaoJ60KSapjN8BvBf4bpKHm22/CVyX5FKggKeAD80lQklarwzoGXhV9XVGbfCl9s8+HEmasSHOppAkzYeXQ0satp5UxiZjScM2lCvwJKm/wvKnvLrHZCxpuHp01zaTsaRhG8rUNknqNytjSWpZu1fVTcNkLGnYbFNIUhdYGUtS+2xTSFLbQl/u+mAyljRcPZpn3I+vDEnqkCS7kjye5EiSm5d5/z80j6R7JMmfJblkrTGtjCUN24wr4yRbgNuBK4GjwENJ9lXVY2O7/QXwr6rqpSRXAXuAy1cb18pY0sBlymVNlwFHqurJqnoFuAu4ZnyHqvqzqlp8AvM3gfPWGtTKWNKAretJH1uTHBhb31NVe8bWdwDPjK0fZfWq94PA/1nroCZjScM2fZvihapaWG3EZbbV8ofOLzNKxr+01kFNxpIGbuazKY4C54+tnwcc+/+Omvw88Bngqqr6q7UGNRlLGrC5PJD0IWBnkouAZ4Frget/4qjJBcA9wHur6vuTDGoyljRcgcx4NkVVnUxyE3AfsAXYW1WHktzYvH8H8J+BnwJ+vzn+yTVaHyZjSUM3+4s+qmo/sH/JtjvGXv868OvTjGkyljRs3rVNktrmM/AkqRt6cm8Kk7GkYetJm6IfUUrSwFkZSxo42xSS1LL+PJB0zTZFktcl+fMk30lyKMnHm+1vTHJ/kiean2fPP1xJmkIY9YynWVoyyZH/HnhnVV0CXArsSvJ24GbggaraCTzQrEtSx8z8FppzsWYyrpGXm9XTm6UY3b/zzmb7ncC75xKhJG1EMt3Skolq8iRbkjwMnADur6pvAedW1XGA5uc5K/zu7iQHkhx4/vmXlttFkuZk8YGk0yztmOjIVfVqVV3K6FZxlyV5y6QHqKo9VbVQVQvbttlWlrTJhlQZL6qqvwa+BuwCnkuyHaD5eWLm0UnSRg0lGSfZluSs5vXrgV8BvgfsA25odrsBuHdeQUrS+vSnTTHJPOPtwJ3NE1FPA+6uqq8k+QZwd5IPAk8DvzbHOCVpfXoyz3jNZFxVjwBvXWb7XwHvmkdQkjQ7A0nGktRbmctjl+aiH1FK0sBZGUsatqH0jCWp30zGktS+nvSMTcaSBsxn4ElSN9gzlqQusE0hSe0KVsaS1L7+XPRhMpY0cFbGktQ+2xSS1LbFW2h2n8lY0rD1pDLux1eGJA2clbGkgetHzWkyljRsPWlTmIwlDVfLDxmdhslY0sDZppCk9lkZS1IXmIwlqWXem0KSOsLKWJLaZ2UsSW3zsUuS1A3OppCklgXbFJLUDf2ojPvxlSFJA7eplfHBg4dfSBb+cmzTVuCFzYxhQsY1va7GZlzT6VJcP7PxIbw3xbKqatv4epIDVbWwmTFMwrim19XYjGs6XY1rY2bfAEiyC/g9YAvwmaq6bcn7ad6/GvgR8P6q+vbmRilJXbJ457ZJlzWHyxbgduAq4GLguiQXL9ntKmBns+wGPrXWuCZjSQO2+Ay8aZY1XQYcqaonq+oV4C7gmiX7XAP8QY18EzgryfbVBm17NsWelo+/EuOaXldjM67pdDWudTl48PB9Oe1tW6f8tdclOTC2vqeqxj+XHcAzY+tHgcuXjLHcPjuA4ysdtNVkvOR/YGcY1/S6GptxTaerca1XVe2aw7DL9TJqHfv8BNsUkjSdo8D5Y+vnAcfWsc9PMBlL0nQeAnYmuSjJa4FrgX1L9tkHvC8jbwf+pqpWbFFAS8k4ya4kjyc5kuTmNmJYSZKnknw3ycNL+kabHcfeJCeSPDq27Y1J7k/yRPPz7I7EdWuSZ5vP7OEkV7cQ1/lJvprkcJJDST7SbG/1M1slri58Zq9L8udJvtPE9vFme+t/Z11WVSeBm4D7gMPA3VV1KMmNSW5sdtsPPAkcAT4N/Me1xk3Vqm2MmWumhXwfuJJRKf8QcF1VPbapgawgyVPAQlW1OvE9yb8EXmZ0RvYtzbb/AbxYVbc1X2JnV9XHOhDXrcDLVfXbmxnLkri2A9ur6ttJzgQOAu8G3k+Ln9kqcf072v/MApxRVS8nOR34OvAR4N/Q8t/ZqaiNyniSaSGnvKp6EHhxyeZrgDub13cy+o96U60QV+uq6vjipPqq+iGjimUHLX9mq8TVumba1cvN6unNUnTg7+xU1EYyXmnKR1cU8CdJDibZ3XYwS5y72Hdqfp7TcjzjbkrySNPGaPWftUkuBN4KfIsOfWZL4oIOfGZJtiR5GDgB3F9VnfrMTiVtJOOpp3xssndU1S8wuoLmw80/y7W6TwFvBi5lNI/yE20FkuQNwJeAj1bVD9qKY6ll4urEZ1ZVr1bVpYzO9l+W5C1txKF2kvHUUz42U1Uda36eAL7MqK3SFc8tXsXT/DzRcjwAVNVzzX/UP2Z0sqKVz6zpe34J+HxV3dNsbv0zWy6urnxmi6rqr4GvAbvowGd2KmojGU8yLaQVSc5oTrKQ5AzgV4FHV/+tTbUPuKF5fQNwb4ux/IMll3m+hxY+s+Zk1GeBw1X1ybG3Wv3MVoqrI5/ZtiRnNa9fD/wK8D06+nc2dJs+mwKgmcbzu4zueLS3qv7bpgexjCQ/y6gahtHViV9oK7YkXwSuYHRLw+eAW4A/Au4GLgCeBn6tqjb1ZNoKcV3B6J/bBTwFfGitOZVziOuXgD8Fvgv8uNn8m4z6s619ZqvEdR3tf2Y/z+gE3RZGhdndVfVfkvwULf+dnYpaScaSpJ/kFXiS1AEmY0nqAJOxJHWAyViSOsBkLEkdYDKWpA4wGUtSB/w/z27zX84PF7EAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "cmap = plt.get_cmap('YlOrRd')\n",
    "cmap.set_bad(color='white')\n",
    "fig, ax = plt.subplots(figsize=(6,6))\n",
    "img = ax.imshow(np.mean(pvalue<=0.05, axis=2), cmap=cmap, vmin=0, vmax=1)\n",
    "plt.colorbar(img, shrink=0.85)"
   ]
  }
 ],
 "metadata": {
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}