{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ad1dfb3c-15c9-4049-b223-dbe4031e678d",
   "metadata": {},
   "source": [
    "# Transfer entropy for Heartbeat process\n",
    "\n",
    "_Copyright (C) 2023-, Joseph T. Lizier\n",
    "Distributed under GNU General Public License v3_\n",
    "\n",
    "In this activity we will compute the local Transfer entropy on the Heartbeat (Poisson switching) process.\n",
    "\n",
    "1. Start by generating some sample data for this process:\n",
    "    1. The function `generateHeartbeatMessages` which will generate the data is defined in the code cell below.\n",
    "    2. Add a line in the next cell after that: `data = generateHeartbeatMessages(0.05, 0.2, 100000);` to generate `N=100000` time steps of sample data for a heartbeat process with transition probabilities (as defined in lecture) `lambda1 = 0.05; lambda0 = 0.2`. And run both code cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "119fe2ba-71ad-47d0-b457-eff9cb0d26e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "\n",
    "\"\"\"function generateHeartbeatMessages(lambda1, lambda0, N)\n",
    "\n",
    "Generate sample time series data, of N time steps, for two processes.\n",
    "The first process s is a Poisson switching process, where the next value s_{n+1}\n",
    " is a simple function of its previous value s_n:\n",
    "\n",
    " p(s_{n+1} = 1 | s_n = 0) = lambda1\n",
    " p(s_{n+1} = 0 | s_n = 0) = 1-lambda1\n",
    " p(s_{n+1} = 1 | s_n = 1) = 1-lambda0\n",
    " p(s_{n+1} = 0 | s_n = 1) = lambda0\n",
    "\n",
    "The second process t simply copies the previous value of s.\n",
    "\n",
    "We have for the steady-state probabilities:\n",
    " p(s_n = 0) = lambda0 / (lambda0 + lambda1)\n",
    " p(s_n = 1) = lambda1 / (lambda0 + lambda1)\n",
    "Proof:\n",
    " p(1) = p(0).p(s_{n+1} = 1 | s_n = 0) + p(1).p(s_{n+1} = 1 | s_n = 1)\n",
    "      = p(0).lambda1 + p(1).(1-lambda0)\n",
    " p(1) = lambda1/lambda0 . p(0)\n",
    "and: p(1) + p(0) = 1\n",
    "Solve these two equations to get the above result.\n",
    "\n",
    "Inputs:\n",
    "- lambda1 - probability of switching to 1 when current value is a 0\n",
    "- lambda0 - probability of switching to 0 when current value is a 1\n",
    "- N - number of time steps to generate\n",
    "\n",
    "Output:\n",
    "- data - Nx2 matrix. First column is process s, second column is process t.\n",
    "\n",
    "Copyright (C) 2017, Joseph T. Lizier\n",
    "Distributed under GNU General Public License v3\n",
    "\"\"\"\n",
    "def generateHeartbeatMessages(lambda1, lambda0, N):\n",
    "\n",
    "    s = numpy.zeros((N));\n",
    "    t = numpy.zeros((N));\n",
    "    # One can show:\n",
    "    p1 = lambda1 / (lambda0 + lambda1);\n",
    "    p0 = lambda0 / (lambda0 + lambda1);\n",
    "    # Now assign s[0], t[0] according to these probabilities:\n",
    "    s[0] = (numpy.random.rand() < p1) * 1;\n",
    "    t[0] = (numpy.random.rand() < p1) * 1;\n",
    "    # And assign the remaining time series activity based on the process\n",
    "    #  dynamics.\n",
    "    # The way I've done it here is rather slow, but fast enough for\n",
    "    #  our purposes ... and v fast to write :)\n",
    "    for n in range(1,N):\n",
    "        # Assign next value conditioned on previous:\n",
    "        if (s[n-1] == 0):\n",
    "            # Prob that s[n] switches to 1 is lambda1\n",
    "            s[n] = (numpy.random.rand() < lambda1) * 1;\n",
    "        else:\n",
    "            # Prob that s[n] stays at 1 is (1-lambda0)\n",
    "            s[n] = (numpy.random.rand() < 1 - lambda0) * 1;\n",
    "    # Now do the delayed copy to the target\n",
    "    t[1:] = s[:-1];\n",
    "\n",
    "    data = numpy.zeros((N,2))\n",
    "    data[:,0] = s;\n",
    "    data[:,1] = t;\n",
    "    return data.astype(numpy.int32) # Makes sure the return values are not floats\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fef20755-7ea5-4214-8aa6-82d81d8fc1d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call the generateHeartbeatMessages() function here:\n",
    "data = generateHeartbeatMessages(0.05, 0.2, 100000);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "090269c0-85e3-4989-ad9a-b2608dfbc034",
   "metadata": {},
   "source": [
    "2. Plot some of the sample data to make sure it's working ok, via the next code cell.\n",
    "\n",
    "    Make sure the target lags the source by one time step, and the transition rates look somewhat ok by eye. A sample result is shown on the tutorial page (note -- yours won't precisely match, since these are generated stochastically):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2ba87024-b259-4896-8258-5efa8ad2e699",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34164 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.scatter(range(1,101), data[:100,0], marker='x', color='red', label='source')\n",
    "plt.scatter(range(1,101), data[:100,1], marker='o', color='blue', facecolors='none', label='target')\n",
    "plt.xlabel('n')\n",
    "plt.ylabel('state')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0b47532-18fe-4be4-8817-5d476b278a51",
   "metadata": {},
   "source": [
    "3. Open the AutoAnalyser and select Transfer Entropy. Select a Discrete estimator, select any discrete data file (e.g. `2coupledBinaryColsUseK2.txt`), leave the source and target at columns 0 and 1 respectively, uncheck `Compute result?`, and click `Generate Code`.\n",
    "4. Copy and paste the generated Python code in the code cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8299519b-055d-4eec-a5a6-91451d2cde57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Place the AutoAnalyser generated code template here, ready for editing.\n",
    "# Put the import lines and JVM startup code here:\n",
    "\n",
    "from jpype import *\n",
    "import numpy\n",
    "import sys\n",
    "# Our python data file readers are a bit of a hack, python users will do better on this:\n",
    "sys.path.append(\"/home/joseph/JIDT/infodynamics-dist-1.6.1/demos/python\")\n",
    "import readIntsFile\n",
    "\n",
    "if (not isJVMStarted()):\n",
    "    # Add JIDT jar library to the path\n",
    "    jarLocation = \"/home/joseph/JIDT/infodynamics-dist-1.6.1/infodynamics.jar\"\n",
    "    # Start the JVM (add the \"-Xmx\" option with say 1024M if you get crashes due to not enough memory space)\n",
    "    startJVM(getDefaultJVMPath(), \"-ea\", \"-Djava.class.path=\" + jarLocation, convertStrings=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9e4b7424-3e52-4d16-a913-97c7a3ac31a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TE_Discrete(col_0 -> col_1) = 0.3715 bits\n",
      "TE_Discrete(col_0 -> col_1) = 0.3715 bits (null: 0.0000 +/- 0.0000 std dev.; p(surrogate > measured)=0.00000 from 100 surrogates)\n"
     ]
    }
   ],
   "source": [
    "# Place the remainder of the auto-generated code here:\n",
    "\n",
    "# 0. Load/prepare the data:\n",
    "data = generateHeartbeatMessages(0.05, 0.2, 100000);\n",
    "source = JArray(JInt, 1)(data[:,0].tolist())\n",
    "destination = JArray(JInt, 1)(data[:,1].tolist())\n",
    "\n",
    "# 1. Construct the calculator:\n",
    "calcClass = JPackage(\"infodynamics.measures.discrete\").TransferEntropyCalculatorDiscrete\n",
    "calc = calcClass(2, 1, 1, 1, 1, 1)\n",
    "# 2. No other properties to set for discrete calculators.\n",
    "# 3. Initialise the calculator for (re-)use:\n",
    "calc.initialise()\n",
    "# 4. Supply the sample data:\n",
    "calc.addObservations(source, destination)\n",
    "# 5. Compute the estimate:\n",
    "result = calc.computeAverageLocalOfObservations()\n",
    "\n",
    "print(\"TE_Discrete(col_0 -> col_1) = %.4f bits\" %\\\n",
    "    (result))\n",
    "\n",
    "# Challenge: Add statistical significance check:\n",
    "# 6. Compute the (statistical significance via) null distribution empirically (e.g. with 100 permutations):\n",
    "measDist = calc.computeSignificance(100)\n",
    "print(\"TE_Discrete(col_0 -> col_1) = %.4f bits (null: %.4f +/- %.4f std dev.; p(surrogate > measured)=%.5f from %d surrogates)\" %\\\n",
    "    (result, measDist.getMeanOfDistribution(), measDist.getStdOfDistribution(), measDist.pValue, 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce831d43-b74c-4f93-a884-77af72a2fa2a",
   "metadata": {},
   "source": [
    "5. Replace the code at step 0 where the data file is loaded with the line we ran above to generate the heartbeat data sample: `data = generateHeartbeatMessages(0.05, 0.2, 100000);` (or you can load a saved sample data file as an alternative)\n",
    "6. Run your code above to check that it computes an average TE ok -- the result should be close to the theoretical value of 0.3735 bits for this large data sample.\n",
    "    * Is this value statistically significant? (_Challenge_: can you check this?)\n",
    "7. In the next code cell, add the following code to compute local TE at each time step: `locals = calc.computeLocalFromPreviousObservations(source, destination);`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0751089a-9af9-4b49-aa9b-d71cba84ec43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the local TE at each time step here:\n",
    "\n",
    "locals = calc.computeLocalFromPreviousObservations(source, destination);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2987c40e-4f68-4be5-8006-1f2b79e4c4b0",
   "metadata": {},
   "source": [
    "8. In the next code cell, we will plot our results. First, replot the source and target time series (since we resampled `data`) -- copy the plotting code from step 2 above in below, and\n",
    "9. Plot the local TE alongside the target variable in the next figure:\n",
    "```python\n",
    "plt.scatter(range(1,101), locals[:100], marker='x', color='red', label='TE (bits)')\n",
    "plt.xlabel('n')\n",
    "plt.ylabel('TE (bits)')\n",
    "plt.legend(loc='upper left')\n",
    "ax2 = plt.twinx()\n",
    "ax2.scatter(range(1,101), data[:100,1], marker='o', color='blue', facecolors='none', label='target')\n",
    "ax2.set_ylabel('state')\n",
    "ax2.legend(loc='upper right')\n",
    "plt.show()\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9ed1746b-5d42-4940-9870-95452834e24e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33812 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 36392 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Add the code to plot the results here:\n",
    "\n",
    "# First replotting the source and target time series:\n",
    "plt.scatter(range(1,101), data[:100,0], marker='x', color='red', label='source')\n",
    "plt.scatter(range(1,101), data[:100,1], marker='o', color='blue', facecolors='none', label='target')\n",
    "plt.xlabel('n')\n",
    "plt.ylabel('state')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# Then plotting the local TE alongside the target variable:\n",
    "plt.scatter(range(1,101), locals[:100], marker='x', color='red', label='TE (bits)')\n",
    "plt.xlabel('n')\n",
    "plt.ylabel('TE (bits)')\n",
    "plt.legend(loc='upper left')\n",
    "ax2 = plt.twinx()\n",
    "ax2.scatter(range(1,101), data[:100,1], marker='o', color='blue', facecolors='none', label='target')\n",
    "ax2.set_ylabel('state')\n",
    "ax2.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ff7d314-170c-4e36-ba38-59bfe2eca9e1",
   "metadata": {},
   "source": [
    "10. Examine the results (you may need to zoom in on the second figure above):\n",
    "    * Where did the source add the _most_ information in the context of the target past? Why?\n",
    "    * Where did the source add the _least_ information in the context of the target past? Why?\n",
    "    * We don't see any negative values of local TE here - can you explain why?\n",
    "11. _Challenge task_: Compute the local lagged mutual information (across one time step from the source to the target) by adding code below (e.g. by adapting the above code). Compare the results, and discuss why they are different. (Note that the local MI values generated by JIDT here will need a 0 padded to the front in order to align them with the local TE values.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7976bec2-4d07-4383-9d98-767f39f2585b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40644 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 1. Creat MI Calc with time-lag 1 (Important to set this!)\n",
    "miCalcClass = JPackage(\"infodynamics.measures.discrete\").MutualInformationCalculatorDiscrete\n",
    "miCalc = miCalcClass(2, 2, 1)\n",
    "# 2. No other properties to set for discrete calculators.\n",
    "# 3. Initialise the calculator for (re-)use:\n",
    "miCalc.initialise()\n",
    "# 4. Supply the sample data:\n",
    "miCalc.addObservations(source, destination)\n",
    "# 5. Compute the estimate:\n",
    "result = miCalc.computeAverageLocalOfObservations()\n",
    "# And compute the local values:\n",
    "localsMI = miCalc.computeLocalFromPreviousObservations(source, destination);\n",
    "\n",
    "# And our localsMI don't have a result for timestep 0 (where there is no history) yet TE does,\n",
    "#  so align these properly:\n",
    "localsMI = numpy.append([0], localsMI);\n",
    "\n",
    "# Now plot them all together:\n",
    "plt.scatter(range(1,101), locals[:100], marker='x', color='red', label='TE (bits)')\n",
    "plt.scatter(range(1,101), localsMI[:100], marker='+', color='green', label='MI (bits)')\n",
    "plt.xlabel('n')\n",
    "plt.ylabel('Information (bits)')\n",
    "plt.legend(loc='upper left')\n",
    "ax2 = plt.twinx()\n",
    "ax2.scatter(range(1,101), data[:100,1], marker='o', color='blue', facecolors='none', label='target')\n",
    "ax2.set_ylabel('state')\n",
    "ax2.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d53392d-f408-4fd1-8375-1e91a944a47b",
   "metadata": {},
   "source": [
    "12. _Optional (and difficult!) Challenge task_: Calculate analytically the local transfer entropy for each configuration, as well as the average transfer entropy, that we expect to get here as a function of the transition probabilities $\\lambda_0$ and $\\lambda_1$. You can start with the transition probabilities $p(x_{n+1} | y_n, x_n)$ and $p(x_{n+1} | x_n)$ - and you will also need to work out the probabilities of each value $x_n = 0, 1$ of the target (the derivation for these is shown in the comments of `generateHeartbeatMessages()` if you get stuck), then plug everything through the TE equation. A full solution is available on the tutorial page.\n",
    "13. _Optional Challenge task_: Can you compute and display how the average TE (not the local values) changes as a function of the transition probabilities $\\lambda_0$ and $\\lambda_1$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdffef2c-644c-45e9-8952-a93dbe8590f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "3c0fb51c-759d-44a7-afae-2a9f69414b08",
   "metadata": {},
   "source": [
    "There is a discussion of the results of this activity in a short video on the tutorial site."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e24c173-4868-4018-bb35-944e499218f4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
