Skip to content

Commit

Permalink
Backport fixes from PR #366. [ci skip]
Browse files Browse the repository at this point in the history
  • Loading branch information
lohedges committed Nov 25, 2024
1 parent 541779c commit cb9814b
Show file tree
Hide file tree
Showing 12 changed files with 429 additions and 122 deletions.
299 changes: 276 additions & 23 deletions python/BioSimSpace/Align/_align.py

Large diffs are not rendered by default.

60 changes: 27 additions & 33 deletions python/BioSimSpace/FreeEnergy/_relative.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ def __init__(
engine : str
The molecular dynamics engine used to run the simulation. Available
options are "GROMACS", or "SOMD". If this argument is omitted then
options are "AMBER", "GROMACS", or "SOMD". If this argument is omitted then
BioSimSpace will choose an appropriate engine for you.
setup_only : bool
Expand Down Expand Up @@ -1158,8 +1158,8 @@ def _preprocess_data(data, estimator, **kwargs):
# Assign defaults in case not passed via kwargs.
auto_eq = False
stat_ineff = False
truncate = False
truncate_keep = "start"
truncate_upper = 100
truncate_lower = 0

# Parse kwargs.
for key, value in kwargs.items():
Expand All @@ -1168,35 +1168,31 @@ def _preprocess_data(data, estimator, **kwargs):
auto_eq = value
if key == "STATISTICALINEFFICIENCY":
stat_ineff = value
if key == "TRUNCATEPERCENTAGE":
truncate = value
if key == "TRUNCATEKEEP":
truncate_keep = value
if key == "TRUNCATEUPPER":
truncate_upper = value
if key == "TRUNCATELOWER":
truncate_lower = value

# First truncate data.
# Copy the data.
raw_data = data
if truncate:
# Get the upper and lower bounds for truncate.
data_len = len(data[0])
data_step = round((data[0].index[-1][0] - data[0].index[-2][0]), 1)
data_kept = data_len * (truncate / 100)
data_time = data_kept * data_step
if truncate_keep == "start":
truncate_lower = 0
truncate_upper = data_time - data_step
if truncate_keep == "end":
truncate_lower = (data_len * data_step) - data_time
truncate_upper = (data_len * data_step) - data_step

try:
data = [
_slicing(i, lower=truncate_lower, upper=truncate_upper)
for i in raw_data
]
except:
_warnings.warn("Could not truncate data.")
data = raw_data
else:
# Data length.
data_len = len(data[0])

# Step size.
data_step = round((data[0].index[-1][0] - data[0].index[-2][0]), 1)

# Get the upper and lower bounds for truncate.
truncate_lower = (data_len * (truncate_lower / 100)) * data_step
truncate_upper = (data_len * (truncate_upper / 100)) * data_step

try:
data = [
_slicing(i, lower=truncate_lower, upper=truncate_upper)
for i in raw_data
]
except:
_warnings.warn("Could not truncate data.")
data = raw_data

# The decorrelate function calls either autoequilibration or statistical_inefficiency
Expand Down Expand Up @@ -1230,10 +1226,7 @@ def _preprocess_data(data, estimator, **kwargs):
)
sampled_data = data

# Concatanate in alchemlyb format.
processed_data = _alchemlyb.concat(sampled_data)

return processed_data
return sampled_data

@staticmethod
def _analyse_internal(files, temperatures, lambdas, engine, estimator, **kwargs):
Expand Down Expand Up @@ -1319,6 +1312,7 @@ def _analyse_internal(files, temperatures, lambdas, engine, estimator, **kwargs)
# Preprocess the data.
try:
processed_data = Relative._preprocess_data(data, estimator, **kwargs)
processed_data = _alchemlyb.concat(processed_data)
except:
_warnings.warn("Could not preprocess the data!")
processed_data = _alchemlyb.concat(data)
Expand Down
8 changes: 8 additions & 0 deletions python/BioSimSpace/Process/_amber.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,14 @@ def _setup(self, **kwargs):
"perturbable molecule!"
)

# Make sure the protocol is valid.
if self._protocol.getPerturbationType() != "full":
raise NotImplementedError(
"AMBER currently only supports the 'full' perturbation "
"type. Please use engine='SOMD' when running multistep "
"perturbation types."
)

# If this is vacuum simulation with pmemd.cuda then
# we need to add a simulation box.
if self._is_vacuum and self._is_pmemd_cuda:
Expand Down
11 changes: 5 additions & 6 deletions python/BioSimSpace/Process/_openmm.py
Original file line number Diff line number Diff line change
Expand Up @@ -1488,13 +1488,12 @@ def getFrame(self, index):

if not type(index) is int:
raise TypeError("'index' must be of type 'int'")
max_index = (
int(
(self._protocol.getRunTime() / self._protocol.getTimeStep())
/ self._protocol.getRestartInterval()
)
- 1

max_index = int(
(self._protocol.getRunTime() / self._protocol.getTimeStep())
/ self._protocol.getRestartInterval()
)

if index < 0 or index > max_index:
raise ValueError(f"'index' must be in range [0, {max_index}].")

Expand Down
44 changes: 40 additions & 4 deletions python/BioSimSpace/Protocol/_equilibration.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,14 +42,15 @@ class Equilibration(_Protocol, _PositionRestraintMixin):
def __init__(
self,
timestep=_Types.Time(2, "femtosecond"),
runtime=_Types.Time(0.2, "nanoseconds"),
runtime=_Types.Time(1, "nanoseconds"),
temperature_start=_Types.Temperature(300, "kelvin"),
temperature_end=_Types.Temperature(300, "kelvin"),
temperature=None,
pressure=None,
thermostat_time_constant=_Types.Time(1, "picosecond"),
report_interval=100,
restart_interval=500,
report_interval=200,
restart_interval=1000,
restart=False,
restraint=None,
force_constant=10 * _Units.Energy.kcal_per_mol / _Units.Area.angstrom2,
):
Expand Down Expand Up @@ -89,6 +90,9 @@ def __init__(
The frequency at which restart configurations and trajectory
frames are saved. (In integration steps.)
restart : bool
Whether this is a continuation of a previous simulation.
restraint : str, [int]
The type of restraint to perform. This should be one of the
following options:
Expand Down Expand Up @@ -155,6 +159,9 @@ def __init__(
# Set the restart interval.
self.setRestartInterval(restart_interval)

# Set the restart flag.
self.setRestart(restart)

# Set the posistion restraint.
_PositionRestraintMixin.__init__(self, restraint, force_constant)

Expand All @@ -169,7 +176,7 @@ def _get_parm(self):
f"thermostat_time_constant={self._thermostat_time_constant}, "
f"report_interval={self._report_interval}, "
f"restart_interval={self._restart_interval}, "
+ _PositionRestraintMixin._get_parm(self)
f"restart={self._restart}, " + _PositionRestraintMixin._get_parm(self)
)

def __str__(self):
Expand Down Expand Up @@ -204,6 +211,7 @@ def __eq__(self, other):
and self._thermostat_time_constant == other._thermostat_time_constant
and self._report_interval == other._report_interval
and self._restart_interval == other._restart_interval
and self._restart == other._restart
and _PositionRestraintMixin.__eq__(self, other)
)

Expand Down Expand Up @@ -484,6 +492,34 @@ def setRestartInterval(self, restart_interval):

self._restart_interval = restart_interval

def isRestart(self):
"""
Return whether this restart simulation.
Returns
-------
is_restart : bool
Whether this is a restart simulation.
"""
return self._restart

def setRestart(self, restart):
"""
Set the restart flag.
Parameters
----------
restart : bool
Whether this is a restart simulation.
"""
if isinstance(restart, bool):
self._restart = restart
else:
_warnings.warn("Non-boolean restart flag. Defaulting to False!")
self._restart = False

def isConstantTemp(self):
"""
Return whether the protocol has a constant temperature.
Expand Down
7 changes: 6 additions & 1 deletion python/BioSimSpace/Protocol/_free_energy_equilibration.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,14 +44,15 @@ def __init__(
max_lam=1.0,
num_lam=11,
timestep=_Types.Time(2, "femtosecond"),
runtime=_Types.Time(0.2, "nanoseconds"),
runtime=_Types.Time(1, "nanoseconds"),
temperature_start=_Types.Temperature(300, "kelvin"),
temperature_end=_Types.Temperature(300, "kelvin"),
temperature=None,
pressure=None,
thermostat_time_constant=_Types.Time(1, "picosecond"),
report_interval=200,
restart_interval=1000,
restart=False,
restraint=None,
force_constant=10 * _Units.Energy.kcal_per_mol / _Units.Area.angstrom2,
perturbation_type="full",
Expand Down Expand Up @@ -106,6 +107,9 @@ def __init__(
The frequency at which restart configurations and trajectory
frames are saved. (In integration steps.)
restart : bool
Whether this is a continuation of a previous simulation.
restraint : str, [int]
The type of restraint to perform. This should be one of the
following options:
Expand Down Expand Up @@ -217,4 +221,5 @@ def _to_regular_protocol(self):
restart_interval=self.getRestartInterval(),
restraint=self.getRestraint(),
force_constant=self.getForceConstant(),
restart=self.isRestart(),
)
4 changes: 2 additions & 2 deletions python/BioSimSpace/Protocol/_production.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,8 @@ def __init__(
temperature=_Types.Temperature(300, "kelvin"),
pressure=_Types.Pressure(1, "atmosphere"),
thermostat_time_constant=_Types.Time(1, "picosecond"),
report_interval=100,
restart_interval=100,
report_interval=200,
restart_interval=1000,
restart=False,
restraint=None,
force_constant=10 * _Units.Energy.kcal_per_mol / _Units.Area.angstrom2,
Expand Down
2 changes: 1 addition & 1 deletion python/BioSimSpace/_Config/_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,6 @@ def steps(self):
steps = self._protocol.getSteps()
else:
steps = _math.ceil(
self._protocol.getRunTime() / self._protocol.getTimeStep()
round(self._protocol.getRunTime() / self._protocol.getTimeStep(), 6)
)
return steps
49 changes: 31 additions & 18 deletions python/BioSimSpace/_Config/_gromacs.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,30 +100,21 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
if not all(isinstance(line, str) for line in extra_lines):
raise TypeError("Lines in 'extra_lines' must be of type 'str'.")

# Make sure the report interval is a multiple of nstcalcenergy.
if isinstance(self._protocol, _FreeEnergyMixin):
nstcalcenergy = 250
else:
nstcalcenergy = 100
report_interval = self.reportInterval()
if report_interval % nstcalcenergy != 0:
report_interval = nstcalcenergy * _math.ceil(
report_interval / nstcalcenergy
)

# Define some miscellaneous defaults.
protocol_dict = {
# Interval between writing to the log file.
"nstlog": report_interval,
"nstlog": self.reportInterval(),
# Interval between writing to the energy file.
"nstenergy": report_interval,
"nstenergy": self.reportInterval(),
# Interval between writing to the trajectory file.
"nstxout-compressed": self.restartInterval(),
}

# Minimisation.
if isinstance(self._protocol, _Protocol.Minimisation):
protocol_dict["integrator"] = "steep"
# Maximum step size in nanometers.
protocol_dict["emstep"] = "0.001"
else:
# Timestep in picoseconds
timestep = self._protocol.getTimeStep().picoseconds().value()
Expand Down Expand Up @@ -212,7 +203,7 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
# Temperature control.
if not isinstance(self._protocol, _Protocol.Minimisation):
if isinstance(self._protocol, _FreeEnergyMixin):
# Langevin dynamics.
# Leap-frog stochastic dynamics integrator.
protocol_dict["integrator"] = "sd"
else:
# Leap-frog molecular dynamics.
Expand Down Expand Up @@ -260,6 +251,24 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
"%.2f" % self._protocol.getTemperature().kelvin().value()
)

# Set as a continuation run.
if self.isRestart():
protocol_dict["continuation"] = "yes"
protocol_dict["gen-vel"] = "no"
# Generate velocities.
else:
protocol_dict["continuation"] = "no"
protocol_dict["gen-vel"] = "yes"
protocol_dict["gen-seed"] = "-1"
if isinstance(self._protocol, _Protocol.Equilibration):
protocol_dict["gen-temp"] = (
"%.2f" % self._protocol.getStartTemperature().kelvin().value()
)
else:
protocol_dict["gen-temp"] = (
"%.2f" % self._protocol.getTemperature().kelvin().value()
)

# Free energies.
if isinstance(self._protocol, _FreeEnergyMixin):
# Extract the lambda array.
Expand All @@ -280,10 +289,14 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
protocol_dict["couple-lambda1"] = "vdw-q"
# Write all lambda values.
protocol_dict["calc-lambda-neighbors"] = -1
# Calculate energies every 250 steps.
protocol_dict["nstcalcenergy"] = 250
# Write gradients every 250 steps.
protocol_dict["nstdhdl"] = 250
# Calculate energies at the report interval.
protocol_dict["nstcalcenergy"] = self.reportInterval()
# Write gradients at the report interval.
protocol_dict["nstdhdl"] = self.reportInterval()
# Soft-core parameters.
protocol_dict["sc-alpha"] = "0.30"
protocol_dict["sc-sigma"] = "0.25"
protocol_dict["sc-coul"] = "yes"

# Put everything together in a line-by-line format.
total_dict = {**protocol_dict, **extra_options}
Expand Down
Loading

0 comments on commit cb9814b

Please sign in to comment.