{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "47696870-1fda-48e0-86e7-ce8f7db2d46f",
   "metadata": {},
   "source": [
    "# Directed functional network inference with TE on Cellular Automata\n",
    "\n",
    "In this activity we will perform pairwise functional network inference using MI and TE on our previous Cellular Automata data set, where we know the structure of the underlying network and have a good understanding of the dynamics.\n",
    "\n",
    "## A. Functional network inference with MI\n",
    "\n",
    "First we will infer the functional network to represent the pairwise statistical relationships between the cells as the nodes in our network.\n",
    "\n",
    "1. Make sure you have a sample data set for CA rule 54 from our tutorial activities in the previous modules -- or to download again see the link on the tutorial website.\n",
    "2. Now open the MI AutoAnalyser. Select a Discrete estimator, select your CA data file, tick the checkbox for `All pairs?` and -- most importantly -- untick the checkbox for `Compute result?` (we don't want to run this computation, it will take a lot of time!), and click `Generate Code`.\n",
    "3. Paste the code into the code cells below.\n",
    "4. Examine how the generated code is computing the MI for each cell pair.\n",
    "5. Make a few changes to the code here for our purposes:\n",
    "    * Once we get to directed network inference, it will make more sense for our outer loop to be over targets, so that we consider the set of sources afterwards for that target. So, swap the order of the loops over `s` and `d` (so that the outer loop is over `d`)\n",
    "    * We also don't want to compute the functional network over the whole 10000 cells - this will take too long and is too large to visualise. Let's just work with the first 100. So, before the loops make a new variable `networkSize=100;`\n",
    "    * Then have the loops over `s` and `d` only run up to `networkSize` instead of all the way up to `10000`.\n",
    "    * Create a 2D array to store the MI values before the for loops, `results = numpy.zeros((networkSize,networkSize));`\n",
    "    * Remove the `print` statement, and instead assign the result to the results array, as `results[s,d] = result;`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e39673b-6d6e-42a5-8520-e54ed6ce4279",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Paste import and JVM startup lines here:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8418503d-4979-43a8-b6c4-da47d216437d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Paste analysis code here:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a913bede-d3ed-4fe7-8190-4ef57c112764",
   "metadata": {},
   "source": [
    "6. Run the above code cells, and then run the next code cell to plot the adjacency matrix of the functional network, as well as say the first 100 time steps of the CA dynamics so we can interpret what's happening:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50626790-8bc7-49cd-a16c-7893f8c90394",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the functional connectivity\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure(figsize=(14,6))\n",
    "plt.subplot(1,2,1) # left subplot\n",
    "plt.imshow(results)\n",
    "plt.xlabel('target');\n",
    "plt.ylabel('source');\n",
    "plt.title('MI between all cell pairs');\n",
    "cbar = plt.colorbar()\n",
    "cbar.set_label('MI (bits)');\n",
    "# Plot the first 100 time steps for the cells we're analysing:\n",
    "plt.subplot(1,2,2) # right subplot\n",
    "plt.imshow(data[:99,:networkSize]) # Plotting the first 100 time steps for first networkSize cells\n",
    "plt.xlabel('target');\n",
    "plt.ylabel('time');\n",
    "plt.title('Raw CA values');\n",
    "cbar = plt.colorbar()\n",
    "cbar.set_label('CA values');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab1b4970-442e-4f11-8524-d3128fe5863d",
   "metadata": {},
   "source": [
    "7. Look at the results for the adjacency matrix for the MI functional network in figure 1 (left):\n",
    "    * What do you observe in terms of the functional network structure? Which sources have the highest MI to each target?\n",
    "    * How does it compare to what you know the underlying causal structure to be? Is that what you expected to see?\n",
    "    * Can you interpret the features of the network via the dynamics that you can see in figure 2 (right)? What generates the modular block structure?\n",
    "\n",
    "**A note** for when you are constructing functional connectivity for other data sets which don't have an obvious ordinal layout like the CA data does: you might consider running a _clustering_ method over the FC adjacency matrix in order to group variables with similar relationships to others. This will give you visual insights into the modularity in the FC. To do this, you would use code like the following:\n",
    "```python\n",
    "fc = results\n",
    "# Import relevant libraries\n",
    "from scipy.cluster.hierarchy import linkage, optimal_leaf_ordering, leaves_list\n",
    "from scipy.spatial.distance import pdist\n",
    "# Perform hierarchical clustering using average linkage\n",
    "distances = numpy.max(fc) - fc\n",
    "Z = linkage(distances, method='average')\n",
    "# Compute the optimal leaf ordering\n",
    "optimal_Z = optimal_leaf_ordering(Z, distances)\n",
    "# Get the order of the leaves\n",
    "optimal_order = leaves_list(optimal_Z)\n",
    "# ...\n",
    "# And utilise this clustering to reorder the variables in your imshow plot call with:\n",
    "plt.imshow(fc[optimal_order,:][:,optimal_order],cmap=cm.coolwarm, norm=mplcolors.CenteredNorm())\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "389f8eb6-8b07-4ad0-9038-05c53c49c272",
   "metadata": {},
   "source": [
    "## B. Directed functional network inference with TE\n",
    "\n",
    "Now, we move on to infer a directed functional network with TE, to show the directed pairwise statistical relationships in the network.\n",
    "\n",
    "1. Start by copying the above two code cells for the functional network with MI (and plots), and pasting it into the code cell below that we will edit for TE.\n",
    "2. Open the TE AutoAnalyser - we're just going to generate some dummy code to grab a few lines from. Select a Discrete estimator, select your CA data file, tick the checkbox for `All pairs?`, set the `k_HISTORY` at 4 (we won't have enough data here for it to be larger, since we're assuming _heterogeneous_ variables here), and -- most importantly -- untick the checkbox for `Compute result?` (we don't want to run this computation, it will take a lot of time!), and click `Generate Code`.\n",
    "3. Copy the lines constructing the TE calculator, and paste it into your code cell below, replacing where the MI calculator was previously constructed.\n",
    "4. Down the bottom of the code where the adjacency matrix is plotted, in the `title()` and `set_label()` functions replace \"MI\" with \"TE\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3f214ba-23ab-4c25-b822-b666e31e46f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the MI analysis and plotting code cells into here; we will edit these to use TE\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a636f3fa-e549-4dea-97ef-3271ba2bf503",
   "metadata": {},
   "source": [
    "5. Run the script, and observe the results for the adjacency matrix for the TE directed functional network in the figure:\n",
    "    * Compare the results to those for the MI functional network. Has the location of the strongest sources for each target changed? Is this what you expected? Is this beginning to inform us of a more sensible model of the dynamics for each target, given who we know the sources for each target are?\n",
    "    * Do you still observe a modular block structure? How might we explain this with reference to the dynamics of the CA and the TE measure?\n",
    "    * Could you select a reasonable threshold to select parent cells for each target cell here? Try replotting the TE results but thresholding for parent selection via, e.g. `plt.imshow(results > 0.1);` for a threshold at 0.1 bits, to see if you can identify a threshold that works well everywhere.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8130ca4-055b-4aa8-abae-802a42270691",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try replotting the TE results with a threshold for parent selection\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9154c103-527c-4dda-8981-82fc10b9d695",
   "metadata": {},
   "source": [
    "7. Let's look into making a more principled threshold selection, following methods described in the lectures. We will add some code to the cell above to compute the p-values for the TE estimates from each source, analytically:\n",
    "    * Go back to the AutoAnalyser, ticking the checkboxes for `Add stat. signif.?` and `analytically?` to generate some template code of how to do this.\n",
    "    * Copy and paste the calculation of the `measDist` object after the TE calculation, and then store its member value `measDist.pValue` into a 2D array `pValues` as we have already done for `results`. Recall - what is the meaning of the p-value result for a given source?\n",
    "    * Where the if statement checks `if (s == d):` set `pValues[s,d] = 1;` otherwise seeing a 0 value there may confuse us.\n",
    "8. Re-run the above TE calculations, and also run the following code cell to plot the p-values.<br/>\n",
    "Take a look at the new figure below plotting the p-values. We will need to zoom in to see which sources have significant p-values here. We'll zoom in to only see sources which pass a Bonferroni corrected threshold `0.05 / (networkSize * (networkSize-1))` (to correct for the multiple comparisons here): uncomment the last line in the next code cell which now zooms in for us, and re-run this code cell.<br/>\n",
    "Any p-value that is below the extreme maximal colour would pass a statistical test here and be inferred as a parent in the directed functional network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27d4d33b-0593-4463-b5e1-3b807719b36e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the p-values for each pairwise calculation:\n",
    "plt.figure()\n",
    "plt.imshow(pValues)\n",
    "plt.xlabel('target');\n",
    "plt.ylabel('source');\n",
    "plt.title('p-value of TE between all cell pairs');\n",
    "cbar = plt.colorbar()\n",
    "cbar.set_label('p-value');\n",
    "# plt.clim(0,0.05/(networkSize*(networkSize-1))); # Can use this to highlight significant p-values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ad83b7f-e951-466d-b25d-11b1b60e22fa",
   "metadata": {},
   "source": [
    "8. Compared to the MI functional network and thresholding the TE value, these directed functional networks capture the direct parents on either side of the target, and we've addressed the issue of a more principled threshold selection. Still, are there many sources which appear to be parents here that you would not expect, given that we know only one source on either side of the target cell are its direct parents? Think about how we could extend the standard algorithm to exclude those other sources from being selected? (This is the subject of the next parts of the lecture)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aed6f74-2b85-484a-b4b9-eef4dc682b5d",
   "metadata": {},
   "source": [
    "<hr/>\n",
    "\n",
    "## C. Effective network inference with multivariate TE\n",
    "\n",
    "I've put together demonstration code to show a rudimentary implementation of the effective network inference on the CA data in the code cell below.\n",
    "\n",
    "1. Inspect the code, comparing it to the pairwise TE directed functional network code that we produced in the previous activity. Try to understand the main features:\n",
    "    * Inside the loop over targets, where we infer the set of parents for each, we have an additional while loop over rounds of adding parents to the set.\n",
    "    * Within each round, we have a loop over the candidate sources.\n",
    "    * Each candidate source is evaluated using a conditional TE, conditioning on previously selected parents.\n",
    "    * After all candidates are evaluated, we identify the candidate with the strongest conditional TE, and check the statistical significance of this value.\n",
    "    * If the source is statistically significant, we continue to another round to try to select another parent. Else we stop for this target.\n",
    "2. Try to run the code, and compare the results to your expectations for a good multivariate model explaining the dynamics, as well as to the functional networks from MI and pairwise TE above.\n",
    "    * Think through why the inference method might still not completely match the underlying structure, including: statistical fluctuations from the short data set effecting who the strongest sources are, and statistical power, hard coded history length used which may leave more information in spurious sources, no use of pruning step here, conservative p-value threshold and use of analytic surrogates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b64330c-585c-48de-b89d-3b2c102f3952",
   "metadata": {},
   "outputs": [],
   "source": [
    "network = numpy.zeros((networkSize,networkSize), dtype=int);\n",
    "# Compute for all targets:\n",
    "for d in range(networkSize):\n",
    "\n",
    "    print('Beginning greedy selection of parents for %d' % d);\n",
    "\n",
    "    destination = JArray(JInt, 1)(data[:, d].tolist())\n",
    "    conditionalSet = [];\n",
    "\n",
    "    while (True):\n",
    "        # Precondition: we have already selected the parents in\n",
    "        #  conditionalSet, now we check if we can add to this:\n",
    "        \n",
    "        results = numpy.zeros((networkSize));\n",
    "        pValues = numpy.zeros((networkSize));\n",
    "\n",
    "        # 1. Construct the calculator:\n",
    "        if (len(conditionalSet) > 0):\n",
    "            calcClass = JPackage(\"infodynamics.measures.discrete\").ConditionalTransferEntropyCalculatorDiscrete\n",
    "            calc = calcClass(2, 4, len(conditionalSet))\n",
    "        else:\n",
    "            calcClass = JPackage(\"infodynamics.measures.discrete\").TransferEntropyCalculatorDiscrete\n",
    "            calc = calcClass(2, 4, 1, 1, 1, 1)\n",
    "        # 2. No other properties to set for discrete calculators.\n",
    "        for s in range(networkSize):\n",
    "            # For each source-dest pair:\n",
    "            \n",
    "            if ( (s == d) or (s in conditionalSet) ):\n",
    "                # If s is the target, or already in the conditioning set, skip evaluating TE from it\n",
    "                pValues[s] = 1;\n",
    "                continue\n",
    "            source = JArray(JInt, 1)(data[:, s].tolist())\n",
    "            conditional = JArray(JInt, 2)(data[:, conditionalSet].tolist())\n",
    "\n",
    "            # 3. Initialise the calculator for (re-)use:\n",
    "            calc.initialise()\n",
    "            # 4. Supply the sample data:\n",
    "            if (len(conditionalSet) > 0):\n",
    "                calc.addObservations(source, destination, conditional)\n",
    "            else:\n",
    "                calc.addObservations(source, destination)\n",
    "\n",
    "            # 5. Compute the estimate:\n",
    "            result = calc.computeAverageLocalOfObservations()\n",
    "            # 6. Compute the (statistical significance via) null distribution analytically:\n",
    "            measDist = calc.computeSignificance()\n",
    "\n",
    "            results[s] = result;\n",
    "            pValues[s] = measDist.pValue;\n",
    "\n",
    "        # Check which was the strongest source:\n",
    "        maxSourceIndex = numpy.argmax(results);\n",
    "        maxTE = results[maxSourceIndex];\n",
    "        if (pValues[maxSourceIndex] < 0.05/(networkSize*(networkSize-1))):\n",
    "            print('Selected source %d, with TE(%d->%d | %s)=%.5f, p-value=%.5f (conditioning on %d parents)' %\\\n",
    "                (maxSourceIndex, maxSourceIndex, d, (','.join(str(x) for x in conditionalSet)), maxTE, \\\n",
    "                pValues[maxSourceIndex], len(conditionalSet) ));\n",
    "            # Add this new source:\n",
    "            conditionalSet.append(maxSourceIndex);\n",
    "        else:\n",
    "            print('-- Max TE source %d was not significant (TE(%d->%d | %s)=%.5f, p-value=%.6f), quitting' %\\\n",
    "                (maxSourceIndex, maxSourceIndex, d, (','.join(str(x) for x in conditionalSet)), maxTE, pValues[maxSourceIndex]));\n",
    "            break\n",
    "    \n",
    "    # Postcondition: conditionalSet holds the parents for target d\n",
    "    network[conditionalSet, d] = 1;\n",
    "\n",
    "# Plot the effective connectivity\n",
    "plt.figure()\n",
    "plt.imshow(network)\n",
    "plt.xlabel('target');\n",
    "plt.ylabel('source');\n",
    "plt.title('Multivariate effective network via TE');\n",
    "cbar = plt.colorbar()\n",
    "cbar.set_label('Connections');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "440026b4-3190-47c1-8982-193648bdfa41",
   "metadata": {},
   "source": [
    "A sample results plot is shown on the tutorial website (yours will not match precisely due to stochastic generation of the dynamics but should have similar features).\n",
    "\n",
    "3. _Optional extensions -- advanced_:\n",
    "    1. Can you incorporate code to address any of those issues to improve the network inference here?\n",
    "    2. E.g. try changing the history length to only 1 and try to understand the changes this makes in terms of what is being modelled, and why this might be a better choice for situations where you want to more closely replicate a causal structure between the observed variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da1fcef8-58e3-453b-944a-f8c73e396e98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the code from the cell above and adapt it here\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
