New Atomgroup _upickle.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Atomgroups Pickling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import MDAnalysis as mda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import multiprocessing\n",
    "from MDAnalysis.tests.datafiles import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = mda.Universe(GRO, XTC)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Current Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before pickle:\n",
      "id u: 140474024000192\n",
      "id ag: 140473304156192\n",
      "id ag_u: 140474024000192\n",
      "\n",
      "after pickle:\n",
      "id ag: 140474066557536\n",
      "id ag._u: 140474024000192\n",
      "ts < Timestep 0 with unit cell dimensions [80.017006 80.017006 80.017006 60.       60.       90.      ] >\n",
      "\n",
      "id ag: 140474066557536\n",
      "id ag._u: 140474024000192\n",
      "ts < Timestep 0 with unit cell dimensions [80.017006 80.017006 80.017006 60.       60.       90.      ] >\n",
      "\n",
      "id ag: 140473753207904\n",
      "id ag._u: 140474024000192\n",
      "ts < Timestep 0 with unit cell dimensions [80.017006 80.017006 80.017006 60.       60.       90.      ] >\n",
      "\n",
      "id ag: 140473753207904\n",
      "id ag._u: 140474024000192\n",
      "ts < Timestep 0 with unit cell dimensions [80.017006 80.017006 80.017006 60.       60.       90.      ] >\n",
      "\n",
      "id ag: 140474066557536\n",
      "id ag._u: 140474024000192\n",
      "ts < Timestep 0 with unit cell dimensions [80.017006 80.017006 80.017006 60.       60.       90.      ] >\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def print_id(ags):\n",
    "    ag = ags[0]\n",
    "    print('id ag:',id(ag))\n",
    "    print('id ag._u:',id(ag._u))\n",
    "    print('ts',ag._u.trajectory.ts)\n",
    "    print('')\n",
    "\n",
    "\n",
    "ag = u.atoms[2:5]\n",
    "print('before pickle:')\n",
    "print('id u:',id(u))\n",
    "print('id ag:',id(ag))\n",
    "print('id ag_u:',id(ag._u))\n",
    "print('')\n",
    "print('after pickle:')\n",
    "\n",
    "\n",
    "p = multiprocessing.Pool(2)\n",
    "res = np.array([p.apply(print_id, args=([ag]))\n",
    "            for ts in u.trajectory[:5]])\n",
    "p.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before pickle:\n",
      "id u: 140474024000192\n",
      "id ag: 140472758431408\n",
      "id ag_u: 140474024000192\n",
      "\n",
      "after pickle:\n",
      "id u: 140473303000496\n",
      "id ag: 140473302206400\n",
      "id ag._u: 140473303000496\n",
      "ts < Timestep 0 with unit cell dimensions [80.017006 80.017006 80.017006 60.       60.       90.      ] >\n",
      "\n",
      "id u: 140473303000736\n",
      "id ag: 140473302210736\n",
      "id ag._u: 140473303000736\n",
      "ts < Timestep 1 with unit cell dimensions [80.13008 80.13008 80.13008 60.      60.      90.     ] >\n",
      "\n",
      "id u: 140474065816112\n",
      "id ag: 140474061773888\n",
      "id ag._u: 140474065816112\n",
      "ts < Timestep 2 with unit cell dimensions [79.98359 79.98359 79.98358 60.      60.      90.     ] >\n",
      "\n",
      "id u: 140474065816352\n",
      "id ag: 140474061774128\n",
      "id ag._u: 140474065816352\n",
      "ts < Timestep 3 with unit cell dimensions [80.00516 80.00516 80.00516 60.      60.      90.     ] >\n",
      "\n",
      "id u: 140474061090480\n",
      "id ag: 140474057052416\n",
      "id ag._u: 140474061090480\n",
      "ts < Timestep 4 with unit cell dimensions [80.13544 80.13544 80.13544 60.      60.      90.     ] >\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def print_id(u, ag):\n",
    "    print('id u:',id(u))\n",
    "\n",
    "    print('id ag:',id(ag))\n",
    "    print('id ag._u:',id(ag._u))\n",
    "    print('ts',ag._u.trajectory.ts)\n",
    "    print('')\n",
    "\n",
    "\n",
    "ag = u.atoms[2:5]\n",
    "print('before pickle:')\n",
    "print('id u:',id(u))\n",
    "print('id ag:',id(ag))\n",
    "print('id ag_u:',id(ag._u))\n",
    "print('')\n",
    "print('after pickle:')\n",
    "\n",
    "\n",
    "p = multiprocessing.Pool(2)\n",
    "i = 0\n",
    "res = np.array([p.apply(print_id, args=(u, ag))\n",
    "            for ts in u.trajectory[:5]])\n",
    "p.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## New Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before pickle:\n",
      "id u: 140667113197040\n",
      "id ag: 140665858102704\n",
      "id ag_u: 140667113197040\n",
      "\n",
      "after pickle:\n",
      "id ag: 140665857901856\n",
      "id ag._u: 140666406389168\n",
      "ts < Timestep 0 with unit cell dimensions [80.017006 80.017006 80.017006 60.       60.       90.      ] >\n",
      "\n",
      "id ag: 140665857902864\n",
      "id ag._u: 140666406389408\n",
      "ts < Timestep 1 with unit cell dimensions [80.13008 80.13008 80.13008 60.      60.      90.     ] >\n",
      "\n",
      "id ag: 140665857903296\n",
      "id ag._u: 140667171789360\n",
      "ts < Timestep 2 with unit cell dimensions [79.98359 79.98359 79.98358 60.      60.      90.     ] >\n",
      "\n",
      "id ag: 140665857903296\n",
      "id ag._u: 140667171793696\n",
      "ts < Timestep 3 with unit cell dimensions [80.00516 80.00516 80.00516 60.      60.      90.     ] >\n",
      "\n",
      "id ag: 140665857901904\n",
      "id ag._u: 140667167063728\n",
      "ts < Timestep 4 with unit cell dimensions [80.13544 80.13544 80.13544 60.      60.      90.     ] >\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def print_id(ags):\n",
    "    ag = ags[0]\n",
    "    print('id ag:',id(ag))\n",
    "    print('id ag._u:',id(ag._u))\n",
    "    print('ts',ag._u.trajectory.ts)\n",
    "    print('')\n",
    "\n",
    "\n",
    "ag = u.atoms[2:5]\n",
    "print('before pickle:')\n",
    "print('id u:',id(u))\n",
    "print('id ag:',id(ag))\n",
    "print('id ag_u:',id(ag._u))\n",
    "print('')\n",
    "print('after pickle:')\n",
    "\n",
    "\n",
    "p = multiprocessing.Pool(2)\n",
    "res = np.array([p.apply(print_id, args=([ag]))\n",
    "            for ts in u.trajectory[:5]])\n",
    "p.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update also works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "atom_sel = u.select_atoms('name OW and around 4 resid 13', updating=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4353 7360 7614 7653 7898 8542 9728 9732 9745]\n",
      "[  626  1471  3176  5927  5983  7653  8542  9403  9667  9745  9818 11078]\n",
      "[  626  1471  5835  5927  7653  7898  8456  8457  9403  9667  9745 11078]\n",
      "[ 5187  5835  6987  7898  8456  8457  9403  9728  9745 10620]\n",
      "[ 3074  3272  3633  5198  6371  7653  7898  8456  8542  9403  9745 10620]\n"
     ]
    }
   ],
   "source": [
    "for ts in u.trajectory[:5]:\n",
    "    print(atom_sel.resids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4353 7360 7614 7653 7898 8542 9728 9732 9745]\n",
      "\n",
      "[  626  1471  3176  5927  5983  7653  8542  9403  9667  9745  9818 11078]\n",
      "\n",
      "[  626  1471  5835  5927  7653  7898  8456  8457  9403  9667  9745 11078]\n",
      "\n",
      "[ 5187  5835  6987  7898  8456  8457  9403  9728  9745 10620]\n",
      "\n",
      "[ 3074  3272  3633  5198  6371  7653  7898  8456  8542  9403  9745 10620]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def print_sel(atom_sel):\n",
    "    print(atom_sel.resids)\n",
    "\n",
    "    print('')\n",
    "\n",
    "p = multiprocessing.Pool(2)\n",
    "res = np.array([p.apply(print_sel, args=([atom_sel]))\n",
    "            for ts in u.trajectory[:5]])\n",
    "p.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gsoc",
   "language": "python",
   "name": "gsoc"
  },
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}