DMFF (Differentiable Molecular Force Field) is a Jax-based python package that provides a full differentiable implementation of molecular force field models.
License: GNU Lesser General Public License v3.0
I think the use of switch function is wrong in DMFF.
https://github.com/deepmodeling/DMFF/blob/dcc9fa5a4516c06b2fd3383372a68bd68b70ffb0/dmff/classical/inter.py#L49C1-L54C21
It should be:
if self.isSwitch:
x = (dr_norm - self.r_switch) / (self.r_cut - self.r_switch)
S = 1 - 6. * x ** 5 + 15. * x ** 4 - 10. * x ** 3
E = jnp.where(dr_norm > self.r_switch, E * S, E)
E = jnp.where(dr_norm >= self.r_cut, 0., E)v1.0.0
0.4.20
8.1.1
Built from source
Value Error.
No need to reproduce.
No response
Moving the jax2tf to HLOModule for the cpp interface of the saved DMFF model.
The current implementation of the cpp interface between the saved DMFF model and MD engine was based on the jax2tf.
The jax2tf was used to convert the the jax function to TensorFlow function.
However, as an experimental feature of JAX, jax2tf does have some limitations for production use.
google/jax#1871
Old solution.
Lack of documentation, more exploration required
No response
The atom with an axistype of 'Zonly' seems to get a shape error when treated by spatial.py in line 105.
dmff devel
0.3.17
7.7.0
pip
reportbug.zip
This zip file contains the following files,
benchmark0.py
EC.pdb
fillmultipole.xml
errfile
python benchmark0.py
then it will report this error,
"""
raise ValueError(msg.format(arr_shape, shape))
ValueError: Incompatible shapes for broadcasting: (10, 3) and requested shape (1, 3)
"""
Could anyone help to have a look at this problem? Thanks a lot.
Hello,
I'm quite interested by using DMFF, but I find myself somewhat puzzled, and I'd greatly appreciate your assistance in clarifying the theoretical aspects.
Specifically, I'd like to understand the process of introducing minor perturbations to the FF parameters in each iteration, followed by the use of MBAR to estimate properties based on these modified parameters, and ultimately, the computation of gradients with respect to each parameter. In essence, does MBAR have the capability to predict the energy and other system properties when presented with a new set of force field parameters?
Any help would be appreciated (@WangXinyan940).
using DMFF
No response
No response
A workflow specially designed for fine-tuing drug-like force field parameters (typically non-bonded ones) to fit against experimental relative protein-ligand binding free energy data.
Accurate prediction of protein-ligand binding free energy is one of the most important tasks of drug-like force field development. Relative protein-ligand binding free energy calculations based on free energy perturbation (FEP) theory are proved to have the ability to reach chemical accuracy. However, the direct fitting against FEP experimental data are rarely reported due to :
Therefore, I would suggest DMFF supports a general workflow (or framework) to fit against experimental FEP data, which will be revolutionary in this realm.
The workflow should take the following data as inputs:
and then
No response
This is the code for runing MD simulations interfaced with i-PI using the fluctuated leading term force field model for water. This is the case of running a classical MD at 295K, with prepared inital structure density_0.03338.pdb. The project could be submitted using the sub.sh under correct environment. Detailed information of MD is given in input.xml file.
When creating neighbor lists, there is inconsistency on which keyword is used for the cutoff distance.
Lines 11 to 20 in 2690192
Lines 113 to 119 in 2690192
In these pieces of code we see r_cutoff, r_cut, and rc. I recommend these be standardized so that issues don't pop up due to confusion. The classical demo notebook, for example, specified rc=4 as an argument to NeighborList but NeighborList takes no argument rc, only r_cutoff. I fixed that issue in PR #77.
Add a flag (e.g., hax_aux=[False|True]) in createPotential.
When the flag is False: keep the original API for the returned function, like
potential(pos, box, pairs, params) -> energy
When the flag is True: introduce a new input/output variable, like:
potential(pos, box, pairs, params, aux_state) -> energy, aux_state
The new tag should work for both ADMP and Qeq functions
For polarizable force field, one often wants to use the dipoles/charges converged in the previous step as the initial guess of the current invokation. So different potential function calls have to find a way to share information. Currently, in ADMP, it is done by storing the induced dipole (U_ind) in the force class. However, this leads to data leaks during jitting. It was not an important issue in python environment, but causes error when trying to invoke the jitted function in C++.
Therefore, I suggest modify the API of potential function, to make the induced dipoles/charges as explicit input/output of the function (i.e., the aux_state). So it can be jitted without a problem. Meanwhile, we do not want to break the existing code, so the old mechanism (implicit data passing via Force class) needs to be maintained.
Another benefit of introducing aux_state as an explicit output is that it can be differentiated directly, and it is especially convenient in Qeq parameters fitting.
Modify the frontend API, use a flag in createPotential to control the API of the returned function.
Have to test it on both Qeq and ADMP.
No response
implement of charmm fluctuating charge model(Charmm-FQ).
Charges are update by extended lagrangian method in Charmm-FQ model,so they are changed with the time.I don't know how to implement the model,because DMFF seems to be static without the concept of the time. Only matrix inversion to solve the charge?
No response
No response
Make DMFF available via pip install dmff.
Currently, the documentation lists too many dependencies to let users install manually, which is too complex. These packages should be added to dependencies and installed automatically.
mdtraj and rdkit are available in PyPI, so they can also listed as dependencies. Users only need to manually install openmm, which is blocked by openmm/openmm#3796.
No response
nbList = NeighborList(box, rc=4, potentials.meta["cov_map"])should be
nbList = NeighborList(box, r_cutoff=4, covalent_map=potentials.meta["cov_map"])pgrad = param_grad_func(positions, box, pairs, params) reported error for pairs
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[12], line 2
1 param_grad_func = jax.grad(nbfunc, argnums=-1)
----> 2 pgrad = param_grad_func(positions, box, pairs, params)
3 pgrad["NonbondedForce"]["charge"]
[... skipping hidden 4 frame]
File ~/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/_src/api.py:1120, in _check_input_dtype_revderiv(name, holomorphic, allow_int, x)
1117 if (dtypes.issubdtype(aval.dtype, np.integer) or
1118 dtypes.issubdtype(aval.dtype, np.bool_)):
1119 if not allow_int:
-> 1120 raise TypeError(f"{name} requires real- or complex-valued inputs (input dtype "
1121 f"that is a sub-dtype of np.inexact), but got {aval.dtype.name}. "
1122 "If you want to use Boolean- or integer-valued inputs, use vjp "
1123 "or set allow_int to True.")
1124 elif not dtypes.issubdtype(aval.dtype, np.inexact):
1125 raise TypeError(f"{name} requires numerical-valued inputs (input dtype that is a "
1126 f"sub-dtype of np.bool_ or np.number), but got {aval.dtype.name}.")
TypeError: grad requires real- or complex-valued inputs (input dtype that is a sub-dtype of np.inexact), but got int32. If you want to use Boolean- or integer-valued inputs, use vjp or set allow_int to True.I attempted to install DMFF by following the guidelines in the installation markdown file. However, there appears to be some confusion regarding the CUDA version and its dependencies. After completing all the installation steps, I encountered an issue: OpenMM requires CUDA Toolkit 11.1, but the JAX version demands CUDA Toolkit 11.8.
1.0
0.4.23
7.7
python=3.9
ubuntu 20.04
cuda on ubuntu =11.4
As the summary said, is the installation tutorial is outdated?
When I run the following comand,
cd ~/software/DMFF/backend
python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100 --pdb ../examples/water_fullpol/water_dimer.pdb --model ./openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/water_dimer.pdb --has_aux True
it reported the error:
2024-03-18 22:02:19.504317: I tensorflow/cc/saved_model/reader.cc:31] Reading SavedModel from: ./openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/output
2024-03-18 22:02:19.505456: I tensorflow/cc/saved_model/reader.cc:54] Reading meta graph with tags { serve }
2024-03-18 22:02:19.506154: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA
2024-03-18 22:02:19.509908: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2112005000 Hz
2024-03-18 22:02:19.511792: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x56436b07f550 executing computations on platform Host. Devices:
2024-03-18 22:02:19.511817: I tensorflow/compiler/xla/service/service.cc:158] StreamExecutor device (0): <undefined>, <undefined>
2024-03-18 22:02:19.519027: I tensorflow/cc/saved_model/loader.cc:182] Restoring SavedModel bundle.
2024-03-18 22:02:19.524285: I tensorflow/cc/saved_model/loader.cc:132] Running initialization op on SavedModel bundle.
2024-03-18 22:02:19.526928: I tensorflow/cc/saved_model/loader.cc:285] SavedModel load for tags { serve }; Status: success. Took 22613 microseconds.
Running dynamics
2024-03-18 22:02:19.541867: W tensorflow/core/kernels/partitioned_function_ops.cc:197] Grappler optimization failed. Error: Op type not registered 'XlaSharding' in binary running on DESKTOP-87V2MP7. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
2024-03-18 22:02:19.542614: W tensorflow/core/framework/op_kernel.cc:1401] OP_REQUIRES failed at partitioned_function_ops.cc:118 : Not found: Op type not registered 'XlaSharding' in binary running on DESKTOP-87V2MP7. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
Traceback (most recent call last):
File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/runpy.py", line 197, in _run_module_as_main
return _run_code(code, main_globals, None,
File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/runpy.py", line 87, in _run_code
exec(code, run_globals)
File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/OpenMMDMFFPlugin/tests/test_dmff_plugin_nve.py", line 109, in <module>
test_dmff_nve(nsteps=nsteps, time_step=time_step, pdb_file=pdb, model_dir=model_dir, platform_name=platform_name, has_aux=args.has_aux)
File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/OpenMMDMFFPlugin/tests/test_dmff_plugin_nve.py", line 64, in test_dmff_nve
sim.step(nsteps)
File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/openmm/app/simulation.py", line 141, in step
self._simulate(endStep=self.currentStep+steps)
File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/openmm/app/simulation.py", line 210, in _simulate
self.integrator.step(stepsToGo)
File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/openmm/openmm.py", line 3092, in step
return _openmm.VerletIntegrator_step(self, steps)
openmm.OpenMMException: Op type not registered 'XlaSharding' in binary running on DESKTOP-87V2MP7. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
[[{{node StatefulPartitionedCall}}]]
[[{{node StatefulPartitionedCall}}]]
Could someone please take a look for me?
1.0.0
0.4.25
7.7.0
No response
the version of tensorflow is 2.16.1
Expand the function of parameter optimization in DMFF for dynamic properties, such as diffusion coefficient, viscosity, etc. This feature requests,
In v0.2.0, DMFF already support parameter optimization for some thermodynamic properties, such as free energy, RDF, etc. However, for more complicated dynamic properties, DMFF can't get the gradient from the trajectory. Expanding the function of parameter optimization for dynamic properties helps DMFF to be a more powerful engine to accelerate forcefield development in real applications.
Many dynamic properties can be generally evaluated using time-correlation function, which utilizes the information of every state in the trajectory. If we simply use jax.grad on the whole trajectory, the memory may be a huge problem. Inspired by Neural ODE, we can evaluate the gradient in backpropagation process , which utilizes the time reversibility of NVE trajectory. As we get the gradient, we can use modern AI optimization algorithm to fit the dynamic properties.
Any suggestions and comments on this solution are welcomed !
No response
Implement the communication of OpenMM and DMFF trained model in c++ interface.
Enhance the production running efficiency with popular MD engine.
Instead of inheriting from the existing interaction forms in MD engine support force fields, implementation of the MD Engine support for calling DMFF models directly have three advantages.
Refer to openmm_deepmd_plugin for openmm_dmff_plugin architecture.
I will create an openmm_dmff_plugin repo with my own account first.
After the implementation of openmm_dmff_plugin, it will be transfer to deepmodeling account and added into the DMFF as an submodule.
The main problem is the neighbor list information fetching from OpenMM in CUDA platform.
Reference platform called an serial function with O(n) complexity to construct the neighbor list, which is easy to be called in plugin.
However, CUDA platform split the atoms into different block region and construct the interacted neighor blocks parallelized.
How to fetch the information of neighbor list from CUDA context without loss of accuracy is required to be solved rigorously.
Any suggestions and comments on this solution are welcomed !
OpenMM v8.0.0 made some changes for the compatible of openmm-tensorflow/openmm-torch. If these changes are suitable to the development of openmm_dmff_plugin is unknown to me until now.
I'm attempting to test MD runs with the AMBER protein force fields and a solvated system as a basis for some future free energy experiments. For very small systems I've generally had good luck getting things working but when I try to test a larger solvated system (# atoms = 37266), I'm having trouble with neighbor list generation causing an out of memory error.
0.2.1
jaxlib[cuda11_cudnn805]==0.3.15
7.7.0
pip
Input file:
import sys
import time
import jax
import jax.numpy as jnp
import numpy as np
import openmm.app as app
import openmm.unit as unit
from dmff import Hamiltonian, NeighborList
from jax_md import space, smap, energy, minimize, quantity, simulate, quantity
from jax.config import config
config.update("jax_enable_x64", True)
prmtop = app.AmberPrmtopFile('../RAMP1_ion.prmtop')
inpcrd = app.AmberInpcrdFile('../RAMP1_ion.inpcrd')
ff = Hamiltonian("amber14/protein.ff14SB.xml", "amber14/tip3p.xml")
def hhbond(bond):
if bond[0].residue.name == 'HOH':
if bond[0].element._symbol == 'H' and bond[1].element._symbol == 'H':
return True
return False
#remove extra H-H bonds found in AMBER format
prmtop.topology._bonds = [bond for bond in prmtop.topology._bonds if not hhbond(bond)]
potentials = ff.createPotential(prmtop.topology, nonbondedMethod=app.PME, nonbondedCutoff=8*unit.angstrom, prm=prmtop)
params = ff.getParameters()
positions = jnp.array(inpcrd.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
positions = positions - jnp.min(positions, axis=0)
#positions range from 0 to ~9.6 in any given direction
box = jnp.array([
[10.0, 0.0, 0.0],
[0.0, 10.0, 0.0],
[0.0, 0.0, 10.0]
])
#8 angstrom cutoff
nbList = NeighborList(box, .8, potentials.meta["cov_map"])
nbList.allocate(positions)
Error log:
Traceback (most recent call last):
File "/mnt/ufs18/home-094/betanc18/DMFF/examples/classical/forces_bench_ramp/jax/jaxrampdebug.py", line 85, in <module>
nbList.allocate(positions)
File "/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py", line 44, in allocate
self.nblist = self.neighborlist_fn.allocate(positions)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax_md/partition.py", line 816, in allocate_fn
return neighbor_list_fn(position, extra_capacity=extra_capacity, **kwargs)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax_md/partition.py", line 803, in neighbor_list_fn
return neighbor_fn((position, False))
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax_md/partition.py", line 772, in neighbor_fn
idx, occupancy = prune_neighbor_list_sparse(position, idx, **kwargs)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/_src/api.py", line 528, in cache_miss
out_flat = xla.xla_call(
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/core.py", line 1963, in bind
return call_bind(self, fun, *args, **params)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/core.py", line 1979, in call_bind
outs = top_trace.process_call(primitive, fun_, tracers, params)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/core.py", line 689, in process_call
return primitive.impl(f, *tracers, **params)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/_src/dispatch.py", line 236, in _xla_call_impl
return compiled_fun(*args)
File "/mnt/home/betanc18/anaconda3/envs/dmff/lib/python3.9/site-packages/jax/_src/dispatch.py", line 837, in _execute_compiled
out_flat = compiled.execute(in_flat)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Failed to allocate request for 93.12GiB (99990344464B) on device ordinal 0
BufferAssignment OOM Debugging.
BufferAssignment stats:
parameter allocation: 10.35GiB
constant allocation: 144B
maybe_live_out allocation: 10.35GiB
preallocated temp allocation: 93.12GiB
preallocated temp fragmentation: 64B (0.00%)
total allocation: 113.82GiB
total fragmentation: 10.35GiB (9.09%)
Peak buffers:
Buffer 1:
Size: 31.04GiB
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/vmap(jit(_einsum))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: custom-call
Shape: f64[3,1388754756]
==========================
Buffer 2:
Size: 31.04GiB
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 3) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: f64[1388754756,3]
==========================
Buffer 3:
Size: 31.04GiB
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 3) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: f64[1388754756,3]
==========================
Buffer 4:
Size: 10.35GiB
Entry Parameter Subshape: s64[37266,37266]
==========================
Buffer 5:
Size: 10.35GiB
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/concatenate[dimension=0]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: s32[2,1388754756]
==========================
Buffer 6:
Size: 873.4KiB
Entry Parameter Subshape: f64[37266,3]
==========================
Buffer 7:
Size: 72B
XLA Label: constant
Shape: f64[3,3]
==========================
Buffer 8:
Size: 72B
XLA Label: constant
Shape: f64[3,3]
==========================
Buffer 9:
Size: 16B
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/jit(_cumulative_reduction)/pad[padding_config=((1, 1, 1),)]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: (s64[10595], s64[10595])
==========================
Buffer 10:
Size: 16B
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/jit(_cumulative_reduction)/pad[padding_config=((1, 0, 1),)]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: (s64[694377378], s64[694377378])
==========================
Buffer 11:
Size: 16B
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/jit(_cumulative_reduction)/pad[padding_config=((1, 0, 1),)]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: (s64[173594344], s64[173594344])
==========================
Buffer 12:
Size: 16B
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/jit(_cumulative_reduction)/pad[padding_config=((1, 0, 1),)]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: (s64[43398586], s64[43398586])
==========================
Buffer 13:
Size: 16B
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/jit(_cumulative_reduction)/pad[padding_config=((1, 0, 1),)]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: (s64[10849646], s64[10849646])
==========================
Buffer 14:
Size: 16B
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/jit(_cumulative_reduction)/pad[padding_config=((1, 1, 1),)]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: (s64[2712411], s64[2712411])
==========================
Buffer 15:
Size: 16B
Operator: op_name="jit(prune_neighbor_list_sparse)/jit(main)/jit(_cumulative_reduction)/pad[padding_config=((1, 0, 1),)]" source_file="/mnt/home/betanc18/.local/lib/python3.9/site-packages/dmff/common/nblist.py" source_line=44
XLA Label: fusion
Shape: (s64[678102], s64[678102])
==========================
Running the above snippet of Python code with the input files below causes this issue. From the above error, it looks like some of the buffers allocated are essentially n^2 in one dimension. I don't understand the neighbor/cell list generation code in JAX MD well enough to figure out why this is happening but my understanding is that a cell neighbor list should avoid these issues with an appropriately set cutoff.
The 2 files are the protein system building examples from the AMBER tutorials of the RAMP1 protein solvated in a water box.
RAMP1_ion.zip
Implement the communication of LAMMPS and DMFF trained model in c++ interface.
Similar with #85 , supporting of popular MD engines would make cutting-edge forcefield easier to be implemented and tested. People can use DMFF with MD engines they are familiar with.
No response
No response
I'm attempting to compare DMFF simulation results in MD runs to some other libraries with AMBER style protein force fields as a starting point for future research into auto differentiable MD tools. I'm utilizing JAX MD as a simulation engine and DMFF for the neighbor list generation and energy computation. For very small molecules (a few dozen atoms), I'm seeing less than ideal agreement with other tools but still seeing good numerical stability for longer simulations. For larger molecules, especially in solvent (thousands of atoms), I'm seeing mixed results in initial agreement for energies and no long term stability for simulations. Although I would expect some of these differences from other packages, it seems like the initial energies should be closer than they are and in a component wise analysis, are the most different in the non bonded forces. I've noticed that varying the neighbor cutoff changes these results considerably and have been using simulation parameters that mimic the defaults found in most simple AMBER runs (8 angstrom, time step choice, etc). I've tried exploring a few different avenues for things I think may be causing these issues, but I've been unable to notice any trends in these irregularities, even when comparing statistics for energies/forces on a per atom basis.
Although I'm not entirely sure if there are some issues compounding to cause this or user error, I've more generally been looking to see if there good examples of using DMFF energy/force generation in traditional point charge model MD simulations as opposed to just instantaneous energies/forces.
Using an example of a 1289 atom system for the RAMP1 protein and comparing to OpenMM:
Initial Pot E: -1550.4843600621414 kJ/mol
"Step","Potential Energy (kJ/mole)"
100,-8029.34375
200,-8204.943359375
300,-8464.1376953125
400,-8592.5234375
500,-8519.421875
600,-8564.0693359375
700,-8481.87890625
800,-8524.376953125
900,-8626.42578125
1000,-8622.3681640625
Assuming units are in nanometers, comparing this to DMFF with what should be an 8 angstrom cutoff (.8 nanometers in the code for internal consistency) yields fairly different initial energies and results that quickly become unstable.
Initial Potential Energy
3388.2755606199653
Statistics after timestep 100
Potential Energy 7426.3045792548555
Statistics after timestep 200
Potential Energy 6702.456583081994
Statistics after timestep 300
Potential Energy 6629.48060030915
Statistics after timestep 400
Potential Energy 360799615965.03326
Comparing this to another DMFF run where the neighbor list cutoff is increased to 8 (this should be 8 nanometers, not angstrom from my understanding), it seems like there's better initial agreement but instantaneous instability.
Initial Potential Energy
1320.4532334809635
Statistics after timestep 1
Potential Energy 3619221529.320178
0.2.1
jaxlib[cuda11_cudnn805]==0.3.15
7.7.0
pip
My attempt at getting a working MD simulation with DMFF energies:
import sys
import time
import jax
import jax.numpy as jnp
import numpy as np
import openmm.app as app
import openmm.unit as unit
from dmff import Hamiltonian, NeighborList, NeighborListFreud
from jax_md import space, smap, energy, minimize, quantity, simulate, quantity
from jax.config import config
config.update("jax_enable_x64", True)
prmtop = app.AmberPrmtopFile('../RAMP1_gas.prmtop')
inpcrd = app.AmberInpcrdFile('../RAMP1_gas.inpcrd')
ff = Hamiltonian("amber14/protein.ff14SB.xml", "amber14/tip3p.xml")
potentials = ff.createPotential(prmtop.topology, nonbondedMethod=app.NoCutoff)
params = ff.getParameters()
positions = jnp.array(inpcrd.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
box = jnp.array([
[20.0, 0.0, 0.0],
[0.0, 20.0, 0.0],
[0.0, 0.0, 20.0]
])
nbList = NeighborListFreud(box, 8, potentials.meta["cov_map"])
nbList.allocate(positions)
efunc = potentials.getPotentialFunc()
masses = jnp.array([a.element.mass.value_in_unit(unit.dalton) for a in prmtop.topology.atoms()]).reshape((-1, 1))
displacement_fn, shift_fn = space.periodic_general(box, fractional_coordinates=False)
key = jax.random.PRNGKey(0)
energy_fn = lambda pos, **kwargs: efunc(pos, box, kwargs["pairs"], params)
init_fn, apply_fn = simulate.nve(energy_fn, shift_fn, 1e-3)
state = init_fn(key, positions, mass=masses, pairs=nbList.pairs, kT=0)
def body_fn(i, stateList):
state, nbList, trimpairs = stateList
nbList.update(state.position)
trimpairs = nbList.pairs
state = apply_fn(state, pairs=trimpairs)
return state, nbList, trimpairs
print("Initial Energy w/ JAX State")
print(energy_fn(state.position, pairs=nbList.pairs)/4.184)
trimmedpair = nbList.pairs
for i in range(10):
print("Statistics after timestep", (i+1)*100)
for j in range(100):
state, jnbList, trimmedpair = body_fn(j, (state, nbList, trimmedpair))
print("Potential Energy ", end="")
nrg = energy_fn(state.position, pairs=trimmedpair)
print(nrg)
Attempting to do something approximately similar with OpenMM:
import parmed as pmd
from openmm.app import *
from openmm import *
from parmed.unit import *
from sys import stdout
import time
import inspect
prmtop = AmberPrmtopFile('../RAMP1_gas.prmtop')
inpcrd = AmberInpcrdFile('../RAMP1_gas.inpcrd')
system = prmtop.createSystem(nonbondedMethod=NoCutoff, constraints=None)
integrator = VerletIntegrator(0.001*picoseconds)
platform = Platform.getPlatformByName('CUDA')
simulation = Simulation(prmtop.topology, system, integrator, platform)
simulation.context.setPositions(inpcrd.positions)
state = simulation.context.getState(getEnergy=True, getVelocities=True, getForces=True, getPositions=True)
energy = state.getPotentialEnergy()
print("Initial Pot E: ", energy/4.184)
simulation.reporters.append(StateDataReporter(stdout, 100, step=True,
potentialEnergy=True, temperature=False))
simulation.step(1000)
Running the above scripts yields the aforementioned results. Based on some of what I've seen with neighbor list configuration, I'm not sure if I'm making fundamental mistakes in my simulation setup or if there's something else going on in terms of unit conversion/more fundamental issues that are causing the results I'm seeing. Any guidance here would be appreciated!
Included is a zip of both a smaller, stable molecule I've been testing with (22 atoms - Alanine Dipeptide) and a larger unsolvated molecule with unstable behavior (1289 atoms - RAMP1). Both were prepared with AMBER tools.
ramp1.zip
This feature can be decomposed to three parts.
In v0.2.0, DMFF can support virtual site in NonbondedForce. However, considering virtual site is somehow a topology based property, a more graceful way is to support patching/adding virtual site in the frontend instead of in potential implementations.
The virtual site information would be saved in Topology object of OpenMM. We need to implement a tool to patch them from parsers and templates.
When generating energy function, the Hamiltonian itself would also generate a function for pre-processing, update the positions of virtual sites.
And the covalent map generation function need also to consider virtual sites, let them inherit the exclusion list of their parents.
No response
Change the unit of the ADMP frontend from Angstrom to nm
The classical potential frontend assumes the input variables in nm. But the ADMP frontend assumes Angstrom. It is inconvenient, especially if one wants to combine ADMP functions with the classical functions.
Change the frontend unit of ADMP
No response
I would like to propose the development of a new feature to implement a reactive force field in the existing codebase. By incorporating reactive force field capabilities, the code would be able to simulate chemical reactions, bond breaking and formation, and other dynamic processes that are essential for studying chemical systems. This would greatly enhance the versatility and applicability of the code for a wide range of scientific research and industrial applications.
The implementation of a reactive force field feature would significantly advance the capabilities and accuracy of the DMFF codebase, making it more relevant and applicable to a wide range of scientific research and industrial applications. Here are some key reasons why this feature is important:
Accurate Modeling of Chemical Reactions: Chemical reactions are fundamental processes in chemistry and play a crucial role in many scientific disciplines. By incorporating a reactive force field, the code would allow researchers to simulate and study chemical reactions with a high level of accuracy, enabling a deeper understanding of reaction mechanisms, kinetics, and thermodynamics.
Studying Complex Molecular Systems: Many systems of scientific interest, such as biological macromolecules, catalytic materials, and dynamic nanomaterials, involve intricate molecular interactions and reactions. A reactive force field would enable the investigation of these systems at a molecular level, providing insights into their behavior, function, and properties.
Industry Applications: The implementation of a reactive force field feature would have significant implications for industrial applications such as materials science, chemical engineering, and drug development. It could facilitate the design and optimization of novel materials with tailored properties, accelerate the discovery of new catalysts, and streamline the development of chemical processes.
While I do not have the capability to directly participate in the implementation of this feature, I am eager to support and contribute to its development in any way possible.
No response
No response
I found that our implementation of Lennard-Jones potential has a numerical problem when eps=0 exists in the system.
The minimal reproduction is shown below:
import jax.numpy as jnp
import jax
sig_label = jnp.array([0, 1, 0, 2], dtype=int)
eps_label = jnp.array([0, 1, 0, 2], dtype=int)
sig_prm = jnp.array([1.0, 1.1, 1.0])
eps_prm = jnp.array([0.4, 0.5, 0.0])
pairs = jnp.array([
[0, 1],
[0, 2],
[0, 3],
[1, 2],
[1, 3],
[2, 3]
], dtype=int)
pos = jnp.array([
[ 0.0, 0.0, 0.0],
[ 0.5, 0.0, 0.0],
[ 0.0, 0.5, 0.0],
[ 0.2, 0.3, 0.1]
])
def getE(pos, sig_prm, eps_prm, sig_label, eps_label, pairs):
sig_i = sig_prm[sig_label[pairs[:,0]]]
sig_j = sig_prm[sig_label[pairs[:,1]]]
eps_i = eps_prm[eps_label[pairs[:,0]]]
eps_j = eps_prm[eps_label[pairs[:,1]]]
eps = jnp.sqrt(eps_i * eps_j)
sig = (sig_i + sig_j) * 0.5
dr_vec = pos[pairs[:,0]] - pos[pairs[:,1]]
dr_norm = jnp.linalg.norm(dr_vec, axis=1)
dr_inv = 1.0 / dr_norm
sig_dr = sig * dr_inv
sig_dr6 = jnp.power(sig_dr, 6)
sig_dr12 = jnp.power(sig_dr6, 2)
E = 4.0 * eps * (sig_dr12 - sig_dr6)
return jnp.sum(E)
print("GRAD of EPS -- eps_2 = 0")
print(jax.grad(getE, argnums=2)(pos, sig_prm, eps_prm, sig_label, eps_label, pairs))
eps_prm = jnp.array([0.4, 0.5, 1e-12])
print("GRAD of EPS -- eps_2 = 1e-12")
print(jax.grad(getE, argnums=2)(pos, sig_prm, eps_prm, sig_label, eps_label, pairs))
On my laptop, the first printed gradient is nan/inf when eps_2 is 0, while the gradient becomes normal when eps_2 is a tiny number.
The following error is raised when DO_JIT is set to True in calculation of large systems. One can execute the run.py in the attach files to reproduce this bug. This bug exists only on GPU (my machine is T4)
2022-07-28 00:19:02.276166: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:61] Constant folding an instruction is taking > 1s:
%reduce.2384 = s32[] reduce(s32[1991,1991]{1,0} %constant.42, s32[] %constant.9), dimensions={0,1}, to_apply=%region_24.2380, metadata={op_name="jit(potential_fn)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/root/miniconda3/lib/python3.9/site-packages/dmff/classical/inter.py" source_line=126}
This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.
If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Segmentation fault (core dumped)
Provide support for machine learning (ML) force field in OpenMM DMFF plugin. Present version does not yet support well, such as EANN and SGNN models.
The OpenMM DMFF plugin uses save_dmff2tf.py script to transform a JAX model trained in DMFF to a tensorflow model used for OpenMM simulations. In this module, both classical and ADMP force field are considered. While the usage of ML force field in ther transformation is not clear. If ML force field are defined in XML file (referenced as https://github.com/deepmodeling/DMFF/blob/master/docs/user_guide/4.4MLForce.md), there will occur errors in running save_dmff2tf.py script (detailed error message can be browsed in Further Information part). The used ML models are EANN and SGNN.
reportbug.zip
test.zip
Above files record two different errors when using ML in this plugin.
electrostatic energy experession:
the parameters X and H describe the electronegativity and hardness of each atomic species.
constant charge constraint:
constant potential constraint:
the charge distribution should be similar to the results from ab inito calculation(DFT):
use Lagrange mutiplier method to fullfill the constraints, CONP or CONQ.
Obtain energy functional expression.
PME is able to calculate the coulombic interaction energy, the second part of the electrostatic energy, but it's based on point charge model, and what wo want to do is gaussian charge model, there are many similarities.
the parameters X and H for different atomic types are unknown, we can use DMFF to learn these parameters.
Any suggestions and comments on this solution are welcomed !
1.Thomas-Fermi model :J. Chem. Phys. 153, 174704 (2020); https://doi.org/10.1063/5.0028232
2.Fully-Periodic CPM MD:J. Chem. Phys. 156, 184101 (2022); https://doi.org/10.1063/5.0086986
Construct neighborlist using non-JITable code (for speed), and use it in downstream nonbonded force calculation.
The construction of nblist is the most time-consuming part of DMFF. This is because the code is written in python and JITted, but automatic differentiation is useless in the usage.
To optimize the code, we can provide a python/c++/CUDA plugin for nblist construction.
Possible solutions are:
No response
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
An Open Source Machine Learning Framework for Everyone
Django
The Web framework for perfectionists with deadlines.
Laravel
A PHP framework for web artisans
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
Some thing interesting about web. New door for the world.
A server is a program made to process requests and deliver data to clients.
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
Some thing interesting about visualization, use data art
Some thing interesting about game, make everyone happy.
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
Microsoft
Open source projects and samples from Microsoft.
Google
Google ❤️ Open Source for everyone.
Alibaba
Alibaba Open Source for everyone
D3
Data-Driven Documents codes.
Tencent
China tencent open source team.