From d8a6f2e70728d8986135380a48570a0fa3251b65 Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 29 Oct 2024 14:35:11 -0500 Subject: [PATCH 01/73] openmm collector reporter. AAE example. coordinate collector for AAE. --- deepdrivewe/ai/__init__.py | 47 +-- deepdrivewe/ai/aae.py | 261 +++++++++++++ deepdrivewe/ai/cvae.py | 46 +++ .../examples/openmm_aae_ddwe/inference.py | 155 ++++++++ deepdrivewe/examples/openmm_aae_ddwe/main.py | 249 +++++++++++++ .../examples/openmm_aae_ddwe/simulate.py | 110 ++++++ deepdrivewe/examples/openmm_aae_ddwe/train.py | 75 ++++ .../examples/openmm_ntl9_ddwe/inference.py | 4 +- .../examples/openmm_ntl9_ddwe/simulate.py | 33 +- deepdrivewe/simulation/openmm.py | 346 +++++++++++++++++- pyproject.toml | 2 +- 11 files changed, 1253 insertions(+), 75 deletions(-) create mode 100644 deepdrivewe/ai/aae.py create mode 100644 deepdrivewe/examples/openmm_aae_ddwe/inference.py create mode 100644 deepdrivewe/examples/openmm_aae_ddwe/main.py create mode 100644 deepdrivewe/examples/openmm_aae_ddwe/simulate.py create mode 100644 deepdrivewe/examples/openmm_aae_ddwe/train.py diff --git a/deepdrivewe/ai/__init__.py b/deepdrivewe/ai/__init__.py index 7625ecf..324290d 100644 --- a/deepdrivewe/ai/__init__.py +++ b/deepdrivewe/ai/__init__.py @@ -2,49 +2,10 @@ from __future__ import annotations -from functools import lru_cache -from pathlib import Path - # Forward imports +from deepdrivewe.ai.aae import AdversarialAE +from deepdrivewe.ai.aae import AdversarialAEConfig +from deepdrivewe.ai.aae import warmstart_aae from deepdrivewe.ai.cvae import ConvolutionalVAE from deepdrivewe.ai.cvae import ConvolutionalVAEConfig -from deepdrivewe.ai.utils import LatentSpaceHistory - - -@lru_cache(maxsize=1) -def warmstart_model( - config_path: Path, - checkpoint_path: Path, -) -> tuple[ConvolutionalVAE, LatentSpaceHistory]: - """Load the model once and then return a cached version. - - Parameters - ---------- - config_path : Path - The path to the model configuration file. - checkpoint_path : Path - The path to the model checkpoint file. - - Returns - ------- - ConvolutionalVAE - The ConvolutionalVAE model. - LatentSpaceHistory - The latent space history. - """ - # Print the warmstart message - print(f'Cold start model from checkpoint {checkpoint_path}') - - # Load the model configuration - model_config = ConvolutionalVAEConfig.from_yaml(config_path) - - # Load the model - model = ConvolutionalVAE( - model_config, - checkpoint_path=checkpoint_path, - ) - - # Initialize the latent space history - history = LatentSpaceHistory() - - return model, history +from deepdrivewe.ai.cvae import warmstart_cvae diff --git a/deepdrivewe/ai/aae.py b/deepdrivewe/ai/aae.py new file mode 100644 index 0000000..91fc31d --- /dev/null +++ b/deepdrivewe/ai/aae.py @@ -0,0 +1,261 @@ +"""Adversarial Autoencoder for Contact Maps.""" + +from __future__ import annotations + +from functools import lru_cache +from pathlib import Path + +import numpy as np +import pandas as pd +from natsort import natsorted +from pydantic import Field + +from deepdrivewe import BaseModel +from deepdrivewe.ai.utils import LatentSpaceHistory + + +class AdversarialAEConfig(BaseModel): + """Settings for mdlearn 3dAAE model.""" + + scalar_dset_names: list[str] = Field( + description='Name of scalar datasets to paint w.r.t.', + ) + num_points: int = Field( + default=3378, # Number of Spike protein residues + description='Number of residues in the protein' + ' (i.e., points in the point cloud).', + ) + num_features: int = Field( + default=0, + description='Number of additional per-point features' + ' in addition to xyz coords.', + ) + latent_dim: int = Field( + default=3, + description='Dimensionality of the latent space.', + ) + encoder_bias: bool = Field( + default=True, + description='Whether to use bias in the encoder.', + ) + encoder_relu_slope: float = Field( + default=0.0, + description='The slope of the ReLU function in the encoder.', + ) + encoder_filters: list[int] = Field( + default=[64, 128, 256, 256, 512], + description='The number of filters in each convolutional layer' + ' of the encoder.', + ) + encoder_kernels: list[int] = Field( + default=[5, 3, 3, 1, 1], + description='The kernel size in each convolutional layer ' + 'of the encoder.', + ) + decoder_bias: bool = Field( + default=True, + description='Whether to use bias in the decoder.', + ) + decoder_relu_slope: float = Field( + default=0.0, + description='The slope of the ReLU function in the decoder.', + ) + decoder_affine_widths: list[int] = Field( + default=[64, 128, 512, 1024], + description='The width of the affine layers in the decoder.', + ) + discriminator_bias: bool = Field( + default=True, + description='Whether to use bias in the discriminator.', + ) + discriminator_relu_slope: float = Field( + default=0.0, + description='The slope of the ReLU function in the discriminator.', + ) + discriminator_affine_widths: list[int] = Field( + default=[512, 512, 128, 64], + description='The width of the affine layers in the discriminator.', + ) + noise_mu: float = Field( + default=0.0, + description='Mean of the prior distribution.', + ) + noise_std: float = Field( + default=0.2, + description='Standard deviation of the prior distribution.', + ) + lambda_gp: float = Field( + default=10.0, + description='Relative weight to put on gradient penalty.', + ) + lambda_rec: float = Field( + default=0.5, + description='Relative weight to put on reconstruction loss.', + ) + num_data_workers: int = Field( + default=0, + description='Number of data loaders for inference.', + ) + batch_size: int = Field( + default=32, + description='Inference batch size.', + ) + inference_batch_size: int = Field( + default=64, + description='Inference batch size.', + ) + + +class AdversarialAE: + """Adversarial autoencoder for protein conformers.""" + + def __init__( + self, + config: AdversarialAEConfig, + checkpoint_path: Path | None = None, + ) -> None: + """Initialize the ConvolutionalVAE. + + Parameters + ---------- + config : AdversarialAEConfig + The configuration settings for the model. + checkpoint_path : Path, optional + The path to the model checkpoint to load, by default None. + """ + # Lazy import to avoid needing torch to load module + from mdlearn.nn.models.aae.point_3d_aae import AAE3dTrainer + + self.config = config + self.checkpoint_path = checkpoint_path + + # We keep the inference_batch_size in the config for convenience + # but exclude it from the model arguments. + model_args = config.model_dump(exclude={'inference_batch_size'}) + self.trainer = AAE3dTrainer(**model_args) + + # Load the model checkpoint if specified + if checkpoint_path is not None: + self.update_model(checkpoint_path) + + def update_model(self, checkpoint_path: Path) -> None: + """Update the model with a new checkpoint. + + Parameters + ---------- + checkpoint_path : Path + The path to the checkpoint to load. + """ + # Skip if the checkpoint path is the same + if checkpoint_path == self.checkpoint_path: + return + + # Lazy import to avoid needing torch to load module + import torch + + # Load the checkpoint + cp = torch.load(checkpoint_path, map_location=self.trainer.device) + + # Load the model state dict + self.trainer.model.load_state_dict(cp['model_state_dict']) + + # Update the checkpoint path + self.checkpoint_path = checkpoint_path + + def fit( + self, + x: np.ndarray, + model_dir: Path, + scalars: dict[str, np.ndarray] | None = None, + ) -> Path: + """Fit the model to the input data. + + Parameters + ---------- + x : np.ndarray + The contact maps to fit the model to. (n_samples, *) where * is a + ragged dimension containing the concatenated row and column indices + of the ones in the contact map. + model_dir : Path + The directory to save the model to. + scalars : dict[str, np.ndarray], optional + The scalars to plot during training, by default None. + + Returns + ------- + Path + The path to the most recent model checkpoint. + """ + # Setup the scalars for plotting if specified + scalars = {} if scalars is None else scalars + + # Fit the model + self.trainer.fit(X=x, scalars=scalars, output_path=model_dir) + + # Log the loss curve to a CSV file + pd.DataFrame(self.trainer.loss_curve_).to_csv(model_dir / 'loss.csv') + + # Get the most recent model checkpoint from the checkpoint directory + checkpoint_dir = model_dir / 'checkpoints' + checkpoint_path = natsorted(list(checkpoint_dir.glob('*.pt')))[-1] + + return checkpoint_path + + def predict(self, x: np.ndarray) -> np.ndarray: + """ + Predicts the latent space coordinates for a given set of coordinates. + + Parameters + ---------- + x: np.ndarray + The contact maps to predict the latent space coordinates for + (n_samples, *) where * is a ragged dimension containing the + concatenated row and column indices of the ones in the contact map. + + Returns + ------- + np.ndarray + The predicted latent space coordinates (n_samples, latent_dim). + """ + # Predict the latent space coordinates + z, _ = self.trainer.predict(x, self.config.inference_batch_size) + return z + + +@lru_cache(maxsize=1) +def warmstart_aae( + config_path: Path, + checkpoint_path: Path, +) -> tuple[AdversarialAE, LatentSpaceHistory]: + """Load the model once and then return a cached version. + + Parameters + ---------- + config_path : Path + The path to the model configuration file. + checkpoint_path : Path + The path to the model checkpoint file. + + Returns + ------- + AdversarialAE + The AdversarialAE model. + LatentSpaceHistory + The latent space history. + """ + # Print the warmstart message + print(f'Cold start model from checkpoint {checkpoint_path}') + + # Load the model configuration + model_config = AdversarialAEConfig.from_yaml(config_path) + + # Load the model + model = AdversarialAE( + model_config, + checkpoint_path=checkpoint_path, + ) + + # Initialize the latent space history + history = LatentSpaceHistory() + + return model, history diff --git a/deepdrivewe/ai/cvae.py b/deepdrivewe/ai/cvae.py index ba04004..241861c 100644 --- a/deepdrivewe/ai/cvae.py +++ b/deepdrivewe/ai/cvae.py @@ -2,6 +2,7 @@ from __future__ import annotations +from functools import lru_cache from pathlib import Path import numpy as np @@ -10,6 +11,7 @@ from pydantic import Field from deepdrivewe import BaseModel +from deepdrivewe.ai.utils import LatentSpaceHistory class ConvolutionalVAEConfig(BaseModel): @@ -123,6 +125,7 @@ def __init__( ) self.config = config + self.checkpoint_path = checkpoint_path # We keep the inference_batch_size in the config for convenience # but exclude it from the model arguments. @@ -141,6 +144,10 @@ def update_model(self, checkpoint_path: Path) -> None: checkpoint_path : Path The path to the checkpoint to load. """ + # Skip if the checkpoint path is the same + if checkpoint_path == self.checkpoint_path: + return + # Lazy import to avoid needing torch to load module import torch @@ -211,3 +218,42 @@ def predict(self, x: np.ndarray) -> np.ndarray: # Predict the latent space coordinates z, *_ = self.trainer.predict(x, self.config.inference_batch_size) return z + + +@lru_cache(maxsize=1) +def warmstart_cvae( + config_path: Path, + checkpoint_path: Path, +) -> tuple[ConvolutionalVAE, LatentSpaceHistory]: + """Load the model once and then return a cached version. + + Parameters + ---------- + config_path : Path + The path to the model configuration file. + checkpoint_path : Path + The path to the model checkpoint file. + + Returns + ------- + ConvolutionalVAE + The ConvolutionalVAE model. + LatentSpaceHistory + The latent space history. + """ + # Print the warmstart message + print(f'Cold start model from checkpoint {checkpoint_path}') + + # Load the model configuration + model_config = ConvolutionalVAEConfig.from_yaml(config_path) + + # Load the model + model = ConvolutionalVAE( + model_config, + checkpoint_path=checkpoint_path, + ) + + # Initialize the latent space history + history = LatentSpaceHistory() + + return model, history diff --git a/deepdrivewe/examples/openmm_aae_ddwe/inference.py b/deepdrivewe/examples/openmm_aae_ddwe/inference.py new file mode 100644 index 0000000..289375d --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/inference.py @@ -0,0 +1,155 @@ +"""Inference module for the LOF strategy.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from pydantic import BaseModel +from pydantic import Field +from sklearn.neighbors import LocalOutlierFactor + +from deepdrivewe import BasisStates +from deepdrivewe import IterationMetadata +from deepdrivewe import SimMetadata +from deepdrivewe import SimResult +from deepdrivewe import TargetState +from deepdrivewe import TrainResult +from deepdrivewe.ai import warmstart_aae +from deepdrivewe.binners import RectilinearBinner +from deepdrivewe.recyclers import LowRecycler +from deepdrivewe.resamplers import LOFLowResampler + + +class InferenceConfig(BaseModel): + """Arguments for the inference module.""" + + # Local outlier factor settings + lof_n_neighbors: int = Field( + default=20, + description='The number of neighbors to use for LOF.', + ) + lof_distance_metric: str = Field( + default='cosine', + description='The distance metric to use for LOF [cosine, minkowski].', + ) + + # Resampling settings + sims_per_bin: int = Field( + default=72, + description='The number of simulations to maintain in each bin.' + ' Default is 72.', + ) + consider_for_resampling: int = Field( + default=12, + description='The number of simulations to consider for resampling.', + ) + max_resamples: int = Field( + default=4, + description='The maximum number of resamples to perform in each ' + 'iteration. Default is 4.', + ) + max_allowed_weight: float = Field( + default=1.0, + description='The maximum allowed weight for a simulation. Default ' + 'is 1.0.', + ) + min_allowed_weight: float = Field( + default=10e-40, + description='The minimum allowed weight for a simulation. Default ' + 'is 10e-40.', + ) + + +def run_inference( + sim_output: list[SimResult], + train_output: TrainResult, + basis_states: BasisStates, + target_states: list[TargetState], + config: InferenceConfig, + output_dir: Path, +) -> tuple[list[SimMetadata], list[SimMetadata], IterationMetadata]: + """Run inference on the input data.""" + # Make the output directory + itetation = sim_output[0].metadata.iteration_id + output_dir = output_dir / f'{itetation:06d}' + output_dir.mkdir(parents=True, exist_ok=True) + + # Extract the rmsd pcoord from the last frame of each simulation + pcoords = [sim.metadata.pcoord[-1][0] for sim in sim_output] + + print(f'Progress coordinates: {pcoords}') + print(f'Best progress coordinate: {min(pcoords)}') + print(f'Num input simulations: {len(sim_output)}') + + # Extract the simulation metadata + cur_sims = [sim.metadata for sim in sim_output] + + # Load the model and history + model, history = warmstart_aae( + train_output.config_path, + train_output.checkpoint_path, + ) + + # Extract the last frame coordinates and rmsd from each simulation + coordinates = [sim.data['coordinates'][-1] for sim in sim_output] + pcoords = [sim.data['pcoords'][-1] for sim in sim_output] + + # Compute the latent space representation + z = model.predict(x=coordinates) + + # Concatenate the latent history + if history: + z = np.concatenate([history.z, z]) + pcoords = np.concatenate([history.pcoords, pcoords]) + + # Run LOF on the latent space + clf = LocalOutlierFactor( + n_neighbors=config.lof_n_neighbors, + metric=config.lof_distance_metric, + ).fit(z) + + # Get the LOF scores + lof_scores = clf.negative_outlier_factor_ + + # Update the latent space history + history.update(z, pcoords) + + # Plot the latent space + history.plot(output_dir / 'pcoord.png') + history.plot( + output_dir / 'pcoord_lof.png', + color=lof_scores, + cblabel='LOF Score', + ) + + # Add the LOF scores to the last frame of each simulation pcoord + for sim, score in zip(cur_sims, lof_scores[-len(cur_sims) :]): + sim_scores = [-1.0 for _ in range(sim.num_frames)] + sim_scores[-1] = float(score) + sim.append_pcoord(sim_scores) + + # Create the binner + binner = RectilinearBinner( + bins=[0.0, 1.0, float('inf')], + bin_target_counts=config.sims_per_bin, + ) + + # Define the recycling policy + recycler = LowRecycler( + basis_states=basis_states, + target_threshold=target_states[0].pcoord[0], + ) + + # Define the resampling policy + resampler = LOFLowResampler( + consider_for_resampling=config.consider_for_resampling, + max_resamples=config.max_resamples, + max_allowed_weight=config.max_allowed_weight, + min_allowed_weight=config.min_allowed_weight, + ) + + # Assign simulations to bins and resample the weighted ensemble + result = resampler.run(cur_sims, binner, recycler) + + return result diff --git a/deepdrivewe/examples/openmm_aae_ddwe/main.py b/deepdrivewe/examples/openmm_aae_ddwe/main.py new file mode 100644 index 0000000..0dfd678 --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/main.py @@ -0,0 +1,249 @@ +"""DDWE example. + +Adapted from: +https://github.com/westpa/westpa2_tutorials/tree/main/tutorial7.7-hamsm +""" + +from __future__ import annotations + +import logging +import sys +from argparse import ArgumentParser +from functools import partial +from functools import update_wrapper +from pathlib import Path + +import MDAnalysis +from colmena.queue.python import PipeQueues +from colmena.task_server import ParslTaskServer +from MDAnalysis.analysis import rms +from proxystore.connectors.file import FileConnector +from proxystore.store import Store +from pydantic import Field +from pydantic import field_validator + +from deepdrivewe import BaseModel +from deepdrivewe import BasisStates +from deepdrivewe import EnsembleCheckpointer +from deepdrivewe import TargetState +from deepdrivewe import WeightedEnsemble +from deepdrivewe.examples.openmm_ntl9_ddwe.inference import InferenceConfig +from deepdrivewe.examples.openmm_ntl9_ddwe.inference import run_inference +from deepdrivewe.examples.openmm_ntl9_ddwe.simulate import run_simulation +from deepdrivewe.examples.openmm_ntl9_ddwe.simulate import SimulationConfig +from deepdrivewe.examples.openmm_ntl9_ddwe.train import run_train +from deepdrivewe.examples.openmm_ntl9_ddwe.train import TrainConfig +from deepdrivewe.parsl import ComputeConfigTypes +from deepdrivewe.workflows.ddwe import DDWEThinker + + +class RMSDBasisStateInitializer(BaseModel): + """RMSD basis state initialization.""" + + reference_file: Path = Field( + description='Reference file for the cpptraj command.', + ) + mda_selection: str = Field( + default='protein and name CA', + description='The MDAnalysis selection string for the atoms to use.', + ) + + def __call__(self, basis_file: str) -> list[float]: + """Initialize the basis state parent coordinates.""" + # Load the basis file + basis = MDAnalysis.Universe(basis_file) + + # Load the reference file + reference = MDAnalysis.Universe(self.reference_file) + + # Get the positions + pos = basis.select_atoms(self.mda_selection).positions + + # Get the reference positions + ref_pos = reference.select_atoms(self.mda_selection).positions + + # Compute the RMSD between the basis and reference structures + rmsd: float = rms.rmsd(pos, ref_pos, superposition=True) + + # Return the RMSD and the zeroed progress coordinate place holder + # for the LOF dimension + return [rmsd, 0.0] + + +class ExperimentSettings(BaseModel): + """Provide a YAML interface to configure the experiment.""" + + output_dir: Path = Field( + description='Directory in which to store the results.', + ) + num_iterations: int = Field( + ge=1, + description='Number of iterations to run the weighted ensemble.', + ) + basis_states: BasisStates = Field( + description='The basis states for the weighted ensemble.', + ) + basis_state_initializer: RMSDBasisStateInitializer = Field( + description='Arguments for initializing the basis states.', + ) + target_states: list[TargetState] = Field( + description='The target threshold for the progress coordinate to be' + ' considered in the target state.', + ) + simulation_config: SimulationConfig = Field( + description='Arguments for the simulation.', + ) + train_config: TrainConfig = Field( + description='Arguments for the training.', + ) + inference_config: InferenceConfig = Field( + description='Arguments for the inference.', + ) + compute_config: ComputeConfigTypes = Field( + description='Config for the compute resources.', + ) + use_stale_model: bool = Field( + default=False, + description='Whether to use the stale model for inference. This will ' + 'be faster but may not be as accurate. It uses the model from the ' + 'previous iteration for inference in the current iteration, which may ' + 'not be updated with new states.', + ) + max_retries: int = Field( + default=2, + description='Number of times to retry a task if it fails.', + ) + + @field_validator('output_dir') + @classmethod + def mkdir_validator(cls, value: Path) -> Path: + """Resolve and make the output directory.""" + value = value.resolve() + value.mkdir(parents=True, exist_ok=True) + return value + + +if __name__ == '__main__': + parser = ArgumentParser() + parser.add_argument('-c', '--config', required=True) + args = parser.parse_args() + cfg = ExperimentSettings.from_yaml(args.config) + cfg.dump_yaml(cfg.output_dir / 'params.yaml') + + # Set up logging + logging.basicConfig( + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + level=logging.INFO, + handlers=[ + logging.FileHandler(cfg.output_dir / 'runtime.log'), + logging.StreamHandler(sys.stdout), + ], + ) + + # Make the store + store = Store( + name='file-store', + register=True, + connector=FileConnector(store_dir=str(cfg.output_dir / 'proxy-store')), + ) + + # Make the queues + queues = PipeQueues( + serialization_method='pickle', + topics=['simulation', 'train', 'inference'], + proxystore_name='file-store', + proxystore_threshold=10000, + ) + + # Define the parsl configuration (this can be done using the + # get_parsl_config for common use cases or by defining your own config.) + parsl_config = cfg.compute_config.get_parsl_config( + cfg.output_dir / 'run-info', + ) + + # Create the checkpoint manager + checkpointer = EnsembleCheckpointer(output_dir=cfg.output_dir) + + # Check if a checkpoint exists + checkpoint = checkpointer.latest_checkpoint() + + if checkpoint is None: + # Initialize the weighted ensemble + ensemble = WeightedEnsemble( + basis_states=cfg.basis_states, + target_states=cfg.target_states, + ) + + # Initialize the simulations with the basis states + ensemble.initialize_basis_states(cfg.basis_state_initializer) + else: + # Load the ensemble from a checkpoint if it exists + ensemble = checkpointer.load(checkpoint) + logging.info(f'Loaded ensemble from checkpoint {checkpoint}') + + # Print the input states + logging.info(f'Basis states: {ensemble.basis_states}') + logging.info(f'Target states: {ensemble.target_states}') + + # Assign constant settings to each task function + my_run_simulation = partial( + run_simulation, + config=cfg.simulation_config, + output_dir=cfg.output_dir / 'simulation', + ) + my_run_train = partial( + run_train, + config=cfg.train_config, + output_dir=cfg.output_dir / 'train', + ) + my_run_inference = partial( + run_inference, + basis_states=ensemble.basis_states, + target_states=ensemble.target_states, + config=cfg.inference_config, + output_dir=cfg.output_dir / 'inference', + ) + update_wrapper(my_run_simulation, run_simulation) + update_wrapper(my_run_train, run_train) + update_wrapper(my_run_inference, run_inference) + + # Create the task server + doer = ParslTaskServer( + [ + (my_run_simulation, {'executors': ['simulation_htex']}), + (my_run_train, {'executors': ['train_htex']}), + (my_run_inference, {'executors': ['inference_htex']}), + ], + queues, + parsl_config, + ) + + # Create the workflow thinker + thinker = DDWEThinker( + queue=queues, + result_dir=cfg.output_dir / 'result', + ensemble=ensemble, + checkpointer=checkpointer, + num_iterations=cfg.num_iterations, + use_stale_model=cfg.use_stale_model, + max_retries=cfg.max_retries, + ) + logging.info('Created the task server and task generator') + + try: + # Launch the servers + doer.start() + thinker.start() + logging.info('Launched the servers') + + # Wait for the task generator to complete + thinker.join() + logging.info('Task generator has completed') + finally: + queues.send_kill_signal() + + # Wait for the task server to complete + doer.join() + + # Clean up proxy store + store.close() diff --git a/deepdrivewe/examples/openmm_aae_ddwe/simulate.py b/deepdrivewe/examples/openmm_aae_ddwe/simulate.py new file mode 100644 index 0000000..9451d66 --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/simulate.py @@ -0,0 +1,110 @@ +"""Simulate a system using OpenMM.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Sequence + +from pydantic import Field + +from deepdrivewe import BaseModel +from deepdrivewe import SimMetadata +from deepdrivewe import SimResult +from deepdrivewe.simulation.openmm import CollectionReporter +from deepdrivewe.simulation.openmm import CoordinatesCollector +from deepdrivewe.simulation.openmm import OpenMMConfig +from deepdrivewe.simulation.openmm import OpenMMSimulation +from deepdrivewe.simulation.openmm import RMSDCollector + + +class SimulationConfig(BaseModel): + """Arguments for the naive resampler.""" + + openmm_config: OpenMMConfig = Field( + description='The configuration for the Amber simulation.', + ) + top_file: Path | None = Field( + default=None, + description='The topology file for the simulation.', + ) + reference_file: Path = Field( + description='The reference PDB file for the analysis.', + ) + cutoff_angstrom: float = Field( + default=8.0, + description='The angstrom cutoff distance for defining contacts.', + ) + mda_selection: str = Field( + default='protein and name CA', + description='The MDAnalysis selection string for the atoms to use.', + ) + openmm_selection: Sequence[str] = Field( + default=('CA',), + description='The OpenMM selection strings for the atoms to use.', + ) + + +def run_simulation( + metadata: SimMetadata, + config: SimulationConfig, + output_dir: Path, +) -> SimResult: + """Run a simulation and return the pcoord and coordinates.""" + # Add performance logging + metadata.mark_simulation_start() + + # Create the simulation output directory + sim_output_dir = output_dir / metadata.simulation_name + + # Remove the directory if it already exists + # (this would be from a task failure) + if sim_output_dir.exists(): + # Remove each file in the directory + for file in sim_output_dir.iterdir(): + file.unlink() + + # Create a fresh output directory + sim_output_dir.mkdir(parents=True, exist_ok=True) + + # Log the yaml config file to this directory + config.dump_yaml(sim_output_dir / 'config.yaml') + + # Initialize the simulation + simulation = OpenMMSimulation( + config=config.openmm_config, + top_file=config.top_file, + output_dir=sim_output_dir, + checkpoint_file=metadata.parent_restart_file, + ) + + # Setup the data collectors and reporter + reporter = CollectionReporter( + report_interval=config.openmm_config.report_steps, + openmm_selection=config.openmm_selection, + collectors=[ + CoordinatesCollector( + reference_file=config.reference_file, + mda_selection=config.mda_selection, + ), + RMSDCollector( + reference_file=config.reference_file, + mda_selection=config.mda_selection, + topic='pcoords', + ), + ], + ) + + # Run the simulation + simulation.run(reporters=[reporter]) + + # Get the collected data + data = reporter.get_collected_data() + + # Update the simulation metadata + metadata.restart_file = simulation.restart_file + metadata.pcoord = data['pcoords'].tolist() + metadata.mark_simulation_end() + + result = SimResult(data=data, metadata=metadata) + + return result diff --git a/deepdrivewe/examples/openmm_aae_ddwe/train.py b/deepdrivewe/examples/openmm_aae_ddwe/train.py new file mode 100644 index 0000000..2157978 --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/train.py @@ -0,0 +1,75 @@ +"""Training module.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from pydantic import BaseModel +from pydantic import Field + +from deepdrivewe import SimResult +from deepdrivewe import TrainResult +from deepdrivewe.ai import AdversarialAE +from deepdrivewe.ai import AdversarialAEConfig + + +class TrainConfig(BaseModel): + """Arguments for the training module.""" + + config_path: Path = Field( + description='The path to the model configuration file.', + ) + checkpoint_path: Path | None = Field( + default=None, + description='The path to the model checkpoint file.' + 'Train from scratch by default.', + ) + + +# TODO: We probably need to store a history of old training data +# to retrain the model. Add a config argument to include a cMD run dataset. +# Contact maps: https://github.com/n-frazee/DL-enhancedWE/blob/main/common_files/train.npy + + +def run_train( + sim_output: list[SimResult], + config: TrainConfig, + output_dir: Path, +) -> TrainResult: + """Train the model on the simulation output.""" + # Make the output directory + itetation = sim_output[0].metadata.iteration_id + output_dir = output_dir / f'{itetation:06d}' + output_dir.mkdir(parents=True, exist_ok=True) + + # Load the model configuration + model_config = AdversarialAEConfig.from_yaml(config.config_path) + + # Load the model + model = AdversarialAE( + model_config, + checkpoint_path=config.checkpoint_path, + ) + + # Extract the last frame contact maps and rmsd from each simulation + coordinates = np.concatenate( + [sim.data['coordinates'] for sim in sim_output], + ) + pcoords = np.concatenate([sim.data['pcoords'] for sim in sim_output]) + pcoords = pcoords.flatten() + + # Fit the model + checkpoint_path = model.fit( + x=coordinates, + model_dir=output_dir / 'model', + scalars={'pcoord': pcoords}, + ) + + # Return the train result + result = TrainResult( + config_path=config.config_path, + checkpoint_path=checkpoint_path, + ) + + return result diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py b/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py index bb422ab..9a28a95 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py @@ -15,7 +15,7 @@ from deepdrivewe import SimResult from deepdrivewe import TargetState from deepdrivewe import TrainResult -from deepdrivewe.ai import warmstart_model +from deepdrivewe.ai import warmstart_cvae from deepdrivewe.binners import RectilinearBinner from deepdrivewe.recyclers import LowRecycler from deepdrivewe.resamplers import LOFLowResampler @@ -86,7 +86,7 @@ def run_inference( cur_sims = [sim.metadata for sim in sim_output] # Load the model and history - model, history = warmstart_model( + model, history = warmstart_cvae( train_output.config_path, train_output.checkpoint_path, ) diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py index cb25d5c..8a2b0f5 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py @@ -10,9 +10,11 @@ from deepdrivewe import BaseModel from deepdrivewe import SimMetadata from deepdrivewe import SimResult -from deepdrivewe.simulation.openmm import ContactMapRMSDReporter +from deepdrivewe.simulation.openmm import CollectionReporter +from deepdrivewe.simulation.openmm import ContactMapCollector from deepdrivewe.simulation.openmm import OpenMMConfig from deepdrivewe.simulation.openmm import OpenMMSimulation +from deepdrivewe.simulation.openmm import RMSDCollector class SimulationConfig(BaseModel): @@ -75,30 +77,33 @@ def run_simulation( checkpoint_file=metadata.parent_restart_file, ) - # Add the contact map and RMSD reporter - reporter = ContactMapRMSDReporter( + # Setup the data collectors and reporter + reporter = CollectionReporter( report_interval=config.openmm_config.report_steps, - reference_file=config.reference_file, - cutoff_angstrom=config.cutoff_angstrom, - mda_selection=config.mda_selection, openmm_selection=config.openmm_selection, + collectors=[ + ContactMapCollector( + cutoff_angstrom=config.cutoff_angstrom, + ), + RMSDCollector( + reference_file=config.reference_file, + mda_selection=config.mda_selection, + topic='pcoords', + ), + ], ) # Run the simulation simulation.run(reporters=[reporter]) - # Run the contact map and RMSD analysis - contact_maps = reporter.get_contact_maps() - pcoord = reporter.get_rmsds() + # Get the collected data + data = reporter.get_collected_data() # Update the simulation metadata metadata.restart_file = simulation.restart_file - metadata.pcoord = pcoord.tolist() + metadata.pcoord = data['pcoords'].tolist() metadata.mark_simulation_end() - result = SimResult( - data={'contact_maps': contact_maps, 'pcoords': pcoord}, - metadata=metadata, - ) + result = SimResult(data=data, metadata=metadata) return result diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index f13290c..67d7749 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -17,6 +17,7 @@ import MDAnalysis import numpy as np +from MDAnalysis.analysis import align from MDAnalysis.analysis import distances from MDAnalysis.analysis import rms from pydantic import BaseModel @@ -34,17 +35,25 @@ class OpenMMReporter(ABC): - """Reporter protocol for OpenMM simulations.""" + """Reporter interface for OpenMM simulations.""" - def __init__(self, report_interval: int) -> None: + def __init__( + self, + report_interval: int, + openmm_selection: Sequence[str] = ('CA',), + ) -> None: """Initialize the reporter. Parameters ---------- report_interval : int The interval at which to write frames. + openmm_selection : Sequence[str] + The OpenMM selection strings for the atoms to use + when reporting positions (default is ('CA',)). """ self.report_interval = report_interval + self.openmm_selection = openmm_selection def describeNextReport( # noqa: N802 self, @@ -70,6 +79,41 @@ def describeNextReport( # noqa: N802 steps = self.report_interval - step_progress return (steps, True, False, False, False, None) + def get_positions( + self, + simulation: app.Simulation, + state: openmm.State, + ) -> np.ndarray: + """Get the atomic positions from the simulation. + + Parameters + ---------- + simulation : Simulation + The Simulation to generate a report for. + state : State + The current state of the simulation. + + Returns + ------- + np.ndarray + The atomic positions from the simulation. + """ + # Get the atom indices for the selection + atom_indices = [ + a.index + for a in simulation.topology.atoms() + if a.name in self.openmm_selection + ] + + # Get the atomic coordinates of the selection + positions = state.getPositions(asNumpy=True) + positions = positions[atom_indices].astype(np.float32) + + # Convert positions from nanometers to angstroms + positions *= 10.0 + + return positions + @abstractmethod def report(self, simulation: app.Simulation, state: openmm.State) -> None: """Generate a report. @@ -84,6 +128,288 @@ def report(self, simulation: app.Simulation, state: openmm.State) -> None: pass +class Collector(ABC): + """Collector interface for OpenMM simulations.""" + + def __init__(self, topic: str) -> None: + """Initialize the collector.""" + self._topic = topic + + @property + def topic(self) -> str: + """The topic of the collector (a string identifier).""" + return self._topic + + @abstractmethod + def get(self) -> np.ndarray: + """Get the collected data from the simulation. + + Returns + ------- + np.ndarray + The collected data from the simulation. + """ + ... + + @abstractmethod + def collect(self, positions: np.ndarray) -> None: + """Collect data from the simulation. + + Parameters + ---------- + positions : np.ndarray + The atomic positions from the simulation. + """ + ... + + +class CoordinatesCollector(Collector): + """Coordinates collector for OpenMM simulations.""" + + def __init__( + self, + reference_file: Path | None = None, + mda_selection: str = 'protein and name CA', + topic: str = 'coordinates', + ) -> None: + """Initialize the coordinates collector. + + Parameters + ---------- + reference_file : Path | None + The reference PDB file for the analysis (if provided, + the coordinates will be aligned to the reference). + Default is None. + mda_selection : str + The MDAnalysis selection string for the atoms to use + for alignment (default is 'protein and name CA'). + topic : str + The topic of the collector, default is 'coordinates'. + """ + super().__init__(topic) + self._coordinates: list[np.ndarray] = [] + self._ref = None + + # If provided, load the reference structure and save the positions + if reference_file is not None: + mda_u = MDAnalysis.Universe(reference_file) + self._ref = mda_u.select_atoms(mda_selection).positions.copy() + + def get(self) -> np.ndarray: + """Get the coordinates from the simulation. + + Returns + ------- + np.ndarray + The atomic positions from each frame of the simulation + (n_frames, n_atoms, 3). Where n_atoms is the number of atoms + in the openmm_selection of the corresponding OpenMMReporter. + """ + return np.array(self._coordinates) + + def _align(self, positions: np.ndarray) -> np.ndarray: + """Align the atomic positions to the reference. + + Parameters + ---------- + positions : np.ndarray + The atomic positions from the simulation. + + Returns + ------- + np.ndarray + The aligned atomic positions. + """ + # Calculate rotation and translation using align.rotation_matrix() + rotation_matrix, _ = align.rotation_matrix(positions, self._ref) + + # Apply the rotation to the raw positions + aligned_positions = np.dot(positions, rotation_matrix.T) + + return aligned_positions + + def collect(self, positions: np.ndarray) -> None: + """Generate a report. + + Parameters + ---------- + positions : np.ndarray + The atomic positions from the simulation. + """ + # Align the coordinates to the reference if provided + pos = positions.copy() if self._ref is None else self._align(positions) + + # Collect the position coordinates + self._coordinates.append(pos) + + +class RMSDCollector(Collector): + """RMSD collector for OpenMM simulations.""" + + def __init__( + self, + reference_file: Path, + mda_selection: str = 'protein and name CA', + topic: str = 'rmsds', + ) -> None: + """Initialize the RMSD collector. + + Parameters + ---------- + reference_file : Path + The reference PDB file for the analysis. + mda_selection : str + The MDAnalysis selection string for the atoms to use + (default is 'protein and name CA'). + topic : str + The topic of the collector, default is 'rmsd'. + """ + super().__init__(topic) + self._rmsd: list[float] = [] + + # Load the reference structure and save the positions + mda_u = MDAnalysis.Universe(reference_file) + self._ref = mda_u.select_atoms(mda_selection).positions.copy() + + def get(self) -> np.ndarray: + """Get the RMSDs from the simulation. + + Returns + ------- + np.ndarray + The RMSDs from the simulation shaped as (n_frames, 1). + """ + return np.array(self._rmsd).reshape(-1, 1) + + def collect(self, positions: np.ndarray) -> None: + """Generate a report. + + Parameters + ---------- + positions : np.ndarray + The atomic positions from the simulation. + """ + # Compute the RMSD + rmsd = rms.rmsd(positions, self._ref, superposition=True) + self._rmsd.append(rmsd) + + +class ContactMapCollector(Collector): + """Contact map collector for OpenMM simulations.""" + + def __init__( + self, + cutoff_angstrom: float = 8.0, + topic: str = 'contact_maps', + ) -> None: + """Initialize the contact map collector.""" + super().__init__(topic) + self.cutoff_angstrom = cutoff_angstrom + self._rows: list[np.ndarray] = [] + self._cols: list[np.ndarray] = [] + + def get(self) -> np.ndarray: + """Get the contact maps from the simulation. + + Returns + ------- + np.ndarray + The contact maps from the simulation as a ragged array + shaped as (n_frames, *). + """ + # Concatenate the row and col indices into a single array + contact_maps = [np.concatenate(x) for x in zip(self._rows, self._cols)] + + # Collect the contact maps in a ragged numpy array + contact_maps = np.array(contact_maps, dtype=object) + + return contact_maps + + def collect(self, positions: np.ndarray) -> None: + """Generate a report. + + Parameters + ---------- + positions : np.ndarray + The atomic positions from the simulation. + """ + # Compute the contact map + contact_map = distances.contact_matrix( + positions, + self.cutoff_angstrom, + returntype='sparse', + ) + + # Convert the contact map to sparse format + coo_matrix = contact_map.tocoo() + + # Append the row and col indices to lists + self._rows.append(coo_matrix.row.astype('int16')) + self._cols.append(coo_matrix.col.astype('int16')) + + +class CollectionReporter(OpenMMReporter): + """Reporter to collect multiple data products from an OpenMM simulation.""" + + def __init__( + self, + report_interval: int, + collectors: list[Collector], + openmm_selection: Sequence[str] = ('CA',), + ) -> None: + """Initialize the reporter. + + Parameters + ---------- + report_interval : int + The interval at which to write frames. + collectors : list[Collector] + The collectors to inject into the simulation. + openmm_selection : Sequence[str] + The OpenMM selection strings for the atoms to use + when reporting positions (default is ('CA',)). + + Raises + ------ + ValueError + If the collectors have duplicate topics. + """ + super().__init__(report_interval, openmm_selection) + + # Check that the collectors have unique topics + if len(collectors) != len({x.topic for x in collectors}): + raise ValueError('Collectors must have unique topics.') + + self.collectors = collectors + + def get_collected_data(self) -> dict[str, np.ndarray]: + """Get the collected data from the simulation. + + Returns + ------- + dict[str, np.ndarray] + The collected data from the simulation. + """ + return {x.topic: x.get() for x in self.collectors} + + def report(self, simulation: app.Simulation, state: openmm.State) -> None: + """Generate a report. + + Parameters + ---------- + simulation : Simulation + The Simulation to generate a report for. + state : State + The current state of the simulation. + """ + # Get the positions + positions = self.get_positions(simulation, state) + + # Collect data from the simulation + for collector in self.collectors: + collector.collect(positions) + + class OpenMMConfig(BaseModel): """Configuration for an OpenMM simulation.""" @@ -532,6 +858,7 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: sim.saveCheckpoint(str(self.output_dir / 'seg.chk')) +# TODO: First test the above implementation, then remove this class. class ContactMapRMSDReporter(OpenMMReporter): """Reporter to compute contact maps and RMSD from an OpenMM simulation.""" @@ -610,19 +937,8 @@ def report(self, simulation: app.Simulation, state: openmm.State) -> None: state : State The current state of the simulation """ - # Get the atom indices for the selection - atom_indices = [ - a.index - for a in simulation.topology.atoms() - if a.name in self.openmm_selection - ] - - # Get the atomic coordinates of the selection - positions = state.getPositions(asNumpy=True) - positions = positions[atom_indices].astype(np.float32) - - # Convert positions from nanometers to angstroms - positions *= 10.0 + # Get the positions + positions = self.get_positions(simulation, state) # Compute the contact map contact_map = distances.contact_matrix( diff --git a/pyproject.toml b/pyproject.toml index 9eb66db..4a9452c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "mdtraj==1.10.0", "MDAnalysis>=2.7.0", "scikit-learn==1.5.1", - "mdlearn==1.0.0", + "mdlearn==1.0.3", "scipy==1.14.0", "natsort>=8.4.0", "matplotlib>=3.9.2", From 9a2ca192918fe1da502835c214a2f707a13cf67d Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 29 Oct 2024 14:36:25 -0500 Subject: [PATCH 02/73] imports --- deepdrivewe/examples/openmm_aae_ddwe/main.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/deepdrivewe/examples/openmm_aae_ddwe/main.py b/deepdrivewe/examples/openmm_aae_ddwe/main.py index 0dfd678..ca3afda 100644 --- a/deepdrivewe/examples/openmm_aae_ddwe/main.py +++ b/deepdrivewe/examples/openmm_aae_ddwe/main.py @@ -27,12 +27,12 @@ from deepdrivewe import EnsembleCheckpointer from deepdrivewe import TargetState from deepdrivewe import WeightedEnsemble -from deepdrivewe.examples.openmm_ntl9_ddwe.inference import InferenceConfig -from deepdrivewe.examples.openmm_ntl9_ddwe.inference import run_inference -from deepdrivewe.examples.openmm_ntl9_ddwe.simulate import run_simulation -from deepdrivewe.examples.openmm_ntl9_ddwe.simulate import SimulationConfig -from deepdrivewe.examples.openmm_ntl9_ddwe.train import run_train -from deepdrivewe.examples.openmm_ntl9_ddwe.train import TrainConfig +from deepdrivewe.examples.openmm_aae_ddwe.inference import InferenceConfig +from deepdrivewe.examples.openmm_aae_ddwe.inference import run_inference +from deepdrivewe.examples.openmm_aae_ddwe.simulate import run_simulation +from deepdrivewe.examples.openmm_aae_ddwe.simulate import SimulationConfig +from deepdrivewe.examples.openmm_aae_ddwe.train import run_train +from deepdrivewe.examples.openmm_aae_ddwe.train import TrainConfig from deepdrivewe.parsl import ComputeConfigTypes from deepdrivewe.workflows.ddwe import DDWEThinker From 0e2d4032ea14ad23086751ab3da7f0f349db8b74 Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 29 Oct 2024 21:40:23 -0500 Subject: [PATCH 03/73] remove old params --- examples/openmm_ntl9_ddwe/config.yaml | 5 ----- 1 file changed, 5 deletions(-) diff --git a/examples/openmm_ntl9_ddwe/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index 9a8ea6b..52dbb71 100644 --- a/examples/openmm_ntl9_ddwe/config.yaml +++ b/examples/openmm_ntl9_ddwe/config.yaml @@ -52,11 +52,6 @@ train_config: # The configuration for the inference inference_config: - # The path to the CVAE model configuration file - ai_model_config_path: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/cvae-config.yaml - # The path to the CVAE model weights file - ai_model_checkpoint_path: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/checkpoint-epoch-100.pt - # The number of neighbors to use for LOF lof_n_neighbors: 20 # The distance metric to use for LOF [cosine, minkowski] From 1123caa13235bf4f213943b420ba52bcf7142875 Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 29 Oct 2024 22:09:56 -0500 Subject: [PATCH 04/73] gpu config --- deepdrivewe/parsl.py | 63 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 62 insertions(+), 1 deletion(-) diff --git a/deepdrivewe/parsl.py b/deepdrivewe/parsl.py index 5ef234b..924d307 100644 --- a/deepdrivewe/parsl.py +++ b/deepdrivewe/parsl.py @@ -128,6 +128,66 @@ def get_parsl_config(self, run_dir: str | Path) -> Config: ) +class WorkstationV2Config(BaseComputeConfig): + """Compute config for a workstation.""" + + name: Literal['workstation_v2'] = 'workstation_v2' # type: ignore[assignment] + + available_accelerators: int | Sequence[str] = Field( + ge=3, + description='Number of GPU accelerators to use.', + ) + worker_port_range: tuple[int, int] = Field( + default=(10000, 20000), + description='Port range for the workers.', + ) + retries: int = Field( + default=1, + description='Number of retries for the task.', + ) + # We have a long idletime to ensure train/inference executors are not + # shut down (to enable warmstarts) while simulations are running. + max_idletime: float = Field( + default=60.0 * 10, + description='The maximum idle time allowed for an executor before ' + 'strategy could shut down unused blocks. Default is 10 minutes.', + ) + + def _get_htex( + self, + label: str, + available_accelerators: Sequence[str], + ) -> HighThroughputExecutor: + return HighThroughputExecutor( + address='localhost', + label=label, + cpu_affinity='block', + available_accelerators=available_accelerators, + worker_port_range=self.worker_port_range, + provider=LocalProvider(init_blocks=1, max_blocks=1), + ) + + def get_parsl_config(self, run_dir: str | Path) -> Config: + """Generate a Parsl configuration.""" + # Handle the case where available_accelerators is an int + accelerators = self.available_accelerators + if isinstance(accelerators, int): + accelerators = [str(i) for i in range(accelerators)] + + return Config( + run_dir=str(run_dir), + retries=self.retries, + max_idletime=self.max_idletime, + executors=[ + # Assign 1 GPU each for training and inference + self._get_htex('train_htex', accelerators[:1]), + self._get_htex('inference_htex', accelerators[1:2]), + # Assign the remaining GPUs to simulation + self._get_htex('simulation_htex', accelerators[2:]), + ], + ) + + class HybridWorkstationConfig(BaseComputeConfig): """Run simulations on CPU and AI models on GPU.""" @@ -273,7 +333,7 @@ def get_parsl_config(self, run_dir: str | Path) -> Config: # Assign 1 node each for training and inference self._get_htex('train_htex', 1), self._get_htex('inference_htex', 1), - # Assign the remaining nodes to the simulation + # Assign the remaining nodes to simulation self._get_htex('simulation_htex', self.num_nodes - 2), ], ) @@ -282,6 +342,7 @@ def get_parsl_config(self, run_dir: str | Path) -> Config: ComputeConfigTypes = Union[ LocalConfig, WorkstationConfig, + WorkstationV2Config, HybridWorkstationConfig, InferenceTrainWorkstationConfig, VistaConfig, From d591991660bfd79ef641819d7a496d137eeffb58 Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 29 Oct 2024 22:30:47 -0500 Subject: [PATCH 05/73] gpu config --- deepdrivewe/parsl.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/deepdrivewe/parsl.py b/deepdrivewe/parsl.py index 924d307..16500fb 100644 --- a/deepdrivewe/parsl.py +++ b/deepdrivewe/parsl.py @@ -134,7 +134,6 @@ class WorkstationV2Config(BaseComputeConfig): name: Literal['workstation_v2'] = 'workstation_v2' # type: ignore[assignment] available_accelerators: int | Sequence[str] = Field( - ge=3, description='Number of GPU accelerators to use.', ) worker_port_range: tuple[int, int] = Field( @@ -153,6 +152,17 @@ class WorkstationV2Config(BaseComputeConfig): 'strategy could shut down unused blocks. Default is 10 minutes.', ) + @model_validator(mode='after') + def validate_available_accelerators(self) -> Self: + """Check there are at least 3 GPUs.""" + min_gpus = 3 + gpus = self.available_accelerators + num_gpus = gpus if isinstance(gpus, int) else len(gpus) + if num_gpus < min_gpus: + raise ValueError('Must use at least 3 GPUs.') + + return self + def _get_htex( self, label: str, From 9b892285607e70d251e04706ff50aeef0a89a386 Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 29 Oct 2024 22:42:28 -0500 Subject: [PATCH 06/73] print --- deepdrivewe/ai/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/ai/utils.py b/deepdrivewe/ai/utils.py index a0af97d..be69c08 100644 --- a/deepdrivewe/ai/utils.py +++ b/deepdrivewe/ai/utils.py @@ -69,7 +69,7 @@ def plot( color = self.pcoords if color is None else color print( - f'Plotting latent space to with {len(self.z)} ' + f'Plotting latent space with {len(self.z)} points ' f'and color with shape {len(color)} frames to {output_path}', ) From d619d977892d81f3707acf4838722c1fb7e71858 Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 30 Oct 2024 10:11:53 -0500 Subject: [PATCH 07/73] set simulation random seed --- deepdrivewe/simulation/openmm.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 67d7749..ab80070 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -455,6 +455,10 @@ class OpenMMConfig(BaseModel): default=True, description='Whether to randomize the basis state initial velocities.', ) + seed: int = Field( + default=np.random.default_rng().integers(2**32, dtype=int), + description='The random seed for the simulation.', + ) hardware_platform: str = Field( default='CUDA', description='The hardware platform to use for the simulation.' @@ -700,6 +704,10 @@ def configure_simulation( ValueError If explicit solvent is selected and no topology file is provided. """ + # Set the random seed + random.seed(self.seed) + np.random.seed(self.seed) + # Select implicit or explicit solvent configuration and load the system if self.solvent_type == 'explicit': if top_file is None: From c3050c3e21d3784963c87e5b5799cbb6a013b20d Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 30 Oct 2024 13:11:13 -0500 Subject: [PATCH 08/73] Add streaming thinker logic. Add proxy stream config. Co-authored-by: Greg Pauloski <18683347+gpauloski@users.noreply.github.com> --- deepdrivewe/examples/openmm_ntl9_ddwe/main.py | 37 ++- deepdrivewe/workflows/ddwe.py | 237 +++++++++++++++--- deepdrivewe/workflows/stream.py | 80 ++++++ pyproject.toml | 2 +- 4 files changed, 318 insertions(+), 38 deletions(-) create mode 100644 deepdrivewe/workflows/stream.py diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py index 0dfd678..4d60eeb 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py @@ -34,7 +34,9 @@ from deepdrivewe.examples.openmm_ntl9_ddwe.train import run_train from deepdrivewe.examples.openmm_ntl9_ddwe.train import TrainConfig from deepdrivewe.parsl import ComputeConfigTypes +from deepdrivewe.workflows.ddwe import DDWEStreamThinker from deepdrivewe.workflows.ddwe import DDWEThinker +from deepdrivewe.workflows.stream import ProxyStreamConfig class RMSDBasisStateInitializer(BaseModel): @@ -102,6 +104,10 @@ class ExperimentSettings(BaseModel): compute_config: ComputeConfigTypes = Field( description='Config for the compute resources.', ) + stream_config: ProxyStreamConfig | None = Field( + None, + description='Stream configuration for simulation data.', + ) use_stale_model: bool = Field( default=False, description='Whether to use the stale model for inference. This will ' @@ -219,15 +225,28 @@ def mkdir_validator(cls, value: Path) -> Path: ) # Create the workflow thinker - thinker = DDWEThinker( - queue=queues, - result_dir=cfg.output_dir / 'result', - ensemble=ensemble, - checkpointer=checkpointer, - num_iterations=cfg.num_iterations, - use_stale_model=cfg.use_stale_model, - max_retries=cfg.max_retries, - ) + if cfg.stream_config is None: + thinker = DDWEThinker( + queue=queues, + result_dir=cfg.output_dir / 'result', + ensemble=ensemble, + checkpointer=checkpointer, + num_iterations=cfg.num_iterations, + use_stale_model=cfg.use_stale_model, + max_retries=cfg.max_retries, + ) + else: + thinker = DDWEStreamThinker( + queue=queues, + result_dir=cfg.output_dir / 'result', + ensemble=ensemble, + checkpointer=checkpointer, + num_iterations=cfg.num_iterations, + use_stale_model=cfg.use_stale_model, + max_retries=cfg.max_retries, + stream_config=cfg.stream_config, + ) + logging.info('Created the task server and task generator') try: diff --git a/deepdrivewe/workflows/ddwe.py b/deepdrivewe/workflows/ddwe.py index 77a9fa1..3989f79 100644 --- a/deepdrivewe/workflows/ddwe.py +++ b/deepdrivewe/workflows/ddwe.py @@ -2,6 +2,7 @@ from __future__ import annotations +import time from pathlib import Path from typing import Any @@ -14,13 +15,14 @@ from deepdrivewe import EnsembleCheckpointer from deepdrivewe import WeightedEnsemble +from deepdrivewe.workflows.stream import ProxyStreamConfig from deepdrivewe.workflows.utils import ResultLogger class DDWEThinker(BaseThinker): """A thinker for the DDWE workflow.""" - def __init__( # noqa: PLR0913 + def __init__( self, queue: ColmenaQueues, result_dir: Path, @@ -28,7 +30,6 @@ def __init__( # noqa: PLR0913 checkpointer: EnsembleCheckpointer, num_iterations: int, use_stale_model: bool = False, - streaming: bool = False, max_retries: int = 2, ) -> None: """Initialize the DDWE workflow thinker. @@ -50,9 +51,6 @@ def __init__( # noqa: PLR0913 This will be faster but may not be as accurate. It uses the model from the previous iteration for inference in the current iteration, which may not be updated with new states. - streaming: bool - Whether to stream the simulation results directly to the - training task (default to False). max_retries: int Number of times to retry a task if it fails (default to 2). """ @@ -62,7 +60,6 @@ def __init__( # noqa: PLR0913 self.checkpointer = checkpointer self.num_iterations = num_iterations self.use_stale_model = use_stale_model - self.streaming = streaming self.max_retries = max_retries self.result_logger = ResultLogger(result_dir) @@ -96,12 +93,6 @@ def start_workflow(self) -> None: for sim in self.ensemble.next_sims: self.submit_task('simulation', sim) - # If we are not streaming, then we need to submit a single train task - # at the start of the workflow - if self.streaming: - self.logger.info('Start streaming train task') - self.submit_task('train') - @result_processor(topic='simulation') def process_simulation_result(self, result: Result) -> None: """Process a simulation result.""" @@ -122,22 +113,13 @@ def process_simulation_result(self, result: Result) -> None: # Note: We need to extract the proxied objects before storing them # to avoid auto-eviction after single use. The return results # are re-proxied before submitting the train/inference tasks. - # If we are streaming, then the simulation results only need - # to be used to submit and inference task, so we don't need to - # extract the proxied objects. The non-streaming case will - # need to extract and re-proxy the objects twice (once for - # the train task and once for the inference task). - output = result.value if self.streaming else extract(result.value) - self.sim_output.append(output) + self.sim_output.append(extract(result.value)) # If we have all the simulation results, submit a train task if len(self.sim_output) == len(self.ensemble.next_sims): - # If we are streaming, then the simulation results are - # directly routed to the training task via ProxyStream. - # So, we don't need to submit an extra training task. - if not self.streaming: - self.submit_task('train', self.sim_output) - self.logger.info('Submitting training task') + # Submit the train task + self.submit_task('train', self.sim_output) + self.logger.info('Submitting training task') # If it's okay to use the stale model, submit the inference task # using the previous iteration's model @@ -160,13 +142,14 @@ def process_train_result(self, result: Result) -> None: self.done.set() return + # See if this is the first training task return value + first_train = self.train_output is None + # Store the training output self.train_output = result.value - # TODO: What should we do in the streaming case? - # Does the process_train_result method even run? # Submit an inference task with the simulation/train task outputs - if not self.streaming: + if first_train or not self.use_stale_model: self.submit_task('inference', self.sim_output, self.train_output) self.logger.info('submitted inference task') @@ -211,3 +194,201 @@ def process_inference_result(self, result: Result) -> None: self.logger.info('Submitting next iteration of simulations') for sim in self.ensemble.next_sims: self.submit_task('simulation', sim) + + +class DDWEStreamThinker(BaseThinker): + """A thinker for the DDWE workflow.""" + + def __init__( + self, + queue: ColmenaQueues, + result_dir: Path, + ensemble: WeightedEnsemble, + checkpointer: EnsembleCheckpointer, + num_iterations: int, + stream_config: ProxyStreamConfig, + use_stale_model: bool = False, + max_retries: int = 2, + ) -> None: + """Initialize the DDWE workflow thinker. + + Parameters + ---------- + queue: ColmenaQueues + Queue used to communicate with the task server. + result_dir: Path + Directory in which to store outputs. + ensemble: WeightedEnsemble + The weighted ensemble to use for the workflow. + checkpointer: EnsembleCheckpointer + Checkpointer for the weighted ensemble. + num_iterations: int + Number of iterations to run the workflow. + stream_config: ProxyStreamConfig + Configuration for the data stream. + use_stale_model: bool + Whether to use the stale model for inference (default to False). + This will be faster but may not be as accurate. It uses the + model from the previous iteration for inference in the current + iteration, which may not be updated with new states. + max_retries: int + Number of times to retry a task if it fails (default to 2). + """ + super().__init__(queue) + + self.ensemble = ensemble + self.checkpointer = checkpointer + self.num_iterations = num_iterations + self.stream_config = stream_config + self.use_stale_model = use_stale_model + self.max_retries = max_retries + self.result_logger = ResultLogger(result_dir) + + # Store the simulation output (the input of both train/inference tasks) + self.sim_output: list[Any] = [] + + # TODO: These two attributes need to be checkpointed and restored + # Store the train output (the input of the inference task) + self.train_output: Any = None + # Keep a counter for the current training iteration + self.train_iteration = ensemble.iteration - 1 + + self.stream_consumer = stream_config.get_consumer(topic='train-output') + + def submit_task(self, topic: str, *inputs: Any) -> None: + """Submit a task to the task server. + + Parameters + ---------- + topic: str + The topic of the task. + inputs: Any + The input args to the task. + """ + # Submit the task to the task server + self.queues.send_inputs( + *inputs, + method=f'run_{topic}', + topic=topic, + max_retries=self.max_retries, + ) + + @agent(startup=True) + def start_workflow(self) -> None: + """Launch the first iteration of simulations to start the workflow.""" + # Submit the next iteration of simulations + for sim in self.ensemble.next_sims: + self.submit_task('simulation', sim) + + # We need to submit a single train task at the start of the workflow + # to kick off the simulation stream consumer + self.logger.info('Start streaming train task') + self.submit_task('train') + + @result_processor(topic='simulation') + def process_simulation_result(self, result: Result) -> None: + """Process a simulation result.""" + # Log simulation job results + self.result_logger.log(result, topic='simulation') + + # Check if the task failed + if not result.success: + self.logger.error( + f'Simulation failed after {result.retries}' + f'/{result.max_retries} attempts, quitting workflow.', + f' result={result}', + ) + self.stop_workflow() + return + + # Collect simulation results for the current iteration + self.sim_output.append(result.value) + + # If we have all the simulation results, submit the inference task + # using the previous iteration's model + if len(self.sim_output) == len(self.ensemble.next_sims): + # We need to wait for the streaming train task to finish + if not self.use_stale_model: + # Wait for the streaming train task to finish + self.logger.info('Waiting for streaming train task to finish') + while self.train_output is None: + time.sleep(10) + + elif self.use_stale_model: + self.logger.info('Waiting for streaming train task to finish') + while self.train_iteration < self.ensemble.iteration: + time.sleep(10) + # This should hold (see train_stream_processor) + assert self.train_output is not None + + # If it's okay to use the stale model, submit the inference task + self.submit_task('inference', self.sim_output, self.train_output) + + @agent() + def train_stream_processor(self) -> None: + """Process the streaming train task.""" + # This for loop will run until the producer closes the topic + # (see stop_workflow) + for result in self.stream_consumer: + # Log a message for each train result + self.logger.info('Received streaming train result') + + # Store the training output + self.train_output = result + + # Increment the training iteration + self.train_iteration += 1 + + def stop_workflow(self) -> None: + """Stop the workflow.""" + # Set the done flag to signal the agents to stop + self.done.set() + + # Close the stream consumer (we use the producer to close the topic) + producer = self.stream_config.get_producer(topic='train-output') + producer.close_topics('train-output') + + # Log a message that the workflow is stopping + self.logger.info('Stopping the workflow') + + @result_processor(topic='inference') + def process_inference_result(self, result: Result) -> None: + """Process an inference result.""" + # Log inference job results + self.result_logger.log(result, topic='inference') + + # Check if the task failed + if not result.success: + self.logger.warning('Inference failed, quitting workflow.') + self.stop_workflow() + return + + # Unpack the output + cur_sims, next_sims, metadata = result.value + + # Update the weighted ensemble with the next iteration + self.ensemble.advance_iteration( + cur_sims=cur_sims, + next_sims=next_sims, + metadata=metadata, + ) + + # Save an ensemble checkpoint + self.checkpointer.save(self.ensemble) + + # Log the current iteration + self.logger.info(f'Current iteration: {self.ensemble.iteration}') + + # Reset the simulation output for the next iteration + self.sim_output = [] + + # Check if the workflow is finished (if so return before submitting) + if self.ensemble.iteration >= self.num_iterations: + self.logger.info('Workflow finished') + self.done.set() + return + + # Submit the next iteration of simulations + self.logger.info('Submitting next iteration of simulations') + for sim in self.ensemble.next_sims: + self.submit_task('simulation', sim) diff --git a/deepdrivewe/workflows/stream.py b/deepdrivewe/workflows/stream.py new file mode 100644 index 0000000..7a4947f --- /dev/null +++ b/deepdrivewe/workflows/stream.py @@ -0,0 +1,80 @@ +"""Streaming configuration for the DeepDriveWE workflow.""" + +from __future__ import annotations + +from typing import Any + +from proxystore.store import get_store +from proxystore.store import register_store +from proxystore.store import Store +from proxystore.store.config import StoreConfig +from proxystore.stream.interface import StreamConsumer +from proxystore.stream.interface import StreamProducer +from proxystore.stream.shims.redis import RedisQueuePublisher +from proxystore.stream.shims.redis import RedisQueueSubscriber + +from deepdrivewe import BaseModel + + +class ProxyStreamConfig(BaseModel): + """Configuration for the proxy stream.""" + + store_config: StoreConfig + redis_host: str = 'localhost' + redis_port: int = 6379 + + def get_store(self) -> Store[Any]: + """Get the store for the proxy stream. + + Returns + ------- + Store + The store for the proxy stream. + """ + store = get_store(self.store_config.name) + if store is None: + store = Store.from_config(self.store_config) + register_store(store, exist_ok=True) + return store + + # The StreamConsumer is generic on the type of the stream items. + def get_consumer(self, topic: str) -> StreamConsumer[Any]: + """Get a consumer for a given topic. + + Parameters + ---------- + topic: str + The topic to consume. + + Returns + ------- + StreamConsumer + The consumer for the given topic. + """ + # The RedisQueueSubscriber is *not* a broadcasting stream. I.e., each + # stream item will only be consumed by one subscriber (the subscriber + # that wins the race). For multi-consumer support, see the + # RedisSubscriber and RedisPublisher. + subscriber = RedisQueueSubscriber( + self.redis_host, + self.redis_port, + topic=topic, + ) + return StreamConsumer(subscriber) + + def get_producer(self, topic: str) -> StreamProducer[Any]: + """Get a producer for a given topic. + + Parameters + ---------- + topic: str + The topic to produce. + + Returns + ------- + StreamProducer + The producer for the given topic. + """ + store = self.get_store() + publisher = RedisQueuePublisher(self.redis_host, self.redis_port) + return StreamProducer(publisher, {topic: store}) diff --git a/pyproject.toml b/pyproject.toml index 4a9452c..9fdde55 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,7 +145,7 @@ select = [ "RUF", ] line-length = 79 -extend-ignore = ["Q001"] +extend-ignore = ["Q001", "SIM102", "PLR0913"] target-version = "py38" [tool.ruff.flake8-pytest-style] From eca6295a745647d28a2433dfabd7a718e45accfe Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 30 Oct 2024 22:23:17 -0500 Subject: [PATCH 09/73] idletime --- deepdrivewe/parsl.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/deepdrivewe/parsl.py b/deepdrivewe/parsl.py index 16500fb..eecdd41 100644 --- a/deepdrivewe/parsl.py +++ b/deepdrivewe/parsl.py @@ -253,6 +253,14 @@ class InferenceTrainWorkstationConfig(BaseComputeConfig): description='Config for the GPU executor to run AI models.', ) + # We have a long idletime to ensure train/inference executors are not + # shut down (to enable warmstarts) while simulations are running. + max_idletime: float = Field( + default=60.0 * 10, + description='The maximum idle time allowed for an executor before ' + 'strategy could shut down unused blocks. Default is 10 minutes.', + ) + @model_validator(mode='after') def validate_htex_labels(self) -> Self: """Ensure that the labels are unique.""" @@ -266,6 +274,7 @@ def get_parsl_config(self, run_dir: str | Path) -> Config: return Config( run_dir=str(run_dir), retries=self.train_gpu_config.retries, + max_idletime=self.max_idletime, executors=[ HighThroughputExecutor( address='localhost', From 90aecb974d5ded61aec54534e2f52eae39a1d596 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 31 Oct 2024 10:08:58 -0500 Subject: [PATCH 10/73] random seed --- deepdrivewe/simulation/openmm.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index ab80070..5681b70 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -455,10 +455,6 @@ class OpenMMConfig(BaseModel): default=True, description='Whether to randomize the basis state initial velocities.', ) - seed: int = Field( - default=np.random.default_rng().integers(2**32, dtype=int), - description='The random seed for the simulation.', - ) hardware_platform: str = Field( default='CUDA', description='The hardware platform to use for the simulation.' @@ -704,9 +700,11 @@ def configure_simulation( ValueError If explicit solvent is selected and no topology file is provided. """ - # Set the random seed - random.seed(self.seed) - np.random.seed(self.seed) + # Set the random seed (we use a different seed for each simulation + # to ensure simulations sample trajectories). + seed = np.random.default_rng().integers(2**32, dtype=int) + random.seed(seed) + np.random.seed(seed) # Select implicit or explicit solvent configuration and load the system if self.solvent_type == 'explicit': From 65977428df42054961a771c8ce0d36e0afc68383 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 31 Oct 2024 12:09:04 -0500 Subject: [PATCH 11/73] Stream logic for openmm, cvae training, workflow. Co-authored-by: Greg Pauloski <18683347+gpauloski@users.noreply.github.com> --- deepdrivewe/examples/openmm_ntl9_ddwe/main.py | 4 +- .../examples/openmm_ntl9_ddwe/simulate.py | 8 ++ .../examples/openmm_ntl9_ddwe/train.py | 102 ++++++++++++++++++ deepdrivewe/simulation/openmm.py | 47 +++++--- deepdrivewe/workflows/ddwe.py | 23 +++- deepdrivewe/workflows/stream.py | 3 + 6 files changed, 169 insertions(+), 18 deletions(-) diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py index 4d60eeb..53f41df 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py @@ -105,7 +105,7 @@ class ExperimentSettings(BaseModel): description='Config for the compute resources.', ) stream_config: ProxyStreamConfig | None = Field( - None, + default=None, description='Stream configuration for simulation data.', ) use_stale_model: bool = Field( @@ -196,11 +196,13 @@ def mkdir_validator(cls, value: Path) -> Path: run_simulation, config=cfg.simulation_config, output_dir=cfg.output_dir / 'simulation', + stream_config=cfg.stream_config, ) my_run_train = partial( run_train, config=cfg.train_config, output_dir=cfg.output_dir / 'train', + stream_config=cfg.stream_config, ) my_run_inference = partial( run_inference, diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py index 8a2b0f5..cd2230c 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py @@ -15,6 +15,7 @@ from deepdrivewe.simulation.openmm import OpenMMConfig from deepdrivewe.simulation.openmm import OpenMMSimulation from deepdrivewe.simulation.openmm import RMSDCollector +from deepdrivewe.workflows.stream import ProxyStreamConfig class SimulationConfig(BaseModel): @@ -48,6 +49,7 @@ def run_simulation( metadata: SimMetadata, config: SimulationConfig, output_dir: Path, + stream_config: ProxyStreamConfig | None = None, ) -> SimResult: """Run a simulation and return the pcoord and coordinates.""" # Add performance logging @@ -91,6 +93,7 @@ def run_simulation( topic='pcoords', ), ], + stream_config=stream_config, ) # Run the simulation @@ -104,6 +107,11 @@ def run_simulation( metadata.pcoord = data['pcoords'].tolist() metadata.mark_simulation_end() + # If we are streaming the data, only keep the last frame + # for use in the inference module. + if stream_config is not None: + data = {key: value[:-1] for key, value in data.items()} + result = SimResult(data=data, metadata=metadata) return result diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py index 680b5e1..6005ddb 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py @@ -2,6 +2,7 @@ from __future__ import annotations +import itertools from pathlib import Path import numpy as np @@ -12,6 +13,9 @@ from deepdrivewe import TrainResult from deepdrivewe.ai import ConvolutionalVAE from deepdrivewe.ai import ConvolutionalVAEConfig +from deepdrivewe.workflows.stream import ProxyStreamConfig +from deepdrivewe.workflows.stream import SIMULATION_TOPIC +from deepdrivewe.workflows.stream import TRAIN_TOPIC class TrainConfig(BaseModel): @@ -25,6 +29,16 @@ class TrainConfig(BaseModel): description='The path to the model checkpoint file.' 'Train from scratch by default.', ) + stream_items_per_train: int = Field( + default=1, + description='The number of items (simulation frames) to train on in ' + 'each stream iteration.', + ) + stream_retrain_interval: int = Field( + default=1, + description='The number of stream training iterations between ' + 're-initializing and re-training the model.', + ) # TODO: We probably need to store a history of old training data @@ -36,8 +50,18 @@ def run_train( sim_output: list[SimResult], config: TrainConfig, output_dir: Path, + stream_config: ProxyStreamConfig | None = None, ) -> TrainResult: """Train the model on the simulation output.""" + # If we are using a stream, run the stream training function + if stream_config is not None: + return run_stream_train( + sim_output=sim_output, + config=config, + output_dir=output_dir, + stream_config=stream_config, + ) + # Make the output directory itetation = sim_output[0].metadata.iteration_id output_dir = output_dir / f'{itetation:06d}' @@ -73,3 +97,81 @@ def run_train( ) return result + + +def run_stream_train( + sim_output: list[SimResult], + config: TrainConfig, + output_dir: Path, + stream_config: ProxyStreamConfig, +) -> TrainResult: + """Train the model on the simulation output.""" + # Make the output directory + itetation = sim_output[0].metadata.iteration_id + output_dir = output_dir / f'{itetation:06d}' + output_dir.mkdir(parents=True, exist_ok=True) + + # Stream consumer for getting new simulation data + stream_consumer = stream_config.get_consumer(topic=SIMULATION_TOPIC) + # Stream producer for sending new trained model weights to the thinker + stream_producer = stream_config.get_producer(topic=TRAIN_TOPIC) + + # TODO: Decide how much data we want to keep in the re-train history. + contact_map_history = [] + pcoord_history = [] + + # Loop indefinitely until we get a stop iteration from the stream consumer + for idx in itertools.count(): + # If we have reached the retrain interval, re-initialize the trainer + # NOTE: This always happens on the first iteration + if idx % config.stream_retrain_interval == 0: + # Load the model configuration + model_config = ConvolutionalVAEConfig.from_yaml(config.config_path) + + # Load the model + model = ConvolutionalVAE( + model_config, + checkpoint_path=config.checkpoint_path, + ) + + # Get the next batch of simulation data from the stream. + # Each item is a dictionary with topic keys defined in the simulation + # module, (e.g. 'contact_maps', 'pcoords', etc.), and values are + # numpy arrays representing a single frame of data. + try: + items = [ + next(stream_consumer) + for _ in range(config.stream_items_per_train) + ] + except StopIteration: + break + + # Extract the contact maps and rmsd from each simulation + contact_maps = np.concatenate([x['contact_maps'] for x in items]) + pcoords = np.concatenate([x['pcoords'] for x in items]) + pcoords = pcoords.flatten() + + # TODO: It might be necessary to put these into a numpy array + # Concatenate the new data with the history + contact_map_history.extend(contact_maps) + pcoord_history.extend(pcoords) + + # Fit the model + checkpoint_path = model.fit( + x=contact_map_history, + model_dir=output_dir / 'model', + scalars={'pcoord': pcoord_history}, + ) + + # Construct the train result + result = TrainResult( + config_path=config.config_path, + checkpoint_path=checkpoint_path, + ) + + # Send the new model weights to the thinker + stream_producer.send(topic=TRAIN_TOPIC, obj=result) + + # NOTE: This final return is not necessary, but it is included + # to keep the function signature consistent with the non-streaming. + return result diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 5681b70..cc2a6ac 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -8,6 +8,7 @@ from abc import ABC from abc import abstractmethod from pathlib import Path +from typing import Any from typing import Sequence if sys.version_info >= (3, 11): # pragma: >=3.11 cover @@ -24,6 +25,8 @@ from pydantic import Field from pydantic import model_validator +from deepdrivewe.workflows.stream import ProxyStreamConfig +from deepdrivewe.workflows.stream import SIMULATION_TOPIC from deepdrivewe.workflows.utils import retry_on_exception try: @@ -152,7 +155,7 @@ def get(self) -> np.ndarray: ... @abstractmethod - def collect(self, positions: np.ndarray) -> None: + def collect(self, positions: np.ndarray) -> Any: """Collect data from the simulation. Parameters @@ -228,7 +231,7 @@ def _align(self, positions: np.ndarray) -> np.ndarray: return aligned_positions - def collect(self, positions: np.ndarray) -> None: + def collect(self, positions: np.ndarray) -> np.ndarray: """Generate a report. Parameters @@ -242,6 +245,8 @@ def collect(self, positions: np.ndarray) -> None: # Collect the position coordinates self._coordinates.append(pos) + return pos + class RMSDCollector(Collector): """RMSD collector for OpenMM simulations.""" @@ -281,7 +286,7 @@ def get(self) -> np.ndarray: """ return np.array(self._rmsd).reshape(-1, 1) - def collect(self, positions: np.ndarray) -> None: + def collect(self, positions: np.ndarray) -> float: """Generate a report. Parameters @@ -292,6 +297,7 @@ def collect(self, positions: np.ndarray) -> None: # Compute the RMSD rmsd = rms.rmsd(positions, self._ref, superposition=True) self._rmsd.append(rmsd) + return rmsd class ContactMapCollector(Collector): @@ -305,8 +311,7 @@ def __init__( """Initialize the contact map collector.""" super().__init__(topic) self.cutoff_angstrom = cutoff_angstrom - self._rows: list[np.ndarray] = [] - self._cols: list[np.ndarray] = [] + self._contact_maps: list[np.ndarray] = [] def get(self) -> np.ndarray: """Get the contact maps from the simulation. @@ -317,15 +322,12 @@ def get(self) -> np.ndarray: The contact maps from the simulation as a ragged array shaped as (n_frames, *). """ - # Concatenate the row and col indices into a single array - contact_maps = [np.concatenate(x) for x in zip(self._rows, self._cols)] - # Collect the contact maps in a ragged numpy array - contact_maps = np.array(contact_maps, dtype=object) + contact_maps = np.array(self._contact_maps, dtype=object) return contact_maps - def collect(self, positions: np.ndarray) -> None: + def collect(self, positions: np.ndarray) -> np.ndarray: """Generate a report. Parameters @@ -343,9 +345,15 @@ def collect(self, positions: np.ndarray) -> None: # Convert the contact map to sparse format coo_matrix = contact_map.tocoo() + # Get the row and col indices and concatenate them + row = coo_matrix.row.astype('int16') + col = coo_matrix.col.astype('int16') + sparse_contact_map = np.concatenate([row, col]) + # Append the row and col indices to lists - self._rows.append(coo_matrix.row.astype('int16')) - self._cols.append(coo_matrix.col.astype('int16')) + self._contact_maps.append(sparse_contact_map) + + return sparse_contact_map class CollectionReporter(OpenMMReporter): @@ -356,6 +364,7 @@ def __init__( report_interval: int, collectors: list[Collector], openmm_selection: Sequence[str] = ('CA',), + stream_config: ProxyStreamConfig | None = None, ) -> None: """Initialize the reporter. @@ -382,6 +391,13 @@ def __init__( self.collectors = collectors + # Initialize the streaming producer + self.producer = None + if stream_config is not None: + self.producer = stream_config.get_producer( + topic=SIMULATION_TOPIC, + ) + def get_collected_data(self) -> dict[str, np.ndarray]: """Get the collected data from the simulation. @@ -406,8 +422,11 @@ def report(self, simulation: app.Simulation, state: openmm.State) -> None: positions = self.get_positions(simulation, state) # Collect data from the simulation - for collector in self.collectors: - collector.collect(positions) + data = {x.topic: x.collect(positions) for x in self.collectors} + + # Stream the data if a stream config is provided + if self.producer is not None: + self.producer.send(topic=SIMULATION_TOPIC, obj=data, evict=True) class OpenMMConfig(BaseModel): diff --git a/deepdrivewe/workflows/ddwe.py b/deepdrivewe/workflows/ddwe.py index 3989f79..82a5393 100644 --- a/deepdrivewe/workflows/ddwe.py +++ b/deepdrivewe/workflows/ddwe.py @@ -12,10 +12,13 @@ from colmena.thinker import BaseThinker from colmena.thinker import result_processor from proxystore.proxy import extract +from proxystore.store.utils import get_key from deepdrivewe import EnsembleCheckpointer from deepdrivewe import WeightedEnsemble from deepdrivewe.workflows.stream import ProxyStreamConfig +from deepdrivewe.workflows.stream import SIMULATION_TOPIC +from deepdrivewe.workflows.stream import TRAIN_TOPIC from deepdrivewe.workflows.utils import ResultLogger @@ -253,7 +256,10 @@ def __init__( # Keep a counter for the current training iteration self.train_iteration = ensemble.iteration - 1 - self.stream_consumer = stream_config.get_consumer(topic='train-output') + # Create a consumer for streaming the training return objects + # to the thinker. + self.stream_config = stream_config + self.stream_consumer = stream_config.get_consumer(topic=TRAIN_TOPIC) def submit_task(self, topic: str, *inputs: Any) -> None: """Submit a task to the task server. @@ -333,6 +339,13 @@ def train_stream_processor(self) -> None: # Log a message for each train result self.logger.info('Received streaming train result') + # Clean up the previous training output from the store + if self.train_output is not None: + # Get the proxy key for the current training output + key = get_key(self.train_output) + # Evict the key from the store to clean up memory + self.stream_config.get_store().evict(key) + # Store the training output self.train_output = result @@ -345,8 +358,12 @@ def stop_workflow(self) -> None: self.done.set() # Close the stream consumer (we use the producer to close the topic) - producer = self.stream_config.get_producer(topic='train-output') - producer.close_topics('train-output') + # NOTE: Closing the train topic, will close the stream_consumer in the + # thinker which will stop the train_stream_processor agent, and closing + # the simulation topic will close the training function consumer + # waiting for new simulation results. + for topic in [TRAIN_TOPIC, SIMULATION_TOPIC]: + self.stream_config.get_producer(topic=topic).close_topics(topic) # Log a message that the workflow is stopping self.logger.info('Stopping the workflow') diff --git a/deepdrivewe/workflows/stream.py b/deepdrivewe/workflows/stream.py index 7a4947f..16235d8 100644 --- a/deepdrivewe/workflows/stream.py +++ b/deepdrivewe/workflows/stream.py @@ -15,6 +15,9 @@ from deepdrivewe import BaseModel +SIMULATION_TOPIC = 'simulation-output' +TRAIN_TOPIC = 'train-output' + class ProxyStreamConfig(BaseModel): """Configuration for the proxy stream.""" From c8b168b7461fd3656e6e62539544490406401eeb Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 31 Oct 2024 14:29:00 -0500 Subject: [PATCH 12/73] random seed --- deepdrivewe/simulation/openmm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index cc2a6ac..cdb285b 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -724,6 +724,7 @@ def configure_simulation( seed = np.random.default_rng().integers(2**32, dtype=int) random.seed(seed) np.random.seed(seed) + print(f'Running simulation with seed: {seed}') # Select implicit or explicit solvent configuration and load the system if self.solvent_type == 'explicit': From 3945c719776745ec3029b04f85052f2707cd1e5c Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 31 Oct 2024 14:29:58 -0500 Subject: [PATCH 13/73] random seed --- deepdrivewe/simulation/openmm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index cdb285b..4df16f7 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -724,7 +724,7 @@ def configure_simulation( seed = np.random.default_rng().integers(2**32, dtype=int) random.seed(seed) np.random.seed(seed) - print(f'Running simulation with seed: {seed}') + print(f'Running simulation with seed: {seed}', flush=True) # Select implicit or explicit solvent configuration and load the system if self.solvent_type == 'explicit': From 4adc14037bd0d68eff8037c43dbcad10f294a41f Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 31 Oct 2024 14:58:50 -0500 Subject: [PATCH 14/73] train output dir, mock input --- deepdrivewe/examples/openmm_ntl9_ddwe/train.py | 9 ++++----- deepdrivewe/workflows/ddwe.py | 5 +++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py index 6005ddb..ea980c6 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py @@ -56,7 +56,6 @@ def run_train( # If we are using a stream, run the stream training function if stream_config is not None: return run_stream_train( - sim_output=sim_output, config=config, output_dir=output_dir, stream_config=stream_config, @@ -100,15 +99,12 @@ def run_train( def run_stream_train( - sim_output: list[SimResult], config: TrainConfig, output_dir: Path, stream_config: ProxyStreamConfig, ) -> TrainResult: """Train the model on the simulation output.""" # Make the output directory - itetation = sim_output[0].metadata.iteration_id - output_dir = output_dir / f'{itetation:06d}' output_dir.mkdir(parents=True, exist_ok=True) # Stream consumer for getting new simulation data @@ -156,10 +152,13 @@ def run_stream_train( contact_map_history.extend(contact_maps) pcoord_history.extend(pcoords) + # Make a new model directory for this iteration + model_dir = output_dir / f'model_{idx:06d}' + # Fit the model checkpoint_path = model.fit( x=contact_map_history, - model_dir=output_dir / 'model', + model_dir=model_dir, scalars={'pcoord': pcoord_history}, ) diff --git a/deepdrivewe/workflows/ddwe.py b/deepdrivewe/workflows/ddwe.py index 82a5393..c293b91 100644 --- a/deepdrivewe/workflows/ddwe.py +++ b/deepdrivewe/workflows/ddwe.py @@ -287,9 +287,10 @@ def start_workflow(self) -> None: self.submit_task('simulation', sim) # We need to submit a single train task at the start of the workflow - # to kick off the simulation stream consumer + # to kick off the simulation stream consumer. We send an empty list + # of simulation outputs to be compatible with the train task signature. self.logger.info('Start streaming train task') - self.submit_task('train') + self.submit_task('train', []) @result_processor(topic='simulation') def process_simulation_result(self, result: Result) -> None: From 73748853c8a58dc9d528258e08ded1460653c781 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:12:12 -0500 Subject: [PATCH 15/73] seed --- deepdrivewe/simulation/openmm.py | 15 ++++++++------- examples/openmm_ntl9_ddwe/config.yaml | 2 +- examples/openmm_ntl9_ddwe_vista/config.yaml | 2 +- examples/openmm_ntl9_hk/config.yaml | 2 +- 4 files changed, 11 insertions(+), 10 deletions(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 4df16f7..3ddee8e 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -719,13 +719,6 @@ def configure_simulation( ValueError If explicit solvent is selected and no topology file is provided. """ - # Set the random seed (we use a different seed for each simulation - # to ensure simulations sample trajectories). - seed = np.random.default_rng().integers(2**32, dtype=int) - random.seed(seed) - np.random.seed(seed) - print(f'Running simulation with seed: {seed}', flush=True) - # Select implicit or explicit solvent configuration and load the system if self.solvent_type == 'explicit': if top_file is None: @@ -877,6 +870,14 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: if openmm_checkpoint.exists(): sim.loadCheckpoint(str(openmm_checkpoint)) + # Set the random seed (we use a different seed for each simulation + # to ensure simulations sample different trajectories). + seed = np.random.default_rng().integers(2**31 - 1, dtype=int) + random.seed(seed) + np.random.seed(seed) + sim.integrator.setRandomNumberSeed(seed) + print(f'Running simulation with seed: {seed}', flush=True) + # Run simulation sim.step(self.config.num_steps) diff --git a/examples/openmm_ntl9_ddwe/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index 52dbb71..846731a 100644 --- a/examples/openmm_ntl9_ddwe/config.yaml +++ b/examples/openmm_ntl9_ddwe/config.yaml @@ -34,7 +34,7 @@ simulation_config: # The time step to use in picoseconds dt_ps: 0.002 # The temperature to run the simulation at - temperature: 300.0 + temperature_kelvin: 300.0 # The solvent type solvent_type: implicit # The hardware platform to run the simulation on diff --git a/examples/openmm_ntl9_ddwe_vista/config.yaml b/examples/openmm_ntl9_ddwe_vista/config.yaml index c59f936..4e7744d 100644 --- a/examples/openmm_ntl9_ddwe_vista/config.yaml +++ b/examples/openmm_ntl9_ddwe_vista/config.yaml @@ -34,7 +34,7 @@ simulation_config: # The time step to use in picoseconds dt_ps: 0.002 # The temperature to run the simulation at - temperature: 300.0 + temperature_kelvin: 300.0 # The solvent type solvent_type: implicit # The hardware platform to run the simulation on diff --git a/examples/openmm_ntl9_hk/config.yaml b/examples/openmm_ntl9_hk/config.yaml index 177a0af..55873a5 100644 --- a/examples/openmm_ntl9_hk/config.yaml +++ b/examples/openmm_ntl9_hk/config.yaml @@ -31,7 +31,7 @@ simulation_config: # The time step to use in picoseconds dt_ps: 0.002 # The temperature to run the simulation at - temperature: 300.0 + temperature_kelvin: 300.0 # The solvent type solvent_type: implicit # The hardware platform to run the simulation on From 2552d8569564daea6c9f2f4af600ba9120134cc2 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:28:29 -0500 Subject: [PATCH 16/73] seed --- deepdrivewe/simulation/openmm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 3ddee8e..d4d19c5 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -876,6 +876,7 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: random.seed(seed) np.random.seed(seed) sim.integrator.setRandomNumberSeed(seed) + sim.context.setParameter('RandomSeed', seed) print(f'Running simulation with seed: {seed}', flush=True) # Run simulation From 4ec63dfa1ce405cbf0f7b62fc2eff947ebbe5655 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:30:01 -0500 Subject: [PATCH 17/73] seed --- deepdrivewe/simulation/openmm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index d4d19c5..7ec9d1a 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -876,7 +876,7 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: random.seed(seed) np.random.seed(seed) sim.integrator.setRandomNumberSeed(seed) - sim.context.setParameter('RandomSeed', seed) + sim.context.reinitialize(preserveState=True) print(f'Running simulation with seed: {seed}', flush=True) # Run simulation From 8c3f5a9c821d9cb56c1b38748368c093f4ac236b Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:32:41 -0500 Subject: [PATCH 18/73] seed --- deepdrivewe/simulation/openmm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 7ec9d1a..a2d37ce 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -878,6 +878,7 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: sim.integrator.setRandomNumberSeed(seed) sim.context.reinitialize(preserveState=True) print(f'Running simulation with seed: {seed}', flush=True) + print(sim.context.getParameters()) # Run simulation sim.step(self.config.num_steps) From abfff3b18a301d6bbd894a7f317994e72cc69e14 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:33:58 -0500 Subject: [PATCH 19/73] seed --- deepdrivewe/simulation/openmm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index a2d37ce..e8d4315 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -878,7 +878,8 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: sim.integrator.setRandomNumberSeed(seed) sim.context.reinitialize(preserveState=True) print(f'Running simulation with seed: {seed}', flush=True) - print(sim.context.getParameters()) + for param_name, param_value in sim.context.getParameters().items(): + print(f'{param_name}: {param_value}') # Run simulation sim.step(self.config.num_steps) From 081d840a2eac2f11e2d7c7c7045035c178ba7ea6 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:36:17 -0500 Subject: [PATCH 20/73] seed --- deepdrivewe/simulation/openmm.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index e8d4315..fc5e29b 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -875,11 +875,15 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: seed = np.random.default_rng().integers(2**31 - 1, dtype=int) random.seed(seed) np.random.seed(seed) - sim.integrator.setRandomNumberSeed(seed) - sim.context.reinitialize(preserveState=True) + # sim.integrator.setRandomNumberSeed(seed) + # sim.context.reinitialize(preserveState=True) print(f'Running simulation with seed: {seed}', flush=True) - for param_name, param_value in sim.context.getParameters().items(): - print(f'{param_name}: {param_value}') + + new_integrator = self.config.configure_integrator() + new_integrator.setRandomNumberSeed(seed) + + # Replace the old integrator in the simulation with the new one + sim.context.setIntegrator(new_integrator) # Run simulation sim.step(self.config.num_steps) From 15b1fe6772e4238798869df6b031dae175f10ccc Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:38:36 -0500 Subject: [PATCH 21/73] seed --- deepdrivewe/simulation/openmm.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index fc5e29b..698d0ed 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -882,14 +882,36 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: new_integrator = self.config.configure_integrator() new_integrator.setRandomNumberSeed(seed) - # Replace the old integrator in the simulation with the new one - sim.context.setIntegrator(new_integrator) + # Step 3: Create a new Simulation with the existing System and Context, + # but with the new integrator + # This effectively applies the new RNG seed + new_simulation = app.Simulation( + sim.topology, + sim.system, + new_integrator, + sim.context.getPlatform(), + ) + + # Step 4: Set the state from the existing context to continue the + # simulation + state = sim.context.getState( + getPositions=True, + getVelocities=True, + getEnergy=True, + getForces=True, + ) + new_simulation.context.setState(state) + + new_simulation.step(self.config.num_steps) + + # Save a checkpoint of the final state + new_simulation.saveCheckpoint(str(self.output_dir / 'seg.chk')) # Run simulation - sim.step(self.config.num_steps) + # sim.step(self.config.num_steps) # Save a checkpoint of the final state - sim.saveCheckpoint(str(self.output_dir / 'seg.chk')) + # sim.saveCheckpoint(str(self.output_dir / 'seg.chk')) # TODO: First test the above implementation, then remove this class. From 19e589347dab7caa75036a6a6427622b7839cf0a Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 1 Nov 2024 12:39:54 -0500 Subject: [PATCH 22/73] seed --- deepdrivewe/simulation/openmm.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 698d0ed..aad3a1a 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -902,6 +902,8 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: ) new_simulation.context.setState(state) + new_simulation.reporters = sim.reporters + new_simulation.step(self.config.num_steps) # Save a checkpoint of the final state From 8a7425996d66eb84ee625a9877184623781775ff Mon Sep 17 00:00:00 2001 From: braceal Date: Sat, 2 Nov 2024 10:29:25 -0500 Subject: [PATCH 23/73] update mdlearn API --- deepdrivewe/ai/aae.py | 8 +++----- deepdrivewe/ai/cvae.py | 8 +++----- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/deepdrivewe/ai/aae.py b/deepdrivewe/ai/aae.py index 91fc31d..ddc5b14 100644 --- a/deepdrivewe/ai/aae.py +++ b/deepdrivewe/ai/aae.py @@ -129,10 +129,8 @@ def __init__( self.config = config self.checkpoint_path = checkpoint_path - # We keep the inference_batch_size in the config for convenience - # but exclude it from the model arguments. - model_args = config.model_dump(exclude={'inference_batch_size'}) - self.trainer = AAE3dTrainer(**model_args) + # Initialize the model + self.trainer = AAE3dTrainer(**config.model_dump()) # Load the model checkpoint if specified if checkpoint_path is not None: @@ -218,7 +216,7 @@ def predict(self, x: np.ndarray) -> np.ndarray: The predicted latent space coordinates (n_samples, latent_dim). """ # Predict the latent space coordinates - z, _ = self.trainer.predict(x, self.config.inference_batch_size) + z, _ = self.trainer.predict(x) return z diff --git a/deepdrivewe/ai/cvae.py b/deepdrivewe/ai/cvae.py index 241861c..38ff402 100644 --- a/deepdrivewe/ai/cvae.py +++ b/deepdrivewe/ai/cvae.py @@ -127,10 +127,8 @@ def __init__( self.config = config self.checkpoint_path = checkpoint_path - # We keep the inference_batch_size in the config for convenience - # but exclude it from the model arguments. - model_args = config.model_dump(exclude={'inference_batch_size'}) - self.trainer = SymmetricConv2dVAETrainer(**model_args) + # Initialize the model + self.trainer = SymmetricConv2dVAETrainer(**config.model_dump()) # Load the model checkpoint if specified if checkpoint_path is not None: @@ -216,7 +214,7 @@ def predict(self, x: np.ndarray) -> np.ndarray: The predicted latent space coordinates (n_samples, latent_dim). """ # Predict the latent space coordinates - z, *_ = self.trainer.predict(x, self.config.inference_batch_size) + z, *_ = self.trainer.predict(x) return z From 6f7cd8d174f812864c8cc50e2868d7e345256f71 Mon Sep 17 00:00:00 2001 From: braceal Date: Sat, 2 Nov 2024 11:10:57 -0500 Subject: [PATCH 24/73] mdlearn version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9fdde55..9e7f367 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "mdtraj==1.10.0", "MDAnalysis>=2.7.0", "scikit-learn==1.5.1", - "mdlearn==1.0.3", + "mdlearn==1.0.4", "scipy==1.14.0", "natsort>=8.4.0", "matplotlib>=3.9.2", From e01778080cecb682f6851a04368be01c2290fce8 Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 5 Nov 2024 11:26:59 -0600 Subject: [PATCH 25/73] fix rng bug --- deepdrivewe/simulation/openmm.py | 82 +++++++++++++------------------- 1 file changed, 34 insertions(+), 48 deletions(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index aad3a1a..81cc5a9 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -697,6 +697,7 @@ def configure_simulation( self, pdb_file: str | Path, top_file: str | Path | None, + checkpoint_file: str | Path | None = None, ) -> app.Simulation: """Configure an OpenMM simulation. @@ -708,6 +709,8 @@ def configure_simulation( top_file : str | Path | None The optional topology file to initialize the systems topology (required for explicit solvent). + checkpoint_file : str | Path | None + The checkpoint file to initialize the simulation. Returns ------- @@ -751,6 +754,30 @@ def configure_simulation( platform_properties, ) + # Load the checkpoint file if provided (skips setting positions + # from PDB, minimization, and randomizing velocities) + if checkpoint_file is not None: + # Create a new Simulation with the existing system context, + # but with the new integrator (which applies the new RNG seed) + new_simulation = app.Simulation( + sim.topology, + sim.system, + self.configure_integrator(), + sim.context.getPlatform(), + *self.configure_hardware(), + ) + + # Set the state from the existing context to continue the sim + state = sim.context.getState( + getPositions=True, + getVelocities=True, + getEnergy=True, + getForces=True, + ) + new_simulation.context.setState(state) + + return new_simulation + # Set the positions if self.set_positions: pdb = app.PDBFile(str(pdb_file)) @@ -831,10 +858,15 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: if self.copy_input_files and self.top_file is not None: self.top_file = shutil.copy(self.top_file, self.output_dir) + # Attempt to locate a checkpoint file + chk_file = self.checkpoint_file.parent / 'seg.chk' + checkpoint_file = chk_file if chk_file.exists() else None + # Initialize an OpenMM simulation sim = self.config.configure_simulation( pdb_file=self.restart_file, top_file=self.top_file, + checkpoint_file=checkpoint_file, ) # Set up a reporter to write a simulation trajectory file @@ -863,57 +895,11 @@ def run(self, reporters: list[OpenMMReporter] | None = None) -> None: if reporters is not None: sim.reporters.extend(reporters) - # Attempt to locate a checkpoint file - openmm_checkpoint = self.checkpoint_file.parent / 'seg.chk' - - # Load the checkpoint file (if it is a OpenMM checkpoint) - if openmm_checkpoint.exists(): - sim.loadCheckpoint(str(openmm_checkpoint)) - - # Set the random seed (we use a different seed for each simulation - # to ensure simulations sample different trajectories). - seed = np.random.default_rng().integers(2**31 - 1, dtype=int) - random.seed(seed) - np.random.seed(seed) - # sim.integrator.setRandomNumberSeed(seed) - # sim.context.reinitialize(preserveState=True) - print(f'Running simulation with seed: {seed}', flush=True) - - new_integrator = self.config.configure_integrator() - new_integrator.setRandomNumberSeed(seed) - - # Step 3: Create a new Simulation with the existing System and Context, - # but with the new integrator - # This effectively applies the new RNG seed - new_simulation = app.Simulation( - sim.topology, - sim.system, - new_integrator, - sim.context.getPlatform(), - ) - - # Step 4: Set the state from the existing context to continue the - # simulation - state = sim.context.getState( - getPositions=True, - getVelocities=True, - getEnergy=True, - getForces=True, - ) - new_simulation.context.setState(state) - - new_simulation.reporters = sim.reporters - - new_simulation.step(self.config.num_steps) - - # Save a checkpoint of the final state - new_simulation.saveCheckpoint(str(self.output_dir / 'seg.chk')) - # Run simulation - # sim.step(self.config.num_steps) + sim.step(self.config.num_steps) # Save a checkpoint of the final state - # sim.saveCheckpoint(str(self.output_dir / 'seg.chk')) + sim.saveCheckpoint(str(self.output_dir / 'seg.chk')) # TODO: First test the above implementation, then remove this class. From 7e6acede3e66278474f6bdad0d6884740568d40c Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 5 Nov 2024 18:47:33 -0600 Subject: [PATCH 26/73] fix rng bug --- deepdrivewe/simulation/openmm.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 81cc5a9..8d8ec00 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -757,6 +757,9 @@ def configure_simulation( # Load the checkpoint file if provided (skips setting positions # from PDB, minimization, and randomizing velocities) if checkpoint_file is not None: + # Load the checkpoint file + sim.loadCheckpoint(str(checkpoint_file)) + # Create a new Simulation with the existing system context, # but with the new integrator (which applies the new RNG seed) new_simulation = app.Simulation( From 7838f5a1c07f60a7b4bbee998c7dbb549e61843d Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 5 Nov 2024 19:35:53 -0600 Subject: [PATCH 27/73] fix rng bug --- deepdrivewe/simulation/openmm.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 8d8ec00..03b0ca6 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -766,8 +766,10 @@ def configure_simulation( sim.topology, sim.system, self.configure_integrator(), - sim.context.getPlatform(), - *self.configure_hardware(), + platform, + platform_properties, + # sim.context.getPlatform(), + # *self.configure_hardware(), ) # Set the state from the existing context to continue the sim From 782986600503b5ab2c76e645abcf71248e8253e9 Mon Sep 17 00:00:00 2001 From: braceal Date: Mon, 11 Nov 2024 15:08:43 -0600 Subject: [PATCH 28/73] mdlearn version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9e7f367..37175d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "mdtraj==1.10.0", "MDAnalysis>=2.7.0", "scikit-learn==1.5.1", - "mdlearn==1.0.4", + "mdlearn==1.0.5", "scipy==1.14.0", "natsort>=8.4.0", "matplotlib>=3.9.2", From 3cb4a334ee54cfcb73d48c4d30fcb0077d495d2e Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 12 Nov 2024 14:29:20 -0600 Subject: [PATCH 29/73] docs --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8a94d0d..2b67eb2 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ pip install -e . ``` To use deep learning models, install the correct version of [PyTorch](https://pytorch.org/get-started/locally/) -for your system and drivers. To use `mdlearn`, you may need an earlier version of PyTorch: +for your system and drivers: ```bash pip install torch==1.12 ``` From 57590b00a75ea179010148a0b4e9a36f62112315 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 16:35:45 -0600 Subject: [PATCH 30/73] parmed explicit topology loader --- deepdrivewe/simulation/openmm.py | 11 +++++++++-- pyproject.toml | 1 + 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 03b0ca6..74c93d6 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -18,6 +18,7 @@ import MDAnalysis import numpy as np +import parmed as pmd from MDAnalysis.analysis import align from MDAnalysis.analysis import distances from MDAnalysis.analysis import rms @@ -526,6 +527,7 @@ def report_steps(self) -> int: def load_explicit_system_from_top( self, top_file: str | Path, + pdb_file: str | Path, ) -> tuple[openmm.System, app.Topology]: """Load an explicit solvent system from a topology file. @@ -533,6 +535,8 @@ def load_explicit_system_from_top( ---------- top_file : str | Path The topology file to load the system from. + pdb_file : str | Path + The PDB file to load the system topology. Returns ------- @@ -540,7 +544,7 @@ def load_explicit_system_from_top( The OpenMM system and topology. """ # Load the topology file - top = app.AmberPrmtopFile(str(top_file)) + top = pmd.load_file(str(top_file), xyz=str(pdb_file)) # Configure system system = top.createSystem( @@ -728,7 +732,10 @@ def configure_simulation( raise ValueError( 'Topology file must be provided for explicit solvent.', ) - system, topology = self.load_explicit_system_from_top(top_file) + system, topology = self.load_explicit_system_from_top( + pdb_file, + top_file, + ) elif top_file is not None: system, topology = self.load_implicit_system_from_top(top_file) else: diff --git a/pyproject.toml b/pyproject.toml index 37175d8..7c8e833 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "h5py==3.11.0", "mdtraj==1.10.0", "MDAnalysis>=2.7.0", + "ParmEd>=4.3.0", "scikit-learn==1.5.1", "mdlearn==1.0.5", "scipy==1.14.0", From 920469456aa658c3bddfddd8283b5b7dfec7953f Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 18:21:23 -0600 Subject: [PATCH 31/73] polaris config --- deepdrivewe/parsl.py | 85 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/deepdrivewe/parsl.py b/deepdrivewe/parsl.py index eecdd41..e9f0c78 100644 --- a/deepdrivewe/parsl.py +++ b/deepdrivewe/parsl.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os import sys from abc import ABC from abc import abstractmethod @@ -358,6 +359,89 @@ def get_parsl_config(self, run_dir: str | Path) -> Config: ) +class PolarisConfig(BaseComputeConfig): + """Compute config for a workstation.""" + + name: Literal['polaris'] = 'polaris' # type: ignore[assignment] + + num_nodes: int = Field( + ge=3, + description='Number of nodes to use (must use at least 3 nodes).', + ) + retries: int = Field( + default=1, + description='Number of retries for the task.', + ) + # We have a long idletime to ensure train/inference executors are not + # shut down (to enable warmstarts) while simulations are running. + max_idletime: float = Field( + default=60.0 * 10, + description='The maximum idle time allowed for an executor before ' + 'strategy could shut down unused blocks. Default is 10 minutes.', + ) + + def _write_nodefiles(self, run_dir: Path) -> None: + """Write nodefiles for the train, inference, and simulation tasks.""" + # Get the nodefile + node_file = os.environ['PBS_NODEFILE'] + with open(node_file) as fp: + hosts = [x.strip() for x in fp] + + # Determine the node files for each task type + labels = ['train', 'inference', 'simulation'] + hostnames = [hosts[0], hosts[1], hosts[2:]] + + # Write the nodefiles for each task type + for label, hnames in zip(labels, hostnames): + nodefile = run_dir / f'{label}.hosts' + nodefile.write_text('\n'.join(hnames)) + + def _get_htex( + self, + label: str, + num_nodes: int, + run_dir: Path, + ) -> HighThroughputExecutor: + hostfile = run_dir / f'{label}.hosts' + return HighThroughputExecutor( + label=label, + cpu_affinity='block-reverse', + available_accelerators=4, + provider=LocalProvider( + launcher=WrappedLauncher( + f'mpiexec -n {num_nodes} --ppn 1 --hostfile ' + f'{hostfile} --depth=64 --cpu-bind depth', + ), + cmd_timeout=120, + nodes_per_block=num_nodes, + init_blocks=1, + max_blocks=1, + ), + ) + + def get_parsl_config(self, run_dir: str | Path) -> Config: + """Generate a Parsl configuration.""" + # Convert run_dir to a Path object + run_dir = Path(run_dir) + + # Write the nodefiles for each task type + self._write_nodefiles(run_dir) + + # Return the Parsl configuration + return Config( + run_dir=str(run_dir), + retries=self.retries, + max_idletime=self.max_idletime, + executors=[ + # Assign 1 GPU each for training and inference + self._get_htex('train_htex', 1, run_dir), + self._get_htex('inference_htex', 1, run_dir), + # Assign the remaining GPUs to simulation + self._get_htex('simulation_htex', self.num_nodes - 2, run_dir), + ], + ) + + ComputeConfigTypes = Union[ LocalConfig, WorkstationConfig, @@ -365,4 +449,5 @@ def get_parsl_config(self, run_dir: str | Path) -> Config: HybridWorkstationConfig, InferenceTrainWorkstationConfig, VistaConfig, + PolarisConfig, ] From 91b7b8e46d33a68d619c43a6f684a93973810eb6 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 20:22:24 -0600 Subject: [PATCH 32/73] docs --- README.md | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/README.md b/README.md index 2b67eb2..22a2766 100644 --- a/README.md +++ b/README.md @@ -20,11 +20,7 @@ conda install conda-forge::openmm==7.7 -y pip install -e . ``` -To use deep learning models, install the correct version of [PyTorch](https://pytorch.org/get-started/locally/) -for your system and drivers: -```bash -pip install torch==1.12 -``` +To use deep learning models, install the correct version of [PyTorch](https://pytorch.org/get-started/locally/). ### Installation on VISTA From f86eb7df4c7a814beecc73cd1f7f41ae80f7cd41 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 20:38:53 -0600 Subject: [PATCH 33/73] docs. parsl address --- README.md | 14 ++++++++++++++ deepdrivewe/parsl.py | 20 ++++++++++++++++++-- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 22a2766..7fd5b73 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ To install the package, run the following command: ```bash git clone git@github.com:braceal/deepdrivewe.git cd deepdrivewe +pip install -U pip setuptools wheel pip install -e . ``` @@ -17,6 +18,7 @@ cd deepdrivewe conda create -n deepdrivewe python=3.10 -y conda install omnia::ambertools -y conda install conda-forge::openmm==7.7 -y +pip install -U pip setuptools wheel pip install -e . ``` @@ -44,6 +46,18 @@ and the YAML config file, and then run the following command: sbatch examples/openmm_ntl9_ddwe_vista/submit.sh ``` +### Installation on Polaris + +To install the package on Polaris@ALCF, run the following commands: +```bash +module use /soft/modulefiles; module load conda +``` + +Follow the full installation instructions above, and install torch via: +```bash +pip install torch +``` + ## Usage To run the example, run the following command: ```bash diff --git a/deepdrivewe/parsl.py b/deepdrivewe/parsl.py index e9f0c78..401db6c 100644 --- a/deepdrivewe/parsl.py +++ b/deepdrivewe/parsl.py @@ -16,7 +16,7 @@ else: # pragma: <3.11 cover from typing_extensions import Self - +from parsl.addresses import address_by_hostname from parsl.config import Config from parsl.executors import HighThroughputExecutor from parsl.launchers import WrappedLauncher @@ -152,6 +152,22 @@ class WorkstationV2Config(BaseComputeConfig): description='The maximum idle time allowed for an executor before ' 'strategy could shut down unused blocks. Default is 10 minutes.', ) + address: str = Field( + default='localhost', + description='Address to bind the executor to [localhost, hostname].', + ) + + @model_validator(mode='after') + def validate_address(self) -> Self: + """Check that the address is valid.""" + if self.address not in ('localhost', 'hostname'): + raise ValueError('Address must be either localhost or hostname.') + + # Get the hostname if the address is 'hostname' + if self.address == 'hostname': + self.address = address_by_hostname() + + return self @model_validator(mode='after') def validate_available_accelerators(self) -> Self: @@ -170,7 +186,7 @@ def _get_htex( available_accelerators: Sequence[str], ) -> HighThroughputExecutor: return HighThroughputExecutor( - address='localhost', + address=self.address, label=label, cpu_affinity='block', available_accelerators=available_accelerators, From b561398d39bd2caeca0be3b10bd4ba610940b830 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 20:50:59 -0600 Subject: [PATCH 34/73] docs --- README.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 7fd5b73..9f32794 100644 --- a/README.md +++ b/README.md @@ -71,8 +71,8 @@ ps -e | grep -E 'sander|python|process_worker|parsl' | awk '{print $1}' | xargs To check if any errors occurred in simulations or inference: ```bash -cat runs/naive_resampler_test_v2/result/inference.json | grep '"success": false' -cat runs/naive_resampler_test_v2/result/simulation.json | grep '"success": false' +cat runs/*/result/inference.json | grep '"success": false' +cat runs/*/result/simulation.json | grep '"success": false' ``` To check the number of iterations completed: @@ -80,6 +80,11 @@ To check the number of iterations completed: h5ls -d runs/naive_resampler_test_v2/west.h5/iterations ``` +To watch the progress of the simulation: +```bash +tail -f runs/*/simulation/*/*/*.log +``` + ### Running with SynD To use the SynD simulation engine, install the following dependencies: ```bash From c57c20ec5579e6747c38f1b79a6f12526ff7126b Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:09:13 -0600 Subject: [PATCH 35/73] cli --- deepdrivewe/cli.py | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index 2a8da7f..f5627f5 100644 --- a/deepdrivewe/cli.py +++ b/deepdrivewe/cli.py @@ -2,9 +2,12 @@ from __future__ import annotations +import json from pathlib import Path import typer +from rich import print +from rich.console import Console app = typer.Typer() @@ -17,6 +20,44 @@ def version() -> None: print(f'deepdrivewe, version {__version__}') +app = typer.Typer() + + +@app.command() +def print_errors( + run_dir: Path = typer.Option( # noqa: B008 + ..., + '--run_dir', + '-r', + help='Path to the run directory.', + ), +) -> None: + """Parse the task result files and print any errors.""" + console = Console() + + # Find all the task result files + results_dir = run_dir / 'results' + + # Read the simulation, train, and inference results + for file_path in results_dir.glob('*.jsonl'): + # Read the entire file as text + file_text = file_path.read_text() + + # Parse each line as JSON + for line in file_text.splitlines(): + data = json.loads(line) + if 'failure_info' in data and 'traceback' in data['failure_info']: + error_message = data['failure_info']['traceback'] + console.print( + f"[bold red]Task ID:[/bold red] {data['task_id']}", + ) + console.print( + f"[bold yellow]Method:[/bold yellow] {data['method']}", + ) + console.print('[bold blue]Traceback:[/bold blue]\n') + console.print(error_message, style='red') + + @app.command() def to_pdb( coordinate_file: Path = typer.Option( # noqa: B008 From eacde114d0613398c91f8712ae907e9c75bbd922 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:10:30 -0600 Subject: [PATCH 36/73] cli --- deepdrivewe/cli.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index f5627f5..c2fa015 100644 --- a/deepdrivewe/cli.py +++ b/deepdrivewe/cli.py @@ -20,9 +20,6 @@ def version() -> None: print(f'deepdrivewe, version {__version__}') -app = typer.Typer() - - @app.command() def print_errors( run_dir: Path = typer.Option( # noqa: B008 From 2bd54d5df3b93cc897ed737569dab146c546219b Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:11:27 -0600 Subject: [PATCH 37/73] cli --- deepdrivewe/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index c2fa015..889e1cd 100644 --- a/deepdrivewe/cli.py +++ b/deepdrivewe/cli.py @@ -36,7 +36,7 @@ def print_errors( results_dir = run_dir / 'results' # Read the simulation, train, and inference results - for file_path in results_dir.glob('*.jsonl'): + for file_path in results_dir.glob('*.json'): # Read the entire file as text file_text = file_path.read_text() From 3b97b4c2a0f389beaf0a72c676ea88da8c1979e2 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:12:27 -0600 Subject: [PATCH 38/73] cli --- deepdrivewe/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index 889e1cd..f7df57a 100644 --- a/deepdrivewe/cli.py +++ b/deepdrivewe/cli.py @@ -33,7 +33,7 @@ def print_errors( console = Console() # Find all the task result files - results_dir = run_dir / 'results' + results_dir = run_dir / 'result' # Read the simulation, train, and inference results for file_path in results_dir.glob('*.json'): From 600221da1e5f967d9969778099a83236ffbb94ce Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:15:11 -0600 Subject: [PATCH 39/73] cli --- deepdrivewe/cli.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index f7df57a..b53c15c 100644 --- a/deepdrivewe/cli.py +++ b/deepdrivewe/cli.py @@ -7,7 +7,6 @@ import typer from rich import print -from rich.console import Console app = typer.Typer() @@ -30,7 +29,7 @@ def print_errors( ), ) -> None: """Parse the task result files and print any errors.""" - console = Console() + # console = Console() # Find all the task result files results_dir = run_dir / 'result' @@ -44,15 +43,21 @@ def print_errors( for line in file_text.splitlines(): data = json.loads(line) if 'failure_info' in data and 'traceback' in data['failure_info']: - error_message = data['failure_info']['traceback'] - console.print( - f"[bold red]Task ID:[/bold red] {data['task_id']}", - ) - console.print( + # error_message = data['failure_info']['traceback'] + # console.print( + # f"[bold red]Task ID:[/bold red] {data['task_id']}", + # ) + # console.print( + # f"[bold yellow]Method:[/bold yellow] {data['method']}", + # ) + # console.print('[bold blue]Traceback:[/bold blue]\n') + # console.print(error_message, style='red') + typer.echo(f"[bold red]Task ID:[/bold red] {data['task_id']}") + typer.echo( f"[bold yellow]Method:[/bold yellow] {data['method']}", ) - console.print('[bold blue]Traceback:[/bold blue]\n') - console.print(error_message, style='red') + typer.echo('[bold blue]Traceback:[/bold blue]\n') + typer.echo(data['failure_info']['traceback']) @app.command() From 6a6708034bb685a1fd60aba8969ad8128bed331a Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:20:11 -0600 Subject: [PATCH 40/73] cli --- deepdrivewe/cli.py | 21 ++++++--------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index b53c15c..9aae6f5 100644 --- a/deepdrivewe/cli.py +++ b/deepdrivewe/cli.py @@ -7,6 +7,7 @@ import typer from rich import print +from rich.console import Console app = typer.Typer() @@ -29,7 +30,7 @@ def print_errors( ), ) -> None: """Parse the task result files and print any errors.""" - # console = Console() + console = Console() # Find all the task result files results_dir = run_dir / 'result' @@ -43,21 +44,11 @@ def print_errors( for line in file_text.splitlines(): data = json.loads(line) if 'failure_info' in data and 'traceback' in data['failure_info']: - # error_message = data['failure_info']['traceback'] - # console.print( - # f"[bold red]Task ID:[/bold red] {data['task_id']}", - # ) - # console.print( - # f"[bold yellow]Method:[/bold yellow] {data['method']}", - # ) - # console.print('[bold blue]Traceback:[/bold blue]\n') - # console.print(error_message, style='red') - typer.echo(f"[bold red]Task ID:[/bold red] {data['task_id']}") - typer.echo( - f"[bold yellow]Method:[/bold yellow] {data['method']}", + console.print( + f"[bold green]Method:[/bold green] {data['method']}", ) - typer.echo('[bold blue]Traceback:[/bold blue]\n') - typer.echo(data['failure_info']['traceback']) + console.print('[bold blue]Traceback:[/bold blue]\n') + console.print(data['failure_info']['traceback'], style='red') @app.command() From 9a7269990713dd01778f446dfa180ef22f11a362 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:22:59 -0600 Subject: [PATCH 41/73] cli --- deepdrivewe/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index 9aae6f5..e6a36ae 100644 --- a/deepdrivewe/cli.py +++ b/deepdrivewe/cli.py @@ -30,6 +30,7 @@ def print_errors( ), ) -> None: """Parse the task result files and print any errors.""" + # Create a console for rich output console = Console() # Find all the task result files @@ -45,9 +46,8 @@ def print_errors( data = json.loads(line) if 'failure_info' in data and 'traceback' in data['failure_info']: console.print( - f"[bold green]Method:[/bold green] {data['method']}", + f"[bold blue]Method:[/bold blue] {data['method']}", ) - console.print('[bold blue]Traceback:[/bold blue]\n') console.print(data['failure_info']['traceback'], style='red') From 42e187785162efa07c584927282a930a7c834a97 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 14 Nov 2024 21:24:56 -0600 Subject: [PATCH 42/73] cli --- README.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/README.md b/README.md index 9f32794..f75e5ec 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,16 @@ To watch the progress of the simulation: tail -f runs/*/simulation/*/*/*.log ``` +To pretty print potential errors: +```bash +deepdrivewe print-errors --run_dir runs/ntl9-v1 +``` + +Run the following, for more information: +```bash +deepdrivewe --help +``` + ### Running with SynD To use the SynD simulation engine, install the following dependencies: ```bash From a12fd6a5d65006df3d9331ceed8d184c2ae727a8 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 15 Nov 2024 08:42:42 -0600 Subject: [PATCH 43/73] cli --- deepdrivewe/parsl.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepdrivewe/parsl.py b/deepdrivewe/parsl.py index 401db6c..5dd6b67 100644 --- a/deepdrivewe/parsl.py +++ b/deepdrivewe/parsl.py @@ -437,8 +437,9 @@ def _get_htex( def get_parsl_config(self, run_dir: str | Path) -> Config: """Generate a Parsl configuration.""" - # Convert run_dir to a Path object + # Convert run_dir to a Path object and create the directory run_dir = Path(run_dir) + run_dir.mkdir(parents=True, exist_ok=True) # Write the nodefiles for each task type self._write_nodefiles(run_dir) From 23d6147ffc397a2c801cd9797c4166a39cd65fa2 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 15 Nov 2024 10:37:12 -0600 Subject: [PATCH 44/73] parsl --- deepdrivewe/parsl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/parsl.py b/deepdrivewe/parsl.py index 5dd6b67..7837e06 100644 --- a/deepdrivewe/parsl.py +++ b/deepdrivewe/parsl.py @@ -404,7 +404,7 @@ def _write_nodefiles(self, run_dir: Path) -> None: hosts = [x.strip() for x in fp] # Determine the node files for each task type - labels = ['train', 'inference', 'simulation'] + labels = ['train_htex', 'inference_htex', 'simulation_htex'] hostnames = [hosts[0], hosts[1], hosts[2:]] # Write the nodefiles for each task type From 35175cd64a17bc8ee9138e4784cbb91cadaf7929 Mon Sep 17 00:00:00 2001 From: braceal Date: Fri, 15 Nov 2024 10:43:40 -0600 Subject: [PATCH 45/73] remove xyx from parmed --- deepdrivewe/simulation/openmm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index 74c93d6..441309a 100644 --- a/deepdrivewe/simulation/openmm.py +++ b/deepdrivewe/simulation/openmm.py @@ -544,7 +544,7 @@ def load_explicit_system_from_top( The OpenMM system and topology. """ # Load the topology file - top = pmd.load_file(str(top_file), xyz=str(pdb_file)) + top = pmd.load_file(str(top_file), str(pdb_file)) # Configure system system = top.createSystem( From 7c61c8fb9dab580fbb7bb9715b9f283a528e38e1 Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 21 Oct 2025 14:23:14 -0500 Subject: [PATCH 46/73] fix proxystore import --- deepdrivewe/workflows/stream.py | 4 ++-- pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepdrivewe/workflows/stream.py b/deepdrivewe/workflows/stream.py index 16235d8..be36aba 100644 --- a/deepdrivewe/workflows/stream.py +++ b/deepdrivewe/workflows/stream.py @@ -8,8 +8,8 @@ from proxystore.store import register_store from proxystore.store import Store from proxystore.store.config import StoreConfig -from proxystore.stream.interface import StreamConsumer -from proxystore.stream.interface import StreamProducer +from proxystore.stream import StreamConsumer +from proxystore.stream import StreamProducer from proxystore.stream.shims.redis import RedisQueuePublisher from proxystore.stream.shims.redis import RedisQueueSubscriber diff --git a/pyproject.toml b/pyproject.toml index 7c8e833..e6d2f89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ classifiers = [ ] dependencies = [ "colmena>=0.7.0", - "proxystore>=0.7.0", + "proxystore>=0.8.3", "parsl>=2024.10.14", "pyyaml>=6.0.1", "typer>=0.12.5", From 6544235f4caddc9165bcb62394168b59ee791c10 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 23 Oct 2025 11:43:26 -0500 Subject: [PATCH 47/73] fix shape issue during concatenation --- deepdrivewe/examples/openmm_ntl9_ddwe/train.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py index ea980c6..1bbd031 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py @@ -75,9 +75,10 @@ def run_train( checkpoint_path=config.checkpoint_path, ) - # Extract the last frame contact maps and rmsd from each simulation - contact_maps = np.concatenate( - [sim.data['contact_maps'] for sim in sim_output], + # Flatten all contact maps and pcoords from all sims into a single array + contact_maps = np.array( + [cm for sim in sim_output for cm in sim.data['contact_maps']], + dtype=object, ) pcoords = np.concatenate([sim.data['pcoords'] for sim in sim_output]) pcoords = pcoords.flatten() From c562eb19a8789527f02bbaa79e00fd40395ab355 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 23 Oct 2025 11:48:54 -0500 Subject: [PATCH 48/73] fix paths in example --- examples/openmm_ntl9_ddwe/config.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/openmm_ntl9_ddwe/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index 846731a..a3bb5ff 100644 --- a/examples/openmm_ntl9_ddwe/config.yaml +++ b/examples/openmm_ntl9_ddwe/config.yaml @@ -9,7 +9,7 @@ num_iterations: 106 # The basis states to use for the ensemble basis_states: # The directory containing the basis states sub directories - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_hk/inputs + basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/examples/openmm_ntl9_ddwe/inputs # The file extension for the basis state files basis_state_ext: .pdb # The number of basis states to use @@ -21,7 +21,7 @@ basis_states: # Strategy for initializing the basis state progress coordinates basis_state_initializer: # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_hk/common_files/reference.pdb + reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/examples/openmm_ntl9_ddwe/common_files/reference.pdb # The configuration for the simulation simulation_config: @@ -41,7 +41,7 @@ simulation_config: hardware_platform: CPU # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_hk/common_files/reference.pdb + reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/examples/openmm_ntl9_ddwe/common_files/reference.pdb # The configuration for training train_config: From 944f49c306974a6ae13b56197035c51f2670c876 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 23 Oct 2025 12:12:44 -0500 Subject: [PATCH 49/73] fix paths in example --- examples/openmm_ntl9_ddwe/config.yaml | 60 +++++++++++++++------------ 1 file changed, 34 insertions(+), 26 deletions(-) diff --git a/examples/openmm_ntl9_ddwe/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index a3bb5ff..03af59f 100644 --- a/examples/openmm_ntl9_ddwe/config.yaml +++ b/examples/openmm_ntl9_ddwe/config.yaml @@ -9,7 +9,7 @@ num_iterations: 106 # The basis states to use for the ensemble basis_states: # The directory containing the basis states sub directories - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/examples/openmm_ntl9_ddwe/inputs + basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/inputs # The file extension for the basis state files basis_state_ext: .pdb # The number of basis states to use @@ -21,7 +21,7 @@ basis_states: # Strategy for initializing the basis state progress coordinates basis_state_initializer: # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/examples/openmm_ntl9_ddwe/common_files/reference.pdb + reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/common_files/reference.pdb # The configuration for the simulation simulation_config: @@ -41,7 +41,7 @@ simulation_config: hardware_platform: CPU # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/examples/openmm_ntl9_ddwe/common_files/reference.pdb + reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/common_files/reference.pdb # The configuration for training train_config: @@ -72,29 +72,37 @@ target_states: pcoord: [1.0] # The settings for the compute environment -compute_config: - # The name of the compute environment to use (CPU + GPU) - name: inference_train_workstation +# compute_config: +# # The name of the compute environment to use (CPU + GPU) +# name: inference_train_workstation + +# # The CPU configuration for simulation +# cpu_config: +# # Specify we want the local parsl configuration +# name: local +# # The maximum number of worker processes to use for parallelization +# max_workers_per_node: 53 - # The CPU configuration for simulation - cpu_config: - # Specify we want the local parsl configuration - name: local - # The maximum number of worker processes to use for parallelization - max_workers_per_node: 53 +# # The GPU configuration for training +# train_gpu_config: +# # Specify we want the workstation parsl configuration +# name: workstation +# # Identify which GPUs to assign tasks to. It's generally recommended to first check +# # nvidia-smi to see which GPUs are available. The numbers below are analogous to +# # setting CUDA_VISIBLE_DEVICES=2 +# available_accelerators: ["2"] - # The GPU configuration for training - train_gpu_config: - # Specify we want the workstation parsl configuration - name: workstation - # Identify which GPUs to assign tasks to. It's generally recommended to first check - # nvidia-smi to see which GPUs are available. The numbers below are analogous to - # setting CUDA_VISIBLE_DEVICES=2 - available_accelerators: ["2"] +# # The GPU configuration for inference +# inference_gpu_config: +# # Specify we want the workstation parsl configuration +# name: workstation +# # Identify which GPUs to assign tasks to +# available_accelerators: ["3"] - # The GPU configuration for inference - inference_gpu_config: - # Specify we want the workstation parsl configuration - name: workstation - # Identify which GPUs to assign tasks to - available_accelerators: ["3"] +compute_config: + # The name of the compute environment to use + name: workstation_v2 + # Set the address + address: hostname + # Identify which GPUs to assign tasks to + available_accelerators: ["1", "2", "3"] From a9c61775f961b709c4673e76d90c13a9e0070332 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 23 Oct 2025 13:42:45 -0500 Subject: [PATCH 50/73] config --- examples/openmm_ntl9_ddwe/config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/openmm_ntl9_ddwe/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index 03af59f..b2d6cb6 100644 --- a/examples/openmm_ntl9_ddwe/config.yaml +++ b/examples/openmm_ntl9_ddwe/config.yaml @@ -38,7 +38,7 @@ simulation_config: # The solvent type solvent_type: implicit # The hardware platform to run the simulation on - hardware_platform: CPU + hardware_platform: CUDA # The path to the reference PDB file reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/common_files/reference.pdb From ff1d118f60c05d7332dd22dc1696d38578612c46 Mon Sep 17 00:00:00 2001 From: jml230 Date: Tue, 28 Oct 2025 15:53:25 -0500 Subject: [PATCH 51/73] multirectilinear binner --- deepdrivewe/binners/__init__.py | 1 + deepdrivewe/binners/multirectilinear.py | 100 ++++++++++++++++++++++++ 2 files changed, 101 insertions(+) create mode 100644 deepdrivewe/binners/multirectilinear.py diff --git a/deepdrivewe/binners/__init__.py b/deepdrivewe/binners/__init__.py index abad4f5..a74bf4d 100644 --- a/deepdrivewe/binners/__init__.py +++ b/deepdrivewe/binners/__init__.py @@ -5,3 +5,4 @@ # Forward imports from deepdrivewe.binners.base import Binner from deepdrivewe.binners.rectilinear import RectilinearBinner +from deepdrivewe.binners.multirectilinear import MultiRectilienarBinner diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py new file mode 100644 index 0000000..8a5224d --- /dev/null +++ b/deepdrivewe/binners/multirectilinear.py @@ -0,0 +1,100 @@ +"""Rectilinear binner.""" + +from __future__ import annotations + +import numpy as np + +from deepdrivewe.binners.base import Binner + + +class MultiRectilinearBinner(Binner): + """Rectilinear binner for the progress coordinate.""" + + def __init__( + self, + bins: list[np.ndarray | list[float]], + bin_target_counts: int | list[int], + target_state_inds: int | list[int] = 0, + ) -> None: + """Initialize the binner. + + Parameters + ---------- + bins : list[np.ndarray | list[float]] + The bin edges for the progress coordinate. + bin_target_counts : int | list[int] + The target counts for each bin. If an integer is provided, + the target counts are assumed to be the same for each bin. + target_state_inds : int | list[int] + The index of the target state. If an integer is provided, then + there is only one target state. If a list of integers is provided, + then there are multiple target states. Default is 0 which + corresponds to the first bin. + """ + self.bins = bins + self.bin_target_counts = bin_target_counts + self.target_state_inds = target_state_inds + + # Check that the bins are sorted + for binbounds in self.bins: + if not np.all(np.diff(binbounds) > 0): + raise ValueError('Bin boundaries must be sorted in ascending order.') + + @property + def nbins(self) -> int: + """The number of bins.""" + nbins_per_dim = np.array([len(dim)-1 for dim in self.bins]) + return np.prod(nbins_per_dim) + + def get_bin_target_counts(self) -> list[int]: + """Get the target counts for each bin. + + Returns + ------- + list[int] + The target counts for each bin. + """ + # Check if the bin target counts is an integer + # If so, then set the target counts for each bin to the same value + # and set the target state bins to 0. Cache the result. + if isinstance(self.bin_target_counts, int): + # Create a list of the bin target counts + bin_target_counts = [self.bin_target_counts] * self.nbins + + # Get the target state indices (convert to a list if an integer) + if isinstance(self.target_state_inds, int): + self.target_state_inds = [self.target_state_inds] + + # Set each of the target state bins to 0 since they are recycled + for i in self.target_state_inds: + bin_target_counts[i] = 0 + + # Cache the result + self.bin_target_counts = bin_target_counts + + # Otherwise, return the list of bin target counts + return self.bin_target_counts + + def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: + """Bin the progress coordinate. + + Parameters + ---------- + pcoords : np.ndarray + The progress coordinates to bin. Shape: (n_simulations, n_dims). + + Returns + ------- + np.ndarray + The bin assignments for each simulation. Shape: (n_simulations,) + """ + # Bin the progress coordinates (make sure the target state + # boundary is included in the target state bin). + _, x_edge, _, bid = binned_statistic_2d(*pcoords.T, values=None, statistic='count', bins=self.bins, expand_binnumbers=True) + + bin_ids = np.array([(ibid[0] -1) * (len(x_edge)-1) + ibid[1] for ibid in bid]) + assert len(bin_ids) == len(pcoords) + + return bin_ids + + # return np.digitize(pcoords[:, self.pcoord_idx], self.bins, right=True) From 70162427df744ee7ca4ea670943228546598e7b9 Mon Sep 17 00:00:00 2001 From: jml230 Date: Tue, 28 Oct 2025 15:53:56 -0500 Subject: [PATCH 52/73] update pyproject.toml, unpin scipy mdtraj --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e6d2f89..11a57b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,13 +25,13 @@ dependencies = [ "pyyaml>=6.0.1", "typer>=0.12.5", "numpy", - "h5py==3.11.0", - "mdtraj==1.10.0", + "h5py", + "mdtraj", "MDAnalysis>=2.7.0", "ParmEd>=4.3.0", "scikit-learn==1.5.1", "mdlearn==1.0.5", - "scipy==1.14.0", + "scipy", "natsort>=8.4.0", "matplotlib>=3.9.2", ] From ce077fde63870827d14d70fb9beffc682cd98bfb Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 28 Oct 2025 16:04:42 -0500 Subject: [PATCH 53/73] pre-commit --- deepdrivewe/binners/__init__.py | 2 +- deepdrivewe/binners/multirectilinear.py | 38 ++++++++++++++++++------- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/deepdrivewe/binners/__init__.py b/deepdrivewe/binners/__init__.py index a74bf4d..b44ba76 100644 --- a/deepdrivewe/binners/__init__.py +++ b/deepdrivewe/binners/__init__.py @@ -4,5 +4,5 @@ # Forward imports from deepdrivewe.binners.base import Binner +from deepdrivewe.binners.multirectilinear import MultiRectilinearBinner from deepdrivewe.binners.rectilinear import RectilinearBinner -from deepdrivewe.binners.multirectilinear import MultiRectilienarBinner diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 8a5224d..d587f67 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -1,14 +1,15 @@ -"""Rectilinear binner.""" +"""Multirectilinear binner.""" from __future__ import annotations import numpy as np +from scipy.stats import binned_statistic_2d from deepdrivewe.binners.base import Binner class MultiRectilinearBinner(Binner): - """Rectilinear binner for the progress coordinate.""" + """Multirectilinear binner for multiple progress coordinates.""" def __init__( self, @@ -21,7 +22,7 @@ def __init__( Parameters ---------- bins : list[np.ndarray | list[float]] - The bin edges for the progress coordinate. + The bin edges for the progress coordinates. bin_target_counts : int | list[int] The target counts for each bin. If an integer is provided, the target counts are assumed to be the same for each bin. @@ -38,12 +39,14 @@ def __init__( # Check that the bins are sorted for binbounds in self.bins: if not np.all(np.diff(binbounds) > 0): - raise ValueError('Bin boundaries must be sorted in ascending order.') + raise ValueError( + 'Bin boundaries must be sorted in ascending order.', + ) @property def nbins(self) -> int: """The number of bins.""" - nbins_per_dim = np.array([len(dim)-1 for dim in self.bins]) + nbins_per_dim = np.array([len(dim) - 1 for dim in self.bins]) return np.prod(nbins_per_dim) def get_bin_target_counts(self) -> list[int]: @@ -90,11 +93,24 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """ # Bin the progress coordinates (make sure the target state # boundary is included in the target state bin). - _, x_edge, _, bid = binned_statistic_2d(*pcoords.T, values=None, statistic='count', bins=self.bins, expand_binnumbers=True) - - bin_ids = np.array([(ibid[0] -1) * (len(x_edge)-1) + ibid[1] for ibid in bid]) - assert len(bin_ids) == len(pcoords) + _, x_edge, _, bid = binned_statistic_2d( + *pcoords.T, + values=None, + statistic='count', + bins=self.bins, + expand_binnumbers=True, + ) + + # Convert the bin edges to bin indices + bin_ids = np.array( + [(ibid[0] - 1) * (len(x_edge) - 1) + ibid[1] for ibid in bid], + ) + + # Check that the number of bin indices is the same as the + # number of simulations + if len(bin_ids) != len(pcoords): + raise ValueError( + 'Number of bin indices must match the number of simulations.', + ) return bin_ids - - # return np.digitize(pcoords[:, self.pcoord_idx], self.bins, right=True) From 6ad4f1ea469804bf977d87d2ae5206a5cfeb45bf Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 28 Oct 2025 16:06:23 -0500 Subject: [PATCH 54/73] comments, typing --- deepdrivewe/binners/multirectilinear.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index d587f67..0e101fe 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -46,8 +46,11 @@ def __init__( @property def nbins(self) -> int: """The number of bins.""" + # Calculate the number of bins per dimension nbins_per_dim = np.array([len(dim) - 1 for dim in self.bins]) - return np.prod(nbins_per_dim) + + # Calculate the total number of bins + return int(np.prod(nbins_per_dim)) def get_bin_target_counts(self) -> list[int]: """Get the target counts for each bin. @@ -101,7 +104,7 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: expand_binnumbers=True, ) - # Convert the bin edges to bin indices + # Calculate the bin indices in row-major order bin_ids = np.array( [(ibid[0] - 1) * (len(x_edge) - 1) + ibid[1] for ibid in bid], ) From b667a8c119b1d72f3cf3627a79b24d0051dfa87f Mon Sep 17 00:00:00 2001 From: braceal Date: Tue, 28 Oct 2025 17:27:56 -0500 Subject: [PATCH 55/73] address #35 --- deepdrivewe/binners/base.py | 57 ++++++++++++++++--- deepdrivewe/binners/multirectilinear.py | 41 ++----------- deepdrivewe/binners/rectilinear.py | 47 +++------------ .../examples/amber_nacl_hk/inference.py | 1 + .../examples/amber_ntl9_hk/inference.py | 1 + .../examples/openmm_aae_ddwe/inference.py | 1 + .../examples/openmm_ntl9_ddwe/inference.py | 1 + .../examples/openmm_ntl9_hk/inference.py | 1 + .../examples/synd_ntl9_hk/inference.py | 1 + .../examples/synd_ntl9_lof/inference.py | 1 + 10 files changed, 71 insertions(+), 81 deletions(-) diff --git a/deepdrivewe/binners/base.py b/deepdrivewe/binners/base.py index 1166794..01cc27f 100644 --- a/deepdrivewe/binners/base.py +++ b/deepdrivewe/binners/base.py @@ -17,16 +17,26 @@ class Binner(ABC): """Binner for the progress coordinate.""" - @abstractmethod - def get_bin_target_counts(self) -> list[int]: - """Get the target counts for each bin. + def __init__( + self, + bin_target_counts: int | list[int], + target_state_inds: int | list[int] | None = None, + ) -> None: + """Initialize the binner. - Returns - ------- - list[int] - The target counts for each bin. + Parameters + ---------- + bin_target_counts : int | list[int] + The target counts for each bin. If an integer is provided, + the target counts are assumed to be the same for each bin. + target_state_inds : int | list[int] | None + The index of the target state. If an integer is provided, then + there is only one target state. If a list of integers is provided, + then there are multiple target states. If None is provided, then + there are no target states. Default is None. """ - ... + self.bin_target_counts = bin_target_counts + self.target_state_inds = target_state_inds @property @abstractmethod @@ -39,6 +49,37 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """Assign the simulation pcoords to bins.""" ... + def get_bin_target_counts(self) -> list[int]: + """Get the target counts for each bin. + + Returns + ------- + list[int] + The target counts for each bin. + """ + # Check if the bin target counts is an integer + # If so, then set the target counts for each bin to the same value + # and set the target state bins to 0. Cache the result. + if isinstance(self.bin_target_counts, int): + # Create a list of the bin target counts + bin_target_counts = [self.bin_target_counts] * self.nbins + + # If there are target states, set the target state bins to 0 + if self.target_state_inds is not None: + # Make sure the target state indices are a list + if isinstance(self.target_state_inds, int): + self.target_state_inds = [self.target_state_inds] + + # Set the target state bins to 0 since they are recycled + for i in self.target_state_inds: + bin_target_counts[i] = 0 + + # Cache the result + self.bin_target_counts = bin_target_counts + + # Otherwise, return the list of bin target counts + return self.bin_target_counts + @property def labels(self) -> list[str]: """The bin labels for WESTPA.""" diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 0e101fe..9a4f944 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -15,7 +15,7 @@ def __init__( self, bins: list[np.ndarray | list[float]], bin_target_counts: int | list[int], - target_state_inds: int | list[int] = 0, + target_state_inds: int | list[int] | None = None, ) -> None: """Initialize the binner. @@ -26,15 +26,15 @@ def __init__( bin_target_counts : int | list[int] The target counts for each bin. If an integer is provided, the target counts are assumed to be the same for each bin. - target_state_inds : int | list[int] + target_state_inds : int | list[int] | None The index of the target state. If an integer is provided, then there is only one target state. If a list of integers is provided, - then there are multiple target states. Default is 0 which - corresponds to the first bin. + then there are multiple target states. If None is provided, then + there are no target states. Default is None. """ + super().__init__(bin_target_counts, target_state_inds) + self.bins = bins - self.bin_target_counts = bin_target_counts - self.target_state_inds = target_state_inds # Check that the bins are sorted for binbounds in self.bins: @@ -52,35 +52,6 @@ def nbins(self) -> int: # Calculate the total number of bins return int(np.prod(nbins_per_dim)) - def get_bin_target_counts(self) -> list[int]: - """Get the target counts for each bin. - - Returns - ------- - list[int] - The target counts for each bin. - """ - # Check if the bin target counts is an integer - # If so, then set the target counts for each bin to the same value - # and set the target state bins to 0. Cache the result. - if isinstance(self.bin_target_counts, int): - # Create a list of the bin target counts - bin_target_counts = [self.bin_target_counts] * self.nbins - - # Get the target state indices (convert to a list if an integer) - if isinstance(self.target_state_inds, int): - self.target_state_inds = [self.target_state_inds] - - # Set each of the target state bins to 0 since they are recycled - for i in self.target_state_inds: - bin_target_counts[i] = 0 - - # Cache the result - self.bin_target_counts = bin_target_counts - - # Otherwise, return the list of bin target counts - return self.bin_target_counts - def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """Bin the progress coordinate. diff --git a/deepdrivewe/binners/rectilinear.py b/deepdrivewe/binners/rectilinear.py index 7244036..0521335 100644 --- a/deepdrivewe/binners/rectilinear.py +++ b/deepdrivewe/binners/rectilinear.py @@ -14,7 +14,7 @@ def __init__( self, bins: list[float], bin_target_counts: int | list[int], - target_state_inds: int | list[int] = 0, + target_state_inds: int | list[int] | None = None, pcoord_idx: int = 0, ) -> None: """Initialize the binner. @@ -23,21 +23,21 @@ def __init__( ---------- bins : list[float] The bin edges for the progress coordinate. - pcoord_idx : int - The index of the progress coordinate to use for binning. - Default is 0. bin_target_counts : int | list[int] The target counts for each bin. If an integer is provided, the target counts are assumed to be the same for each bin. - target_state_inds : int | list[int] + target_state_inds : int | list[int] | None The index of the target state. If an integer is provided, then there is only one target state. If a list of integers is provided, - then there are multiple target states. Default is 0 which - corresponds to the first bin. + then there are multiple target states. If None is provided, then + there are no target states. Default is None. + pcoord_idx : int + The index of the progress coordinate to use for binning. + Default is 0. """ + super().__init__(bin_target_counts, target_state_inds) + self.bins = bins - self.bin_target_counts = bin_target_counts - self.target_state_inds = target_state_inds self.pcoord_idx = pcoord_idx # Check that the bins are sorted @@ -49,35 +49,6 @@ def nbins(self) -> int: """The number of bins.""" return len(self.bins) - 1 - def get_bin_target_counts(self) -> list[int]: - """Get the target counts for each bin. - - Returns - ------- - list[int] - The target counts for each bin. - """ - # Check if the bin target counts is an integer - # If so, then set the target counts for each bin to the same value - # and set the target state bins to 0. Cache the result. - if isinstance(self.bin_target_counts, int): - # Create a list of the bin target counts - bin_target_counts = [self.bin_target_counts] * self.nbins - - # Get the target state indices (convert to a list if an integer) - if isinstance(self.target_state_inds, int): - self.target_state_inds = [self.target_state_inds] - - # Set each of the target state bins to 0 since they are recycled - for i in self.target_state_inds: - bin_target_counts[i] = 0 - - # Cache the result - self.bin_target_counts = bin_target_counts - - # Otherwise, return the list of bin target counts - return self.bin_target_counts - def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """Bin the progress coordinate. diff --git a/deepdrivewe/examples/amber_nacl_hk/inference.py b/deepdrivewe/examples/amber_nacl_hk/inference.py index f143f9b..d61fc2a 100644 --- a/deepdrivewe/examples/amber_nacl_hk/inference.py +++ b/deepdrivewe/examples/amber_nacl_hk/inference.py @@ -80,6 +80,7 @@ def run_inference( float('inf'), ], bin_target_counts=config.sims_per_bin, + target_state_inds=0, # The first bin is the target (folded) state ) # Define the recycling policy diff --git a/deepdrivewe/examples/amber_ntl9_hk/inference.py b/deepdrivewe/examples/amber_ntl9_hk/inference.py index 7da45c8..6700739 100644 --- a/deepdrivewe/examples/amber_ntl9_hk/inference.py +++ b/deepdrivewe/examples/amber_ntl9_hk/inference.py @@ -60,6 +60,7 @@ def run_inference( + [6.60 + 0.6 * i for i in range(6)] + [float('inf')], bin_target_counts=config.sims_per_bin, + target_state_inds=0, # The first bin is the target (folded) state ) # Define the recycling policy diff --git a/deepdrivewe/examples/openmm_aae_ddwe/inference.py b/deepdrivewe/examples/openmm_aae_ddwe/inference.py index 289375d..5521458 100644 --- a/deepdrivewe/examples/openmm_aae_ddwe/inference.py +++ b/deepdrivewe/examples/openmm_aae_ddwe/inference.py @@ -133,6 +133,7 @@ def run_inference( binner = RectilinearBinner( bins=[0.0, 1.0, float('inf')], bin_target_counts=config.sims_per_bin, + target_state_inds=0, # The first bin is the target (folded) state ) # Define the recycling policy diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py b/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py index 9a28a95..d57faf7 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/inference.py @@ -136,6 +136,7 @@ def run_inference( binner = RectilinearBinner( bins=[0.0, 1.0, float('inf')], bin_target_counts=config.sims_per_bin, + target_state_inds=0, # The first bin is the target (folded) state ) # Define the recycling policy diff --git a/deepdrivewe/examples/openmm_ntl9_hk/inference.py b/deepdrivewe/examples/openmm_ntl9_hk/inference.py index d9ae2c2..866c129 100644 --- a/deepdrivewe/examples/openmm_ntl9_hk/inference.py +++ b/deepdrivewe/examples/openmm_ntl9_hk/inference.py @@ -60,6 +60,7 @@ def run_inference( + [6.60 + 0.6 * i for i in range(6)] + [float('inf')], bin_target_counts=config.sims_per_bin, + target_state_inds=0, # The first bin is the target (folded) state ) # Define the recycling policy diff --git a/deepdrivewe/examples/synd_ntl9_hk/inference.py b/deepdrivewe/examples/synd_ntl9_hk/inference.py index 7da45c8..6700739 100644 --- a/deepdrivewe/examples/synd_ntl9_hk/inference.py +++ b/deepdrivewe/examples/synd_ntl9_hk/inference.py @@ -60,6 +60,7 @@ def run_inference( + [6.60 + 0.6 * i for i in range(6)] + [float('inf')], bin_target_counts=config.sims_per_bin, + target_state_inds=0, # The first bin is the target (folded) state ) # Define the recycling policy diff --git a/deepdrivewe/examples/synd_ntl9_lof/inference.py b/deepdrivewe/examples/synd_ntl9_lof/inference.py index 752076c..d692399 100644 --- a/deepdrivewe/examples/synd_ntl9_lof/inference.py +++ b/deepdrivewe/examples/synd_ntl9_lof/inference.py @@ -142,6 +142,7 @@ def run_inference( binner = RectilinearBinner( bins=[0.0, 1.0, float('inf')], bin_target_counts=config.sims_per_bin, + target_state_inds=0, # The first bin is the target (folded) state ) # Define the recycling policy From 9d91a0a20c942c3089bc7b5ebe38514ccc02a431 Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Wed, 29 Oct 2025 14:04:27 -0500 Subject: [PATCH 56/73] binners fix + tests --- deepdrivewe/binners/multirectilinear.py | 13 ++++++-- deepdrivewe/binners/rectilinear.py | 9 +++++- tests/test_binner.py | 41 +++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 3 deletions(-) create mode 100644 tests/test_binner.py diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 0e101fe..829d9df 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -4,6 +4,7 @@ import numpy as np from scipy.stats import binned_statistic_2d +import warnings from deepdrivewe.binners.base import Binner @@ -96,7 +97,7 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """ # Bin the progress coordinates (make sure the target state # boundary is included in the target state bin). - _, x_edge, _, bid = binned_statistic_2d( + _, x_edge, y_edge, bid = binned_statistic_2d( *pcoords.T, values=None, statistic='count', @@ -104,9 +105,17 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: expand_binnumbers=True, ) + # Clip the bin indices so any index outside of defined bins are moved to nearest defined bin + for idx, ibid in enumerate(bid): + if not np.all(ibid>= 0) or not np.all(ibid < len(self.bins[idx])): + warnings.warn("Simulations with progress coordinates outside the bin boundaries definitions are automatically placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") + + bid[0] = np.clip(bid[0], 1, len(x_edge)-1) + bid[1] = np.clip(bid[1], 1, len(y_edge)-1) + # Calculate the bin indices in row-major order bin_ids = np.array( - [(ibid[0] - 1) * (len(x_edge) - 1) + ibid[1] for ibid in bid], + [(ibid[0] - 1) * (len(x_edge) - 1) + (ibid[1] - 1) for ibid in bid.T], ) # Check that the number of bin indices is the same as the diff --git a/deepdrivewe/binners/rectilinear.py b/deepdrivewe/binners/rectilinear.py index 7244036..41c16ec 100644 --- a/deepdrivewe/binners/rectilinear.py +++ b/deepdrivewe/binners/rectilinear.py @@ -93,4 +93,11 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """ # Bin the progress coordinates (make sure the target state # boundary is included in the target state bin). - return np.digitize(pcoords[:, self.pcoord_idx], self.bins, right=True) + bin_id = np.digitize(pcoords[:, self.pcoord_idx], self.bins) -1 + + if not np.all(bin_id >= 0) or not np.all(bin_id < len(self.bins)): + warnings.warn("Simulations with progress coordinates outside the bin boundaries definitions are placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") + + # This ensures our bin index is >=0 and < len(self.bins) + return np.clip(bin_id, 0, len(self.bins)-1) + diff --git a/tests/test_binner.py b/tests/test_binner.py new file mode 100644 index 0000000..7ba70dd --- /dev/null +++ b/tests/test_binner.py @@ -0,0 +1,41 @@ +import os +import pytest + +import numpy as np + +from deepdrivewe.binners import RectilinearBinner, MultiRectilinearBinner + +class TestRectilinearBinner: + def test1dAssign(self): + bounds = [0.0, 1.0, 2.0, 3.0] + coords = np.array([0, 0.5, 1.5, 1.6, 2.0, 2.0, 2.9])[:, None] + + assigner = RectilinearBinner(bins=bounds, bin_target_counts = 3, target_state_inds=[None], pcoord_idx=0) + assert (assigner.assign_bins(coords) == [0, 0, 1, 1, 2, 2, 2]).all() + + def test2dAssign(self): + boundaries = [(-1, -0.5, 0, 0.5, 1), (-1, -0.5, 0, 0.5, 1)] + coords = np.array([(-0.75, -0.75), (-0.25, -0.25), (0, 0), (0.25, 0.25), (0.75, 0.75), (-0.25, 0.75), (0.25, -0.75)]) + assigner = MultiRectilinearBinner(boundaries, bin_target_counts=3, target_state_inds=[None]) + + """bin structure: [(a,b), (c,d)] => x in [a,b), y in [c, d) + 0:[(-1, -0.5), (-1, -0.5)] + 1:[(-1, -0.5), (-0.5, 0)] + 2:[(-1, -0.5), (0, 0.5)] + 3:[(-1, -0.5), (0.5, 1)] + 4:[(-0.5, 0), (-1, -0.5)] + 5:[(-0.5, 0), (-0.5, 0)] + 6:[(-0.5, 0), (0, 0.5)] + 7:[(-0.5, 0), (0.5, 1)] + 8:[(0, 0.5), (-1, -0.5)] + 9:[(0, 0.5), (-0.5, 0)] + 10:[(0, 0.5), (0, 0.5)] + 11:[(0, 0.5), (0.5, 1)] + 12:[(0.5, 1), (-1, -0.5)] + 13:[(0.5, 1), (-0.5, 0)] + 14:[(0.5, 1), (0, 0.5)] + 15:[(0.5, 1), (0.5, 1)]""" + + assert (assigner.assign_bins(coords) == [0, 5, 10, 10, 15, 7, 8]).all() + + From d40a9afc096600baed1073b0645b0cb7ba9d5e69 Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 29 Oct 2025 14:12:01 -0500 Subject: [PATCH 57/73] Add path validation to basis states #37 --- deepdrivewe/api.py | 38 +++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/deepdrivewe/api.py b/deepdrivewe/api.py index 13daa4f..e50fc7f 100644 --- a/deepdrivewe/api.py +++ b/deepdrivewe/api.py @@ -18,6 +18,7 @@ import yaml # type: ignore[import-untyped] from pydantic import BaseModel as _BaseModel from pydantic import Field +from pydantic import field_validator T = TypeVar('T') @@ -284,6 +285,16 @@ class BasisStates(BaseModel): description='The basis states for the weighted ensemble.', ) + @field_validator('basis_state_dir') + @classmethod + def validate_basis_state_dir(cls, value: Path) -> Path: + """Validate and resolve the basis state directory.""" + if not value.is_dir(): + raise NotADirectoryError( + f'The basis state directory {value} is not a directory.', + ) + return value.resolve() + @property def unique_basis_states(self) -> list[SimMetadata]: """Return the unique basis states.""" @@ -306,10 +317,35 @@ def load_basis_states( self, basis_state_initializer: BasisStateInitializer, ) -> None: - """Load the basis states for the weighted ensemble.""" + """Load the basis states for the weighted ensemble. + + Parameters + ---------- + basis_state_initializer : BasisStateInitializer + The initializer for the basis states (e.g., a function that + reads the progress coordinate from a file and computes and + returns a progress coordinate). + + Raises + ------ + NotADirectoryError + If the basis state directory is not a directory. + FileNotFoundError + If no basis state files are found in the input directory. + ValueError + If no basis states are found in the basis_state_dir. + """ # Collect the basis state files basis_files = self._glob_basis_states() + # Raise an error if there are no basis states + if not basis_files: + raise ValueError( + 'No basis states found in the basis state directory. ' + 'Please check that the basis_state_dir exists and contains ' + 'the correct files with the correct extension.', + ) + # Compute the pcoord for each basis state basis_pcoords = [ basis_state_initializer(basis_file.as_posix()) From 8cb13a55f263deadf6c2f684d62b69f6351002cd Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Wed, 29 Oct 2025 14:42:03 -0500 Subject: [PATCH 58/73] generalize multirectilinear to multidimensional, fix warning --- deepdrivewe/binners/multirectilinear.py | 24 +++++++++++++----------- deepdrivewe/binners/rectilinear.py | 3 ++- tests/test_binner.py | 12 ++++++++---- 3 files changed, 23 insertions(+), 16 deletions(-) diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 829d9df..5503f26 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -3,7 +3,7 @@ from __future__ import annotations import numpy as np -from scipy.stats import binned_statistic_2d +from scipy.stats import binned_statistic_dd import warnings from deepdrivewe.binners.base import Binner @@ -97,8 +97,8 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """ # Bin the progress coordinates (make sure the target state # boundary is included in the target state bin). - _, x_edge, y_edge, bid = binned_statistic_2d( - *pcoords.T, + _, bin_edges, bid = binned_statistic_dd( + np.asarray(pcoords), values=None, statistic='count', bins=self.bins, @@ -106,17 +106,19 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: ) # Clip the bin indices so any index outside of defined bins are moved to nearest defined bin + nbins_per_dim = [len(edges)-1 for edges in bin_edges] + for idx, ibid in enumerate(bid): - if not np.all(ibid>= 0) or not np.all(ibid < len(self.bins[idx])): - warnings.warn("Simulations with progress coordinates outside the bin boundaries definitions are automatically placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") - - bid[0] = np.clip(bid[0], 1, len(x_edge)-1) - bid[1] = np.clip(bid[1], 1, len(y_edge)-1) + if not np.all(ibid> 0) or not np.all(ibid < len(self.bins[idx])): + warnings.warn(f"Simulations with progress coordinates outside the bin boundaries definition of dimension {idx} are automatically placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") + bid[idx] = np.clip(ibid, 1, nbins_per_dim[idx]) # Calculate the bin indices in row-major order - bin_ids = np.array( - [(ibid[0] - 1) * (len(x_edge) - 1) + (ibid[1] - 1) for ibid in bid.T], - ) + bin_ids = np.zeros(len(pcoords), dtype=int) + for idx, ibid in enumerate(bid.T): + for idim in range(len(nbins_per_dim)-1): + bin_ids[idx] += (ibid[idim] -1) * nbins_per_dim[idim] + bin_ids[idx] += (ibid[-1] - 1) # Check that the number of bin indices is the same as the # number of simulations diff --git a/deepdrivewe/binners/rectilinear.py b/deepdrivewe/binners/rectilinear.py index 41c16ec..a9f0ed6 100644 --- a/deepdrivewe/binners/rectilinear.py +++ b/deepdrivewe/binners/rectilinear.py @@ -3,6 +3,7 @@ from __future__ import annotations import numpy as np +import warnings from deepdrivewe.binners.base import Binner @@ -95,7 +96,7 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: # boundary is included in the target state bin). bin_id = np.digitize(pcoords[:, self.pcoord_idx], self.bins) -1 - if not np.all(bin_id >= 0) or not np.all(bin_id < len(self.bins)): + if not np.all(bin_id > 0) or not np.all(bin_id < len(self.bins)): warnings.warn("Simulations with progress coordinates outside the bin boundaries definitions are placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") # This ensures our bin index is >=0 and < len(self.bins) diff --git a/tests/test_binner.py b/tests/test_binner.py index 7ba70dd..15364bc 100644 --- a/tests/test_binner.py +++ b/tests/test_binner.py @@ -8,14 +8,17 @@ class TestRectilinearBinner: def test1dAssign(self): bounds = [0.0, 1.0, 2.0, 3.0] - coords = np.array([0, 0.5, 1.5, 1.6, 2.0, 2.0, 2.9])[:, None] + coords = np.array([-1, 0, 0.5, 1.5, 1.6, 2.0, 2.0, 2.9])[:, None] assigner = RectilinearBinner(bins=bounds, bin_target_counts = 3, target_state_inds=[None], pcoord_idx=0) - assert (assigner.assign_bins(coords) == [0, 0, 1, 1, 2, 2, 2]).all() + + with pytest.warns(UserWarning): + assert (assigner.assign_bins(coords) == [0, 0, 0, 1, 1, 2, 2, 2]).all() def test2dAssign(self): boundaries = [(-1, -0.5, 0, 0.5, 1), (-1, -0.5, 0, 0.5, 1)] - coords = np.array([(-0.75, -0.75), (-0.25, -0.25), (0, 0), (0.25, 0.25), (0.75, 0.75), (-0.25, 0.75), (0.25, -0.75)]) + coords = np.array([(-2, -2), (-0.75, -0.75), (-0.25, -0.25), (0, 0), (0.25, 0.25), (0.75, 0.75), (-0.25, 0.75), (0.25, -0.75)]) + assigner = MultiRectilinearBinner(boundaries, bin_target_counts=3, target_state_inds=[None]) """bin structure: [(a,b), (c,d)] => x in [a,b), y in [c, d) @@ -36,6 +39,7 @@ def test2dAssign(self): 14:[(0.5, 1), (0, 0.5)] 15:[(0.5, 1), (0.5, 1)]""" - assert (assigner.assign_bins(coords) == [0, 5, 10, 10, 15, 7, 8]).all() + with pytest.warns(UserWarning): + assert (assigner.assign_bins(coords) == [0, 0, 5, 10, 10, 15, 7, 8]).all() From 38074f31d72eb9882acf8693a2162664c11cc090 Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 29 Oct 2025 14:45:03 -0500 Subject: [PATCH 59/73] add path validation to everything --- deepdrivewe/__init__.py | 1 + deepdrivewe/api.py | 30 +++++++++++++++++++ deepdrivewe/examples/amber_nacl_hk/main.py | 7 +++++ .../examples/amber_nacl_hk/simulate.py | 8 +++++ deepdrivewe/examples/amber_ntl9_hk/main.py | 7 +++++ .../examples/amber_ntl9_hk/simulate.py | 8 +++++ deepdrivewe/examples/openmm_aae_ddwe/main.py | 7 +++++ .../examples/openmm_aae_ddwe/simulate.py | 8 +++++ deepdrivewe/examples/openmm_aae_ddwe/train.py | 8 +++++ deepdrivewe/examples/openmm_ntl9_ddwe/main.py | 10 +++++++ .../examples/openmm_ntl9_ddwe/simulate.py | 8 +++++ .../examples/openmm_ntl9_ddwe/train.py | 8 +++++ deepdrivewe/examples/openmm_ntl9_hk/main.py | 7 +++++ .../examples/openmm_ntl9_hk/simulate.py | 8 +++++ .../examples/synd_ntl9_lof/simulate.py | 14 +++++++++ deepdrivewe/simulation/amber.py | 9 ++++++ 16 files changed, 148 insertions(+) diff --git a/deepdrivewe/__init__.py b/deepdrivewe/__init__.py index ff6c4fc..b99c677 100644 --- a/deepdrivewe/__init__.py +++ b/deepdrivewe/__init__.py @@ -14,6 +14,7 @@ from deepdrivewe.api import SimResult from deepdrivewe.api import TargetState from deepdrivewe.api import TrainResult +from deepdrivewe.api import validate_and_resolve_file from deepdrivewe.api import WeightedEnsemble from deepdrivewe.binners import Binner from deepdrivewe.checkpoint import EnsembleCheckpointer diff --git a/deepdrivewe/api.py b/deepdrivewe/api.py index e50fc7f..d16efe1 100644 --- a/deepdrivewe/api.py +++ b/deepdrivewe/api.py @@ -23,6 +23,36 @@ T = TypeVar('T') +def validate_and_resolve_file(value: Path | None) -> Path | None: + """Validate and resolve a file path. + + Parameters + ---------- + value : Path | None + The file path to validate and resolve. + + Returns + ------- + Path | None + The validated and resolved file path. + + Raises + ------ + FileNotFoundError + If the file path is not a file. + """ + # Return None if the file path is None + if value is None: + return None + + # Raise an error if the file path is not a file + if not value.is_file(): + raise FileNotFoundError(f'The file {value} is not a file.') + + # Resolve the file path + return value.resolve() + + class BaseModel(_BaseModel): """Provide an easy interface to read/write YAML files.""" diff --git a/deepdrivewe/examples/amber_nacl_hk/main.py b/deepdrivewe/examples/amber_nacl_hk/main.py index 408308f..cac8998 100644 --- a/deepdrivewe/examples/amber_nacl_hk/main.py +++ b/deepdrivewe/examples/amber_nacl_hk/main.py @@ -24,6 +24,7 @@ from deepdrivewe import BasisStates from deepdrivewe import EnsembleCheckpointer from deepdrivewe import TargetState +from deepdrivewe import validate_and_resolve_file from deepdrivewe import WeightedEnsemble from deepdrivewe.examples.amber_nacl_hk.inference import InferenceConfig from deepdrivewe.examples.amber_nacl_hk.inference import run_inference @@ -44,6 +45,12 @@ class CustomBasisStateInitializer(BaseModel): description='Reference file for the cpptraj command.', ) + @field_validator('top_file', 'reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def __call__(self, basis_file: str) -> list[float]: """Initialize the basis state parent coordinates.""" # Create the cpptraj command file diff --git a/deepdrivewe/examples/amber_nacl_hk/simulate.py b/deepdrivewe/examples/amber_nacl_hk/simulate.py index 3b80ba2..16dea51 100644 --- a/deepdrivewe/examples/amber_nacl_hk/simulate.py +++ b/deepdrivewe/examples/amber_nacl_hk/simulate.py @@ -8,10 +8,12 @@ import numpy as np from pydantic import Field +from pydantic import field_validator from deepdrivewe import BaseModel from deepdrivewe import SimMetadata from deepdrivewe import SimResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.simulation.amber import AmberConfig from deepdrivewe.simulation.amber import AmberSimulation from deepdrivewe.simulation.amber import AmberTrajAnalyzer @@ -28,6 +30,12 @@ class SimulationConfig(BaseModel): description='The reference PDB file for the cpptraj analysis.', ) + @field_validator('reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + class DistanceAnalyzer(AmberTrajAnalyzer): """Analyze Amber simulations using cpptraj.""" diff --git a/deepdrivewe/examples/amber_ntl9_hk/main.py b/deepdrivewe/examples/amber_ntl9_hk/main.py index 3fe28bb..7b15d43 100644 --- a/deepdrivewe/examples/amber_ntl9_hk/main.py +++ b/deepdrivewe/examples/amber_ntl9_hk/main.py @@ -24,6 +24,7 @@ from deepdrivewe import BasisStates from deepdrivewe import EnsembleCheckpointer from deepdrivewe import TargetState +from deepdrivewe import validate_and_resolve_file from deepdrivewe import WeightedEnsemble from deepdrivewe.examples.amber_ntl9_hk.inference import InferenceConfig from deepdrivewe.examples.amber_ntl9_hk.inference import run_inference @@ -44,6 +45,12 @@ class CustomBasisStateInitializer(BaseModel): description='Reference file for the cpptraj command.', ) + @field_validator('top_file', 'reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def __call__(self, basis_file: str) -> list[float]: """Initialize the basis state parent coordinates.""" # Create the cpptraj command file diff --git a/deepdrivewe/examples/amber_ntl9_hk/simulate.py b/deepdrivewe/examples/amber_ntl9_hk/simulate.py index 3627299..baf9007 100644 --- a/deepdrivewe/examples/amber_ntl9_hk/simulate.py +++ b/deepdrivewe/examples/amber_ntl9_hk/simulate.py @@ -8,10 +8,12 @@ import numpy as np from pydantic import Field +from pydantic import field_validator from deepdrivewe import BaseModel from deepdrivewe import SimMetadata from deepdrivewe import SimResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.simulation.amber import AmberConfig from deepdrivewe.simulation.amber import AmberSimulation from deepdrivewe.simulation.amber import AmberTrajAnalyzer @@ -28,6 +30,12 @@ class SimulationConfig(BaseModel): description='The reference PDB file for the cpptraj analysis.', ) + @field_validator('reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + class BackboneRMSDAnalyzer(AmberTrajAnalyzer): """Analyze Amber simulations using cpptraj.""" diff --git a/deepdrivewe/examples/openmm_aae_ddwe/main.py b/deepdrivewe/examples/openmm_aae_ddwe/main.py index ca3afda..a92d582 100644 --- a/deepdrivewe/examples/openmm_aae_ddwe/main.py +++ b/deepdrivewe/examples/openmm_aae_ddwe/main.py @@ -26,6 +26,7 @@ from deepdrivewe import BasisStates from deepdrivewe import EnsembleCheckpointer from deepdrivewe import TargetState +from deepdrivewe import validate_and_resolve_file from deepdrivewe import WeightedEnsemble from deepdrivewe.examples.openmm_aae_ddwe.inference import InferenceConfig from deepdrivewe.examples.openmm_aae_ddwe.inference import run_inference @@ -48,6 +49,12 @@ class RMSDBasisStateInitializer(BaseModel): description='The MDAnalysis selection string for the atoms to use.', ) + @field_validator('reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def __call__(self, basis_file: str) -> list[float]: """Initialize the basis state parent coordinates.""" # Load the basis file diff --git a/deepdrivewe/examples/openmm_aae_ddwe/simulate.py b/deepdrivewe/examples/openmm_aae_ddwe/simulate.py index 9451d66..6d2ef08 100644 --- a/deepdrivewe/examples/openmm_aae_ddwe/simulate.py +++ b/deepdrivewe/examples/openmm_aae_ddwe/simulate.py @@ -6,10 +6,12 @@ from typing import Sequence from pydantic import Field +from pydantic import field_validator from deepdrivewe import BaseModel from deepdrivewe import SimMetadata from deepdrivewe import SimResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.simulation.openmm import CollectionReporter from deepdrivewe.simulation.openmm import CoordinatesCollector from deepdrivewe.simulation.openmm import OpenMMConfig @@ -43,6 +45,12 @@ class SimulationConfig(BaseModel): description='The OpenMM selection strings for the atoms to use.', ) + @field_validator('top_file', 'reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def run_simulation( metadata: SimMetadata, diff --git a/deepdrivewe/examples/openmm_aae_ddwe/train.py b/deepdrivewe/examples/openmm_aae_ddwe/train.py index 2157978..93633ee 100644 --- a/deepdrivewe/examples/openmm_aae_ddwe/train.py +++ b/deepdrivewe/examples/openmm_aae_ddwe/train.py @@ -7,9 +7,11 @@ import numpy as np from pydantic import BaseModel from pydantic import Field +from pydantic import field_validator from deepdrivewe import SimResult from deepdrivewe import TrainResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.ai import AdversarialAE from deepdrivewe.ai import AdversarialAEConfig @@ -26,6 +28,12 @@ class TrainConfig(BaseModel): 'Train from scratch by default.', ) + @field_validator('config_path', 'checkpoint_path') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + # TODO: We probably need to store a history of old training data # to retrain the model. Add a config argument to include a cMD run dataset. diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py index 53f41df..dc10437 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py @@ -50,6 +50,16 @@ class RMSDBasisStateInitializer(BaseModel): description='The MDAnalysis selection string for the atoms to use.', ) + @field_validator('reference_file') + @classmethod + def validate_reference_file(cls, value: Path) -> Path: + """Validate and resolve the reference file.""" + if not value.is_file(): + raise FileNotFoundError( + f'The reference file {value} is not a file.', + ) + return value.resolve() + def __call__(self, basis_file: str) -> list[float]: """Initialize the basis state parent coordinates.""" # Load the basis file diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py index cd2230c..ec73137 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py @@ -6,10 +6,12 @@ from typing import Sequence from pydantic import Field +from pydantic import field_validator from deepdrivewe import BaseModel from deepdrivewe import SimMetadata from deepdrivewe import SimResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.simulation.openmm import CollectionReporter from deepdrivewe.simulation.openmm import ContactMapCollector from deepdrivewe.simulation.openmm import OpenMMConfig @@ -44,6 +46,12 @@ class SimulationConfig(BaseModel): description='The OpenMM selection strings for the atoms to use.', ) + @field_validator('top_file', 'reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def run_simulation( metadata: SimMetadata, diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py index 1bbd031..6e8abf5 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py @@ -8,9 +8,11 @@ import numpy as np from pydantic import BaseModel from pydantic import Field +from pydantic import field_validator from deepdrivewe import SimResult from deepdrivewe import TrainResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.ai import ConvolutionalVAE from deepdrivewe.ai import ConvolutionalVAEConfig from deepdrivewe.workflows.stream import ProxyStreamConfig @@ -40,6 +42,12 @@ class TrainConfig(BaseModel): 're-initializing and re-training the model.', ) + @field_validator('config_path', 'checkpoint_path') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + # TODO: We probably need to store a history of old training data # to retrain the model. Add a config argument to include a cMD run dataset. diff --git a/deepdrivewe/examples/openmm_ntl9_hk/main.py b/deepdrivewe/examples/openmm_ntl9_hk/main.py index 76d02dd..0efbeb3 100644 --- a/deepdrivewe/examples/openmm_ntl9_hk/main.py +++ b/deepdrivewe/examples/openmm_ntl9_hk/main.py @@ -26,6 +26,7 @@ from deepdrivewe import BasisStates from deepdrivewe import EnsembleCheckpointer from deepdrivewe import TargetState +from deepdrivewe import validate_and_resolve_file from deepdrivewe import WeightedEnsemble from deepdrivewe.examples.openmm_ntl9_hk.inference import InferenceConfig from deepdrivewe.examples.openmm_ntl9_hk.inference import run_inference @@ -46,6 +47,12 @@ class RMSDBasisStateInitializer(BaseModel): description='The MDAnalysis selection string for the atoms to use.', ) + @field_validator('reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def __call__(self, basis_file: str) -> list[float]: """Initialize the basis state parent coordinates.""" # Load the basis file diff --git a/deepdrivewe/examples/openmm_ntl9_hk/simulate.py b/deepdrivewe/examples/openmm_ntl9_hk/simulate.py index a62b9f1..0017e88 100644 --- a/deepdrivewe/examples/openmm_ntl9_hk/simulate.py +++ b/deepdrivewe/examples/openmm_ntl9_hk/simulate.py @@ -6,10 +6,12 @@ from typing import Sequence from pydantic import Field +from pydantic import field_validator from deepdrivewe import BaseModel from deepdrivewe import SimMetadata from deepdrivewe import SimResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.simulation.openmm import ContactMapRMSDReporter from deepdrivewe.simulation.openmm import OpenMMConfig from deepdrivewe.simulation.openmm import OpenMMSimulation @@ -41,6 +43,12 @@ class SimulationConfig(BaseModel): description='The OpenMM selection strings for the atoms to use.', ) + @field_validator('top_file', 'reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def run_simulation( metadata: SimMetadata, diff --git a/deepdrivewe/examples/synd_ntl9_lof/simulate.py b/deepdrivewe/examples/synd_ntl9_lof/simulate.py index fe3f01e..697415a 100644 --- a/deepdrivewe/examples/synd_ntl9_lof/simulate.py +++ b/deepdrivewe/examples/synd_ntl9_lof/simulate.py @@ -10,11 +10,13 @@ import numpy as np from pydantic import BaseModel from pydantic import Field +from pydantic import field_validator from scipy.sparse import coo_matrix from scipy.spatial import distance_matrix from deepdrivewe import SimMetadata from deepdrivewe import SimResult +from deepdrivewe import validate_and_resolve_file from deepdrivewe.simulation.synd import SynDConfig from deepdrivewe.simulation.synd import SynDSimulation from deepdrivewe.simulation.synd import SynDTrajAnalyzer @@ -35,6 +37,12 @@ class SimulationConfig(SynDConfig): description='The mdtraj selection string for the atoms to use.', ) + @field_validator('reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + class ContactMapAnalyzer(BaseModel, SynDTrajAnalyzer): """Analyze SynD simulations using contact maps.""" @@ -55,6 +63,12 @@ class ContactMapAnalyzer(BaseModel, SynDTrajAnalyzer): description='Whether to convert the coordinates to angstroms.', ) + @field_validator('reference_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def get_contact_maps(self, sim: SynDSimulation) -> np.ndarray: """Compute contact maps from the trajectory. diff --git a/deepdrivewe/simulation/amber.py b/deepdrivewe/simulation/amber.py index 0c6cdc5..ab8aaaa 100644 --- a/deepdrivewe/simulation/amber.py +++ b/deepdrivewe/simulation/amber.py @@ -13,6 +13,9 @@ import numpy as np from pydantic import BaseModel from pydantic import Field +from pydantic import field_validator + +from deepdrivewe import validate_and_resolve_file class AmberConfig(BaseModel): @@ -29,6 +32,12 @@ class AmberConfig(BaseModel): description='The prmtop file for the Amber simulation.', ) + @field_validator('input_file', 'top_file') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + class AmberSimulation(BaseModel): """Run an Amber simulation.""" From da05532a826f8d782156082cc9fb60b4f2ea895c Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 29 Oct 2025 14:46:47 -0500 Subject: [PATCH 60/73] remove absolute paths --- examples/openmm_ntl9_ddwe/config.yaml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/openmm_ntl9_ddwe/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index b2d6cb6..7dff556 100644 --- a/examples/openmm_ntl9_ddwe/config.yaml +++ b/examples/openmm_ntl9_ddwe/config.yaml @@ -9,7 +9,7 @@ num_iterations: 106 # The basis states to use for the ensemble basis_states: # The directory containing the basis states sub directories - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/inputs + basis_state_dir: examples/openmm_ntl9_ddwe/inputs # The file extension for the basis state files basis_state_ext: .pdb # The number of basis states to use @@ -21,7 +21,7 @@ basis_states: # Strategy for initializing the basis state progress coordinates basis_state_initializer: # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/common_files/reference.pdb + reference_file: examples/openmm_ntl9_ddwe/common_files/reference.pdb # The configuration for the simulation simulation_config: @@ -41,14 +41,14 @@ simulation_config: hardware_platform: CUDA # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/common_files/reference.pdb + reference_file: examples/openmm_ntl9_ddwe/common_files/reference.pdb # The configuration for training train_config: # The path to the CVAE model configuration file - config_path: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/cvae-config.yaml + config_path: examples/openmm_ntl9_ddwe/cvae-config.yaml # The path to the CVAE model weights file - checkpoint_path: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe/checkpoint-epoch-100.pt + checkpoint_path: examples/openmm_ntl9_ddwe/checkpoint-epoch-100.pt # The configuration for the inference inference_config: From c42987e75ad3f372321a75b0d17250f4096ec562 Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 29 Oct 2025 14:47:05 -0500 Subject: [PATCH 61/73] remove absolute paths --- examples/openmm_ntl9_ddwe/config.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/openmm_ntl9_ddwe/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index 7dff556..f21b2f4 100644 --- a/examples/openmm_ntl9_ddwe/config.yaml +++ b/examples/openmm_ntl9_ddwe/config.yaml @@ -1,4 +1,5 @@ # Configuration file for the NTL9 folding example using OpenMM +# The inputs and common files are relative to the root of the deepdrivewe repository. # The output directory for the runs output_dir: runs/ntl9-v1 From 5bd614ec944fc645c47e2ea3a7be43e3844d0519 Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Wed, 29 Oct 2025 14:47:37 -0500 Subject: [PATCH 62/73] lint --- deepdrivewe/binners/multirectilinear.py | 4 ++-- deepdrivewe/binners/rectilinear.py | 1 - tests/{test_binner.py => binner_test.py} | 4 +--- 3 files changed, 3 insertions(+), 6 deletions(-) rename tests/{test_binner.py => binner_test.py} (99%) diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 5503f26..a473f47 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -105,9 +105,9 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: expand_binnumbers=True, ) - # Clip the bin indices so any index outside of defined bins are moved to nearest defined bin + # Clip the bin indices so any index outside of defined bins are moved to nearest defined bin nbins_per_dim = [len(edges)-1 for edges in bin_edges] - + for idx, ibid in enumerate(bid): if not np.all(ibid> 0) or not np.all(ibid < len(self.bins[idx])): warnings.warn(f"Simulations with progress coordinates outside the bin boundaries definition of dimension {idx} are automatically placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") diff --git a/deepdrivewe/binners/rectilinear.py b/deepdrivewe/binners/rectilinear.py index a9f0ed6..96ee1f4 100644 --- a/deepdrivewe/binners/rectilinear.py +++ b/deepdrivewe/binners/rectilinear.py @@ -101,4 +101,3 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: # This ensures our bin index is >=0 and < len(self.bins) return np.clip(bin_id, 0, len(self.bins)-1) - diff --git a/tests/test_binner.py b/tests/binner_test.py similarity index 99% rename from tests/test_binner.py rename to tests/binner_test.py index 15364bc..8a98fe4 100644 --- a/tests/test_binner.py +++ b/tests/binner_test.py @@ -11,7 +11,7 @@ def test1dAssign(self): coords = np.array([-1, 0, 0.5, 1.5, 1.6, 2.0, 2.0, 2.9])[:, None] assigner = RectilinearBinner(bins=bounds, bin_target_counts = 3, target_state_inds=[None], pcoord_idx=0) - + with pytest.warns(UserWarning): assert (assigner.assign_bins(coords) == [0, 0, 0, 1, 1, 2, 2, 2]).all() @@ -41,5 +41,3 @@ def test2dAssign(self): with pytest.warns(UserWarning): assert (assigner.assign_bins(coords) == [0, 0, 5, 10, 10, 15, 7, 8]).all() - - From cd0ad468ee96d4a137807ce3dff3fb06bd14b064 Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Wed, 29 Oct 2025 14:51:04 -0500 Subject: [PATCH 63/73] lint part 2 --- tests/binner_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/binner_test.py b/tests/binner_test.py index 8a98fe4..174305a 100644 --- a/tests/binner_test.py +++ b/tests/binner_test.py @@ -6,16 +6,16 @@ from deepdrivewe.binners import RectilinearBinner, MultiRectilinearBinner class TestRectilinearBinner: - def test1dAssign(self): + def test1dAssign(self) -> None: bounds = [0.0, 1.0, 2.0, 3.0] coords = np.array([-1, 0, 0.5, 1.5, 1.6, 2.0, 2.0, 2.9])[:, None] - assigner = RectilinearBinner(bins=bounds, bin_target_counts = 3, target_state_inds=[None], pcoord_idx=0) + assigner = RectilinearBinner(bins=bounds, bin_target_counts=3, target_state_inds=[None], pcoord_idx=0) with pytest.warns(UserWarning): assert (assigner.assign_bins(coords) == [0, 0, 0, 1, 1, 2, 2, 2]).all() - def test2dAssign(self): + def test2dAssign(self) -> None: boundaries = [(-1, -0.5, 0, 0.5, 1), (-1, -0.5, 0, 0.5, 1)] coords = np.array([(-2, -2), (-0.75, -0.75), (-0.25, -0.25), (0, 0), (0.25, 0.25), (0.75, 0.75), (-0.25, 0.75), (0.25, -0.75)]) From 27033afd95b4b2b6a8b0c9b06dfeabd076fe4798 Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 29 Oct 2025 15:15:58 -0500 Subject: [PATCH 64/73] remove path prefix --- examples/amber_nacl_hk/config.yaml | 14 ++++++++------ examples/amber_ntl9_hk/config.yaml | 14 ++++++++------ examples/openmm_ntl9_ddwe_vista/config.yaml | 15 ++++++++------- examples/openmm_ntl9_hk/config.yaml | 7 ++++--- examples/synd_ntl9_hk/config.yaml | 6 ++++-- examples/synd_ntl9_lof/config.yaml | 12 +++++++----- 6 files changed, 39 insertions(+), 29 deletions(-) diff --git a/examples/amber_nacl_hk/config.yaml b/examples/amber_nacl_hk/config.yaml index 3a11297..ac58f1c 100644 --- a/examples/amber_nacl_hk/config.yaml +++ b/examples/amber_nacl_hk/config.yaml @@ -1,22 +1,24 @@ +# The inputs files are relative to the root of the deepdrivewe repository. + output_dir: runs/naive_resampler_test_v2 basis_states: - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_nacl_hk/inputs + basis_state_dir: examples/amber_nacl_hk/inputs basis_state_ext: .ncrst initial_ensemble_members: 5 basis_state_initializer: - top_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_nacl_hk/common_files/nacl.parm7 - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_nacl_hk/inputs/bstates/bstate.ncrst + top_file: examples/amber_nacl_hk/common_files/nacl.parm7 + reference_file: examples/amber_nacl_hk/inputs/bstates/bstate.ncrst num_iterations: 50 simulation_config: amber_config: amber_exe: sander - input_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_nacl_hk/common_files/md.in - top_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_nacl_hk/common_files/nacl.parm7 - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_nacl_hk/inputs/bstates/bstate.ncrst + input_file: examples/amber_nacl_hk/common_files/md.in + top_file: examples/amber_nacl_hk/common_files/nacl.parm7 + reference_file: examples/amber_nacl_hk/inputs/bstates/bstate.ncrst inference_config: # The number of simulations to maintain per bin diff --git a/examples/amber_ntl9_hk/config.yaml b/examples/amber_ntl9_hk/config.yaml index 0e034dc..1b3fb3b 100644 --- a/examples/amber_ntl9_hk/config.yaml +++ b/examples/amber_ntl9_hk/config.yaml @@ -1,22 +1,24 @@ +# The inputs files are relative to the root of the deepdrivewe repository. + output_dir: runs/ntl9-v1 basis_states: - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_ntl9_hk/inputs + basis_state_dir: examples/amber_ntl9_hk/inputs basis_state_ext: .rst7 initial_ensemble_members: 4 basis_state_initializer: - top_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_ntl9_hk/common_files/ntl9.prmtop - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_ntl9_hk/common_files/reference.pdb + top_file: examples/amber_ntl9_hk/common_files/ntl9.prmtop + reference_file: examples/amber_ntl9_hk/common_files/reference.pdb num_iterations: 106 simulation_config: amber_config: amber_exe: sander - input_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_ntl9_hk/common_files/md.in - top_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_ntl9_hk/common_files/ntl9.prmtop - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/amber_ntl9_hk/common_files/reference.pdb + input_file: examples/amber_ntl9_hk/common_files/md.in + top_file: examples/amber_ntl9_hk/common_files/ntl9.prmtop + reference_file: examples/amber_ntl9_hk/common_files/reference.pdb inference_config: # The number of simulations to maintain per bin diff --git a/examples/openmm_ntl9_ddwe_vista/config.yaml b/examples/openmm_ntl9_ddwe_vista/config.yaml index 4e7744d..256e8cf 100644 --- a/examples/openmm_ntl9_ddwe_vista/config.yaml +++ b/examples/openmm_ntl9_ddwe_vista/config.yaml @@ -1,4 +1,5 @@ # Configuration file for the NTL9 folding example using OpenMM +# The inputs files are relative to the root of the deepdrivewe repository. # The output directory for the runs output_dir: runs/ntl9-v1 @@ -9,7 +10,7 @@ num_iterations: 106 # The basis states to use for the ensemble basis_states: # The directory containing the basis states sub directories - basis_state_dir: /scratch/08288/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe_vista/inputs + basis_state_dir: examples/openmm_ntl9_ddwe_vista/inputs # The file extension for the basis state files basis_state_ext: .pdb # The number of basis states to use @@ -21,7 +22,7 @@ basis_states: # Strategy for initializing the basis state progress coordinates basis_state_initializer: # The path to the reference PDB file - reference_file: /scratch/08288/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe_vista/common_files/reference.pdb + reference_file: examples/openmm_ntl9_ddwe_vista/common_files/reference.pdb # The configuration for the simulation simulation_config: @@ -41,21 +42,21 @@ simulation_config: hardware_platform: CUDA # The path to the reference PDB file - reference_file: /scratch/08288/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe_vista/common_files/reference.pdb + reference_file: examples/openmm_ntl9_ddwe_vista/common_files/reference.pdb # The configuration for training train_config: # The path to the CVAE model configuration file - config_path: /scratch/08288/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe_vista/cvae-config.yaml + config_path: examples/openmm_ntl9_ddwe_vista/cvae-config.yaml # The path to the CVAE model weights file - checkpoint_path: /scratch/08288/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe_vista/checkpoint-epoch-100.pt + checkpoint_path: examples/openmm_ntl9_ddwe_vista/checkpoint-epoch-100.pt # The configuration for the inference inference_config: # The path to the CVAE model configuration file - ai_model_config_path: /scratch/08288/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe_vista/cvae-config.yaml + ai_model_config_path: examples/openmm_ntl9_ddwe_vista/cvae-config.yaml # The path to the CVAE model weights file - ai_model_checkpoint_path: /scratch/08288/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_ddwe_vista/checkpoint-epoch-100.pt + ai_model_checkpoint_path: examples/openmm_ntl9_ddwe_vista/checkpoint-epoch-100.pt # The number of neighbors to use for LOF lof_n_neighbors: 20 diff --git a/examples/openmm_ntl9_hk/config.yaml b/examples/openmm_ntl9_hk/config.yaml index 55873a5..de6a31c 100644 --- a/examples/openmm_ntl9_hk/config.yaml +++ b/examples/openmm_ntl9_hk/config.yaml @@ -1,4 +1,5 @@ # Configuration file for the NTL9 folding example using OpenMM +# The inputs files are relative to the root of the deepdrivewe repository. # The output directory for the runs output_dir: runs/ntl9-v1 @@ -9,7 +10,7 @@ num_iterations: 106 # The basis states to use for the ensemble basis_states: # The directory containing the basis states sub directories - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_hk/inputs + basis_state_dir: examples/openmm_ntl9_hk/inputs # The file extension for the basis state files basis_state_ext: .pdb # The number of basis states to use @@ -18,7 +19,7 @@ basis_states: # Strategy for initializing the basis state progress coordinates basis_state_initializer: # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_hk/common_files/reference.pdb + reference_file: examples/openmm_ntl9_hk/common_files/reference.pdb # The configuration for the simulation simulation_config: @@ -38,7 +39,7 @@ simulation_config: hardware_platform: CPU # The path to the reference PDB file - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_hk/common_files/reference.pdb + reference_file: examples/openmm_ntl9_hk/common_files/reference.pdb # The configuration for the inference inference_config: diff --git a/examples/synd_ntl9_hk/config.yaml b/examples/synd_ntl9_hk/config.yaml index 1550def..be4d5cd 100644 --- a/examples/synd_ntl9_hk/config.yaml +++ b/examples/synd_ntl9_hk/config.yaml @@ -1,3 +1,5 @@ +# The inputs files are relative to the root of the deepdrivewe repository. + # The output directory for the run output_dir: runs/ntl9-synd-v1 @@ -7,7 +9,7 @@ num_iterations: 150 # The basis states for the simulation basis_states: # The nested directory containing the basis state files - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/synd_ntl9_hk/bstates + basis_state_dir: examples/synd_ntl9_hk/bstates # The extension for the basis state files basis_state_ext: .npy # The number of initial ensemble members to use @@ -25,7 +27,7 @@ target_states: # The configuration for the simulation simulation_config: # The path to the synd model file - synd_model_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/synd_ntl9_hk/ntl9_folding.synd + synd_model_file: examples/synd_ntl9_hk/ntl9_folding.synd # The number of steps to run the simulation for (this includes the initial step) n_steps: 2 diff --git a/examples/synd_ntl9_lof/config.yaml b/examples/synd_ntl9_lof/config.yaml index 8e334a9..d0dec5d 100644 --- a/examples/synd_ntl9_lof/config.yaml +++ b/examples/synd_ntl9_lof/config.yaml @@ -1,3 +1,5 @@ +# The inputs files are relative to the root of the deepdrivewe repository. + # The output directory for the run output_dir: runs/ntl9-synd-v1 @@ -7,7 +9,7 @@ num_iterations: 150 # The basis states for the simulation basis_states: # The nested directory containing the basis state files - basis_state_dir: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/synd_ntl9_lof/bstates + basis_state_dir: examples/synd_ntl9_lof/bstates # The extension for the basis state files basis_state_ext: .npy # The number of initial ensemble members to use (should be the same as sims_per_bin in this use case) @@ -25,17 +27,17 @@ target_states: # The configuration for the simulation simulation_config: # The path to the synd model file - synd_model_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/synd_ntl9_lof/ntl9_folding.synd + synd_model_file: examples/synd_ntl9_lof/ntl9_folding.synd # The number of steps to run the simulation for (this includes the initial step) n_steps: 2 # The reference structure for computing contact maps (getting the CA atom indices) - reference_file: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/synd_ntl9_lof/ntl9_reference.pdb + reference_file: examples/synd_ntl9_lof/ntl9_reference.pdb inference_config: # The path to the CVAE model configuration file - ai_model_config_path: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/synd_ntl9_lof/cvae-config.yaml + ai_model_config_path: examples/synd_ntl9_lof/cvae-config.yaml # The path to the CVAE model weights file - ai_model_checkpoint_path: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/synd_ntl9_lof/checkpoint-epoch-100.pt + ai_model_checkpoint_path: examples/synd_ntl9_lof/checkpoint-epoch-100.pt # The number of neighbors to use for LOF lof_n_neighbors: 20 From d3eb24cf50387f1d3aa62499ca1c85fe6bafd37d Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 29 Oct 2025 15:16:44 -0500 Subject: [PATCH 65/73] add path validator --- deepdrivewe/examples/synd_ntl9_lof/inference.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/deepdrivewe/examples/synd_ntl9_lof/inference.py b/deepdrivewe/examples/synd_ntl9_lof/inference.py index d692399..805f4ad 100644 --- a/deepdrivewe/examples/synd_ntl9_lof/inference.py +++ b/deepdrivewe/examples/synd_ntl9_lof/inference.py @@ -7,6 +7,7 @@ import numpy as np from pydantic import BaseModel from pydantic import Field +from pydantic import field_validator from sklearn.neighbors import LocalOutlierFactor from deepdrivewe import BasisStates @@ -14,6 +15,7 @@ from deepdrivewe import SimMetadata from deepdrivewe import SimResult from deepdrivewe import TargetState +from deepdrivewe import validate_and_resolve_file from deepdrivewe.ai import warmstart_model from deepdrivewe.binners import RectilinearBinner from deepdrivewe.recyclers import LowRecycler @@ -67,6 +69,12 @@ class InferenceConfig(BaseModel): 'is 10e-40.', ) + @field_validator('ai_model_config_path', 'ai_model_checkpoint_path') + @classmethod + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) + def run_inference( sim_output: list[SimResult], From 40ebf8be469b5c465032fe6317aeccaadf285af7 Mon Sep 17 00:00:00 2001 From: braceal Date: Wed, 29 Oct 2025 15:21:46 -0500 Subject: [PATCH 66/73] path validation --- deepdrivewe/examples/openmm_ntl9_ddwe/main.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py index dc10437..a58cb7b 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/main.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py @@ -26,6 +26,7 @@ from deepdrivewe import BasisStates from deepdrivewe import EnsembleCheckpointer from deepdrivewe import TargetState +from deepdrivewe import validate_and_resolve_file from deepdrivewe import WeightedEnsemble from deepdrivewe.examples.openmm_ntl9_ddwe.inference import InferenceConfig from deepdrivewe.examples.openmm_ntl9_ddwe.inference import run_inference @@ -52,13 +53,9 @@ class RMSDBasisStateInitializer(BaseModel): @field_validator('reference_file') @classmethod - def validate_reference_file(cls, value: Path) -> Path: - """Validate and resolve the reference file.""" - if not value.is_file(): - raise FileNotFoundError( - f'The reference file {value} is not a file.', - ) - return value.resolve() + def validate_and_resolve_file(cls, value: Path | None) -> Path | None: + """Validate and resolve the file path.""" + return validate_and_resolve_file(value) def __call__(self, basis_file: str) -> list[float]: """Initialize the basis state parent coordinates.""" From 91196c4ec964e9b9ad3b65b7173162a0b4bbc2b9 Mon Sep 17 00:00:00 2001 From: braceal Date: Thu, 30 Oct 2025 10:00:17 -0500 Subject: [PATCH 67/73] format --- deepdrivewe/binners/multirectilinear.py | 26 +++++++++++++++++-------- deepdrivewe/binners/rectilinear.py | 18 ++++++++++++----- 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 227e4be..053fc73 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -2,9 +2,10 @@ from __future__ import annotations +import warnings + import numpy as np from scipy.stats import binned_statistic_dd -import warnings from deepdrivewe.binners.base import Binner @@ -76,20 +77,29 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: expand_binnumbers=True, ) - # Clip the bin indices so any index outside of defined bins are moved to nearest defined bin - nbins_per_dim = [len(edges)-1 for edges in bin_edges] + # Clip the bin indices so any index outside of defined bins are moved + # to nearest defined bin + nbins_per_dim = [len(edges) - 1 for edges in bin_edges] for idx, ibid in enumerate(bid): - if not np.all(ibid> 0) or not np.all(ibid < len(self.bins[idx])): - warnings.warn(f"Simulations with progress coordinates outside the bin boundaries definition of dimension {idx} are automatically placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") + if not np.all(ibid > 0) or not np.all(ibid < len(self.bins[idx])): + warnings.warn( + 'Simulations with progress coordinates outside the bin ' + f'boundaries definition of dimension {idx} are ' + 'automatically placed into the nearest terminal bins. ' + 'Consider modifying your bin boundaries by adding ' + "'np.inf' or '-np.inf' on either end of your bin " + 'definitions.', + stacklevel=2, + ) bid[idx] = np.clip(ibid, 1, nbins_per_dim[idx]) # Calculate the bin indices in row-major order bin_ids = np.zeros(len(pcoords), dtype=int) for idx, ibid in enumerate(bid.T): - for idim in range(len(nbins_per_dim)-1): - bin_ids[idx] += (ibid[idim] -1) * nbins_per_dim[idim] - bin_ids[idx] += (ibid[-1] - 1) + for idim in range(len(nbins_per_dim) - 1): + bin_ids[idx] += (ibid[idim] - 1) * nbins_per_dim[idim] + bin_ids[idx] += ibid[-1] - 1 # Check that the number of bin indices is the same as the # number of simulations diff --git a/deepdrivewe/binners/rectilinear.py b/deepdrivewe/binners/rectilinear.py index d583bd5..2310c03 100644 --- a/deepdrivewe/binners/rectilinear.py +++ b/deepdrivewe/binners/rectilinear.py @@ -2,9 +2,10 @@ from __future__ import annotations -import numpy as np import warnings +import numpy as np + from deepdrivewe.binners.base import Binner @@ -65,10 +66,17 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: """ # Bin the progress coordinates (make sure the target state # boundary is included in the target state bin). - bin_id = np.digitize(pcoords[:, self.pcoord_idx], self.bins) -1 + bin_ids = np.digitize(pcoords[:, self.pcoord_idx], self.bins) - 1 - if not np.all(bin_id > 0) or not np.all(bin_id < len(self.bins)): - warnings.warn("Simulations with progress coordinates outside the bin boundaries definitions are placed into the nearest terminal bins. Consider modifying your bin boundaries by adding 'np.inf' or '-np.inf' on either end of your bin definitions.") + # Check that the bin indices are within the valid range + if not np.all(bin_ids > 0) or not np.all(bin_ids < len(self.bins)): + warnings.warn( + 'Simulations with progress coordinates outside the bin ' + 'boundaries definitions are placed into the nearest terminal ' + 'bins. Consider modifying your bin boundaries by adding ' + "'np.inf' or '-np.inf' on either end of your bin definitions.", + stacklevel=2, + ) # This ensures our bin index is >=0 and < len(self.bins) - return np.clip(bin_id, 0, len(self.bins)-1) + return np.clip(bin_ids, 0, len(self.bins) - 1) From e21fbf7d3ed4b76f8b8af129e2ca36142496bb8e Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Thu, 29 Jan 2026 11:39:54 -0500 Subject: [PATCH 68/73] bug fix for binning 1D pcoord in multirectilinear binner --- deepdrivewe/binners/multirectilinear.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 053fc73..3534419 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -81,6 +81,9 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: # to nearest defined bin nbins_per_dim = [len(edges) - 1 for edges in bin_edges] + # If binning a 1D coordinate, a 1D array will be returned. + bid = np.atleast_2d(bid) + for idx, ibid in enumerate(bid): if not np.all(ibid > 0) or not np.all(ibid < len(self.bins[idx])): warnings.warn( From 2f1dd9139fae48141a52eb6425497836c62e355d Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Tue, 11 Aug 2026 15:33:24 -0400 Subject: [PATCH 69/73] fix bin id output for MultiRectilinearBinner --- deepdrivewe/binners/multirectilinear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index 3534419..f5ec418 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -101,7 +101,7 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: bin_ids = np.zeros(len(pcoords), dtype=int) for idx, ibid in enumerate(bid.T): for idim in range(len(nbins_per_dim) - 1): - bin_ids[idx] += (ibid[idim] - 1) * nbins_per_dim[idim] + bin_ids[idx] += (ibid[idim] - 1) * nbins_per_dim[idim + 1] bin_ids[idx] += ibid[-1] - 1 # Check that the number of bin indices is the same as the From d36eacfa757b11fdcbe9dc1531e5ed63d16d2d94 Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Tue, 11 Aug 2026 15:33:54 -0400 Subject: [PATCH 70/73] new test for MultiRectilinearBinner --- tests/binner_test.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/binner_test.py b/tests/binner_test.py index 174305a..71fa6eb 100644 --- a/tests/binner_test.py +++ b/tests/binner_test.py @@ -41,3 +41,12 @@ def test2dAssign(self) -> None: with pytest.warns(UserWarning): assert (assigner.assign_bins(coords) == [0, 0, 5, 10, 10, 15, 7, 8]).all() + + def test2dAssign_v2(self) -> None: + boundaries = [(0, 1, 2, 3), (0, 1, 2)] + coords = np.array([(0.5, 0.5), (0.5, 1.5), (1.5, 0.5), (1.5, 1.5), (2.5, 0.5), (2.5, 1.5), (3.5, 1.5)]) + + assigner = MultiRectilinearBinner(boundaries, bin_target_counts=3, target_state_inds=[None]) + + with pytest.warns(UserWarning): + assert (assigner.assign_bins(coords) == [0, 1, 2, 3, 4, 5, 5]).all() From b40ffed84d6f3fb041f19b07693cc03f13f3119f Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Tue, 11 Aug 2026 16:05:03 -0400 Subject: [PATCH 71/73] further fix for 3d+ and also 3D test --- deepdrivewe/binners/multirectilinear.py | 2 +- tests/binner_test.py | 34 +++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/deepdrivewe/binners/multirectilinear.py b/deepdrivewe/binners/multirectilinear.py index f5ec418..78ac364 100644 --- a/deepdrivewe/binners/multirectilinear.py +++ b/deepdrivewe/binners/multirectilinear.py @@ -101,7 +101,7 @@ def assign_bins(self, pcoords: np.ndarray) -> np.ndarray: bin_ids = np.zeros(len(pcoords), dtype=int) for idx, ibid in enumerate(bid.T): for idim in range(len(nbins_per_dim) - 1): - bin_ids[idx] += (ibid[idim] - 1) * nbins_per_dim[idim + 1] + bin_ids[idx] += (ibid[idim] - 1) * np.prod(nbins_per_dim[idim + 1:]) bin_ids[idx] += ibid[-1] - 1 # Check that the number of bin indices is the same as the diff --git a/tests/binner_test.py b/tests/binner_test.py index 71fa6eb..b835d90 100644 --- a/tests/binner_test.py +++ b/tests/binner_test.py @@ -49,4 +49,38 @@ def test2dAssign_v2(self) -> None: assigner = MultiRectilinearBinner(boundaries, bin_target_counts=3, target_state_inds=[None]) with pytest.warns(UserWarning): + # first 6 points are in bins [0, 5]. The last point locate outside the bounds but will be clipped to bin 5 assert (assigner.assign_bins(coords) == [0, 1, 2, 3, 4, 5, 5]).all() + + def test3dAssign(self) -> None: + boundaries = [(0, 1, 2), (0, 1, 2, 3, 4, 5), (0, 1, 2)] + coords = np.array([(0.5, 0.5, 0.5), + (0.5, 0.5, 1.5), + (0.5, 1.5, 0.5), + (0.5, 1.5, 1.5), + (0.5, 2.5, 0.5), + (0.5, 2.5, 1.5), + (0.5, 3.5, 0.5), + (0.5, 3.5, 1.5), + (0.5, 4.5, 0.5), + (0.5, 4.5, 1.5), + (1.5, 0.5, 0.5), + (1.5, 0.5, 1.5), + (1.5, 1.5, 0.5), + (1.5, 1.5, 1.5), + (1.5, 2.5, 0.5), + (1.5, 2.5, 1.5), + (1.5, 3.5, 0.5), + (1.5, 3.5, 1.5), + (1.5, 4.5, 0.5), + (1.5, 4.5, 1.5), + (2.5, 4.5, 1.5), + (1.5, 5.5, 1.5),]) + + assigner = MultiRectilinearBinner(boundaries, bin_target_counts=3, target_state_inds=[None]) + + with pytest.warns(UserWarning): + # first 20 points are in bins [0, 19]. The last two locate outside the bounds but will be clipped to bin 19 + assert (assigner.assign_bins(coords) == list(range(20)) + [19, 19]).all() + + From e4d71ef18145883bc24c06fcedbb526174081e1c Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Tue, 11 Aug 2026 16:10:59 -0400 Subject: [PATCH 72/73] clean up test a little bit --- tests/binner_test.py | 25 +++---------------------- 1 file changed, 3 insertions(+), 22 deletions(-) diff --git a/tests/binner_test.py b/tests/binner_test.py index b835d90..07073cc 100644 --- a/tests/binner_test.py +++ b/tests/binner_test.py @@ -1,5 +1,6 @@ import os import pytest +from itertools import product import numpy as np @@ -54,28 +55,8 @@ def test2dAssign_v2(self) -> None: def test3dAssign(self) -> None: boundaries = [(0, 1, 2), (0, 1, 2, 3, 4, 5), (0, 1, 2)] - coords = np.array([(0.5, 0.5, 0.5), - (0.5, 0.5, 1.5), - (0.5, 1.5, 0.5), - (0.5, 1.5, 1.5), - (0.5, 2.5, 0.5), - (0.5, 2.5, 1.5), - (0.5, 3.5, 0.5), - (0.5, 3.5, 1.5), - (0.5, 4.5, 0.5), - (0.5, 4.5, 1.5), - (1.5, 0.5, 0.5), - (1.5, 0.5, 1.5), - (1.5, 1.5, 0.5), - (1.5, 1.5, 1.5), - (1.5, 2.5, 0.5), - (1.5, 2.5, 1.5), - (1.5, 3.5, 0.5), - (1.5, 3.5, 1.5), - (1.5, 4.5, 0.5), - (1.5, 4.5, 1.5), - (2.5, 4.5, 1.5), - (1.5, 5.5, 1.5),]) + coords = list(product([0.5, 1.5], [0.5, 1.5, 2.5, 3.5, 4.5], [0.5, 1.5])) # One point per bin, in row-major order + coords += [(2.5, 4.5, 1.5), (1.5, 5.5, 1.5)] # Two points that are located outside the bin boundaries assigner = MultiRectilinearBinner(boundaries, bin_target_counts=3, target_state_inds=[None]) From cb19ac4ac9e5bce87dd924400e21347ad5e86640 Mon Sep 17 00:00:00 2001 From: "Jeremy M. G. Leung" Date: Tue, 11 Aug 2026 16:24:56 -0400 Subject: [PATCH 73/73] make sure test array is actually an array --- tests/binner_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/binner_test.py b/tests/binner_test.py index 07073cc..97ccae9 100644 --- a/tests/binner_test.py +++ b/tests/binner_test.py @@ -57,7 +57,8 @@ def test3dAssign(self) -> None: boundaries = [(0, 1, 2), (0, 1, 2, 3, 4, 5), (0, 1, 2)] coords = list(product([0.5, 1.5], [0.5, 1.5, 2.5, 3.5, 4.5], [0.5, 1.5])) # One point per bin, in row-major order coords += [(2.5, 4.5, 1.5), (1.5, 5.5, 1.5)] # Two points that are located outside the bin boundaries - + coords = np.asarray(coords) + assigner = MultiRectilinearBinner(boundaries, bin_target_counts=3, target_state_inds=[None]) with pytest.warns(UserWarning):