From 70640aa17f168d0b1b9404d9b062f576f2bb6b3d Mon Sep 17 00:00:00 2001
From: Trung Nguyen <ndactrung@gmail.com>
Date: Thu, 3 Aug 2023 14:58:14 -0500
Subject: [PATCH] Add lammps example

Co-authored-by: Pablo Zubieta <8410335+pabloferz@users.noreply.github.com>
---
 examples/lammps/unbiased/lj.lmp      | 21 +++++++
 examples/lammps/unbiased/unbiased.py | 93 ++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+)
 create mode 100644 examples/lammps/unbiased/lj.lmp
 create mode 100644 examples/lammps/unbiased/unbiased.py

diff --git a/examples/lammps/unbiased/lj.lmp b/examples/lammps/unbiased/lj.lmp
new file mode 100644
index 00000000..00cf15d8
--- /dev/null
+++ b/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
diff --git a/examples/lammps/unbiased/unbiased.py b/examples/lammps/unbiased/unbiased.py
new file mode 100644
index 00000000..5c37075a
--- /dev/null
+++ b/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:])