diff --git a/README.md b/README.md index 8a94d0d..f75e5ec 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,14 +18,11 @@ 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 . ``` -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: -```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 @@ -48,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 @@ -61,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: @@ -70,6 +80,21 @@ 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 +``` + +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 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/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..ddc5b14 --- /dev/null +++ b/deepdrivewe/ai/aae.py @@ -0,0 +1,259 @@ +"""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 + + # Initialize the model + self.trainer = AAE3dTrainer(**config.model_dump()) + + # 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) + 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..38ff402 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,11 +125,10 @@ 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: @@ -141,6 +142,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 @@ -209,5 +214,44 @@ 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 + + +@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/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}', ) diff --git a/deepdrivewe/api.py b/deepdrivewe/api.py index 13daa4f..d16efe1 100644 --- a/deepdrivewe/api.py +++ b/deepdrivewe/api.py @@ -18,10 +18,41 @@ import yaml # type: ignore[import-untyped] from pydantic import BaseModel as _BaseModel from pydantic import Field +from pydantic import field_validator 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.""" @@ -284,6 +315,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 +347,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()) diff --git a/deepdrivewe/binners/__init__.py b/deepdrivewe/binners/__init__.py index abad4f5..b44ba76 100644 --- a/deepdrivewe/binners/__init__.py +++ b/deepdrivewe/binners/__init__.py @@ -4,4 +4,5 @@ # Forward imports from deepdrivewe.binners.base import Binner +from deepdrivewe.binners.multirectilinear import MultiRectilinearBinner from deepdrivewe.binners.rectilinear import RectilinearBinner 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 new file mode 100644 index 0000000..78ac364 --- /dev/null +++ b/deepdrivewe/binners/multirectilinear.py @@ -0,0 +1,114 @@ +"""Multirectilinear binner.""" + +from __future__ import annotations + +import warnings + +import numpy as np +from scipy.stats import binned_statistic_dd + +from deepdrivewe.binners.base import Binner + + +class MultiRectilinearBinner(Binner): + """Multirectilinear binner for multiple progress coordinates.""" + + def __init__( + self, + bins: list[np.ndarray | list[float]], + bin_target_counts: int | list[int], + target_state_inds: int | list[int] | None = None, + ) -> None: + """Initialize the binner. + + Parameters + ---------- + bins : list[np.ndarray | list[float]] + 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. + 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. + """ + super().__init__(bin_target_counts, target_state_inds) + + self.bins = bins + + # 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.""" + # Calculate the number of bins per dimension + nbins_per_dim = np.array([len(dim) - 1 for dim in self.bins]) + + # Calculate the total number of bins + return int(np.prod(nbins_per_dim)) + + 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). + _, bin_edges, bid = binned_statistic_dd( + np.asarray(pcoords), + values=None, + statistic='count', + bins=self.bins, + 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] + + # 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( + '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) * 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 + # number of simulations + if len(bin_ids) != len(pcoords): + raise ValueError( + 'Number of bin indices must match the number of simulations.', + ) + + return bin_ids diff --git a/deepdrivewe/binners/rectilinear.py b/deepdrivewe/binners/rectilinear.py index 7244036..2310c03 100644 --- a/deepdrivewe/binners/rectilinear.py +++ b/deepdrivewe/binners/rectilinear.py @@ -2,6 +2,8 @@ from __future__ import annotations +import warnings + import numpy as np from deepdrivewe.binners.base import Binner @@ -14,7 +16,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 +25,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 +51,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. @@ -93,4 +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). - return np.digitize(pcoords[:, self.pcoord_idx], self.bins, right=True) + bin_ids = np.digitize(pcoords[:, self.pcoord_idx], self.bins) - 1 + + # 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_ids, 0, len(self.bins) - 1) diff --git a/deepdrivewe/cli.py b/deepdrivewe/cli.py index 2a8da7f..e6a36ae 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,37 @@ def version() -> None: print(f'deepdrivewe, version {__version__}') +@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.""" + # Create a console for rich output + console = Console() + + # Find all the task result files + results_dir = run_dir / 'result' + + # Read the simulation, train, and inference results + for file_path in results_dir.glob('*.json'): + # 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']: + console.print( + f"[bold blue]Method:[/bold blue] {data['method']}", + ) + console.print(data['failure_info']['traceback'], style='red') + + @app.command() def to_pdb( coordinate_file: Path = typer.Option( # noqa: B008 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_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/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/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/inference.py b/deepdrivewe/examples/openmm_aae_ddwe/inference.py new file mode 100644 index 0000000..5521458 --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/inference.py @@ -0,0 +1,156 @@ +"""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, + target_state_inds=0, # The first bin is the target (folded) state + ) + + # 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..a92d582 --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/main.py @@ -0,0 +1,256 @@ +"""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 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 +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 + + +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.', + ) + + @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 + 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..6d2ef08 --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/simulate.py @@ -0,0 +1,118 @@ +"""Simulate a system using OpenMM.""" + +from __future__ import annotations + +from pathlib import Path +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 +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.', + ) + + @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, + 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..93633ee --- /dev/null +++ b/deepdrivewe/examples/openmm_aae_ddwe/train.py @@ -0,0 +1,83 @@ +"""Training module.""" + +from __future__ import annotations + +from pathlib import Path + +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 + + +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.', + ) + + @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. +# 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..d57faf7 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, ) @@ -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_ddwe/main.py b/deepdrivewe/examples/openmm_ntl9_ddwe/main.py index 0dfd678..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 @@ -34,7 +35,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): @@ -48,6 +51,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 @@ -102,6 +111,10 @@ class ExperimentSettings(BaseModel): compute_config: ComputeConfigTypes = Field( description='Config for the compute resources.', ) + stream_config: ProxyStreamConfig | None = Field( + default=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 ' @@ -190,11 +203,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, @@ -219,15 +234,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/examples/openmm_ntl9_ddwe/simulate.py b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py index cb25d5c..ec73137 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/simulate.py @@ -6,13 +6,18 @@ 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.simulation.openmm import ContactMapRMSDReporter +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 from deepdrivewe.simulation.openmm import OpenMMSimulation +from deepdrivewe.simulation.openmm import RMSDCollector +from deepdrivewe.workflows.stream import ProxyStreamConfig class SimulationConfig(BaseModel): @@ -41,11 +46,18 @@ 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, config: SimulationConfig, output_dir: Path, + stream_config: ProxyStreamConfig | None = None, ) -> SimResult: """Run a simulation and return the pcoord and coordinates.""" # Add performance logging @@ -75,30 +87,39 @@ 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', + ), + ], + stream_config=stream_config, ) # 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, - ) + # 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..6e8abf5 100644 --- a/deepdrivewe/examples/openmm_ntl9_ddwe/train.py +++ b/deepdrivewe/examples/openmm_ntl9_ddwe/train.py @@ -2,16 +2,22 @@ from __future__ import annotations +import itertools from pathlib import Path 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 +from deepdrivewe.workflows.stream import SIMULATION_TOPIC +from deepdrivewe.workflows.stream import TRAIN_TOPIC class TrainConfig(BaseModel): @@ -25,6 +31,22 @@ 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.', + ) + + @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 @@ -36,8 +58,17 @@ 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( + 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}' @@ -52,9 +83,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() @@ -73,3 +105,81 @@ def run_train( ) return result + + +def run_stream_train( + config: TrainConfig, + output_dir: Path, + stream_config: ProxyStreamConfig, +) -> TrainResult: + """Train the model on the simulation output.""" + # Make the output directory + 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) + + # 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=model_dir, + 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/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/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_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..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], @@ -142,6 +150,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/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/parsl.py b/deepdrivewe/parsl.py index 5ef234b..7837e06 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 @@ -15,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 @@ -128,6 +129,92 @@ 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( + 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.', + ) + 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: + """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, + available_accelerators: Sequence[str], + ) -> HighThroughputExecutor: + return HighThroughputExecutor( + address=self.address, + 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.""" @@ -183,6 +270,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.""" @@ -196,6 +291,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', @@ -273,16 +369,102 @@ 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), ], ) +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_htex', 'inference_htex', 'simulation_htex'] + 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 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) + + # 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, + WorkstationV2Config, HybridWorkstationConfig, InferenceTrainWorkstationConfig, VistaConfig, + PolarisConfig, ] 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.""" diff --git a/deepdrivewe/simulation/openmm.py b/deepdrivewe/simulation/openmm.py index f13290c..441309a 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 @@ -17,12 +18,16 @@ 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 from pydantic import BaseModel 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: @@ -34,17 +39,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 +83,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 +132,304 @@ 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) -> Any: + """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) -> np.ndarray: + """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) + + return 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) -> float: + """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) + return 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._contact_maps: 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, *). + """ + # Collect the contact maps in a ragged numpy array + contact_maps = np.array(self._contact_maps, dtype=object) + + return contact_maps + + def collect(self, positions: np.ndarray) -> np.ndarray: + """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() + + # 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._contact_maps.append(sparse_contact_map) + + return sparse_contact_map + + +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',), + stream_config: ProxyStreamConfig | None = None, + ) -> 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 + + # 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. + + 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 + 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): """Configuration for an OpenMM simulation.""" @@ -181,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. @@ -188,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 ------- @@ -195,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), str(pdb_file)) # Configure system system = top.createSystem( @@ -352,6 +701,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. @@ -363,6 +713,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 ------- @@ -380,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: @@ -406,6 +761,35 @@ 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: + # 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( + sim.topology, + sim.system, + self.configure_integrator(), + platform, + platform_properties, + # 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)) @@ -486,10 +870,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 @@ -518,13 +907,6 @@ 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)) - # Run simulation sim.step(self.config.num_steps) @@ -532,6 +914,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 +993,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/deepdrivewe/workflows/ddwe.py b/deepdrivewe/workflows/ddwe.py index 77a9fa1..c293b91 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 @@ -11,16 +12,20 @@ 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 class DDWEThinker(BaseThinker): """A thinker for the DDWE workflow.""" - def __init__( # noqa: PLR0913 + def __init__( self, queue: ColmenaQueues, result_dir: Path, @@ -28,7 +33,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 +54,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 +63,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 +96,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 +116,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 +145,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 +197,216 @@ 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 + + # 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. + + 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. 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', []) + + @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') + + # 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 + + # 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) + # 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') + + @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..be36aba --- /dev/null +++ b/deepdrivewe/workflows/stream.py @@ -0,0 +1,83 @@ +"""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 import StreamConsumer +from proxystore.stream import StreamProducer +from proxystore.stream.shims.redis import RedisQueuePublisher +from proxystore.stream.shims.redis import RedisQueueSubscriber + +from deepdrivewe import BaseModel + +SIMULATION_TOPIC = 'simulation-output' +TRAIN_TOPIC = 'train-output' + + +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/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/config.yaml b/examples/openmm_ntl9_ddwe/config.yaml index 9a8ea6b..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 @@ -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_ddwe/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: /nfs/lambda_stor_01/homes/abrace/projects/ddwe/src/deepdrivewe/examples/openmm_ntl9_hk/common_files/reference.pdb + reference_file: examples/openmm_ntl9_ddwe/common_files/reference.pdb # The configuration for the simulation simulation_config: @@ -34,29 +35,24 @@ 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 - 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_hk/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: - # 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] @@ -77,29 +73,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"] diff --git a/examples/openmm_ntl9_ddwe_vista/config.yaml b/examples/openmm_ntl9_ddwe_vista/config.yaml index c59f936..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: @@ -34,28 +35,28 @@ 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 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 177a0af..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: @@ -31,14 +32,14 @@ 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 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 diff --git a/pyproject.toml b/pyproject.toml index 9eb66db..11a57b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,17 +20,18 @@ 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", "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.0", - "scipy==1.14.0", + "mdlearn==1.0.5", + "scipy", "natsort>=8.4.0", "matplotlib>=3.9.2", ] @@ -145,7 +146,7 @@ select = [ "RUF", ] line-length = 79 -extend-ignore = ["Q001"] +extend-ignore = ["Q001", "SIM102", "PLR0913"] target-version = "py38" [tool.ruff.flake8-pytest-style] diff --git a/tests/binner_test.py b/tests/binner_test.py new file mode 100644 index 0000000..97ccae9 --- /dev/null +++ b/tests/binner_test.py @@ -0,0 +1,68 @@ +import os +import pytest +from itertools import product + +import numpy as np + +from deepdrivewe.binners import RectilinearBinner, MultiRectilinearBinner + +class TestRectilinearBinner: + 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) + + with pytest.warns(UserWarning): + assert (assigner.assign_bins(coords) == [0, 0, 0, 1, 1, 2, 2, 2]).all() + + 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)]) + + 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)]""" + + 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): + # 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 = 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): + # 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() + +