2017-06-07_mg_dicke1_QSL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "1"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from os.path import join\n",
    "from io import StringIO\n",
    "from textwrap import dedent\n",
    "\n",
    "import sympy\n",
    "from sympy import Symbol, sqrt, cos, pi, symbols\n",
    "import numpy as np\n",
    "import matplotlib.pylab as plt\n",
    "import pandas as pd\n",
    "from collections import defaultdict\n",
    "\n",
    "import qnet\n",
    "from qnet.algebra import *\n",
    "\n",
    "import QDYN\n",
    "\n",
    "from src.notebook_plots_v1 import plot_bs_decay, display_hamiltonian, display_eq, show_summary_dicke\n",
    "from src.single_sided_network_v1 import network_slh\n",
    "from src.dicke_single_model_v1 import write_dicke_single_model, err_dicke_single\n",
    "\n",
    "from doit.tools import register_doit_as_IPython_magic\n",
    "import clusterjob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "2"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "qnet.init_printing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "3"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "DOIT_CONFIG = {\n",
    "    'backend': 'json',\n",
    "    'dep_file': '.doit_db/2017-06-07_mg_dicke1_QSL.json',\n",
    "}\n",
    "register_doit_as_IPython_magic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "4"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "clusterjob.JobScript.read_defaults('./config/copper.ini')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\n",
    "\\newcommand{ket}[1]{\\vert #1 \\rangle}\n",
    "\\newcommand{bra}[1]{\\langle #1 \\vert}\n",
    "\\newcommand{Op}[1]{\\hat{#1}}\n",
    "$\n",
    "\n",
    "# QSL for Single-Excitation Dicke state\n",
    "\n",
    "## action wrappers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*This notebook used a buggy version of QDYN where the MCWF trajectory optimization would use a simple non-Hermtian propagation for the forward propagation*\n",
    "\n",
    "Data is now located in `./data/dicke1_QSL_xxx/`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "5"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def update_config(rf, lambda_a, iter_stop):\n",
    "    config = QDYN.config.read_config_file(join(rf, 'config'))\n",
    "    config['oct']['iter_stop'] = iter_stop\n",
    "    for pulse_config in config['pulse']:\n",
    "        pulse_config['oct_lambda_a'] = lambda_a\n",
    "    QDYN.config.write_config(config, join(rf, 'config'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "6"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def submit_optimization(rf, n_trajs, task):\n",
    "    \"\"\"Asynchronously run optimization\"\"\"\n",
    "    if os.path.isfile(join(rf, 'oct.job.dump')):\n",
    "        return\n",
    "    body = dedent(r'''\n",
    "    {module_load}\n",
    "\n",
    "    cd $CLUSTERJOB_WORKDIR\n",
    "    pwd\n",
    "\texport MPICH_NO_BUFFER_ALIAS_CHECK=1\n",
    "    OMP_NUM_THREADS=1 {runner} qdyn_optimize --n-trajs={n_trajs} \\\n",
    "         --J_T=J_T_sm .\n",
    "    ''')\n",
    "    taskname = \"oct_%s\" % task.name.split(\":\")[-1]\n",
    "    jobscript = clusterjob.JobScript(\n",
    "        body=body, jobname=taskname, workdir=rf,\n",
    "        stdout='oct.log')\n",
    "    runner = {\n",
    "        'copper': 'aprun -B',\n",
    "        'mlhpc': 'mpirun -n %d' % n_trajs}\n",
    "    if jobscript.cluster_name != 'copper':\n",
    "        jobscript.ppn = int(n_trajs)\n",
    "        if n_trajs == 1:\n",
    "            runner['mlhpc'] = ''\n",
    "    jobscript.runner = runner[jobscript.cluster_name]\n",
    "    jobscript.n_trajs = str(int(n_trajs))\n",
    "    run = jobscript.submit(cache_id=taskname)\n",
    "    run.dump(join(rf, 'oct.job.dump'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "7"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def wait_for_clusterjob(dumpfile):\n",
    "    \"\"\"Wait until the clusterjob.AsyncResult cached in the given dumpfile ends\"\"\"\n",
    "    try:\n",
    "        run = clusterjob.AsyncResult.load(dumpfile)\n",
    "        run.wait()\n",
    "        os.unlink(dumpfile)\n",
    "        return run.successful()\n",
    "    except OSError:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## custom uptodate routines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "8"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from src.qdyn_model_v1 import pulses_uptodate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## task definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "9"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def runfolder(row):\n",
    "    int_part, frac_part = (\"%.2f\" % row['T']).split('.')\n",
    "    root = './data/dicke1_QSL_xxx/'\n",
    "    if frac_part == '00':\n",
    "        rf = 'rf_%dnodes_T%s' % (row['n_nodes'], int_part)\n",
    "    else:\n",
    "        rf = 'rf_%dnodes_T%s_%s' % (row['n_nodes'], int_part, frac_part)\n",
    "    return join(root, rf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "10"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_create_runfolder():\n",
    "    \"\"\"Create all necessary runfolders for the runs defined in params_df\"\"\"\n",
    "    jobs = {}\n",
    "    slh_models = {}\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = runfolder(row)\n",
    "        n_nodes = int(row['n_nodes'])\n",
    "        try:\n",
    "            slh = slh_models[n_nodes]\n",
    "        except KeyError:\n",
    "            slh = network_slh(\n",
    "                n_cavity=2, n_nodes=n_nodes, topology='driven_bs_fb')\n",
    "            slh_models[n_nodes] = slh\n",
    "        if n_nodes not in slh_models:\n",
    "            slh_mode\n",
    "        if rf in jobs:\n",
    "            continue\n",
    "        jobs[rf] = {\n",
    "            'name': str(rf),\n",
    "            'actions': [\n",
    "                (write_dicke_single_model, [slh, ], dict(\n",
    "                    rf=rf, T=row['T'], theta=0, nt=500,\n",
    "                    kappa=1.0, E0_cycles=2, mcwf=True, non_herm=True,\n",
    "                    lambda_a=row['lambda_a'],\n",
    "                    iter_stop=int(row['iter_stop'])))],\n",
    "            'targets': [join(rf, 'config')],\n",
    "            'uptodate': [True, ] # up to date if target exists\n",
    "        }\n",
    "    for job in jobs.values():\n",
    "        yield job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "11"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_update_runfolder():\n",
    "    \"\"\"For every row in params_df, update the config file in the appropriate\n",
    "    runfolder with the value in that row\"\"\"\n",
    "    rf_jobs = defaultdict(list)\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = runfolder(row)\n",
    "        # we only update the config after any earlier optimization has finished\n",
    "        task_dep = ['wait_for_optimization:%s' % ind2 for ind2 in rf_jobs[rf]]\n",
    "        rf_jobs[rf].append(ind)\n",
    "        yield {\n",
    "            'name': str(ind),\n",
    "            'actions': [\n",
    "                (update_config, [], dict(\n",
    "                    rf=rf, lambda_a=row['lambda_a'],\n",
    "                    iter_stop=int(row['iter_stop'])))],\n",
    "            'file_dep': [join(rf, 'config')],\n",
    "            'uptodate': [False, ],  # always run task\n",
    "            'task_dep': task_dep}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "12"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_submit_optimization():\n",
    "    \"\"\"Run optimization for every runfolder from params_df\"\"\"\n",
    "    rf_jobs = defaultdict(list)\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = runfolder(row)\n",
    "        task_dep = ['wait_for_optimization:%s' % ind2 for ind2 in rf_jobs[rf]]\n",
    "        task_dep.append('update_runfolder:%s' % ind)\n",
    "        yield {\n",
    "            'name': str(ind),\n",
    "            'actions': [\n",
    "                (submit_optimization, [rf, ], dict(n_trajs=row['n_trajs']))],\n",
    "                # 'task' keyword arg is added automatically\n",
    "            'task_dep': task_dep,\n",
    "            'uptodate': [(pulses_uptodate, [], {'rf': rf}), ],\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "13"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_wait_for_optimization():\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = runfolder(row)\n",
    "        yield {\n",
    "            'name': str(ind),\n",
    "            'task_dep': ['submit_optimization:%d' % ind],\n",
    "            'actions': [\n",
    "                (wait_for_clusterjob, [join(rf, 'oct.job.dump')], {}),]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  OCT Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "14"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "params_data_str = r'''\n",
    "# n_nodes   T     lambda_a  n_trajs   iter_stop\n",
    "        2   1.00     0.001       32        1000\n",
    "        2   1.00    0.0005       32       10000\n",
    "        2   1.00    0.0005       32       15000\n",
    "        2   2.00     0.001       32        1000\n",
    "        2   2.00    0.0005       32       10000\n",
    "        2   2.00    0.0005       32       15000\n",
    "        2   3.00     0.001       32        1000\n",
    "        2   3.00    0.0005       32       10000\n",
    "        2   3.00    0.0005       32       15000\n",
    "        2   4.00     0.001       32        1000\n",
    "        2   4.00    0.0005       32       10000\n",
    "        2   5.00     0.001       32        1000\n",
    "        2   7.00     0.001       32        1000\n",
    "        2  10.00     0.001       32        1000\n",
    "        2  10.00    0.0005       32        2000\n",
    "        3   1.00     0.001       32        1000\n",
    "        3   1.00    0.0005       32       10000\n",
    "        3   2.00     0.001       32        1000\n",
    "        3   2.00    0.0005       32       10000\n",
    "        3   3.00     0.001       32        1000\n",
    "        3   3.00    0.0005       32       10000\n",
    "        3   4.00     0.001       32        1000\n",
    "        3   4.00    0.0005       32        4000\n",
    "        3   5.00     0.001       32        1000\n",
    "        3   5.00    0.0005       32        4000\n",
    "        3   8.00     0.001       32        1000\n",
    "        3   8.00    0.0005       32        2000\n",
    "        3   9.00     0.001       32        1000\n",
    "        3   9.00    0.0005       32        2000\n",
    "        3  10.00     0.001       32        1000\n",
    "        3  10.00    0.0005       32        2000\n",
    "        4   0.50     0.001       32        1000\n",
    "        4   1.00     0.001       32        1000\n",
    "        4   1.00    0.0005       32       10000\n",
    "        4   1.00    0.0005       32       15000\n",
    "        4   2.00     0.001       32        1000\n",
    "        4   2.00    0.0005       32       10000\n",
    "        4   2.00    0.0005       32       15000\n",
    "        4   3.00     0.001       32        1000\n",
    "        4   3.00    0.0005       32       10000\n",
    "        4   3.00    0.0005       32       15000\n",
    "        4   4.00     0.001       32        1000\n",
    "        4   4.00    0.0005       32       10000\n",
    "        4   5.00     0.001       32        1000\n",
    "        4   5.00    0.0005       32       10000\n",
    "        4   6.00     0.001       32        1000\n",
    "        4   6.00    0.0005       32       10000\n",
    "        4   7.00     0.001       32        1000\n",
    "        4   7.00    0.0005       32       10000\n",
    "        4   8.00     0.001       32        1000\n",
    "        4   8.00    0.0005       32        4000\n",
    "        4   9.00     0.001       32        1000\n",
    "        4   9.00    0.0005       32        4000\n",
    "        4  10.00     0.001       32        1000\n",
    "        4  10.00    0.0005       32        4000\n",
    "        5   1.00     0.001       32        1000\n",
    "        5   1.00    0.0005       32       10000\n",
    "        5   1.00    0.0005       32       15000\n",
    "        5   2.00     0.001       32        1000\n",
    "        5   2.00    0.0005       32       10000\n",
    "        5   2.00    0.0005       32       15000\n",
    "        5   3.00     0.001       32        1000\n",
    "        5   3.00    0.0005       32       10000\n",
    "        5   3.00    0.0005       32       15000\n",
    "        5   4.00     0.001       32        1000\n",
    "        5   4.00    0.0005       32       10000\n",
    "        5   4.00    0.0005       32       15000\n",
    "        5   5.00     0.001       32        1000\n",
    "        5   5.00    0.0005       32       10000\n",
    "        5   5.00    0.0005       32       15000\n",
    "        5   6.00     0.001       32        1000\n",
    "        5   6.00    0.0005       32       10000\n",
    "        5   7.00     0.001       32        1000\n",
    "        5   7.00    0.0005       32       10000\n",
    "        5   8.00     0.001       32        1000\n",
    "        5   8.00    0.0005       32       10000\n",
    "        5   9.00     0.001       32        1000\n",
    "        5   9.00    0.0005       32       10000\n",
    "        5   9.00    0.0005       32       20000\n",
    "        5  10.00     0.001       32        1000\n",
    "        5  10.00    0.0005       32        8000\n",
    "        6   1.00     0.001       32        1000\n",
    "        6   1.00    0.0005       32       10000\n",
    "        6   2.00     0.001       32        1000\n",
    "        6   2.00    0.0005       32       10000\n",
    "        6   3.00     0.001       32        1000\n",
    "        6   3.00    0.0005       32       10000\n",
    "        6   4.00     0.001       32        1000\n",
    "        6   4.00    0.0005       32       10000\n",
    "        6   5.00     0.001       32        1000\n",
    "        6   5.00    0.0005       32       10000\n",
    "        6   5.00    0.0005       32       15000\n",
    "        6   6.00    0.0005       32       10000\n",
    "        6   6.00    0.0005       32       15000\n",
    "        6   7.00    0.0005       32       10000\n",
    "        6   7.00    0.0005       32       15000\n",
    "        6   8.00     0.001       32        1000\n",
    "        6   8.00    0.0005       32       10000\n",
    "        6   8.00    0.0005       32       15000\n",
    "        6   9.00     0.001       32        1000\n",
    "        6   9.00    0.0005       32       10000\n",
    "        6   9.00    0.0005       32       15000\n",
    "        6  10.00     0.001       32        1000\n",
    "        6  10.00    0.0005       32       10000\n",
    "        7   5.00     0.001       32        1000\n",
    "        7   5.00    0.0005       32       10000\n",
    "        7   8.00     0.001       32        1000\n",
    "        7   8.00    0.0005       32       10000\n",
    "        7   9.00     0.001       32        1000\n",
    "        7   9.00    0.0005       32       10000\n",
    "        7  10.00     0.001       32        1000\n",
    "        7  10.00    0.0005       32       10000\n",
    "'''\n",
    "params_df = pd.read_fwf(\n",
    "        StringIO(params_data_str), comment='#', header=1,\n",
    "        names=['n_nodes', 'T', 'lambda_a', 'n_trajs', 'iter_stop'],\n",
    "        converters={'n_nodes': int, 'T': float, 'lambda_a': float,\n",
    "                    'n_trajs': int, 'iter_stop': int})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "15"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(params_df.to_string())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "16"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "root = logging.getLogger()\n",
    "for handler in root.handlers[:]:\n",
    "    root.removeHandler(handler)\n",
    "logging.basicConfig(level=logging.DEBUG, filename='./data/dicke1_QSL_clusterjob.log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "22"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%doit -n 20 wait_for_optimization"
   ]
  }
 ],
 "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}