Build_TAA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import argparse\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('-l', '--length', type = int, help = 'Alkyl chain length')\n",
    "parser.add_argument('-n', '--name', help = 'name of the molecule')\n",
    "\n",
    "try:\n",
    "    __IPYTHON__\n",
    "    args = parser.parse_args([])\n",
    "except NameError:\n",
    "    args = parser.parse_args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Atom(object):\n",
    "    index = 0\n",
    "    cgroup = 0\n",
    "    def __init__(self, atomname):\n",
    "        '''Initialize an atom with a unique index and empty neighbor list and no charge group'''\n",
    "        self.name = atomname\n",
    "        Atom.index += 1\n",
    "        self.index = Atom.index\n",
    "        self.neighbors = []\n",
    "        self.cgroup = None\n",
    "    def assign_cgroup(self, n = None):\n",
    "        if n is None:\n",
    "            Atom.cgroup += 1\n",
    "            self.cgroup = Atom.cgroup\n",
    "        else:\n",
    "            self.cgroup = n\n",
    "    def add_neighbor(self, other):\n",
    "        '''Add neighbor to an atom'''\n",
    "        self.neighbors.append(other.index)\n",
    "        other.neighbors.append(self.index)\n",
    "    def add_neighbors(self, iterator):\n",
    "        '''Add neighbors from an iterator object'''\n",
    "        for item in iterator:\n",
    "            self.add_neighbor(item)\n",
    "    def __str__(self):\n",
    "        return self.name\n",
    "    def __repr__(self):\n",
    "        return self.name\n",
    "    def __lt__(self, other):\n",
    "        return self.index < other.index\n",
    "    @staticmethod\n",
    "    def clear_index():\n",
    "        '''Clean the atom index'''\n",
    "        Atom.index = 0\n",
    "        Atom.cgroup = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TAA(object):\n",
    "    '''Tetraalkylammonium'''\n",
    "    def __init__(self, name = 'TAA'):\n",
    "        '''Initialize an molecule with an N atom'''\n",
    "        self.name = name\n",
    "        self.atoms = {}\n",
    "        self.root = Atom('N3')\n",
    "        self.root.assign_cgroup()\n",
    "        self.add_atom(self.root)\n",
    "        self.bonds = []\n",
    "        self.angles = []\n",
    "        self.dihedrals = []\n",
    "    def get_number(self):\n",
    "        '''Return the '''\n",
    "        return len(self.atoms)\n",
    "    def add_atom(self, atom):\n",
    "        self.atoms[atom.index] = atom\n",
    "    def add_bond(self, atom1, atom2):\n",
    "        atom1.add_neighbor(atom2)\n",
    "        self.bonds.append((min(atom1.index, atom2.index), max(atom1.index, atom2.index)))\n",
    "    def add_alkyl(self, n = 2):\n",
    "        for rep in range(4):\n",
    "            prev = self.root\n",
    "            for i in range(n):\n",
    "                if i == 0:\n",
    "                    Cname = 'C1'\n",
    "                    Hname = 'H1'\n",
    "                else:\n",
    "                    Hname = 'HC'\n",
    "                    if i == 1:\n",
    "                        Cname = 'C2'\n",
    "                    elif i == n-1:\n",
    "                        Cname = 'CT'\n",
    "                    else:\n",
    "                        Cname = 'CS'\n",
    "                C = Atom(Cname)\n",
    "                C.assign_cgroup()\n",
    "                self.add_atom(C)\n",
    "                self.add_bond(prev, C)\n",
    "                prev = C\n",
    "                for j in range(2):\n",
    "                    H = Atom(Hname)\n",
    "                    self.add_atom(H)\n",
    "                    self.add_bond(C, H)\n",
    "                    H.assign_cgroup(C.cgroup)\n",
    "            H = Atom('HC')\n",
    "            self.add_atom(H)\n",
    "            self.add_bond(C, H)\n",
    "            H.assign_cgroup(C.cgroup)\n",
    "        self.add_angles()\n",
    "        self.add_dihedrals()\n",
    "    def add_angles(self):\n",
    "        for i in range(1, self.get_number()+1):\n",
    "            center = self.atoms[i]\n",
    "            if len(center.neighbors) > 1:\n",
    "                for index1, index2 in itertools.combinations(center.neighbors, 2):\n",
    "                    self.angles.append((min(index1, index2), center.index, max(index1, index2)))\n",
    "                    \n",
    "    def add_dihedrals(self):\n",
    "        '''Only proper dihedrals are needed for this series of compounds. Add dihedrals according to the\n",
    "        bond information.\n",
    "        '''\n",
    "        for bond in self.bonds:\n",
    "            for i, j in itertools.product(self.atoms[bond[0]].neighbors, self.atoms[bond[1]].neighbors):\n",
    "                if i != bond[1] and j != bond[0] and i != j: # avoid triangle loop like cyclopropane\n",
    "                    self.dihedrals.append((i, bond[0], bond[1], j))\n",
    "    def write_itp(self, para, filename = 'TAA.itp'):\n",
    "        '''write all atoms to an itp file'''\n",
    "        myfile = open(filename, 'w')\n",
    "        myfile.write('[ moleculetype ]\\n')\n",
    "        myfile.write('%s\\t3\\n' %self.name)\n",
    "        \n",
    "        myfile.write('[ atoms ]\\n')\n",
    "        for i in range(1, self.get_number()+1):\n",
    "            atomname = self.atoms[i].name\n",
    "            myfile.write('%5d%5s%5d%5s%5s%5d%12.4f%12.4f\\n' %(i, atomname, 1, self.name, atomname[0]+str(i),\n",
    "                                        self.atoms[i].cgroup, para[atomname][0], para[atomname][1]))\n",
    "        \n",
    "        myfile.write('[ bonds ]\\n')\n",
    "        for bond in self.bonds:\n",
    "            myfile.write('%5d%5d%5d\\n' %(bond+(1,)))\n",
    "        \n",
    "        myfile.write('[ angles ]\\n')\n",
    "        for angle in self.angles:\n",
    "            myfile.write('%5d%5d%5d%5d\\n' %(angle+(1,)))\n",
    "            \n",
    "        myfile.write('[ dihedrals ]\\n')\n",
    "        for dih in self.dihedrals:\n",
    "            myfile.write('%5d%5d%5d%5d%5d\\n' %(dih+(1,)))\n",
    "        myfile.close()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%save gen_TAA.py 1-1000\n",
    "para = {'N3': (0.12, 14.007),\n",
    "        'C1': (-0.17, 12.011),\n",
    "        'H1': (0.13, 1.008),\n",
    "        'C2': (0.01, 12.011),\n",
    "        'CS': (-0.12, 12.011),\n",
    "        'CT': (-0.18, 12.011),\n",
    "        'HC': (0.06, 1.008)\n",
    "       }\n",
    "\n",
    "Atom.clear_index()\n",
    "TEA = TAA('TPA')\n",
    "TEA.add_alkyl(3)\n",
    "\n",
    "TEA.write_itp(para)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}