Rydberg_Reconstruction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First we import the libraries\n",
    "# Please make sure you have the required libraries.\n",
    "# You can use pip or conda to get them.\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from qucumber.nn_states import PositiveWaveFunction\n",
    "from qucumber.callbacks import MetricEvaluator\n",
    "\n",
    "import qucumber.utils.training_statistics as ts\n",
    "import qucumber.utils.data as data\n",
    "import qucumber\n",
    "\n",
    "import torch\n",
    "\n",
    "# set random seed on cpu but not gpu, as gpu is not used\n",
    "qucumber.set_random_seed(1234, cpu=True, gpu=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10000, 8])"
      ]
     },
     "execution_count": 183,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Now we get the training data\n",
    "# Training data is given in the data/nY=8/ directory for each delta\n",
    "# in the '...samples.csv' files.\n",
    "# So the delta needs to be specified below\n",
    "\n",
    "#-------------------- Specify delta and nY ---------------------#\n",
    "delta = \"1.14\" \n",
    "nY = \"8\" # nY refers to the number of atoms in the array\n",
    "# For the nY = 8 system case:\n",
    "# delta is any of 1.00, 1.02, 1.04, 1.06, 1.08, 1.10, 1.12, 1.14, 1.16, 1.18, 1.20\n",
    "# For nY > 8:\n",
    "# delta is any of 1.00, 1.04, 1.08, 1.12, 1.16, 1.20, 1.28\n",
    "#---------------------------------------------------------------#\n",
    "\n",
    "train_path = \"data/nY=\"+nY+\"/δ=\"+delta+\"_samples.csv\"\n",
    "train_data = data.load_data(train_path)\n",
    "\n",
    "# The training data is stored in train_data[0]\n",
    "# and the dimension of the data is:\n",
    "train_data[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [],
   "source": [
    "# So now we have the training data. Our goal is to train a\n",
    "# Restricted Boltzman Machine using this training data. We use\n",
    "# QuCumber to create an instance of an RBM. First we need to specify the \n",
    "# number of visible nodes (nv) and number of higgen nodes (nh). \n",
    "# As we have an array of 8 atoms, we have 8 inputs and so nv = 8\n",
    "nv = train_data[0].shape[-1]\n",
    "\n",
    "# Number of hidden nodes of an RBM is a hyperparameter which depends on the \n",
    "# data we are using and which can be varied to get optimal result. \n",
    "# For delta=1.14, nh was varied from 4 to 12 and nh = 12 gave quite good results.\n",
    "# Hence, we set nh = 12.\n",
    "# Note that for different delta, nh=12 may not give good results.\n",
    "nh = 12\n",
    "\n",
    "# Finally we create an RBM with nv visible nodes and nh  hidden nodes.\n",
    "nn_state = PositiveWaveFunction(num_visible=nv, num_hidden=nh, gpu=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total time elapsed during training: 121.132 s\n"
     ]
    }
   ],
   "source": [
    "# Below we have more hyperparameters of the RBM model. Like the number of \n",
    "# hidden nodes, the parameters below can be varied to get optimal results.\n",
    "# The following hyperparameters seemed to work quite well. Note that when\n",
    "# you have different data, then different values of the hyperparameter will give\n",
    "# better solutions. \n",
    "# Further description of the parameters (if you are curious) is in:\n",
    "# https://qucumber.readthedocs.io/en/stable/quantum_states.html?highlight=fit#qucumber.nn_states.PositiveWaveFunction.fit\n",
    "\n",
    "epochs = 500\n",
    "pbs = 100\n",
    "nbs = pbs\n",
    "lr = 0.001 # For nY > 8 systems lr = 0.0065 can be used as it works quite well.\n",
    "k = 10\n",
    "\n",
    "# Now we train our RBM using the above parameters. \n",
    "# Note that the training process will take 2 - 3 minutes\n",
    "\n",
    "nn_state.fit(\n",
    "    train_data[0],\n",
    "    epochs=epochs,\n",
    "    pos_batch_size=pbs,\n",
    "    neg_batch_size=nbs,\n",
    "    lr=lr,\n",
    "    k=k,\n",
    "    time=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'rbm_am': OrderedDict([('weights',\n",
       "               tensor([[ 0.0436,  0.1365, -0.1789, -0.2058,  0.0070, -0.2799,  0.0170, -0.0827],\n",
       "                       [ 0.5144, -0.4623,  0.2377,  0.0121, -0.3839, -0.3526,  0.2013, -0.4161],\n",
       "                       [ 0.1809, -1.1153,  0.9791, -1.1769,  1.1068, -0.9989,  0.5834, -0.6266],\n",
       "                       [-0.9534,  0.8799, -0.9635,  0.8807, -0.9633,  0.5620, -0.7005,  0.7425],\n",
       "                       [-0.7777,  0.6887, -0.5053,  0.4842, -0.8278,  0.4021, -1.1933,  0.9725],\n",
       "                       [-0.1604,  0.1566, -0.7861,  0.0782,  0.2232, -0.3297,  0.3641, -0.7745],\n",
       "                       [-0.9716,  0.7356, -1.1741,  0.6283, -1.0942,  1.0014, -0.7571,  0.4562],\n",
       "                       [-0.1474, -0.2230, -0.0836,  0.6210, -0.9229,  0.6951, -0.2905, -0.1504],\n",
       "                       [ 0.4037, -0.9583,  0.5166, -0.3403, -0.0837, -0.2000,  0.4310, -0.5927],\n",
       "                       [ 0.0930, -0.5438,  0.2563, -0.5043,  0.3958, -0.6337,  0.3522, -0.3941],\n",
       "                       [ 1.1114, -0.9445,  0.7091, -1.1020,  1.1863, -1.0498,  0.5771, -1.1580],\n",
       "                       [-0.8273,  0.2133, -0.6560,  0.3653, -0.2026,  0.4236, -1.1612,  0.9219]],\n",
       "                      dtype=torch.float64)),\n",
       "              ('visible_bias',\n",
       "               tensor([-0.1893, -0.3280, -0.0948, -0.7597, -0.2868, -0.4847, -0.2894, -0.3721],\n",
       "                      dtype=torch.float64)),\n",
       "              ('hidden_bias',\n",
       "               tensor([ 0.0263,  0.0339,  0.1548, -0.1065,  0.0265,  0.0317,  0.1420, -0.0583,\n",
       "                        0.0944,  0.1240, -0.0400,  0.0784], dtype=torch.float64))])}"
      ]
     },
     "execution_count": 186,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# After training is complete, we save the parameters of the trained RBM below\n",
    "# as rydberg_data.pt in the output directory.\n",
    "nn_state.save(\"output/rydberg_data.pt\")\n",
    "torch.load(\"output/rydberg_data.pt\")\n",
    "# Below we have the parameters  (weights and biases) of the trained RBM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we have our trained RBM. Here we reconstruct the wavefunction by sampling from the RBM. \n",
    "# Let us see step by step how it works -\n",
    "\n",
    "# The first step is to sample from our trained RBM using QuCumber.\n",
    "# We then get num_samples samples which is stored in the variable 'samples'\n",
    "num_samples = 10000\n",
    "samples = nn_state.sample(num_samples = num_samples , k = 10)\n",
    "\n",
    "# Now we print out the samples in 'output/reconstructedSample.txt'\n",
    "sampleList = samples.tolist() # convert to list for convenience\n",
    "sampleList\n",
    "with open('output/reconstructedSample.txt', 'w') as fp:\n",
    "    for item in sampleList:\n",
    "        fp.write(\"%s\\n\" % item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "## -------------------------------------------------------------------- ##\n",
    "## This block of code computes amplitudes. For nY > 8, it is not necessary\n",
    "## to run this block as computation takes a lot of time.\n",
    "## -------------------------------------------------------------------- ##\n",
    "\n",
    "# From the samples we obtain the amplitudes of each of the basis states.\n",
    "# Note that here we have an 8 atom array. Hence, there are 2^8 basis states\n",
    "# which are: |00000000>, |00000001>, ..., |11111111>\n",
    "\n",
    "# Now say X is any of the states from {|00000000>, |00000001>, ..., |11111111>}\n",
    "# Now, the amplitude of a state X is computed by taking the\n",
    "# square root of the number of states X present in \n",
    "# our produced sample in 'output/reconstructedSample.txt'.\n",
    "\n",
    "# Therefore the reconstructed sample is like: \n",
    "#    |psi> = amplitudeList[0] |00000000> + amplitudeList[1] |00000001> + ... + amplitudeList[255] |11111111>  \n",
    "# where amplitudeList is defined below.\n",
    "\n",
    "# The following functions below finds the amplitudes in the way described above, stores the \n",
    "# amplitudes in amplitudeList and then also \n",
    "# prints out the amplitudes in 'output/reconstructedSample.txt'\n",
    "\n",
    "def amplitude(sample):\n",
    "    return np.sqrt(sampleList.count(sample)/num_samples)\n",
    "\n",
    "def getBinaryString(i):\n",
    "    sites = nv # nv = 8\n",
    "    getbinary = lambda x, n: format(x, 'b').zfill(n)\n",
    "    tempStr = getbinary(i, sites)\n",
    "    return tempStr\n",
    "\n",
    "# Python code to convert string to list character-wise\n",
    "def ConvertToList(string):\n",
    "    list1=[]\n",
    "    list1[:0]=string\n",
    "    return list1\n",
    "\n",
    "amplitudeList = []\n",
    "def getAmplitude(sample):\n",
    "    for i in range (2**nv):\n",
    "        binaryString = getBinaryString(i)\n",
    "        strList = ConvertToList(binaryString)\n",
    "        tempStr = str(amplitude(list(map(int, strList))))\n",
    "        amplitudeList.append(tempStr)\n",
    "    return 0\n",
    "\n",
    "getAmplitude(sampleList)\n",
    "with open('output/reconstructedStateAmplitudes.txt', 'w') as fp:\n",
    "    for item in amplitudeList:\n",
    "        fp.write(\"%s\\n\" % item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.399, 0.3969, 0.4121, 0.4132, 0.3944, 0.4195, 0.3856, 0.4128]"
      ]
     },
     "execution_count": 189,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# So now we have obtained the reconstructed wavefunction.\n",
    "# Now in the 'data' directory for each delta we have 1_pt_fun and 2_pt_fn\n",
    "# 1_pt_fn is the one point function and \n",
    "# 2_pt_fn is the 2 point function.\n",
    "\n",
    "# 1_pt_fn is the average occupation of each site. Note that\n",
    "# there are 8 sites as there are 8 atoms. Occupation of each site\n",
    "# refers to the proportion of atoms in the up state (state 1) \n",
    "# in each state.\n",
    "\n",
    "# Now we get average occupation of each site by finding the \n",
    "# number of atoms in the 1 state in each site and then\n",
    "# by dividing by the number of the produced samples \n",
    "occupation = [0]*nv\n",
    "for i in range(len(sampleList)):\n",
    "    j = 0\n",
    "    for j in range(nv):\n",
    "        occupation[j] = sampleList[i][j] + occupation[j]\n",
    "for i in range(nv):\n",
    "    occupation[i] = occupation[i]/len(sampleList)\n",
    "\n",
    "# Thus the average occupation per site of each site for the reconstructed state is given below\n",
    "# This data can be compared to data/nY=8/δ=delta_1_pt_fn.csv for the relevant delta\n",
    "\n",
    "# (Note: The list in data/nY=8/δ=delta_1_pt_fn.csv was obtained from \n",
    "# the sample in data/nY=8/δ=delta_samples.csv)\n",
    "\n",
    "# Thus the data in 'occupation' and 'data/nY=8/δ=delta_1_pt_fn.csv' are expected to be close (which they are)\n",
    "occupation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.239799  , -0.1195631 ,  0.1284721 , -0.1196668 ,  0.1225344 ,\n",
       "        -0.1231805 ,  0.1236456 , -0.1207072 ],\n",
       "       [-0.1195631 ,  0.23937039, -0.12536249,  0.12710092, -0.12193736,\n",
       "         0.12710045, -0.11804464,  0.12655968],\n",
       "       [ 0.1284721 , -0.12536249,  0.24227359, -0.12507972,  0.12896776,\n",
       "        -0.12637595,  0.12369424, -0.12291488],\n",
       "       [-0.1196668 ,  0.12710092, -0.12507972,  0.24246576, -0.12536608,\n",
       "         0.1329626 , -0.11992992,  0.12683104],\n",
       "       [ 0.1225344 , -0.12193736,  0.12896776, -0.12536608,  0.23884864,\n",
       "        -0.1266508 ,  0.12311936, -0.12180832],\n",
       "       [-0.1231805 ,  0.12710045, -0.12637595,  0.1329626 , -0.1266508 ,\n",
       "         0.24351975, -0.1207592 ,  0.1309304 ],\n",
       "       [ 0.1236456 , -0.11804464,  0.12369424, -0.11992992,  0.12311936,\n",
       "        -0.1207592 ,  0.23691264, -0.12347568],\n",
       "       [-0.1207072 ,  0.12655968, -0.12291488,  0.12683104, -0.12180832,\n",
       "         0.1309304 , -0.12347568,  0.24239616]])"
      ]
     },
     "execution_count": 190,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Now we look at the 2 point function\n",
    "# 2 point function is the covariance matrix for occupations.\n",
    "# It is found by 2_pt_fn = outerProdT1 - OuterProdT2\n",
    "\n",
    "# outerProdT1 is the average of the outer product of the spin vectors with themselves\n",
    "# So we first find the sum of the outerproducts of the spin vectors with themselves.\n",
    "# Then we take the average\n",
    "outerProdT1 = np.zeros((nv,nv))\n",
    "for i in range(len(sampleList)):\n",
    "    outerProdT1 = outerProdT1 + np.outer(sampleList[i], sampleList[i])\n",
    "outerProdT1 = outerProdT1/len(sampleList)\n",
    "\n",
    "# outerProd2 is simply the outer product of the occupation vector\n",
    "# with itself. Note that we already found the occupation vector for the \n",
    "# one point function\n",
    "\n",
    "outerProdT2 = np.outer(occupation,occupation)\n",
    "outerProdT2\n",
    "\n",
    "#Hence we have:\n",
    "TwoPointFunc = outerProdT1 - outerProdT2\n",
    "\n",
    "# The TwoPointFunction data for the reconstructed state is given below.\n",
    "# This matrix can be compared to data/nY=8/δ=delta_2_pt_fn.csv for the relevant delta\n",
    "\n",
    "# (Note: The matrix in data/nY=8/δ=delta_2_pt_fn.csv was obtained from \n",
    "# the sample in data/nY=8/δ=delta_samples.csv)\n",
    "\n",
    "# Thus the matrix in 'TwoFuncPoint' and 'data/nY=8/δ=delta_2_pt_fn.csv' are expected \n",
    "# to be nearly equal (which they are)\n",
    "TwoPointFunc"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}