Add files via upload

examples/openmm/funnel_abf/run.py

@@ -1,85 +1,59 @@
 #!/usr/bin/env python3
 
+import dill as pickle
+import matplotlib.pyplot as plt
+import numpy as np
+from parmed import load_file
+from parmed.openmm import StateDataReporter
 
+import pysages
 from pysages.backends import SamplingContext
 from pysages.colvars import Projection_on_Axis_mobile
 
 # %%
 from pysages.methods import CVRestraints, Funnel_ABF, Funnel_Logger, get_funnel_force
 from pysages.utils import try_import
 
-# from pysages.methods import Metadynamics
 # %%
-# from IPython import get_ipython
 openmm = try_import("openmm", "simtk.openmm")
 unit = try_import("openmm.unit", "simtk.unit")
 app = try_import("openmm.app", "simtk.openmm.app")
-import math
-
-import dill as pickle
-import jax.numpy as np
-
-# import numpy as np
-import matplotlib.pyplot as plt
-import numpy
-from openmm import *
-
-# from openmm.app import *
-from openmm.unit import *
-
-# ParmEd Imports
-# ParmEd Imports
-from parmed import load_file
-from parmed import unit as u
-from parmed.openmm import NetCDFReporter, StateDataReporter
-
-import pysages
-
-# %%
-
-# %%
-
 
 # %%
-pi = numpy.pi
-kB = 0.008314462618
+pi = np.pi
+kB = 0.008314462618  # kJ/mol
 
 
 # %%
-def generate_simulation(T=298.15 * kelvin, dt=1.0 * femtoseconds):
+def generate_simulation(T=298.15 * unit.kelvin, dt=1.0 * unit.femtoseconds):
     print("Loading AMBER files...")
     ala2_solv = load_file("complex-wat.prmtop", "complex-wat-prod.rst")
-    pdb = app.PDBFile("last.pdb")
     system = ala2_solv.createSystem(
         nonbondedMethod=app.PME,
         rigidWater=True,
-        switchDistance=1.0 * nanometer,
-        nonbondedCutoff=1.2 * nanometer,
+        switchDistance=1.0 * unit.nanometer,
+        nonbondedCutoff=1.2 * unit.nanometer,
         constraints=app.HBonds,
     )
     # Create the integrator to do Langevin dynamics
     integrator = openmm.LangevinIntegrator(
-        300 * kelvin,  # Temperature of heat bath
-        1.0 / picoseconds,  # Friction coefficient
-        1.0 * femtoseconds,  # Time step
+        300 * unit.kelvin,  # Temperature of heat bath
+        1.0 / unit.picoseconds,  # Friction coefficient
+        1.0 * unit.femtoseconds,  # Time step
     )
 
     # Define the platform to use; CUDA, OpenCL, CPU, or Reference. Or do not specify
-    platform = Platform.getPlatformByName("CPU")
+    platform = openmm.Platform.getPlatformByName("CPU")
     # Create the Simulation object
 
     sim = app.Simulation(ala2_solv.topology, system, integrator, platform)
 
     # Set the particle positions
-    #    sim.context.setPositions(ala2_solv.positions)
-    sim.context.setPositions(pdb.getPositions(frame=-1))
+    sim.context.setPositions(ala2_solv.positions)
     sim.reporters.append(app.PDBReporter("output.pdb", 200000))
     sim.reporters.append(app.DCDReporter("output.dcd", 200000))
     sim.reporters.append(StateDataReporter("data.txt", 20000, step=True, separator=" "))
 
-    # Minimize the energy
-    #    print('Minimizing energy')
-    #    sim.minimizeEnergy()
     return sim
 
 
@@ -91,8 +65,8 @@ def plot_energy(result):
     ax.set_xlabel("CV")
     ax.set_ylabel("Free energy $[\\epsilon]$")
 
-    free_energy = numpy.asarray(result["free_energy"])
-    x = numpy.asarray(result["mesh"])
+    free_energy = np.asarray(result["free_energy"])
+    x = np.asarray(result["mesh"])
     ax.plot(x, free_energy, color="teal")
 
     fig.savefig("energy.png")
@@ -104,8 +78,8 @@ def plot_forces(result):
     ax.set_xlabel("CV")
     ax.set_ylabel("Forces $[\\epsilon]$")
 
-    forces = numpy.asarray(result["mean_force"])
-    x = numpy.asarray(result["mesh"])
+    forces = np.asarray(result["mean_force"])
+    x = np.asarray(result["mesh"])
     ax.plot(x, forces, color="teal")
 
     fig.savefig("forces.png")
@@ -117,21 +91,21 @@ def plot_histogram(result):
     ax.set_xlabel("CV")
     ax.set_ylabel("Histogram $[\\epsilon]$")
 
-    hist = numpy.asarray(result["histogram"])
-    x = numpy.asarray(result["mesh"])
+    hist = np.asarray(result["histogram"])
+    x = np.asarray(result["mesh"])
     ax.plot(x, hist, color="teal")
 
     fig.savefig("histogram2.png")
 
 
 def save_energy_forces(result):
-    Energy = numpy.asarray(result["free_energy"])
-    Forces = numpy.asarray(result["mean_force"])
-    Grid = numpy.asarray(result["mesh"])
-    hist = numpy.asarray(result["histogram"])
-    numpy.savetxt("FES.csv", numpy.column_stack([Grid, Energy]))
-    numpy.savetxt("Forces.csv", numpy.column_stack([Grid, Forces]))
-    numpy.savetxt("Histogram.csv", numpy.column_stack([Grid, hist]))
+    Energy = np.asarray(result["free_energy"])
+    Forces = np.asarray(result["mean_force"])
+    Grid = np.asarray(result["mesh"])
+    hist = np.asarray(result["histogram"])
+    np.savetxt("FES.csv", np.column_stack([Grid, Energy]))
+    np.savetxt("Forces.csv", np.column_stack([Grid, Forces]))
+    np.savetxt("Histogram.csv", np.column_stack([Grid, hist]))
 
 
 #    numpy.savetxt("Forces.csv", numpy.column_stack([Grid, Forces]))
@@ -212,14 +186,13 @@ def save_energy_forces(result):
 funnel_file = "funnel.dat"
 callback = Funnel_Logger(funnel_file, 10)
 sampling_context = SamplingContext(methodo, generate_simulation, callback)
-# la linea de abajo se comenta para cuando se carga el pickle
-# state = pysages.run(sampling_context, timesteps)
-# estas lineas se descomentan para usar el pickle para continuar la simulacion
-with open("restart1.pickle", "rb") as f:
-    state = pickle.load(f)
-state = pysages.run(state, generate_simulation, timesteps)
-# se le puede cambiar e nombre del restart para no sobreescribirlo
-with open("restart2.pickle", "wb") as f:
+state = pysages.run(sampling_context, timesteps)
+# for restart simulation
+# with open("restart1.pickle", "rb") as f:
+#    state = pickle.load(f)
+# state = pysages.run(state, generate_simulation, timesteps)
+# save pickle file
+with open("restart1.pickle", "wb") as f:
     pickle.dump(state, f)
 topology = (8, 8)
 result = pysages.analyze(state, topology=topology)