LAMMPS backend

.github/workflows/ci.yml

@@ -32,7 +32,7 @@ jobs:
  - name: Install python dependecies
  run: |
  python -m pip install --upgrade pip
- pip install dill "e3nn-jax<0.19.2" jaxlib jax-md jaxopt pytest matplotlib
+ pip install dill jaxlib "jax-md>=0.2.7" jaxopt pytest matplotlib
 
  - name: Install pysages
  run: pip install .
@@ -60,7 +60,7 @@ jobs:
  - name: Install python dependecies
  run: |
  python -m pip install --upgrade pip
- pip install dill "e3nn-jax<0.19.2" jaxlib jax-md jaxopt pytest pylint flake8
+ pip install dill jaxlib "jax-md>=0.2.7" jaxopt pytest pylint flake8
  pip install -r docs/requirements.txt
  - name: Install pysages
  run: pip install .

CONTRIBUTORS.md

@@ -9,6 +9,7 @@
 
 - Pablo Zubieta
 - Ludwig Schneider
+- [Trung Nguyen](https:/ndtrung81)
 
 ## Collective Variables
 
@@ -38,4 +39,4 @@ details. Specific contributions to this repository are listed below
 ## Other
 
 For other contributions such as bugfixes or performance improvements, take a look at
-https:/SSAGESLabs/PySAGES/graphs/contributors
+https:/SSAGESLabs/PySAGES/graphs/contributors.

Dockerfile

@@ -7,7 +7,7 @@ RUN python -m pip install ase gsd matplotlib "pyparsing<3"
 
 # Install JAX and JAX-MD
 RUN python -m pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-RUN python -m pip install --upgrade "e3nn-jax<0.19.2" jax-md jaxopt
+RUN python -m pip install --upgrade "jax-md>=0.2.7" jaxopt
 
 COPY . /PySAGES
 RUN pip install /PySAGES/

examples/lammps/unbiased/lj.lmp

@@ -0,0 +1,21 @@
+# 3d Lennard-Jones melt
+
+units lj
+atom_style atomic
+atom_modify map yes
+
+lattice fcc 0.8442
+region box block 0 20 0 20 0 20
+create_box 1 box
+create_atoms 1 box
+mass 1 1.0
+
+velocity all create 1.44 87287 loop geom
+
+pair_style lj/cut 2.5
+pair_coeff 1 1 1.0 1.0 2.5
+
+neighbor 0.3 bin
+neigh_modify delay 5 every 1
+
+fix 1 all nve

examples/lammps/unbiased/unbiased.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+
+"""
+Example unbiased simulation with pysages and lammps.
+
+For a list of possible options for running the script pass `-h` as argument from the
+command line, or call `get_args(["-h"])` if the module was loaded interactively.
+"""
+
+# %%
+import argparse
+import sys
+
+from lammps import lammps
+
+import pysages
+from pysages.backends import SamplingContext
+from pysages.colvars import Component
+from pysages.methods import Unbiased
+
+
+# %%
+def generate_context(args="", script="lj.lmp", store_freq=1):
+ """
+ Returns a lammps simulation defined by the contents of `script` using `args` as
+ initialization arguments.
+ """
+ context = lammps(cmdargs=args.split())
+ context.file(script)
+ # Allow for the retrieval of the unwrapped positions
+ context.command(f"fix unwrap all store/state {store_freq} xu yu zu")
+ return context
+
+
+def get_args(argv):
+ """Process the command-line arguments to this script."""
+
+ available_args = [
+ ("time-steps", "t", int, 1e2, "Number of simulation steps"),
+ ("kokkos", "k", bool, True, "Whether to use Kokkos acceleration"),
+ ]
+ parser = argparse.ArgumentParser(description="Example script to run pysages with lammps")
+
+ for name, short, T, val, doc in available_args:
+ if T is bool:
+ action = "store_" + str(val).lower()
+ parser.add_argument("--" + name, "-" + short, action=action, help=doc)
+ else:
+ convert = (lambda x: int(float(x))) if T is int else T
+ parser.add_argument("--" + name, "-" + short, type=convert, default=T(val), help=doc)
+
+ return parser.parse_args(argv)
+
+
+def main(argv):
+ """Example simulation with pysages and lammps."""
+ args = get_args(argv)
+
+ context_args = {"store_freq": args.time_steps}
+ if args.kokkos:
+ # Passed to the lammps constructor as `cmdargs` when running the script
+ # with the --kokkos (or -k) option
+ context_args["args"] = "-k on g 1 -sf kk -pk kokkos newton on neigh half"
+
+ # Setting the collective variable, method, and running the simulation
+ cvs = [Component([0], i) for i in range(3)]
+ method = Unbiased(cvs)
+ sampling_context = SamplingContext(method, generate_context, context_args=context_args)
+ result = pysages.run(sampling_context, args.time_steps)
+
+ # Post-run analysis
+ # -----------------
+ context = sampling_context.context
+ nlocal = sampling_context.sampler.view.local_particle_number()
+ snapshot = result.snapshots[0]
+ state = result.states[0]
+
+ # Retrieve the pointer to the unwrapped positions,
+ ptr = context.extract_fix("unwrap", 1, 2)
+ # and make them available as a numpy ndarray
+ positions = context.numpy.darray(ptr, nlocal, dim=3)
+ # Get the map to sort the atoms since they can be reordered during the simulation
+ ids = context.numpy.extract_atom("id").argsort()
+
+ # The ids of the final snapshot in pysages and lammps should be the same
+ assert (snapshot.ids == ids).all()
+ # For our example, the last value of the CV should match
+ # the unwrapped position of the zeroth atom
+ assert (state.xi.flatten() == positions[ids[0]]).all()
+
+
+if __name__ == "__main__":
+ main(sys.argv[1:])

pysages/backends/core.py

@@ -82,10 +82,12 @@ def __init__(
  self._backend_name = "ase"
  elif module_name.startswith("hoomd"):
  self._backend_name = "hoomd"
- elif module_name.startswith("simtk.openmm") or module_name.startswith("openmm"):
- self._backend_name = "openmm"
  elif isinstance(context, JaxMDContext):
  self._backend_name = "jax-md"
+ elif module_name.startswith("lammps"):
+ self._backend_name = "lammps"
+ elif module_name.startswith("simtk.openmm") or module_name.startswith("openmm"):
+ self._backend_name = "openmm"
 
  if self._backend_name is None:
  backends = ", ".join(supported_backends())
@@ -113,14 +115,15 @@ def __enter__(self):
  """
  if hasattr(self.context, "__enter__"):
  return self.context.__enter__()
+ return self.context
 
  def __exit__(self, exc_type, exc_value, exc_traceback):
  """
  Trampoline 'with statements' to the wrapped context when the backend supports it.
  """
  if hasattr(self.context, "__exit__"):
- return self.context.__exit__(exc_type, exc_value, exc_traceback)
+ self.context.__exit__(exc_type, exc_value, exc_traceback)
 
 
 def supported_backends():
- return ("ase", "hoomd", "jax-md", "openmm")
+ return ("ase", "hoomd", "jax-md", "lammps", "openmm")