Merge pull request #9 from mattwthompson/typing-abstraction

Typing abstraction

.gitignore

@@ -92,6 +92,9 @@ ENV/
 .spyderproject
 .spyproject
 
+# PyCharm
+.idea
+
 # Rope project settings
 .ropeproject
 
@@ -100,3 +103,5 @@ ENV/
 
 # mypy
 .mypy_cache/
+
+scratch/

devtools/conda-envs/test_env.yaml

@@ -7,6 +7,9 @@ dependencies:
     # Base depends
   - python
   - pip
+  - pydantic
+  - pint
+  - sympy
 
     # Testing
   - pytest
@@ -15,3 +18,4 @@ dependencies:
   - sympy
   - pint
   - openmm
+  - openforcefield

examples/README.md

@@ -0,0 +1,3 @@
+### Examples
+
+Jupyter notebooks demonstrating use cases

examples/smirnoff_argon.ipynb

@@ -0,0 +1,361 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openforcefield.typing.engines.smirnoff import ForceField\n",
+    "from openforcefield.topology.molecule import Molecule\n",
+    "from openforcefield.topology.topology import Topology\n",
+    "\n",
+    "from system.typing.smirnoff import build_smirnoff_map, build_smirnoff_collection\n",
+    "from system.collections import PotentialHandler, PotentialCollection\n",
+    "from system.system import System\n",
+    "from system.utils import get_test_file_path"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load in a minimal SMIRNOFF forcefield and Argon topology\n",
+    "argon_ff = ForceField(get_test_file_path('ar.offxml'))\n",
+    "\n",
+    "mol = Molecule.from_smiles('[#18]')\n",
+    "mol.generate_conformers(n_conformers=1)\n",
+    "\n",
+    "argon_top =  Topology.from_molecules(10 * [mol])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'vdW': {(0,): '[#18:1]',\n",
+       "  (1,): '[#18:1]',\n",
+       "  (2,): '[#18:1]',\n",
+       "  (3,): '[#18:1]',\n",
+       "  (4,): '[#18:1]',\n",
+       "  (5,): '[#18:1]',\n",
+       "  (6,): '[#18:1]',\n",
+       "  (7,): '[#18:1]',\n",
+       "  (8,): '[#18:1]',\n",
+       "  (9,): '[#18:1]'}}"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Generate a mapping of topology slots and SMIRNOFF SMIRKS\n",
+    "smirks_map = build_smirnoff_map(forcefield=argon_ff, topology=argon_top)\n",
+    "# {\n",
+    "#   handler_name1: {\n",
+    "#     {\n",
+    "#       slot1: SMIRKS1,\n",
+    "#       slot2: SMIRKS2,\n",
+    "#       slot3: SMIRKS3,\n",
+    "#       ...,\n",
+    "#       slotN: SMIRKSN,\n",
+    "#     },\n",
+    "#   handler_name2: { ... },\n",
+    "#   handler_name3: { ... },\n",
+    "#   ...\n",
+    "#   handler_nameN: { ... },\n",
+    "# }\n",
+    "smirks_map"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "smirnoff_collection = build_smirnoff_collection(forcefield=argon_ff)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "argon_system = System(\n",
+    "    potential_collection=smirnoff_collection,\n",
+    "    topology=argon_top,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "openforcefield.topology.topology.Topology"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# The topology attribute simply stores the toolkit topology\n",
+    "type(argon_system.topology)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "system.collections.PotentialCollection"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# The potential_collection is effectively a force field for SMIRNOFF systems,\n",
+    "# but is more flexible for future typing schemes that are much less \"1:1\" than\n",
+    "# traiditional atom-typing or SMIRNOFF, which generally map 1 parameter to a slot\n",
+    "type(argon_system.potential_collection)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "dict_keys(['vdW'])"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# A PotentialCollection object stores \"types\" of potentials, i.e. one section for vdW,\n",
+    "# one section for bonds, etc. in order to separate parameters from instructions; defining\n",
+    "# a SMIRKS pattern and a slot is not sufficient instructions to apply the corresponding \n",
+    "# parameter. This system only has vdW, but this dict is intended to store other non-bonded\n",
+    "# terms, valence terms, wild cross-coupling terms, etc.\n",
+    "argon_system.potential_collection.handlers.keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\\[0.3\\ nanometer\\]"
+      ],
+      "text/latex": [
+       "$0.3\\ \\mathrm{nanometer}$"
+      ],
+      "text/plain": [
+       "0.3 <Unit('nanometer')>"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# A PotentialCollection is more or less nested dictionaries, so you can drill down\n",
+    "# to find the value of a particular parameter in a potential, returning a \n",
+    "# unit-bearing quantity (via pint).\n",
+    "argon_system.potential_collection['vdW'].potentials['[#18:1]'].parameters['sigma']\n",
+    "# The path is convoluted, but goes something like this\n",
+    "# sys_name.potential_collection['vdW'].potentials[smirks].parameters['sigma']\n",
+    "#    ^                            ^                 ^                   ^\n",
+    "#    |                            |                 |                   | \n",
+    "#    |                            |                 |                   | \n",
+    "#    |                \"type\" (?) of potential       |     finally, the actual param\n",
+    "#    |                                              |\n",
+    "# system I care about                SMIRKS key, mapping to a potential"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ParametrizedAnalyticalPotential(name='n1', smirks='[#18:1]', expression='4*epsilon*((sigma/r)**12-(sigma/r)**6)', independent_variables={'r'}, parameters={'sigma': <Quantity(0.3, 'nanometer')>, 'epsilon': <Quantity(0.1, 'kilojoule / mole')>})"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# There's other information included in each potential, notably the\n",
+    "# analytical expression of the potential and the SMIRKS\n",
+    "argon_system.potential_collection['vdW'].potentials['[#18:1]']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{(0,): '[#18:1]',\n",
+       " (1,): '[#18:1]',\n",
+       " (2,): '[#18:1]',\n",
+       " (3,): '[#18:1]',\n",
+       " (4,): '[#18:1]',\n",
+       " (5,): '[#18:1]',\n",
+       " (6,): '[#18:1]',\n",
+       " (7,): '[#18:1]',\n",
+       " (8,): '[#18:1]',\n",
+       " (9,): '[#18:1]'}"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Going back to the smirks_map, which is a dictionary mapping \"slots\" to potentials\n",
+    "# In the simple case, this is just atom indicies : SMKIRKS\n",
+    "smirks_map['vdW']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ParametrizedAnalyticalPotential(name='n1', smirks='[#18:1]', expression='4*epsilon*((sigma/r)**12-(sigma/r)**6)', independent_variables={'r'}, parameters={'sigma': <Quantity(0.3, 'nanometer')>, 'epsilon': <Quantity(0.1, 'kilojoule / mole')>})"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# The smirks_map and potential_collection can be used to look up potentials\n",
+    "# using the 'vdW' and SMIRKS as sufficient information\n",
+    "argon_system.potential_collection['vdW'][smirks_map['vdW'][(7,)]]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'topology': <openforcefield.topology.topology.Topology at 0x11f658f50>,\n",
+       " 'potential_collection': {'handlers': {'vdW': {'name': 'vdW',\n",
+       "    'potentials': {'[#18:1]': {'name': 'n1',\n",
+       "      'smirks': '[#18:1]',\n",
+       "      'expression': '4*epsilon*((sigma/r)**12-(sigma/r)**6)',\n",
+       "      'independent_variables': {'r'},\n",
+       "      'parameters': {'sigma': 0.3 <Unit('nanometer')>,\n",
+       "       'epsilon': 0.1 <Unit('kilojoule / mole')>}}}}}},\n",
+       " 'positions': None,\n",
+       " 'box': None,\n",
+       " 'slots_map': None}"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Being built off of pydantic, we get a serializable representation for free,\n",
+    "# caveats being that some components are not actually serializable now\n",
+    "argon_system.dict()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'name': 'n1',\n",
+       " 'smirks': '[#18:1]',\n",
+       " 'expression': '4*epsilon*((sigma/r)**12-(sigma/r)**6)',\n",
+       " 'independent_variables': {'r'},\n",
+       " 'parameters': {'sigma': 0.3 <Unit('nanometer')>,\n",
+       "  'epsilon': 0.1 <Unit('kilojoule / mole')>}}"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Being a nested dictionary, you could also grab a component and make that a dict\n",
+    "argon_system.potential_collection.handlers['vdW'].potentials['[#18:1]'].dict()"
+   ]
+  }
+ ],
+ "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.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}