Classical Spin Reservoir Demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "23ce7bc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m  Activating\u001b[22m\u001b[39m project at `c:\\Users\\mkornjaca\\Documents\\GitHub\\QRC-tutorials`\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Initialize the environment\n",
    "using Pkg\n",
    "Pkg.activate(\".\")\n",
    "Pkg.instantiate()\n",
    "ENV[\"DATADEPS_ALWAYS_ACCEPT\"] = true # accept the download of the MNIST dataset\n",
    "SHOW_PROGRESS_BAR = true # turn off progress bars"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d45256c4",
   "metadata": {},
   "source": [
    "In this demo, we process the MNIST dataset with classical spin reservior. The steps should be very familiar to the QRC expert, the only change is the simulation itself that integrates differential equations governing the spin reservoir evolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "386c6cb1-6d57-47fe-b8eb-22c7d0a9b7a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import required libraries\n",
    "using MLDatasets\n",
    "using MultivariateStats\n",
    "using OneHotArrays\n",
    "using Flux\n",
    "using Bloqade\n",
    "using Colors\n",
    "using ProgressBars\n",
    "using JLD2\n",
    "using LIBSVM\n",
    "using Statistics\n",
    "using CSV\n",
    "using DataFrames\n",
    "using PythonCall\n",
    "using PythonPlot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ff3d8f8c-5c2d-4be8-854f-550e4318446b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dataset MNIST:\n",
       "  metadata  =>    Dict{String, Any} with 3 entries\n",
       "  split     =>    :test\n",
       "  features  =>    28×28×10000 Array{Float32, 3}\n",
       "  targets   =>    10000-element Vector{Int64}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Download the MNIST dataset\n",
    "\n",
    "# training data set (each image has 28x28 pixels and there are 60000 training samples)\n",
    "data_train = MNIST(:train)\n",
    "# test data set (10000 test samples)\n",
    "data_test = MNIST(:test) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "31a01230",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ClassicalDetuningLayer"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Define the struct for CRC layer\n",
    "Base.@kwdef struct ClassicalDetuningLayer\n",
    "    atoms # atom positions\n",
    "    Ω::Real # Rabi frequency\n",
    "    t_start::Real # evolution starting time\n",
    "    t_end::Real # evolution ending time\n",
    "    step::Real  # readout time step\n",
    "    state::Vector{<:Real} # classical state\n",
    "    readout::String #readout (\"s\" or \"ss\")\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2398271b",
   "metadata": {},
   "source": [
    " The CRC is derived by promoting all the qubits (spin $1/2$'s) to $S\\rightarrow\\infty$, thus making them classical unit vectors ($\\hat{S}$).  The dequantized dynamics of the Rydberg Hamiltonian is described by:\n",
    "\\begin{equation}\n",
    "    \\frac{\\mathrm{d} \\hat{S}_i}{\\mathrm{d} t}= \\frac{\\partial H[\\hat{S}]}{\\partial \\hat{S}_i} \\times \\hat{S},\n",
    "\\end{equation}\n",
    "where the effective instantaneous magnetic field acting on the $\\hat{S}_i$ is:\n",
    "\\begin{equation}\n",
    "    \\frac{\\partial H[\\hat{S}]}{\\partial \\hat{S}_i}= \\frac{\\Omega(t)}{2} \\left[\\hat{x}\\cos{\\phi(t)}+\\hat{y}\\sin{\\phi(t)}\\right] + \\left[-\\frac{\\Delta_i(t)}{2} + \\frac{1}{4} \\sum_{j\\neq i}V_{ij}\\left(1 + \\hat{S}_j^{(z)} \\right)\\right] \\hat{z}.\n",
    "\\end{equation}\n",
    "Thus, the dynamics of $N_q$-classical spins is efficiently simulated by integrating a system of $3N_q$ equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "51be6363",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "train_svm (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function generate_Vmat(locs::Vector{Vector{Float64}}, C6=862690*2π)\n",
    "    nsites=length(locs)\n",
    "    Vmat=zeros(nsites,nsites)\n",
    "    VDel=zeros(nsites)\n",
    "    for i in 1:nsites\n",
    "        for j in 1:nsites\n",
    "            if i != j\n",
    "                Vmat[i,j]=C6/(norm(locs[i,:]-locs[j,:]))^6\n",
    "                VDel[i]+=Vmat[i,j]\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    return Vmat, VDel\n",
    "end\n",
    "\n",
    "function deriv!(du, u, p::Tuple{Int, Vector{<:Real}, <:Real, Matrix{<:Real}, Vector{<:Real}},t)\n",
    "    #calculating du/dt\n",
    "    nsites, Delta, Omega, Vmat, VDel =p\n",
    "    Bv=Vmat*u[3:3:3*nsites]\n",
    "    for i in 1:nsites\n",
    "        Bi = [Omega, 0, -Delta[i]+Bv[i]/2+VDel[i]/2]/2\n",
    "        s = u[3*i-2:3*i]\n",
    "        du[3*(i-1)+1] = Bi[3]*s[2] - Bi[2]*s[3]\n",
    "        du[3*(i-1)+2] = Bi[1]*s[3] - Bi[3]*s[1]\n",
    "        du[3*(i-1)+3] = Bi[2]*s[1] - Bi[1]*s[2] \n",
    "    end\n",
    "    return du\n",
    "end\n",
    "\n",
    "function apply_classical_layer(layer::ClassicalDetuningLayer, x::Vector{<:Real})\n",
    "    # define Rydberg hamiltonian, detunings parameterized in terms of pca values (x)  \n",
    "    locs=layer.atoms\n",
    "    readout=layer.readout\n",
    "    nsites=length(locs)\n",
    "    Vmat, VDel =generate_Vmat(locs) \n",
    "    # At the start of the simulation, all atoms of the system are in ground states\n",
    "    u0 = layer.state\n",
    "\n",
    "    t_start = layer.t_start\n",
    "    t_end = layer.t_end\n",
    "    t_step = layer.step\n",
    "    \n",
    "    # initialize output vector\n",
    "    steps = floor(Int, (t_end - t_start) / t_step)\n",
    "    out2=zeros(3*nsites, 3*nsites, steps)\n",
    "    out3=zeros(nsites, nsites, steps)\n",
    "    out31=zeros(div(nsites*(nsites-1),2), steps)\n",
    "    # Numerically simulate the classical evolution with and store the results\n",
    "    Omega=layer.Ω\n",
    "    timespan=(t_start,t_end)\n",
    "    prob = ODEProblem(deriv!,u0,timespan,(nsites,x,Omega,Vmat, VDel))\n",
    "    sol = solve(prob, RK4(), saveat= t_step, adaptive=false, dt=1e-3, save_start=false)\n",
    "    sol3=sol[3:3:3*nsites,:]\n",
    "    if readout==\"s\"\n",
    "        out=reduce(vcat, sol)\n",
    "    elseif readout==\"ss\"\n",
    "        for i in 1:steps\n",
    "            out2[:,:,i]=sol[:,i] .* sol[:,i]'\n",
    "        end\n",
    "        out=reduce(vcat,(reduce(vcat,sol),reduce(vcat,out2)))\n",
    "    elseif readout==\"zz\"\n",
    "        for i in 1:steps\n",
    "            ind=1\n",
    "            for i1 in 1:nsites\n",
    "                for i2 in (i1+1):nsites\n",
    "                    out31[ind, i]= sol3[i1,i]* sol3[i2,i]\n",
    "                    ind+=1\n",
    "                end\n",
    "            end\n",
    "            #out3[:,:,i]=sol3[:,i] .* sol3[:,i]'\n",
    "        end\n",
    "        out=reduce(vcat,(reduce(vcat,sol3),reduce(vcat,out31)))\n",
    "    elseif readout==\"szz\"\n",
    "        ind=1\n",
    "        for i in 1:steps\n",
    "            for i1 in 1:nsites\n",
    "                for i2 in (i1+1):nsites\n",
    "                    out31[ind, i]= sol3[i1,i]* sol3[i2,i]\n",
    "                    ind+=1\n",
    "                end\n",
    "            end\n",
    "            #out3[:,:,i]=sol3[:,i] .* sol3[:,i]'\n",
    "        end\n",
    "        out=reduce(vcat,(reduce(vcat,sol),reduce(vcat,out3)))\n",
    "    elseif readout==\"z\"\n",
    "        out=reduce(vcat, sol3)\n",
    "    end\n",
    "    \n",
    "    return out\n",
    "end\n",
    "\n",
    "# implement functions that apply a `DetuningLayer` to a matrix containing scaled detunings for each image\n",
    "function apply_classical_layer(layer::ClassicalDetuningLayer, x::Matrix{<:Real})\n",
    "    iter = SHOW_PROGRESS_BAR ? ProgressBar(1:size(x, 2)) : 1:size(x, 2)\n",
    "    outs = [apply_classical_layer(layer, x[:, i][:]) for i in iter]\n",
    "    return hcat(outs...)\n",
    "end\n",
    "\n",
    "\n",
    "accuracy(model, xs, targets) = sum(onecold(model(xs), 0:9) .== targets)/length(targets)\n",
    "function train_linear_nn!(xs_train, ys_train, xs_test, ys_test; \n",
    "    regularization::Float64 = 0.0, nepochs::Int = 100, batchsize::Int = 100, \n",
    "    opt = Flux.Adam(0.01), verbose::Bool, nonlinear::Bool=false)\n",
    "    \n",
    "    model = Chain(\n",
    "    Dense(length(xs_train[:, 1]), 10),\n",
    "    softmax\n",
    "    )\n",
    "\n",
    "    if nonlinear\n",
    "        model = Chain(\n",
    "            Dense(length(xs_train[:, 1]), 100, relu),\n",
    "            Dense(100, 100, relu),\n",
    "            Dense(100, 10),\n",
    "            softmax\n",
    "        )\n",
    "    end\n",
    "\n",
    "    loader = Flux.DataLoader((data = xs_train, label = ys_train); batchsize, shuffle=true);\n",
    "    ps = Flux.params(model)\n",
    "\n",
    "    verbose && println(\"Training...\")\n",
    "    losses = zeros(nepochs)\n",
    "    accs_train = zeros(nepochs)\n",
    "    accs_test = zeros(nepochs)\n",
    "    for epoch in (verbose ? ProgressBar(1:nepochs) : 1:nepochs)\n",
    "    l = 1.0\n",
    "    for (x, y) in loader\n",
    "        grads = Flux.gradient(ps) do\n",
    "            ŷ = model(x)\n",
    "            if iszero(regularization)\n",
    "                l = Flux.crossentropy(ŷ, y)\n",
    "            else\n",
    "                l = Flux.crossentropy(ŷ, y) + regularization * sum(sum(abs, p) for p in ps)\n",
    "            end\n",
    "        end\n",
    "        Flux.update!(opt, ps, grads)\n",
    "    end\n",
    "    losses[epoch] = Flux.crossentropy(model(xs_train), ys_train)\n",
    "    accs_train[epoch] = accuracy(model, xs_train, onecold(ys_train, 0:9))\n",
    "    accs_test[epoch] = accuracy(model, xs_test, ys_test)\n",
    "    end\n",
    "    return losses, accs_train, accs_test\n",
    "end\n",
    "\n",
    "function train_svm(xs, ys, test_features, test_targets)\n",
    "    model = svmtrain(xs, onecold(ys, 0:9); svmtype = SVC, kernel = Kernel.Linear)\n",
    "    train_ŷ, train_decision_values = svmpredict(model, xs);\n",
    "    acc_train = mean(train_ŷ .== onecold(ys, 0:9))\n",
    "    ŷ, decision_values = svmpredict(model, test_features);\n",
    "    acc_test = mean(ŷ .== test_targets)\n",
    "    return acc_train, acc_test\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e12c0795",
   "metadata": {},
   "source": [
    "Let's now implement CRC in a standard reservoir pipeline. Note that in contrast to QRC, simulations with large numbers of spin are possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ba217106",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We first use PCA to downsample the data into 10-dimensional vectors\n",
    "dim_pca = 8\n",
    "\n",
    "# Use the `fit` function from the `MultivariateStats` package to generate the projection operator for PCA\n",
    "model_pca = fit(PCA, reshape(data_train.features, (784, 60000)); maxoutdim=dim_pca);\n",
    "\n",
    "# Use the `predict` function from `MultivariateStats` package to compute the first 10 principal components\n",
    "x = MultivariateStats.predict(model_pca, reshape(data_train.features, (784, 60000)))\n",
    "\n",
    "# Let us see how it looks\n",
    "num_examples = 10000\n",
    "xs = x[:, 1:num_examples]\n",
    "\n",
    "#normalizing features\n",
    "Δ_max = 6\n",
    "spectral = max(abs(maximum(xs)), abs(minimum(xs)))\n",
    "xs = xs/spectral * Δ_max # to make sure values to be between [-6.0, 6.0]\n",
    "\n",
    "#one hot encoding of labels \n",
    "y = onehotbatch(data_train.targets, 0:9)\n",
    "# Let us see how labels look\n",
    "ys = y[:, 1:num_examples]\n",
    "\n",
    "d = 10.0\n",
    "nsites=dim_pca\n",
    "locs = [[i*d,0.0] for i in 1:nsites] # put atoms in a chain with 10 micron spacing\n",
    "u0 =zeros(3*nsites)\n",
    "for i in 1:nsites\n",
    "    u0[3*i] = -1 \n",
    "end\n",
    "pre_layer = ClassicalDetuningLayer(;\n",
    "    atoms=locs, \n",
    "    Ω = 2π, \n",
    "    t_start = 0.0, \n",
    "    t_end = 4.0, \n",
    "    step = 0.5, \n",
    "    state=u0,\n",
    "    readout=\"zz\"\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3dc34f18",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "288×10000 Matrix{Float64}:\n",
       " -0.0025347   -0.016056    -0.00667874   …  -0.0200377   -0.0196264\n",
       " -0.0394597   -0.037505    -0.00609092      -0.00435202  -0.0030882\n",
       " -0.0190444   -0.0621152   -0.0124669       -0.00530816  -0.0586101\n",
       " -0.059114    -0.0499508   -0.00379331      -0.0359914   -0.0512428\n",
       " -0.021711    -0.0844998   -0.0394741       -0.00952299  -0.0428961\n",
       " -0.0341243   -0.0304166   -0.0206318    …  -0.0250836   -0.0117516\n",
       " -0.0339379   -0.0181282   -0.00327114      -0.00976858  -0.041323\n",
       " -0.0019961   -0.00830646  -0.000632081     -0.0106883   -0.00456156\n",
       "  0.805169     0.996427     0.656976         0.434101     0.432186\n",
       " -0.209162    -0.141048     0.303297         0.368926     0.517966\n",
       "  ⋮                                      ⋱               \n",
       "  0.180948    -0.146288     0.623413        -0.0312721    0.00124957\n",
       " -0.396762    -0.20815     -0.0549247    …  -0.156231    -0.0291297\n",
       "  0.425304    -0.181323    -0.155715         0.0326403    0.0386768\n",
       "  0.00929519  -0.153269    -0.0059803        0.103712     0.0192941\n",
       " -0.0203814   -0.218083     0.000526884      0.518129    -0.449779\n",
       "  0.0218476   -0.189976     0.00149375      -0.108249     0.597191\n",
       " -0.19553      0.0484924   -0.0831955    …   0.119434    -0.0163808\n",
       "  0.209596     0.0422427   -0.235864        -0.0249526    0.0217495\n",
       " -0.459578     0.060106     0.0207804       -0.12466     -0.507019"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# uncomment the next line to see progress bar\n",
    "SHOW_PROGRESS_BAR = false \n",
    "embeddings = apply_classical_layer(pre_layer, xs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eb20f35d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "288×1000 Matrix{Float64}:\n",
       " -0.0176704   -0.00440874  -0.0600804   …  -0.00608695  -0.0115501\n",
       " -0.00250291  -0.0982701   -0.0476215      -0.00277866  -0.00277022\n",
       " -0.0184618   -0.013301    -0.023997       -0.0308231   -0.00327755\n",
       " -0.0394076   -0.00606576  -0.00966865     -0.0327782   -0.00718883\n",
       " -0.0538571   -0.0231886   -0.0573744      -0.00343375  -0.0184842\n",
       " -0.0101416   -0.0516895   -0.0162518   …  -0.020267    -0.0187568\n",
       " -0.0454267   -0.0368324   -0.0192229      -0.0704857   -0.0145644\n",
       " -0.00109283  -0.00116435  -0.00847021     -0.0066797   -0.000695645\n",
       "  0.45184      0.811446     0.0902105       0.974704     0.554119\n",
       "  0.577258    -0.595967    -0.210033        0.468598     0.445898\n",
       "  ⋮                                     ⋱               \n",
       " -0.340072    -0.0216938    0.101647        0.160414     0.588281\n",
       " -0.0816118    0.014536     0.143924    …  -0.0423724    0.394969\n",
       " -0.530955     0.0572977    0.0661621       0.155324    -0.0879907\n",
       "  0.422837    -0.0657711   -0.159792        0.0268686    0.0258371\n",
       "  0.101474     0.0440703   -0.226253       -0.00709715   0.0173469\n",
       "  0.660178     0.173715    -0.104009        0.0260159   -0.00386452\n",
       "  0.0983421   -0.061909     0.136308    …  -0.170919     0.509216\n",
       "  0.6398      -0.244031     0.0626612       0.626534    -0.113443\n",
       "  0.153542     0.163514     0.0887232      -0.165495    -0.0761648"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "num_test_examples = 1000\n",
    "\n",
    "# using the same PCA model we fit to the test set\n",
    "test_features = MultivariateStats.predict(model_pca, reshape(data_test.features, (784, 10000)))[:, 1:num_test_examples] \n",
    "\n",
    "# same rescaling we applied on train input get reasonable detuning values \n",
    "test_features_qrc = test_features/spectral * Δ_max \n",
    "\n",
    "# quantum embeddings for 100 test samples\n",
    "test_embeddings = apply_classical_layer(pre_layer, test_features_qrc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abbe64c8",
   "metadata": {},
   "source": [
    "And finally, let's train the linear SVM on the CRC embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "82058076",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.8241, 0.797)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tr_svm, ts_svm=train_svm(xs, ys, test_features_qrc, data_test.targets[1:num_test_examples])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "738293fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.9192, 0.871)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tr_svm, ts_svm=train_svm(embeddings, ys, test_embeddings, data_test.targets[1:num_test_examples])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf254000",
   "metadata": {},
   "source": [
    "The results are similar to the QRC with the equivalent parameters on this data."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.2",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}