Aiida_db_query_examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AiiDA DB Query"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I'm just collecting notes and simple examples of how to build queries in aiida. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "from aiida import load_dbenv\n",
    "load_dbenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "collapsed": true,
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "from aiida.orm import Calculation, Code, Computer, Data, Node\n",
    "from aiida.orm import CalculationFactory, DataFactory\n",
    "from aiida.backends.djsite.db import models\n",
    "from aiida.orm.querybuilder import QueryBuilder\n",
    "from datetime import datetime\n",
    "from aiida.orm.utils import load_node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## StructureData"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This subclass of Node contains structure data.\n",
    "\n",
    "Before we build a query, let's see what attributes this subclass contains. I happen to know that node 375 is of this class (containing a cubic SrTiO3 structure)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 221,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[u'kinds', u'pbc3', u'pbc2', u'pbc1', u'sites', u'cell']\n",
      "\n",
      "[{u'symbols': [u'Sr'], u'mass': 87.62, u'weights': [1.0], u'name': u'Sr'}, {u'symbols': [u'Ti'], u'mass': 47.867, u'weights': [1.0], u'name': u'Ti'}, {u'symbols': [u'O'], u'weights': [1.0], u'mass': 15.9994, u'name': u'O'}]\n"
     ]
    }
   ],
   "source": [
    "strucdata = load_node(375)\n",
    "print strucdata.get_attrs().keys()\n",
    "print ''\n",
    "print strucdata.get_attrs()['kinds']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of results returned:  51\n"
     ]
    }
   ],
   "source": [
    "qb = QueryBuilder()\n",
    "#We shall use this date as a cut-off date in a filter\n",
    "date=datetime(2016,11,1)\n",
    "StructureData = DataFactory('structure')\n",
    "\n",
    "qb.append(\n",
    "    StructureData,\n",
    "    # We project out only the ID and 'kinds' columns of the DB\n",
    "    project=[\"id\",\"attributes.kinds\"],\n",
    "    # Filter by date created\n",
    "    filters={\"ctime\": {\">\":date}})\n",
    "\n",
    "print \"Number of results returned: \", qb.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's have a look at the first 10 results. We may be interested for example in what elements are present in the structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22 Ba Ti O\n",
      "112 Ba Ti O\n",
      "118 Ba Ti O\n",
      "142 Ba Ti O\n",
      "129 Ba Ti O\n",
      "133 Ba Ti O\n",
      "137 Ba Ti O\n",
      "169 Ba Ti O\n",
      "147 Ba Ti O\n",
      "153 Ba Ti O\n"
     ]
    }
   ],
   "source": [
    "results = [i for i in qb.iterall()][:10]\n",
    "for iD, item in results:\n",
    "    elements = [basis['symbols'] for basis in item]\n",
    "    elements = [val for sublist in elements for val in sublist]\n",
    "    print iD, \" \".join(elements)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ParameterData"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This subclass of Node contains parameter data, which can be either input parameters, or output parameters (results!).\n",
    "\n",
    "Before we build a query, let's see what attributes this subclass contains. I happen to know that nodes 108 and 296 are input and output parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{u'CONTROL': {u'calculation': u'scf',\n",
       "  u'restart_mode': u'from_scratch',\n",
       "  u'wf_collect': True},\n",
       " u'ELECTRONS': {u'conv_thr': 1e-06},\n",
       " u'SYSTEM': {u'ecutrho': 240.0, u'ecutwfc': 30.0}}"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#For example an input ParameterData:\n",
    "inputparamnode = load_node(108)\n",
    "inputparamnode.get_attrs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[u'xml_warnings',\n",
       " u'monkhorst_pack_offset',\n",
       " u'non_colinear_calculation',\n",
       " u'symmetries',\n",
       " u'format_version',\n",
       " u'number_of_atoms',\n",
       " u'energy',\n",
       " u'rho_cutoff_units',\n",
       " u'energy_hartree',\n",
       " u'inversion_symmetry',\n",
       " u'energy_one_electron',\n",
       " u'magnetization_angle2',\n",
       " u'number_of_spin_components',\n",
       " u'magnetization_angle1',\n",
       " u'number_of_bands',\n",
       " u'fixed_occupations',\n",
       " u'wfc_cutoff_units',\n",
       " u'monkhorst_pack_grid',\n",
       " u'number_of_symmetries',\n",
       " u'spin_orbit_domag',\n",
       " u'pointgroup_international',\n",
       " u'fft_grid',\n",
       " u'fermi_energy',\n",
       " u'energy_one_center_paw',\n",
       " u'lsda',\n",
       " u'energy_accuracy_units',\n",
       " u'pp_check_flag',\n",
       " u'energy_ewald',\n",
       " u'no_time_rev_operations',\n",
       " u'wall_time',\n",
       " u'smearing_method',\n",
       " u'energy_ewald_units',\n",
       " u'non_local_df',\n",
       " u'number_of_species',\n",
       " u'wfc_cutoff',\n",
       " u'number_of_bravais_symmetries',\n",
       " u'parser_info',\n",
       " u'k_points_units',\n",
       " u'energy_one_electron_units',\n",
       " u'smooth_fft_grid',\n",
       " u'beta_real_space',\n",
       " u'constraint_mag',\n",
       " u'warnings',\n",
       " u'do_not_use_time_reversal',\n",
       " u'energy_threshold',\n",
       " u'energy_one_center_paw_units',\n",
       " u'init_wall_time_seconds',\n",
       " u'lda_plus_u_calculation',\n",
       " u'scf_iterations',\n",
       " u'has_dipole_correction',\n",
       " u'charge_density',\n",
       " u'dft_exchange_correlation',\n",
       " u'parser_warnings',\n",
       " u'wall_time_seconds',\n",
       " u'number_of_k_points',\n",
       " u'starting_magnetization',\n",
       " u'symmetries_units',\n",
       " u'volume',\n",
       " u'total_number_of_scf_iterations',\n",
       " u'spin_orbit_calculation',\n",
       " u'creator_name',\n",
       " u'pointgroup_schoenflies',\n",
       " u'number_of_atomic_wfc',\n",
       " u'tetrahedron_method',\n",
       " u'time_reversal_flag',\n",
       " u'energy_hartree_units',\n",
       " u'rho_cutoff',\n",
       " u'format_name',\n",
       " u'energy_xc',\n",
       " u'q_real_space',\n",
       " u'energy_units',\n",
       " u'has_electric_field',\n",
       " u'fermi_energy_units',\n",
       " u'energy_accuracy',\n",
       " u'lkpoint_dir',\n",
       " u'number_of_electrons',\n",
       " u'energy_xc_units',\n",
       " u'creator_version']"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#For example an output ParameterData:\n",
    "outputparamnode = load_node(296)\n",
    "outputparamnode.get_attrs().keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here let's try append a second thing to our query in order to only select those ParameterDatas for which a calculation was an input (i.e. output params)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of results returned:  40\n"
     ]
    }
   ],
   "source": [
    "qb = QueryBuilder()\n",
    "date=datetime(2016,11,1)\n",
    "ParameterData = DataFactory('parameter')\n",
    "\n",
    "qb.append(\n",
    "    ParameterData,\n",
    "    # We need to tag this 'append' so we can refer to it in the second 'append'\n",
    "    tag=\"outputparams\",\n",
    "    # Let's project out the ID and the total energy\n",
    "    project=[\"id\",\"attributes.energy\"],\n",
    "    filters={\"ctime\": {\">\":date}})\n",
    "\n",
    "# Defining relationships: select those ParameterData that for which Computer is an input\n",
    "qb.append(Calculation, input_of=\"outputparams\")\n",
    "print \"Number of results returned: \", qb.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "163   -3646.0172\n",
      "198   -3646.0172\n",
      "201   -3644.3705\n",
      "195   -3644.1018\n",
      "204   -3641.2883\n",
      "296   -14619.8863\n",
      "290   -14633.8570\n",
      "212   -3646.0298\n",
      "192   -3634.6463\n",
      "293   -14632.7079\n"
     ]
    }
   ],
   "source": [
    "# Let's print out the first 10 results:\n",
    "results = [i for i in qb.iterall()][:10]\n",
    "for iD, energy in results:\n",
    "    print \"{0}   {1:7.4f}\".format(iD, energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KpointsData "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Self-explanatory... I know that node 377 has k-points data. Let's see what attributes it has:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{u'mesh': [4, 4, 4], u'offset': [0.0, 0.0, 0.0]}"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kptsnode = load_node(377)\n",
    "kptsnode.get_attrs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of results returned:  42\n"
     ]
    }
   ],
   "source": [
    "qb = QueryBuilder()\n",
    "date=datetime(2016,11,1)\n",
    "KpointsData = DataFactory('array.kpoints')\n",
    "\n",
    "qb.append(\n",
    "    KpointsData,\n",
    "    # We need to tag this 'append' so we can refer to it in the second 'append'\n",
    "    tag=\"outputparams\",\n",
    "    # Let's project out the ID and the total energy\n",
    "    project=[\"id\",\"attributes.mesh\"],\n",
    "    filters={\"ctime\": {\">\":date}})\n",
    "\n",
    "# Defining relationships: select those KpointsData that were inputs for a Computer\n",
    "qb.append(Calculation, output_of=\"outputparams\")\n",
    "print \"Number of results returned: \", qb.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 218,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "157   [4, 4, 4]\n",
      "175   [4, 4, 4]\n",
      "179   [4, 4, 4]\n",
      "171   [4, 4, 4]\n",
      "183   [4, 4, 4]\n",
      "215   [4, 4, 4]\n",
      "223   [4, 4, 4]\n",
      "207   [4, 4, 4]\n",
      "167   [4, 4, 4]\n",
      "227   [4, 4, 4]\n"
     ]
    }
   ],
   "source": [
    "# Let's print out the first 10 results:\n",
    "results = [i for i in qb.iterall()][:10]\n",
    "\n",
    "for iD, mesh in results:\n",
    "    print \"{0}   {1}\".format(iD, mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UpfData "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This contains the pseudopotential information. For example, node 98 contains a pseudopotential for Ti."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{u'filename': u'ti_pbesol_v1.4.uspp.F.UPF', u'md5': u'ebdc10ccce51265d8c030443845e563a', u'element': u'Ti'}\n",
      "[u'SSSP']\n"
     ]
    }
   ],
   "source": [
    "upfnode = load_node(98)\n",
    "print upfnode.get_attrs()\n",
    "print upfnode.get_upf_family_names()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of results returned:  126\n"
     ]
    }
   ],
   "source": [
    "qb = QueryBuilder()\n",
    "date=datetime(2016,11,1)\n",
    "UpfData = DataFactory('upf')\n",
    "\n",
    "qb.append(\n",
    "    UpfData,\n",
    "    # We need to tag this 'append' so we can refer to it in the second 'append'\n",
    "    tag=\"inputparams\",\n",
    "    # Let's project out the ID and the total energy\n",
    "    project=[\"id\",\"attributes.filename\"],\n",
    "    filters={\"ctime\": {\">\":date}})\n",
    "\n",
    "# Defining relationships: select those KpointsData that were inputs for a Computer\n",
    "qb.append(Calculation, output_of=\"inputparams\")\n",
    "print \"Number of results returned: \", qb.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is quite a lot of info, but for this purpose, let's just print the unique pseudopotential filenames:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[u'Sr.pbesol-spn-rrkjus_psl.1.0.0.UPF',\n",
       " u'Ba_ONCV_PBEsol-1.0.oncvpsp.upf',\n",
       " u'Pb.pbesol-dn-kjpaw_psl.0.2.2.UPF',\n",
       " u'o_pbesol_v1.2.uspp.F.UPF',\n",
       " u'ti_pbesol_v1.4.uspp.F.UPF']"
      ]
     },
     "execution_count": 168,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fnames = [filename for iD, filename in qb.iterall()]\n",
    "list(set(fnames))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JobCalculation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Node 378 was a calculation job. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[u'retrieve_list',\n",
       " u'last_jobinfo',\n",
       " u'job_id',\n",
       " u'scheduler_state',\n",
       " u'parser',\n",
       " u'max_wallclock_seconds',\n",
       " u'jobresource_params',\n",
       " u'remote_workdir',\n",
       " u'state',\n",
       " u'linkname_retrieved',\n",
       " u'retrieve_singlefile_list',\n",
       " u'scheduler_lastchecktime']"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calcnode = load_node(378)\n",
    "calcnode.get_attrs().keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of results returned:  46\n"
     ]
    }
   ],
   "source": [
    "qb = QueryBuilder()\n",
    "date=datetime(2016,11,1)\n",
    "\n",
    "qb.append(\n",
    "    Calculation,\n",
    "    # Let's project out the ID and the state\n",
    "    project=[\"id\",\"attributes.state\", \"attributes.job_id\"],\n",
    "    filters={\"ctime\": {\">\":date}})\n",
    "\n",
    "print \"Number of results returned: \", qb.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "159   FINISHED    4078535.sdb\n",
      "176   FINISHED    4079024.sdb\n",
      "180   FINISHED    4079026.sdb\n",
      "172   FINISHED    4079023.sdb\n",
      "184   FINISHED    4079027.sdb\n",
      "216   FINISHED    4079127.sdb\n",
      "224   FINISHED    4079125.sdb\n",
      "208   FINISHED    4079043.sdb\n",
      "168   FINISHED    4079021.sdb\n",
      "228   FINISHED    4079126.sdb\n",
      "220   FINISHED    4079131.sdb\n",
      "252   FINISHED    4079137.sdb\n",
      "244   FINISHED    4079135.sdb\n",
      "260   FINISHED    4079139.sdb\n",
      "248   FINISHED    4079136.sdb\n",
      "256   FINISHED    4079138.sdb\n",
      "268   FINISHED    4079141.sdb\n",
      "232   FINISHED    4079130.sdb\n",
      "272   FINISHED    4079142.sdb\n",
      "236   FINISHED    4079145.sdb\n",
      "264   FINISHED    4079140.sdb\n",
      "240   FINISHED    4079133.sdb\n",
      "324   FINISHED    4079213.sdb\n",
      "334   FINISHED    4079214.sdb\n",
      "390   FINISHED    4079434.sdb\n",
      "360   SUBMISSIONFAILED    None\n",
      "362   SUBMISSIONFAILED    None\n",
      "370   RETRIEVALFAILED    4079430.sdb\n",
      "398   FINISHED    4079436.sdb\n",
      "344   FINISHED    4079215.sdb\n",
      "374   RETRIEVALFAILED    4079429.sdb\n",
      "382   FINISHED    4079432.sdb\n",
      "358   SUBMISSIONFAILED    None\n",
      "394   FINISHED    4079435.sdb\n",
      "366   FINISHED    4079428.sdb\n",
      "378   FINISHED    4079431.sdb\n",
      "414   FINISHED    4079440.sdb\n",
      "418   FINISHED    4079441.sdb\n",
      "422   FINISHED    4079442.sdb\n",
      "386   FINISHED    4079443.sdb\n",
      "410   FINISHED    4079439.sdb\n",
      "406   FINISHED    4079437.sdb\n",
      "402   FINISHED    4079438.sdb\n",
      "477   FINISHED    4079615.sdb\n",
      "484   FINISHED    4079617.sdb\n",
      "494   NEW    None\n"
     ]
    }
   ],
   "source": [
    "for iD, state, jobid in qb.iterall():\n",
    "    print \"{0}   {1}    {2}\".format(iD, state, jobid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Initialisation Cell",
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  },
  "latex_envs": {
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 0
  },
  "nav_menu": {
   "height": "169px",
   "width": "177px"
  },
  "toc": {
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": true,
   "threshold": 6,
   "toc_cell": false,
   "toc_section_display": "block",
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}