diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py new file mode 100644 index 0000000..c4ca98a --- /dev/null +++ b/Examples/LJ_langevin.py @@ -0,0 +1,82 @@ +from openmmtools.testsystems import LennardJonesFluid + +# Use the LennardJonesFluid example from openmmtools to initialize particle positions and topology +# For this example, the topology provides the masses for the particles +# The default LennardJonesFluid example considers the system to be Argon with 39.9 amu +lj_fluid = LennardJonesFluid(reduced_density=0.1, nparticles=1000) + + +from chiron.potential import LJPotential +from openmm import unit + +# initialize the LennardJones potential in chiron +# +sigma = 0.34 * unit.nanometer +epsilon = 0.238 * unit.kilocalories_per_mole +cutoff = 3.0 * sigma + +lj_potential = LJPotential( + lj_fluid.topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff +) + +from chiron.states import SamplerState, ThermodynamicState + +# define the sampler state +sampler_state = SamplerState( + x0=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors() +) + +# define the thermodynamic state +thermodynamic_state = ThermodynamicState( + potential=lj_potential, temperature=300 * unit.kelvin +) + +from chiron.neighbors import NeighborListNsqrd, OrthogonalPeriodicSpace + +# define the neighbor list for an orthogonal periodic space +skin = 0.5 * unit.nanometer + +nbr_list = NeighborListNsqrd( + OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180 +) +# build the neighbor list from the sampler state +nbr_list.build_from_state(sampler_state) + +from chiron.reporters import SimulationReporter + +# initialize a reporter to save the simulation data +filename = "test_lj.h5" +import os + +if os.path.isfile(filename): + os.remove(filename) +reporter = SimulationReporter("test_lj.h5", lj_fluid.topology, 1) + +from chiron.integrators import LangevinIntegrator + +# initialize the Langevin integrator +integrator = LangevinIntegrator(reporter=reporter, save_frequency=100) + +integrator.run( + sampler_state, + thermodynamic_state, + n_steps=2000, + nbr_list=nbr_list, + progress_bar=True, +) + +import h5py + +# read the data from the reporter +with h5py.File("test_lj.h5", "r") as f: + energies = f["energy"][:] + steps = f["step"][:] + + +# plot the energy +import matplotlib.pyplot as plt + +plt.plot(steps, energies) +plt.xlabel("Step (fs)") +plt.ylabel("Energy (kj/mol)") +plt.show() diff --git a/chiron/potential.py b/chiron/potential.py index 5d3daf5..abb3a03 100644 --- a/chiron/potential.py +++ b/chiron/potential.py @@ -25,7 +25,8 @@ def compute_force(self, positions, nbr_list=None) -> jnp.ndarray: def compute_pairlist(self, positions, cutoff) -> jnp.array: """ - Compute the pairlist for a given set of positions and a cutoff distance. + Compute the pairlist for a given set of positions and a cutoff distance + without using periodic boundary conditions or any specific optimizations. Parameters ---------- @@ -89,13 +90,21 @@ def __init__( if not isinstance(topology, Topology): if not isinstance(topology, property): if topology is not None: - raise TypeError(f"Topology must be a Topology object or None, type(topology) = {type(topology)}") + raise TypeError( + f"Topology must be a Topology object or None, type(topology) = {type(topology)}" + ) if not isinstance(sigma, unit.Quantity): - raise TypeError(f"sigma must be a unit.Quantity, type(sigma) = {type(sigma)}") + raise TypeError( + f"sigma must be a unit.Quantity, type(sigma) = {type(sigma)}" + ) if not isinstance(epsilon, unit.Quantity): - raise TypeError(f"epsilon must be a unit.Quantity, type(epsilon) = {type(epsilon)}") + raise TypeError( + f"epsilon must be a unit.Quantity, type(epsilon) = {type(epsilon)}" + ) if not isinstance(cutoff, unit.Quantity): - raise TypeError(f"cutoff must be a unit.Quantity, type(cutoff) = {type(cutoff)}") + raise TypeError( + f"cutoff must be a unit.Quantity, type(cutoff) = {type(cutoff)}" + ) if not sigma.unit.is_compatible(unit.angstrom): raise ValueError(f"sigma must have units of distance, got {sigma.unit}") @@ -104,8 +113,6 @@ def __init__( if not cutoff.unit.is_compatible(unit.nanometer): raise ValueError(f"cutoff must have units of distance, got {cutoff.unit}") - - self.sigma = sigma.value_in_unit_system( unit.md_unit_system ) # The distance at which the potential is zero @@ -141,11 +148,7 @@ def _compute_energy_masked(self, distance, mask): ) return energy.sum() - def compute_energy( - self, - positions: jnp.array, - nbr_list=None, - ): + def compute_energy(self, positions: jnp.array, nbr_list=None, debug_mode=False): """ Compute the LJ energy. @@ -153,10 +156,9 @@ def compute_energy( ---------- positions : jnp.array The positions of the particles in the system - nbr_list : NeighborList, optional - Instance of the neighborlist class to use. By default, set to None, which will use an N^2 pairlist - shift : bool, optional - Whether to shift the potential energy at the cutoff, by default False + nbr_list : NeighborList, default=None + Instance of a neighbor list or pair list class to use. + If None, an unoptimized N^2 pairlist will be used without PBC conditions. Returns ------- potential_energy : float @@ -166,6 +168,9 @@ def compute_energy( # Compute the pair distances and displacement vectors if nbr_list is None: + log.debug( + "nbr_list is None, computing pairlist using N^2 method without PBC." + ) # Compute the pairlist for a given set of positions and a cutoff distance # Note in this case, we do not need the pairs or displacement vectors # Since we already calculate the distance in the pairlist computation @@ -194,8 +199,10 @@ def compute_energy( raise ValueError("Neighborlist must be built before use") # ensure that the cutoff in the neighbor list is the same as the cutoff in the potential - if nbr_list.cutoff != self.cutoff: - raise ValueError(f"Neighborlist cutoff ({nbr_list.cutoff}) must be the same as the potential cutoff ({self.cutoff})") + if nbr_list.cutoff != self.cutoff: + raise ValueError( + f"Neighborlist cutoff ({nbr_list.cutoff}) must be the same as the potential cutoff ({self.cutoff})" + ) n_neighbors, mask, dist, displacement_vectors = nbr_list.calculate( positions @@ -273,9 +280,13 @@ def __init__( U0: unit.Quantity = 0.0 * unit.kilocalories_per_mole, ): if not isinstance(topology, Topology): - if not isinstance(topology, property): #importing from the topology from the model results in it being a property object + if not isinstance( + topology, property + ): # importing from the topology from the model results in it being a property object if topology is not None: - raise TypeError(f"Topology must be a Topology object or None, type(topology) = {type(topology)}") + raise TypeError( + f"Topology must be a Topology object or None, type(topology) = {type(topology)}" + ) if not isinstance(k, unit.Quantity): raise TypeError(f"k must be a unit.Quantity, type(k) = {type(k)}") if not isinstance(x0, unit.Quantity): @@ -284,11 +295,17 @@ def __init__( raise TypeError(f"U0 must be a unit.Quantity, type(U0) = {type(U0)}") if not k.unit.is_compatible(unit.kilocalories_per_mole / unit.angstrom**2): - raise ValueError(f"k must be a unit.Quantity with units of energy per distance squared, k.unit = {k.unit}") + raise ValueError( + f"k must be a unit.Quantity with units of energy per distance squared, k.unit = {k.unit}" + ) if not x0.unit.is_compatible(unit.angstrom): - raise ValueError(f"x0 must be a unit.Quantity with units of distance, x0.unit = {x0.unit}") + raise ValueError( + f"x0 must be a unit.Quantity with units of distance, x0.unit = {x0.unit}" + ) if not U0.unit.is_compatible(unit.kilocalories_per_mole): - raise ValueError(f"U0 must be a unit.Quantity with units of energy, U0.unit = {U0.unit}") + raise ValueError( + f"U0 must be a unit.Quantity with units of energy, U0.unit = {U0.unit}" + ) log.info("Initializing HarmonicOscillatorPotential") log.info(f"k = {k}")