Merge pull request #149 from InnocentBug/parallel2

Reimplement parallelization API

Author: Ludwig Schneider

.github/workflows/docker-ci.yml

@@ -84,6 +84,30 @@ jobs:
           docker load --input /tmp/pysages.tar
           docker run -t pysages bash -c "cd PySAGES/examples/openmm/metad/ && python3 ./alanine-dipeptide.py --time-steps=25"
 
+  alanine-dipeptide-openmm-mpi:
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Download artifact
+        uses: actions/download-artifact@v2
+        with:
+          name: pysages
+          path: /tmp
+      - name: Set up script
+        id: example
+        run: |
+          SCRIPT="cd PySAGES/examples/openmm/umbrella_integration &&"
+          SCRIPT="${SCRIPT} mpirun --allow-run-as-root --oversubscribe -n 6"
+          SCRIPT="${SCRIPT} python3 -m mpi4py.futures"
+          SCRIPT="${SCRIPT} integration.py --replicas=5 --time-steps=1000 --mpi"
+          echo "::set-output name=script::$SCRIPT"
+      - name: Load and run test
+        run: |
+          docker load --input /tmp/pysages.tar
+          docker run -t pysages bash -c "${{ steps.example.outputs.script }}"
+
   metad-hoomd:
     runs-on: ubuntu-latest
     needs: build
@@ -143,3 +167,35 @@ jobs:
           name: umbrella-integration-plots
           path: /tmp/plots
           retention-days: 1
+
+  umbrella-integration-hoomd-mpi:
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Download artifact
+        uses: actions/download-artifact@v2
+        with:
+          name: pysages
+          path: /tmp
+      - name: Set up script
+        id: setup
+        run: |
+          SCRIPT="cd PySAGES/examples/hoomd-blue/umbrella_integration &&"
+          SCRIPT="${SCRIPT} python3 ./gen_gsd.py &&"
+          SCRIPT="${SCRIPT} mpirun --allow-run-as-root --oversubscribe -n 6"
+          SCRIPT="${SCRIPT} python3 -m mpi4py.futures"
+          SCRIPT="${SCRIPT} integration.py --replicas=5 --time-steps=1000 --mpi &&"
+          SCRIPT="${SCRIPT} mkdir /tmp/plots && mv *.pdf /tmp/plots/"
+          echo "::set-output name=script::$SCRIPT"
+      - name: Load and run test
+        run: |
+          docker load --input /tmp/pysages.tar
+          docker run -v /tmp:/tmp -t pysages bash -c "${{ steps.setup.outputs.script }}"
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v2
+        with:
+          name: umbrella-integration-mpi-plots
+          path: /tmp/plots
+          retention-days: 1

.gitignore

@@ -1,9 +1,10 @@
 __pycache__/
-pysages/_version.py
+**.pdf
 build/
 dist/
 docs/build
+examples/**/*.dat
+examples/**/*.gsd
+examples/**/.ipynb_checkpoints/
+pysages/_version.py
 pysages.egg-info/
-examples/*/*pdf
-examples/*/*dat
-examples/*/*gsd

examples/hoomd-blue/umbrella_integration/integration.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import sys
 import argparse
+import importlib
 import numpy as np
 import matplotlib.pyplot as plt
 
@@ -10,14 +11,19 @@
 
 import pysages
 from pysages.colvars import Component
-from pysages.methods import UmbrellaIntegration
+from pysages.methods import UmbrellaIntegration, SerialExecutor
 
 
 params = {"A": 0.5, "w": 0.2, "p": 2}
 
 
 def generate_context(**kwargs):
-    hoomd.context.initialize("")
+    if kwargs.get("mpi_enabled"):
+        MPI = importlib.import_module("mpi4py.MPI")
+        init_kwargs = {"mpi_comm": MPI.COMM_SELF}
+    else:
+        init_kwargs = {}
+    hoomd.context.initialize("--single-mpi", **init_kwargs)
     context = hoomd.context.SimulationContext()
 
     with context:
@@ -106,6 +112,7 @@ def get_args(argv):
     parser = argparse.ArgumentParser(description="Example script to run umbrella integration")
     for (name, short, T, val, doc) in available_args:
         parser.add_argument("--" + name, "-" + short, type=T, default=T(val), help=doc)
+    parser.add_argument("--mpi", action="store_true", help="Use MPI executor")
     args = parser.parse_args(argv)
     return args
 
@@ -119,6 +126,13 @@ def post_run_action(**kwargs):
     )
 
 
+def get_executor(args):
+    if args.mpi:
+        futures = importlib.import_module("mpi4py.futures")
+        return futures.MPIPoolExecutor()
+    return SerialExecutor()
+
+
 def main(argv):
     args = get_args(argv)
 
@@ -127,8 +141,15 @@ def main(argv):
     centers = list(np.linspace(args.start_path, args.end_path, args.replicas))
     method = UmbrellaIntegration(cvs, args.k_spring, centers, args.log_period, args.log_delay)
 
+    context_args = {"mpi_enabled": args.mpi}
+
     raw_result = pysages.run(
-        method, generate_context, args.time_steps, post_run_action=post_run_action
+        method,
+        generate_context,
+        args.time_steps,
+        context_args=context_args,
+        post_run_action=post_run_action,
+        executor=get_executor(args),
     )
     result = pysages.analyze(raw_result)
 

examples/hoomd-blue/umbrella_integration/slurm.sbatch

@@ -0,0 +1,13 @@
+#!/bin/sh
+
+#SBATCH --job-name=umbrella
+#SBATCH --partition=gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=5
+#SBATCH --time=01:00:00
+#SBATCH --output=out.txt
+#SBATCH --gres=gpu:4
+#SBATCH --mem=20G
+
+#python harmonic_bias.py
+mpirun -n 5 python -m mpi4py.futures integration.py

examples/openmm/Harmonic_Bias.ipynb

@@ -437,17 +437,26 @@
    "provenance": []
   },
   "jupytext": {
-   "formats": "ipynb,md",
-   "main_language": "python"
+   "formats": "ipynb,md"
   },
   "kernelspec": {
    "display_name": "Python 3",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.2"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 0
+ "nbformat_minor": 1
 }

examples/openmm/Harmonic_Bias.md

@@ -2,14 +2,14 @@
 jupyter:
   jupytext:
     formats: ipynb,md
-    main_language: python
     text_representation:
       extension: .md
       format_name: markdown
       format_version: '1.3'
       jupytext_version: 1.14.0
   kernelspec:
     display_name: Python 3
+    language: python
     name: python3
 ---
 

examples/openmm/umbrella_integration/adp-explicit.pdb

@@ -0,0 +1 @@
+../../inputs/alanine-dipeptide/adp-explicit.pdb

examples/openmm/umbrella_integration/integration.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+
+import sys
+import argparse
+import importlib
+import numpy as np
+
+from pysages.colvars import DihedralAngle
+from pysages.methods import UmbrellaIntegration, SerialExecutor
+from pysages.utils import try_import
+
+import pysages
+
+openmm = try_import("openmm", "simtk.openmm")
+unit = try_import("openmm.unit", "simtk.unit")
+app = try_import("openmm.app", "simtk.openmm.app")
+
+
+def generate_simulation(**kwargs):
+    """
+    Generates a simulation context, we pass this function to `pysages.run`.
+    """
+    pdb_filename = "adp-explicit.pdb"
+    T = 298.15 * unit.kelvin
+    dt = 2.0 * unit.femtoseconds
+    pdb = app.PDBFile(pdb_filename)
+
+    ff = app.ForceField("amber99sb.xml", "tip3p.xml")
+    cutoff_distance = 1.0 * unit.nanometer
+    topology = pdb.topology
+    system = ff.createSystem(
+        topology,
+        constraints=app.HBonds,
+        nonbondedMethod=app.NoCutoff,
+        nonbondedCutoff=cutoff_distance,
+    )
+
+    positions = pdb.getPositions(asNumpy=True)
+
+    integrator = openmm.LangevinIntegrator(T, 1 / unit.picosecond, dt)
+
+    simulation = app.Simulation(topology, system, integrator)
+    simulation.context.setPositions(positions)
+    simulation.minimizeEnergy()
+
+    return simulation
+
+
+def get_args(argv):
+    available_args = [
+        ("k-spring", "k", float, 50, "Spring constant for each replica"),
+        ("replicas", "N", int, 25, "Number of replicas along the path"),
+        ("time-steps", "t", int, 1e4, "Number of simulation steps for each replica"),
+        ("log-period", "l", int, 100, "Frequency of logging the CVs into each histogram"),
+        ("log-delay", "d", int, 0, "Number of timesteps to discard before logging"),
+    ]
+    parser = argparse.ArgumentParser(description="Example script to run umbrella integration")
+    for (name, short, T, val, doc) in available_args:
+        parser.add_argument("--" + name, "-" + short, type=T, default=T(val), help=doc)
+    parser.add_argument("--mpi", action="store_true", help="Use MPI executor")
+    args = parser.parse_args(argv)
+    return args
+
+
+def get_executor(args):
+    if args.mpi:
+        futures = importlib.import_module("mpi4py.futures")
+        return futures.MPIPoolExecutor()
+    return SerialExecutor()
+
+
+def main(argv):
+    args = get_args(argv)
+
+    cvs = (DihedralAngle((4, 6, 8, 14)), DihedralAngle((6, 8, 14, 16)))
+    centers = []
+    center_pos = np.linspace(+0.45 * np.pi, -0.45 * np.pi, args.replicas)
+    for pos in center_pos:
+        centers.append((pos, pos))
+    method = UmbrellaIntegration(cvs, args.k_spring, centers, args.log_period, args.log_delay)
+
+    raw_result = pysages.run(
+        method,
+        generate_simulation,
+        args.time_steps,
+        # post_run_action=post_run_action,
+        executor=get_executor(args),
+    )
+    result = pysages.analyze(raw_result)
+    print(result)
+
+    return result
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])