"""A streaming LAMMPS external MD engine.
This module defines the canonical LAMMPS engine, selected by
``class = "lammps"``. LAMMPS is run as a single persistent process and
its dump trajectory is read *on-the-fly*, so the path is filled
incrementally and LAMMPS is terminated as soon as a stopping interface
is crossed. The LAMMPS run is built from a self-contained input template
(``lammps.input``) whose ``pyretis_*`` placeholders -- timestep, number
of steps, subcycles, initial configuration, data file, temperature and
the per-segment thermostat seed -- are substituted for each propagation.
The relaunch (start-and-stop) engine, which builds its input by
appending PyRETIS run commands to a force-field-only template, remains
available as ``class = "lammps_steps"``
(:py:mod:`pyretis.engines.lammps_steps`).
Important classes defined here
------------------------------
LAMMPSEngine (:py:class:`.LAMMPSEngine`)
The streaming LAMMPS engine.
"""
from __future__ import annotations
import logging
import os
import shlex
import signal
import subprocess # nosec B404
from pathlib import Path
from time import sleep
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import numpy as np
from pyretis.engines.external import ExternalMDEngine
from pyretis.core.systemhelp import update_system_box
# NOTE: ``pyretis.engines._parts`` is imported lazily inside the methods
# that need it (``_propagate_from`` / ``modify_velocities``). A top-level
# import would pull in ``pyretis.setup`` while ``pyretis.engines`` is still
# initialising (this module is imported from ``engines/__init__``), which
# is a circular import. The streaming/velocity paths are runtime-only, so
# importing at call time is safe.
# Re-export the LAMMPS input reader from the relaunch module so that
# ``from pyretis.engines.lammps import read_lammps_input`` keeps working
# now that the streaming engine owns the canonical module name. The
# relaunch module does not import this one, so there is no import cycle.
from pyretis.engines.lammps_steps import (
LAMMPSEngineSteps,
PYRETIS_ORDER_MODE,
read_lammps_input,
)
__all__ = [
'LAMMPSEngine',
'read_lammps_input',
'read_lammpstrj',
'write_lammpstrj',
'read_energies',
'get_atom_masses',
]
if TYPE_CHECKING:
from pyretis.core.path import Path as InfPath
from pyretis.core._system_inf import System
from pyretis.inout.fileio import FileIO
# ``FileIO``, ``InfPath``, ``System`` references stay as forward
# strings via PEP 563 (``from __future__ import annotations`` above).
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())
[docs]def write_lammpstrj(
outfile: str,
id_type: np.ndarray,
pos: np.ndarray,
vel: np.ndarray,
box: Optional[np.ndarray],
append: bool = False,
triclinic: bool = False,
) -> None:
"""Write a LAMMPS trajectory frame in .lammpstrj format.
This method assumes a dump format from LAMMPS of:
`dump id type x y z vx vy vz`.
Args:
outfile: Path to the file to write.
id_type:
pos: The positions to write.
vel: The velocities to write.
box: The simulation box to write (if available).
append: If True, the `outfile` will be appended to,
otherwise, `outfile` will be overwritten.
triclinic: If true, write in triclinic box format.
"""
filemode = "a" if append else "w"
box_header = "xy xz yz " if triclinic else ""
with open(outfile, filemode) as writefile:
to_write = f"ITEM: TIMESTEP\n0\nITEM: NUMBER OF ATOMS\n\
{pos.shape[0]}\nITEM: BOX BOUNDS {box_header}pp pp pp\n"
if box is not None:
for box_vector in box:
to_write += " ".join(box_vector.astype(str)) + "\n"
to_write += "ITEM: ATOMS id type x y z vx vy vz\n"
for t, x, v in zip(id_type, pos, vel):
to_write += (
" ".join(t.astype(int).astype(str))
+ " "
+ " ".join(x.astype(str))
+ " "
+ " ".join(v.astype(str))
+ "\n"
)
writefile.write(to_write)
[docs]def read_lammpstrj(
infile: str, frame: int, n_atoms: int
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Read a single frame from a LAMMPS trajectory.
Args:
infile: The path to the file to read.
frame: The index of the frame to read from `infile`.
n_atoms: The number of atoms to read.
Returns:
A tuple containing:
- The id type for the particles.
- The positions.
- The velocities.
- The simulation box.
Note:
- The atoms are sorted according to their index, so that
(index-based) order calculations are correct.
"""
block_size = n_atoms + 9
box = np.genfromtxt(infile, skip_header=block_size * frame + 5, max_rows=3)
posvel = np.genfromtxt(
infile, skip_header=block_size * frame + 9, max_rows=n_atoms
)
id_sorted = np.argsort(posvel[:, 0])
pos = posvel[id_sorted, 2:5]
vel = posvel[id_sorted, 5:8]
id_type = posvel[id_sorted, :2]
return id_type, pos, vel, box
[docs]def read_energies(filename: str) -> Dict[str, np.ndarray]:
"""Read energies from a LAMMPS log file.
This method is used to read the thermodynamic output from a
simulation (e.g., potential and kinetic energies).
Args:
filename: The path to the LAMMPS log file to read.
Returns:
A dictionary containing the data we found in the file.
"""
energy_keys = []
energy_data: Dict[str, List[Union[float, int]]] = {}
read_energy = False
with open(filename, encoding="utf-8") as logfile:
for lines in logfile:
if "Step" in lines.split():
# Assume that this is the start of the thermo output.
energy_keys = [i.strip() for i in lines.strip().split()]
for key in energy_keys:
# Note: This will discard the previously read
# thermodynamic data. This is because we only want
# to read the final section of thermodynamic data,
# which, by construction of the LAMMPS input file,
# correspond to the output of the MD run.
energy_data[key] = []
read_energy = True
continue
if lines.startswith("Loop time"):
# Assume this marks the end of the thermo output.
read_energy = False
energy_keys = []
continue
if read_energy and energy_keys:
# Assume that we are reading energies.
try:
data = [float(i) for i in lines.strip().split()]
except ValueError:
# Some text has snuck into the thermo output,
# ignore it:
continue
for key, value in zip(energy_keys, data):
if key == "Step":
energy_data[key].append(int(value))
else:
energy_data[key].append(value)
return {key: np.array(val) for key, val in energy_data.items()}
[docs]def get_atom_masses(lammps_data: Union[str, Path], atom_style) -> np.ndarray:
"""Read atom masses from a LAMMPS data file.
TODO: atom style dependent!!
Args:
lammps_data: Path to the LAMMPS data file to read.
atom_style: The style used in lammps. Supports 'full' and 'charge'.
Returns:
An array with the masses read from the file.
"""
col = {"full": 2, "charge": 1}
if atom_style not in col:
raise NotImplementedError(f"Style {atom_style}' not supported yet.")
n_atoms = 0
n_atom_types = 0
atom_type_masses = np.empty((0, 2))
atoms = np.empty(0)
with open(lammps_data) as readfile:
for i, line in enumerate(readfile):
spl = line.split()
# skip empty lines
if not spl:
continue
# get number of atoms
if len(spl) == 2 and "atoms" == spl[1]:
n_atoms = int(spl[0])
elif len(spl) == 3 and spl[1] == "atom" and spl[2] == "types":
n_atom_types = int(spl[0])
elif spl and spl[0] == "Masses":
atom_type_masses = np.genfromtxt(
lammps_data, skip_header=i + 1, max_rows=n_atom_types
).reshape(-1, 2)
elif spl[0] == "Atoms":
atoms = np.genfromtxt(
lammps_data, skip_header=i + 1, max_rows=n_atoms
)
# sort atoms according to index
idx = np.argsort(atoms[:, 0])
atoms = atoms[idx]
# if we did not find all of the information
if n_atoms == 0 or n_atom_types == 0:
raise ValueError(f"Could not read atom masses from {lammps_data}. \
Found {n_atoms} atoms and {n_atom_types} atom_types.")
if np.shape(atoms)[0] == 0 or np.any(np.isnan(atoms)):
raise ValueError(f"Could not read 'Atoms' from {lammps_data}")
if np.shape(atom_type_masses)[0] == 0 or np.any(
np.isnan(atom_type_masses)
):
raise ValueError(f"Could not read 'Masses' from {lammps_data}")
if np.shape(atoms)[0] != n_atoms:
msg = (
f"Found {n_atoms} atoms but read {np.shape(atoms)[0]} "
f"from 'Atoms' in {lammps_data}"
)
raise ValueError(msg)
if np.shape(atom_type_masses)[0] != n_atom_types:
msg = (
f"Found {n_atom_types} atom types but read "
f"{np.shape(atom_type_masses)[0]} from 'Masses' in {lammps_data}"
)
raise ValueError(msg)
masses = np.zeros((n_atoms, 1))
for atom_type in range(1, n_atom_types + 1):
idx = np.where(atoms[:, col[atom_style]] == atom_type)[0]
masses[idx] = atom_type_masses[atom_type - 1, 1]
return masses
def shift_boxbounds(
xyz: np.ndarray, box: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
"""Shift positions so that the lower box bounds are 0.
Args:
xyz: The positions to shift.
box: The box boundaries.
Returns:
A tuple containing:
- The shifted positions.
- A flattended version of the upper box bounds.
"""
xyz -= box[:, 0]
box[:, 1] -= box[:, 0]
return xyz, box[:, 1].flatten()
def check_lammps_data(lmp_data):
"""Perform checks on a lammps.data file.
E.g. to retrieve the box type (triclinic or not).
"""
important_info = {"triclinic": False}
with open(lmp_data) as rfile:
for line in rfile:
processed_line = " ".join(line.split())
if "xy xz yz" in processed_line:
important_info["triclinic"] = True
return important_info
[docs]class LAMMPSEngine(LAMMPSEngineSteps):
"""Streaming LAMMPS engine.
This engine shares its configuration and input convention with the
relaunch engine :py:class:`.LAMMPSEngineSteps` (same ``[engine]``
settings, same ``lammps.in`` / ``system.data`` / ``order.in`` /
``extra_files`` input directory) so a user can switch
``class = lammps`` <-> ``class = lammps_steps`` and have it just
work -- as for the two GROMACS engines. The two engines build the
*identical* LAMMPS input (via the inherited
:py:meth:`.LAMMPSEngineSteps._prepare_lammps_run`); they differ only
in how the trajectory is consumed: this engine runs LAMMPS as a
single persistent process and reads the trajectory on the fly,
terminating LAMMPS early when an interface is crossed.
TO DO:
* Add possibility to use LAMMPS internally to modify the velocities
such that all constraints (e.g. between bonds) are fulfilled
(``external_orderparameter : boolean``).
* Velocity generation depends on units... only 'real' atm
* Fix "exe_path" and "exe_dir", use one but not both...
* Fix mpi in cp2k?
* For small systems the engine is so fast that whole trajectory
finishes before we managed to check for a crossing, leading to
large files.
* Make LAMMPS remove and replace commands in input
instead of making variables. Much easier to run md with
the same input file
* Periodicity in LAMMPS has to take into account hi and lo box bounds:
As of now, we subtract the lower bounds from the positions and the box
vectors to create a regular box with 0 lower bounds. This is hard coded
in _propagate_from() when calculating, and also in _read_configuration().
* external_orderparameter. If this is set to true, we read
in the orderparameter from and external file. May be useful when the
orderparameter is calculated internally in LAMMPS for expensive
calculations such as in nucleation, or any other fancy stuff.
"""
# LAMMPS is an external real-binary engine whose order parameter is
# computed by PyRETIS from the dumped configuration (the engine calls
# ``self.calculate_order`` with ``self.order_function``). The inf base
# declared neither attribute, so the effective values came from the
# native ``ExternalMDEngine``/``EngineBase`` defaults; we set them
# explicitly here to match: external engine, needs an order parameter.
engine_type = "external"
needs_order = True
default_units = None
[docs] @classmethod
def get_default_units(cls, settings=None):
"""Return the default unit system for the LAMMPS engine.
The unit system is read from the ``units`` command in the
LAMMPS input file. This mirrors how the engine time step is
recovered at settings-parse time
(:py:func:`pyretis.core.engine_time._engine_timestep_from_input`),
which also inspects ``lammps.in``, so unit and time-step
auto-detection stay consistent.
Parameters
----------
settings : dict, optional
Full simulation settings or engine settings.
Returns
-------
out : string
The default unit system for the LAMMPS engine.
Raises
------
ValueError
If the LAMMPS input file can not be read, or if no
``units`` command is found in that file.
"""
if settings is None:
return super().get_default_units(settings)
engine_settings = settings.get('engine', settings)
input_path = engine_settings.get('input_path', '.')
exe_path = engine_settings.get('exe_path')
if 'simulation' in settings:
exe_path = settings['simulation'].get('exe_path', exe_path)
if exe_path is None:
exe_path = os.path.abspath('.')
template = os.path.join(exe_path, input_path, 'lammps.in')
if not os.path.isfile(template):
msg = f'Could not determine LAMMPS units: Missing "{template}"'
raise ValueError(msg)
for key, value in read_lammps_input(template):
if key.lower() == 'units':
return value.strip()
msg = f'Could not determine LAMMPS units from "{template}"'
raise ValueError(msg)
[docs] def __init__(self, lmp, input_path, subcycles, extra_files=None,
order_mode=PYRETIS_ORDER_MODE, timestep=None,
velocity_generation='engine', temperature=None,
atom_style="full", sleep=0.1):
"""Set up the streaming LAMMPS engine.
The signature matches :py:class:`.LAMMPSEngineSteps`, so the same
``[engine]`` settings and input directory drive either engine
(switch ``class = lammps`` <-> ``class = lammps_steps`` and it
just works). ``atom_style`` and ``sleep`` are the only
streaming-specific extras.
Parameters
----------
lmp : string
The LAMMPS executable.
input_path : string
Directory with the LAMMPS input files (``lammps.in``,
``system.data``, optionally ``order.in`` and ``extra_files``).
subcycles : integer
The number of LAMMPS steps between stored snapshots.
extra_files : list of strings, optional
Additional files required to run LAMMPS.
order_mode : string, optional
Where the order parameter is evaluated. The streaming engine
reads dumped frames and evaluates the PyRETIS order function
in Python; ``"pyretis"`` (the default) reflects that.
timestep : float, optional
Cross-check only; the time step is read from ``lammps.in``.
velocity_generation : string, optional
``"engine"`` (LAMMPS) or ``"maxwell"`` (Python draw); see
:py:class:`.LAMMPSEngineSteps`.
temperature : float, optional
Temperature (K) for Maxwell-Boltzmann velocity generation.
atom_style : string, optional
The LAMMPS ``atom_style`` used to read masses from the data
file. Default ``"full"``.
sleep : float, optional
Seconds to wait between polls while reading the trajectory on
the fly. Default ``0.1``.
"""
super().__init__(
lmp, input_path, subcycles, extra_files=extra_files,
order_mode=order_mode, timestep=timestep,
velocity_generation=velocity_generation, temperature=temperature,
)
# Streaming-specific I/O state. The trajectory is read on the fly
# in LAMMPS-dump format, so we need the per-atom masses, the atom
# count and the box type up front (the relaunch engine reads these
# lazily from each stored frame instead).
self.exe_dir = "."
self.name = "lammps"
self.sleep = sleep
self.atom_style = atom_style
self.ext = "lammpstrj"
self.mass = get_atom_masses(self.input_files['conf'], self.atom_style)
self.n_atoms = self.mass.shape[0]
self.triclinic = check_lammps_data(
self.input_files['conf']
)['triclinic']
[docs] def step(self, system, name):
"""Perform a single step with the engine.
LAMMPS is propagated via ``_propagate_from``; a single
:py:meth:`step` is not supported. The method exists only so the
class can be instantiated (the native external base declares it
as an abstract method).
"""
raise NotImplementedError(
"LAMMPSEngine does not support single-step integration "
'("step()"); it is propagated via "_propagate_from".'
)
[docs] def calculate_order(self, system, xyz=None, vel=None, box=None):
"""Calculate the order parameter of the current system.
This engine's ``_read_configuration`` returns the box as its THIRD
element (``pos, vel, box, names``), unlike the native
``ExternalMDEngine.calculate_order`` which expects the box first.
We therefore keep the box-third unpacking here.
"""
if any((xyz is None, vel is None, box is None)):
out = self._read_configuration(system.config[0])
xyz = out[0]
vel = out[1]
box = out[2]
if xyz is not None:
system.pos = xyz
if vel is not None:
system.vel = vel * -1.0 if system.vel_rev else vel
if box is not None:
update_system_box(system, box)
if self.order_function is None:
raise ValueError("Order parameter is not defined!")
return self.order_function.calculate(system)
[docs] @staticmethod
def add_to_path(path, phase_point, left, right):
"""Add a phase point and perform some checks."""
status = "Running propagate..."
success = False
stop = False
path.append(phase_point)
if path.phasepoints[-1].order[0] < left:
status = "Crossed left interface!"
success = True
stop = True
elif path.phasepoints[-1].order[0] > right:
status = "Crossed right interface!"
success = True
stop = True
elif path.maxlen is not None and path.length >= path.maxlen:
status = "Max. path length exceeded"
success = False
stop = True
return status, success, stop
[docs] def dump_phasepoint(self, phasepoint, deffnm="conf"):
"""Dump the current configuration from a system object to a file."""
pos_file = self.dump_config(phasepoint.config, deffnm=deffnm)
phasepoint.set_pos((pos_file, 0))
[docs] def draw_maxwellian_velocities(self, vel, mass, beta, rgen, sigma_v=None):
"""Draw velocities from a Gaussian distribution."""
if sigma_v is None or np.any(sigma_v < 0.0):
kbt = 1.0 / beta
sigma_v = np.sqrt(kbt * (1 / mass))
npart, dim = vel.shape
vel = rgen.normal(loc=0.0, scale=sigma_v, size=(npart, dim))
return vel, sigma_v
[docs] def propagate(self, path, ens_set, system, reverse=False):
"""Propagate by streaming a single persistent LAMMPS run.
The relaunch base (:py:class:`.LAMMPSEngineSteps`) overrides
``propagate`` with its relaunch-per-step driver. The streaming
engine instead uses the generic external-engine driver, which
dispatches to
:py:meth:`pyretis.engines.lammps.LAMMPSEngine._propagate_from`
-- a single persistent LAMMPS process whose trajectory is read
on the fly and stopped early on an interface crossing.
Parameters
----------
path : object like :py:class:`.PathBase`
The path to fill with phase points.
ens_set : dict
The ensemble settings (interfaces, etc.).
system : object like :py:class:`.System`
The system to propagate.
reverse : boolean, optional
If True, propagate the path in reverse.
Returns
-------
success : boolean
True if the propagation produced an acceptable path.
status : string
A textual status of the propagation.
"""
return ExternalMDEngine.propagate(
self, path, ens_set, system, reverse=reverse
)
[docs] def dump_config(self, config, deffnm="conf"):
"""Extract a configuration frame, honouring the data-file case.
The shared convention stores the initial configuration as a
LAMMPS *data* file (``system.data``), which is a single
configuration and cannot be sliced frame-wise like a trajectory
dump. In that case we pass the data file through unchanged (LAMMPS
reads it via ``read_data`` in the run input); otherwise we defer to
the generic trajectory-frame extraction.
Parameters
----------
config : tuple
The configuration as ``(filename, index)``.
deffnm : string, optional
The base name for the dumped file.
Returns
-------
out : string
The file name we dumped to.
"""
pos_file = config[0]
if str(pos_file).endswith(".data"):
# The initial configuration is often given as the bare conf
# name ("system.data") before it has been staged into exe_dir;
# resolve it against exe_dir and, failing that, the input conf.
source = pos_file
if not os.path.isfile(source):
staged = os.path.join(self.exe_dir, os.path.basename(source))
source = staged if os.path.isfile(staged) \
else self.input_files["conf"]
out_file = os.path.join(self.exe_dir, f"{deffnm}.data")
if os.path.abspath(source) != os.path.abspath(out_file):
self._copyfile(source, out_file)
return out_file
return ExternalMDEngine.dump_config(self, config, deffnm=deffnm)
def _propagate_from(
self,
name: str,
path: InfPath,
system: System,
ens_set: Dict[str, Any],
msg_file: FileIO,
reverse: bool = False,
) -> Tuple[bool, str]:
# Lazy import (see module-level note) to avoid the _parts -> setup
# circular import at module load.
from pyretis.engines._parts import (
ReadAndProcessOnTheFly,
lammpstrj_reader,
)
status = f"propagating with LAMMPS (reverse = {reverse})"
interfaces = ens_set["interfaces"]
logger.debug(status)
success = False
left, _, right = interfaces
initial_conf = system.config[0]
# A LAMMPS data file is a single configuration that cannot be sliced
# frame-wise in Python; there the initial order parameter is captured
# from frame 0 of the on-the-fly dump below (the same frame the
# relaunch engine reads). We only pre-compute it here when the initial
# configuration is a trajectory frame.
if not str(initial_conf).endswith(".data"):
xyzi, veli, boxi, _ = self._read_configuration(initial_conf)
order = self.calculate_order(system, xyz=xyzi, vel=veli, box=boxi)
msg_file.write(
"# Initial order parameter: "
+ " ".join([str(i) for i in order])
)
# Build the per-run LAMMPS input via the shared relaunch helper so
# the streaming and relaunch engines ask LAMMPS to do the *same*
# thing (same template, same reads, same dump); only the trajectory
# consumption below differs (we read it on the fly and stop early).
md_settings = {
"steps_subcycles": path.maxlen * self.subcycles,
"reverse_velocities": reverse != system.vel_rev,
"interfaces": interfaces,
}
cmd = self._prepare_lammps_run(system, md_settings, name)
run_input = os.path.join(self.exe_dir, f"{name}.in")
traj_file = md_settings["traj"]
msg_file.write(f"# Trajectory file is: {traj_file}")
out_name = os.path.join(self.exe_dir, "stdout.txt")
err_name = os.path.join(self.exe_dir, "stderr.txt")
cwd = self.exe_dir
cmd2 = " ".join(cmd)
# fire off lammps
return_code = None
lammps_was_terminated = False
step_nr = 0
with open(out_name, "wb") as fout, open(err_name, "wb") as ferr:
exe = subprocess.Popen(
cmd,
stdin=subprocess.PIPE,
stdout=fout,
stderr=ferr,
shell=False, # nosec B603
cwd=cwd,
preexec_fn=os.setsid,
)
# wait for trajectories to appear
while not os.path.exists(traj_file):
sleep(self.sleep)
if exe.poll() is not None:
logger.debug("LAMMPS execution stopped")
break
# LAMMPS may have finished after last processing the files
# or it may have crashed without writing to the files
if exe.poll() is None or exe.returncode == 0:
traj_reader = ReadAndProcessOnTheFly(
traj_file, lammpstrj_reader
)
# start reading on the fly as LAMMPS is still running
# if it stops, perform one more iteration to read
# the remaining content in the files.
iterations_after_stop = 0
step_nr = 0
trajectory: list[np.ndarray] = []
box_trajectory: list[np.ndarray] = []
while exe.poll() is None or iterations_after_stop <= 1:
# we may still have some data in the trajectory
# so use += here
frames = traj_reader.read_and_process_content()
trajectory += frames[0]
box_trajectory += frames[1]
# loop over the frames that are ready
for frame in range(len(trajectory)):
posvel = trajectory.pop(0)
box = box_trajectory.pop()
pos = posvel[:, :3]
vel = posvel[:, 3:]
# shift the box bounds
pos, box = shift_boxbounds(pos, box)
# calculate order, check for crossings, etc
order = self.calculate_order(
system, xyz=pos, vel=vel, box=box
)
msg_file.write(
f'{step_nr} {" ".join([str(j) for j in order])}'
)
snapshot = {
"order": order,
# Index frames by LAMMPS time step (frame *
# subcycles), matching the relaunch engine so a
# config produced by either engine refers to the
# same snapshot.
"config": (traj_file, step_nr * self.subcycles),
"vel_rev": reverse,
}
phase_point = self.snapshot_to_system(system, snapshot)
status, success, stop = self.add_to_path(
path, phase_point, left, right
)
if stop:
# process may have terminated since we last checked
if exe.poll() is None:
logger.debug("Terminating LAMMPS execution")
os.killpg(os.getpgid(exe.pid), signal.SIGTERM)
# wait for process to die, necessary for mpi
exe.wait(timeout=360)
logger.debug(
"LAMMPS propagation ended at %i. Reason: %s",
step_nr,
status,
)
# exit while loop without reading additional data
iterations_after_stop = 2
lammps_was_terminated = True
break
step_nr += 1
sleep(self.sleep)
# if LAMMPS finished, we run one more loop
if exe.poll() is not None and iterations_after_stop <= 1:
iterations_after_stop += 1
return_code = exe.returncode
if return_code != 0 and not lammps_was_terminated:
logger.error(
"Execution of external program (%s) failed!",
self.description,
)
logger.error("Attempted command: %s", cmd2)
logger.error("Execution directory: %s", cwd)
logger.error(
"Return code from external program: %i", return_code
)
logger.error("STDOUT, see file: %s", out_name)
logger.error("STDERR, see file: %s", err_name)
msg = (
f"Execution of external program ({self.description}) "
f"failed with command:\n {cmd2}.\n"
f"Return code: {return_code}"
)
raise RuntimeError(msg)
if (return_code is not None) and (
return_code == 0 or lammps_was_terminated
):
self._removefile(out_name)
self._removefile(err_name)
msg_file.write("# Propagation done.")
energy_file = os.path.join(self.exe_dir, f"{name}.log")
msg_file.write(f"# Reading energies from: {energy_file}")
energy = read_energies(energy_file)
end = (step_nr + 1) * self.subcycles
ekin: Union[np.ndarray, list[float]] = energy.get("KinEng", [])
vpot: Union[np.ndarray, list[float]] = energy.get("PotEng", [])
etot: Union[np.ndarray, list[float]] = energy.get("TotEng", [])
temp: Union[np.ndarray, list[float]] = energy.get("Temp", [])
path.update_energies(ekin[:end], vpot[:end], etot[:end], temp[:end])
self._removefile(run_input)
# Mirror the inf base ``propagate``: accumulate the MD steps this
# segment consumed so the scheduler's subcycle accounting stays
# byte-identical.
self.steps += path.length
return success, status
[docs] def _frame_config_index(self, frame: int) -> int:
"""LAMMPS indexes configurations by MD time step.
The on-the-fly dump stores one frame every ``subcycles`` steps and
snapshots are indexed by time step (``frame * subcycles``), the
shared convention with the relaunch engine. This is the inverse of
the ``idx // subcycles`` mapping in :py:meth:`._extract_frame`.
"""
return frame * self.subcycles
[docs] def _recompute_partial_orders(self, partial_traj, reverse):
"""Order parameters per frame of a partial ``.lammpstrj`` dump.
``recalculate_order`` does not handle the LAMMPS trajectory
format, so the dump frames are read directly. Only complete frames
are read: a final frame left half-written by an interruption (an
incomplete block) is skipped, keeping the surviving frames -- the
same "use the last good frame" recovery used for the other engines.
"""
# Lazy import (see module-level note) to avoid import cycles.
# pylint: disable=import-outside-toplevel
from pyretis.core import System, ParticlesExt
block_size = self.n_atoms + 9
with open(partial_traj, encoding='utf-8') as fileh:
n_complete = sum(1 for _ in fileh) // block_size
system = System(box=None)
system.particles = ParticlesExt(dim=3)
system.vel_rev = reverse
orders = []
for frame in range(n_complete):
_, pos, vel, box = read_lammpstrj(partial_traj, frame,
self.n_atoms)
pos, box = shift_boxbounds(pos, box)
orders.append(
self.calculate_order(system, xyz=pos, vel=vel, box=box)
)
return orders
[docs] def _read_configuration(
self, filename: str
) -> Tuple[
np.ndarray, np.ndarray, Optional[np.ndarray], Optional[list[str]]
]:
"""Read a configuration (a single frame trajectory)."""
id_type, pos, vel, box = read_lammpstrj(filename, 0, self.n_atoms)
pos, box = shift_boxbounds(pos, box)
return pos, vel, box, None
[docs] def _reverse_velocities(self, filename: str, outfile: str) -> None:
"""Reverse the velocities of a configuration."""
id_type, pos, vel, box = read_lammpstrj(filename, 0, self.n_atoms)
vel *= -1.0
write_lammpstrj(
outfile, id_type, pos, vel, box, triclinic=self.triclinic
)
[docs] def modify_velocities(
self, system: System, vel_settings: Dict[str, Any], rgen
) -> Tuple[float, float]:
"""Modify the velocities of all particles.
Args:
system: The system whose particle velocities are to be modified.
vel_settings: A dict containing the key ``'zero_momentum'``
(boolean): if true we reset the linear momentum to zero
after generating velocities internally.
Returns:
A tuple containing:
- dek: The change in kinetic energy as a result of the
velocity modification.
- kin_new: The new kinetic energy of the system.
Note:
This method does **not** take care of constraints.
"""
# Lazy import (see module-level note).
from pyretis.engines._parts import kinetic_energy, reset_momentum
mass = self.mass
# energy is in units kcal/mol which we want to convert
# to units (g/mol)*Å^2/fs (units of m*v^2), the velocity
# units of lammps.
# using Unitful:
# uconvert((u"kcal/g")^0.5, 1u"Å/fs") = 48.88821290839617
# so we need to scale the velocities by this factor
scale = 48.88821290839617
pos = self.dump_frame(system)
id_type, xyz, vel, box = read_lammpstrj(pos, 0, self.n_atoms)
kin_old = kinetic_energy(vel, mass)[0]
vel, _ = self.draw_maxwellian_velocities(vel, mass, self.beta, rgen)
# convert to correct units
vel /= scale
# reset momentum is not the default in LAMMPS
if vel_settings.get("zero_momentum", False):
vel = reset_momentum(vel, mass)
conf_out = os.path.join(self.exe_dir, f"genvel.{self.ext}")
write_lammpstrj(
conf_out, id_type, xyz, vel, box, triclinic=self.triclinic
)
kin_new = kinetic_energy(vel, mass)[0]
system.config = (conf_out, 0)
system.ekin = kin_new
if kin_old == 0.0:
dek = float("inf")
logger.debug(
"Kinetic energy not found for previous point."
"\n(This happens when the initial configuration "
"does not contain energies.)"
)
else:
dek = kin_new - kin_old
return dek, kin_new
[docs] def set_mdrun(self, md_items: Dict[str, Any]) -> None:
"""Set the executional directory for workers."""
self.exe_dir = md_items["exe_dir"]