diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 38120c6..7b157ba 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -75,6 +75,7 @@ class Config: # A dictionary of nargs for the various options. _nargs = { "lambda_values": "+", + "lambda_energy": "+", } def __init__( @@ -92,6 +93,7 @@ def __init__( h_mass_factor=1.5, num_lambda=11, lambda_values=None, + lambda_energy=None, lambda_schedule="standard_morph", charge_scale_factor=0.2, swap_end_states=False, @@ -163,6 +165,11 @@ def __init__( A list of lambda values. When specified, this takes precedence over the 'num_lambda' option. + lambda_energy: [float] + A list of lambda values at which to output energy data. If not set, + then this will be set to the same as 'lambda_values', or the values + defined by 'num_lambda' if 'lambda_values' is not set. + lambda_schedule: str Lambda schedule to use for alchemical free energy simulations. @@ -292,6 +299,7 @@ def __init__( self.timestep = timestep self.num_lambda = num_lambda self.lambda_values = lambda_values + self.lambda_energy = lambda_energy self.lambda_schedule = lambda_schedule self.charge_scale_factor = charge_scale_factor self.swap_end_states = swap_end_states @@ -650,6 +658,27 @@ def lambda_values(self, lambda_values): self._lambda_values = lambda_values + @property + def lambda_energy(self): + return self._lambda_energy + + @lambda_energy.setter + def lambda_energy(self, lambda_energy): + if lambda_energy is not None: + if not isinstance(lambda_energy, _Iterable): + raise ValueError("'lambda_energy' must be an iterable") + try: + lambda_energy = [float(x) for x in lambda_energy] + except: + raise ValueError("'lambda_energy' must be an iterable of floats") + + if not all(0 <= x <= 1 for x in lambda_energy): + raise ValueError( + "All entries in 'lambda_energy' must be between 0 and 1" + ) + + self._lambda_energy = lambda_energy + @property def lambda_schedule(self): return self._lambda_schedule diff --git a/src/somd2/runner/_dynamics.py b/src/somd2/runner/_dynamics.py index 1ce4933..cfaf9ff 100644 --- a/src/somd2/runner/_dynamics.py +++ b/src/somd2/runner/_dynamics.py @@ -43,6 +43,7 @@ def __init__( system, lambda_val, lambda_array, + lambda_energy, config, increment=0.001, device=None, @@ -61,8 +62,11 @@ def __init__( Lambda value for the simulation lambda_array : list - List of lambda values to be used for perturbation, if none won't return - reduced perturbed energies + List of lambda values to be used for simulation. + + lambda_energy: list + List of lambda values to be used for sampling energies. If None, then we + won't return reduced perturbed energies. increment : float Increment of lambda value - used for calculating the gradient @@ -107,18 +111,22 @@ def __init__( self._lambda_val = lambda_val self._lambda_array = lambda_array + self._lambda_energy = lambda_energy self._increment = increment self._device = device self._has_space = has_space self._filenames = self.create_filenames( self._lambda_array, self._lambda_val, + self._lambda_energy, self._config.output_directory, self._config.restart, ) @staticmethod - def create_filenames(lambda_array, lambda_value, output_directory, restart=False): + def create_filenames( + lambda_array, lambda_value, lambda_energy, output_directory, restart=False + ): # Create incremental file name for current restart. def increment_filename(base_filename, suffix): file_number = 0 @@ -348,10 +356,10 @@ def generate_lam_vals(lambda_base, increment): # Work out the lambda values for finite-difference gradient analysis. self._lambda_grad = generate_lam_vals(self._lambda_val, self._increment) - if self._lambda_array is None: + if self._lambda_energy is None: lam_arr = self._lambda_grad else: - lam_arr = self._lambda_array + self._lambda_grad + lam_arr = self._lambda_energy + self._lambda_grad _logger.info(f"Running dynamics at {_lam_sym} = {self._lambda_val}") diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 1a1efde..6264ee8 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -174,6 +174,12 @@ def __init__(self, system, config): for i in range(0, self._config.num_lambda) ] + # Set the lambda energy list. + if self._config.lambda_energy is not None: + self._lambda_energy = self._config.lambda_energy + else: + self._lambda_energy = self._config.lambda_values + # Work out the current hydrogen mass factor. h_mass_factor, has_hydrogen = self._get_h_mass_factor(self._system) @@ -312,6 +318,7 @@ def _check_directory(self): files = Dynamics.create_filenames( self._lambda_values, lambda_value, + self._lambda_energy, self._config.output_directory, self._config.restart, ) @@ -679,6 +686,7 @@ def _initialise_simulation(self, system, lambda_value, device=None): system, lambda_val=lambda_value, lambda_array=self._lambda_values, + lambda_energy=self._lambda_energy, config=self._config, device=device, has_space=self._has_space, @@ -939,6 +947,12 @@ def _run(sim, is_restart=False): from somd2 import __version__, _sire_version, _sire_revisionid + # Add the current lambda value to the list of lambda values and sort. + lambda_array = self._lambda_energy.copy() + if lambda_value not in lambda_array: + lambda_array.append(lambda_value) + lambda_array = sorted(lambda_array) + # Write final dataframe for the system to the energy trajectory file. # Note that sire s3 checkpoint files contain energy trajectory data, so this works even for restarts. _ = _dataframe_to_parquet( @@ -948,7 +962,7 @@ def _run(sim, is_restart=False): "somd2 version": __version__, "sire version": f"{_sire_version}+{_sire_revisionid}", "lambda": str(lambda_value), - "lambda_array": self._lambda_values, + "lambda_array": lambda_array, "lambda_grad": lambda_grad, "speed": speed, "temperature": str(self._config.temperature.value()), diff --git a/tests/runner/test_lambda_values.py b/tests/runner/test_lambda_values.py new file mode 100644 index 0000000..e28eca4 --- /dev/null +++ b/tests/runner/test_lambda_values.py @@ -0,0 +1,90 @@ +from pathlib import Path + +import tempfile +import pytest + +import sire as sr + +from somd2.runner import Runner +from somd2.config import Config +from somd2.io import * + + +def test_lambda_values(ethane_methanol): + """ + Validate that a simulation can be run with a custom list of lambda values. + """ + + with tempfile.TemporaryDirectory() as tmpdir: + mols = ethane_methanol.clone() + + config = { + "runtime": "12fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "CPU", + "max_threads": 1, + "lambda_values": [0.0, 0.5, 1.0], + } + + # Instantiate a runner using the config defined above. + runner = Runner(mols, Config(**config)) + + # Run the simulation. + runner.run() + + # Load the energy trajectory. + energy_traj, meta = parquet_to_dataframe(Path(tmpdir) / "energy_traj_0.parquet") + + # Make sure the lambda_array in the metadata is correct. This is the + # lambda_values list in the config. + assert meta["lambda_array"] == [0.0, 0.5, 1.0] + + # Make sure the second dimension of the energy trajectory is the correct + # size. This is one for the current lambda value, one for its gradient, + # and two for the additional values in the lambda_values list. + assert energy_traj.shape[1] == 4 + + +def test_lambda_energy(ethane_methanol): + """ + Validate that a simulation can sample energies at a different set of + lambda values. + """ + + with tempfile.TemporaryDirectory() as tmpdir: + mols = ethane_methanol.clone() + + config = { + "runtime": "12fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "CPU", + "max_threads": 1, + "lambda_values": [0.0, 1.0], + "lambda_energy": [0.5], + } + + # Instantiate a runner using the config defined above. + runner = Runner(mols, Config(**config)) + + # Run the simulation. + runner.run() + + # Load the energy trajectory. + energy_traj, meta = parquet_to_dataframe(Path(tmpdir) / "energy_traj_0.parquet") + + # Make sure the lambda_array in the metadata is correct. This is the + # sampled lambda plus the lambda_energy values in the config. + assert meta["lambda_array"] == [0.0, 0.5] + + # Make sure the second dimension of the energy trajectory is the correct + # size. This is one for the current lambda value, one for its gradient, + # and one for the length of lambda_energy. + assert energy_traj.shape[1] == 3