QGML.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum Geometric Machine Learning: Notebook\n",
    "\n",
    "This notebook contains code for the QMLcontrolmodel class, the class that instantiates the various SubRiemannian, GRU RNN Greybox and Fully-connected Greybox models from the paper <i>Quantum Geometric Machine Learning for Quantum Circuits and Control </i> by Perrier et al (2020).\n",
    "\n",
    "It also contains instructions on how to call the subRiemannian normal geodesic data generating simulation.py code to generate training (and test) data and how to extract control pulses and unitary estimates from each model.\n",
    "It also includes example code for generalisation with respect to various test sets. It contains extensive comments to assist researchers. The code is built using TensorFlow 2.2 and Qutip > 4.0 among other packages.\n",
    "\n",
    "Training times vary, with the FC Greybox model taking the longest on a Dell XPS 9850 i7 laptop (without using a GPU)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers,optimizers,constraints,initializers,Model,backend\n",
    "import pickle\n",
    "import time\n",
    "import zipfile    \n",
    "import os\n",
    "from itertools import combinations, accumulate\n",
    "from operator import itemgetter\n",
    "import qutip as qt\n",
    "import scipy as sc\n",
    "import random as rn\n",
    "import operator\n",
    "from functools import reduce\n",
    "import csv\n",
    "\n",
    "'''Optional for plotting using Qutip'''\n",
    "from matplotlib import pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from mpl_toolkits import mplot3d\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML, Image\n",
    "%matplotlib inline\n",
    "\n",
    "\n",
    "from flattenunitary import flattenunitary\n",
    "from flattenunitary import realise\n",
    "\n",
    "\n",
    "'''simulation.py contains the code for generating subRiemannian geodesic training data'''\n",
    "from simulation import GenData\n",
    "'''customlayers.py contains the various customised greybox layers used in the models'''\n",
    "from customlayers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data generation\n",
    "\n",
    "Description:\n",
    "\n",
    "GenData: a class from simulation.py that generates approximations to geodesics for use\n",
    "as training data. \n",
    "\n",
    "Inputs are:\n",
    "\n",
    "GenData(su2ndim, nseg, ntrain, time, holo, su2ngen, hscale)\n",
    "\n",
    "with the following arguments:\n",
    "* su2ndim: set n, used to set the dimension the Lie group SU($2^n$) from which unitaries are drawn;\n",
    "* nseg: the number of segments of the approximate geodesic sequences $(U_j)$;\n",
    "* ntrain: the number of training examples (number of approximate geodesic sequences);\n",
    "* time: option in simulation where given number of segments nseg, can set\n",
    "h = delta_t (the time-step for each unitary in the independent approximation\n",
    "to Schrodinger's equation) to be time/nseg;\n",
    "* holo: option to as to whether first unitary in loop generating $(U_j)$ is to be unitary generated\n",
    "via the implementation of method in Boozer (2012) (used in Appendix to paper to compare \n",
    "subRiemannian method of generating geodesic approximations to the analytic method in Boozer (2012));\n",
    "* su2ngen: Specify generators for initial condition $\\Lambda_0$: \n",
    "    * 0 $: \\Lambda_0 \\in \\Delta$ (the distribution) and \\\\\n",
    "    * 1$: \\Lambda_0 \\in su(2^n)$ (full Lie algebra)\n",
    "* hscale: Set $\\Delta t_j$ (i.e. h in code or duration of pulse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''========================Example: generate and save data============================'''\n",
    "'''Generate geodesic dataset by specifying hyperparameters and save as pickle file'''\n",
    "'''Set dimension via n for SU(2^n).'''\n",
    "nsu2dim = 1\n",
    "'''Set the number of segements of the geodesic approximation'''\n",
    "nseg = 5\n",
    "'''Set the number of training examples'''\n",
    "ntrain = 50\n",
    "'''Set \\delta tj (i.e. h in code or duration of pulse)'''\n",
    "hscale = 0.1\n",
    "'''Specify generators for initial condition \\Lambda_0:\n",
    "0: \\Lambda_0 in \\Delta (the distribution); \n",
    "1: \\Lambda_0 \\in su(2^n) (full Lie algebra)'''\n",
    "su2gen = None #\n",
    "'''Specify whether want to initialise unitary sequence using Boozer (2012) unitary, usually set to None'''\n",
    "holo = None #whether training data to be generated from Boozer\n",
    "\n",
    "'''Generate data'''\n",
    "gentest = GenData(nsu2dim,nseg,ntrain,hscale,holo=holo, su2ngen=su2gen, hscale=hscale)\n",
    "\n",
    "'''Convert data to form to save as pickle file'''\n",
    "a = vars(gentest)\n",
    "\n",
    "'''Note: using above variables allows automatic saving of name according to hyperparameters above'''\n",
    "with open('su{0}_ntrain{1}_nseg{2}_hscale{3}_su2gen{4}.pickle'.format(\n",
    "gentest.su2n_dim,gentest.ntraining, gentest.segments, gentest.h, gentest.su2ngen), 'wb') as handle:\n",
    "    pickle.dump(a, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QMLcontrol model\n",
    "\n",
    "Code below constructs a class QMLcontrolmodel to implement the three greybox quantum machine learning (QML) models that are described in the paper. Comments on functionality are in the code.\n",
    "\n",
    "The class contains the following:\n",
    "* Attributes such as the number of segments, training examples, distribution (control subalgebra generators) and so on\n",
    "* Methods for preprocessing data\n",
    "* Three primary models as methods: the SubRiemannian model, the Fully-connected (FC) Greybox model and the GRU RNN (LSTM) Greybox model\n",
    "* methods for training and predicting using such models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QMLcontrolmodel():\n",
    "    \"\"\"\n",
    "    This class defines the machine learning models determining controls (generator coefficients).\n",
    "    Arguments: \n",
    "    - gentest: the dataset class (see above) generated from simulation.py containing datasets in \n",
    "    various forms as attributes\n",
    "    - genseg: set to 1\n",
    "    \"\"\" \n",
    "    def __init__(self,gentest, genseg):\n",
    "        '''Number of training examples'''\n",
    "        self.ntraining = gentest.ntraining\n",
    "        '''Convert identity in SU(2^n) to a tensor'''\n",
    "        self.su2nidentity = tf.convert_to_tensor(gentest.su_dim_identity)\n",
    "        '''Set dimension for SU(2^n)'''\n",
    "        self.su2n_dim = gentest.su2n_dim\n",
    "        '''Set dimension of realised SU(2^n) (i.e. dim(SU(2^n) * 2))'''\n",
    "        self.su2n_realdim = gentest.su2n_dim * 2\n",
    "        '''Set dimension of flattened realised unitary (for use in inputs to models below)'''\n",
    "        self.flatrealdim = self.su2n_realdim ** 2\n",
    "        '''Number of generators and coefficients used for prediction (legacy)'''\n",
    "        self.genseg = genseg\n",
    "        '''Set \\Delta (distribution i.e. control subalgebra of Lie algebra) as attribute'''\n",
    "        self.su_2n_delta = gentest.su_2n_delta\n",
    "        '''Set as attribute hyperparameter determining if generators for \\Lambda_0 drawn \n",
    "        from distribution or su(2^n)'''\n",
    "        self.su2ngen = gentest.su2ngen\n",
    "        '''Set dimension of control subalgebra (distribution \\Delta)'''\n",
    "        self.deltadim = np.asarray(gentest.su_2n_delta).shape[0]\n",
    "        '''Set as attribute full Lie algebra su(2^n)'''\n",
    "        self.su_2n = gentest.su_2n\n",
    "        \n",
    "        '''Conditional logic to ensure correct coefficients and generators for \\Lambda_0 \n",
    "        used for SubRiemannian model depending on su2ngen value (distribution or full su(2^n)) '''\n",
    "        if self.su2ngen == 1:\n",
    "            self.genlist = self.su_2n * self.genseg\n",
    "            \n",
    "            self.coef_numb = len(self.su_2n) * self.genseg\n",
    "        else:\n",
    "            self.genlist = self.su_2n_delta * self.genseg\n",
    "            \n",
    "            self.coef_numb = len(self.su_2n_delta) * self.genseg\n",
    "        \n",
    "        '''Set number of segments for approximate geodesic sequence (U)j'''\n",
    "        self.segments = gentest.segments\n",
    "        '''Set attribute number of segments x number of training examples, used in models below'''\n",
    "        self.nseg = gentest.segments * gentest.ntraining\n",
    "        '''self.tot is for tensor-based calculations where we have \n",
    "        segments x deltadim generators etc'''\n",
    "        self.tot = self.deltadim * self.segments\n",
    "        #self.coefzeros = np.asarray([0] * self.coef_numb)\n",
    "        '''For main SubRiemannian model, set list for training history'''\n",
    "        self.training_history = []\n",
    "        '''For main SubRiemannian model, set list for validation history'''\n",
    "        self.val_history = []\n",
    "        '''Set vector of ones used for MSE with batch fidelity calculations'''\n",
    "        self.fidones = gentest.fidones\n",
    "        \n",
    "        #self.fidones2 = np.asarray([1] * 3 * self.ntraining)\n",
    "        '''Set vector of ones used for MSE with batch fidelity calculations in certain models'''\n",
    "        self.fidones_seg = tf.ones((self.ntraining,self.segments))\n",
    "        \n",
    "        '''Set Ujarray (a list containing realised sequences (Uj)) as an array'''\n",
    "        self.Ujarray = np.asarray(gentest.Uj_master_realised)\n",
    "        '''Set Ujarray (a list containing sequences (Uj)) as an array'''\n",
    "        self.Ujmasterarray = np.asarray(gentest.Uj_master_array)\n",
    "        \n",
    "        '''Set attribute that calls xtrainmethod(), a data preprocessing step for inputs into models.\n",
    "        Specifically, the method takes as inputs (U_j) and outputs the final target unitary U_T\n",
    "        by accumulating such that U_n @ ... @ U_0 = U_T'''\n",
    "        self.xtrain = self.xtrainmethod()\n",
    "        \n",
    "        self.Ujm1 = np.asarray(gentest.Uj_master_realised).astype('float64')\n",
    "        \n",
    "        '''Set real and imaginary parts of (Uj) for inputs into models'''\n",
    "        self.Ujseq_real = (tf.math.real(np.asarray(gentest.Uj_master_array)))\n",
    "        self.Ujseq_imag = (tf.math.imag(np.asarray(gentest.Uj_master_array)))\n",
    "        '''Set time-step h = \\delta t'''\n",
    "        self.h = gentest.h \n",
    "        '''Set attribute - number of coefficients equal to number of segments x number of generators'''\n",
    "        self.coef_numb_seg = self.segments * tf.convert_to_tensor(gentest.su_2n_delta).shape[0]\n",
    "        \n",
    "        \n",
    "        \n",
    "        '''===Initialise models from methods'''\n",
    "        '''SubRiemannian model'''\n",
    "        self.mod1 = self.mod1()\n",
    "        '''GRU RNN (LSTM) model'''\n",
    "        self.mod_GRU_grey = self.LSTMGRUGrey()\n",
    "        '''Fully-connected (FC) Greybox model'''\n",
    "        self.modfc = self.modfc()\n",
    "    \n",
    "    \n",
    "    '''===============Model: SubRiemannian model====================='''\n",
    "    def mod1(self):\n",
    "        '''Description: takes as inputs U_T, together with real and imaginary parts of the\n",
    "        sequence of unitaries (U_j) as inputs and outputs estimated (\\hat(U)_j) which are\n",
    "        then compared with the true (U_j) using fidelity. This fidelity is then compared with\n",
    "        a vector of ones (which are technically the labels in this model). The output of the model\n",
    "        is a sequence of fidelities, one fidelity for each Uj for each training example. \n",
    "        Intermediate outputs, such as control pulses, Hamiltonian estimates etc are accessible via \n",
    "        constructing intermediate models in TensorFlow (see below)'''\n",
    "        \n",
    "        '''=====================Custom layer initialisation======================'''\n",
    "        '''Custom layers are contained in customlayers.py'''\n",
    "        \n",
    "        '''===Hamiltonian estimation layers\n",
    "        Description: Initialise a layer that takes controls for each generator in control subalgebra and\n",
    "        multiplies them together for form an estimate of \\Lambda_0\n",
    "        (which is actually a Hamiltonian \\hat{H}_j). Two layers\n",
    "        are used for this purpose, one contains positive control values, the other negative.\n",
    "        Two Hamiltonian estimates for \\Lambda_0, a positive estimate \\Lambda_0+ and negative \\Lambda_0-\n",
    "        are then constructed for each \\Lambda_0, then added together (as are the controls).'''\n",
    "        '''Positive layer'''\n",
    "        ham_est_layer = ham_est(self.genlist,self.genseg, name='ham_layer')\n",
    "        '''Negative layer'''\n",
    "        ham_est_layer_neg = ham_est_neg(self.genlist,self.genseg, name='ham_layer_neg')\n",
    "        \n",
    "        '''========Fidelity layer\n",
    "        Initialise fidelity layer (see customlayers.py for a description)'''\n",
    "        fidelity_layer = fidelity(self.su2n_dim)\n",
    "        \n",
    "        \n",
    "        '''========Unitary output and projection layer\n",
    "        Description: layer recreates the method for generating approximations to normal \n",
    "        subRiemannian geodesics (\\hat{U}_j) using the variational methods described in the paper drawn from\n",
    "        Swaddle (2017). Inputs are \\Lambda_0 and generators, outputs a sequence (\\hat{U}_j).'''\n",
    "        projev = identitylayer(self.su_2n_delta, self.su2nidentity, self.segments, self.h, name='Ujseq_est')\n",
    "        \n",
    "        '''Set dimensions for input'''\n",
    "        origdim = (self.su2n_realdim,self.su2n_realdim)\n",
    "        \n",
    "        '''Set generic number of dense layer inputs (can be modified)'''\n",
    "        intermediate_dim = 640\n",
    "        \n",
    "        '''Original inputs are the realised, flattened, target unitaries U_T. Note, these form the \n",
    "        training examples input into the model.'''\n",
    "        original_inputs = tf.keras.Input(shape=origdim, name='flatinput')\n",
    "        original_inputs1 = layers.Flatten()(original_inputs)\n",
    "        \n",
    "        '''Second inputs are the real and imaginary part of the target unitary,\n",
    "        these are used to compare the estimated unitary against the target unitary.\n",
    "        These are unflattened. The idea is to learn how to accurately input a flattened unitary\n",
    "        constructed from delta generators using learnt weights (controls) and those generators.\n",
    "        The targets are the same unitaries which are then used to calculate fidelities.'''\n",
    "        targetreal = tf.keras.Input(shape=(self.su2n_dim,self.su2n_dim), name='target_real')\n",
    "        targetimag = tf.keras.Input(shape=(self.su2n_dim,self.su2n_dim), name='target_imag')\n",
    "        \n",
    "        '''Clipped relu (customised)'''\n",
    "        clip_relu = 'relu'\n",
    "        '''=====================Dense layers======================\n",
    "        Description: a fully-connected feed-forward intermediate neural network\n",
    "        that takes as input U_T and outputs the control pulses for generating \\Lambda_0 from\n",
    "        its generating set.'''\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(original_inputs1)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        \n",
    "        '''=====================Greybox layers======================'''\n",
    "        '''Output positive controls'''\n",
    "        output = layers.Dense(self.coef_numb,activation='tanh', name='controls')(x)\n",
    "        '''Input positive controls to construct \\Lambda_0+'''\n",
    "        output1 = ham_est_layer(output)\n",
    "        '''Output negative controls'''\n",
    "        output2 = layers.Dense(self.coef_numb,activation='tanh', name='controls2')(x)\n",
    "        '''Input negative controls to construct \\Lambda_0-'''\n",
    "        output2 = ham_est_layer_neg(output2)\n",
    "        \n",
    "        '''Reconstruct \\Lambda_0 from \\Lambda_0+ and \\Lambda_0-'''\n",
    "        Lambda0 = hamcomb(name='hamcomb1')([output1,output2])\n",
    "        '''Input \\Lambda_0 into projection/evolution layer to output (\\hat{U}_j).\n",
    "        Note: this layer outputs a list, first element is (\\hat{U}_j). Second is controls i.e. x1[1]\n",
    "        are the controls generated by the projection function used in the subRiemannian equations\n",
    "        (see paper).'''\n",
    "        x1 = projev(Lambda0)\n",
    "        \n",
    "        '''=====================Ground truth layer======================'''\n",
    "        '''Input real and imaginary part of true (U_j) for batch fidelity layer'''\n",
    "        inputUre = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim), name='input2', batch_size=None)\n",
    "        inputUim = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim), name='input3', batch_size=None)\n",
    "        \n",
    "        '''=====================Output layer======================'''\n",
    "        '''Input the real and imaginary part of true (U_j) and estimate (\\hat{U}_j) into the fidelity function'''\n",
    "        y = fidmulti_unitary(self.su2n_dim, name='finalfidelity')([inputUre,inputUim,x1[0]])\n",
    "        \n",
    "        '''Define model'''\n",
    "        mod1 = tf.keras.Model(inputs=[original_inputs, inputUre, inputUim], outputs=y, name='mod1')\n",
    "        '''Set optimizer (Adam)'''\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
    "        '''Compile, set loss as MSE'''\n",
    "        mod1.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())\n",
    "        '''Print summary (for verification)'''\n",
    "        mod1.summary()\n",
    "        \n",
    "        \n",
    "        '''Return the SubRiemannian model'''\n",
    "        return mod1\n",
    "    \n",
    "    \n",
    "    '''===============Training: SubRiemannian model====================='''\n",
    "    '''Description: trains SubRiemannian model using in-class data.\n",
    "    For other data, the SubRiemannian model method or attribute can be called and \n",
    "    training of it can be done manually.'''\n",
    "    def train_model_mod1(self, epochs):\n",
    "        csv_logger = CSVLogger('train_mod1_su{0}_ntrain{1}_nseg{2}_hscale{3}_su2gen{4}.csv'\n",
    "                       .format(self.su2n_dim,self.ntraining, self.segments, self.h, self.su2ngen))\n",
    "        h = self.mod1.fit([self.xtrain,self.Ujseq_real,self.Ujseq_imag],self.fidones_seg, epochs=epochs, verbose=2, batch_size=10, validation_split = 0.25, shuffle=True, callbacks=[csv_logger])\n",
    "        self.training_history = h.history[\"loss\"]\n",
    "        self.val_history       = h.history[\"val_loss\"]\n",
    "        \n",
    "    '''===============Prediction: SubRiemannian model====================='''\n",
    "    '''Description: generates predictions from trained model.'''\n",
    "    def predict_mod1(self,inputs):\n",
    "        return self.mod1.predict(inputs)\n",
    "    \n",
    "    \n",
    "    '''===============SubModel: SubRiemannian model controls====================='''\n",
    "    '''Description: outputs controls from trained model. Call after training.''' \n",
    "    '''Note: inputs needs to be [xtrain, Ujseq_real, Ujseq_imag] for a particular dataset.'''\n",
    "    def controls_mod1(self,inputs):\n",
    "        new_model = self.mod1\n",
    "        controlmod = Model(inputs=new_model.input, outputs=new_model.get_layer('Ujseq_est').output[1])\n",
    "        return controlmod.predict(inputs)\n",
    "    \n",
    "    \n",
    "    '''===============SubModel: SubRiemannian model unitaries====================='''\n",
    "    '''Description: outputs unitary estimates (\\hat{U}_j) from trained model. Call after training.''' \n",
    "    '''Note: inputs needs to be [xtrain, Ujseq_real, Ujseq_imag] for a particular dataset.'''\n",
    "    def unitaries_mod1(self,inputs):\n",
    "        new_model = self.mod1\n",
    "        unimod = Model(inputs=new_model.input, outputs=new_model.get_layer('Ujseq_est').output[0])\n",
    "        return unimod.predict(inputs)\n",
    "    \n",
    "    def modgeneral(self,inputs):\n",
    "        un = Model(inputs=self.mod1.input, outputs=self.mod1.outputs)\n",
    "        un = un.predict(inputs)\n",
    "        return un\n",
    "    \n",
    "    '''===============Model: Fully-connected (FC) Greybox model====================='''\n",
    "    def modfc(self):\n",
    "        \n",
    "        '''Description: learns controls for sequence U = \\prod_i exp(-\\tau_i)\n",
    "        where \\tau_i are the elements of \\Delta (with imaginary unit already\n",
    "        included). See paper.'''\n",
    "        \n",
    "        '''=====================Custom layer initialisation======================'''\n",
    "        '''Hamiltonian estimation layer, takes as inputs controls and generators, outputs\n",
    "        estimated Hamiltonian sequence (\\hat{H}_j).'''\n",
    "        ham_est_layer_simple = ham_est_simple(self.su_2n_delta,self.segments, name='ham_layer')\n",
    "        '''Quantum evolution layer, takes as input estimated Hamiltonian sequence (\\hat{H}_j)\n",
    "        and outputs estimated unitary sequence (\\hat{U}_j).'''\n",
    "        unitary_simp = unitary_simple(self.segments, self.su2nidentity, self.su2n_dim, name='unitaryoutput')\n",
    "        \n",
    "        '''See description in SubRiemannian layer for below'''\n",
    "        origdim = (self.su2n_realdim,self.su2n_realdim)\n",
    "        intermediate_dim = 640\n",
    "        \n",
    "        original_inputs = tf.keras.Input(shape=origdim, name='flatinput')\n",
    "        original_inputs1 = layers.Flatten()(original_inputs)\n",
    "        \n",
    "        targetreal = tf.keras.Input(shape=(self.su2n_dim,self.su2n_dim), name='target_real')\n",
    "        targetimag = tf.keras.Input(shape=(self.su2n_dim,self.su2n_dim), name='target_imag')\n",
    "        clip_relu = 'relu'\n",
    "        '''=====================Dense layers======================\n",
    "        Description: a fully-connected feed-forward intermediate neural network\n",
    "        that takes as input U_T and outputs the control pulses for generating \n",
    "        estimated Hamiltonian sequence (\\hat{H}_j) using the control subalgebra (\\Delta).'''\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(original_inputs1)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(intermediate_dim, activation=clip_relu)(x)\n",
    "        \n",
    "        '''=====================Greybox layers======================'''\n",
    "        '''Output estimated control pulses'''\n",
    "        output = layers.Dense(self.coef_numb_seg,activation='tanh', name='controls')(x)\n",
    "        '''Combine the controls and generators from \\Delta to construct \n",
    "        estimated Hamiltonian sequence (\\hat{H}_j)'''\n",
    "        output1 = ham_est_layer_simple(output)\n",
    "        '''Quantum evolution layer that outputs estimated unitary sequence (\\hat{U}_j)\n",
    "        from estimated Hamiltonian sequence (\\hat{H}_j)'''\n",
    "        x1 = unitary_simp(output1)\n",
    "        \n",
    "        '''=====================Ground truth layer======================'''\n",
    "        '''Input real and imaginary part of true (U_j) for batch fidelity layer'''\n",
    "        inputUre = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim), name='input2', batch_size=None)\n",
    "        inputUim = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim), name='input3', batch_size=None)\n",
    "        \n",
    "        '''=====================Output layer======================'''\n",
    "        '''Input the real and imaginary part of true (U_j) and estimate (\\hat{U}_j) into the fidelity function'''\n",
    "        y = fidmulti_unitary(self.su2n_dim, name='finalfidelity')([inputUre,inputUim,x1])\n",
    "        \n",
    "        modfc = tf.keras.Model(inputs=[original_inputs, inputUre, inputUim], outputs=y, name='modfc')\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
    "\n",
    "        modfc.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())\n",
    "        \n",
    "        print(modfc.summary())\n",
    "        \n",
    "        return modfc\n",
    "    \n",
    "    \n",
    "    '''===============Training: FC Greybox model====================='''\n",
    "    '''Description: trains FC Greybox model using in-class data.\n",
    "    For other data, the FC Greybox model method or attribute can be called and \n",
    "    training of it can be done manually.'''\n",
    "    def train_model_modfc(self, epochs):\n",
    "        csv_logger = CSVLogger('train_modfc_su{0}_ntrain{1}_nseg{2}_hscale{3}_su2gen{4}.csv'\n",
    "                       .format(self.su2n_dim,self.ntraining, self.segments, self.h, self.su2ngen))\n",
    "        h = self.modfc.fit([self.xtrain,self.Ujseq_real,self.Ujseq_imag],self.fidones_seg, epochs=epochs, verbose=2, batch_size=10, validation_split = 0.25, shuffle=True, callbacks=[csv_logger])\n",
    "        self.modfc_training_history = h.history[\"loss\"]\n",
    "        self.modfc_val_history       = h.history[\"val_loss\"]\n",
    "        \n",
    "        \n",
    "\n",
    "        \n",
    "    '''===============Prediction: FC Greybox model====================='''\n",
    "    '''Description: generates predictions from trained model.'''\n",
    "    def predict_modfc(self,inputs):\n",
    "        #self.LSTMGRUGrey().predict(inputs)\n",
    "        return self.modfc.predict(inputs)\n",
    "    \n",
    "    '''===============SubModel: FC Greybox model controls====================='''\n",
    "    '''Description: outputs controls from trained model. Call after training.''' \n",
    "    '''Note: inputs needs to be [xtrain, Ujseq_real, Ujseq_imag] for a particular dataset.'''\n",
    "    def controls_modfc(self,inputs):\n",
    "        new_model = self.modfc\n",
    "        controlmod = Model(inputs=new_model.input, outputs=new_model.get_layer('controls').output)\n",
    "        return controlmod.predict(inputs)\n",
    "    \n",
    "    '''===============SubModel: FC Greybox model unitaries====================='''\n",
    "    '''Description: outputs unitary estimates (\\hat{U}_j) from trained model. Call after training.''' \n",
    "    '''Note: inputs needs to be [xtrain, Ujseq_real, Ujseq_imag] for a particular dataset.'''\n",
    "    def unitaries_modfc(self,inputs):\n",
    "        new_model = self.modfc\n",
    "        unimod = Model(inputs=new_model.input, outputs=new_model.get_layer('unitaryoutput').output)\n",
    "        return unimod.predict(inputs)\n",
    "    \n",
    "    \n",
    "    '''===============Model: GRU RNN (LSTM) Greybox model====================='''\n",
    "    def LSTMGRUGrey(self):\n",
    "        '''Description: learns controls c_i for sequence U = \\prod_i exp(-c_i \\tau_i)\n",
    "        via a GRU RNN. See paper.'''\n",
    "        \n",
    "        '''=====================Custom layer initialisation======================'''\n",
    "        origdim = (self.su2n_realdim,self.su2n_realdim)\n",
    "        \n",
    "        '''Hamiltonian estimation layer, takes as inputs controls and generators, outputs\n",
    "        estimated Hamiltonian sequence (\\hat{H}_j). In this case a Hamiltonian GRU layer\n",
    "        uses the output of the GRU c^k to construct the Hamiltonian estimates using each generator \\tau_k\n",
    "        then the actual specific \\hat{H}_j = \\sum c^k \\tau_k (imaginary unit incorporate in generators)'''\n",
    "        ham1 = ham_est_GRU(self.su_2n_delta,self.segments, name='ham_GRU_layer')\n",
    "        hsum = ham_sum_GRU(self.su_2n_delta,self.segments, name='ham_sum_GRU')\n",
    "        \n",
    "        inputUre = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim), name='GRU_input2', batch_size=None)\n",
    "        inputUim = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim), name='GRU_input3', batch_size=None)\n",
    "        original_inputs = tf.keras.Input(shape=origdim, name='GRU_flatinput')\n",
    "        \n",
    "        '''=====================GRU and Greybox layers======================'''\n",
    "        '''Note: reshape -1 is to add None dimension so input dimensions for GRU are \n",
    "        batch, time, dimensions. Here each segment of (U_j) is associated with a time step,\n",
    "        that is, the index j is effectively time t in the GRU.'''\n",
    "        y = layers.Reshape((-1,self.flatrealdim,))(original_inputs)\n",
    "        '''Input (U_T) into GRU to output estimated controls'''\n",
    "        y = layers.GRU(self.tot, return_sequences=True, activation='tanh', name='GRUcoef')(y)\n",
    "        '''Input controls into Hamiltonian estimation layer, one Hamiltonian per generator'''\n",
    "        y = ham1(y)\n",
    "        y = layers.Reshape((self.segments,self.deltadim,self.su2n_dim,self.su2n_dim))(y)\n",
    "        '''Sum all the generator Hamiltonians to obtain \\hat{H}_j = \\sum c^k \\tau_k and output\n",
    "        a sequence (\\hat{H}_j)'''\n",
    "        y = hsum(y)\n",
    "        '''Quantum evolution layer to generate unitary sequence (\\hat{U}_j) from \\hat{H}_j'''\n",
    "        y = unitary_GRU(name='GRUunitary')(y)\n",
    "        \n",
    "        '''=====================Output layer======================'''\n",
    "        '''Input the real and imaginary part of true (U_j) and estimate (\\hat{U}_j) into the fidelity function'''\n",
    "        y = fidmulti_unitary(self.su2n_dim, name='GRU_grey_fidelity')([inputUre,inputUim,y])\n",
    "        gru_grey_mod = tf.keras.Model(inputs=[original_inputs, inputUre, inputUim], outputs=y, name='gru_grey_mod')\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
    "\n",
    "        gru_grey_mod.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())\n",
    "        \n",
    "        print(gru_grey_mod.summary())\n",
    "        return gru_grey_mod\n",
    "    \n",
    "    \n",
    "    '''===============Method: preprocess unitary targets=====================\n",
    "    Description: xtrainmethod takes as input the sequence of ground-truth unitaries\n",
    "    (U_j), multiplies them together to obtain a corresponding sequence of target unitaries (U_T).'''\n",
    "    def xtrainmethod(self):\n",
    "        evolve = lambda U_j,U: U @ U_j\n",
    "        acclist = []\n",
    "        for list1 in self.Ujarray:\n",
    "            acclist.append(list(accumulate(list1, evolve))[-1])\n",
    "        xtrain = tf.convert_to_tensor(np.asarray(acclist))\n",
    "        return xtrain\n",
    "\n",
    "    \n",
    "    '''===============Training: GRU RNN (LSTM) model====================='''\n",
    "    '''Description: trains GRU RNN (LSTM) model using in-class data.\n",
    "    For other data, the GRU RNN (LSTM) model method or attribute can be called and \n",
    "    training of it can be done manually.'''\n",
    "    def train_model_LSTMGRUGrey(self, epochs):\n",
    "        csv_logger = CSVLogger('train_GRUgrey_su{0}_ntrain{1}_nseg{2}_hscale{3}_su2gen{4}.csv'\n",
    "                       .format(self.su2n_dim,self.ntraining, self.segments, self.h, self.su2ngen))\n",
    "        '''Note: can modify batch_size as need be'''\n",
    "        h = self.mod_GRU_grey.fit([self.xtrain,self.Ujseq_real,self.Ujseq_imag],self.fidones_seg, epochs=epochs, verbose=2, batch_size=None, validation_split = 0.25, shuffle=True, callbacks=[csv_logger])\n",
    "        self.GRU_training_history = h.history[\"loss\"]\n",
    "        self.GRU_val_history = h.history[\"val_loss\"]\n",
    "        \n",
    "    '''===============Prediction: GRU RNN (LSTM) model====================='''\n",
    "    '''Description: outputs controls for predictions.'''\n",
    "    '''Note: inputs needs to be [xtrain, Ujseq_real, Ujseq_imag] for a particular dataset.'''\n",
    "    def predict_LSTMGRUGrey(self,inputs):\n",
    "        #self.LSTMGRUGrey().predict(inputs)\n",
    "        return self.mod_GRU_grey.predict(inputs)\n",
    "        \n",
    "    '''===============SubModel: GRU RNN (LSTM) model controls====================='''\n",
    "    '''Description: outputs controls from trained model. Call after training.''' \n",
    "    '''Note: inputs needs to be [xtrain, Ujseq_real, Ujseq_imag] for a particular dataset.'''\n",
    "    def controls_LSTMGRUGrey(self,inputs):\n",
    "        new_model = self.mod_GRU_grey\n",
    "        controlmod = Model(inputs=new_model.input, outputs=new_model.get_layer('GRUcoef').output)\n",
    "        return controlmod.predict(inputs)\n",
    "    \n",
    "    '''===============SubModel: GRU RNN (LSTM) model unitaries====================='''\n",
    "    '''Description: outputs unitary estimates (\\hat{U}_j) from trained model. Call after training.''' \n",
    "    '''Note: inputs needs to be [xtrain, Ujseq_real, Ujseq_imag] for a particular dataset.'''\n",
    "    def unitaries_GRU(self,inputs):\n",
    "        new_model = self.mod_GRU_grey\n",
    "        unimod = Model(inputs=new_model.input, outputs=new_model.get_layer('GRUunitary').output)\n",
    "        return unimod.predict(inputs)\n",
    "    \n",
    "    \n",
    "    '''===============Model: Simple GRU model====================='''\n",
    "    '''Note: not used in paper but was one of the models trialled. May be of use for researchers.'''\n",
    "    def LSTMsimple(self,inputs):\n",
    "        '''Description: generates estimated sequence of subunitaries (U_j) from\n",
    "        input U_T using only variable tensors (not generators) using GRU LSTM.\n",
    "        Aim is global decomposition of U_T into subunitaries. It in effect is a greybox version\n",
    "        of a GRU RNN global decomposition algorithm (finding (U_j) from (U_T)) as distinct\n",
    "        from the main GRU RNN (LSTM) greybox model above which uses the GRU layer to estimate\n",
    "        the control pulses as part of the greybox approach.'''\n",
    "        origdim = (self.segments,self.su2n_realdim,self.su2n_realdim)\n",
    "        totalflat = self.segments*self.su2n_dim*self.su2n_dim\n",
    "        original_inputs = tf.keras.Input(shape=origdim, name='Ujseqinput', batch_size=None)\n",
    "        y1 = layers.Reshape((-1,self.segments * self.su2n_realdim * self.su2n_realdim,))(original_inputs)\n",
    "        GRUreal = layers.GRU(totalflat, return_sequences=True, activation='tanh', name='GRUreal')(y1)\n",
    "        GRUimag = layers.GRU(totalflat, return_sequences=True, activation='tanh', name='GRUimag')(y1)\n",
    "        Reshapereal = layers.Reshape((self.segments,self.su2n_dim,self.su2n_dim,))(GRUreal)\n",
    "        Reshapeimag = layers.Reshape((self.segments,self.su2n_dim,self.su2n_dim,))(GRUimag)\n",
    "        inputUre = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim,), name='input2', batch_size=None)\n",
    "        inputUim = layers.Input(shape=(self.segments,self.su2n_dim,self.su2n_dim,), name='input3', batch_size=None)\n",
    "        y = recombcomplex(name='unitaryoutput')([Reshapereal,Reshapeimag])\n",
    "        y_out = fidmulti_unitary(gentest._su2n_dim, name='finalfidelity')([inputUre,inputUim,y])\n",
    "        LSTMsimpmod = Model(inputs=[original_inputs, inputUre, inputUim], outputs=y_out, name='LSTMsimple')\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
    "        MSE = tf.keras.losses.MeanSquaredError()\n",
    "        LSTMsimpmod.compile(optimizer, loss=MSE)\n",
    "        print(LSTMsimpmod.summary())\n",
    "        \n",
    "        return LSTMsimpmod\n",
    "    \n",
    "    \n",
    "    '''===============Method: flatten unitary=====================\n",
    "    Description: Method takes target unitary in form of numpy arrayas input, flattens and also calculates \n",
    "        real and imaginary parts. Output is a 3-tuple of flattened (realised) unitary, \n",
    "        real part of operator and imaginary part of unitary.'''\n",
    "    def flatun(self, unitary):\n",
    "        Ujm_re = np.real(unitary)\n",
    "        Ujm_im = np.imag(unitary)\n",
    "        Ujm_top = np.concatenate((Ujm_re, -Ujm_im), axis=1)\n",
    "        Ujm_bottom = np.concatenate((Ujm_im, Ujm_re), axis=1)\n",
    "        Ujm = np.vstack((Ujm_top, Ujm_bottom)).astype('float64')\n",
    "        Ujm = Ujm.flatten()\n",
    "        Ujm = np.expand_dims(Ujm,0)\n",
    "        Ujm_re = np.expand_dims(Ujm_re,0)\n",
    "        Ujm_im = np.expand_dims(Ujm_im,0)\n",
    "        return [Ujm,Ujm_re,Ujm_im]\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Worked example\n",
    "\n",
    "In the example below, we demonstrate how to load saved datasets (see above and simulation.py for a description of dataset generation), training and generalisation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load datasets\n",
    "\n",
    "Datasets are saved as pickle files. They are generated via simulation.py which provides code that simulates methods of generating discretised approximations to normal subRiemannian geodesics. \n",
    "\n",
    "The simulation.py code creates a class with attributes and methods specified by certain hyperparameters (such as n in SU$(2^n)$). That class data is saved using pickle.\n",
    "\n",
    "To load pickle files in a way that enables access to attributes and methods as if the data were generated directly by the simulation.py class, we need to convert the pickle input using the\n",
    "namedtuple method below. \n",
    "\n",
    "The example below loads a dataset with the following key parameters:\n",
    "* Lie group SU$(2)$\n",
    "* segments: 10\n",
    "* training examples: 1000\n",
    "* $\\Lambda_0$ is drawn from $\\frak{su}(2)$\n",
    "* $h = 0.1$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''========================Example: load saved data============================'''\n",
    "\n",
    "'''Note: here need to specify su2dim = 2^n, not just n'''\n",
    "su2dim = 2\n",
    "nseg = 10\n",
    "ntrain = 1000\n",
    "hscale = 0.1 #\\delta tj (i.e. h in code or duration of pulse)\n",
    "su2ngen = 1 #0: \\Lambda_0 in \\Delta; 1: \\Lambda_0 \\in su2n\n",
    "holo = None #whether training data to be generated from Boozer\n",
    "\n",
    "with open('su{0}_ntrain{1}_nseg{2}_hscale{3}_su2gen{4}.pickle'\n",
    "          .format(su2dim,ntrain, nseg, hscale, su2ngen), 'rb') as handle:\n",
    "    gentest = pickle.load(handle)\n",
    "\n",
    "'''Convert dictionary to tuple so behaves like an object once loaded'''\n",
    "from collections import namedtuple\n",
    "'''Note: need to save as namedtuple to recreate class attributes/methods for input into\n",
    "code if loading from pickle'''\n",
    "gentest = namedtuple('Struct', gentest.keys())(*gentest.values())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialise QMLcontrol model\n",
    "\n",
    "Here we initialise the QMLcontrolmodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"mod1\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "flatinput (InputLayer)          [(None, 4, 4)]       0                                            \n",
      "__________________________________________________________________________________________________\n",
      "flatten_20 (Flatten)            (None, 16)           0           flatinput[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_140 (Dense)               (None, 640)          10880       flatten_20[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_120 (Dropout)           (None, 640)          0           dense_140[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_141 (Dense)               (None, 640)          410240      dropout_120[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_121 (Dropout)           (None, 640)          0           dense_141[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_142 (Dense)               (None, 640)          410240      dropout_121[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_122 (Dropout)           (None, 640)          0           dense_142[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_143 (Dense)               (None, 640)          410240      dropout_122[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_123 (Dropout)           (None, 640)          0           dense_143[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_144 (Dense)               (None, 640)          410240      dropout_123[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_124 (Dropout)           (None, 640)          0           dense_144[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_145 (Dense)               (None, 640)          410240      dropout_124[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_125 (Dropout)           (None, 640)          0           dense_145[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_146 (Dense)               (None, 640)          410240      dropout_125[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "controls (Dense)                (None, 2)            1282        dense_146[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "controls2 (Dense)               (None, 2)            1282        dense_146[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "ham_layer (ham_est)             (None, 2, 2)         0           controls[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "ham_layer_neg (ham_est_neg)     (None, 2, 2)         0           controls2[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "hamcomb1 (hamcomb)              (None, 2, 2)         0           ham_layer[0][0]                  \n",
      "                                                                 ham_layer_neg[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "input2 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input3 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "Ujseq_est (identitylayer)       [(None, 10, 2, 2), ( 0           hamcomb1[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "finalfidelity (fidmulti_unitary (None, 10)           0           input2[0][0]                     \n",
      "                                                                 input3[0][0]                     \n",
      "                                                                 Ujseq_est[0][0]                  \n",
      "==================================================================================================\n",
      "Total params: 2,474,884\n",
      "Trainable params: 2,474,884\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "Model: \"gru_grey_mod\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "GRU_flatinput (InputLayer)      [(None, 4, 4)]       0                                            \n",
      "__________________________________________________________________________________________________\n",
      "reshape_20 (Reshape)            (None, None, 16)     0           GRU_flatinput[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "GRUcoef (GRU)                   (None, None, 20)     2280        reshape_20[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "ham_GRU_layer (ham_est_GRU)     (None, None, 20, 2,  0           GRUcoef[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "reshape_21 (Reshape)            (None, 10, 2, 2, 2)  0           ham_GRU_layer[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "ham_sum_GRU (ham_sum_GRU)       (None, 10, 2, 2)     0           reshape_21[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "GRU_input2 (InputLayer)         [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "GRU_input3 (InputLayer)         [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "GRUunitary (unitary_GRU)        (None, 10, 2, 2)     0           ham_sum_GRU[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "GRU_grey_fidelity (fidmulti_uni (None, 10)           0           GRU_input2[0][0]                 \n",
      "                                                                 GRU_input3[0][0]                 \n",
      "                                                                 GRUunitary[0][0]                 \n",
      "==================================================================================================\n",
      "Total params: 2,280\n",
      "Trainable params: 2,280\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "None\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"modfc\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "flatinput (InputLayer)          [(None, 4, 4)]       0                                            \n",
      "__________________________________________________________________________________________________\n",
      "flatten_21 (Flatten)            (None, 16)           0           flatinput[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_147 (Dense)               (None, 640)          10880       flatten_21[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_126 (Dropout)           (None, 640)          0           dense_147[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_148 (Dense)               (None, 640)          410240      dropout_126[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_127 (Dropout)           (None, 640)          0           dense_148[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_149 (Dense)               (None, 640)          410240      dropout_127[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_128 (Dropout)           (None, 640)          0           dense_149[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_150 (Dense)               (None, 640)          410240      dropout_128[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_129 (Dropout)           (None, 640)          0           dense_150[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_151 (Dense)               (None, 640)          410240      dropout_129[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_130 (Dropout)           (None, 640)          0           dense_151[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_152 (Dense)               (None, 640)          410240      dropout_130[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "dropout_131 (Dropout)           (None, 640)          0           dense_152[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_153 (Dense)               (None, 640)          410240      dropout_131[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "controls (Dense)                (None, 20)           12820       dense_153[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "ham_layer (ham_est_simple)      (None, 20, 2, 2)     0           controls[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "input2 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input3 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "unitaryoutput (unitary_simple)  (None, 10, 2, 2)     0           ham_layer[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "finalfidelity (fidmulti_unitary (None, 10)           0           input2[0][0]                     \n",
      "                                                                 input3[0][0]                     \n",
      "                                                                 unitaryoutput[0][0]              \n",
      "==================================================================================================\n",
      "Total params: 2,485,140\n",
      "Trainable params: 2,485,140\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "None\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Note: set the second argument, genseg, to 1'"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genmod = QMLcontrolmodel(gentest,1)\n",
    "'''Note: set the second argument, genseg, to 1'''\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train model\n",
    "\n",
    "Here we show an example of training the SubRiemannian model on the pickle data loaded above.\n",
    "\n",
    "Doing so will output data into both history attributes in the class and also CSVs via the callback in the model code (using CSVLogger). Note: verbose is set to 2 in the class above, this can be modified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "24/24 - 1s - loss: 5.9419e-04 - val_loss: 3.3978e-04\n",
      "Epoch 2/2\n",
      "24/24 - 0s - loss: 2.4539e-04 - val_loss: 1.6527e-04\n"
     ]
    }
   ],
   "source": [
    "'''Select epochs - generally between 200-500 is sufficient'''\n",
    "epochs = 2\n",
    "\n",
    "'''Select model (comment out unused models). Note that:\n",
    "- mod1 is the SubRiemannian modeltribute as an attribute\n",
    "- mod_GRU_grey is the GRU RNN Greybox model as an attribute\n",
    "- modfc is the FC Greybox model as an attribute'''\n",
    "#genmod.train_model_mod1(epochs)\n",
    "genmod.train_model_LSTMGRUGrey(epochs)\n",
    "# genmod.train_model_modfc(epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract controls\n",
    "\n",
    "To extract controls, we call the control submodels for the respective model we have selected.\n",
    "This should be done <i>after</i> training the model in question. The controls are those that can be applied to the generators to recreate $(\\hat{U}_j)$. \n",
    "\n",
    "In particular, the GenData class will contain an attribute genmod.su_2n_delta which is the distribution (control subalgebra) in question. The order of the distribution is fixed e.g. for SU(2) it may be X,Y.\n",
    "\n",
    "The controls output a single sequence. So if there are two generators and ten segements, then the control output will be a vector of 20 elements. The first two are the controls for generators X,Y respectively for generating the estimated Hamiltonian $\\hat{H}_1$, the third and fourth elements are the two controls for X and Y to generate $\\hat{H}_2$ and so on. \n",
    "\n",
    "The estimated unitaries $\\hat{U)_j$ are then obtained via application of the time-independent approximation to Schrodinger's equation i.e:\n",
    "\n",
    "$$\\hat{U}_j = \\exp(-h \\hat{H}_j)$$\n",
    "\n",
    "where the imaginary unit $i$ is already included in $\\hat{H}_j$ because it is incorporated into the distribution (see paper)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1000, 1, 20)"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'''Controls - '''\n",
    "genmod.controls_LSTMGRUGrey([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "# genmod.controls_modfc([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "#genmod.controls_mod1([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract unitaries\n",
    "\n",
    "To extract unitary estimates $(\\hat{U}_j)$, we call the unitary submodels for the respective model we have selected.\n",
    "This should be done <i>after</i> training the model in question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[[ 9.99994023e-01-3.45731721e-03j,\n",
       "          -9.53507068e-06-3.29659610e-08j],\n",
       "         [ 9.53507068e-06-3.29659610e-08j,\n",
       "           9.99994023e-01+3.45731721e-03j]],\n",
       "\n",
       "        [[ 9.99899461e-01-1.37292660e-02j,\n",
       "          -3.54587342e-03-4.86871345e-05j],\n",
       "         [ 3.54587342e-03-4.86871345e-05j,\n",
       "           9.99899461e-01+1.37292660e-02j]],\n",
       "\n",
       "        [[ 9.99955699e-01+7.66408257e-03j,\n",
       "           5.46435570e-03-4.18811286e-05j],\n",
       "         [-5.46435570e-03-4.18811286e-05j,\n",
       "           9.99955699e-01-7.66408257e-03j]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 9.99975149e-01+2.56536523e-03j,\n",
       "           6.56654586e-03-1.68460070e-05j],\n",
       "         [-6.56654586e-03-1.68460070e-05j,\n",
       "           9.99975149e-01-2.56536523e-03j]],\n",
       "\n",
       "        [[ 9.99972219e-01+6.19963551e-03j,\n",
       "           4.13829586e-03-2.56566387e-05j],\n",
       "         [-4.13829586e-03-2.56566387e-05j,\n",
       "           9.99972219e-01-6.19963551e-03j]],\n",
       "\n",
       "        [[ 9.99992442e-01-2.25826381e-03j,\n",
       "          -3.16480075e-03-7.14700901e-06j],\n",
       "         [ 3.16480075e-03-7.14700901e-06j,\n",
       "           9.99992442e-01+2.25826381e-03j]]],\n",
       "\n",
       "\n",
       "       [[[ 9.99988037e-01-4.82430128e-03j,\n",
       "          -8.07394401e-04-3.89516044e-06j],\n",
       "         [ 8.07394401e-04-3.89516044e-06j,\n",
       "           9.99988037e-01+4.82430128e-03j]],\n",
       "\n",
       "        [[ 9.99949990e-01-8.31504674e-03j,\n",
       "          -5.55659037e-03-4.62056193e-05j],\n",
       "         [ 5.55659037e-03-4.62056193e-05j,\n",
       "           9.99949990e-01+8.31504674e-03j]],\n",
       "\n",
       "        [[ 9.99913141e-01+6.57462825e-03j,\n",
       "           1.14227678e-02-7.51069758e-05j],\n",
       "         [-1.14227678e-02-7.51069758e-05j,\n",
       "           9.99913141e-01-6.57462825e-03j]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 9.99981162e-01+5.95488841e-03j,\n",
       "           1.48837897e-03-8.86329766e-06j],\n",
       "         [-1.48837897e-03-8.86329766e-06j,\n",
       "           9.99981162e-01-5.95488841e-03j]],\n",
       "\n",
       "        [[ 9.99983472e-01+5.67477686e-03j,\n",
       "           9.23299305e-04-5.23960413e-06j],\n",
       "         [-9.23299305e-04-5.23960413e-06j,\n",
       "           9.99983472e-01-5.67477686e-03j]],\n",
       "\n",
       "        [[ 9.99993591e-01-3.37986421e-03j,\n",
       "          -1.18104390e-03-3.99179361e-06j],\n",
       "         [ 1.18104390e-03-3.99179361e-06j,\n",
       "           9.99993591e-01+3.37986421e-03j]]],\n",
       "\n",
       "\n",
       "       [[[ 9.99982175e-01-4.83935803e-03j,\n",
       "           3.49700502e-03+1.69235610e-05j],\n",
       "         [-3.49700502e-03+1.69235610e-05j,\n",
       "           9.99982175e-01+4.83935803e-03j]],\n",
       "\n",
       "        [[ 9.99929680e-01-1.07487765e-02j,\n",
       "          -5.00956433e-03-5.38504742e-05j],\n",
       "         [ 5.00956433e-03-5.38504742e-05j,\n",
       "           9.99929680e-01+1.07487765e-02j]],\n",
       "\n",
       "        [[ 9.99908706e-01+3.28239019e-03j,\n",
       "           1.31074215e-02-4.30275998e-05j],\n",
       "         [-1.31074215e-02-4.30275998e-05j,\n",
       "           9.99908706e-01-3.28239019e-03j]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 9.99956191e-01+9.29002470e-03j,\n",
       "           1.14490257e-03-1.06366392e-05j],\n",
       "         [-1.14490257e-03-1.06366392e-05j,\n",
       "           9.99956191e-01-9.29002470e-03j]],\n",
       "\n",
       "        [[ 9.99985379e-01+5.19053391e-03j,\n",
       "          -1.51673359e-03+7.87277226e-06j],\n",
       "         [ 1.51673359e-03+7.87277226e-06j,\n",
       "           9.99985379e-01-5.19053391e-03j]],\n",
       "\n",
       "        [[ 9.99975106e-01-2.43927710e-03j,\n",
       "          -6.62102599e-03-1.61509191e-05j],\n",
       "         [ 6.62102599e-03-1.61509191e-05j,\n",
       "           9.99975106e-01+2.43927710e-03j]]],\n",
       "\n",
       "\n",
       "       ...,\n",
       "\n",
       "\n",
       "       [[[ 9.99988377e-01-3.95960787e-03j,\n",
       "           2.75071077e-03+1.08918626e-05j],\n",
       "         [-2.75071077e-03+1.08918626e-05j,\n",
       "           9.99988377e-01+3.95960787e-03j]],\n",
       "\n",
       "        [[ 9.99943002e-01-2.03678464e-03j,\n",
       "          -1.04806059e-02-2.13479538e-05j],\n",
       "         [ 1.04806059e-02-2.13479538e-05j,\n",
       "           9.99943002e-01+2.03678464e-03j]],\n",
       "\n",
       "        [[ 9.99917203e-01+7.14635721e-03j,\n",
       "           1.07009662e-02-7.64792592e-05j],\n",
       "         [-1.07009662e-02-7.64792592e-05j,\n",
       "           9.99917203e-01-7.14635721e-03j]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 9.99999765e-01+6.83664634e-04j,\n",
       "           4.28346882e-05-2.92845683e-08j],\n",
       "         [-4.28346882e-05-2.92845683e-08j,\n",
       "           9.99999765e-01-6.83664634e-04j]],\n",
       "\n",
       "        [[ 9.99993496e-01-1.44770588e-03j,\n",
       "          -3.30339828e-03-4.78238022e-06j],\n",
       "         [ 3.30339828e-03-4.78238022e-06j,\n",
       "           9.99993496e-01+1.44770588e-03j]],\n",
       "\n",
       "        [[ 9.99892799e-01-1.44482088e-02j,\n",
       "          -2.37465597e-03-3.43132038e-05j],\n",
       "         [ 2.37465597e-03-3.43132038e-05j,\n",
       "           9.99892799e-01+1.44482088e-02j]]],\n",
       "\n",
       "\n",
       "       [[[ 9.99958137e-01-7.96089857e-04j,\n",
       "           9.11542419e-03+7.25700054e-06j],\n",
       "         [-9.11542419e-03+7.25700054e-06j,\n",
       "           9.99958137e-01+7.96089857e-04j]],\n",
       "\n",
       "        [[ 9.99948261e-01-9.46372745e-03j,\n",
       "          -3.72984115e-03-3.53000265e-05j],\n",
       "         [ 3.72984115e-03-3.53000265e-05j,\n",
       "           9.99948261e-01+9.46372745e-03j]],\n",
       "\n",
       "        [[ 9.99910076e-01+5.55284161e-03j,\n",
       "           1.22066176e-02-6.77875097e-05j],\n",
       "         [-1.22066176e-02-6.77875097e-05j,\n",
       "           9.99910076e-01-5.55284161e-03j]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 9.99829313e-01+1.18333267e-02j,\n",
       "           1.41876369e-02-1.67915603e-04j],\n",
       "         [-1.41876369e-02-1.67915603e-04j,\n",
       "           9.99829313e-01-1.18333267e-02j]],\n",
       "\n",
       "        [[ 9.99986174e-01+3.64114460e-03j,\n",
       "          -3.79398685e-03+1.38146458e-05j],\n",
       "         [ 3.79398685e-03+1.38146458e-05j,\n",
       "           9.99986174e-01-3.64114460e-03j]],\n",
       "\n",
       "        [[ 9.99991855e-01-2.12742198e-03j,\n",
       "          -3.42984402e-03-7.29678500e-06j],\n",
       "         [ 3.42984402e-03-7.29678500e-06j,\n",
       "           9.99991855e-01+2.12742198e-03j]]],\n",
       "\n",
       "\n",
       "       [[[ 9.99990852e-01-4.24081602e-03j,\n",
       "          -5.58839561e-04-2.36995744e-06j],\n",
       "         [ 5.58839561e-04-2.36995744e-06j,\n",
       "           9.99990852e-01+4.24081602e-03j]],\n",
       "\n",
       "        [[ 9.99930485e-01-1.08204527e-02j,\n",
       "          -4.68404174e-03-5.06869754e-05j],\n",
       "         [ 4.68404174e-03-5.06869754e-05j,\n",
       "           9.99930485e-01+1.08204527e-02j]],\n",
       "\n",
       "        [[ 9.99944553e-01+7.03170179e-03j,\n",
       "           7.83856558e-03-5.51215120e-05j],\n",
       "         [-7.83856558e-03-5.51215120e-05j,\n",
       "           9.99944553e-01-7.03170179e-03j]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 9.99976879e-01+6.17630087e-03j,\n",
       "           2.84515216e-03-1.75729220e-05j],\n",
       "         [-2.84515216e-03-1.75729220e-05j,\n",
       "           9.99976879e-01-6.17630087e-03j]],\n",
       "\n",
       "        [[ 9.99962676e-01+8.31712595e-03j,\n",
       "           2.33904368e-03-1.94548470e-05j],\n",
       "         [-2.33904368e-03-1.94548470e-05j,\n",
       "           9.99962676e-01-8.31712595e-03j]],\n",
       "\n",
       "        [[ 9.99995622e-01-2.17739082e-03j,\n",
       "          -2.00373658e-03-4.36293673e-06j],\n",
       "         [ 2.00373658e-03-4.36293673e-06j,\n",
       "           9.99995622e-01+2.17739082e-03j]]]])"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genmod.unitaries_mod1([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "# genmod.unitaries_modfc([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "# genmod.unitaries_GRU([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: SubR\\assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: SubR\\assets\n"
     ]
    }
   ],
   "source": [
    "'''Save model'''\n",
    "genmod.mod1.save('SubR')\n",
    "'''Generate predictions to output a list of fidelities for each U_j for each training example.'''\n",
    "mod_pred =  genmod.mod1.predict([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "'''Save the predictions'''\n",
    "np.savetxt(\"SubR_fidelity.csv\", mod_pred, delimiter=\",\")\n",
    "print('Save completed.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Access intermediate layers of model\n",
    "\n",
    "The code below demonstrates how to construct submodels to access intermediate outputs of the model. For controls and unitary sequences, this is hard-coded into the QMLcontrolmodel class (see submodel sections).\n",
    "\n",
    "This is necessary because the sequences of control pulses or estimated Hamiltonians or unitaries are themselves only intermediate outputs of the models.\n",
    "\n",
    "In the case of the SubRiemannian model below, the layer projev outputs a list of length two:\n",
    "* the first entry is an array (tensor) of the sequence of estimated unitaries\n",
    "* the second is a list of the controls generated via the trace of the projection function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''===========SubRiemannian: Load coefficient and unitary sub-models=============='''\n",
    "'''Note: is for main model mod1 in QML model class'''\n",
    "'''Load a pretrained model'''\n",
    "new_model = tf.keras.models.load_model('SubR')\n",
    "\n",
    "'''===========Unitary submodel: extracts output from unitary output layer==========='''\n",
    "'''Description: create a submodel which takes the same inputs but outputs the layer which\n",
    "undertakes the quantum evolution of Hamiltonians i.e. the layer that outputs the estimated unitaries.\n",
    "Here use output index (0)'''\n",
    "unimod = Model(inputs=new_model.input, outputs=new_model.get_layer('Ujseq_est').output[0])\n",
    "unimodpred = unimod.predict([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "unimodpred\n",
    "\n",
    "#np.savetxt(\"SubRoutput.csv\", unimodpred, delimiter=\",\")\n",
    "\n",
    "'''===========Control pulse submodel: extracts control pulses==========='''\n",
    "'''Description: extracts controls that are generated in the SubRiemannian model (index 1)'''\n",
    "controlmod = Model(inputs=new_model.input, outputs=new_model.get_layer('Ujseq_est').output[1])\n",
    "controlpred = controlmod.predict([genmod.xtrain, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "#np.savetxt(\"SubRcontrols.csv\", controlpred, delimiter=\",\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing/generalisation against estimates\n",
    "\n",
    "Below we provide an example of different methods of generalisation of the models by using the models to predict using out-of-sample datasets. There are a variety of ways of doing this, including by reinitialising the GenData geodesic generating code and inputting the unitary targets etc from that new class of data, or by randomly generating unitary targets $U_T$, generating sequences $(\\hat{U}_j)$ and then comparing how well $\\hat{U}_T = \\hat{U}_n...\\hat{U}_0$ approximates $U_T$, or by generating $U_T$ as rotations about the $z$-axis by a random angle, $\\theta$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"mod1\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "flatinput (InputLayer)          [(None, 4, 4)]       0                                            \n",
      "__________________________________________________________________________________________________\n",
      "flatten_6 (Flatten)             (None, 16)           0           flatinput[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_42 (Dense)                (None, 640)          10880       flatten_6[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dropout_36 (Dropout)            (None, 640)          0           dense_42[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_43 (Dense)                (None, 640)          410240      dropout_36[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_37 (Dropout)            (None, 640)          0           dense_43[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_44 (Dense)                (None, 640)          410240      dropout_37[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_38 (Dropout)            (None, 640)          0           dense_44[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_45 (Dense)                (None, 640)          410240      dropout_38[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_39 (Dropout)            (None, 640)          0           dense_45[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_46 (Dense)                (None, 640)          410240      dropout_39[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_40 (Dropout)            (None, 640)          0           dense_46[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_47 (Dense)                (None, 640)          410240      dropout_40[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_41 (Dropout)            (None, 640)          0           dense_47[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_48 (Dense)                (None, 640)          410240      dropout_41[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "controls (Dense)                (None, 2)            1282        dense_48[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "controls2 (Dense)               (None, 2)            1282        dense_48[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "ham_layer (ham_est)             (None, 2, 2)         0           controls[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "ham_layer_neg (ham_est_neg)     (None, 2, 2)         0           controls2[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "hamcomb1 (hamcomb)              (None, 2, 2)         0           ham_layer[0][0]                  \n",
      "                                                                 ham_layer_neg[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "input2 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input3 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "Ujseq_est (identitylayer)       [(None, 10, 2, 2), ( 0           hamcomb1[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "finalfidelity (fidmulti_unitary (None, 10)           0           input2[0][0]                     \n",
      "                                                                 input3[0][0]                     \n",
      "                                                                 Ujseq_est[0][0]                  \n",
      "==================================================================================================\n",
      "Total params: 2,474,884\n",
      "Trainable params: 2,474,884\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "Model: \"gru_grey_mod\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "GRU_flatinput (InputLayer)      [(None, 4, 4)]       0                                            \n",
      "__________________________________________________________________________________________________\n",
      "reshape_6 (Reshape)             (None, None, 16)     0           GRU_flatinput[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "GRUcoef (GRU)                   (None, None, 20)     2280        reshape_6[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "ham_GRU_layer (ham_est_GRU)     (None, None, 20, 2,  0           GRUcoef[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "reshape_7 (Reshape)             (None, 10, 2, 2, 2)  0           ham_GRU_layer[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "ham_sum_GRU (ham_sum_GRU)       (None, 10, 2, 2)     0           reshape_7[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "GRU_input2 (InputLayer)         [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "GRU_input3 (InputLayer)         [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "GRUunitary (unitary_GRU)        (None, 10, 2, 2)     0           ham_sum_GRU[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "GRU_grey_fidelity (fidmulti_uni (None, 10)           0           GRU_input2[0][0]                 \n",
      "                                                                 GRU_input3[0][0]                 \n",
      "                                                                 GRUunitary[0][0]                 \n",
      "==================================================================================================\n",
      "Total params: 2,280\n",
      "Trainable params: 2,280\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "None\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"modfc\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "flatinput (InputLayer)          [(None, 4, 4)]       0                                            \n",
      "__________________________________________________________________________________________________\n",
      "flatten_7 (Flatten)             (None, 16)           0           flatinput[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_49 (Dense)                (None, 640)          10880       flatten_7[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dropout_42 (Dropout)            (None, 640)          0           dense_49[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_50 (Dense)                (None, 640)          410240      dropout_42[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_43 (Dropout)            (None, 640)          0           dense_50[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_51 (Dense)                (None, 640)          410240      dropout_43[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_44 (Dropout)            (None, 640)          0           dense_51[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_52 (Dense)                (None, 640)          410240      dropout_44[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_45 (Dropout)            (None, 640)          0           dense_52[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_53 (Dense)                (None, 640)          410240      dropout_45[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_46 (Dropout)            (None, 640)          0           dense_53[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_54 (Dense)                (None, 640)          410240      dropout_46[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_47 (Dropout)            (None, 640)          0           dense_54[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_55 (Dense)                (None, 640)          410240      dropout_47[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "controls (Dense)                (None, 20)           12820       dense_55[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "ham_layer (ham_est_simple)      (None, 20, 2, 2)     0           controls[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "input2 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input3 (InputLayer)             [(None, 10, 2, 2)]   0                                            \n",
      "__________________________________________________________________________________________________\n",
      "unitaryoutput (unitary_simple)  (None, 10, 2, 2)     0           ham_layer[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "finalfidelity (fidmulti_unitary (None, 10)           0           input2[0][0]                     \n",
      "                                                                 input3[0][0]                     \n",
      "                                                                 unitaryoutput[0][0]              \n",
      "==================================================================================================\n",
      "Total params: 2,485,140\n",
      "Trainable params: 2,485,140\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "None\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\elija\\Anaconda3\\lib\\site-packages\\ipykernel_launcher.py:65: DeprecationWarning: Importing functions/classes of the qip submodule directly from the namespace qutip is deprecated. Please import them from the submodule instead, e.g.\n",
      "from qutip.qip.operations import cnot\n",
      "from qutip.qip.circuit import QubitCircuit\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SpearmanrResult(correlation=0.02883754083754084, pvalue=0.3623101241284554) (0.028418322332049557, 0.36933293101126446) [[1.         0.02841832]\n",
      " [0.02841832 1.        ]]\n"
     ]
    }
   ],
   "source": [
    "'''=================Generalisation==================='''\n",
    "'''Step 1: select hyperparameters for datsets (load existing or generate new ones)'''\n",
    "su2dim = 2\n",
    "nseg = 10\n",
    "ntrain = 1000\n",
    "hscale = 0.1 #\\delta tj (i.e. h in code or duration of pulse)\n",
    "su2ngen = None #0: \\Lambda_0 in \\Delta; 1: \\Lambda_0 \\in su2n\n",
    "holo = None #whether training data to be generated from Boozer\n",
    "\n",
    "'''Load pre-saved model'''\n",
    "modstring = 'SubR'\n",
    "\n",
    "'''Option: random generation of z-rotation unitaries as test set of \n",
    "target unitaries U_T. Set to 1 to do so.'''\n",
    "testz = 1 #1\n",
    "\n",
    "'''==========Load dataset========='''\n",
    "'''Note: even if random unitaries or some other test set is being used, the prediction for\n",
    "out of sample code will still require input of the ground-truth real and imaginary parts of (U_j)\n",
    "though these are not used in prediction, they must be input because of how the models are constructed.'''\n",
    "with open('su{0}_ntrain{1}_nseg{2}_hscale{3}_su2gen{4}.pickle'\n",
    "          .format(su2dim,ntrain, nseg, hscale, su2ngen), 'rb') as handle:\n",
    "    gentest = pickle.load(handle)\n",
    "\n",
    "'''Convert dictionary to tuple so behaves like an object once loaded'''\n",
    "from collections import namedtuple\n",
    "gentest = namedtuple('Struct', gentest.keys())(*gentest.values())\n",
    "\n",
    "genmod = QMLcontrolmodel(gentest,1)\n",
    "\n",
    "'''============FINAL TESTING/GENERALISATION CODE==============='''\n",
    "\n",
    "'''===========SubRiemannian: Load coefficient and unitary sub-models=============='''\n",
    "'''Note: is for main model mod1 in QML model class'''\n",
    "\n",
    "'''Load the existing model'''\n",
    "new_model = tf.keras.models.load_model(modstring)\n",
    "\n",
    "'''Create a submodel to extract unitaries'''\n",
    "unimod = Model(inputs=new_model.input, outputs=new_model.get_layer('Ujseq_est').output)\n",
    "\n",
    "'''=============Testing against random unitaries==========='''\n",
    "\n",
    "pi = np.pi\n",
    "'''=============Function: realise unitaries'''\n",
    "'''Description: converts complex unitaries into realised form for input into model'''\n",
    "def realise1(U1):\n",
    "    U1_re = U1.real\n",
    "    U1_im = U1.imag\n",
    "    '''Create realisation of U as: [[X, -Y],[Y,X]]'''\n",
    "    complex_U1_top = np.concatenate((U1_re, -U1_im), axis=1)\n",
    "    complex_U1_bottom = np.concatenate((U1_im, U1_re), axis=1)\n",
    "    '''Stack them both together'''\n",
    "    complex_U1 = np.vstack((complex_U1_top, complex_U1_bottom))\n",
    "    return complex_U1\n",
    "\n",
    "'''=============Method: looping to generate random unitaries============='''\n",
    "'''Description: generates random unitaries for test set via loop,\n",
    "here we generate random angle to rotate about z-axis by. The stack lists are to preprocess\n",
    "them into a form so that they can be input as target unitaries U_T in the models.'''\n",
    "realised_U1stack = []\n",
    "complex_U1stack = []\n",
    "'''List to save random angles'''\n",
    "anglelist = []\n",
    "for i in range(gentest.ntraining):\n",
    "    if testz == 1:\n",
    "        rnangle = np.random.uniform(-2*pi,2*pi)\n",
    "        anglelist.append(rnangle)\n",
    "        U1 = qt.rz(rnangle).data.toarray()\n",
    "    else:\n",
    "        U1 = qt.rand_unitary(gentest.su2n_dim).data.toarray()\n",
    "    \n",
    "    \n",
    "    complex_U1stack.append(U1)\n",
    "    realised_U1 = realise1(U1)\n",
    "    realised_U1stack.append(realised_U1)\n",
    "    \n",
    "realisedU1array = np.asarray(realised_U1stack)\n",
    "complexU1array = np.asarray(complex_U1stack)\n",
    "#testfid = new_model.predict([realisedU1array, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "testunitary = unimod.predict([realisedU1array, genmod.Ujseq_real,genmod.Ujseq_imag])\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "'''=============Function: xtrainmethod============='''\n",
    "'''Description: replicates xtrainmethod from QMLcontrolmodel class above, that is\n",
    "it creates U_T by multiplying each U_j in a sequence together (cumulant)'''\n",
    "from itertools import accumulate\n",
    "Ujarray = np.asarray(testunitary[0])[::-1]\n",
    "def xtrainmethod(Ujarray):\n",
    "    evolve = lambda U_j,U: U @ U_j\n",
    "    acclist = []\n",
    "    for list1 in Ujarray:\n",
    "        acclist.append(list(accumulate(list1, evolve))[-1])\n",
    "    xtrain = tf.convert_to_tensor(np.asarray(acclist))\n",
    "    return xtrain\n",
    "xtrain_est = xtrainmethod(Ujarray)\n",
    "xtrain_est.shape\n",
    "\n",
    "'''Convert to tensor for input into models'''\n",
    "complexU1tensor = tf.convert_to_tensor(complexU1array)\n",
    "\n",
    "\n",
    "'''=============Function: fidelity function============='''\n",
    "'''Description: estimates fidelities of the true U_T and the estimate \\hat{U}_T'''\n",
    "def fidelityfunction(UTtrue,UTestimate):\n",
    "    fid = tf.square(tf.abs(tf.linalg.trace(tf.matmul(UTtrue,UTestimate, adjoint_b=True))))/(UTtrue.shape[1]**2)\n",
    "    return fid\n",
    "\n",
    "'''=============Function: bulk fidelity function============='''\n",
    "'''Description: estimates fidelities of the true U_T and the estimate \\hat{U}_T but en mass\n",
    "(in bulk) to avoid loops (faster implementation using tensors rather than loops)'''\n",
    "def fidelitybulk(UTtrue,UTestimate):\n",
    "    A1 = tf.einsum('aijk,aikm->aijm', UTtrue,tf.transpose(UTestimate,perm=[0,1,3,2],conjugate=True))\n",
    "    fid = tf.square(tf.abs(tf.linalg.trace(A1)))/(UTtrue.shape[2]**2)\n",
    "    return fid\n",
    "\n",
    "'''=============Method: fidelity list============='''\n",
    "'''Description: calculates fidelity of estimated UT v true UT and inputs into list named fidlist.\n",
    "- complexU1tensor: the (true) randomly generated UT targets.\n",
    "- xtrain_est: the estimates of UT i.e. UThat generated by the applicable model.\n",
    "- fidelityfunction: comparison of fidelity F(UT,UThat). Note the estimate is second.\n",
    "'''\n",
    "fidlist = (fidelityfunction(complexU1tensor, xtrain_est))\n",
    "\n",
    "'''=============Method: fidelity comparison============='''\n",
    "'''Description: compares each random (true) UT with UTtrain (UT from the training data)\n",
    "to estimate how similar UT is to the class UTtrain. The way this code part works below\n",
    "is described in the paper. In short, it takes a particular target U_T. We have above\n",
    "the means of assessing the fidelity of U_T versus the estimate \\hat{U}_T i.e.\n",
    "F(U_T,\\hat{U}_T). But we want to assess the extent to which higher F(U_T,\\hat{U}_T) is as a result\n",
    "of the randomly generated U_T being more similar to the training set upon which the model was trained.\n",
    "To do this, we calculate the fidelity of U_T with each U_T in the training set and average it.\n",
    "This gives us an indication of how similar U_T was to the U_T from the training set. See paper for\n",
    "a discussion.'''\n",
    "\n",
    "xtrain_orig = xtrainmethod(genmod.Ujmasterarray)\n",
    "xtrainstack = tf.stack([xtrain_orig] * xtrain_orig.shape[0], axis=0)\n",
    "Ustack = tf.stack([complexU1tensor] * complexU1tensor.shape[0],axis=1)\n",
    "\n",
    "fbulk = fidelitybulk(Ustack,xtrainstack)\n",
    "\n",
    "'''Bulk average operator fidelity i.e. comparing each estimated UT \n",
    "to all UTs in xtrain (random target UTs)'''\n",
    "fbulkav = tf.math.reduce_mean(fbulk, axis=1)\n",
    "\n",
    "fidconcat = tf.transpose(tf.stack([fidlist,fbulkav], axis=0, name='concat'),perm=[1,0])\n",
    "fidarray = fidconcat.numpy()\n",
    "agmx = tf.argmax(fidconcat[:,0]).numpy()\n",
    "A = fidarray[fidarray[:,1].argsort()]\n",
    "\n",
    "'''Another copy to save but columns switched'''\n",
    "angletens = tf.cast(tf.transpose(tf.convert_to_tensor(anglelist)), dtype=tf.float64)\n",
    "\n",
    "\n",
    "'''Note: in the case of test set being randomly generated z-rotation unitaries.'''\n",
    "if testz == 1:\n",
    "    fidconcat1 = tf.transpose(tf.stack([fbulkav,fidlist,angletens], axis=0))\n",
    "    \n",
    "    fidarray1 = fidconcat1.numpy()\n",
    "    A1 = fidarray1[fidarray1[:,1].argsort()]\n",
    "    np.savetxt('modangle{0}_angle_su{1}_ntrain{2}_nseg{3}_hscale{4}_su2gen{5}_z{6}.csv'\n",
    "              .format(modstring,su2dim,ntrain, nseg, hscale, su2ngen,testz), A1, delimiter=\",\")\n",
    "\n",
    "\n",
    "fidconcat1 = tf.transpose(tf.stack([fbulkav,fidlist], axis=0, name='concat'),perm=[1,0])\n",
    "\n",
    "fidarray1 = fidconcat1.numpy()\n",
    "\n",
    "A1 = fidarray1[fidarray1[:,1].argsort()]\n",
    "\n",
    "np.savetxt('mod{0}_su{1}_ntrain{2}_nseg{3}_hscale{4}_su2gen{5}_z{6}.csv'\n",
    "              .format(modstring,su2dim,ntrain, nseg, hscale, su2ngen,testz), A1, delimiter=\",\")\n",
    "\n",
    "\n",
    "'''Correlation analysis to determine how significant the similarity of U_T to those in the test set is.'''\n",
    "from scipy.stats.stats import pearsonr, spearmanr  \n",
    "a = A[:,0]\n",
    "b = A[:,1]   \n",
    "spearman = spearmanr(a,b)\n",
    "pearson = pearsonr(a,b)\n",
    "corrmx = np.corrcoef(a,b)\n",
    "print(spearman,pearson,corrmx)\n",
    "np.savetxt('modcorr{0}_su{1}_ntrain{2}_nseg{3}_hscale{4}_su2gen{5}.csv'\n",
    "              .format(modstring,su2dim,ntrain, nseg, hscale, su2ngen), [spearman,pearson], delimiter=\",\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}