From 3384ca75512b83dfc2f326bdb8076eb24f0f2b4e Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 28 Aug 2026 10:52:33 +0100 Subject: [PATCH] Accept stream file paths for restraints and lambda schedules. --- README.md | 60 ++++++++++++ src/somd2/config/_config.py | 180 ++++++++++++++++++++++++++---------- tests/runner/test_config.py | 103 +++++++++++++++++++++ 3 files changed, 296 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index c9f1e65..828a9f0 100644 --- a/README.md +++ b/README.md @@ -205,6 +205,48 @@ the `--max-gpus` option can be set, for example setting `--max-gpus 2` while `CUDA_VISIBLE_DEVICES` are set as above would restrict SOMD2 to using only GPUs 0 and 1. +## Restarting + +A simulation can be continued from the files in its output directory using the +`--restart` option: + +``` +somd2 perturbable_system.bss --restart --output-directory output +``` + +Each λ window (or replica) resumes from its most recent checkpoint. The +configuration used for the original run is written to `config.yaml` in the +output directory, controlled by `--write-config`, which is enabled by default. +This file is required in order to restart, since the current configuration is +validated against it. + +Only a limited set of options may be changed on restart. Broadly, anything that +would change the perturbation or the Hamiltonian is fixed, whereas options +controlling how long to run for, what to write out, and which hardware to use +can be varied. The most useful of these is `--runtime`, which allows a completed +simulation to be extended. SOMD2 will tell you which option is at fault if you +change one that isn't allowed. + +> [!NOTE] +> If the most recent checkpoint files are incomplete or corrupt, for example +> when recovering from a crash, pass `--use-backup` to restart from the last +> but one checkpoint instead. + +## Hydrogen mass repartitioning + +By default SOMD2 applies hydrogen mass repartitioning (HMR), scaling hydrogen +masses by the factor given by `--h-mass-factor` (default 1.5). This is what +allows the default `--timestep` of 4 fs. + +If the masses of your input system have already been repartitioned, or you want +to use a different repartitioning scheme, pass `--no-hmr` so that the masses of +the input system are used as they are. + +> [!NOTE] +> A 4 fs timestep is not stable without repartitioning, so if you disable HMR +> you will need to reduce `--timestep` accordingly, or supply a system that has +> already been repartitioned. + ## Replica exchange SOMD2 supports Hamiltonian replica exchange (HREX) simulations, which can be @@ -597,6 +639,24 @@ More details on MPS, including tuning options, can be found in the following SOMD2 can also be used as a Python API, allowing it to be embedded within other Python scripts. +A few options take objects rather than values, so cannot be set directly on the +command line. A custom lambda schedule can be passed to `lambda_schedule` as a +`sire.cas.LambdaSchedule`, rather than one of the named schedules, and +user-defined restraints can be passed to `restraints`. + +Both options can also be set via a YAML configuration file, where they are +stored as a hex string of the serialised object. This is the form written to +`config.yaml`, so the simplest way to obtain one is to configure the option in +Python, run a simulation, and re-use the value from the resulting file. + +Alternatively, both accept a path to a [Sire](https://github.com/OpenBioSim/sire) +stream file containing the serialised object, which can be written with +`sire.stream.save`: + +``` +somd2 perturbable_system.bss --lambda-schedule my_schedule.s3 --restraints my_restraints.s3 +``` + ## Known issues If using the regular `Runner` class via the Python API, then you will need to diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 28f2470..cfa0a99 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -78,11 +78,17 @@ class Config: "precision": ["single", "mixed", "double"], } + # Options that advertise a set of choices, but which also accept other + # forms, e.g. the path to a stream file. These are validated by the setter, + # rather than by argparse. + _open_choices = ["lambda_schedule"] + # A dictionary of nargs for the various options. _nargs = { "lambda_values": "+", "lambda_energy": "+", "rest2_scale": "+", + "restraints": "+", } def __init__( @@ -239,8 +245,10 @@ def __init__( 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. + lambda_schedule: str, sire.cas.LambdaSchedule + Lambda schedule to use for alchemical free energy simulations. This + can be the name of one of the standard schedules, or the path to a + Sire stream file containing a custom LambdaSchedule. charge_scale_factor: float Factor by which to scale charges for charge scaled morph. @@ -256,9 +264,10 @@ def __init__( The soft-core shift-delta parameter. This is used to soften the Lennard-Jones interaction. - restraints: sire.mm._MM.Restraints - A single set of restraints, or a list of sets of restraints that - will be applied to the atoms during the simulation. + restraints: str, sire.mm._MM.Restraints + One or more paths to Sire stream files containing the sets of + restraints that will be applied to the atoms during the simulation. + A stream file may hold a single set, or a list of sets. constraint: str Constraint type to use for non-perturbable molecules. @@ -1180,44 +1189,51 @@ def lambda_schedule(self, lambda_schedule): "'lambda_schedule' must be of type 'str' or 'LambdaSchedule' object" ) if isinstance(lambda_schedule, str): - # Strip whitespace and convert to lower case. - lambda_schedule = lambda_schedule.strip().lower() - if lambda_schedule == "standard_morph": + # Strip whitespace. The keyword comparison is made against a + # lower case copy, since the string may also be a path, which + # is case sensitive. + lambda_schedule = lambda_schedule.strip() + keyword = lambda_schedule.lower() + if keyword == "standard_morph": self._lambda_schedule = _LambdaSchedule.standard_morph() self._lambda_schedule_name = "standard_morph" - elif lambda_schedule == "charge_scaled_morph": + elif keyword == "charge_scaled_morph": self._lambda_schedule = _LambdaSchedule.charge_scaled_morph(0.2) self._lambda_schedule_name = "charge_scaled_morph" - elif lambda_schedule == "ring_break_morph": + elif keyword == "ring_break_morph": from .._utils._schedules import ( ring_break_morph as _ring_break_morph, ) self._lambda_schedule = _ring_break_morph() self._lambda_schedule_name = "ring_break_morph" - elif lambda_schedule == "reverse_ring_break_morph": + elif keyword == "reverse_ring_break_morph": from .._utils._schedules import ( reverse_ring_break_morph as _reverse_ring_break_morph, ) self._lambda_schedule = _reverse_ring_break_morph() self._lambda_schedule_name = "reverse_ring_break_morph" - elif lambda_schedule == "annihilate": + elif keyword == "annihilate": self._lambda_schedule = None self._lambda_schedule_name = "annihilate" - elif lambda_schedule == "decouple": + elif keyword == "decouple": self._lambda_schedule = None self._lambda_schedule_name = "decouple" else: - try: - self._lambda_schedule = self._from_hex(lambda_schedule) - self._lambda_schedule_name = None - except Exception: + schedule = self._from_string( + lambda_schedule, + "lambda_schedule", + hint=", or one of the following strings: " + f"{', '.join(self._choices['lambda_schedule'])}", + ) + if not isinstance(schedule, _LambdaSchedule): raise ValueError( - "Unable to deserialise 'lambda_schedule'. Ensure that this is a " - "hex string representation of a valid LambdaSchedule object, or " - f"one of the following strings: {', '.join(self._choices['lambda_schedule'])}" + f"'lambda_schedule' deserialised to a " + f"'{type(schedule).__name__}', not a 'LambdaSchedule'." ) + self._lambda_schedule = schedule + self._lambda_schedule_name = None else: self._lambda_schedule = lambda_schedule self._lambda_schedule_name = None @@ -1303,32 +1319,34 @@ def restraints(self): @restraints.setter def restraints(self, restraints): - # If not supplied as a list, convert to a list. + # If not supplied as a list, convert to a list. Note that a string is + # itself iterable, so must be wrapped explicitly. if restraints is not None: - if not isinstance(restraints, _Iterable): + if isinstance(restraints, str) or not isinstance(restraints, _Iterable): restraints = [restraints] - # Check that all restraints are of the correct type. - deserialised_restraints = [] + # Resolve each entry, keeping objects and deserialised strings in + # the order they were given. + resolved_restraints = [] for restraint in restraints: - if isinstance(restraint, _sr.mm._MM.Restraints): - continue - elif isinstance(restraint, str): - try: - restraint = self._from_hex(restraint) - except Exception: - raise ValueError( - "Unable to deserialise restraint. Ensure that this " - "is a hex string representation of a valid sire.mm._MM.Restraints object." - ) - deserialised_restraints.append(restraint) + if isinstance(restraint, str): + restraint = self._from_string(restraint.strip(), "restraints") + + # A stream file may hold a list of sets of restraints, e.g. the + # pair used for a ring-breaking perturbation. + if isinstance(restraint, _Iterable): + resolved_restraints.extend(restraint) else: + resolved_restraints.append(restraint) + + # Check that all restraints are of the correct type. + for restraint in resolved_restraints: + if not isinstance(restraint, _sr.mm._MM.Restraints): raise ValueError( "'restraints' must be a sire.mm._MM.Restraints object, or a list of these objects." ) - if len(deserialised_restraints) > 0: - restraints = deserialised_restraints + restraints = resolved_restraints self._restraints = restraints @@ -2622,6 +2640,61 @@ def _from_hex(hex): return obj + @classmethod + def _from_string(cls, string, name, hint=""): + """ + Internal method to deserialise a Sire object from a string, which can + either be the path to a stream file, or the hex string representation + of the serialised object. + + Parameters + ---------- + + string: str + The path to a stream file, or a hex string representation of the + Sire object. + + name: str + The name of the option being set, used for error messages. + + hint: str + An additional clause appended to the error message, e.g. listing + the keywords that the option also accepts. + + Returns + ------- + + obj: + The deserialised Sire object. + """ + from pathlib import Path as _Path + + # Work out whether this is a path to an existing file. A hex string can + # exceed the maximum filename length, which raises rather than simply + # returning False on some platforms. + try: + is_file = _Path(string).is_file() + except Exception: + is_file = False + + if is_file: + from sire.stream import load + + try: + return load(string) + except Exception as e: + raise ValueError( + f"Unable to load '{name}' from stream file '{string}': {e}" + ) + else: + try: + return cls._from_hex(string) + except Exception: + raise ValueError( + f"Unable to interpret '{name}'. Expected the path to a Sire " + f"stream file, or a hex string of a serialised object{hint}." + ) + def __getstate__(self): """ Hex-encode the same fields that to_yaml()/from_yaml() already @@ -2658,6 +2731,7 @@ def _create_parser(cls): import argparse import inspect + import re # Inspect the signature to get the parameters. sig = inspect.signature(Config.__init__) @@ -2665,7 +2739,7 @@ def _create_parser(cls): params = { key: value for key, value in params.items() - if key not in ["self", "args", "kwargs", "restraints"] + if key not in ["self", "args", "kwargs"] } # Get the docstring. @@ -2682,7 +2756,7 @@ def _create_parser(cls): # Loop over all lines in the docstring until we find the parameter. for line in doc: line = line.strip() - if line.startswith(param): + if re.match(rf"{re.escape(param)}\s*:", line): found_param = True elif found_param: if line == "": @@ -2724,14 +2798,26 @@ def _create_parser(cls): # This parameter has choices. if param in cls._choices: - parser.add_argument( - f"--{cli_param}", - type=typ, - default=params[param].default, - choices=cls._choices[param], - help=help[param], - required=False, - ) + # Other forms are also accepted, so advertise the choices in the + # help text, but leave the validation to the setter. + if param in cls._open_choices: + parser.add_argument( + f"--{cli_param}", + type=typ, + default=params[param].default, + metavar="{" + ",".join(cls._choices[param]) + "}", + help=help[param], + required=False, + ) + else: + parser.add_argument( + f"--{cli_param}", + type=typ, + default=params[param].default, + choices=cls._choices[param], + help=help[param], + required=False, + ) # This is a standard parameter. else: if typ == bool: diff --git a/tests/runner/test_config.py b/tests/runner/test_config.py index 27d8e4a..ce0168e 100644 --- a/tests/runner/test_config.py +++ b/tests/runner/test_config.py @@ -123,3 +123,106 @@ def test_morse_restraint_options(): with pytest.raises(ValueError, match="units are invalid"): Config(morse_soft_force_constant="125 kcal mol-1") + + +def test_lambda_schedule_input_forms(): + """Validate that all supported lambda schedule input forms are accepted.""" + import os + + import pytest + + schedule = sr.cas.LambdaSchedule.standard_morph() + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "schedule.s3") + sr.stream.save(schedule, path) + + # A named schedule, which is case insensitive. + config = Config(lambda_schedule="DECOUPLE") + assert config._lambda_schedule_name == "decouple" + + # The path to a stream file. + config = Config(lambda_schedule=path) + assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule) + assert config._lambda_schedule_name is None + + # A hex string of the serialised object. + config = Config(lambda_schedule=Config._to_hex(schedule)) + assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule) + + # The object itself. + config = Config(lambda_schedule=schedule) + assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule) + + # Anything else is rejected. + with pytest.raises(ValueError, match="Unable to interpret"): + Config(lambda_schedule="not_a_schedule") + + # A stream file holding the wrong type of object. + wrong_path = os.path.join(tmpdir, "wrong.s3") + sr.stream.save(sr.cas.Symbol("x"), wrong_path) + with pytest.raises(ValueError, match="not a 'LambdaSchedule'"): + Config(lambda_schedule=wrong_path) + + +def test_restraints_input_forms(): + """Validate that all supported restraint input forms are accepted.""" + import os + + import pytest + + mols = sr.load_test_files("ala.top", "ala.crd") + restraint0 = sr.restraints.positional(mols, atoms="atomidx 0") + restraint1 = sr.restraints.positional(mols, atoms="atomidx 1") + + with tempfile.TemporaryDirectory() as tmpdir: + path0 = os.path.join(tmpdir, "restraint0.s3") + both_path = os.path.join(tmpdir, "both.s3") + sr.stream.save(restraint0, path0) + sr.stream.save([restraint0, restraint1], both_path) + + # A single object, or a list of objects. + assert len(Config(restraints=restraint0).restraints) == 1 + assert len(Config(restraints=[restraint0, restraint1]).restraints) == 2 + + # The path to a stream file, or a list of paths. + assert len(Config(restraints=path0).restraints) == 1 + assert len(Config(restraints=[path0, path0]).restraints) == 2 + + # A stream file holding a list of sets of restraints. + assert len(Config(restraints=both_path).restraints) == 2 + + # A hex string of the serialised object. + assert len(Config(restraints=Config._to_hex(restraint0)).restraints) == 1 + + # Objects and paths can be mixed, and all are retained. + config = Config(restraints=[restraint0, path0]) + assert len(config.restraints) == 2 + assert all( + isinstance(restraint, sr.mm._MM.Restraints) + for restraint in config.restraints + ) + + # Anything else is rejected. + with pytest.raises(ValueError, match="Unable to interpret"): + Config(restraints="not_a_restraint") + + # A stream file holding the wrong type of object. + wrong_path = os.path.join(tmpdir, "wrong.s3") + sr.stream.save(sr.cas.LambdaSchedule.standard_morph(), wrong_path) + with pytest.raises(ValueError, match="must be a sire.mm._MM.Restraints"): + Config(restraints=wrong_path) + + +def test_help_text_scraping(): + """Validate that help text isn't truncated by the parameter name.""" + parser = Config._create_parser() + + for action in parser._actions: + if action.dest == "restraints": + break + + # The description wraps onto a line starting with the parameter name, which + # must not be mistaken for the start of the next parameter. + assert "applied to the atoms" in action.help + assert "a list of sets" in action.help