2017-04-21_mg_doit_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "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 collections import defaultdict\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",
    "import psutil\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, config_changed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "qnet.init_printing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "register_doit_as_IPython_magic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if not 'MODULEPATH' in os.environ:\n",
    "    f = open(os.environ['MODULESHOME'] + \"/init/.modulespath\", \"r\")\n",
    "    path = []\n",
    "    for line in f.readlines():\n",
    "        line = re.sub(\"#.*$\", '', line)\n",
    "        if line is not '':\n",
    "            path.append(line)\n",
    "    os.environ['MODULEPATH'] = ':'.join(path)\n",
    "\n",
    "if 'LOADEDMODULES' not in os.environ:\n",
    "    os.environ['LOADEDMODULES'] = ''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def module(*args):\n",
    "    if isinstance(args[0], list):\n",
    "        args = args[0]\n",
    "    else:\n",
    "        args = list(args)\n",
    "    (output, error) = subprocess.Popen(\n",
    "            ['/usr/bin/modulecmd', 'python'] +\n",
    "            args, stdout=subprocess.PIPE).communicate()\n",
    "    exec(output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "module('load', 'intel/xe2011')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "module('load', 'mpi')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "module('list')"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## action wrappers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "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": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def run_traj_oct(rf, n_trajs):\n",
    "    \"\"\"Asynchronously run optimization, storing the pid of the oct process\"\"\"\n",
    "    env = os.environ.copy()\n",
    "    env['OMP_NUM_THREADS'] = '1'\n",
    "    cmd = ['mpirun', '-n', str(int(n_trajs)), 'qdyn_optimize',\n",
    "           '--n-trajs=%s' % int(n_trajs), '--J_T=J_T_sm', '.']\n",
    "    with open(join(rf, 'oct.log'), 'wb') as log_fh:\n",
    "        proc = subprocess.Popen(\n",
    "            cmd, env=env, stdout=log_fh, stderr=subprocess.STDOUT, cwd=rf)\n",
    "        with open(join(rf, 'oct.pid'), 'w') as pid_fh:\n",
    "            pid_fh.write(str(proc.pid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def run_traj_prop(rf, n_trajs):\n",
    "    \"\"\"Run propagation\"\"\"\n",
    "    env = os.environ.copy()\n",
    "    env['OMP_NUM_THREADS'] = '1'\n",
    "    cmd = ['mpirun', '-n', str(int(n_trajs)), 'qdyn_prop_traj',\n",
    "           '--n-trajs=%s' % int(n_trajs), '--use-oct-pulses',\n",
    "           '--write-final-state=state_final.dat', '.']\n",
    "    with open(join(rf, 'prop.log'), 'wb') as log_fh:\n",
    "        subprocess.call(\n",
    "            cmd, env=env, stdout=log_fh, stderr=subprocess.STDOUT, cwd=rf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def wait_for_oct(rf):\n",
    "    \"\"\"Wait until the OCT process with the given pid ends\"\"\"\n",
    "    try:\n",
    "        with open(join(rf, 'oct.pid')) as pid_fh:\n",
    "            pid = int(pid_fh.read())\n",
    "            proc = psutil.Process(pid)\n",
    "            proc.wait()\n",
    "    except (psutil.NoSuchProcess, OSError):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## custom uptodate routines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def pulses_uptodate(rf, verbose=False):\n",
    "    \"\"\"Return True if all optimized pulses in the given runfolder\n",
    "    are converged or have reached the max number of iterations, False\n",
    "    otherwise\"\"\"\n",
    "    config = QDYN.config.read_config_file(join(rf, 'config'))\n",
    "    for pulse_config in config['pulse']:\n",
    "        pulse_oct_file = join(rf, pulse_config['oct_outfile'])\n",
    "        try:\n",
    "            p = QDYN.pulse.Pulse.read(pulse_oct_file)\n",
    "        except OSError:\n",
    "            if verbose:\n",
    "                print(\"    pulse %s does not exist\" % pulse_oct_file)\n",
    "            return False\n",
    "        if p.oct_iter < config['oct']['iter_stop']:\n",
    "            m = re.search(r'J_T\\s*=\\s*([0-9E.+-]+)', \"\".join(p.preamble))\n",
    "            if verbose:\n",
    "                print(\"    pulse %s has too foo iterations %d\"\n",
    "                      % (pulse_oct_file, p.oct_iter))\n",
    "            if float(m.group(1)) > config['oct']['J_T_conv']:\n",
    "                if verbose:\n",
    "                    print(\"    pulse %s is not converged\" % pulse_oct_file)\n",
    "                return False\n",
    "            else:\n",
    "                if verbose:\n",
    "                    print(\"    pulse %s is converged\" % pulse_oct_file)\n",
    "        if verbose:\n",
    "            print(\"    pulse %s is up to date\" % pulse_oct_file)\n",
    "    return True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## task definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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 = './data/doit/rf%d' % row['T']\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": 16,
   "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 = './data/doit/rf%d' % row['T']\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": 17,
   "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 = './data/doit/rf%d' % row['T']\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",
    "                (run_traj_oct, [rf, ], dict(n_trajs=row['n_trajs']))],\n",
    "            'task_dep': task_dep,\n",
    "            'uptodate': [(pulses_uptodate, [], {'rf': rf}), ],\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_wait_for_optimization():\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = './data/doit/rf%d' % row['T']\n",
    "        yield {\n",
    "            'name': str(ind),\n",
    "            'task_dep': ['submit_optimization:%d' % ind],\n",
    "            'actions': [\n",
    "                (wait_for_oct, [rf, ], {}),\n",
    "                'rm -f %s' % join(rf, 'oct.pid')]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def task_propagate():\n",
    "    \"\"\"Run optimization for every runfolder from params_df\"\"\"\n",
    "    jobs = {}\n",
    "    for ind, row in params_df.iterrows():\n",
    "        rf = './data/doit/rf%d' % row['T']\n",
    "        jobs[rf] = {\n",
    "            'name': str(rf),\n",
    "            'actions': [\n",
    "                (run_traj_prop, [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": "markdown",
   "metadata": {},
   "source": [
    "##  Bringing it all together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "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": 21,
   "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": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "slh = network_slh(n_cavity=2, n_nodes=4, topology='driven_bs_fb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- create_runfolder:./data/doit/rf10\n",
      "-- create_runfolder:./data/doit/rf20\n",
      "-- create_runfolder:./data/doit/rf50\n",
      "-- create_runfolder:./data/doit/rf70\n",
      ".  update_runfolder:0\n",
      ".  update_runfolder:3\n",
      ".  update_runfolder:4\n",
      ".  update_runfolder:5\n",
      "-- submit_optimization:0\n",
      "-- submit_optimization:3\n",
      "-- submit_optimization:4\n",
      "-- submit_optimization:5\n",
      ".  wait_for_optimization:0\n",
      ".  wait_for_optimization:3\n",
      ".  wait_for_optimization:4\n",
      ".  wait_for_optimization:5\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 wait_for_optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- propagate:./data/doit/rf10\n",
      "-- propagate:./data/doit/rf20\n",
      "-- propagate:./data/doit/rf50\n",
      "-- propagate:./data/doit/rf70\n"
     ]
    }
   ],
   "source": [
    "%doit propagate"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "459px",
    "width": "492px"
   },
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": true,
   "threshold": 4,
   "toc_cell": false,
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}