PyTensor_for_RL.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNGWZSl05FJCLhSo9Vj3xSy",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/TheLemonPig/RL-SSM/blob/main/PyTensor_for_RL.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##PyTensor for RL"
      ],
      "metadata": {
        "id": "Kqtk2k0CamtN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pytensor.tensor as pt\n",
        "from pytensor import function, scan\n",
        "import numpy as np"
      ],
      "metadata": {
        "id": "_YyaiHPCPMR3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Reinforcment Learning is inherently sequential, meaning that we will need to make use of `for` loops.\n",
        "\n",
        "In PyTensor, this can be acheived in a very narrow sense using the `scan` function."
      ],
      "metadata": {
        "id": "NbaBoPG5KYdC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "###2a. Using scan"
      ],
      "metadata": {
        "id": "pRO7jUZYKxjf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For more, see: https://pytensor.readthedocs.io/en/latest/library/scan.html"
      ],
      "metadata": {
        "id": "8kwH_PWAK6g0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For simplicity, we will start with a numerical example, before moving on to a matrix example."
      ],
      "metadata": {
        "id": "n2qVNkrPLH-q"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "####2ai. Using scan with numerical inputs/outputs"
      ],
      "metadata": {
        "id": "FQGG6TOALOG9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Suppose we want to calculate the first $n$ terms of the series using the recursion formula: $s_{n+1} = \\frac{s_{n}}{n+1} + m$ for $s_0=1$ and any $m$"
      ],
      "metadata": {
        "id": "J_E7zawpL1b-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "A python function for this could look like:"
      ],
      "metadata": {
        "id": "GB1BK-QPNtwN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def formula(sn,n,m):\n",
        "    return sn/(n+1) + m\n",
        "\n",
        "def get_term(n,m):\n",
        "    sn = 1\n",
        "    series = np.zeros(n)\n",
        "    series[0] = sn\n",
        "    for i in range(n):\n",
        "        sn = formula(sn,i,m)\n",
        "        series[i] = sn\n",
        "    return series"
      ],
      "metadata": {
        "id": "QftkqLvSK577"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This is how we would write it in PyTensor:"
      ],
      "metadata": {
        "id": "NVlDnsnBOqlJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# We have to slightly adapt our function because the order of arguments in the function is incredibly important for getting PyTensor to work correctly\n",
        "# They should be organized like this: sequences, outputs_info, non_sequences\n",
        "def formula(n,sn,m):\n",
        "    return sn/(n+1) + m\n",
        "\n",
        "def get_term_pytensor(n,m):\n",
        "    sn = pt.constant(1.0, dtype='float64')  # you will sometimes need to define the type like this. You will need to be very careful about type. PyTensor generally defaults to 32-bit types\n",
        "    ns = pt.arange(n, dtype='float64')\n",
        "    term, _ = scan(fn=formula, sequences=ns, non_sequences=m, outputs_info=sn)  # The ordering of\n",
        "    return term"
      ],
      "metadata": {
        "id": "1tynn8FMc2TP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "n_compile = pt.scalar(\"n\", dtype='int64')  # note n must be an integer (you can't have the 3.5th term of a series)\n",
        "m_compile = pt.scalar(\"m\", dtype='float64')\n",
        "\n",
        "output = get_term_pytensor(n_compile, m_compile)\n",
        "my_func = function(inputs=[n_compile,m_compile],outputs=output)  # note the LACK of parentheses for the outputs. You will need these if you have multiple outputs but otherwise you do not want your function unknowingly wrapping your output in a list"
      ],
      "metadata": {
        "id": "Wx9ZG2q0QJXX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "It is always important to compare the outputs of your PyTensor function with the original Python function."
      ],
      "metadata": {
        "id": "iniPynF7RXTN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "get_term(10,1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5QIR8kTBRVUv",
        "outputId": "9f6f3903-f909-4e3d-b724-6ea74faaa239"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([1.        , 1.5       , 1.8       , 2.07142857, 2.30232558,\n",
              "       2.51408451, 2.70741483, 2.88810811, 3.05755596, 3.21808401])"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "my_func(10,1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QOP1lVwCRgL0",
        "outputId": "1f2035e6-15ad-4e27-925d-6b4c8efd990f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([2.        , 2.        , 1.66666667, 1.41666667, 1.28333333,\n",
              "       1.21388889, 1.1734127 , 1.14667659, 1.12740851, 1.11274085])"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###2b.Using scan for RL"
      ],
      "metadata": {
        "id": "Phegn_r6WlsA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now because you cannot index PyTensor arrays using arrays, we must one-hot the data so that indexing is not necessary. This saves us having to scan over trials to extract the relevant Q-values for the likelihood. (If you can find another way, then a lot of the code dealing with one-hotting from here on isn't necessary).\n",
        "\n",
        "**Note:** It is important to know that the first dimension of the tensors that you place in `sequences` should be the dimension that is missing from the tensor you pass to `outputs_info`. `scan` will iterate over subtensors of the sequence tensor by indexing from the first dimension. This dimension will be added to the output of your scan function. In this way `outputs_info` allows you to place an initial value to iterate over, where the intial value is used to calculate the second value and this repeats over each iteration of the scan function."
      ],
      "metadata": {
        "id": "2nrujwrgW2TW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "####2bi.Using scan for single participant RL"
      ],
      "metadata": {
        "id": "tAIoMmF9exEd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For simplicty we will go over what it looks like for an individual participant, and then extend the code to multiple participants"
      ],
      "metadata": {
        "id": "A8WgaHiCe1Dd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In all of these functions you will notice that we include the dimension for choice. This is the dimension which is being one-hotted: taking the choice 0 is represented as [1, 0], and taking the choice 1 is represented as [0, 1]. We need as many numbers as there are choices if we want to place a 1 at the index of the value of the choice. This is why we need to create an extra dimension which is the size of the number of choices."
      ],
      "metadata": {
        "id": "pEZyM05xkX7C"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Simulating RL in Python**"
      ],
      "metadata": {
        "id": "SCYhJXsKmfVg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from functools import partial\n",
        "from typing import Callable, List"
      ],
      "metadata": {
        "id": "u8OKLXPsnU19"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def softmax(qs, beta):\n",
        "  return np.exp(qs*beta) / np.exp(qs*beta).sum()\n",
        "\n",
        "class Distribution:\n",
        "\n",
        "  def __init__(self, func: Callable, kwargs):\n",
        "    self.func: Callable = partial(func, **kwargs)\n",
        "\n",
        "  def __call__(self):\n",
        "    return self.func.__call__()\n",
        "\n",
        "class SimpleRL:\n",
        "\n",
        "  def __init__(self, n_trials: int, distributions: List[Distribution]):\n",
        "    self.n_choices: int = len(distributions)\n",
        "    self.n_trials: int = n_trials\n",
        "    self.distributions: List[Distribution] = distributions\n",
        "    self.qs: np.array = np.ones((self.n_choices,)) * 0.5\n",
        "    self.q_trace: np.array = np.ones((self.n_trials,self.n_choices))\n",
        "    self.rewards: np.array = np.zeros((self.n_trials))\n",
        "    self.choices: np.array = np.zeros((self.n_trials),dtype=np.int32)\n",
        "\n",
        "  def simulate(self, alpha, beta):\n",
        "    for i in range(self.n_trials):\n",
        "      # Q-values are recorded to trace\n",
        "      self.q_trace[i] = self.qs\n",
        "      # softmax decision function\n",
        "      ps = softmax(self.qs,beta)\n",
        "      # choice made based on weighted probabilities of Q-values\n",
        "      choice = np.random.choice(a=self.n_choices,size=1,p=ps)[0]\n",
        "      # choice is recorded to trace\n",
        "      self.choices[i] = choice\n",
        "      # reward calculated\n",
        "      dist = self.distributions[choice]  # supply a list of distributions to choose from\n",
        "      reward = dist()  # sample from distribution by calling it\n",
        "      # Q-values updated\n",
        "      self.rewards[i] = reward\n",
        "      self.qs[choice] = self.qs[choice] + alpha * (reward - self.qs[choice])\n",
        "      self.q_trace[i] = self.qs\n",
        "    # Q-values trace returned\n",
        "    # main data to be returned (basis for fits), is choices and rewards per trial\n",
        "    return self.rewards, self.choices, self.q_trace"
      ],
      "metadata": {
        "id": "fjQja5XOm_Rh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "seed = 0\n",
        "np.random.seed(seed)\n",
        "mean_rewards = [-1.0,1.0]\n",
        "dists = [Distribution(np.random.normal,{\"loc\":mn, \"scale\":1.0}) for mn in mean_rewards]\n",
        "n_trials = 100\n",
        "\n",
        "alpha, beta = 0.1, 1.0\n",
        "\n",
        "simple_rl = SimpleRL(n_trials, dists)\n",
        "Rs, Cs, Qs = simple_rl.simulate(alpha, beta)"
      ],
      "metadata": {
        "id": "1aCw8zIfnMsB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Recovering Qs in PyTensor**"
      ],
      "metadata": {
        "id": "AwVGTYU_m3g_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def rl_step(C, R, Q, A):\n",
        "    \"\"\"\n",
        "    rl_step: function for a single RL step\n",
        "\n",
        "    C: Choice -- vector (choices)\n",
        "    R: Reward -- vector (choices)\n",
        "    Q: Q-Values at the previous time step -- vector (choices)\n",
        "    A: Alpha parameter -- scalar ()\n",
        "    \"\"\"\n",
        "    return Q + A * (R - Q) * C"
      ],
      "metadata": {
        "id": "9XXBzm81Xx7a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def rl_scan(Cs,Rs,A):\n",
        "    \"\"\"\n",
        "    rl_scan: scan function over RL steps\n",
        "\n",
        "    Cs: Choices -- matrix (trials, choices)\n",
        "    Rs: Rewards -- matrix (trials, choices)\n",
        "    A: Alpha parameter -- vector (choices)\n",
        "\n",
        "    Qs: Q-values -- matrix (trials, choices)\n",
        "    \"\"\"\n",
        "    Qs = pt.ones(Cs.shape[1]) * 0.5 # note Qs shape changes over scan:  vector (choices) --> matrix (trials, choices)\n",
        "    Qs, _ = scan(fn=rl_step, sequences=[Cs,Rs], non_sequences = [A], outputs_info=Qs)\n",
        "    return Qs"
      ],
      "metadata": {
        "id": "9SkAHsV-X36S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Defining the inputs of the to-be-compiled function\n",
        "Cs = pt.imatrix(\"Cs\")\n",
        "Rs = pt.dmatrix(\"Rs\")\n",
        "A = pt.dvector(\"A\")\n",
        "\n",
        "# Compiling function\n",
        "output = rl_scan(Cs,Rs,A)\n",
        "rl_func = function(inputs=[Cs,Rs,A], outputs=output)"
      ],
      "metadata": {
        "id": "vapxdRdYYBsl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "seed = 0\n",
        "np.random.seed(seed)\n",
        "mean_rewards = [-1.0,1.0]\n",
        "dists = [Distribution(np.random.normal,{\"loc\":mn, \"scale\":1.0}) for mn in mean_rewards]\n",
        "n_trials = 100\n",
        "\n",
        "alpha, beta = 0.1, 1.0\n",
        "\n",
        "simple_rl = SimpleRL(n_trials, dists)\n",
        "Python_Rs, Python_Cs, Python_Qs = simple_rl.simulate(alpha, beta)\n",
        "n_choices = len(mean_rewards)\n",
        "\n",
        "# Testing compiled function\n",
        "Rs = np.repeat(Python_Rs.reshape(-1,1), repeats = n_choices, axis=1)\n",
        "Cs = np.zeros((Python_Cs.shape[0], n_choices))\n",
        "for n in range(n_choices):\n",
        "  Cs[:, n] = (Python_Cs == n)\n",
        "Cs = np.array(Cs, dtype=np.int32)\n",
        "A = np.ones(n_choices) * alpha\n",
        "PyTensor_Qs = rl_func(Cs,Rs,A)\n",
        "PyTensor_Qs[:10]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Y9s2S6JAZQPv",
        "outputId": "ffc87914-9dcd-4552-96b7-9bad69c8583c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0.5       , 0.62415917],\n",
              "       [0.5       , 0.81703463],\n",
              "       [0.5       , 0.75105876],\n",
              "       [0.5       , 0.97294533],\n",
              "       [0.5       , 1.10226265],\n",
              "       [0.5       , 1.04144873],\n",
              "       [0.5       , 1.08573507],\n",
              "       [0.5       , 1.13507561],\n",
              "       [0.5       , 1.10340979],\n",
              "       [0.5       , 1.23408928]])"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "Python_Qs[:10]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HCMzhDQkpLS9",
        "outputId": "754d8895-db63-402b-c20a-6a499b073596"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0.5       , 0.62415917],\n",
              "       [0.5       , 0.81703463],\n",
              "       [0.5       , 0.75105876],\n",
              "       [0.5       , 0.97294533],\n",
              "       [0.5       , 1.10226265],\n",
              "       [0.5       , 1.04144873],\n",
              "       [0.5       , 1.08573507],\n",
              "       [0.5       , 1.13507561],\n",
              "       [0.5       , 1.10340979],\n",
              "       [0.5       , 1.23408928]])"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "####2bii. Using scan over multiple participants"
      ],
      "metadata": {
        "id": "ojxZVxmggVXm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Simulating multi-participant RL in Python**"
      ],
      "metadata": {
        "id": "aaHf1tWv8ZmJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiRL:\n",
        "\n",
        "  def __init__(self, n_trials: int, n_participants: int, distributions: List[Distribution]):\n",
        "    self.participants: List[SimpleRL] = [\n",
        "        SimpleRL(n_trials, distributions) for _ in range(n_participants)\n",
        "    ]\n",
        "    self.n_trials = n_trials\n",
        "    self.n_participants = n_participants\n",
        "    self.n_choices = len(distributions)\n",
        "    self.alphas = np.zeros(n_participants)\n",
        "    self.betas = np.zeros(n_participants)\n",
        "\n",
        "  def simulate(self, alpha_a, alpha_b, beta_a, beta_b):\n",
        "    group_data = []\n",
        "    group_rewards = np.zeros((self.n_trials,self.n_participants))\n",
        "    group_choices = np.zeros((self.n_trials,self.n_participants))\n",
        "    group_Qs = np.zeros((self.n_trials,self.n_participants,self.n_choices))\n",
        "    for idx, participant_model in enumerate(self.participants):\n",
        "        # sample participant parameters\n",
        "        alpha = np.random.beta(alpha_a, alpha_b)\n",
        "        beta = np.random.beta(beta_a, beta_b)\n",
        "        self.alphas[idx] = alpha\n",
        "        self.betas[idx] = beta\n",
        "        # run RL for participant\n",
        "        participant_rewards, participant_choices, participant_Qs = participant_model.simulate(alpha, beta)\n",
        "        group_rewards[:, idx] = participant_rewards\n",
        "        group_choices[:, idx] = participant_choices\n",
        "        group_Qs[:, idx] = participant_Qs\n",
        "    return group_rewards, group_choices, group_Qs\n",
        "\n",
        "  def get_params(self):\n",
        "    return self.alphas, self.betas\n"
      ],
      "metadata": {
        "id": "lIGS8alM8ZAn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Recovering multi-participant Qs in PyTensor**"
      ],
      "metadata": {
        "id": "qY4p7WpB8hcN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Notice that this function is written in the same way as for a single participant\n",
        "\n",
        "def rl_multi_step(C, R, A, Q):  # NOTE: we made A a sequence parameter. This will give us more flexibility in the future\n",
        "    \"\"\"\n",
        "    rl_step: function for a single RL step\n",
        "\n",
        "    C: Choice -- matrix (participants, choices)\n",
        "    R: Reward -- matrix (participants, choices)\n",
        "    Q: Q-Values at the previous time step -- matrix (participants, choices)\n",
        "    A: Alpha parameter -- matrix (participants, choices)\n",
        "    \"\"\"\n",
        "    return Q + A * (R - Q) * C"
      ],
      "metadata": {
        "id": "A0_o50K_gd2X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Notice that this function is written in ALMOST the same way as for a single participant\n",
        "# Tensor3 is how PyTensor defines 3-Dimensional arrays\n",
        "\n",
        "def rl_multi_scan(CMs,RMs,AM):\n",
        "    \"\"\"\n",
        "    rl_scan: scan function over RL steps\n",
        "\n",
        "    CMs: Choices -- tensor3 (trials, participants, choices)\n",
        "    RMs: Rewards -- tensor3 (trials, participants, choices)\n",
        "    AM: Alpha parameter -- matrix (participants, choices)\n",
        "\n",
        "    Qs: Q-values -- tensor3 (trials, participants, choices)\n",
        "    \"\"\"\n",
        "    QMs = pt.ones((CMs.shape[1], CMs.shape[2])) * 0.5 # note Qs shape changes over scan:  matrix (participants, choices) --> tensor3 (trials, participants, choices)\n",
        "    QTs, _ = scan(fn=rl_multi_step, sequences=[CMs,RMs,AM], non_sequences = [], outputs_info=QMs)  # NOTICE THAT WE DO NOT PLACE BRACKETS AROUND THE OUTPUT!!\n",
        "    return QTs"
      ],
      "metadata": {
        "id": "saB55Ka0iBlt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Defining the inputs of the to-be-compiled function\n",
        "CMs = pt.itensor3(\"CMs\")\n",
        "RMs = pt.dtensor3(\"RMs\")\n",
        "AM = pt.dtensor3(\"AM\")\n",
        "\n",
        "# Compiling function\n",
        "output = rl_multi_scan(CMs,RMs,AM)\n",
        "rl_func = function(inputs=[CMs,RMs,AM], outputs=output)"
      ],
      "metadata": {
        "id": "YTB8fxXQjdhE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "seed = 0\n",
        "np.random.seed(seed)\n",
        "mean_rewards = [-1.0,1.0]\n",
        "dists = [Distribution(np.random.normal,{\"loc\":mn, \"scale\":1.0}) for mn in mean_rewards]\n",
        "n_trials = 100\n",
        "n_participants = 3\n",
        "n_choices = len(mean_rewards)\n",
        "\n",
        "alpha_a, alpha_b, beta_a, beta_b = 2, 2, 2, 2\n",
        "# alpha, beta = 0.1, 1.0 -- For MultiRL the parameters vary with a distribution\n",
        "\n",
        "multi_rl = MultiRL(n_trials, n_participants, dists)\n",
        "Python_Rs, Python_Cs, Python_Qs = multi_rl.simulate(alpha_a, alpha_b, beta_a, beta_b)\n",
        "\n",
        "alphas, betas = multi_rl.get_params()\n",
        "\n",
        "# Testing compiled function\n",
        "Rs = np.repeat(Python_Rs.reshape(n_trials,n_participants,1), n_choices, axis=2)\n",
        "Cs = np.zeros((n_trials, n_participants, n_choices))\n",
        "for n in range(n_choices):\n",
        "    Cs[:, :, n] = (Python_Cs == n)\n",
        "Cs = np.array(Cs, dtype=np.int32)\n",
        "A = np.tile(alphas.reshape(1,-1,1),[n_trials,1,n_choices])\n",
        "PyTensor_Qs = rl_func(Cs,Rs,A)\n",
        "PyTensor_Qs[:5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gGnFmy2P8s26",
        "outputId": "558aaf8d-6b29-4e36-cb9d-42e07e3360aa"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[[ 0.5       ,  1.73022911],\n",
              "        [ 0.21225973,  0.5       ],\n",
              "        [ 0.07119759,  0.5       ]],\n",
              "\n",
              "       [[ 0.5       ,  0.86919225],\n",
              "        [-0.03987525,  0.5       ],\n",
              "        [-0.52974563,  0.5       ]],\n",
              "\n",
              "       [[ 0.5       ,  1.29766769],\n",
              "        [-0.03987525,  0.75788475],\n",
              "        [-0.52974563,  0.61445943]],\n",
              "\n",
              "       [[ 0.5       ,  1.49373382],\n",
              "        [-0.26587443,  0.75788475],\n",
              "        [-0.52974563, -0.30926341]],\n",
              "\n",
              "       [[ 0.5       ,  1.02332718],\n",
              "        [-0.4027839 ,  0.75788475],\n",
              "        [-2.22131283, -0.30926341]]])"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "Python_Qs[:5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VFjvOoQe8vAM",
        "outputId": "553b4ff9-4482-4fdd-82d6-bccaa077a3ab"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[[ 0.5       ,  1.73022911],\n",
              "        [ 0.21225973,  0.5       ],\n",
              "        [ 0.07119759,  0.5       ]],\n",
              "\n",
              "       [[ 0.5       ,  0.86919225],\n",
              "        [-0.03987525,  0.5       ],\n",
              "        [-0.52974563,  0.5       ]],\n",
              "\n",
              "       [[ 0.5       ,  1.29766769],\n",
              "        [-0.03987525,  0.75788475],\n",
              "        [-0.52974563,  0.61445943]],\n",
              "\n",
              "       [[ 0.5       ,  1.49373382],\n",
              "        [-0.26587443,  0.75788475],\n",
              "        [-0.52974563, -0.30926341]],\n",
              "\n",
              "       [[ 0.5       ,  1.02332718],\n",
              "        [-0.4027839 ,  0.75788475],\n",
              "        [-2.22131283, -0.30926341]]])"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###2c. Getting RL Likelihoods in PyTensor"
      ],
      "metadata": {
        "id": "0E_a4q4mMca4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now that we can generate synthetic RL data and we can simulate the internal state of an RL agent used the information we have about choices and rewards, we can now calculate the log-likelihood of our simulated Q-values.\n",
        "\n",
        "This will be necessary because for human subjects we will not know what their \"parameters\" are. Therefore, we will need a metric to assess how well different parameter values replicate their behavior. The ultimate goal is to locate the parameters which best replicate the behavior of human subjects, and we use the Log-likelihood as a relative measure of how well we are doing. Likelihood is like standard deviation in that it lacks absolute meaning, and so it is only meaningful relative to other measurements. Log-likelihoods are always negative, and less negative values are more likely.\n",
        "\n",
        "**Note:** This is the step of the process which requires us to one-hot the data. In Python we can index arrays with arrays but we cannot do this in PyTensor. So instead what we do is we create a mask over the Q-values using the one-hotted Choices, which allows us to efficiently choose all of the correct Q-values."
      ],
      "metadata": {
        "id": "WMCAm-jJMl9U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def pytensor_likelihood(Qs, B):\n",
        "    \"\"\"\n",
        "    pytensor_softmax: calculate loglikelihoods using a tempered softmax over Q-Values\n",
        "\n",
        "    Qs: Q-Values (data)\n",
        "    pB: Betas (parameter)\n",
        "    \"\"\"\n",
        "    shape = Qs.shape\n",
        "    tempered_qs = pt.mul(Qs,B)\n",
        "    qs_max = pt.max(tempered_qs,axis=2)\n",
        "    qs_max = pt.repeat(qs_max.reshape((shape[0], shape[1], 1)), shape[2], axis=2)\n",
        "    numerator = pt.exp(tempered_qs - qs_max)\n",
        "    denominator = pt.sum(numerator, axis=2)\n",
        "    denominator = pt.repeat(denominator.reshape((shape[0], shape[1], 1)), shape[2], axis=2)\n",
        "    Ps = (numerator / denominator).sum(axis=2)\n",
        "    ll = pt.log(Ps)\n",
        "    return ll.flatten()"
      ],
      "metadata": {
        "id": "4nBxbNsIEdUX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Qs_compile = pt.dtensor3('Qs_compile')\n",
        "B_compile = pt.dtensor3('B_compile')\n",
        "\n",
        "output = pytensor_likelihood(Qs_compile, B_compile)\n",
        "likelihood_func = function(inputs=[Qs_compile, B_compile], outputs=output)"
      ],
      "metadata": {
        "id": "yKR9w1F0Ep78"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "B = np.tile(betas.reshape(1,-1,1),[n_trials,1,n_choices])\n",
        "\n",
        "likelihood_func(PyTensor_Qs, B)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H-JUg8VQF7ag",
        "outputId": "28109d31-d4ba-4c4b-b933-5d290aeccb7b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([ 0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        2.22044605e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16, -1.11022302e-16,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16, -1.11022302e-16,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        2.22044605e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  2.22044605e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  2.22044605e-16, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  2.22044605e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "       -1.11022302e-16,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        2.22044605e-16, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        2.22044605e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00])"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###2d. Making your differentiable blackbox likelihood function"
      ],
      "metadata": {
        "id": "4-6L0oMkRMJK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def pytensor_blackbox_rl_likelihood(CMs,RMs,AM,BM):\n",
        "    Qs = rl_multi_scan(CMs,RMs,AM)\n",
        "    ll = pytensor_likelihood(Qs,BM)\n",
        "    return ll"
      ],
      "metadata": {
        "id": "qyz5oF9RGPCY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "C_compile = pt.itensor3(\"C_compile\")\n",
        "R_compile = pt.dtensor3(\"R_compile\")\n",
        "A_compile = pt.dtensor3(\"A_compile\")\n",
        "B_compile = pt.dtensor3('B_compile')\n",
        "\n",
        "output = pytensor_blackbox_rl_likelihood(C_compile, R_compile, A_compile, B_compile)\n",
        "pt_bb_rl_ll_func = function(inputs=[C_compile, R_compile, A_compile, B_compile], outputs=output)"
      ],
      "metadata": {
        "id": "BiMnMdkVI15T"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "alphas, betas = multi_rl.get_params()\n",
        "\n",
        "# Testing compiled function\n",
        "Rs = np.repeat(Python_Rs.reshape(n_trials,n_participants,1), n_choices, axis=2)\n",
        "Cs = np.zeros((n_trials, n_participants, n_choices))\n",
        "for n in range(n_choices):\n",
        "    Cs[:, :, n] = (Python_Cs == n)\n",
        "Cs = np.array(Cs, dtype=np.int32)\n",
        "\n",
        "pt_bb_rl_ll_func(Cs, Rs, A, B)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A0NRazmzJXaz",
        "outputId": "95a64e86-6ef8-4728-bea0-bf698d447b0a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([ 0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        2.22044605e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16, -1.11022302e-16,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16, -1.11022302e-16,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        2.22044605e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  2.22044605e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  2.22044605e-16, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  2.22044605e-16,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00, -1.11022302e-16,\n",
              "       -1.11022302e-16,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        2.22044605e-16, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        2.22044605e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "       -1.11022302e-16,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
              "        0.00000000e+00,  0.00000000e+00, -1.11022302e-16,  0.00000000e+00,\n",
              "        0.00000000e+00, -1.11022302e-16,  0.00000000e+00,  0.00000000e+00])"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    }
  ]
}