numerics.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26678bbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import device_put"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0065b86",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pennylane as qml\n",
    "import numpy as np\n",
    "from IPython import display\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "from math import sqrt,pi\n",
    "from tqdm import tqdm\n",
    "from itertools import product\n",
    "import matplotlib.pyplot as plt\n",
    "import optax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15a41be5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#seeds used for random functions\n",
    "from jax import random\n",
    "key = random.PRNGKey(999)\n",
    "np.random.seed(999)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e122d62",
   "metadata": {},
   "source": [
    "# data generation: rps model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d629fc9",
   "metadata": {},
   "source": [
    "Below we generate the RPS data set as described in the manuscript."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36646bbc",
   "metadata": {},
   "source": [
    "first we construct a 'points' matricies for each pair of players, where a point is awarded if a player beats another and deducted if they are beaten: `A[i,j]` gives the points awarded to player i when playing player j for each of the 9 possible pairs of actions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9052694",
   "metadata": {},
   "outputs": [],
   "source": [
    "A=np.zeros([3,3,3,3])\n",
    "A[0,1]=np.array([[1,-1,1],[1,-1,-1],[-1,1,0]])\n",
    "A[0,2]=np.array([[1,-1,1],[1,0,-1],[-1,1,-1]])\n",
    "A[1,2]=np.array([[0,-1,1],[1,1,-1],[-1,1,-1]])\n",
    "A[1,0]=-A[0,1].T\n",
    "A[2,0]=-A[0,2].T\n",
    "A[2,1]=-A[1,2].T\n",
    "A = jnp.array(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c511716d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_payoff_probs(X):\n",
    "    \"\"\"\n",
    "    get the payoff probabilities for each player given a strategy X\n",
    "    \"\"\"\n",
    "    points0 = jnp.matmul(jnp.matmul(X[0],A[0,1]),X[1])+jnp.matmul(jnp.matmul(X[0],A[0,2]),X[2])\n",
    "    points1 = jnp.matmul(jnp.matmul(X[1],A[1,0]),X[0])+jnp.matmul(jnp.matmul(X[1],A[1,2]),X[2])\n",
    "    points2 = jnp.matmul(jnp.matmul(X[2],A[2,0]),X[0])+jnp.matmul(jnp.matmul(X[2],A[2,1]),X[1])\n",
    "    probs = jnp.array([points0,points1,points2])\n",
    "    return (probs/2+1)/2\n",
    "\n",
    "#JAX vectorisation\n",
    "vpayoff_probs = jax.vmap(get_payoff_probs,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bd03bf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_strat_mat(N):\n",
    "    \"\"\"\n",
    "    This generates N strategy matrices, normalised by row\n",
    "    \"\"\"\n",
    "    pmat = np.random.rand(N,3,3)\n",
    "    for i in range(N):\n",
    "        for k in range(3):\n",
    "            pmat[i,k]=pmat[i,k]/np.sum(pmat[i,k])\n",
    "    return pmat\n",
    "\n",
    "def generate_data(N):\n",
    "    X = get_strat_mat(N) #strategies\n",
    "    P = vpayoff_probs(X) #payoff probabilities\n",
    "    r=np.random.rand(*P.shape) \n",
    "    Y = np.where(P>r,0,1) #sampled payoffs for data labels\n",
    "    return X, Y, P\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f8e5f98",
   "metadata": {},
   "source": [
    "# Define learning models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b712eb6d",
   "metadata": {},
   "source": [
    "Here we define the quantum and surrogate models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94460c81",
   "metadata": {},
   "outputs": [],
   "source": [
    "#some functions used to cosntruct the quantum models\n",
    "\n",
    "def swap_rot(x,wires):\n",
    "    #two qubit rotation with swap matrix as generator \n",
    "    #(identiy is not included since it adds a global phase)\n",
    "    qml.PauliRot(x,'XX',wires=wires)\n",
    "    qml.PauliRot(x,'YY',wires=wires)\n",
    "    qml.PauliRot(x,'ZZ',wires=wires)\n",
    "    \n",
    "def data_encoding(x):\n",
    "    #S_x^1 in paper\n",
    "    for q in range(3):\n",
    "        qml.RZ(x[q],wires=q)\n",
    "\n",
    "def data_encoding_pairs(x):\n",
    "    #S_x^2 in paper\n",
    "    qml.PauliRot(x[0]*x[1],'ZZ',wires=[0,1])\n",
    "    qml.PauliRot(x[1]*x[2],'ZZ',wires=[1,2])\n",
    "    qml.PauliRot(x[0]*x[2],'ZZ',wires=[0,2])\n",
    "    \n",
    "def state_prep(alpha):\n",
    "    #V in paper for biased model\n",
    "    qml.RY(alpha[0],wires=0)\n",
    "    qml.RY(alpha[0]+pi,wires=1) \n",
    "    \n",
    "def layer_q(weights,blocks):\n",
    "    #U in paper for the biased model\n",
    "    for b in range(blocks):\n",
    "        for q in range(3):\n",
    "            qml.RZ(weights[b,q],wires=q)\n",
    "            qml.PauliRot(weights[b,3],'ZZ',wires=[0,1])\n",
    "            qml.PauliRot(weights[b,4],'ZZ',wires=[0,2])\n",
    "            qml.PauliRot(weights[b,5],'ZZ',wires=[1,2])\n",
    "        swap_rot(weights[b,6],wires=[0,1])\n",
    "        swap_rot(weights[b,7],wires=[1,2])\n",
    "        swap_rot(weights[b,8],wires=[0,2])\n",
    "        \n",
    "def layer_u(weights,blocks):\n",
    "    #U in paper for the unbiased model\n",
    "    for b in range(blocks):\n",
    "        for q in range(3):\n",
    "            qml.Rot(weights[b,3*q],weights[b,3*q+1],weights[b,3*q+2],wires=q)\n",
    "        qml.CNOT(wires=[0,1])\n",
    "        qml.CNOT(wires=[0,1])\n",
    "        qml.CNOT(wires=[0,1])\n",
    "       "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb34c90f",
   "metadata": {},
   "source": [
    "## Quantum models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a20f38e6",
   "metadata": {},
   "source": [
    "Define the two quantum model classes (biased and unbiased)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e57bd1ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_q_model(basic=False):\n",
    "    #The biased quantum model\n",
    "    dev = qml.device('default.qubit',wires=3)\n",
    "    @qml.qnode(dev, interface=\"jax\")\n",
    "    def model(weights,x,layers,blocks):\n",
    "        qml.Hadamard(wires=0)\n",
    "        qml.Hadamard(wires=1)\n",
    "        qml.Hadamard(wires=2)\n",
    "        #we use the last layer of the weights array to store the parameter V\n",
    "        state_prep(weights[2*layers+1,0])\n",
    "        x1 = jnp.array([x[0,0],x[1,1],x[2,2]])\n",
    "        x2 = jnp.array(([x[0,1]-x[0,2],x[1,2]-x[1,0],x[2,0]-x[2,1]]))\n",
    "        for l in range(0,2*layers,2):\n",
    "            layer_q(weights[l],blocks)\n",
    "            data_encoding(x1)\n",
    "            layer_q(weights[l+1],blocks)\n",
    "            data_encoding_pairs(x2)\n",
    "        layer_q(weights[2*layers],blocks)\n",
    "        return [qml.expval(qml.PauliZ(0)),qml.expval(qml.PauliZ(1)),qml.expval(qml.PauliZ(2))]\n",
    "    #jax vectorisation, we vectorise over the data input (the second argument)\n",
    "    vmodel = jax.vmap(model,(None,0,None,None))\n",
    "    return vmodel\n",
    "\n",
    "def get_q_model_unstr(basic=False):\n",
    "    #The unbiased quantum model\n",
    "    dev = qml.device('default.qubit',wires=3)\n",
    "    @qml.qnode(dev, interface=\"jax\")\n",
    "    def model(weights,x,layers,blocks):\n",
    "        x1 = jnp.array([x[0,0],x[1,1],x[2,2]])\n",
    "        x2 = jnp.array([x[0,1]-x[0,2],x[1,2]-x[1,0],x[2,0]-x[2,1]])\n",
    "        for l in range(0,layers*2,2):\n",
    "            layer_u(weights[l],blocks)\n",
    "            data_encoding(x1)\n",
    "            layer_u(weights[l+1],blocks)\n",
    "            data_encoding_pairs(x2)\n",
    "        layer_u(weights[2*layers],blocks)\n",
    "        return [qml.expval(qml.PauliZ(0)),qml.expval(qml.PauliZ(1)),qml.expval(qml.PauliZ(2))]\n",
    "    vmodel = jax.vmap(model,(None,0,None,None))\n",
    "    return vmodel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18ea8b88",
   "metadata": {},
   "source": [
    "## Surroagte models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ff2d84e",
   "metadata": {},
   "source": [
    "Define the surrogate model class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6d8773f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_freqs(nvars,spectrum):\n",
    "        \"\"\"\n",
    "        Get the array of all possible frequencies omega appearing in the truncated fourier series\n",
    "        Where the number of data variables is nvars, and the largest value of any component of \n",
    "        omega is spectrum. \n",
    "        \"\"\"\n",
    "        nlambda = 2*spectrum+1\n",
    "        freqs= [jnp.unravel_index(i,[nlambda for __ in range(nvars)]) for i in range(nlambda**nvars)]\n",
    "        freqs = jnp.array(freqs)\n",
    "        freqs=freqs-spectrum\n",
    "        return freqs\n",
    "\n",
    "def get_surrogate_model():\n",
    "    def surrogate(weights,x,freqs,arg2=0):\n",
    "        x1 = jnp.array([x[0,0],x[1,1],x[2,2]])\n",
    "        x2 = jnp.array([x[0,1]-x[0,2],x[1,2]-x[1,0],x[2,0]-x[2,1]])\n",
    "        variables = jnp.concatenate((x1,\n",
    "                                    jnp.array([x2[0]*x2[1] ,x2[1]*x2[2], x2[0]*x2[2]])))\n",
    "        #get all the omega*x values appearing in the fourier sum\n",
    "        z = jnp.sum(freqs*variables.flatten(),axis=1)\n",
    "        waves_cos = jnp.cos(z)\n",
    "        waves_sin = jnp.sin(z)\n",
    "        #build the linear models for each component via matrix multiplication\n",
    "        y=jnp.matmul(weights,jnp.hstack([waves_cos,waves_sin]))\n",
    "        #below implements the hardtanh function\n",
    "        signs = jnp.sign(y)\n",
    "        y=jnp.abs(y)\n",
    "        y=jnp.vstack([y,jnp.ones(3)])\n",
    "        y=jnp.min(y,axis=0)\n",
    "        y=y*signs\n",
    "        return  y\n",
    "    #JAX vectorisation\n",
    "    vsurrogate = jax.vmap(surrogate,(None,0,None,None))\n",
    "\n",
    "    return vsurrogate\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "128fecaf",
   "metadata": {},
   "source": [
    "Some functions used in training an evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ded2e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def kl_div(p,q):\n",
    "    \"\"\"\n",
    "    Get the KL divergence between two probability distribtuions\n",
    "    \"\"\"\n",
    "    p=jnp.vstack([p,jnp.ones(len(p))*10**(-8)]) #lower cutoff of prob values of 10e-8\n",
    "    p=jnp.max(p,axis=0)\n",
    "    return jnp.sum(q*jnp.log(q/p)) #forward kl div \n",
    "\n",
    "def kl_mean(probs,probs0):\n",
    "    \"\"\"\n",
    "    get the mean KL divergence of the three marginal distributions\n",
    "    \"\"\"\n",
    "    kl = 0\n",
    "    for t in range(3):\n",
    "        kl=kl+kl_div(probs[t,:],probs0[t,:])\n",
    "    return kl/3\n",
    "\n",
    "#vectorise kl_mean function, makes estimating the average KL diverence of a model faster.\n",
    "vkl_mean = jax.vmap(kl_mean,(0,0))\n",
    "\n",
    "def model_probs(model,X_test,weights,arg1,arg2):\n",
    "    \"\"\"\n",
    "    Returns the marginal probabilties of a given model for a test data set X_test\n",
    "    \"\"\"\n",
    "    probs = np.zeros([len(X_test),3,2])\n",
    "    expvals = model(weights,X_test,arg1,arg2)\n",
    "    for t in range(3):\n",
    "        probs[:,t,0] = (1+expvals[:,t])/2\n",
    "        probs[:,t,1] = (1-expvals[:,t])/2\n",
    "    return probs\n",
    "\n",
    "def likelihood(weights,X,Y,model,arg1,arg2,reg):\n",
    "    \"\"\"\n",
    "    The cost function. Returns the negative log likelihood plus a L2  weight regularisation penalty\n",
    "    \"\"\"\n",
    "    expvals = model(weights,X,arg1,arg2)\n",
    "    probs = (1+(1-2*Y)*expvals)/2\n",
    "    probs = jnp.log(probs)\n",
    "    llh = jnp.sum(probs)/len(X)/3\n",
    "    return -llh + jnp.sum(weights**2)*reg\n",
    "\n",
    "def gen_test_set(N_test):\n",
    "    \"\"\"\n",
    "    generates a test set of strategies\n",
    "    \"\"\"\n",
    "    X_test = get_strat_mat(N_test)\n",
    "    return X_test\n",
    "\n",
    "def get_av_test_kl(model,weights,probs0,X_test,arg1,arg2):\n",
    "    \"\"\"\n",
    "    returns the average KL divergence of a model for a test set X_test. \n",
    "    This is the figure of merit used to evaluate the generalisation performance.\n",
    "    \"\"\"\n",
    "    N_test=len(X_test)\n",
    "    probs = np.zeros(probs0.shape)\n",
    "    expvals = model(weights,X_test,arg1,arg2)\n",
    "    for t in range(3):\n",
    "        probs[:,t,0] = (1+expvals[:,t])/2\n",
    "        probs[:,t,1] = (1-expvals[:,t])/2\n",
    "    return np.sum(vkl_mean(probs,probs0))/N_test\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1753073b",
   "metadata": {},
   "source": [
    "We use the optax library to optimise the models. For the quanutm model arg1 and arg2 correspond to L and B in the paper. For the surrogate model arg1 is an array containing all values of omega, and arg2 is unused. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4c4801c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimise_model_adam(model,nstep,lr,weights,batchfrac=1.0,sgd=False,arg1=1,arg2=1,reg=0.):\n",
    "    plot=[[],[],[]]\n",
    "    X_test = gen_test_set(10000) #the test set for evaluation\n",
    "    N_test=len(X_test)\n",
    "    probs0 = np.zeros([N_test,3,2])\n",
    "    probs0[:,:,0] = vpayoff_probs(X_test) #the true probabilities for the test set\n",
    "    probs0[:,:,1] = 1-probs0[:,:,0]\n",
    "    probs0=jnp.array(probs0)\n",
    "    if sgd:\n",
    "        optimizer = optax.sgd(lr)\n",
    "    else:\n",
    "        optimizer = optax.adam(lr) #this is default\n",
    "    opt_state = optimizer.init(weights)\n",
    "    steps = tqdm(range(nstep))\n",
    "    for step in steps:\n",
    "        #use optax to update parameters\n",
    "        llh, grads = jax.value_and_grad(likelihood)(weights, X, Y, model,arg1,arg2,reg)\n",
    "        updates, opt_state = optimizer.update(grads, opt_state, weights)\n",
    "        weights = optax.apply_updates(weights, updates)\n",
    "        if step%1==0:\n",
    "            kl = get_av_test_kl(model,weights,probs0,X_test,arg1,arg2)\n",
    "            steps.set_description(\"Current divergence: %s\" % str(kl)+ \" :::: \"+\n",
    "                                  \"Current likelihood: %s\" % str(llh))\n",
    "        plot[0].append(step)\n",
    "        plot[1].append(float(llh))\n",
    "        plot[2].append(float(kl))\n",
    "    return weights, llh, kl, plot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67d26733",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d8f083",
   "metadata": {},
   "outputs": [],
   "source": [
    "#generate data\n",
    "N=1500 #number of data points\n",
    "\n",
    "X = np.load('Xdata.npy')\n",
    "Y = np.load('Ydata.npy')\n",
    "\n",
    "X = jnp.array(X)\n",
    "Y = jnp.array(Y)\n",
    "# P = jnp.array(P)\n",
    "\n",
    "X_copy = jnp.array(X)\n",
    "\n",
    "#scale the data by pi/2\n",
    "scaling = pi/2\n",
    "X = scaling*jnp.array(X_copy)\n",
    "\n",
    "#put the data on the GPU if there is one\n",
    "X = device_put(X)\n",
    "Y = device_put(Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0272b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = 2000 #number of optimisation steps\n",
    "n_runs = 15 #number of runs (or trials)\n",
    "\n",
    "layers_list = [1,2] #the possible values of L we consider for the quantum models\n",
    "spectrum_list = [1,2] #the possible spectra we consider for the surrogate\n",
    "reg = 0. #the L2 weight regularisation penalty coefficient\n",
    "nvars=6 #the number of data variables in the fourier series of the surrogate model\n",
    "blocks = 2 #the value of B in the paper\n",
    "\n",
    "#arrays for storing the plot data\n",
    "data_q = [np.zeros([n_runs,n_steps,3]) for __ in range(len(layers_list))]\n",
    "data_u = [np.zeros([n_runs,n_steps,3]) for __ in range(len(layers_list))]\n",
    "data_s = [np.zeros([n_runs,n_steps,3]) for __ in range(len(spectrum_list))] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eea7f4a5",
   "metadata": {},
   "source": [
    "## Train the models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eee4df23",
   "metadata": {},
   "source": [
    "Quatum biased bodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6b3d18d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i,lay in enumerate(layers_list):\n",
    "    layers=lay\n",
    "    vmodel = get_q_model()\n",
    "    vmodel=jax.jit(vmodel,static_argnums=[2,3])\n",
    "    for run in range(n_runs):\n",
    "        weights_q = np.random.rand(2*layers+2,blocks,9)*2*pi\n",
    "        weights_q = device_put(weights_q)\n",
    "        weights_q, llh, kl, plot = optimise_model_adam(vmodel,n_steps,.001,weights_q,sgd=False,arg1=layers,arg2=blocks)\n",
    "        data_q[i][run,:,0]=plot[0]\n",
    "        data_q[i][run,:,1]=plot[1]\n",
    "        data_q[i][run,:,2]=plot[2]\n",
    "                \n",
    "np.save('data_q',np.array(data_q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66ed3638",
   "metadata": {},
   "source": [
    "Quantum unbiased model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e647b050",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i,lay in enumerate(layers_list):\n",
    "    layers=lay\n",
    "    vmodel_unstr = get_q_model_unstr()\n",
    "    vmodel_unstr = jax.jit(vmodel_unstr,static_argnums=[2,3])\n",
    "    for run in range(n_runs):\n",
    "        weights_u = np.random.rand(2*layers+1,blocks,9)*2*pi\n",
    "        weights_u = device_put(weights_u)\n",
    "        weights_u, llh, kl, plot = optimise_model_adam(vmodel_unstr,n_steps,.001,weights_u,sgd=False,arg1=layers,arg2=blocks)\n",
    "        data_u[i][run,:,0]=plot[0]\n",
    "        data_u[i][run,:,1]=plot[1]\n",
    "        data_u[i][run,:,2]=plot[2]\n",
    "        \n",
    "np.save('data_u',np.array(data_u))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8d020d2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for i,spectrum in enumerate(spectrum_list):\n",
    "    n_freq = 2*(2*spectrum+1)**6\n",
    "    freqs= get_freqs(nvars,spectrum)\n",
    "    if spectrum==1:\n",
    "        lr=0.001\n",
    "    else:\n",
    "        lr=0.0001\n",
    "    vsurrogate = get_surrogate_model()\n",
    "    vsurrogate=jax.jit(vsurrogate)\n",
    "    for run in range(n_runs):\n",
    "        weights_s = (2*np.random.rand(3,n_freq)-1)*1/n_freq\n",
    "        weights_s = device_put(weights_s)\n",
    "        weights_s, llh, kl, plot = optimise_model_adam(vsurrogate,n_steps,lr,weights_s,sgd=True,arg1=freqs,arg2=0,reg=reg)\n",
    "        data_s[i][run,:,0]=plot[0]\n",
    "        data_s[i][run,:,1]=plot[1]\n",
    "        data_s[i][run,:,2]=plot[2]\n",
    "\n",
    "np.save('data_s',np.array(data_s))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e40cb89b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#calculate average values over all trials for each model\n",
    "data_q_av = [np.mean(data_q[l],axis=0) for l in range(1)]\n",
    "data_u_av = [np.mean(data_u[l],axis=0) for l in range(1)]\n",
    "data_s_av = [np.mean(data_s[0],axis=0),np.mean(data_s[1],axis=0)] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5142514",
   "metadata": {},
   "source": [
    "Plot the KL divergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f1574d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "steplist = np.arange(0,n_steps,1)\n",
    "\n",
    "for run in range(n_runs):\n",
    "    plt.plot(data_q[0][run,:,0],data_q[0][run,:,2],alpha=0.2,color='red')\n",
    "    plt.plot(data_u[0][run,:,0],data_u[0][run,:,2],alpha=0.2,color='blue')\n",
    "    plt.plot(data_s[0][run,:,0],data_s[0][run,:,2],alpha=0.2,color='green')\n",
    "    plt.plot(data_s[1][run,:,0],data_s[1][run,:,2],alpha=0.2,color='black')\n",
    "    \n",
    "plt.plot(data_q_av[0][:,0],data_q_av[0][:,2],alpha=1.0,color='red',label='q_biased')\n",
    "plt.plot(data_u_av[0][:,0],data_u_av[0][:,2],alpha=1.0,color='blue',label='q_unbiased')\n",
    "plt.plot(data_s_av[0][:,0],data_s_av[0][:,2],alpha=1.0,color='green',label='surrogate $\\Omega$=1')\n",
    "\n",
    "plt.plot(data_s_av[0][:,0],data_s_av[1][:,2],alpha=1.0,color='black',label='surrogate $\\Omega$=2')\n",
    "\n",
    "plt.yscale(\"log\") \n",
    "plt.legend()\n",
    "plt.xlabel(\"optimisation step\")\n",
    "plt.ylabel(\"av kl div\")\n",
    "plt.legend(loc='upper right')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adef6897",
   "metadata": {},
   "source": [
    "plot the negative log likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68fc174e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for run in range(n_runs):\n",
    "    plt.plot(data_q[0][run,:,0],data_q[0][run,:,1],alpha=0.2,color='red')\n",
    "    plt.plot(data_u[0][run,:,0],data_u[0][run,:,1],alpha=0.2,color='blue')\n",
    "    plt.plot(data_s[0][run,:,0],data_s[0][run,:,1],alpha=0.2,color='green')\n",
    "    plt.plot(data_s[1][run,:,0],data_s[1][run,:,1],alpha=0.2,color='black')\n",
    "    \n",
    "plt.plot(data_q_av[0][:,0],data_q_av[0][:,1],alpha=1.0,color='red',label='q_biased')\n",
    "plt.plot(data_u_av[0][:,0],data_u_av[0][:,1],alpha=1.0,color='blue',label='q_unbiased')\n",
    "plt.plot(data_s_av[0][:,0],data_s_av[0][:,1],alpha=1.0,color='green',label='surrogate $\\Omega$=1')\n",
    "\n",
    "plt.plot(data_s_av[0][:,0],data_s_av[1][:,1],alpha=1.0,color='black',label='surrogate $\\Omega$=2')\n",
    "\n",
    "plt.yscale(\"log\") \n",
    "plt.legend()\n",
    "plt.xlabel(\"optimisation step\")\n",
    "plt.ylabel(\"negative log likelihood\")\n",
    "plt.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1de068c4",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}