diff --git a/.gitignore b/.gitignore index bd40dbd8a..8b6f1c54d 100644 --- a/.gitignore +++ b/.gitignore @@ -778,4 +778,7 @@ FodyWeavers.xsd .vscode/** # acpype cache -*.acpype/ \ No newline at end of file +*.acpype/ + +*/_date.py +*/_version.py \ No newline at end of file diff --git a/dmff/api.py b/dmff/api.py index dfd7fa6de..713a76951 100644 --- a/dmff/api.py +++ b/dmff/api.py @@ -4,6 +4,7 @@ from collections import defaultdict import xml.etree.ElementTree as ET from copy import deepcopy +import warnings import numpy as np import jax.numpy as jnp @@ -1189,10 +1190,11 @@ def __init__(self, hamiltonian): def registerBondType(self, bond): typetxt = findAtomTypeTexts(bond, 2) types = self.ff._findAtomTypes(bond, 2) - self.types.append(types) - self.typetexts.append(typetxt) - self.params["k"].append(float(bond["k"])) - self.params["length"].append(float(bond["length"])) # length := r0 + if None not in types: + self.types.append(types) + self.typetexts.append(typetxt) + self.params["k"].append(float(bond["k"])) + self.params["length"].append(float(bond["length"])) # length := r0 @staticmethod def parseElement(element, hamiltonian): @@ -1287,9 +1289,10 @@ def __init__(self, hamiltonian): def registerAngleType(self, angle): types = self.ff._findAtomTypes(angle, 3) - self.types.append(types) - self.params["k"].append(float(angle["k"])) - self.params["angle"].append(float(angle["angle"])) + if None not in types: + self.types.append(types) + self.params["k"].append(float(angle["k"])) + self.params["angle"].append(float(angle["angle"])) @staticmethod def parseElement(element, hamiltonian): @@ -1302,8 +1305,12 @@ def parseElement(element, hamiltonian): <\HarmonicAngleForce> """ - generator = HarmonicAngleJaxGenerator(hamiltonian) - hamiltonian.registerGenerator(generator) + existing = [f for f in hamiltonian._forces if isinstance(f, HarmonicAngleJaxGenerator)] + if len(existing) == 0: + generator = HarmonicAngleJaxGenerator(hamiltonian) + hamiltonian.registerGenerator(generator) + else: + generator = existing[0] for angletype in element.findall("Angle"): generator.registerAngleType(angletype.attrib) @@ -1342,7 +1349,7 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args): n_angles += 1 break if not ifFound: - print( + warnings.warn( "No parameter for angle %i - %i - %i" % (idx1, idx2, idx3) ) @@ -1994,7 +2001,7 @@ def renderXML(self): class NonbondJaxGenerator: - SCALETOL = 1e-5 + SCALETOL = 1e-3 def __init__(self, hamiltionian, coulomb14scale, lj14scale): @@ -2019,7 +2026,8 @@ def __init__(self, hamiltionian, coulomb14scale, lj14scale): def registerAtom(self, atom): # use types in nb cards or resname+atomname in residue cards types = self.ff._findAtomTypes(atom, 1)[0] - self.types.append(types) + if None not in types: + self.types.append(types) for key in ["sigma", "epsilon", "charge"]: if key not in self.useAttributeFromResidue: @@ -2065,11 +2073,6 @@ def parseElement(element, ff): generator.n_atoms = len(element.findall("Atom")) - # jax it! - for k in generator.params.keys(): - generator.params[k] = jnp.array(generator.params[k]) - generator.types = np.array(generator.types) - def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args): methodMap = { app.NoCutoff: "NoCutoff", diff --git a/dmff/classical/inter.py b/dmff/classical/inter.py index 83a62d366..d7c16317f 100644 --- a/dmff/classical/inter.py +++ b/dmff/classical/inter.py @@ -41,11 +41,6 @@ def get_LJ_energy(dr_vec, sig, eps, box): if self.ifPBC: dr_vec = v_pbc_shift(dr_vec, box, jnp.linalg.inv(box)) dr_norm = jnp.linalg.norm(dr_vec, axis=1) - if not self.ifNoCut: - msk = dr_norm <= self.r_cut - sig = sig[msk] - eps = eps[msk] - dr_norm = dr_norm[msk] dr_inv = 1.0 / dr_norm sig_dr = sig * dr_inv diff --git a/examples/classical/demo.ipynb b/examples/classical/demo.ipynb new file mode 100644 index 000000000..0de988b5b --- /dev/null +++ b/examples/classical/demo.ipynb @@ -0,0 +1,221 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Classical Force Field in DMFF" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "DMFF implements classcial molecular mechanics force fields with the following forms:\n", + "\n", + "$$\\begin{align*}\n", + " V(\\mathbf{R}) &= V_{\\mathrm{bond}} + V_{\\mathrm{angle}} + V_\\mathrm{torsion} + V_\\mathrm{vdW} + V_\\mathrm{Coulomb} \\\\\n", + " &= \\sum_{\\mathrm{bonds}}\\frac{1}{2}k_b(r - r_0)^2 + \\sum_{\\mathrm{angles}}\\frac{1}{2}k_\\theta (\\theta -\\theta_0)^2 + \\sum_{\\mathrm{torsion}}\\sum_{n=1}^4 V_n[1+\\cos(n\\phi - \\phi_s)] \\\\\n", + " &\\quad+ \\sum_{ij}4\\varepsilon_{ij}\\left[\\left(\\frac{\\sigma_{ij}}{r_{ij}}\\right)^{12} - \\left(\\frac{\\sigma_{ij}}{r_{ij}}\\right)^6\\right] + \\sum_{ij}\\frac{q_iq_j}{4\\pi\\varepsilon_0r_{ij}}\n", + "\\end{align*}$$" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Import necessary packages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import jax\n", + "import jax.numpy as jnp\n", + "import openmm.app as app\n", + "import openmm.unit as unit\n", + "from dmff import Hamiltonian, NeighborList" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Compute energy" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "DMFF uses **OpenMM** to parse input files, including coordinates files, topology specification files. Class `Hamiltonian` inherited from `openmm.ForceField` will be initialized and used to parse force field parameters in XML format. Take parametrzing an organic moleclue with GAFF2 force field as an example.\n", + "\n", + "- `lig_top.xml`: Define bond connections (topology). Not necessary if such information is specified in pdb with `CONNECT` keyword.\n", + "- `gaff-2.11.xml`: GAFF2 force field parameters: bonds, angles, torsions and vdW params\n", + "- `lig-prm.xml`: Atomic charges" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "app.Topology.loadBondDefinitions(\"lig-top.xml\")\n", + "pdb = app.PDBFile(\"lig.pdb\")\n", + "ff = Hamiltonian(\"gaff-2.11.xml\", \"lig-prm.xml\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The method `Hamiltonian.createPotential` will be called to create differentiable potential energy functions for different energy terms. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "potentials = ff.createPotential(\n", + " pdb.topology,\n", + " nonbondedMethod=app.NoCutoff\n", + ")\n", + "for pot in potentials:\n", + " print(pot)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The force field parameters are stored as a Python dict in the `param` attribute of force generators." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "nbparam = ff.getGenerators()[3].params\n", + "nbparam[\"charge\"] # also \"epsilon\", \"sigma\" etc. keys" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each generated function will read **coordinates, box, pairs** and force field parameters as inputs. `pairs` is a integer array in which each row specifying atoms condsidered as neighbors within rcut. This can be calculated with `dmff.NeighborList` class which is supported by `jax_md`.\n", + "\n", + "The potential energy function will give energy (a scalar, in kJ/mol) as output:\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "positions = jnp.array(pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))\n", + "box = jnp.array([\n", + " [10.0, 0.0, 0.0], \n", + " [0.0, 10.0, 0.0],\n", + " [0.0, 0.0, 10.0]\n", + "])\n", + "nbList = NeighborList(box, rc=4)\n", + "nbList.allocate(positions)\n", + "pairs = nbList.pairs\n", + "nbfunc = potentials[-1]\n", + "energy = nbfunc(positions, box, pairs, ff.getGenerators()[-1].params)\n", + "print(energy)\n", + "print(pairs)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can also obtain the whole potential energy function and force field parameter set, instead of seperated functions for different energy terms." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "efunc = ff.getPotentialFunc()\n", + "params = ff.getParameters()\n", + "totene = efunc(positions, box, pairs, params)\n", + "totene" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Compute forces and parametric gradients" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use `jax.grad` to compute forces and parametric gradients automatically" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pos_grad_func = jax.grad(efunc, argnums=0)\n", + "force = -pos_grad_func(positions, box, pairs, params)\n", + "force.shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "param_grad_func = jax.grad(nbfunc, argnums=-1)\n", + "pgrad = param_grad_func(positions, box, pairs, nbparam)\n", + "pgrad[\"charge\"]" + ] + } + ], + "metadata": { + "interpreter": { + "hash": "44fe82502fda871be637af1aa98d2b3ddaac01204dd30f1519cbec4e95000815" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/fluctuated_leading_term_waterff/demo.ipynb b/examples/fluctuated_leading_term_waterff/demo.ipynb new file mode 100644 index 000000000..8453e1da0 --- /dev/null +++ b/examples/fluctuated_leading_term_waterff/demo.ipynb @@ -0,0 +1,274 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "bc65a9bf-b20f-473f-ab61-e74e6aaaf48e", + "metadata": {}, + "source": [ + "# Mutipolar polarizable force field with fluctuating charges" + ] + }, + { + "cell_type": "markdown", + "id": "9cccf29a", + "metadata": {}, + "source": [ + "In this demo, we show how to implement a **multipolar polarizable potential with fluctuating charges** with DMFF API.\n", + "\n", + "In conventional models, atomic charges are pre-defined and remain unchanged during the simulation. Here, we want to implement a model that considers atomic charges as *conformer-dependent*, so that the charges can vary during a molecular dynamics simulation. This will give better description of the system's behavior.\n", + "\n", + "## System preparation\n", + "Load the coordinates, box and compute neighbor list. Note that conventionally in multipolar polarizable models, the length unit is **angstrom**." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db18fde6-94bf-4e58-8b18-de85d5f15c6a", + "metadata": {}, + "outputs": [], + "source": [ + "import jax\n", + "import jax.numpy as jnp\n", + "import openmm.app as app\n", + "import openmm.unit as unit\n", + "from dmff import Hamiltonian, NeighborList\n", + "\n", + "app.Topology.loadBondDefinitions(\"residues.xml\")\n", + "pdb = app.PDBFile(\"water_dimer.pdb\")\n", + "rc = 4 # cutoff, in angstrom\n", + "positions = jnp.array(pdb.getPositions(asNumpy=True).value_in_unit(unit.angstrom))\n", + "box = jnp.array(\n", + " [vec.value_in_unit(unit.angstrom) for vec in pdb.topology.getPeriodicBoxVectors()]\n", + ")\n", + "nbList = NeighborList(box, rc=rc)\n", + "nbList.allocate(positions)\n", + "pairs = nbList.pairs" + ] + }, + { + "cell_type": "markdown", + "id": "43a9afc9-e01c-49e5-9ddc-006fd44a3a31", + "metadata": {}, + "source": [ + "## Genearate auto-differentiable multipolar polarizable (ADMP) forces" + ] + }, + { + "cell_type": "markdown", + "id": "cb817059-1b3a-42a5-a494-02fcabe6bc75", + "metadata": {}, + "source": [ + "First, we will use the `dmff` to create a multipolar polarizable potential with **fixed** atomic charges.\n", + "\n", + "Here, we have two types of force: \n", + "\n", + "- Dispersion force\n", + "- Multipolar polarizable PME force.\n", + "\n", + "We will focus on the PME force." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2dad4351", + "metadata": {}, + "outputs": [], + "source": [ + "H = Hamiltonian('forcefield.xml')\n", + "disp_pot, pme_pot = H.createPotential(pdb.topology, nonbondedCutoff=rc*unit.angstrom, step_pol=5)\n", + "disp_generator, pme_generator = H.getGenerators()\n", + "print(pme_generator)\n", + "print(pme_pot)\n", + "pme_generator.params" + ] + }, + { + "cell_type": "markdown", + "id": "97892dfe", + "metadata": {}, + "source": [ + "The function `pme_pot` takes the following actions:\n", + "\n", + "- Expand **force field parameters** (oxygen and hydrogen charges, polarizabilites, etc.) pre-defined in `forcefield.xml` to each atom, which we called **atomic parameters**\n", + "- Calls the real PME kernel function to evaluate energy.\n", + "\n", + "The force field parameters are stored in `pme_generator.params`. And the expansion is implemented with the *broadcast* feature of `jax.numpy.ndarray`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33a6af99-886b-4934-97fd-6e8b4fa0ecbf", + "metadata": {}, + "outputs": [], + "source": [ + "params = pme_generator.params[\"Q_local\"]\n", + "params[pme_generator.map_atomtype]" + ] + }, + { + "cell_type": "markdown", + "id": "f78b9b80-fdfb-45ba-90a4-b09475c5feff", + "metadata": {}, + "source": [ + "## Implement fluctuating charges" + ] + }, + { + "cell_type": "markdown", + "id": "a17ba838", + "metadata": {}, + "source": [ + "Since this expansion process is done internally within `pme_pot`, it is **not flexible** enough for us to specify atom-specific charges, i.e. **fluctuating charges**. \n", + "\n", + "As a result, we must re-write `pme_pot` to enable modifying the atomic charges after force field parameter expansion. \n", + "\n", + "Benifiting from the flexible APIs in DMFF, we will reuse most of the functions and variables in the `pme_generator`, only modify charges in the input parameters, i.e. the `Q_local` argument in `pme_generator.pme_force.get_energy` function." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bae85400", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "from dmff.utils import jit_condition\n", + "from dmff.admp.pme import trim_val_0\n", + "from dmff.admp.spatial import v_pbc_shift\n", + "\n", + "\n", + "@jit_condition(static_argnums=())\n", + "def compute_leading_terms(positions, box):\n", + " n_atoms = len(positions)\n", + " c0 = jnp.zeros(n_atoms)\n", + " c6_list = jnp.zeros(n_atoms)\n", + " box_inv = jnp.linalg.inv(box)\n", + " O = positions[::3]\n", + " H1 = positions[1::3]\n", + " H2 = positions[2::3]\n", + " ROH1 = H1 - O\n", + " ROH2 = H2 - O\n", + " ROH1 = v_pbc_shift(ROH1, box, box_inv)\n", + " ROH2 = v_pbc_shift(ROH2, box, box_inv)\n", + " dROH1 = jnp.linalg.norm(ROH1, axis=1)\n", + " dROH2 = jnp.linalg.norm(ROH2, axis=1)\n", + " costh = jnp.sum(ROH1 * ROH2, axis=1) / (dROH1 * dROH2)\n", + " angle = jnp.arccos(costh) * 180 / jnp.pi\n", + " dipole = -0.016858755 + 0.002287251 * angle + 0.239667591 * dROH1 + (-0.070483437) * dROH2\n", + " charge_H = dipole / dROH1\n", + " charge_O = charge_H * (-2)\n", + " C6_H = (-2.36066199 + (-0.007049238) * angle + 1.949429648 * dROH1+ 2.097120784 * dROH2) * 0.529**6 * 2625.5\n", + " C6_O = (-8.641301261 + 0.093247893 * angle + 11.90395358 * (dROH1+ dROH2)) * 0.529**6 * 2625.5\n", + " C6_H = trim_val_0(C6_H)\n", + " c0 = c0.at[::3].set(charge_O)\n", + " c0 = c0.at[1::3].set(charge_H)\n", + " c0 = c0.at[2::3].set(charge_H)\n", + " c6_list = c6_list.at[::3].set(jnp.sqrt(C6_O))\n", + " c6_list = c6_list.at[1::3].set(jnp.sqrt(C6_H))\n", + " c6_list = c6_list.at[2::3].set(jnp.sqrt(C6_H))\n", + " return c0, c6_list\n", + "\n", + "\n", + "def generate_calculator(pme_generator):\n", + " def admp_calculator(positions, box, pairs):\n", + " c0, c6_list = compute_leading_terms(positions,box) # compute fluctuated charges\n", + " Q_local = pme_generator.params[\"Q_local\"][pme_generator.map_atomtype]\n", + " Q_local = Q_local.at[:,0].set(c0) # change fixed charge into fluctuated one\n", + " pol = pme_generator.params[\"pol\"][pme_generator.map_atomtype]\n", + " tholes = pme_generator.params[\"tholes\"][pme_generator.map_atomtype]\n", + " mScales = pme_generator.params[\"mScales\"]\n", + " pScales = pme_generator.params[\"pScales\"]\n", + " dScales = pme_generator.params[\"dScales\"]\n", + " E_pme = pme_generator.pme_force.get_energy(\n", + " positions, \n", + " box, \n", + " pairs, \n", + " Q_local, \n", + " pol, \n", + " tholes, \n", + " mScales, \n", + " pScales, \n", + " dScales\n", + " )\n", + " return E_pme \n", + " return jax.jit(admp_calculator)\n" + ] + }, + { + "cell_type": "markdown", + "id": "9aac71bd-c759-4423-8441-b802ea213722", + "metadata": {}, + "source": [ + "**Finally, compute the energy and force!**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "560a8b7b-1dd7-4368-8154-99240aed9d81", + "metadata": {}, + "outputs": [], + "source": [ + "potential_fn = generate_calculator(pme_generator)\n", + "ene = potential_fn(positions, box, pairs)\n", + "print(ene)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "632d4024-76a2-4b84-a7a4-e59eecdfde20", + "metadata": {}, + "outputs": [], + "source": [ + "force_fn = jax.grad(potential_fn, argnums=(0))\n", + "force = -force_fn(positions, box, pairs)\n", + "print(force)" + ] + }, + { + "cell_type": "markdown", + "id": "c74539d4", + "metadata": {}, + "source": [ + "The running speed of the first pass is slow because JAX is trying to track the data flow and compile the code. Once the code is compiled, it is run much faster, until the shapes of the input parameters change, trigerring a recompilation. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d7a0fd1b-a363-40a0-9017-febe0fe76f01", + "metadata": {}, + "outputs": [], + "source": [ + "print(-force_fn(positions, box, pairs))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/fluctuated_leading_term_waterff/fluctuated_water_ff.ipynb b/examples/fluctuated_leading_term_waterff/fluctuated_water_ff.ipynb deleted file mode 100644 index 0fa7dedec..000000000 --- a/examples/fluctuated_leading_term_waterff/fluctuated_water_ff.ipynb +++ /dev/null @@ -1,384 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "9cccf29a", - "metadata": {}, - "source": [ - "In this demo we show how to implement a multipolar polarizable potential with fluctuating charges.\n", - "First, we will use the DMFF API to create a multipolar polarizable potential with fixed atomic charges:\n", - "\n", - "Read input pdb and xml files, create the Hamiltonian object H." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "28b869ac", - "metadata": {}, - "outputs": [], - "source": [ - "import sys\n", - "import jax\n", - "import jax.numpy as jnp\n", - "import openmm.app as app\n", - "import openmm.unit as unit\n", - "from dmff.api import Hamiltonian\n", - "\n", - "H = Hamiltonian('forcefield.xml')\n", - "app.Topology.loadBondDefinitions(\"residues.xml\")\n", - "pdb = app.PDBFile(\"water_dimer.pdb\")\n", - "rc = 4" - ] - }, - { - "cell_type": "markdown", - "id": "f8bf6f57", - "metadata": {}, - "source": [ - "For each force component, the Hamiltonian will create a generator. Here, we have two components: dispersion force and multipolar polarizable pme force. We will focus on the pme force. In the generator, you can find the force field parameters (charges, multipoles) stored in generator.params:" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "2dad4351", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "dict_keys(['mScales', 'pScales', 'dScales', 'Q_local', 'pol', 'tholes'])\n" - ] - } - ], - "source": [ - "disp_generator, pme_generator = H.getGenerators()\n", - "print(pme_generator.params.keys())" - ] - }, - { - "cell_type": "markdown", - "id": "0ea2ec1a", - "metadata": {}, - "source": [ - "Then use the `createPotential` function to create potential functions. When `createPotentials` is invoked, it will use the topology information of the system to create potential functions, and creates many useful information (e.g., local axis definitions) as well. This information is stored in generator.pme_force and will be used later" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "18d5650f", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "True\n", - "[1 0 0 1 0 0]\n", - "[[ 1 2 -1]\n", - " [ 0 2 -1]\n", - " [ 0 1 -1]\n", - " [ 4 5 -1]\n", - " [ 3 5 -1]\n", - " [ 3 4 -1]]\n", - "5\n" - ] - } - ], - "source": [ - "potentials = H.createPotential(pdb.topology, nonbondedCutoff=rc*unit.angstrom, step_pol=5)\n", - "pot_disp = potentials[-2]\n", - "pot_pme = potentials[-1]\n", - "\n", - "print(pme_generator.pme_force.lpol)\n", - "print(pme_generator.pme_force.axis_type)\n", - "print(pme_generator.pme_force.axis_indices)\n", - "print(pme_generator.pme_force.steps_pol)" - ] - }, - { - "cell_type": "markdown", - "id": "97892dfe", - "metadata": {}, - "source": [ - "What `pot_pme` actually does is to take force field parameters (i.e., oxygen and hydrogen charges), and expand it into atomic parameterss (i.e., charge of each atom), then calls the real pme kernel function to evaluate energy.\n", - "\n", - "Therefore, the input interface of `pot_pme` is not flexible enough to specify atom-specific charges, which is what we intend to do in this demo. So we will disregard `pot_pme`, but to wrap the real pme kernel directly in a way different to `pot_pme`. The pme kernel function can be accessed via generator:" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "39f6d776", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Help on function get_energy in module dmff.admp.pme:\n", - "\n", - "get_energy(positions, box, pairs, Q_local, pol, tholes, mScales, pScales, dScales, U_init=DeviceArray([[0., 0., 0.],\n", - " [0., 0., 0.],\n", - " [0., 0., 0.],\n", - " [0., 0., 0.],\n", - " [0., 0., 0.],\n", - " [0., 0., 0.]], dtype=float32))\n", - " # this is the wrapper that include a Uind optimizer\n", - "\n" - ] - } - ], - "source": [ - "help(pme_generator.pme_force.get_energy)" - ] - }, - { - "cell_type": "markdown", - "id": "a17ba838", - "metadata": {}, - "source": [ - "Let's wrap it using the following function:" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "bae85400", - "metadata": { - "scrolled": true - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - ".admp_calculator at 0x7f59d034e830>\n" - ] - } - ], - "source": [ - "from run import compute_leading_terms\n", - "\n", - "def generate_calculator(get_energy_pme, pme_generator):\n", - " def admp_calculator(positions, box, pairs):\n", - " c0, c6_list = compute_leading_terms(positions,box) # compute fluctuated charges\n", - " Q_local = pme_generator.params[\"Q_local\"][pme_generator.map_atomtype]\n", - " Q_local = Q_local.at[:,0].set(c0) # change fixed charge into fluctuated one\n", - " pol = pme_generator.params[\"pol\"][pme_generator.map_atomtype]\n", - " tholes = pme_generator.params[\"tholes\"][pme_generator.map_atomtype]\n", - " mScales = pme_generator.params[\"mScales\"]\n", - " pScales = pme_generator.params[\"pScales\"]\n", - " dScales = pme_generator.params[\"dScales\"]\n", - " E_pme = pme_generator.pme_force.get_energy(\n", - " positions, \n", - " box, \n", - " pairs, \n", - " Q_local, \n", - " pol, \n", - " tholes, \n", - " mScales, \n", - " pScales, \n", - " dScales\n", - " )\n", - " return E_pme \n", - " return admp_calculator\n", - "\n", - "potential_fn = generate_calculator(pme_generator.pme_force.get_energy, pme_generator)\n", - "print(potential_fn)" - ] - }, - { - "cell_type": "markdown", - "id": "6171bac9", - "metadata": {}, - "source": [ - "Note in here, we use a function called `compute_leading_terms` to compute the atomic charges, and feed it into `pme_generator.pme_force.get_energy`, so geometry-dependent charge fluctuations are realized.\n", - "\n", - "Let's see how the wrapped `potential_fn` works:" - ] - }, - { - "cell_type": "markdown", - "id": "bddf9b49", - "metadata": {}, - "source": [ - "First, prepare the inputs" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "8680d936", - "metadata": {}, - "outputs": [], - "source": [ - "positions = jnp.array(pdb.positions._value) * 10\n", - "a, b, c = pdb.topology.getPeriodicBoxVectors()\n", - "box = jnp.array([a._value, b._value, c._value]) * 10" - ] - }, - { - "cell_type": "markdown", - "id": "94c4f89e", - "metadata": {}, - "source": [ - "Consturct neighbor list " - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "bb2af638", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[[0 1]\n", - " [0 2]\n", - " [1 2]\n", - " [0 3]\n", - " [1 3]\n", - " [2 3]\n", - " [0 4]\n", - " [1 4]\n", - " [2 4]\n", - " [3 4]\n", - " [0 5]\n", - " [1 5]\n", - " [2 5]\n", - " [3 5]\n", - " [4 5]\n", - " [6 6]\n", - " [6 6]\n", - " [6 6]]\n" - ] - } - ], - "source": [ - "import dmff.common.nblist as nblist\n", - "nbl = nblist.NeighborList(box, rc)\n", - "pairs = nbl.allocate(positions).idx.T\n", - "print(pairs)" - ] - }, - { - "cell_type": "markdown", - "id": "698c029e", - "metadata": {}, - "source": [ - "Call the potential function" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "f74f9f6c", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "-41.352478\n" - ] - } - ], - "source": [ - "E = potential_fn(positions, box, pairs)\n", - "print(E)" - ] - }, - { - "cell_type": "markdown", - "id": "6e59f678", - "metadata": {}, - "source": [ - "You can use the auto differentiation function in jax to compute force:" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "131a8f3b", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[[ 7.6653423 -11.747272 7.999815 ]\n", - " [-75.19945 58.282196 25.169626 ]\n", - " [ 1.8799052 4.969598 -14.643274 ]\n", - " [ 67.67523 -38.337395 -20.55421 ]\n", - " [ 2.5347214 5.280134 -4.1624823]\n", - " [ -4.551193 -18.451092 6.1247015]]\n" - ] - } - ], - "source": [ - "force_fn = jax.jit(jax.grad(potential_fn, argnums=(0)))\n", - "F = force_fn(positions, box, pairs)\n", - "print(F)" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "id": "e723ed6f", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[[ 7.665777 -11.743576 8.00062 ]\n", - " [-75.19835 58.281765 25.170254 ]\n", - " [ 1.8792243 4.9697785 -14.642523 ]\n", - " [ 67.675674 -38.338062 -20.55489 ]\n", - " [ 2.5349426 5.279976 -4.1608496]\n", - " [ -4.5494537 -18.451893 6.1248045]]\n" - ] - } - ], - "source": [ - "F = force_fn(positions, box, pairs)\n", - "print(F)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7187faa9", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "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.7" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/setup.py b/setup.py index e8ef2b134..ccfb42817 100644 --- a/setup.py +++ b/setup.py @@ -1,22 +1,58 @@ -from setuptools import setup, find_packages +from os import path +import setuptools +import datetime +from setuptools import find_packages + +NAME = "dmff" + +readme_file = path.join(path.dirname(path.abspath(__file__)), 'README.md') +try: + from m2r import parse_from_file + + readme = parse_from_file(readme_file) +except ImportError: + with open(readme_file) as f: + readme = f.read() + +today = datetime.date.today().strftime("%b-%d-%Y") +with open(path.join('dmff', '_date.py'), 'w') as fp: + fp.write('date = \'%s\'' % today) install_requires = [ "numpy", - "jax_md", + "jax_md>=0.1.28", "openmm", ] -setup( - name='dmff', - version='0.0.1', - author='dptech.net', - # author_email='hermite@dptech.net', - description=('DMFF.'), - url='https://github.com/deepmodeling/DMFF', - license=None, - keywords='Differentiable', - install_requires=install_requires, - packages=find_packages(), - zip_safe=False, - #packages=packages, - entry_points={}, - include_package_data=True) \ No newline at end of file + + +def setup(scm=None): + packages = find_packages(exclude=["tests"]) + + setuptools.setup( + name=NAME, + use_scm_version=scm, + setup_requires=['setuptools_scm'], + author="DeepModeling", + author_email="windwhisper.yu@gmail.com", + description="Differentiable Molecular Force Field", + long_description=readme, + long_description_content_type="text/markdown", + url="https://github.com/deepmodeling/DMFF", + python_requires="~=3.6", + packages=packages, + data_files=[], + package_data={}, + classifiers=[ + "Programming Language :: Python :: 3.7", + "License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)", + ], + keywords='DMFF', + install_requires=install_requires, + entry_points={} + ) + + +try: + setup(scm={'write_to': 'dmff/_version.py'}) +except: + setup(scm=None) \ No newline at end of file diff --git a/tests/data/methane.xml b/tests/data/methane.xml new file mode 100644 index 000000000..4bb3fea3c --- /dev/null +++ b/tests/data/methane.xml @@ -0,0 +1,44 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/tests/test_api.py b/tests/test_api.py index 8e27e47f7..f576de147 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -150,4 +150,13 @@ def test_PeriodicTorsion_renderXML(self, generators): assert xml[0].name == 'Proper' assert xml[0]['type1'] == 'n1' assert xml[1].name == 'Improper' - assert xml[1]['type1'] == 'n1' \ No newline at end of file + assert xml[1]['type1'] == 'n1' + + def test_parse_multiple_files(self): + pdb = app.PDBFile("tests/data/methane_water.pdb") + h = Hamiltonian("tests/data/methane.xml", "tip3p.xml") + potentials = h.createPotential(pdb.topology) + npt.assert_allclose( + h.getGenerators()[-1].params["charge"], + [-0.1068, 0.0267, 0.0267, 0.0267, 0.0267, -0.834, 0.417] + )