2017-05-01_mg_doit_clusterjob_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "1"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import re\n",
    "from os.path import join\n",
    "from io import StringIO\n",
    "from textwrap import dedent\n",
    "from collections import defaultdict\n",
    "\n",
    "import pandas as pd\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_half_model_v2 import write_dicke_half_model, err_dicke_half\n",
    "\n",
    "from doit.tools import register_doit_as_IPython_magic\n",
    "import clusterjob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "2"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "3"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "qnet.init_printing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "4"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "register_doit_as_IPython_magic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "clusterjob.JobScript.read_defaults('./config/mlhpc_cluster.ini')\n",
    "clusterjob.JobScript.cache_folder = './data/doit_clusterjobs/cache/'"
   ]
  },
  {
   "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",
    "$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the doit package for generating run data\n",
    "\n",
    "## action wrappers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def submit_optimization(rf, n_trajs, task):\n",
    "    \"\"\"Asynchronously run optimization\"\"\"\n",
    "    body = dedent(r'''\n",
    "    {module_load}\n",
    "\n",
    "    cd {rf}\n",
    "    OMP_NUM_THREADS=1 mpirun -n {n_trajs} qdyn_optimize --n-trajs={n_trajs} \\\n",
    "         --J_T=J_T_sm .\n",
    "    ''')\n",
    "    taskname = \"oct_%s\" % task.name.replace(\":\", '_')\n",
    "    jobscript = clusterjob.JobScript(\n",
    "        body=body, filename=join(rf, 'oct.slr'),\n",
    "        jobname=taskname, nodes=1, ppn=int(n_trajs), threads=1,\n",
    "        stdout=join(rf, 'oct.log'))\n",
    "    jobscript.rf = rf\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": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def submit_propagation(rf, n_trajs):\n",
    "    \"\"\"Run propagation\"\"\"\n",
    "    body = dedent(r'''\n",
    "    {module_load}\n",
    "\n",
    "    cd {rf}\n",
    "    OMP_NUM_THREADS=1 mpirun -n {n_trajs} qdyn_prop_traj --n-trajs={n_trajs} \\\n",
    "        --use-oct-pulses --write-final-state=state_final.dat .\n",
    "    ''')\n",
    "    taskname = \"prop_\" + os.path.split(rf)[-1]\n",
    "    jobscript = clusterjob.JobScript(\n",
    "        body=body, filename=join(rf, 'prop.slr'),\n",
    "        jobname=taskname, nodes=1, ppn=int(n_trajs), threads=1,\n",
    "        stdout=join(rf, 'prop.log'))\n",
    "    jobscript.rf = rf\n",
    "    jobscript.n_trajs = str(int(n_trajs))\n",
    "    run = jobscript.submit(cache_id=taskname, force=True)\n",
    "    run.dump(join(rf, 'prop.job.dump'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "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": 10,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "6"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from src.qdyn_model_v1 import pulses_uptodate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## task definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def runfolder(row):\n",
    "    return './data/doit_clusterjobs/rf%d' % row['T']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_create_runfolder():\n",
    "    \"\"\"Create all necessary runfolders for the runs defined in params_df\"\"\"\n",
    "    jobs = {}\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = runfolder(row)\n",
    "        if rf in jobs:\n",
    "            continue\n",
    "        jobs[rf] = {\n",
    "            'name': str(rf),\n",
    "            'actions': [\n",
    "                (write_dicke_half_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": 13,
   "metadata": {
    "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": 14,
   "metadata": {
    "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": 15,
   "metadata": {
    "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": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_submit_propagation():\n",
    "    \"\"\"Run optimization for every runfolder from params_df\"\"\"\n",
    "    jobs = {}\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = runfolder(row)\n",
    "        jobs[rf] = {\n",
    "            'name': str(rf),\n",
    "            'actions': [\n",
    "                (submit_propagation, [rf, ], dict(n_trajs=row['n_trajs']))],\n",
    "            'file_dep': [join(rf, 'pulse1.oct.dat'),],}\n",
    "    for job in jobs.values():\n",
    "        yield job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_wait_for_propagation():\n",
    "    \"\"\"Run optimization for every runfolder from params_df\"\"\"\n",
    "    jobs = {}\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = runfolder(row)\n",
    "        jobs[rf] = {\n",
    "            'name': str(rf),\n",
    "            'task_dep': ['submit_propagation:%s' % rf],\n",
    "            'actions': [\n",
    "                (wait_for_clusterjob, [join(rf, 'prop.job.dump')], {}),]}\n",
    "    for job in jobs.values():\n",
    "        yield job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Bringing it all together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "params_data_str = r'''\n",
    "#    T  lambda_a  n_trajs   iter_stop\n",
    "    10     0.001       10          10\n",
    "    10    0.0005       10          15\n",
    "    10    0.0001       10          30\n",
    "    20     0.001       10          20\n",
    "    50     0.001       10          30\n",
    "    70     0.001       10          30\n",
    "'''\n",
    "params_df = pd.read_fwf(\n",
    "        StringIO(params_data_str), comment='#', header=1,\n",
    "        names=['T', 'lambda_a', 'n_trajs', 'iter_stop'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>T</th>\n",
       "      <th>lambda_a</th>\n",
       "      <th>n_trajs</th>\n",
       "      <th>iter_stop</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>10</td>\n",
       "      <td>0.0010</td>\n",
       "      <td>10</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>10</td>\n",
       "      <td>0.0005</td>\n",
       "      <td>10</td>\n",
       "      <td>15</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>10</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>10</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>20</td>\n",
       "      <td>0.0010</td>\n",
       "      <td>10</td>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>50</td>\n",
       "      <td>0.0010</td>\n",
       "      <td>10</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>70</td>\n",
       "      <td>0.0010</td>\n",
       "      <td>10</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    T  lambda_a  n_trajs  iter_stop\n",
       "0  10    0.0010       10         10\n",
       "1  10    0.0005       10         15\n",
       "2  10    0.0001       10         30\n",
       "3  20    0.0010       10         20\n",
       "4  50    0.0010       10         30\n",
       "5  70    0.0010       10         30"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "slh = network_slh(n_cavity=2, n_nodes=4, topology='driven_bs_fb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "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/doit_clusterjobs/debug_clusterjob.log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- create_runfolder:./data/doit_clusterjobs/rf10\n",
      "-- create_runfolder:./data/doit_clusterjobs/rf20\n",
      "-- create_runfolder:./data/doit_clusterjobs/rf50\n",
      "-- create_runfolder:./data/doit_clusterjobs/rf70\n",
      ".  update_runfolder:0\n",
      ".  update_runfolder:3\n",
      ".  update_runfolder:4\n",
      "-- submit_optimization:0\n",
      ".  update_runfolder:5\n",
      "-- submit_optimization:3\n",
      "-- submit_optimization:5\n",
      "-- submit_optimization:4\n",
      ".  wait_for_optimization:0\n",
      ".  wait_for_optimization:3\n",
      ".  wait_for_optimization:5\n",
      ".  wait_for_optimization:4\n",
      ".  update_runfolder:1\n",
      "-- submit_optimization:1\n",
      ".  wait_for_optimization:1\n",
      ".  update_runfolder:2\n",
      "-- submit_optimization:2\n",
      ".  wait_for_optimization:2\n"
     ]
    }
   ],
   "source": [
    "%doit -n 4 wait_for_optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- submit_propagation:./data/doit_clusterjobs/rf10\n",
      "-- submit_propagation:./data/doit_clusterjobs/rf20\n",
      "-- submit_propagation:./data/doit_clusterjobs/rf50\n",
      "-- submit_propagation:./data/doit_clusterjobs/rf70\n",
      ".  wait_for_propagation:./data/doit_clusterjobs/rf10\n",
      ".  wait_for_propagation:./data/doit_clusterjobs/rf20\n",
      ".  wait_for_propagation:./data/doit_clusterjobs/rf50\n",
      ".  wait_for_propagation:./data/doit_clusterjobs/rf70\n"
     ]
    }
   ],
   "source": [
    "%doit -n 4 wait_for_propagation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "default",
   "language": "python",
   "name": "default"
  },
  "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
}