From fc13b8358cef5e081b28c32b913031eed6911bfe Mon Sep 17 00:00:00 2001 From: EdoardoRolando Date: Thu, 19 Jun 2025 17:15:16 +0200 Subject: [PATCH 1/3] Added parameter atoms_batch_size to allow batched atoms processing when computing constraints. --- input_generator/raw_dataset.py | 3 +- input_generator/utils.py | 44 ++++++++++++++++++++++-- scripts/gen_input_data.py | 63 +++++++++++++++++++--------------- 3 files changed, 80 insertions(+), 30 deletions(-) diff --git a/input_generator/raw_dataset.py b/input_generator/raw_dataset.py index d047707..5edb38f 100644 --- a/input_generator/raw_dataset.py +++ b/input_generator/raw_dataset.py @@ -296,6 +296,7 @@ def process_coords_forces( filter_cis: bool = False, force_stride: int = 100, batch_size: Optional[int] = None, + atoms_batch_size: Optional[int] = None, ) -> Tuple[np.ndarray, np.ndarray]: """ Maps coordinates and forces to CG resolution @@ -341,7 +342,7 @@ def process_coords_forces( break cg_coords, cg_forces, cg_map, force_map = slice_coord_forces( - coords, forces, self.cg_map, mapping, force_stride, batch_size + coords, forces, self.cg_map, mapping, force_stride, batch_size, atoms_batch_size ) # update the entries with the sparse version self.cg_map = cg_map diff --git a/input_generator/utils.py b/input_generator/utils.py index fbc5e05..5b31317 100644 --- a/input_generator/utils.py +++ b/input_generator/utils.py @@ -242,6 +242,7 @@ def slice_coord_forces( mapping: str = "slice_aggregate", force_stride: int = 100, batch_size: Optional[int] = None, + atoms_batch_size: Optional[int] = None, ) -> Tuple: """ Parameters @@ -261,15 +262,54 @@ def slice_coord_forces( batch_size: Optional length of batch in which divide the AA mapping of coords and forces to CG ones + atoms_batch_size: + Optional length of batch in which divide the atoms in coords to compute the pairwise constraints. Returns ------- Coarse-grained coordinates and forces """ + # Original hard coded values + n_frames = 100 # taking only first 100 frames gives same results in ~1/15th of time + threshold = 5e-3 # threshold for pairwise constraints + config_map = LinearMap(cg_map) config_map_matrix = config_map.standard_matrix - # taking only first 100 frames gives same results in ~1/15th of time - constraints = guess_pairwise_constraints(coords[:100], threshold=5e-3) + n_sites = coords.shape[1] # number of atomistic sites + + if atoms_batch_size is None or atoms_batch_size >= n_sites: + # No batching: process all atoms at once + constraints = guess_pairwise_constraints(coords[:n_frames], threshold=threshold) + + else: + # Batching mode + batches = [(range(i, min(i + atoms_batch_size, n_sites))) for i in range(0, n_sites, atoms_batch_size)] + constraints = set() + + # Within-batch constraints + for batch in batches: + xyz_batch = coords[:n_frames, batch, :] + local_constraints = guess_pairwise_constraints(xyz_batch, threshold=threshold) + global_constraints = {frozenset([batch[i] for i in pair]) for pair in local_constraints} + constraints.update(global_constraints) + + # Cross-batch constraints + # To significantly reduce computational cost, we assume residues are ordered in the structure. + # Therefore, we only compute constraints between consecutive batches, rather than all pairs of batches. + # For even greater efficiency, one could restrict this to just the first and last (e.g., 30) atoms of each batch, + # which scales as O(1) basically, but computing all pairs between consecutive batches is generally still efficient. + # This tecnically can be extended to the case with no batches (smaller molecules), assuming again ordered resiues, + # treating all molecules at the same way and getting rid of the atoms_batch_size parameters. + for i in range(len(batches) - 1): + b1 = batches[i] + b2 = batches[i + 1] + xyz1 = coords[:n_frames, b1, :] + xyz2 = coords[:n_frames, b2, :] + local_constraints = guess_pairwise_constraints(xyz1, cross_xyz=xyz2, threshold=threshold) + # guess_pairwise_constraints returns ordered pairs (i, j) where i indexes into cross_xyz (b2) and j indexes into xyz (b1) + global_constraints = {frozenset([b1[j], b2[i]]) for i, j in local_constraints} + constraints.update(global_constraints) + if isinstance(mapping, str): if mapping == "slice_aggregate": method = constraint_aware_uni_map diff --git a/scripts/gen_input_data.py b/scripts/gen_input_data.py index 3348d1c..ca86562 100644 --- a/scripts/gen_input_data.py +++ b/scripts/gen_input_data.py @@ -37,50 +37,55 @@ def process_raw_dataset( filter_cis: Optional[bool] = False, batch_size: Optional[int] = None, mol_num_batches: Optional[int] = 1, + atoms_batch_size: Optional[int] = None, ): """ - Applies coarse-grained mapping to coordinates and forces using input sample - topology and specified mapping strategies + Processes a raw dataset by applying coarse-grained (CG) mapping to atomic coordinates and forces, + using the provided topology and mapping strategies. Parameters ---------- dataset_name : str - Name given to specific dataset + Identifier for the dataset. names : List[str] - List of sample names + List of molecule or sample names to process. sample_loader : DatasetLoader - Loader object defined for specific dataset + Loader object for retrieving trajectories and data. raw_data_dir : str - Path to coordinate and force files + Directory containing raw coordinate and force files. tag : str - Label given to all output files produced from dataset + Label to append to all output files generated from this dataset. pdb_template_fn : str - Template file location of atomistic structure to be used for topology + Path to the template PDB file for atomistic topology. save_dir : str - Path to directory in which output will be saved + Directory where processed outputs will be saved. cg_atoms : List[str] - List of atom names to preserve in coarse-grained resolution + Atom names to retain in the coarse-grained representation. embedding_map : CGEmbeddingMap - Mapping object + Object defining the mapping from atomistic to CG representations. embedding_func : Callable - Function which will be used to apply CG mapping + Function to apply the CG mapping. skip_residues : List[str] - List of residues to skip, can be None + Residues to exclude from processing (can be empty). cg_mapping_strategy : str - Strategy to use for coordinate and force mappings; - currently only "slice_aggregate" and "slice_optimize" are implemented - stride : int - Interval by which to stride loaded data - force_stride : int - stride for inferring the force maps in aggforce - filter_cis : bool - if True, frames with cis-configurations will be filtered out from the dataset - batch_size : int - Optional size in which performing batches of AA mapping to CG, to avoid - memory overhead in large AA dataset - mol_num_batches : int - If greater than 1, will save each molecule data into the specified number of batches - that will be treated as different samples + Strategy for mapping coordinates and forces (e.g., "slice_aggregate", "slice_optimize"). + stride : int, optional + Interval for subsampling loaded data (default: 1). + force_stride : int, optional + Interval for subsampling forces during mapping (default: 100). + filter_cis : Optional[bool], optional + If True, filters out frames with cis-configurations (default: False). + batch_size : Optional[int], optional + Number of frames to process in each batch to reduce memory usage (default: None). + mol_num_batches : Optional[int], optional + Number of batches to split each molecule's data into (default: 1). + atoms_batch_size : Optional[int], optional + Batch size for atoms when processing large molecules (default: None). + + Returns + ------- + None + Saves processed CG data and mapping files to the specified directory. """ dataset = RawDataset(dataset_name, names, tag, n_batches=mol_num_batches) for samples in tqdm(dataset, f"Processing CG data for {dataset_name} dataset..."): @@ -117,6 +122,7 @@ def process_raw_dataset( force_stride=force_stride, batch_size=batch_size, filter_cis=filter_cis, + atoms_batch_size=atoms_batch_size, ) samples.save_cg_output(save_dir, save_coord_force=True, save_cg_maps=True) @@ -147,6 +153,7 @@ def build_neighborlists( filter_cis: bool = False, batch_size: Optional[int] = None, mol_num_batches: Optional[int] = 1, + atoms_batch_size: Optional[int] = None, ): """ Generates neighbour lists for all samples in dataset using prior term information @@ -193,6 +200,8 @@ def build_neighborlists( mol_num_batches : int unused in this function present to allow the use of the same .yaml config for process_raw_dataset and build_neighborlists + atoms_batch_size : Optional[int], optional + Batch size for atoms when processing large molecules (default: None). """ dataset = RawDataset(dataset_name, names, tag) for samples in tqdm(dataset, f"Building NL for {dataset_name} dataset..."): From fda2f248fe0f5b639fa26e2e2b1d73fb6baf6848 Mon Sep 17 00:00:00 2001 From: EdoardoRolando Date: Fri, 20 Jun 2025 18:36:22 +0200 Subject: [PATCH 2/3] Added documentation --- examples/README.md | 4 ++ input_generator/raw_dataset.py | 2 + input_generator/utils.py | 12 +++--- scripts/gen_input_data.py | 68 +++++++++++++++++----------------- 4 files changed, 47 insertions(+), 39 deletions(-) diff --git a/examples/README.md b/examples/README.md index c06da1f..0f9d189 100644 --- a/examples/README.md +++ b/examples/README.md @@ -37,6 +37,10 @@ Note if you are using a custom dataset: If your program gets killed after the loading of the all-atom data succeeded (tqdm bar finished) but before `process_raw_dataset` saved the CG output, try to set `batch_size` in your `trpcage.yaml` file. This will batch the matrix multiplication between atomistic coordinates/forces, which is the most memory-consuming part of the coarse-graining at this stage. +##### Batch processing for large molecules: + +If the dataset loads into memory successfully (the tqdm bar completes), but the program fails before saving the CG output, consider setting atoms_batch_size in your trpcage.yaml file. This optional parameter specifies the batch size for processing atoms in large molecules. When set, constraints among atoms for coordinate and force mappings will be computed in batches of this size to reduce memory usage. If atoms_batch_size is larger than the total number of atoms in the molecule, all atoms will be processed at once (the default behavior). + ##### Batch processing for large datasets: Should your dataset be too big to be loaded into memory at once (the tqdm bar doesn't finish before it fails), you can set the `mol_num_batches` in your `trpcage.yaml` file as well as your `trpcage_stats.yaml`, `trpcage_delta_forces.yaml` and `trpcage_packaging.yaml` file. This will seperate the trajectories in your dataset into `mol_num_batches` chunks that will be treated as separate molecules for the coarse-graining and statistics computing stages (see 2 below) and the statistics of the different batches will be automatically accumulated to get only one prior object in the end. Note that in this case, the force map will be only computed on the first batch and re-used for all subsequent batches to ensure consistency in the case of optimized force maps. diff --git a/input_generator/raw_dataset.py b/input_generator/raw_dataset.py index 5edb38f..e95d5bf 100644 --- a/input_generator/raw_dataset.py +++ b/input_generator/raw_dataset.py @@ -317,6 +317,8 @@ def process_coords_forces( Striding to use for force projection results batch_size: Batching the coords and forces projection to CG + atoms_batch_size: + Batch size for processing atoms when inferring constrained atoms Returns ------- diff --git a/input_generator/utils.py b/input_generator/utils.py index 5b31317..823fe53 100644 --- a/input_generator/utils.py +++ b/input_generator/utils.py @@ -263,7 +263,7 @@ def slice_coord_forces( Optional length of batch in which divide the AA mapping of coords and forces to CG ones atoms_batch_size: - Optional length of batch in which divide the atoms in coords to compute the pairwise constraints. + Optional batch size for dividing atoms in coordinates to estimate pairwise constraints Returns ------- @@ -295,11 +295,11 @@ def slice_coord_forces( # Cross-batch constraints # To significantly reduce computational cost, we assume residues are ordered in the structure. - # Therefore, we only compute constraints between consecutive batches, rather than all pairs of batches. - # For even greater efficiency, one could restrict this to just the first and last (e.g., 30) atoms of each batch, - # which scales as O(1) basically, but computing all pairs between consecutive batches is generally still efficient. - # This tecnically can be extended to the case with no batches (smaller molecules), assuming again ordered resiues, - # treating all molecules at the same way and getting rid of the atoms_batch_size parameters. + # Therefore, constraints are computed only between consecutive batches rather than all pairs of batches. + # For even greater efficiency, this could be further limited to just the first and last (e.g., 30) atoms of each batch, + # which scales approximately as O(1). However, computing all pairs between consecutive batches is generally still efficient. + # This approach can also be extended to the case with no batching (for smaller molecules), + # again assuming ordered residues, treating all molecules uniformly and eliminating the need for the atoms_batch_size parameter. for i in range(len(batches) - 1): b1 = batches[i] b2 = batches[i + 1] diff --git a/scripts/gen_input_data.py b/scripts/gen_input_data.py index ca86562..bff4442 100644 --- a/scripts/gen_input_data.py +++ b/scripts/gen_input_data.py @@ -40,52 +40,53 @@ def process_raw_dataset( atoms_batch_size: Optional[int] = None, ): """ - Processes a raw dataset by applying coarse-grained (CG) mapping to atomic coordinates and forces, - using the provided topology and mapping strategies. + Applies coarse-grained mapping to coordinates and forces using input sample + topology and specified mapping strategies Parameters ---------- dataset_name : str - Identifier for the dataset. + Name given to specific dataset names : List[str] - List of molecule or sample names to process. + List of sample names sample_loader : DatasetLoader - Loader object for retrieving trajectories and data. + Loader object defined for specific dataset raw_data_dir : str - Directory containing raw coordinate and force files. + Path to coordinate and force files tag : str - Label to append to all output files generated from this dataset. + Label given to all output files produced from dataset pdb_template_fn : str - Path to the template PDB file for atomistic topology. + Template file location of atomistic structure to be used for topology save_dir : str - Directory where processed outputs will be saved. + Path to directory in which output will be saved cg_atoms : List[str] - Atom names to retain in the coarse-grained representation. + List of atom names to preserve in coarse-grained resolution embedding_map : CGEmbeddingMap - Object defining the mapping from atomistic to CG representations. + Mapping object embedding_func : Callable - Function to apply the CG mapping. + Function which will be used to apply CG mapping skip_residues : List[str] - Residues to exclude from processing (can be empty). + List of residues to skip, can be None cg_mapping_strategy : str - Strategy for mapping coordinates and forces (e.g., "slice_aggregate", "slice_optimize"). - stride : int, optional - Interval for subsampling loaded data (default: 1). - force_stride : int, optional - Interval for subsampling forces during mapping (default: 100). - filter_cis : Optional[bool], optional - If True, filters out frames with cis-configurations (default: False). - batch_size : Optional[int], optional - Number of frames to process in each batch to reduce memory usage (default: None). - mol_num_batches : Optional[int], optional - Number of batches to split each molecule's data into (default: 1). - atoms_batch_size : Optional[int], optional - Batch size for atoms when processing large molecules (default: None). - - Returns - ------- - None - Saves processed CG data and mapping files to the specified directory. + Strategy to use for coordinate and force mappings; + currently only "slice_aggregate" and "slice_optimize" are implemented + stride : int + Interval by which to stride loaded data + force_stride : int + stride for inferring the force maps in aggforce + filter_cis : bool + if True, frames with cis-configurations will be filtered out from the dataset + batch_size : int + Optional size in which performing batches of AA mapping to CG, to avoid + memory overhead in large AA dataset + mol_num_batches : int + If greater than 1, will save each molecule data into the specified number of batches + that will be treated as different samples + atoms_batch_size : int + Optional batch size for processing atoms in large molecules (default: None). If specified, constraints among atoms for coordinate and + force mappings (as defined by `cg_mapping_strategy`) will be computed in batches of this size to reduce memory usage. If + `atoms_batch_size` exceeds the total number of atoms in the molecule, all atoms will be processed at once (default behaviour). + """ dataset = RawDataset(dataset_name, names, tag, n_batches=mol_num_batches) for samples in tqdm(dataset, f"Processing CG data for {dataset_name} dataset..."): @@ -200,8 +201,9 @@ def build_neighborlists( mol_num_batches : int unused in this function present to allow the use of the same .yaml config for process_raw_dataset and build_neighborlists - atoms_batch_size : Optional[int], optional - Batch size for atoms when processing large molecules (default: None). + atoms_batch_size : int + unused in this function + present to allow the use of the same .yaml config for process_raw_dataset and build_neighborlists """ dataset = RawDataset(dataset_name, names, tag) for samples in tqdm(dataset, f"Building NL for {dataset_name} dataset..."): From 9edd5a6bee6f13bcd506c10101018d4d4b175686 Mon Sep 17 00:00:00 2001 From: EdoardoRolando Date: Fri, 20 Jun 2025 18:50:20 +0200 Subject: [PATCH 3/3] Improved documentation --- examples/README.md | 2 +- scripts/gen_input_data.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/README.md b/examples/README.md index 0f9d189..2bba696 100644 --- a/examples/README.md +++ b/examples/README.md @@ -39,7 +39,7 @@ If your program gets killed after the loading of the all-atom data succeeded (tq ##### Batch processing for large molecules: -If the dataset loads into memory successfully (the tqdm bar completes), but the program fails before saving the CG output, consider setting atoms_batch_size in your trpcage.yaml file. This optional parameter specifies the batch size for processing atoms in large molecules. When set, constraints among atoms for coordinate and force mappings will be computed in batches of this size to reduce memory usage. If atoms_batch_size is larger than the total number of atoms in the molecule, all atoms will be processed at once (the default behavior). +If the dataset loads into memory successfully (the tqdm bar completes), but the program fails before saving the CG output, consider setting atoms_batch_size in your trpcage.yaml file. This optional parameter specifies the batch size for processing atoms in large molecules. When set, constraints among atoms for coordinate and force mappings will be computed in batches of this size to reduce memory usage. To improve computational efficiency, it is assumed that the molecular structures have ordered residues. If atoms_batch_size is larger than the total number of atoms in the molecule, all atoms will be processed at once (the default behavior). ##### Batch processing for large datasets: diff --git a/scripts/gen_input_data.py b/scripts/gen_input_data.py index bff4442..2e74ac3 100644 --- a/scripts/gen_input_data.py +++ b/scripts/gen_input_data.py @@ -82,10 +82,11 @@ def process_raw_dataset( mol_num_batches : int If greater than 1, will save each molecule data into the specified number of batches that will be treated as different samples - atoms_batch_size : int - Optional batch size for processing atoms in large molecules (default: None). If specified, constraints among atoms for coordinate and - force mappings (as defined by `cg_mapping_strategy`) will be computed in batches of this size to reduce memory usage. If - `atoms_batch_size` exceeds the total number of atoms in the molecule, all atoms will be processed at once (default behaviour). + atoms_batch_size : int, optional + Optional batch size for processing atoms in large molecules (default: None). If specified, constraints among atoms for coordinate and + force mappings (as defined by `cg_mapping_strategy`) will be computed in batches of this size. To significantly improve + computational efficiency, it is assumed that structures have ordered residues. If `atoms_batch_size` exceeds the total number of atoms + in the molecule, all atoms will be processed at once (default behavior). """ dataset = RawDataset(dataset_name, names, tag, n_batches=mol_num_batches)