2DLattice_MJ_QPUAncillaEncoding.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aux import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The sequence of protein to fold\n",
    "# prot_seq = \"YYDPETGTWY\"\n",
    "prot_seq = \"YYDPET\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prot_seq, q, energy_expr, ene_terms_expr = prepare_quantum(prot_seq, return_Hs=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_func = lambdify(q, energy_expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sched = np.power(2, np.linspace(4, -1, 300))\n",
    "max(sched), min(sched), len(sched)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(sched)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.exp(-10/sched))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cores = 12\n",
    "\n",
    "res = Parallel(n_jobs=num_cores)(delayed(simulate_quantum)(prot_seq, q, energy_expr, sched, jobid) for jobid in range(720))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tra, H, E, seeds, results = zip(*res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tra = np.array(tra)\n",
    "history = np.array(H)\n",
    "energy = np.array(E)\n",
    "results = np.array(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acp = history.cumsum(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(acp.mean(axis=0))\n",
    "plt.fill_between(\n",
    "    range(acp.shape[1]),\n",
    "    acp.mean(axis=0) - acp.std(axis=0),\n",
    "    acp.mean(axis=0) + acp.std(axis=0), alpha=0.2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enes = np.array([energy_func(*qs) for qs in results])\n",
    "best = np.argmin(enes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(enes)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "correct = results[enes.round(4) == -6.20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_viz(tra[best], prot_seq, schedule=sched, movie=False, energy_func=energy_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(correct) > 0:\n",
    "    make_viz(correct[[0]], prot_seq, schedule=sched, movie=False, energy_func=energy_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_best =enes == min(enes)\n",
    "(all_best).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_correct = enes.round(2) == -6.20\n",
    "(all_correct).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.sort(enes))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rnd_best = np.random.choice(np.argwhere(all_best).reshape((-1,)))\n",
    "print(rnd_best)\n",
    "plt.plot(energy[rnd_best])\n",
    "plt.ylim(min(enes)-5, 200)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now Using DWave to find solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.system.samplers import DWaveSampler\n",
    "from dwave.system.composites import EmbeddingComposite, FixedEmbeddingComposite\n",
    "import dwave_networkx as dnx\n",
    "from minorminer import find_embedding\n",
    "import dimod\n",
    "from sympy.parsing.sympy_parser import parse_expr\n",
    "\n",
    "import json\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the solver and get its correspondiing adjacency graph for embedding\n",
    "# other sampler parameters are in dwave config file\n",
    "solver = DWaveSampler(solver=\"DW_2000Q_5\")\n",
    "solver_G = nx.Graph(solver.edgelist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Maximum anneal schedule points: {}\".format(solver.properties[\"max_anneal_schedule_points\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annealing_range = solver.properties[\"annealing_time_range\"]\n",
    "max_slope = 1.0/annealing_range[0]\n",
    "print(\"Annealing time range: {}\".format(solver.properties[\"annealing_time_range\"]))\n",
    "print(\"Maximum slope:\", max_slope)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fix first 3 qbits to avoid redundant conformations\n",
    "# for multiple problems in same graph\n",
    "energy_expr_less_00 = preprocess_expr(energy_expr.subs({'q0000':0, 'q0001':1, 'q0002':0}), q)\n",
    "energy_expr_less_00_func = lambdify(q[3:], energy_expr_less_00)\n",
    "energy_expr_ot_00 = energy_expr_less_00.as_ordered_terms()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hubo_00 = {}\n",
    "for coeff, qqss in [term.as_coeff_mul() for term in energy_expr_ot_00]:\n",
    "    try:\n",
    "        hubo_00[qqss[1:]] = float(qqss[0]*coeff)\n",
    "    except TypeError:\n",
    "        hubo_00[qqss] = float(coeff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poly_00 = dimod.BinaryPolynomial(hubo_00, dimod.BINARY)\n",
    "poly_00.normalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mk_quaq_strength = 1.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bqm_00 = dimod.make_quadratic(poly_00, mk_quaq_strength, dimod.BINARY)\n",
    "bqm_00.normalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bqm_spin_00 = bqm_00.change_vartype(dimod.SPIN, inplace=False)\n",
    "bqm_spin_00.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Will setup several copies of the problem in a single embedding, this way several solutions are retrieved from a single run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_subproblem(i):\n",
    "    qi = symbols([f'q{i:02d}{j:04d}' for j in range(len(q))])\n",
    "    energy_expr_less = energy_expr_less_00.subs(dict(zip(q, qi)))\n",
    "    return qi, energy_expr_less"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Nduplicates = 170 # good for setup with 4 residues\n",
    "Nduplicates = 7 # good for setup with 6 residues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cores = 12\n",
    "par_res = Parallel(n_jobs=num_cores)(delayed(create_subproblem)(i) for i in range(Nduplicates))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qii, energy_expr_less = zip(*par_res)\n",
    "qii = dict(zip(range(len(qii)), qii))\n",
    "energy_expr_less = sum(energy_expr_less)\n",
    "energy_expr_less = energy_expr_less.evalf()\n",
    "energy_expr_ot = energy_expr_less.as_ordered_terms()\n",
    "\n",
    "hubo = {}\n",
    "for coeff, qqss in [term.as_coeff_mul() for term in energy_expr_ot]:\n",
    "    try:\n",
    "        hubo[qqss[1:]] = float(qqss[0]*coeff)\n",
    "    except TypeError:\n",
    "        hubo[qqss] = float(coeff)\n",
    "\n",
    "poly = dimod.BinaryPolynomial(hubo, dimod.BINARY)\n",
    "poly.normalize()\n",
    "bqm = dimod.make_quadratic(poly, mk_quaq_strength, dimod.BINARY)\n",
    "bqm.normalize()\n",
    "\n",
    "bqm_spin = bqm.change_vartype(dimod.SPIN, inplace=False)\n",
    "bqm_spin.normalize()\n",
    "\n",
    "# Now make sure an embedding can be found\n",
    "problem = nx.Graph()\n",
    "for k, v in dict(bqm.linear).items():\n",
    "    problem.add_node(k)\n",
    "for k, v in dict(bqm.quadratic).items():\n",
    "    problem.add_edge(*k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cores = 12\n",
    "embedding_list = Parallel(n_jobs=num_cores)(delayed(find_embedding)(problem, solver_G) for jobid in range(num_cores))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assigned_qbits = np.array([len(sum(emb.values(), [])) if len(emb) > 0 else np.nan for emb in embedding_list])\n",
    "assert bool(sum(assigned_qbits[~np.isnan(assigned_qbits)])), \"No embedding has been found\"\n",
    "print(np.sort(assigned_qbits[~np.isnan(assigned_qbits)]))\n",
    "\n",
    "embedding = embedding_list[np.nanargmin(assigned_qbits)]\n",
    "\n",
    "all_assigned_qbits = sum(embedding.values(), [])\n",
    "print(f\"Using {len(all_assigned_qbits)}/{solver.properties['num_qubits']} qbits\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = FixedEmbeddingComposite(solver, embedding=embedding)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find energy of optimal solution\n",
    "bst_00 = {}\n",
    "for j, qj in enumerate(q):\n",
    "    bst_00[qj] = results[best, j]\n",
    "\n",
    "smpl_00 = dict(zip(list(bqm_00.variables), [None]*len(bqm_00.variables)))\n",
    "for key in bqm_00.variables:\n",
    "    if type(key) == str:\n",
    "        k1, k2 = parse_expr(key).as_coeff_mul()[1]\n",
    "        smpl_00[key] = bst_00[k1] * bst_00[k2]\n",
    "    else:\n",
    "        smpl_00[key] = bst_00[key]\n",
    "\n",
    "smpl_spin_00 = smpl_00.copy()\n",
    "for key, val in smpl_00.items():\n",
    "    smpl_spin_00[key] = 2*val -1\n",
    "\n",
    "mene_00 = bqm_00.energy(smpl_00)\n",
    "mene_spin_00 = bqm_spin_00.energy(smpl_spin_00)\n",
    "print((mene_00, mene_spin_00))\n",
    "# viz_short_smpl(smpl_00, prot_seq, q, energy_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Corroborate energy of optimal solution in tandem sample\n",
    "bst = {}\n",
    "for i, qi in qii.items():\n",
    "    for j, qij in enumerate(qi):\n",
    "        bst[qij] = results[best, j]\n",
    "\n",
    "smpl = dict(zip(list(bqm.variables), [None]*len(bqm.variables)))\n",
    "smpl_spin = smpl.copy()\n",
    "for key, val in smpl.items():\n",
    "    if type(key) == str:\n",
    "        k1, k2 = parse_expr(key).as_coeff_mul()[1]\n",
    "        smpl[key] = bst[k1] * bst[k2]\n",
    "    else:\n",
    "        smpl[key] = bst[key]\n",
    "    smpl_spin[key] = 2 * smpl[key] - 1\n",
    "\n",
    "mene = bqm.energy(smpl)\n",
    "mene_spin = bqm_spin.energy(smpl_spin)\n",
    "\n",
    "# after normalizations, the energy landscape has been scaled\n",
    "# this is the target energy for the true minimum (as found by conventional simulation)\n",
    "\n",
    "# should be equal to the min energy for a single problem\n",
    "assert round(mene / Nduplicates - mene_00, 10) == 0, \"Found a problem with tandem model energies\"\n",
    "\n",
    "# should be equal to the min energy for a single problem\n",
    "assert round(mene_spin / Nduplicates - mene_spin_00, 10) == 0, \"Found a problem with tandem model energies\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subsdic = {}\n",
    "keybook = {}\n",
    "for i, qi in qii.items():\n",
    "    tmp = dict(zip(qi, q))\n",
    "\n",
    "    for key in bqm.variables:\n",
    "        if key in qi:\n",
    "            keybook[key] = i\n",
    "            subsdic[key] = tmp[key]\n",
    "        elif type(key) == str:\n",
    "            k1, k2 = parse_expr(key).as_coeff_mul()[1]\n",
    "            if k1 in qi:\n",
    "                keybook[key] = i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def separate_subsamples(smpl, occurrences=1, vartype=dimod.BINARY):\n",
    "\n",
    "    # separate samples in batch\n",
    "    samples = [{} for i in range(Nduplicates)]\n",
    "\n",
    "    for key, val in smpl.items():\n",
    "        i = keybook[key]\n",
    "        if type(key) == str:\n",
    "            k1, k2 = parse_expr(key).as_coeff_mul()[1]\n",
    "            k1 = k1.subs(k1, subsdic[k1])\n",
    "            k2 = k2.subs(k2, subsdic[k2])\n",
    "            nkey = f'{k1}*{k2}'\n",
    "            if nkey not in bqm_00.variables:\n",
    "                nkey = f'{k2}*{k1}'\n",
    "        else:\n",
    "            nkey = key.subs(key, subsdic[key])\n",
    "\n",
    "        if vartype == dimod.SPIN:\n",
    "            val = int(val/2 + 0.5)\n",
    "\n",
    "        samples[i][nkey] = val\n",
    "\n",
    "    return samples, [occurrences] * Nduplicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_valid_bqsample(sample):\n",
    "    valid = True\n",
    "    for k, v in sample.items():\n",
    "        if type(k) == str:\n",
    "            k1, k2 = parse_expr(k).as_coeff_mul()[1]\n",
    "            valid = valid and (sample[k1] * sample[k2] == sample[k])\n",
    "    return valid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_reads = 1000\n",
    "\n",
    "response = sampler.sample(bqm_spin, num_reads=num_reads,\n",
    "                          annealing_time=20,\n",
    "                          num_spin_reversal_transforms=2,\n",
    "                          postprocess='optimization',\n",
    "                          answer_mode='raw')\n",
    "\n",
    "tused = response.info['timing']['qpu_access_time']\n",
    "print(f\"QPU time used: {tused/1e3:.0f} ms or {tused/num_reads/1e3:.3f} ms/read\")\n",
    "\n",
    "dat = response.record.energy\n",
    "# plt.hist(dat)\n",
    "print(f\"mean: {np.mean(dat):.2f}, sd: {np.std(dat):.2f}\")\n",
    "# plt.show()\n",
    "\n",
    "num_cores = 12\n",
    "par_res = Parallel(n_jobs=num_cores)(delayed(separate_subsamples)(sample, occurrences, response.vartype) for sample, occurrences in response.data(['sample', 'num_occurrences']))\n",
    "subsamp_lst, occurre_lst = zip(*par_res)\n",
    "\n",
    "all_subsamp = np.array(sum(subsamp_lst, []))\n",
    "all_occurre = np.array(sum(occurre_lst, []))\n",
    "\n",
    "all_sbsvald = np.array(Parallel(n_jobs=num_cores)(delayed(check_valid_bqsample)(smpl) for smpl in all_subsamp))\n",
    "all_sbsener = np.array(Parallel(n_jobs=num_cores)(delayed(bqm_00.energy)(smpl) for smpl in all_subsamp))\n",
    "all_sbsstat = np.array([[smpl[k] for k in bqm_00.variables] for smpl in all_subsamp])\n",
    "\n",
    "sbsorder = np.argsort(all_sbsener)\n",
    "all_subsamp = all_subsamp[sbsorder]\n",
    "all_occurre = all_occurre[sbsorder]\n",
    "all_sbsvald = all_sbsvald[sbsorder]\n",
    "all_sbsener = all_sbsener[sbsorder]\n",
    "all_sbsstat = all_sbsstat[sbsorder]\n",
    "updated_tab = False\n",
    "\n",
    "print(all_sbsener[:9].round(4))\n",
    "print(all_sbsener[all_sbsvald][:9].round(4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz_short_smpl(all_subsamp[all_sbsvald][0], prot_seq, q, energy_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(all_sbsener[all_sbsvald])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(all_sbsener[all_sbsvald&(all_sbsener<0)])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(all_sbsener[all_sbsvald], [-0.05 + 0.005*i for i in range(60)])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tab_sbsstat, cnt_sbsstat = np.unique(all_sbsstat, axis=0, return_counts=True)\n",
    "tab_subsamp = np.array([dict(zip(bqm_00.variables, stat)) for stat in tab_sbsstat])\n",
    "tab_sbsvald = np.array(Parallel(n_jobs=num_cores)(delayed(check_valid_bqsample)(smpl) for smpl in tab_subsamp))\n",
    "tab_sbsener = np.array(Parallel(n_jobs=num_cores)(delayed(bqm_00.energy)(smpl) for smpl in tab_subsamp))\n",
    "\n",
    "tab_sbsordr = np.argsort(tab_sbsener)\n",
    "tab_subsamp = tab_subsamp[tab_sbsordr]\n",
    "tab_sbsvald = tab_sbsvald[tab_sbsordr]\n",
    "tab_sbsener = tab_sbsener[tab_sbsordr]\n",
    "tab_sbsstat = tab_sbsstat[tab_sbsordr]\n",
    "cnt_sbsstat = cnt_sbsstat[tab_sbsordr]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qq = [qi for qi in q if qi in bqm_00.variables]\n",
    "ee = energy_expr.subs({q[0]: 0, q[1]: 1, q[2]: 0})\n",
    "ee_func = lambdify(qq, ee)\n",
    "tt = [t.subs({q[0]: 0, q[1]: 1, q[2]: 0}) for t in ene_terms_expr]\n",
    "hb_f, ho_f, hi_f = [lambdify(qq, h) for h in tt] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame({\n",
    "    \"state\":[''.join([str(v) for k, v in s.items() if type(k) != str]) for s in tab_subsamp], \n",
    "    \"isvalid\": tab_sbsvald,\n",
    "    \"energy\": tab_sbsener,\n",
    "    \"Hb\": [hb_f(**{str(k): v for k, v in sample.items() if type(k) != str}) for sample in tab_subsamp],\n",
    "    \"Ho\": [ho_f(**{str(k): v for k, v in sample.items() if type(k) != str}) for sample in tab_subsamp],\n",
    "    \"Hi\": [hi_f(**{str(k): v for k, v in sample.items() if type(k) != str}) for sample in tab_subsamp],\n",
    "    \"count\": cnt_sbsstat})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[(df.isvalid & (df.energy < 0))]"
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}