Added parameter atoms_batch_size

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. 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:
 
 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.

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
@@ -316,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
         -------
@@ -341,7 +344,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

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 batch size for dividing atoms in coordinates to estimate 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, 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]
+            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

scripts/gen_input_data.py

@@ -37,6 +37,7 @@ 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
@@ -81,6 +82,12 @@ 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
+        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)
     for samples in tqdm(dataset, f"Processing CG data for {dataset_name} dataset..."):
@@ -117,6 +124,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 +155,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 +202,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 : 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..."):