From 5e73871556f763da4587679d42296e09c5d888f7 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 17 May 2024 10:59:13 +0100 Subject: [PATCH 1/8] Add support for saving energy contribution from each force. --- src/somd2/config/_config.py | 16 +++++++++++ src/somd2/runner/_dynamics.py | 54 ++++++++++++++++++++++++++++++++++- src/somd2/runner/_runner.py | 4 ++- 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 6f6afbc..51062f1 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -124,6 +124,7 @@ def __init__( overwrite=False, somd1_compatibility=False, pert_file=None, + save_energy_components=False, ): """ Constructor. @@ -281,6 +282,10 @@ def __init__( pert_file: str The path to a SOMD1 perturbation file to apply to the reference system. When set, this will automatically set 'somd1_compatibility' to True. + + save_energy_components: bool + Whether to save the energy contribution for each force when checkpointing. + This is useful when debugging crashes. """ # Setup logger before doing anything else @@ -327,6 +332,7 @@ def __init__( self.restart = restart self.somd1_compatibility = somd1_compatibility self.pert_file = pert_file + self.save_energy_components = save_energy_components self.write_config = write_config @@ -1201,6 +1207,16 @@ def pert_file(self, pert_file): if pert_file is not None: self._somd1_compatibility = True + @property + def save_energy_components(self): + return self._save_energy_components + + @save_energy_components.setter + def save_energy_components(self, save_energy_components): + if not isinstance(save_energy_components, bool): + raise ValueError("'save_energy_components' must be of type 'bool'") + self._save_energy_components = save_energy_components + @property def output_directory(self): return self._output_directory diff --git a/src/somd2/runner/_dynamics.py b/src/somd2/runner/_dynamics.py index 2a54d27..e7eeabb 100644 --- a/src/somd2/runner/_dynamics.py +++ b/src/somd2/runner/_dynamics.py @@ -127,6 +127,9 @@ def __init__( self._config.restart, ) + self._nrg_sample = 0 + self._nrg_file = "energy_components.txt" + @staticmethod def create_filenames(lambda_array, lambda_value, output_directory, restart=False): # Create incremental file name for current restart. @@ -153,6 +156,7 @@ def increment_filename(base_filename, suffix): filenames["energy_traj"] = f"energy_traj_{lam}.parquet" filenames["trajectory"] = f"traj_{lam}.dcd" filenames["trajectory_chunk"] = f"traj_{lam}_" + filenames["energy_components"] = f"energy_components_{lam}.txt" if restart: filenames["config"] = increment_filename("config", "yaml") else: @@ -371,7 +375,7 @@ def generate_lam_vals(lambda_base, increment): ) if self._config.checkpoint_frequency.value() > 0.0: - # Calculate the number of blocks and the remaineder time. + # Calculate the number of blocks and the remainder time. frac = ( self._config.runtime.value() / self._config.checkpoint_frequency.value() ) @@ -409,6 +413,10 @@ def generate_lam_vals(lambda_base, increment): # Checkpoint. try: + # Save the energy contribution for each force. + if self._config.save_energy_components: + self._save_energy_components() + # Set to the current block number if this is a restart. if x == 0: x = self._current_block @@ -584,3 +592,47 @@ def get_timing(self): def _cleanup(self): del self._dyn + + def _save_energy_components(self): + + from copy import deepcopy + import openmm + + # Get the current context and system. + context = self._dyn._d._omm_mols + system = deepcopy(context.getSystem()) + + # Add each force to a unique group. + for i, f in enumerate(system.getForces()): + f.setForceGroup(i) + + # Create a new context. + new_context = openmm.Context(system, deepcopy(context.getIntegrator())) + new_context.setPositions(context.getState(getPositions=True).getPositions()) + + header = f"{'# Sample':>10}" + record = f"{self._nrg_sample:>10}" + + # Process the records. + for i, f in enumerate(system.getForces()): + state = new_context.getState(getEnergy=True, groups={i}) + header += f"{f.getName():>25}" + record += f"{state.getPotentialEnergy().value_in_unit(openmm.unit.kilocalories_per_mole):>25.2f}" + + # Write to file. + if self._nrg_sample == 0: + with open( + self._config.output_directory / self._filenames["energy_components"], + "w", + ) as f: + f.write(header + "\n") + f.write(record + "\n") + else: + with open( + self._config.output_directory / self._filenames["energy_components"], + "a", + ) as f: + f.write(record + "\n") + + # Increment the sample number. + self._nrg_sample += 1 diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index d97a4ed..30e4dad 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -969,5 +969,7 @@ def _run(sim, is_restart=False): filename=self._fnames[lambda_value]["energy_traj"], ) del system - _logger.success(f"{_lam_sym} = {lambda_value} complete, speed = {speed:.2f} ns day-1") + _logger.success( + f"{_lam_sym} = {lambda_value} complete, speed = {speed:.2f} ns day-1" + ) return True From 3e51375433af8840faf69c4234b581dff466cade Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 21 May 2024 12:27:02 +0100 Subject: [PATCH 2/8] Add functionality to fix LJ sigmas for ghost atoms. --- src/somd2/runner/_runner.py | 4 +- src/somd2/runner/_somd1.py | 91 ++++++++++++++++++++++++++++++------- 2 files changed, 78 insertions(+), 17 deletions(-) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 30e4dad..40637c4 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -151,12 +151,14 @@ def __init__(self, system, config): # Only check for light atoms by the maxium end state mass if running # in SOMD1 compatibility mode. Ghost atoms are considered light when - # adding bond constraints. + # adding bond constraints. Also fix the LJ sigma for ghost atoms so + # it isn't scaled to zero. self._config._extra_args["ghosts_are_light"] = True self._config._extra_args["check_for_h_by_max_mass"] = True self._config._extra_args["check_for_h_by_mass"] = False self._config._extra_args["check_for_h_by_element"] = False self._config._extra_args["check_for_h_by_ambertype"] = False + self._config._extra_args["fix_ghost_sigmas"] = True # Check for a periodic space. self._check_space() diff --git a/src/somd2/runner/_somd1.py b/src/somd2/runner/_somd1.py index d6444b4..24c1278 100644 --- a/src/somd2/runner/_somd1.py +++ b/src/somd2/runner/_somd1.py @@ -59,8 +59,9 @@ def _make_compatible(system): except KeyError: raise KeyError("No perturbable molecules in the system") - # Store a dummy element. - dummy = _SireMol.Element("Xx") + # Store a dummy element and ambertype. + element_dummy = _SireMol.Element("Xx") + amber_dummy = "du" for mol in pert_mols: # Store the molecule info. @@ -69,9 +70,43 @@ def _make_compatible(system): # Get an editable version of the molecule. edit_mol = mol.edit() - ########################## - # First process the bonds. - ########################## + ##################################### + # First fix the ghost atom LJ sigmas. + ##################################### + + for atom in mol.atoms(): + # Lambda = 0 state is a dummy, use sigma from the lambda = 1 state. + if ( + atom.property("element0") == element_dummy + or atom.property("ambertype0") == amber_dummy + ): + lj0 = atom.property("LJ0") + lj1 = atom.property("LJ1") + edit_mol = ( + edit_mol.atom(atom.index()) + .set_property( + "LJ0", _SireMM.LJParameter(lj1.sigma(), lj0.epsilon()) + ) + .molecule() + ) + # Lambda = 1 state is a dummy, use sigma from the lambda = 0 state. + elif ( + atom.property("element1") == element_dummy + or atom.property("ambertype1") == amber_dummy + ): + lj0 = atom.property("LJ0") + lj1 = atom.property("LJ1") + edit_mol = ( + edit_mol.atom(atom.index()) + .set_property( + "LJ1", _SireMM.LJParameter(lj0.sigma(), lj1.epsilon()) + ) + .molecule() + ) + + ######################## + # Now process the bonds. + ######################## new_bonds0 = _SireMM.TwoAtomFunctions(mol.info()) new_bonds1 = _SireMM.TwoAtomFunctions(mol.info()) @@ -534,17 +569,26 @@ def _has_dummy(mol, idxs, is_lambda1=False): Whether a dummy atom is present. """ - # Set the element property associated with the end state. + # We need to check by ambertype too since this molecule may have been + # created via sire.morph.create_from_pertfile, in which case the element + # property will have been set to the end state with the largest mass, i.e. + # may no longer by a dummy. if is_lambda1: - prop = "element1" + element_prop = "element1" + ambertype_prop = "ambertype1" else: - prop = "element0" + element_prop = "element0" + ambertype_prop = "ambertype0" - dummy = _SireMol.Element(0) + element_dummy = _SireMol.Element(0) + ambertype_dummy = "du" # Check whether an of the atoms is a dummy. for idx in idxs: - if mol.atom(idx).property(prop) == dummy: + if ( + mol.atom(idx).property(element_prop) == element_dummy + or mol.atom(idx).property(ambertype_prop) == ambertype_dummy + ): return True return False @@ -573,21 +617,36 @@ def _is_dummy(mol, idxs, is_lambda1=False): Whether each atom is a dummy. """ - # Set the element property associated with the end state. + # We need to check by ambertype too since this molecule may have been + # created via sire.morph.create_from_pertfile, in which case the element + # property will have been set to the end state with the largest mass, i.e. + # may no longer by a dummy. if is_lambda1: - prop = "element1" + element_prop = "element1" + ambertype_prop = "ambertype1" else: - prop = "element0" + element_prop = "element0" + ambertype_prop = "ambertype0" - # Store a dummy element. - dummy = _SireMol.Element(0) + if is_lambda1: + element_prop = "element1" + ambertype_prop = "ambertype1" + else: + element_prop = "element0" + ambertype_prop = "ambertype0" + + element_dummy = _SireMol.Element(0) + ambertype_dummy = "du" # Initialise a list to store the state of each atom. is_dummy = [] # Check whether each of the atoms is a dummy. for idx in idxs: - is_dummy.append(mol.atom(idx).property(prop) == dummy) + is_dummy.append( + mol.atom(idx).property(element_prop) == element_dummy + or mol.atom(idx).property(ambertype_prop) == ambertype_dummy + ) return is_dummy From e7717ad6e3e3c1d01d3846d64b43c29b0a671271 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 23 May 2024 14:24:20 +0100 Subject: [PATCH 3/8] Need to re-link properties to the reference state since LJ is lost. --- src/somd2/runner/_runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 40637c4..8adf5a1 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -114,6 +114,7 @@ def __init__(self, system, config): _logger.info("Applying SOMD1 perturbation compatibility.") self._system = _make_compatible(self._system) + self._system = _morph.link_to_reference(self._system) # Next, swap the water topology so that it is in AMBER format. From c0507311ce69bc3231fa3bd2bf3084f472d2633e Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 23 May 2024 15:30:46 +0100 Subject: [PATCH 4/8] Formatting tweak. [ci skip]:wq --- src/somd2/runner/_somd1.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/somd2/runner/_somd1.py b/src/somd2/runner/_somd1.py index 24c1278..2c19737 100644 --- a/src/somd2/runner/_somd1.py +++ b/src/somd2/runner/_somd1.py @@ -61,7 +61,7 @@ def _make_compatible(system): # Store a dummy element and ambertype. element_dummy = _SireMol.Element("Xx") - amber_dummy = "du" + ambertype_dummy = "du" for mol in pert_mols: # Store the molecule info. @@ -78,7 +78,7 @@ def _make_compatible(system): # Lambda = 0 state is a dummy, use sigma from the lambda = 1 state. if ( atom.property("element0") == element_dummy - or atom.property("ambertype0") == amber_dummy + or atom.property("ambertype0") == ambertype_dummy ): lj0 = atom.property("LJ0") lj1 = atom.property("LJ1") @@ -92,7 +92,7 @@ def _make_compatible(system): # Lambda = 1 state is a dummy, use sigma from the lambda = 0 state. elif ( atom.property("element1") == element_dummy - or atom.property("ambertype1") == amber_dummy + or atom.property("ambertype1") == ambertype_dummy ): lj0 = atom.property("LJ0") lj1 = atom.property("LJ1") From e0686bec6667a3e1133630ff9c8ec1e1c68cdfbc Mon Sep 17 00:00:00 2001 From: Julien Michel Date: Wed, 29 May 2024 16:19:35 +0100 Subject: [PATCH 5/8] Update README.md fixed a typo in the installation instructions --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 77d9bbe..f08accb 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [Ope First create a conda environment using the provided environment file: ``` -mamba create -f environment.yaml +mamba env create -f environment.yaml ``` (We recommend using [Miniforge](https://github.com/conda-forge/miniforge).) From 882ec42cfe58c7372f39d28d0d9dcc8bb1eb5bf7 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 30 May 2024 09:12:09 +0100 Subject: [PATCH 6/8] Switch to conda over mamba. [ci skip] --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index f08accb..c414630 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [Ope First create a conda environment using the provided environment file: ``` -mamba env create -f environment.yaml +conda env create -f environment.yaml ``` (We recommend using [Miniforge](https://github.com/conda-forge/miniforge).) @@ -19,7 +19,7 @@ mamba env create -f environment.yaml Now install `somd2` into the environment: ``` -mamba activate somd2 +conda activate somd2 pip install --editable . ``` From d1c5a279d83ac6c49f5407f611d966c46c5d9ea3 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 30 May 2024 09:12:32 +0100 Subject: [PATCH 7/8] Make sure search returns a SelectorMol iterable. --- src/somd2/runner/_somd1.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/somd2/runner/_somd1.py b/src/somd2/runner/_somd1.py index 2c19737..063a76e 100644 --- a/src/somd2/runner/_somd1.py +++ b/src/somd2/runner/_somd1.py @@ -681,7 +681,7 @@ def _apply_pert(system, pert_file): from sire import morph as _morph # Get the non-water molecules in the system. - non_waters = system["not water"] + non_waters = system["not water"].molecules() # Try to apply the perturbation to each non-water molecule. is_pert = False From 8cfc3c27544e77e393a2b698a73d0bdbb9d1ff4d Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 30 May 2024 09:18:46 +0100 Subject: [PATCH 8/8] Make sure all LJ sigma values are non-zero. --- src/somd2/runner/_somd1.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/src/somd2/runner/_somd1.py b/src/somd2/runner/_somd1.py index 063a76e..b21469c 100644 --- a/src/somd2/runner/_somd1.py +++ b/src/somd2/runner/_somd1.py @@ -70,18 +70,21 @@ def _make_compatible(system): # Get an editable version of the molecule. edit_mol = mol.edit() - ##################################### - # First fix the ghost atom LJ sigmas. - ##################################### + ################################## + # First fix zero LJ sigmas values. + ################################## + + # Create a null LJParameter. + null_lj = _SireMM.LJParameter() for atom in mol.atoms(): - # Lambda = 0 state is a dummy, use sigma from the lambda = 1 state. - if ( - atom.property("element0") == element_dummy - or atom.property("ambertype0") == ambertype_dummy - ): - lj0 = atom.property("LJ0") - lj1 = atom.property("LJ1") + # Get the end state LJ sigma values. + lj0 = atom.property("LJ0") + lj1 = atom.property("LJ1") + + # Lambda = 0 state has a zero sigma value. + if lj0.sigma() == null_lj.sigma(): + # Use the sigma value from the lambda = 1 state. edit_mol = ( edit_mol.atom(atom.index()) .set_property( @@ -89,13 +92,10 @@ def _make_compatible(system): ) .molecule() ) - # Lambda = 1 state is a dummy, use sigma from the lambda = 0 state. - elif ( - atom.property("element1") == element_dummy - or atom.property("ambertype1") == ambertype_dummy - ): - lj0 = atom.property("LJ0") - lj1 = atom.property("LJ1") + + # Lambda = 1 state has a zero sigma value. + if lj1.sigma() == null_lj.sigma(): + # Use the sigma value from the lambda = 0 state. edit_mol = ( edit_mol.atom(atom.index()) .set_property(