center coordinates at 0

openfold/data/mmcif_parsing.py

@@ -429,7 +429,7 @@ def _is_set(data: str) -> bool:
 
 
 def get_atom_coords(
-    mmcif_object: MmcifObject, chain_id: str
+    mmcif_object: MmcifObject, chain_id: str, zero_center: bool = True
 ) -> Tuple[np.ndarray, np.ndarray]:
     # Locate the right chain
     chains = list(mmcif_object.structure.get_chains())
@@ -474,4 +474,9 @@ def get_atom_coords(
         all_atom_positions[res_index] = pos
         all_atom_mask[res_index] = mask
 
+    if zero_center:
+        binary_mask = all_atom_mask.astype(np.bool)
+        translation_vec = all_atom_positions[binary_mask].mean(axis=0)
+        all_atom_positions[binary_mask] -= translation_vec
+
     return all_atom_positions, all_atom_mask