fix(io): Correct handling of residue indices in bio topology conversion

Author: TKanX

src/io/util.rs

@@ -186,7 +186,14 @@ pub fn to_bio_topology(system: &System) -> Result<bf::Topology, ConversionError>
     type ResidueKey = (String, i32, Option<char>);
     let mut chains: IndexMap<String, IndexMap<ResidueKey, bf::Residue>> = IndexMap::new();
 
-    for (atom, info) in system.atoms.iter().zip(metadata.atom_info.iter()) {
+    let mut residue_members: IndexMap<ResidueKey, Vec<usize>> = IndexMap::new();
+
+    for (i, (atom, info)) in system
+        .atoms
+        .iter()
+        .zip(metadata.atom_info.iter())
+        .enumerate()
+    {
         let bio_atom = bf::Atom::new(
             &info.atom_name,
             convert_element_to_bf(atom.element)?,
@@ -196,7 +203,7 @@ pub fn to_bio_topology(system: &System) -> Result<bf::Topology, ConversionError>
         let residue_key = (info.chain_id.clone(), info.residue_id, info.insertion_code);
         let residues = chains.entry(info.chain_id.clone()).or_default();
 
-        let residue = match residues.entry(residue_key) {
+        let residue = match residues.entry(residue_key.clone()) {
             Entry::Occupied(e) => e.into_mut(),
             Entry::Vacant(e) => {
                 let mut res = bf::Residue::new(
@@ -212,6 +219,7 @@ pub fn to_bio_topology(system: &System) -> Result<bf::Topology, ConversionError>
         };
 
         residue.add_atom(bio_atom);
+        residue_members.entry(residue_key).or_default().push(i);
     }
 
     let mut bio_struct = bf::Structure::new();
@@ -225,13 +233,18 @@ pub fn to_bio_topology(system: &System) -> Result<bf::Topology, ConversionError>
         bio_struct.add_chain(chain);
     }
 
+    let mut old_to_new = vec![0usize; system.atoms.len()];
+    for (flat, &orig) in residue_members.values().flatten().enumerate() {
+        old_to_new[orig] = flat;
+    }
+
     let bio_bonds = system
         .bonds
         .iter()
         .map(|bond| {
             Ok(bf::Bond::new(
-                bond.i,
-                bond.j,
+                old_to_new[bond.i],
+                old_to_new[bond.j],
                 convert_bond_order_to_bf(bond.order)?,
             ))
         })
@@ -572,6 +585,43 @@ mod tests {
         assert_eq!(roundtrip.bio_metadata.as_ref(), Some(&metadata));
     }
 
+    #[test]
+    fn to_bio_topology_remaps_bonds_for_non_contiguous_atoms() {
+        let atoms = vec![
+            Atom::new(Element::C, [0.0, 0.0, 0.0]),
+            Atom::new(Element::N, [3.0, 0.0, 0.0]),
+            Atom::new(Element::C, [0.0, 1.5, 0.0]),
+        ];
+        let metadata = BioMetadata {
+            atom_info: vec![
+                AtomResidueInfo::builder("CA", "ALA", 1, "A").build(),
+                AtomResidueInfo::builder("N", "ALA", 2, "A").build(),
+                AtomResidueInfo::builder("CB", "ALA", 1, "A").build(),
+            ],
+            target_ph: None,
+        };
+        let bonds = vec![Bond::new(0, 2, BondOrder::Single)];
+        let system = System {
+            atoms,
+            bonds,
+            box_vectors: None,
+            bio_metadata: Some(metadata),
+        };
+
+        let bio_topo = to_bio_topology(&system).expect("conversion should succeed");
+
+        let bio_bonds = bio_topo.bonds();
+        assert_eq!(bio_bonds.len(), 1);
+        assert_eq!(
+            bio_bonds[0].a1_idx, 0,
+            "bond endpoint a should be flat 0 (CA)"
+        );
+        assert_eq!(
+            bio_bonds[0].a2_idx, 1,
+            "bond endpoint b should be flat 1 (CB)"
+        );
+    }
+
     #[test]
     fn converts_clean_config_fields() {
         let mut remove_res = HashSet::new();