diff --git a/src/somd2/__init__.py b/src/somd2/__init__.py index 5b24bb5..6a45080 100644 --- a/src/somd2/__init__.py +++ b/src/somd2/__init__.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/_utils/__init__.py b/src/somd2/_utils/__init__.py index f2fea7a..ec25a36 100644 --- a/src/somd2/_utils/__init__.py +++ b/src/somd2/_utils/__init__.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/_utils/_somd1.py b/src/somd2/_utils/_somd1.py index 64144b0..aa379f7 100644 --- a/src/somd2/_utils/_somd1.py +++ b/src/somd2/_utils/_somd1.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/app/__init__.py b/src/somd2/app/__init__.py index 5bbf9e5..150b1c1 100644 --- a/src/somd2/app/__init__.py +++ b/src/somd2/app/__init__.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/app/_cli.py b/src/somd2/app/_cli.py index 5f74c84..27b165c 100644 --- a/src/somd2/app/_cli.py +++ b/src/somd2/app/_cli.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/config/__init__.py b/src/somd2/config/__init__.py index e5fb3f0..3507337 100644 --- a/src/somd2/config/__init__.py +++ b/src/somd2/config/__init__.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 7bb5332..faac9ad 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/io/__init__.py b/src/somd2/io/__init__.py index 9000525..efd2cf2 100644 --- a/src/somd2/io/__init__.py +++ b/src/somd2/io/__init__.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/io/_io.py b/src/somd2/io/_io.py index 0b4518d..6b66497 100644 --- a/src/somd2/io/_io.py +++ b/src/somd2/io/_io.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/runner/__init__.py b/src/somd2/runner/__init__.py index 1e755b4..69dd52a 100644 --- a/src/somd2/runner/__init__.py +++ b/src/somd2/runner/__init__.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index b45f9ac..fa857c7 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # @@ -1646,151 +1646,175 @@ def _checkpoint( is_final_block: bool Whether this is the final block of the simulation. - """ - from somd2 import __version__, _sire_version, _sire_revisionid + Returns + ------- - # Get the lambda value. - lam = self._lambda_values[index] + result: bool + Whether the checkpoint was successful. - # Get the energy trajectory. - df = system.energy_trajectory(to_alchemlyb=True, energy_unit="kT") + index: int + The index of the window or replica. - # Set the lambda values at which energies were sampled. - if lambda_energy is None: - lambda_energy = self._lambda_values + error: Exception + The exception raised during the checkpoint, if any. + """ - # Create the metadata. - metadata = { - "attrs": df.attrs, - "somd2 version": __version__, - "sire version": f"{_sire_version}+{_sire_revisionid}", - "lambda": str(lam), - "speed": speed, - "temperature": str(self._config.temperature.value()), - } + try: + from somd2 import __version__, _sire_version, _sire_revisionid + + # Get the lambda value. + lam = self._lambda_values[index] + + # Get the energy trajectory. + df = system.energy_trajectory(to_alchemlyb=True, energy_unit="kT") + + # Set the lambda values at which energies were sampled. + if lambda_energy is None: + lambda_energy = self._lambda_values + + # Create the metadata. + metadata = { + "attrs": df.attrs, + "somd2 version": __version__, + "sire version": f"{_sire_version}+{_sire_revisionid}", + "lambda": str(lam), + "speed": speed, + "temperature": str(self._config.temperature.value()), + } - # Add the lambda gradient if available. - if lambda_grad is not None: - metadata["lambda_grad"] = lambda_grad - - if is_final_block: - # Save the end-state GCMC topologies for trajectory analysis and visualisation. - # This topology contains additional water molecules that are used for GCMC - # insertion moves. - if self._config.gcmc: - mols0 = _sr.morph.link_to_reference(system) - mols1 = _sr.morph.link_to_perturbed(system) - - # Save to AMBER format. - _sr.save(mols0, self._filenames["topology0"]) - _sr.save(mols1, self._filenames["topology1"]) - - # Save to PDB format. - _sr.save( - mols0, - self._filenames["topology0"].replace(".prm7", ".pdb"), - ) - _sr.save( - mols1, - self._filenames["topology1"].replace(".prm7", ".pdb"), - ) + # Add the lambda gradient if available. + if lambda_grad is not None: + metadata["lambda_grad"] = lambda_grad + + if is_final_block: + # Save the end-state GCMC topologies for trajectory analysis and visualisation. + # This topology contains additional water molecules that are used for GCMC + # insertion moves. + if self._config.gcmc: + mols0 = _sr.morph.link_to_reference(system) + mols1 = _sr.morph.link_to_perturbed(system) - # Assemble and save the final trajectory. - if self._config.save_trajectories: - # Save the final trajectory chunk to file. - if self._save_frames and system.num_frames() > 0: - traj_filename = ( - self._filenames[index]["trajectory_chunk"] + f"{block:05d}.dcd" + # Save to AMBER format. + _sr.save(mols0, self._filenames["topology0"]) + _sr.save(mols1, self._filenames["topology1"]) + + # Save to PDB format. + _sr.save( + mols0, + self._filenames["topology0"].replace(".prm7", ".pdb"), ) _sr.save( - system.trajectory(), - traj_filename, - format=["DCD"], + mols1, + self._filenames["topology1"].replace(".prm7", ".pdb"), ) - # Create the final topology file name. - topology0 = self._filenames["topology0"] - - # Create the final trajectory file name. - traj_filename = self._filenames[index]["trajectory"] + # Assemble and save the final trajectory. + if self._config.save_trajectories: + # Save the final trajectory chunk to file. + if self._save_frames and system.num_frames() > 0: + traj_filename = ( + self._filenames[index]["trajectory_chunk"] + + f"{block:05d}.dcd" + ) + _sr.save( + system.trajectory(), + traj_filename, + format=["DCD"], + ) - # Glob for the trajectory chunks. - traj_chunks = sorted( - _glob(f"{self._filenames[index]['trajectory_chunk']}*") - ) + # Create the final topology file name. + topology0 = self._filenames["topology0"] - # If this is a restart, then we need to check for an existing - # trajectory file with the same name. If it exists and is non-empty, - # then copy it to a backup file and prepend it to the list of chunks. - if self._config.restart: - path = _Path(traj_filename) - if path.exists() and path.stat().st_size > 0: - _copyfile(traj_filename, f"{traj_filename}.prev") - traj_chunks = [f"{traj_filename}.prev"] + traj_chunks - - # Make sure there are trajectory chunks to process. - if len(traj_chunks) > 0: - # Load the topology and chunked trajectory files. - mols = _sr.load([topology0] + traj_chunks) - - # Save the final trajectory to a single file. - _sr.save(mols.trajectory(), traj_filename, format=["DCD"]) - - # Now remove the chunked trajectory files. - for chunk in traj_chunks: - _Path(chunk).unlink() - - # Add config and lambda value to the system properties. - system.set_property("config", self._config.as_dict(sire_compatible=True)) - system.set_property("lambda", lam) - - # Stream the final system to file. - _sr.stream.save(system, self._filenames[index]["checkpoint"]) - - # Create the final parquet file. - _dataframe_to_parquet( - df, - metadata=metadata, - filename=self._filenames[index]["energy_traj"], - ) + # Create the final trajectory file name. + traj_filename = self._filenames[index]["trajectory"] - else: - # Update the starting block if necessary. - if block == 0: - block = self._start_block - - # Save the current trajectory chunk to file. - if self._config.save_trajectories: - if self._save_frames and system.num_frames() > 0: - traj_filename = ( - self._filenames[index]["trajectory_chunk"] + f"{block:05d}.dcd" - ) - _sr.save( - system.trajectory(), - traj_filename, - format=["DCD"], + # Glob for the trajectory chunks. + traj_chunks = sorted( + _glob(f"{self._filenames[index]['trajectory_chunk']}*") ) - # Encode the configuration and lambda value as system properties. - system.set_property("config", self._config.as_dict(sire_compatible=True)) - system.set_property("lambda", lam) + # If this is a restart, then we need to check for an existing + # trajectory file with the same name. If it exists and is non-empty, + # then copy it to a backup file and prepend it to the list of chunks. + if self._config.restart: + path = _Path(traj_filename) + if path.exists() and path.stat().st_size > 0: + _copyfile(traj_filename, f"{traj_filename}.prev") + traj_chunks = [f"{traj_filename}.prev"] + traj_chunks + + # Make sure there are trajectory chunks to process. + if len(traj_chunks) > 0: + # Load the topology and chunked trajectory files. + mols = _sr.load([topology0] + traj_chunks) + + # Save the final trajectory to a single file. + _sr.save(mols.trajectory(), traj_filename, format=["DCD"]) + + # Now remove the chunked trajectory files. + for chunk in traj_chunks: + _Path(chunk).unlink() + + # Add config and lambda value to the system properties. + system.set_property( + "config", self._config.as_dict(sire_compatible=True) + ) + system.set_property("lambda", lam) - # Stream the checkpoint to file. - _sr.stream.save(system, self._filenames[index]["checkpoint"]) + # Stream the final system to file. + _sr.stream.save(system, self._filenames[index]["checkpoint"]) - # Create the parquet file name. - filename = self._filenames[index]["energy_traj"] + # Create the final parquet file. + _dataframe_to_parquet( + df, + metadata=metadata, + filename=self._filenames[index]["energy_traj"], + ) - # Create the parquet file. - if block == self._start_block: - _dataframe_to_parquet(df, metadata=metadata, filename=filename) - # Append to the parquet file. else: - _parquet_append( - filename, - df.iloc[-self._energy_per_block :], + # Update the starting block if necessary. + if block == 0: + block = self._start_block + + # Save the current trajectory chunk to file. + if self._config.save_trajectories: + if self._save_frames and system.num_frames() > 0: + traj_filename = ( + self._filenames[index]["trajectory_chunk"] + + f"{block:05d}.dcd" + ) + _sr.save( + system.trajectory(), + traj_filename, + format=["DCD"], + ) + + # Encode the configuration and lambda value as system properties. + system.set_property( + "config", self._config.as_dict(sire_compatible=True) ) + system.set_property("lambda", lam) + + # Stream the checkpoint to file. + _sr.stream.save(system, self._filenames[index]["checkpoint"]) + + # Create the parquet file name. + filename = self._filenames[index]["energy_traj"] + + # Create the parquet file. + if block == self._start_block: + _dataframe_to_parquet(df, metadata=metadata, filename=filename) + # Append to the parquet file. + else: + _parquet_append( + filename, + df.iloc[-self._energy_per_block :], + ) + + except Exception as e: + return index, e + + return index, None def _backup_checkpoint(self, index): """ @@ -1801,6 +1825,9 @@ def _backup_checkpoint(self, index): index : int The index of the window or replica. + + error: Exception + The exception raised during the backup, if any. """ try: @@ -1813,7 +1840,7 @@ def _backup_checkpoint(self, index): ) traj_filename = self._filenames[index]["trajectory"] except Exception as e: - return False, e + return index, e try: # Backup the existing energy trajectory file, if it exists. @@ -1824,9 +1851,9 @@ def _backup_checkpoint(self, index): str(self._filenames[index]["energy_traj"]) + ".bak", ) except Exception as e: - return False, e + return index, e - return True, None + return index, None def _save_energy_components(self, index, context): """ diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index b459676..a408c59 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # @@ -281,6 +281,14 @@ def _create_dynamics( f"Created GCMC sampler for lambda {lam:.5f} on device {device}" ) + # Log the initial position of the GCMC sphere. + if self._gcmc_samplers[i]._reference is not None: + positions = _sr.io.get_coords_array(mols) + target = self._gcmc_samplers[i]._get_target_position(positions) + _logger.info( + f"Initial GCMC sphere centre: [{target[0]:.3f}, {target[1]:.3f}, {target[2]:.3f}] A" + ) + # Create the dynamics object. try: dynamics = mols.dynamics(**dynamics_kwargs) @@ -940,7 +948,7 @@ def run(self): ] with ThreadPoolExecutor(max_workers=num_workers) as executor: try: - for result, error in executor.map( + for index, error in executor.map( self._backup_checkpoint, replicas, ): @@ -964,7 +972,7 @@ def run(self): ] with ThreadPoolExecutor(max_workers=num_workers) as executor: try: - for result, error in executor.map( + for index, error in executor.map( self._checkpoint, replicas, repeat(self._lambda_values), @@ -1525,6 +1533,15 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): is_final_block: bool Whether this is the final block. + + Returns + ------- + + index: int + The index of the replica. + + exception: Exception + The exception if the checkpoint failed. """ try: # Get the lambda value. @@ -1545,10 +1562,13 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): # Call the base class checkpoint method to save the system state. with self._lock: - super()._checkpoint( + index, error = super()._checkpoint( system, index, block, speed, is_final_block=is_final_block ) + if error is not None: + return index, error + # Delete all trajectory frames from the Sire system within the # dynamics object. dynamics._d._sire_mols.delete_all_frames() @@ -1574,10 +1594,10 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): if is_final_block: _logger.success(f"{_lam_sym} = {lam:.5f} complete") - return True, None + return index, None except Exception as e: - return False, e + return index, e @staticmethod @_njit diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index da7f3fa..b6ac952 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -1,7 +1,7 @@ ###################################################################### # SOMD2: GPU accelerated alchemical free-energy engine. # -# Copyright: 2023-2025 +# Copyright: 2023-2026 # # Authors: The OpenBioSim Team # @@ -431,6 +431,14 @@ def generate_lam_vals(lambda_base, increment=0.001): # Get the GCMC system. system = gcmc_sampler.system() + # Log the initial position of the GCMC sphere. + if gcmc_sampler._reference is not None: + positions = _sr.io.get_coords_array(system) + target = gcmc_sampler._get_target_position(positions) + _logger.info( + f"Initial GCMC sphere centre: [{target[0]:.3f}, {target[1]:.3f}, {target[2]:.3f}] A" + ) + else: gcmc_sampler = None @@ -770,10 +778,13 @@ def generate_lam_vals(lambda_base, increment=0.001): # in a consistent state if read by another process. with lock.acquire(timeout=self._config.timeout.to("seconds")): # Backup any existing checkpoint files. - self._backup_checkpoint(index) + index, error = self._backup_checkpoint(index) + + if error is not None: + raise error # Write the checkpoint files. - self._checkpoint( + index, error = self._checkpoint( system, index, block, @@ -783,6 +794,9 @@ def generate_lam_vals(lambda_base, increment=0.001): is_final_block=is_final_block, ) + if error is not None: + raise error + # Delete all trajectory frames from the Sire system within the # dynamics object. dynamics._d._sire_mols.delete_all_frames() @@ -957,10 +971,15 @@ def generate_lam_vals(lambda_base, increment=0.001): # in a consistent state if read by another process. with lock.acquire(timeout=self._config.timeout.to("seconds")): # Backup any existing checkpoint files. - self._backup_checkpoint(index) + index, error = self._backup_checkpoint(index) + + if error is not None: + msg = f"Checkpoint backup failed for {_lam_sym} = {lambda_value:.5f}: {error}" + _logger.error(msg) + raise RuntimeError(msg) # Write the checkpoint files. - self._checkpoint( + index, error = self._checkpoint( system, index, 0, @@ -970,6 +989,11 @@ def generate_lam_vals(lambda_base, increment=0.001): is_final_block=True, ) + if error is not None: + msg = f"Checkpoint failed for {_lam_sym} = {lambda_value:.5f}: {error}" + _logger.error(msg) + raise RuntimeError(msg) + _logger.success( f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1" )