From 685036772eca793a4ce97f0f1dd6c16cd50d7299 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 18:27:14 -0700 Subject: [PATCH 01/20] chore(package): Add slotmap dependency for enhanced data structure support --- Cargo.lock | 16 ++++++++++++++++ crates/scream-core/Cargo.toml | 1 + 2 files changed, 17 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index fefdab0c..bec7995f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -445,6 +445,7 @@ version = "0.1.0-alpha.1" dependencies = [ "bitflags", "nalgebra", + "slotmap", "thiserror", ] @@ -461,6 +462,15 @@ dependencies = [ "wide", ] +[[package]] +name = "slotmap" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbff4acf519f630b3a3ddcfaea6c06b42174d9a44bc70c620e9ed1649d58b82a" +dependencies = [ + "version_check", +] + [[package]] name = "strsim" version = "0.11.1" @@ -534,6 +544,12 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + [[package]] name = "wasm-bindgen" version = "0.2.100" diff --git a/crates/scream-core/Cargo.toml b/crates/scream-core/Cargo.toml index aebf3e54..31eb4e87 100644 --- a/crates/scream-core/Cargo.toml +++ b/crates/scream-core/Cargo.toml @@ -10,4 +10,5 @@ edition.workspace = true [dependencies] bitflags = "2.9.1" nalgebra = "0.33.2" +slotmap = "1.0.7" thiserror = "2.0.12" From cf36a4918eb81e1fe5a68f1414c03de0b87d775b Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 18:50:32 -0700 Subject: [PATCH 02/20] feat(core): Introduce new key types for AtomId, ResidueId, and ChainId using slotmap --- crates/scream-core/src/core/models/ids.rs | 7 +++++++ crates/scream-core/src/core/models/mod.rs | 1 + 2 files changed, 8 insertions(+) create mode 100644 crates/scream-core/src/core/models/ids.rs diff --git a/crates/scream-core/src/core/models/ids.rs b/crates/scream-core/src/core/models/ids.rs new file mode 100644 index 00000000..378e9328 --- /dev/null +++ b/crates/scream-core/src/core/models/ids.rs @@ -0,0 +1,7 @@ +use slotmap::new_key_type; + +new_key_type! { + pub struct AtomId; + pub struct ResidueId; + pub struct ChainId; +} diff --git a/crates/scream-core/src/core/models/mod.rs b/crates/scream-core/src/core/models/mod.rs index fa400a17..d8a24ce3 100644 --- a/crates/scream-core/src/core/models/mod.rs +++ b/crates/scream-core/src/core/models/mod.rs @@ -1,6 +1,7 @@ pub mod atom; pub mod builder; pub mod chain; +pub mod ids; pub mod residue; pub mod system; pub mod topology; From 116d8e390fdfb43e3722a832f983518ffdf6c33f Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 19:02:36 -0700 Subject: [PATCH 03/20] refactor(core): Simplify Atom struct by removing unused fields and updating residue_id type --- crates/scream-core/src/core/models/atom.rs | 109 ++++----------------- 1 file changed, 17 insertions(+), 92 deletions(-) diff --git a/crates/scream-core/src/core/models/atom.rs b/crates/scream-core/src/core/models/atom.rs index d9e081a7..feb69ad6 100644 --- a/crates/scream-core/src/core/models/atom.rs +++ b/crates/scream-core/src/core/models/atom.rs @@ -1,3 +1,4 @@ +use super::ids::ResidueId; use bitflags::bitflags; use nalgebra::Point3; @@ -14,12 +15,9 @@ bitflags! { #[derive(Debug, Clone, PartialEq)] pub struct Atom { // --- Identity & Topology --- - pub index: usize, // Global index in the system's atom list - pub serial: usize, // Atom serial number from source file - pub name: String, // Atom name (e.g., "CA", "N", "O") - pub res_name: String, // Residue name (e.g., "ALA", "GLY") - pub res_id: isize, // Residue sequence number from source file - pub chain_id: char, // Chain identifier (e.g., 'A', 'B') + pub serial: usize, // Atom serial number from source file + pub name: String, // Atom name (e.g., "CA", "N", "O") + pub residue_id: ResidueId, // Stable ID of the parent residue // --- Physicochemical Properties --- pub force_field_type: String, // Force field atom type (e.g., "C.3", "N.2") @@ -36,93 +34,20 @@ pub struct Atom { pub hbond_type_id: i8, // Hydrogen bond type identifier } -#[cfg(test)] -mod tests { - use super::*; - use nalgebra::Point3; - - #[test] - fn atom_flags_set_and_check_individual_flags() { - let mut flags = AtomFlags::empty(); - flags.insert(AtomFlags::IS_FIXED_ROLE); - assert!(flags.contains(AtomFlags::IS_FIXED_ROLE)); - assert!(!flags.contains(AtomFlags::IS_TREATED_AS_FIXED)); - flags.insert(AtomFlags::IS_TREATED_AS_FIXED); - assert!(flags.contains(AtomFlags::IS_TREATED_AS_FIXED)); - } - - #[test] - fn atom_flags_combined_and_removed() { - let mut flags = AtomFlags::IS_VISIBLE_INTERACTION | AtomFlags::IS_VISIBLE_LATTICE; - assert!(flags.contains(AtomFlags::IS_VISIBLE_INTERACTION)); - assert!(flags.contains(AtomFlags::IS_VISIBLE_LATTICE)); - flags.remove(AtomFlags::IS_VISIBLE_INTERACTION); - assert!(!flags.contains(AtomFlags::IS_VISIBLE_INTERACTION)); - assert!(flags.contains(AtomFlags::IS_VISIBLE_LATTICE)); - } - - #[test] - fn atom_flags_default_is_empty() { - let flags = AtomFlags::default(); - assert!(!flags.contains(AtomFlags::IS_FIXED_ROLE)); - assert_eq!(flags.bits(), 0); - } - - #[test] - fn atom_struct_fields_are_set_correctly() { - let atom = Atom { - index: 0, - serial: 42, - name: "CA".to_string(), - res_name: "GLY".to_string(), - res_id: 5, - chain_id: 'A', - force_field_type: "C.3".to_string(), - partial_charge: -0.123, - position: Point3::new(1.0, 2.0, 3.0), - flags: AtomFlags::IS_FIXED_ROLE | AtomFlags::IS_VISIBLE_LATTICE, - delta: 0.5, - vdw_radius: 1.7, - vdw_well_depth: 0.2, - hbond_type_id: 3, - }; - assert_eq!(atom.index, 0); - assert_eq!(atom.serial, 42); - assert_eq!(atom.name, "CA"); - assert_eq!(atom.res_name, "GLY"); - assert_eq!(atom.res_id, 5); - assert_eq!(atom.chain_id, 'A'); - assert_eq!(atom.force_field_type, "C.3"); - assert_eq!(atom.partial_charge, -0.123); - assert_eq!(atom.position, Point3::new(1.0, 2.0, 3.0)); - assert!(atom.flags.contains(AtomFlags::IS_FIXED_ROLE)); - assert!(atom.flags.contains(AtomFlags::IS_VISIBLE_LATTICE)); - assert!(!atom.flags.contains(AtomFlags::IS_TREATED_AS_FIXED)); - assert_eq!(atom.delta, 0.5); - assert_eq!(atom.vdw_radius, 1.7); - assert_eq!(atom.vdw_well_depth, 0.2); - assert_eq!(atom.hbond_type_id, 3); - } - - #[test] - fn atom_struct_equality_and_clone() { - let atom1 = Atom { - index: 1, - serial: 2, - name: "N".to_string(), - res_name: "ALA".to_string(), - res_id: 10, - chain_id: 'B', - force_field_type: "N.3".to_string(), +impl Atom { + pub fn new(serial: usize, name: &str, residue_id: ResidueId, position: Point3) -> Self { + Self { + serial, + name: name.to_string(), + residue_id, + position, + force_field_type: "".to_string(), partial_charge: 0.0, - position: Point3::new(0.0, 0.0, 0.0), - flags: AtomFlags::empty(), + flags: AtomFlags::default(), delta: 0.0, - vdw_radius: 1.0, - vdw_well_depth: 0.1, - hbond_type_id: 0, - }; - let atom2 = atom1.clone(); - assert_eq!(atom1, atom2); + vdw_radius: 0.0, + vdw_well_depth: 0.0, + hbond_type_id: -1, + } } } From 00e8b71347353430abe6dfd72f0fd84abdbdfdb1 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 19:10:42 -0700 Subject: [PATCH 04/20] test(core): Add unit tests for Atom struct and its flags operations --- crates/scream-core/src/core/models/atom.rs | 44 ++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/crates/scream-core/src/core/models/atom.rs b/crates/scream-core/src/core/models/atom.rs index feb69ad6..8275554b 100644 --- a/crates/scream-core/src/core/models/atom.rs +++ b/crates/scream-core/src/core/models/atom.rs @@ -51,3 +51,47 @@ impl Atom { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::models::ids::ResidueId; + use nalgebra::Point3; + + #[test] + fn new_atom_has_expected_default_fields() { + let residue_id = ResidueId::default(); + let atom = Atom::new(42, "CA", residue_id, Point3::new(1.0, 2.0, 3.0)); + assert_eq!(atom.serial, 42); + assert_eq!(atom.name, "CA"); + assert_eq!(atom.residue_id, residue_id); + assert_eq!(atom.position, Point3::new(1.0, 2.0, 3.0)); + assert_eq!(atom.force_field_type, ""); + assert_eq!(atom.partial_charge, 0.0); + assert_eq!(atom.flags, AtomFlags::default()); + assert_eq!(atom.delta, 0.0); + assert_eq!(atom.vdw_radius, 0.0); + assert_eq!(atom.vdw_well_depth, 0.0); + assert_eq!(atom.hbond_type_id, -1); + } + + #[test] + fn atom_flags_bitwise_operations_work() { + let mut flags = AtomFlags::IS_FIXED_ROLE | AtomFlags::IS_VISIBLE_LATTICE; + assert!(flags.contains(AtomFlags::IS_FIXED_ROLE)); + assert!(flags.contains(AtomFlags::IS_VISIBLE_LATTICE)); + assert!(!flags.contains(AtomFlags::IS_VISIBLE_INTERACTION)); + flags.insert(AtomFlags::IS_VISIBLE_INTERACTION); + assert!(flags.contains(AtomFlags::IS_VISIBLE_INTERACTION)); + flags.remove(AtomFlags::IS_FIXED_ROLE); + assert!(!flags.contains(AtomFlags::IS_FIXED_ROLE)); + } + + #[test] + fn atom_equality_and_clone_works() { + let residue_id = ResidueId::default(); + let atom1 = Atom::new(1, "N", residue_id, Point3::new(0.0, 0.0, 0.0)); + let atom2 = atom1.clone(); + assert_eq!(atom1, atom2); + } +} From bfd30487944c7325467cfb06e9e2d7c5ed38aea5 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 19:14:53 -0700 Subject: [PATCH 05/20] refactor(core): Update Residue struct to include chain_id and change atom storage to AtomId --- crates/scream-core/src/core/models/atom.rs | 2 +- crates/scream-core/src/core/models/residue.rs | 79 +++++-------------- 2 files changed, 21 insertions(+), 60 deletions(-) diff --git a/crates/scream-core/src/core/models/atom.rs b/crates/scream-core/src/core/models/atom.rs index 8275554b..8194917d 100644 --- a/crates/scream-core/src/core/models/atom.rs +++ b/crates/scream-core/src/core/models/atom.rs @@ -17,7 +17,7 @@ pub struct Atom { // --- Identity & Topology --- pub serial: usize, // Atom serial number from source file pub name: String, // Atom name (e.g., "CA", "N", "O") - pub residue_id: ResidueId, // Stable ID of the parent residue + pub residue_id: ResidueId, // ID of the parent residue // --- Physicochemical Properties --- pub force_field_type: String, // Force field atom type (e.g., "C.3", "N.2") diff --git a/crates/scream-core/src/core/models/residue.rs b/crates/scream-core/src/core/models/residue.rs index 6abdcf7b..21a893a2 100644 --- a/crates/scream-core/src/core/models/residue.rs +++ b/crates/scream-core/src/core/models/residue.rs @@ -1,80 +1,41 @@ +use super::ids::{AtomId, ChainId}; use std::collections::HashMap; -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct Residue { - pub id: isize, // Residue sequence number from source file - pub name: String, // Name of the residue (e.g., "ALA", "GLY") - pub atom_indices: Vec, // Indices of atoms belonging to this residue - atom_name_map: HashMap, // Map from atom name to its global index + pub id: isize, // Residue sequence number from source file + pub name: String, // Name of the residue (e.g., "ALA", "GLY") + pub chain_id: ChainId, // ID of the parent chain + pub(crate) atoms: Vec, // Indices of atoms belonging to this residue + atom_name_map: HashMap, // Map from atom name to its stable ID } impl Residue { - pub(crate) fn new(id: isize, name: &str) -> Self { + pub(crate) fn new(id: isize, name: &str, chain_id: ChainId) -> Self { Self { id, name: name.to_string(), - atom_indices: Vec::new(), + chain_id, + atoms: Vec::new(), atom_name_map: HashMap::new(), } } - pub(crate) fn add_atom(&mut self, atom_name: &str, atom_idx: usize) { - self.atom_indices.push(atom_idx); - self.atom_name_map.insert(atom_name.to_string(), atom_idx); + pub(crate) fn add_atom(&mut self, atom_name: &str, atom_id: AtomId) { + self.atoms.push(atom_id); + self.atom_name_map.insert(atom_name.to_string(), atom_id); } - pub fn get_atom_index_by_name(&self, name: &str) -> Option { - self.atom_name_map.get(name).copied() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn creates_residue_with_correct_id_and_name() { - let residue = Residue::new(42, "GLY"); - assert_eq!(residue.id, 42); - assert_eq!(residue.name, "GLY"); - assert!(residue.atom_indices.is_empty()); - } - - #[test] - fn adds_atom_and_retrieves_index_by_name() { - let mut residue = Residue::new(1, "ALA"); - residue.add_atom("CA", 5); - residue.add_atom("CB", 6); - - assert_eq!(residue.atom_indices, vec![5, 6]); - assert_eq!(residue.get_atom_index_by_name("CA"), Some(5)); - assert_eq!(residue.get_atom_index_by_name("CB"), Some(6)); + pub(crate) fn remove_atom(&mut self, atom_name: &str, atom_id: AtomId) { + self.atoms.retain(|&id| id != atom_id); + self.atom_name_map.remove(atom_name); } - #[test] - fn get_atom_index_by_name_returns_none_for_missing_atom() { - let mut residue = Residue::new(2, "SER"); - residue.add_atom("OG", 10); - - assert_eq!(residue.get_atom_index_by_name("CA"), None); - assert_eq!(residue.get_atom_index_by_name("OG"), Some(10)); + pub fn atoms(&self) -> &[AtomId] { + &self.atoms } - #[test] - fn add_atom_overwrites_existing_atom_name() { - let mut residue = Residue::new(3, "THR"); - residue.add_atom("CG2", 20); - residue.add_atom("CG2", 21); - - assert_eq!(residue.atom_indices, vec![20, 21]); - assert_eq!(residue.get_atom_index_by_name("CG2"), Some(21)); - } - - #[test] - fn add_atom_with_empty_name_and_retrieve() { - let mut residue = Residue::new(4, "VAL"); - residue.add_atom("", 30); - - assert_eq!(residue.get_atom_index_by_name(""), Some(30)); + pub fn get_atom_id_by_name(&self, name: &str) -> Option { + self.atom_name_map.get(name).copied() } } From 7614a6b03aa81a728155f237c220b3c61b6634a9 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 19:27:33 -0700 Subject: [PATCH 06/20] test(core): Add unit tests for Residue struct methods --- crates/scream-core/src/core/models/residue.rs | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/crates/scream-core/src/core/models/residue.rs b/crates/scream-core/src/core/models/residue.rs index 21a893a2..87675034 100644 --- a/crates/scream-core/src/core/models/residue.rs +++ b/crates/scream-core/src/core/models/residue.rs @@ -39,3 +39,85 @@ impl Residue { self.atom_name_map.get(name).copied() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::models::ids::{AtomId, ChainId}; + use slotmap::KeyData; + use std::collections::HashSet; + + fn dummy_atom_id(n: u64) -> AtomId { + AtomId::from(KeyData::from_ffi(n)) + } + + fn dummy_chain_id(n: u64) -> ChainId { + ChainId::from(KeyData::from_ffi(n)) + } + + #[test] + fn new_residue_initializes_fields_correctly() { + let chain_id = dummy_chain_id(1); + let residue = Residue::new(10, "GLY", chain_id); + assert_eq!(residue.id, 10); + assert_eq!(residue.name, "GLY"); + assert_eq!(residue.chain_id, chain_id); + assert!(residue.atoms().is_empty()); + assert!(residue.get_atom_id_by_name("CA").is_none()); + } + + #[test] + fn add_atom_adds_atom_and_maps_name() { + let chain_id = dummy_chain_id(2); + let mut residue = Residue::new(5, "ALA", chain_id); + let atom_id = dummy_atom_id(42); + residue.add_atom("CA", atom_id); + assert_eq!(residue.atoms(), &[atom_id]); + assert_eq!(residue.get_atom_id_by_name("CA"), Some(atom_id)); + } + + #[test] + fn add_atom_allows_multiple_atoms_with_different_names() { + let chain_id = dummy_chain_id(3); + let mut residue = Residue::new(7, "SER", chain_id); + let atom_id1 = dummy_atom_id(1); + let atom_id2 = dummy_atom_id(2); + residue.add_atom("CA", atom_id1); + residue.add_atom("CB", atom_id2); + let atom_set: HashSet<_> = residue.atoms().iter().copied().collect(); + assert!(atom_set.contains(&atom_id1)); + assert!(atom_set.contains(&atom_id2)); + assert_eq!(residue.get_atom_id_by_name("CA"), Some(atom_id1)); + assert_eq!(residue.get_atom_id_by_name("CB"), Some(atom_id2)); + } + + #[test] + fn remove_atom_removes_atom_and_name_mapping() { + let chain_id = dummy_chain_id(4); + let mut residue = Residue::new(8, "THR", chain_id); + let atom_id = dummy_atom_id(100); + residue.add_atom("OG1", atom_id); + residue.remove_atom("OG1", atom_id); + assert!(residue.atoms().is_empty()); + assert!(residue.get_atom_id_by_name("OG1").is_none()); + } + + #[test] + fn remove_atom_does_nothing_if_atom_not_present() { + let chain_id = dummy_chain_id(5); + let mut residue = Residue::new(9, "VAL", chain_id); + let atom_id = dummy_atom_id(200); + residue.add_atom("CG1", atom_id); + residue.remove_atom("CG2", dummy_atom_id(201)); + assert_eq!(residue.atoms(), &[atom_id]); + assert_eq!(residue.get_atom_id_by_name("CG1"), Some(atom_id)); + } + + #[test] + fn get_atom_id_by_name_returns_none_for_unknown_name() { + let chain_id = dummy_chain_id(6); + let mut residue = Residue::new(11, "LEU", chain_id); + residue.add_atom("CD1", dummy_atom_id(300)); + assert!(residue.get_atom_id_by_name("CD2").is_none()); + } +} From ac616358733253565aa2908c93d39b97a29ac966 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 19:34:27 -0700 Subject: [PATCH 07/20] refactor(core): Update Chain struct to use ResidueId and simplify residue management --- crates/scream-core/src/core/models/chain.rs | 101 ++------------------ 1 file changed, 6 insertions(+), 95 deletions(-) diff --git a/crates/scream-core/src/core/models/chain.rs b/crates/scream-core/src/core/models/chain.rs index dfd2d1e1..7b49df1d 100644 --- a/crates/scream-core/src/core/models/chain.rs +++ b/crates/scream-core/src/core/models/chain.rs @@ -1,5 +1,4 @@ -use super::residue::Residue; -use std::collections::HashMap; +use super::ids::ResidueId; use std::fmt; use std::str::FromStr; use thiserror::Error; @@ -49,12 +48,11 @@ impl fmt::Display for ChainType { } } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct Chain { - pub id: char, // Chain identifier (e.g., 'A', 'B') - pub chain_type: ChainType, // Type of the chain - pub(crate) residues: Vec, // List of residues in the chain - pub(crate) residue_map: HashMap, // Map from residue ID to its index in the `residues` vector + pub id: char, // Chain identifier (e.g., 'A', 'B') + pub chain_type: ChainType, // Type of the chain + pub(crate) residues: Vec, // Ordered list of residue IDs belonging to this chain } impl Chain { @@ -63,97 +61,10 @@ impl Chain { id, chain_type, residues: Vec::new(), - residue_map: HashMap::new(), } } - pub fn residues(&self) -> &[Residue] { + pub fn residues(&self) -> &[ResidueId] { &self.residues } - - pub fn get_residue(&self, index: usize) -> Option<&Residue> { - self.residues.get(index) - } - - pub fn get_residue_by_id(&self, id: isize) -> Option<&Residue> { - self.residue_map - .get(&id) - .and_then(|&index| self.residues.get(index)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::core::models::residue::Residue; - - #[test] - fn parses_valid_chain_types() { - assert_eq!(ChainType::from_str("protein").unwrap(), ChainType::Protein); - assert_eq!(ChainType::from_str("dna").unwrap(), ChainType::DNA); - assert_eq!(ChainType::from_str("rna").unwrap(), ChainType::RNA); - assert_eq!(ChainType::from_str("ligand").unwrap(), ChainType::Ligand); - assert_eq!(ChainType::from_str("water").unwrap(), ChainType::Water); - assert_eq!(ChainType::from_str("other").unwrap(), ChainType::Other); - } - - #[test] - fn parses_chain_type_case_insensitive() { - assert_eq!(ChainType::from_str("PrOtEiN").unwrap(), ChainType::Protein); - assert_eq!(ChainType::from_str("DNA").unwrap(), ChainType::DNA); - assert_eq!(ChainType::from_str("rNa").unwrap(), ChainType::RNA); - assert_eq!(ChainType::from_str("LiGaNd").unwrap(), ChainType::Ligand); - assert_eq!(ChainType::from_str("WATER").unwrap(), ChainType::Water); - } - - #[test] - fn parses_unknown_chain_type_as_other() { - assert_eq!( - ChainType::from_str("carbohydrate").unwrap(), - ChainType::Other - ); - assert_eq!(ChainType::from_str("").unwrap(), ChainType::Other); - assert_eq!(ChainType::from_str("123").unwrap(), ChainType::Other); - } - - #[test] - fn displays_chain_type_correctly() { - assert_eq!(ChainType::Protein.to_string(), "Protein"); - assert_eq!(ChainType::DNA.to_string(), "DNA"); - assert_eq!(ChainType::RNA.to_string(), "RNA"); - assert_eq!(ChainType::Ligand.to_string(), "Ligand"); - assert_eq!(ChainType::Water.to_string(), "Water"); - assert_eq!(ChainType::Other.to_string(), "Other"); - } - - #[test] - fn creates_chain_with_correct_id_and_type() { - let chain = Chain::new('A', ChainType::Protein); - assert_eq!(chain.id, 'A'); - assert_eq!(chain.chain_type, ChainType::Protein); - assert!(chain.residues().is_empty()); - } - - #[test] - fn get_residue_returns_none_for_empty_chain() { - let chain = Chain::new('B', ChainType::DNA); - assert_eq!(chain.get_residue(0), None); - } - - #[test] - fn get_residue_by_id_returns_none_for_missing_residue() { - let chain = Chain::new('C', ChainType::RNA); - assert_eq!(chain.get_residue_by_id(10), None); - } - - #[test] - fn get_residue_and_get_residue_by_id_return_correct_residue() { - let mut chain = Chain::new('D', ChainType::Protein); - let residue = Residue::new(5, "GLY"); - chain.residues.push(residue.clone()); - chain.residue_map.insert(5, 0); - - assert_eq!(chain.get_residue(0), Some(&residue)); - assert_eq!(chain.get_residue_by_id(5), Some(&residue)); - } } From c18bbab7d3bd260c6b396103b4901d3f13011325 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 19:34:36 -0700 Subject: [PATCH 08/20] test(core): Add unit tests for ChainType parsing and Chain creation --- crates/scream-core/src/core/models/chain.rs | 38 +++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/crates/scream-core/src/core/models/chain.rs b/crates/scream-core/src/core/models/chain.rs index 7b49df1d..855564d5 100644 --- a/crates/scream-core/src/core/models/chain.rs +++ b/crates/scream-core/src/core/models/chain.rs @@ -68,3 +68,41 @@ impl Chain { &self.residues } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_valid_chain_types() { + assert_eq!(ChainType::from_str("protein").unwrap(), ChainType::Protein); + assert_eq!(ChainType::from_str("dna").unwrap(), ChainType::DNA); + assert_eq!(ChainType::from_str("water").unwrap(), ChainType::Water); + } + + #[test] + fn parses_chain_type_case_insensitive() { + assert_eq!(ChainType::from_str("PrOtEiN").unwrap(), ChainType::Protein); + } + + #[test] + fn parses_unknown_chain_type_as_other() { + assert_eq!( + ChainType::from_str("carbohydrate").unwrap(), + ChainType::Other + ); + } + + #[test] + fn displays_chain_type_correctly() { + assert_eq!(ChainType::Protein.to_string(), "Protein"); + } + + #[test] + fn creates_chain_with_correct_id_and_type() { + let chain = Chain::new('A', ChainType::Protein); + assert_eq!(chain.id, 'A'); + assert_eq!(chain.chain_type, ChainType::Protein); + assert!(chain.residues().is_empty()); + } +} From 62fe7a882ccd56bd9909aa173ec10cf85122c324 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 20:04:47 -0700 Subject: [PATCH 09/20] refactor(core): Update Bond struct to use AtomId for atom identification --- .../scream-core/src/core/models/topology.rs | 102 ++---------------- 1 file changed, 10 insertions(+), 92 deletions(-) diff --git a/crates/scream-core/src/core/models/topology.rs b/crates/scream-core/src/core/models/topology.rs index 8f23ce2e..8b1978d4 100644 --- a/crates/scream-core/src/core/models/topology.rs +++ b/crates/scream-core/src/core/models/topology.rs @@ -1,3 +1,4 @@ +use super::ids::AtomId; use std::fmt; use std::str::FromStr; use thiserror::Error; @@ -51,104 +52,21 @@ impl fmt::Display for BondOrder { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct Bond { - pub atom1_idx: usize, // Index of the first atom in the global atom vector - pub atom2_idx: usize, // Index of the second atom in the global atom vector + pub atom1_id: AtomId, // ID of the first atom + pub atom2_id: AtomId, // ID of the second atom pub order: BondOrder, // Bond order (e.g., single, double, etc.) } impl Bond { - pub fn new(atom1_idx: usize, atom2_idx: usize, order: BondOrder) -> Self { - if atom1_idx < atom2_idx { - Self { - atom1_idx, - atom2_idx, - order, - } - } else { - Self { - atom1_idx: atom2_idx, - atom2_idx: atom1_idx, - order, - } + pub fn new(atom1_id: AtomId, atom2_id: AtomId, order: BondOrder) -> Self { + Self { + atom1_id, + atom2_id, + order, } } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parses_valid_bond_orders_from_str() { - assert_eq!(BondOrder::from_str("1").unwrap(), BondOrder::Single); - assert_eq!(BondOrder::from_str("s").unwrap(), BondOrder::Single); - assert_eq!(BondOrder::from_str("single").unwrap(), BondOrder::Single); - assert_eq!(BondOrder::from_str("2").unwrap(), BondOrder::Double); - assert_eq!(BondOrder::from_str("d").unwrap(), BondOrder::Double); - assert_eq!(BondOrder::from_str("double").unwrap(), BondOrder::Double); - assert_eq!(BondOrder::from_str("3").unwrap(), BondOrder::Triple); - assert_eq!(BondOrder::from_str("t").unwrap(), BondOrder::Triple); - assert_eq!(BondOrder::from_str("triple").unwrap(), BondOrder::Triple); - assert_eq!(BondOrder::from_str("ar").unwrap(), BondOrder::Aromatic); - assert_eq!( - BondOrder::from_str("aromatic").unwrap(), - BondOrder::Aromatic - ); - } - - #[test] - fn parses_bond_order_case_insensitive() { - assert_eq!(BondOrder::from_str("SINGLE").unwrap(), BondOrder::Single); - assert_eq!(BondOrder::from_str("DoUbLe").unwrap(), BondOrder::Double); - assert_eq!(BondOrder::from_str("TrIpLe").unwrap(), BondOrder::Triple); - assert_eq!( - BondOrder::from_str("AROMATIC").unwrap(), - BondOrder::Aromatic - ); - } - - #[test] - fn fails_to_parse_invalid_bond_order() { - assert!(BondOrder::from_str("quadruple").is_err()); - assert!(BondOrder::from_str("").is_err()); - assert!(BondOrder::from_str("x").is_err()); - assert!(BondOrder::from_str("0").is_err()); - } - - #[test] - fn displays_bond_order_correctly() { - assert_eq!(BondOrder::Single.to_string(), "Single"); - assert_eq!(BondOrder::Double.to_string(), "Double"); - assert_eq!(BondOrder::Triple.to_string(), "Triple"); - assert_eq!(BondOrder::Aromatic.to_string(), "Aromatic"); - } - - #[test] - fn bond_order_default_is_single() { - assert_eq!(BondOrder::default(), BondOrder::Single); - } - - #[test] - fn creates_bond_with_lower_index_first() { - let bond = Bond::new(2, 5, BondOrder::Double); - assert_eq!(bond.atom1_idx, 2); - assert_eq!(bond.atom2_idx, 5); - assert_eq!(bond.order, BondOrder::Double); - } - - #[test] - fn creates_bond_with_indices_swapped_if_needed() { - let bond = Bond::new(7, 3, BondOrder::Aromatic); - assert_eq!(bond.atom1_idx, 3); - assert_eq!(bond.atom2_idx, 7); - assert_eq!(bond.order, BondOrder::Aromatic); - } - #[test] - fn creates_bond_with_equal_indices() { - let bond = Bond::new(4, 4, BondOrder::Single); - assert_eq!(bond.atom1_idx, 4); - assert_eq!(bond.atom2_idx, 4); - assert_eq!(bond.order, BondOrder::Single); + pub fn contains(&self, atom_id: AtomId) -> bool { + self.atom1_id == atom_id || self.atom2_id == atom_id } } From 3daa7118e2ec8bed9b35c05a00c5988bb3a70ad7 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 20:07:35 -0700 Subject: [PATCH 10/20] test(core): Add comprehensive tests for Bond and BondOrder functionality --- .../scream-core/src/core/models/topology.rs | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/crates/scream-core/src/core/models/topology.rs b/crates/scream-core/src/core/models/topology.rs index 8b1978d4..2828911e 100644 --- a/crates/scream-core/src/core/models/topology.rs +++ b/crates/scream-core/src/core/models/topology.rs @@ -70,3 +70,80 @@ impl Bond { self.atom1_id == atom_id || self.atom2_id == atom_id } } + +#[cfg(test)] +mod tests { + use super::*; + use slotmap::KeyData; + + fn dummy_atom_id(n: u64) -> AtomId { + AtomId::from(KeyData::from_ffi(n)) + } + + #[test] + fn bond_order_from_str_parses_valid_strings() { + assert_eq!("1".parse::().unwrap(), BondOrder::Single); + assert_eq!("single".parse::().unwrap(), BondOrder::Single); + assert_eq!("S".parse::().unwrap(), BondOrder::Single); + assert_eq!("2".parse::().unwrap(), BondOrder::Double); + assert_eq!("double".parse::().unwrap(), BondOrder::Double); + assert_eq!("D".parse::().unwrap(), BondOrder::Double); + assert_eq!("3".parse::().unwrap(), BondOrder::Triple); + assert_eq!("triple".parse::().unwrap(), BondOrder::Triple); + assert_eq!("T".parse::().unwrap(), BondOrder::Triple); + assert_eq!("ar".parse::().unwrap(), BondOrder::Aromatic); + assert_eq!( + "aromatic".parse::().unwrap(), + BondOrder::Aromatic + ); + } + + #[test] + fn bond_order_from_str_rejects_invalid_strings() { + assert!("".parse::().is_err()); + assert!("quadruple".parse::().is_err()); + assert!("unknown".parse::().is_err()); + assert!("0".parse::().is_err()); + } + + #[test] + fn bond_order_display_outputs_expected_strings() { + assert_eq!(BondOrder::Single.to_string(), "Single"); + assert_eq!(BondOrder::Double.to_string(), "Double"); + assert_eq!(BondOrder::Triple.to_string(), "Triple"); + assert_eq!(BondOrder::Aromatic.to_string(), "Aromatic"); + } + + #[test] + fn bond_order_default_is_single() { + assert_eq!(BondOrder::default(), BondOrder::Single); + } + + #[test] + fn bond_new_initializes_fields_correctly() { + let a1 = dummy_atom_id(1); + let a2 = dummy_atom_id(2); + let bond = Bond::new(a1, a2, BondOrder::Double); + assert_eq!(bond.atom1_id, a1); + assert_eq!(bond.atom2_id, a2); + assert_eq!(bond.order, BondOrder::Double); + } + + #[test] + fn bond_contains_returns_true_for_both_atoms() { + let a1 = dummy_atom_id(10); + let a2 = dummy_atom_id(20); + let bond = Bond::new(a1, a2, BondOrder::Single); + assert!(bond.contains(a1)); + assert!(bond.contains(a2)); + } + + #[test] + fn bond_contains_returns_false_for_unrelated_atom() { + let a1 = dummy_atom_id(100); + let a2 = dummy_atom_id(200); + let unrelated = dummy_atom_id(300); + let bond = Bond::new(a1, a2, BondOrder::Aromatic); + assert!(!bond.contains(unrelated)); + } +} From f118a9055daf689047c3cb3d9113689cfa1ded61 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 20:15:08 -0700 Subject: [PATCH 11/20] refactor(core): Revise MolecularSystem structure for improved data management and performance --- crates/scream-core/src/core/models/system.rs | 332 ++++++++----------- 1 file changed, 133 insertions(+), 199 deletions(-) diff --git a/crates/scream-core/src/core/models/system.rs b/crates/scream-core/src/core/models/system.rs index 719dcea8..2b74c638 100644 --- a/crates/scream-core/src/core/models/system.rs +++ b/crates/scream-core/src/core/models/system.rs @@ -1,247 +1,181 @@ -use super::atom::Atom; -use super::chain::Chain; +use super::atom::{Atom, AtomFlags}; +use super::chain::{Chain, ChainType}; +use super::ids::{AtomId, ChainId, ResidueId}; use super::residue::Residue; -use super::topology::Bond; +use super::topology::{Bond, BondOrder}; +use nalgebra::Point3; +use slotmap::{SecondaryMap, SlotMap}; +use std::collections::HashMap; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct MolecularSystem { - pub(crate) atoms: Vec, - pub(crate) chains: Vec, - pub(crate) bonds: Vec, - - pub(crate) atom_serial_map: std::collections::HashMap, - pub(crate) chain_id_map: std::collections::HashMap, + // --- Primary Data Stores (Source of Truth) --- + atoms: SlotMap, + residues: SlotMap, + chains: SlotMap, + bonds: Vec, + + // --- Lookup Maps for Fast Access by External Identifiers --- + atom_serial_map: HashMap, + residue_id_map: HashMap<(ChainId, isize), ResidueId>, + chain_id_map: HashMap, + + // --- Adjacency information (cache for performance) --- + bond_adjacency: SecondaryMap>, } impl MolecularSystem { - pub fn atoms(&self) -> &[Atom] { - &self.atoms + pub fn new() -> Self { + Self::default() } - pub fn chains(&self) -> &[Chain] { - &self.chains + pub fn atom(&self, id: AtomId) -> Option<&Atom> { + self.atoms.get(id) } - - pub fn bonds(&self) -> &[Bond] { - &self.bonds + pub fn atom_mut(&mut self, id: AtomId) -> Option<&mut Atom> { + self.atoms.get_mut(id) } - - pub fn bonds_mut(&mut self) -> &mut Vec { - &mut self.bonds + pub fn atoms_iter(&self) -> impl Iterator { + self.atoms.iter() } - pub fn get_atom(&self, index: usize) -> Option<&Atom> { - self.atoms.get(index) + pub fn residue(&self, id: ResidueId) -> Option<&Residue> { + self.residues.get(id) } - - pub fn get_atom_mut(&mut self, index: usize) -> Option<&mut Atom> { - self.atoms.get_mut(index) + pub fn residue_mut(&mut self, id: ResidueId) -> Option<&mut Residue> { + self.residues.get_mut(id) + } + pub fn residues_iter(&self) -> impl Iterator { + self.residues.iter() } - pub fn get_atom_by_serial(&self, serial: usize) -> Option<&Atom> { - self.atom_serial_map - .get(&serial) - .and_then(|&index| self.atoms.get(index)) + pub fn chain(&self, id: ChainId) -> Option<&Chain> { + self.chains.get(id) + } + pub fn chain_mut(&mut self, id: ChainId) -> Option<&mut Chain> { + self.chains.get_mut(id) + } + pub fn chains_iter(&self) -> impl Iterator { + self.chains.iter() } - pub fn get_chain(&self, index: usize) -> Option<&Chain> { - self.chains.get(index) + pub fn bonds(&self) -> &[Bond] { + &self.bonds } - pub fn get_chain_by_id(&self, id: char) -> Option<&Chain> { - self.chain_id_map - .get(&id) - .and_then(|&index| self.chains.get(index)) + pub fn find_atom_by_serial(&self, serial: usize) -> Option { + self.atom_serial_map.get(&serial).copied() + } + pub fn find_chain_by_id(&self, id: char) -> Option { + self.chain_id_map.get(&id).copied() + } + pub fn find_residue_by_id(&self, chain_id: ChainId, res_seq: isize) -> Option { + self.residue_id_map.get(&(chain_id, res_seq)).copied() } - pub fn atoms_in_residue<'a>(&'a self, residue: &'a Residue) -> impl Iterator { - residue - .atom_indices - .iter() - .map(move |&idx| &self.atoms[idx]) + pub fn add_chain(&mut self, id: char, chain_type: ChainType) -> ChainId { + *self.chain_id_map.entry(id).or_insert_with(|| { + let chain = Chain::new(id, chain_type); + self.chains.insert(chain) + }) } -} -#[cfg(test)] -mod tests { - use super::*; - use crate::core::models::atom::{Atom, AtomFlags}; - use crate::core::models::chain::Chain; - use crate::core::models::residue::Residue; - use crate::core::models::topology::Bond; - use nalgebra::Point3; - - fn mock_system() -> MolecularSystem { - let atom1 = Atom { - index: 0, - serial: 10, - name: "CA".to_string(), - res_name: "GLY".to_string(), - res_id: 5, - chain_id: 'A', - force_field_type: "C.3".to_string(), - partial_charge: -0.123, - position: Point3::new(1.0, 2.0, 3.0), - flags: AtomFlags::IS_FIXED_ROLE | AtomFlags::IS_VISIBLE_LATTICE, - delta: 0.5, - vdw_radius: 1.7, - vdw_well_depth: 0.2, - hbond_type_id: 3, - }; - let atom2 = Atom { - index: 1, - serial: 20, - name: "CB".to_string(), - res_name: "GLY".to_string(), - res_id: 5, - chain_id: 'A', - force_field_type: "C.3".to_string(), - partial_charge: -0.456, - position: Point3::new(2.0, 3.0, 4.0), - flags: AtomFlags::IS_VISIBLE_INTERACTION, - delta: 0.6, - vdw_radius: 1.8, - vdw_well_depth: 0.3, - hbond_type_id: 2, - }; - let mut residue = Residue::new(5, "GLY"); - residue.atom_indices.push(0); - residue.atom_indices.push(1); - let mut chain = Chain::new('A', crate::core::models::chain::ChainType::Protein); - chain.residues.push(residue); - let bond = Bond::new(0, 1, crate::core::models::topology::BondOrder::Single); - - MolecularSystem { - atoms: vec![atom1, atom2], - chains: vec![chain], - bonds: vec![bond], - atom_serial_map: vec![(10, 0), (20, 1)].into_iter().collect(), - chain_id_map: vec![('A', 0)].into_iter().collect(), + pub fn add_residue( + &mut self, + chain_id: ChainId, + res_seq: isize, + name: &str, + ) -> Option { + let chain = self.chains.get_mut(chain_id)?; + let key = (chain_id, res_seq); + let residue_id = *self.residue_id_map.entry(key).or_insert_with(|| { + let residue = Residue::new(res_seq, name, chain_id); + self.residues.insert(residue) + }); + + let chain_residues = &mut self.chains.get_mut(chain_id).unwrap().residues; + if !chain_residues.contains(&residue_id) { + chain_residues.push(residue_id); } - } - #[test] - fn atoms_returns_all_atoms() { - let system = mock_system(); - assert_eq!(system.atoms().len(), 2); + Some(residue_id) } - #[test] - fn chains_returns_all_chains() { - let system = mock_system(); - assert_eq!(system.chains().len(), 1); - } + pub fn add_atom_to_residue(&mut self, residue_id: ResidueId, atom: Atom) -> Option { + if !self.residues.contains_key(residue_id) { + return None; + } - #[test] - fn bonds_returns_all_bonds() { - let system = mock_system(); - assert_eq!(system.bonds().len(), 1); - } + let serial = atom.serial; + let name = atom.name.clone(); - #[test] - fn bonds_mut_allows_modification() { - let mut system = mock_system(); - system.bonds_mut().clear(); - assert_eq!(system.bonds().len(), 0); - } + let atom_id = self.atoms.insert(atom); + self.bond_adjacency.insert(atom_id, Vec::new()); + self.atom_serial_map.insert(serial, atom_id); - #[test] - fn get_atom_returns_correct_atom() { - let system = mock_system(); - assert_eq!(system.get_atom(0).unwrap().serial, 10); - assert_eq!(system.get_atom(1).unwrap().serial, 20); - assert!(system.get_atom(2).is_none()); + let residue = self.residues.get_mut(residue_id).unwrap(); + residue.add_atom(&name, atom_id); + + Some(atom_id) } - #[test] - fn get_atom_mut_allows_modification() { - let mut system = mock_system(); - if let Some(atom) = system.get_atom_mut(0) { - atom.serial = 99; + pub fn add_bond(&mut self, atom1_id: AtomId, atom2_id: AtomId, order: BondOrder) -> Option<()> { + if self.atoms.contains_key(atom1_id) && self.atoms.contains_key(atom2_id) { + self.bonds.push(Bond::new(atom1_id, atom2_id, order)); + self.bond_adjacency[atom1_id].push(atom2_id); + self.bond_adjacency[atom2_id].push(atom1_id); + Some(()) + } else { + None } - assert_eq!(system.get_atom(0).unwrap().serial, 99); } - #[test] - fn get_atom_by_serial_returns_correct_atom() { - let system = mock_system(); - assert_eq!(system.get_atom_by_serial(10).unwrap().index, 0); - assert_eq!(system.get_atom_by_serial(20).unwrap().index, 1); - assert!(system.get_atom_by_serial(999).is_none()); - } + pub fn remove_atom(&mut self, atom_id: AtomId) -> Option { + let atom = self.atoms.remove(atom_id)?; - #[test] - fn get_chain_returns_correct_chain() { - let system = mock_system(); - assert!(system.get_chain(0).is_some()); - assert!(system.get_chain(1).is_none()); - } + // 1. Remove from parent residue + let residue = self.residues.get_mut(atom.residue_id).unwrap(); + residue.remove_atom(&atom.name, atom_id); - #[test] - fn get_chain_by_id_returns_correct_chain() { - let system = mock_system(); - assert!(system.get_chain_by_id('A').is_some()); - assert!(system.get_chain_by_id('B').is_none()); - } + // 2. Remove from serial map + self.atom_serial_map.remove(&atom.serial); - #[test] - fn get_chain_by_id_returns_none_for_invalid_id() { - let system = mock_system(); - assert!(system.get_chain_by_id('Z').is_none()); - } + // 3. Remove all bonds connected to this atom + let original_bonds = std::mem::take(&mut self.bonds); + self.bonds = original_bonds + .into_iter() + .filter(|bond| !bond.contains(atom_id)) + .collect(); - #[test] - fn atoms_in_residue_iterates_correct_atoms() { - let system = mock_system(); - let residue = &system.chains()[0].residues()[0]; - let indices: Vec<_> = system.atoms_in_residue(residue).map(|a| a.index).collect(); - assert_eq!(indices, vec![0, 1]); - } + // 4. Clean up adjacency list + let neighbors = self.bond_adjacency.remove(atom_id).unwrap_or_default(); + for neighbor_id in neighbors { + if let Some(adjacency) = self.bond_adjacency.get_mut(neighbor_id) { + adjacency.retain(|&id| id != atom_id); + } + } - #[test] - fn atoms_in_residue_empty_for_residue_with_no_atoms() { - let system = mock_system(); - let empty_residue = Residue::new(6, "ALA"); - let atoms: Vec<_> = system.atoms_in_residue(&empty_residue).collect(); - assert!(atoms.is_empty()); + Some(atom) } - #[test] - fn get_atom_by_serial_returns_none_for_empty_atoms() { - let mut system = mock_system(); - system.atoms.clear(); - assert!(system.get_atom_by_serial(10).is_none()); - } + pub fn remove_residue(&mut self, residue_id: ResidueId) -> Option { + let residue = self.residues.get(residue_id)?.clone(); // Clone to avoid borrow checker issues - #[test] - fn get_atom_returns_none_for_empty_atoms() { - let mut system = mock_system(); - system.atoms.clear(); - assert!(system.get_atom(0).is_none()); - } + // 1. Remove all atoms within the residue + for atom_id in residue.atoms().to_vec() { + // to_vec to avoid borrow issues + self.remove_atom(atom_id); + } - #[test] - fn get_chain_returns_none_for_empty_chains() { - let mut system = mock_system(); - system.chains.clear(); - assert!(system.get_chain(0).is_none()); - } + // 2. Remove from parent chain + if let Some(chain) = self.chains.get_mut(residue.chain_id) { + chain.residues.retain(|&id| id != residue_id); + } - #[test] - fn bonds_returns_empty_when_no_bonds() { - let mut system = mock_system(); - system.bonds.clear(); - assert!(system.bonds().is_empty()); - } + // 3. Remove from residue maps + self.residue_id_map.remove(&(residue.chain_id, residue.id)); - #[test] - fn atoms_in_residue_panics_on_invalid_index() { - let mut system = mock_system(); - let mut residue = Residue::new(5, "GLY"); - residue.atom_indices.push(99); - let result = std::panic::catch_unwind(|| { - let _ = system.atoms_in_residue(&residue).count(); - }); - assert!(result.is_err()); + // 4. Finally, remove the residue itself + self.residues.remove(residue_id) } } From bc3af3edc3a531e4938cc5ee752e3ddc98f6a654 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 20:15:41 -0700 Subject: [PATCH 12/20] test(core): Add unit tests for MolecularSystem functionality including creation, atom and residue removal --- crates/scream-core/src/core/models/system.rs | 115 +++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/crates/scream-core/src/core/models/system.rs b/crates/scream-core/src/core/models/system.rs index 2b74c638..c8e89132 100644 --- a/crates/scream-core/src/core/models/system.rs +++ b/crates/scream-core/src/core/models/system.rs @@ -179,3 +179,118 @@ impl MolecularSystem { self.residues.remove(residue_id) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn create_test_system() -> MolecularSystem { + let mut system = MolecularSystem::new(); + + let chain_a_id = system.add_chain('A', ChainType::Protein); + + let gly_id = system.add_residue(chain_a_id, 1, "GLY").unwrap(); + let n_atom = Atom::new(1, "N", gly_id, Point3::new(0.0, 0.0, 0.0)); + let ca_atom = Atom::new(2, "CA", gly_id, Point3::new(1.4, 0.0, 0.0)); + let n_id = system.add_atom_to_residue(gly_id, n_atom).unwrap(); + let ca_id = system.add_atom_to_residue(gly_id, ca_atom).unwrap(); + system.add_bond(n_id, ca_id, BondOrder::Single); + + let ala_id = system.add_residue(chain_a_id, 2, "ALA").unwrap(); + let ala_ca_atom = Atom::new(3, "CA", ala_id, Point3::new(2.0, 1.0, 0.0)); + system.add_atom_to_residue(ala_id, ala_ca_atom); + + system + } + + #[test] + fn system_creation_and_access() { + let system = create_test_system(); + + assert_eq!(system.atoms_iter().count(), 3); + assert_eq!(system.residues_iter().count(), 2); + assert_eq!(system.chains_iter().count(), 1); + assert_eq!(system.bonds.len(), 1); + + let chain_a_id = system.find_chain_by_id('A').unwrap(); + assert!(system.find_chain_by_id('B').is_none()); + + let gly_id = system.find_residue_by_id(chain_a_id, 1).unwrap(); + let ala_id = system.find_residue_by_id(chain_a_id, 2).unwrap(); + + assert_eq!(system.residue(gly_id).unwrap().name, "GLY"); + assert_eq!(system.residue(ala_id).unwrap().name, "ALA"); + + let atom_n_id = system.find_atom_by_serial(1).unwrap(); + assert_eq!(system.atom(atom_n_id).unwrap().name, "N"); + } + + #[test] + fn atom_removal_updates_system_correctly() { + let mut system = create_test_system(); + let atom_n_id = system.find_atom_by_serial(1).unwrap(); + let atom_ca_id = system.find_atom_by_serial(2).unwrap(); + + assert_eq!(system.bonds.len(), 1); + assert!(system.bonds[0].contains(atom_n_id)); + assert!(system.atom(atom_n_id).is_some()); + let gly_id = system + .residue(system.atom(atom_n_id).unwrap().residue_id) + .unwrap() + .id; + assert_eq!(gly_id, 1); + + let removed_atom = system.remove_atom(atom_n_id).unwrap(); + assert_eq!(removed_atom.name, "N"); + + assert_eq!(system.atoms_iter().count(), 2); + assert!(system.atom(atom_n_id).is_none()); + assert!(system.find_atom_by_serial(1).is_none()); + + assert!(system.bonds.is_empty()); + + assert!(system.bond_adjacency.get(atom_n_id).is_none()); + assert!(!system.bond_adjacency[atom_ca_id].contains(&atom_n_id)); + + let chain_a_id = system.find_chain_by_id('A').unwrap(); + let gly_id = system.find_residue_by_id(chain_a_id, 1).unwrap(); + assert_eq!(system.residue(gly_id).unwrap().atoms().len(), 1); + assert!(!system.residue(gly_id).unwrap().atoms().contains(&atom_n_id)); + } + + #[test] + fn residue_removal_updates_system_correctly() { + let mut system = create_test_system(); + + let chain_a_id = system.find_chain_by_id('A').unwrap(); + let gly_id = system.find_residue_by_id(chain_a_id, 1).unwrap(); + + assert_eq!(system.atoms_iter().count(), 3); + assert_eq!(system.residues_iter().count(), 2); + assert_eq!(system.bonds.len(), 1); + assert_eq!(system.chain(chain_a_id).unwrap().residues().len(), 2); + + let removed_residue = system.remove_residue(gly_id).unwrap(); + assert_eq!(removed_residue.name, "GLY"); + + assert_eq!(system.residues_iter().count(), 1); + assert!(system.residue(gly_id).is_none()); + assert!(system.find_residue_by_id(chain_a_id, 1).is_none()); + + assert_eq!(system.atoms_iter().count(), 1); + assert!(system.find_atom_by_serial(1).is_none()); + assert!(system.find_atom_by_serial(2).is_none()); + assert!(system.find_atom_by_serial(3).is_some()); + + assert!(system.bonds.is_empty()); + + assert_eq!(system.chain(chain_a_id).unwrap().residues().len(), 1); + assert!( + !system + .chain(chain_a_id) + .unwrap() + .residues() + .contains(&gly_id) + ); + } +} From 3aeede39c90fc9c204150c4f803f717189d32b68 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 20:19:55 -0700 Subject: [PATCH 13/20] refactor(core): Remove unused AtomFlags import and adjust test module imports --- crates/scream-core/src/core/models/system.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/scream-core/src/core/models/system.rs b/crates/scream-core/src/core/models/system.rs index c8e89132..75effcc2 100644 --- a/crates/scream-core/src/core/models/system.rs +++ b/crates/scream-core/src/core/models/system.rs @@ -1,9 +1,8 @@ -use super::atom::{Atom, AtomFlags}; +use super::atom::Atom; use super::chain::{Chain, ChainType}; use super::ids::{AtomId, ChainId, ResidueId}; use super::residue::Residue; use super::topology::{Bond, BondOrder}; -use nalgebra::Point3; use slotmap::{SecondaryMap, SlotMap}; use std::collections::HashMap; @@ -183,6 +182,7 @@ impl MolecularSystem { #[cfg(test)] mod tests { use super::*; + use nalgebra::Point3; fn create_test_system() -> MolecularSystem { let mut system = MolecularSystem::new(); From d3b8428cd5fe22f105e69392206d5d9f6b47bf65 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 20:22:01 -0700 Subject: [PATCH 14/20] refactor(core): Remove builder module from core models --- crates/scream-core/src/core/models/builder.rs | 256 ------------------ crates/scream-core/src/core/models/mod.rs | 1 - 2 files changed, 257 deletions(-) delete mode 100644 crates/scream-core/src/core/models/builder.rs diff --git a/crates/scream-core/src/core/models/builder.rs b/crates/scream-core/src/core/models/builder.rs deleted file mode 100644 index 297ada3b..00000000 --- a/crates/scream-core/src/core/models/builder.rs +++ /dev/null @@ -1,256 +0,0 @@ -use super::atom::{Atom, AtomFlags}; -use super::chain::{Chain, ChainType}; -use super::residue::Residue; -use super::system::MolecularSystem; -use super::topology::{Bond, BondOrder}; -use nalgebra::Point3; -use std::collections::HashMap; - -pub struct MolecularSystemBuilder { - system: MolecularSystem, - - // --- Builder-specific state for efficient construction --- - atom_serial_map: HashMap, - chain_id_map: HashMap, - current_chain_idx: Option, - current_residue_idx: Option, -} - -impl Default for MolecularSystemBuilder { - fn default() -> Self { - Self::new() - } -} - -impl MolecularSystemBuilder { - pub fn new() -> Self { - Self { - system: MolecularSystem { - atoms: Vec::new(), - chains: Vec::new(), - bonds: Vec::new(), - atom_serial_map: HashMap::new(), - chain_id_map: HashMap::new(), - }, - atom_serial_map: HashMap::new(), - chain_id_map: HashMap::new(), - current_chain_idx: None, - current_residue_idx: None, - } - } - - pub fn start_chain(&mut self, id: char, chain_type: ChainType) -> &mut Self { - let idx = *self.chain_id_map.entry(id).or_insert_with(|| { - let index = self.system.chains.len(); - self.system.chains.push(Chain::new(id, chain_type)); - index - }); - self.current_chain_idx = Some(idx); - self.current_residue_idx = None; - self - } - - pub fn start_residue(&mut self, id: isize, name: &str) -> &mut Self { - let chain_idx = self - .current_chain_idx - .expect("Must start a chain before starting a residue"); - let chain = &mut self.system.chains[chain_idx]; - - let res_idx = *chain.residue_map.entry(id).or_insert_with(|| { - let index = chain.residues.len(); - chain.residues.push(Residue::new(id, name)); - index - }); - self.current_residue_idx = Some(res_idx); - self - } - - pub fn add_atom( - &mut self, - serial: usize, - name: &str, - res_name: &str, - position: Point3, - charge: Option, - ff_type: Option<&str>, - ) -> &mut Self { - let chain_idx = self - .current_chain_idx - .expect("Cannot add atom without a current chain"); - let res_idx = self - .current_residue_idx - .expect("Cannot add atom without a current residue"); - - let chain_id = self.system.chains[chain_idx].id; - let res_id = self.system.chains[chain_idx].residues[res_idx].id; - let atom_idx = self.system.atoms.len(); - - let atom = Atom { - index: atom_idx, - serial, - name: name.to_string(), - res_name: res_name.to_string(), - res_id, - chain_id, - position, - partial_charge: charge.unwrap_or(0.0), - force_field_type: ff_type.unwrap_or("").to_string(), - flags: AtomFlags::default(), - delta: 0.0, - vdw_radius: 0.0, - vdw_well_depth: 0.0, - hbond_type_id: -1, - }; - - self.system.atoms.push(atom); - self.atom_serial_map.insert(serial, atom_idx); - self.system.chains[chain_idx].residues[res_idx].add_atom(name, atom_idx); - self - } - - pub fn add_bond(&mut self, serial1: usize, serial2: usize, order: BondOrder) -> &mut Self { - let idx1 = *self - .atom_serial_map - .get(&serial1) - .expect("Atom 1 for bond not found"); - let idx2 = *self - .atom_serial_map - .get(&serial2) - .expect("Atom 2 for bond not found"); - self.system.bonds.push(Bond::new(idx1, idx2, order)); - self - } - - pub fn build(mut self) -> MolecularSystem { - MolecularSystem { - atoms: self.system.atoms, - chains: self.system.chains, - bonds: self.system.bonds, - atom_serial_map: self.atom_serial_map, - chain_id_map: self.chain_id_map, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use nalgebra::Point3; - - #[test] - fn builder_creates_system_with_single_chain_residue_atom() { - let mut builder = MolecularSystemBuilder::new(); - builder - .start_chain('A', ChainType::Protein) - .start_residue(1, "GLY") - .add_atom( - 100, - "CA", - "GLY", - Point3::new(1.0, 2.0, 3.0), - Some(-0.1), - Some("C.3"), - ); - let system = builder.build(); - assert_eq!(system.chains().len(), 1); - assert_eq!(system.chains()[0].residues().len(), 1); - assert_eq!(system.atoms().len(), 1); - assert_eq!(system.atoms()[0].serial, 100); - assert_eq!(system.atoms()[0].name, "CA"); - assert_eq!(system.atoms()[0].res_name, "GLY"); - assert_eq!(system.atoms()[0].partial_charge, -0.1); - assert_eq!(system.atoms()[0].force_field_type, "C.3"); - } - - #[test] - fn builder_adds_multiple_chains_and_residues() { - let mut builder = MolecularSystemBuilder::new(); - builder - .start_chain('A', ChainType::Protein) - .start_residue(1, "GLY") - .add_atom(1, "CA", "GLY", Point3::new(0.0, 0.0, 0.0), None, None) - .start_residue(2, "ALA") - .add_atom(2, "CB", "ALA", Point3::new(1.0, 1.0, 1.0), None, None) - .start_chain('B', ChainType::Protein) - .start_residue(1, "SER") - .add_atom(3, "OG", "SER", Point3::new(2.0, 2.0, 2.0), None, None); - let system = builder.build(); - assert_eq!(system.chains().len(), 2); - assert_eq!(system.chains()[0].residues().len(), 2); - assert_eq!(system.chains()[1].residues().len(), 1); - assert_eq!(system.atoms().len(), 3); - } - - #[test] - fn builder_adds_bond_between_atoms() { - let mut builder = MolecularSystemBuilder::new(); - builder - .start_chain('A', ChainType::Protein) - .start_residue(1, "GLY") - .add_atom(1, "CA", "GLY", Point3::new(0.0, 0.0, 0.0), None, None) - .add_atom(2, "CB", "GLY", Point3::new(1.0, 1.0, 1.0), None, None) - .add_bond(1, 2, BondOrder::Single); - let system = builder.build(); - assert_eq!(system.bonds().len(), 1); - let bond = &system.bonds()[0]; - assert_eq!(bond.atom1_idx, 0); - assert_eq!(bond.atom2_idx, 1); - } - - #[test] - fn builder_add_atom_without_charge_and_ff_type_defaults() { - let mut builder = MolecularSystemBuilder::new(); - builder - .start_chain('A', ChainType::Protein) - .start_residue(1, "GLY") - .add_atom(1, "CA", "GLY", Point3::new(0.0, 0.0, 0.0), None, None); - let system = builder.build(); - let atom = &system.atoms()[0]; - assert_eq!(atom.partial_charge, 0.0); - assert_eq!(atom.force_field_type, ""); - } - - #[test] - #[should_panic(expected = "Must start a chain before starting a residue")] - fn builder_panics_if_start_residue_without_chain() { - let mut builder = MolecularSystemBuilder::new(); - builder.start_residue(1, "GLY"); - } - - #[test] - #[should_panic(expected = "Cannot add atom without a current chain")] - fn builder_panics_if_add_atom_without_chain() { - let mut builder = MolecularSystemBuilder::new(); - builder.add_atom(1, "CA", "GLY", Point3::new(0.0, 0.0, 0.0), None, None); - } - - #[test] - #[should_panic(expected = "Cannot add atom without a current residue")] - fn builder_panics_if_add_atom_without_residue() { - let mut builder = MolecularSystemBuilder::new(); - builder.start_chain('A', ChainType::Protein); - builder.add_atom(1, "CA", "GLY", Point3::new(0.0, 0.0, 0.0), None, None); - } - - #[test] - #[should_panic(expected = "Atom 1 for bond not found")] - fn builder_panics_if_add_bond_with_invalid_serial1() { - let mut builder = MolecularSystemBuilder::new(); - builder - .start_chain('A', ChainType::Protein) - .start_residue(1, "GLY") - .add_atom(1, "CA", "GLY", Point3::new(0.0, 0.0, 0.0), None, None) - .add_bond(99, 1, BondOrder::Single); - } - - #[test] - #[should_panic(expected = "Atom 2 for bond not found")] - fn builder_panics_if_add_bond_with_invalid_serial2() { - let mut builder = MolecularSystemBuilder::new(); - builder - .start_chain('A', ChainType::Protein) - .start_residue(1, "GLY") - .add_atom(1, "CA", "GLY", Point3::new(0.0, 0.0, 0.0), None, None) - .add_bond(1, 99, BondOrder::Single); - } -} diff --git a/crates/scream-core/src/core/models/mod.rs b/crates/scream-core/src/core/models/mod.rs index d8a24ce3..2f1a8ff1 100644 --- a/crates/scream-core/src/core/models/mod.rs +++ b/crates/scream-core/src/core/models/mod.rs @@ -1,5 +1,4 @@ pub mod atom; -pub mod builder; pub mod chain; pub mod ids; pub mod residue; From 6b7a564578906fea87d5f855204b44de36414c56 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 23:13:28 -0700 Subject: [PATCH 15/20] refactor(core): Use BufWriter for improved file writing in MolecularFile trait --- crates/scream-core/src/core/io/traits.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/crates/scream-core/src/core/io/traits.rs b/crates/scream-core/src/core/io/traits.rs index fb40dd38..dccce5f9 100644 --- a/crates/scream-core/src/core/io/traits.rs +++ b/crates/scream-core/src/core/io/traits.rs @@ -1,7 +1,7 @@ use crate::core::models::system::MolecularSystem; use std::error::Error; use std::fs::File; -use std::io::{self, BufRead, BufReader, Write}; +use std::io::{self, BufRead, BufReader, BufWriter, Write}; use std::path::Path; pub trait MolecularFile { @@ -36,15 +36,17 @@ pub trait MolecularFile { metadata: &Self::Metadata, path: P, ) -> Result<(), Self::Error> { - let mut file = File::create(path)?; - Self::write_to(system, metadata, &mut file) + let file = File::create(path)?; + let mut writer = BufWriter::new(file); + Self::write_to(system, metadata, &mut writer) } fn write_system_to_path>( system: &MolecularSystem, path: P, ) -> Result<(), Self::Error> { - let mut file = File::create(path)?; - Self::write_system_to(system, &mut file) + let file = File::create(path)?; + let mut writer = BufWriter::new(file); + Self::write_system_to(system, &mut writer) } } From 8ba1d14c25b7357b5f2bb5537e770568486908f6 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 23:24:21 -0700 Subject: [PATCH 16/20] refactor(core): Improve bond addition logic and enhance atom removal safety in MolecularSystem --- crates/scream-core/src/core/models/system.rs | 33 +++++++++++--------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/crates/scream-core/src/core/models/system.rs b/crates/scream-core/src/core/models/system.rs index 75effcc2..0c002fff 100644 --- a/crates/scream-core/src/core/models/system.rs +++ b/crates/scream-core/src/core/models/system.rs @@ -119,22 +119,30 @@ impl MolecularSystem { } pub fn add_bond(&mut self, atom1_id: AtomId, atom2_id: AtomId, order: BondOrder) -> Option<()> { - if self.atoms.contains_key(atom1_id) && self.atoms.contains_key(atom2_id) { - self.bonds.push(Bond::new(atom1_id, atom2_id, order)); - self.bond_adjacency[atom1_id].push(atom2_id); - self.bond_adjacency[atom2_id].push(atom1_id); - Some(()) - } else { - None + if !self.atoms.contains_key(atom1_id) || !self.atoms.contains_key(atom2_id) { + return None; + } + + if let Some(neighbors) = self.bond_adjacency.get(atom1_id) { + if neighbors.contains(&atom2_id) { + // Bond already exists, operation is successful (idempotent) + return Some(()); + } } + + self.bonds.push(Bond::new(atom1_id, atom2_id, order)); + self.bond_adjacency[atom1_id].push(atom2_id); + self.bond_adjacency[atom2_id].push(atom1_id); + Some(()) } pub fn remove_atom(&mut self, atom_id: AtomId) -> Option { let atom = self.atoms.remove(atom_id)?; // 1. Remove from parent residue - let residue = self.residues.get_mut(atom.residue_id).unwrap(); - residue.remove_atom(&atom.name, atom_id); + if let Some(residue) = self.residues.get_mut(atom.residue_id) { + residue.remove_atom(&atom.name, atom_id); + } // 2. Remove from serial map self.atom_serial_map.remove(&atom.serial); @@ -234,11 +242,8 @@ mod tests { assert_eq!(system.bonds.len(), 1); assert!(system.bonds[0].contains(atom_n_id)); assert!(system.atom(atom_n_id).is_some()); - let gly_id = system - .residue(system.atom(atom_n_id).unwrap().residue_id) - .unwrap() - .id; - assert_eq!(gly_id, 1); + let residue_id = system.atom(atom_n_id).unwrap().residue_id; + assert_eq!(system.residue(residue_id).unwrap().id, 1); let removed_atom = system.remove_atom(atom_n_id).unwrap(); assert_eq!(removed_atom.name, "N"); From dbdf8738948fb2d9006555977d6761eaa803ba15 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 23:28:04 -0700 Subject: [PATCH 17/20] refactor(core): Replace MolecularSystemBuilder with direct MolecularSystem usage in BgfFile read_from method --- crates/scream-core/src/core/io/bgf.rs | 473 +++++++++----------------- 1 file changed, 167 insertions(+), 306 deletions(-) diff --git a/crates/scream-core/src/core/io/bgf.rs b/crates/scream-core/src/core/io/bgf.rs index ac1da7aa..400b8229 100644 --- a/crates/scream-core/src/core/io/bgf.rs +++ b/crates/scream-core/src/core/io/bgf.rs @@ -1,6 +1,7 @@ use super::traits::MolecularFile; -use crate::core::models::builder::MolecularSystemBuilder; +use crate::core::models::atom::Atom; use crate::core::models::chain::ChainType; +use crate::core::models::ids::{ChainId, ResidueId}; use crate::core::models::system::MolecularSystem; use crate::core::models::topology::BondOrder; use nalgebra::Point3; @@ -24,6 +25,8 @@ pub enum BgfError { Io(#[from] io::Error), #[error("Parse error on line {line_num}: {message}")] Parse { line_num: usize, message: String }, + #[error("Logic error: {0}")] + Logic(String), } impl MolecularFile for BgfFile { @@ -33,16 +36,15 @@ impl MolecularFile for BgfFile { fn read_from( reader: &mut impl BufRead, ) -> Result<(MolecularSystem, Self::Metadata), Self::Error> { - let mut builder = MolecularSystemBuilder::new(); + let mut system = MolecularSystem::new(); let mut metadata = BgfMetadata::default(); let mut atom_lines = Vec::new(); let mut conect_lines = Vec::new(); - let mut atom_section_started = false; - let lines: Vec = reader.lines().collect::>()?; - for line in lines { + for (i, line_result) in reader.lines().enumerate() { + let line = line_result?; if line.trim().is_empty() { continue; } @@ -50,64 +52,120 @@ impl MolecularFile for BgfFile { break; } - if line.starts_with("ATOM") || line.starts_with("HETATM") { - atom_section_started = true; - atom_lines.push(line); - } else if line.starts_with("CONECT") { - conect_lines.push(line); - } else if line.starts_with("ORDER") { - metadata.order_lines.push(line); - } else if !line.trim().starts_with("FORMAT") { - if !atom_section_started { - metadata.header_lines.push(line); - } else { - metadata.footer_lines.push(line); + let record = line.get(0..6).unwrap_or("").trim(); + match record { + "ATOM" | "HETATM" => { + atom_section_started = true; + atom_lines.push((i + 1, line)); + } + "CONECT" => conect_lines.push((i + 1, line)), + "ORDER" => metadata.order_lines.push(line), + "FORMAT" => {} + _ => { + if !atom_section_started { + metadata.header_lines.push(line); + } else { + metadata.footer_lines.push(line); + } } } } - let mut current_chain_id = '\0'; - let mut current_res_id = isize::MIN; + let mut current_chain_id: Option = None; + let mut current_residue_id: Option = None; - for (i, line) in atom_lines.iter().enumerate() { - let (record_type, serial, name, res_name, chain_id, res_id, pos, charge, ff_type) = - parse_atom_line(line).map_err(|msg| BgfError::Parse { - line_num: i + 1, - message: msg, - })?; + for (line_num, line) in atom_lines { + let atom_info = parse_atom_line(&line).map_err(|msg| BgfError::Parse { + line_num, + message: msg, + })?; + + let chain_char = atom_info.chain_char; + let mut new_chain = false; + if let Some(id) = current_chain_id { + if system.chain(id).unwrap().id != chain_char { + new_chain = true; + } + } else { + new_chain = true; + } - if chain_id != current_chain_id { - let chain_type = match record_type.as_str() { + if new_chain { + let chain_type = match atom_info.record_type.as_str() { "ATOM" => ChainType::Protein, - "HETATM" => match res_name.as_str() { + "HETATM" => match atom_info.res_name.as_str() { "HOH" | "TIP3" => ChainType::Water, _ => ChainType::Ligand, }, _ => ChainType::Other, }; - builder.start_chain(chain_id, chain_type); - current_chain_id = chain_id; - current_res_id = isize::MIN; + current_chain_id = Some(system.add_chain(chain_char, chain_type)); + current_residue_id = None; } - if res_id != current_res_id { - builder.start_residue(res_id, &res_name); - current_res_id = res_id; + let active_chain_id = current_chain_id.unwrap(); + + let res_seq = atom_info.res_seq; + let mut new_residue = false; + if let Some(id) = current_residue_id { + if system.residue(id).unwrap().id != res_seq { + new_residue = true; + } + } else { + new_residue = true; + } + + if new_residue { + current_residue_id = + system.add_residue(active_chain_id, res_seq, &atom_info.res_name); } - builder.add_atom(serial, &name, &res_name, pos, Some(charge), Some(&ff_type)); + let active_residue_id = current_residue_id + .ok_or_else(|| BgfError::Logic("Failed to create or find residue".to_string()))?; + + let mut atom = Atom::new( + atom_info.serial, + &atom_info.name, + active_residue_id, + atom_info.pos, + ); + atom.partial_charge = atom_info.charge; + atom.force_field_type = atom_info.ff_type; + + system + .add_atom_to_residue(active_residue_id, atom) + .ok_or_else(|| BgfError::Logic("Failed to add atom".to_string()))?; } - for (i, line) in conect_lines.iter().enumerate() { + for (line_num, line) in conect_lines { let (base_serial, connected_serials) = - parse_conect_line(line).map_err(|msg| BgfError::Parse { - line_num: i + 1, + parse_conect_line(&line).map_err(|msg| BgfError::Parse { + line_num, message: msg, })?; + + let base_id = + system + .find_atom_by_serial(base_serial) + .ok_or_else(|| BgfError::Parse { + line_num, + message: format!("Atom serial {} not found for CONECT record", base_serial), + })?; + for &connected_serial in &connected_serials { - builder.add_bond(base_serial, connected_serial, BondOrder::Single); + let connected_id = + system + .find_atom_by_serial(connected_serial) + .ok_or_else(|| BgfError::Parse { + line_num, + message: format!( + "Atom serial {} not found for CONECT record", + connected_serial + ), + })?; + system.add_bond(base_id, connected_id, BondOrder::Single); } } - Ok((builder.build(), metadata)) + Ok((system, metadata)) } fn write_to( @@ -119,17 +177,25 @@ impl MolecularFile for BgfFile { writeln!(writer, "{}", line)?; } - for chain in system.chains() { + for (_chain_id, chain) in system.chains_iter() { let record_type = match chain.chain_type { ChainType::Protein | ChainType::DNA | ChainType::RNA => "ATOM ", _ => "HETATM", }; - for residue in chain.residues() { - for &atom_index in &residue.atom_indices { - let atom = system.get_atom(atom_index).unwrap(); + for &residue_id in chain.residues() { + let residue = system.residue(residue_id).unwrap(); + for &atom_id in residue.atoms() { + let atom = system.atom(atom_id).unwrap(); let atoms_connected = 0; // TODO: Replace with actual logic (placeholder) let lone_pair = 0; // TODO: Replace with actual logic (placeholder) - let atom_line = format_atom_line(record_type, atom, atoms_connected, lone_pair); + let atom_line = format_atom_line( + record_type, + atom, + residue, + chain, + atoms_connected, + lone_pair, + ); writeln!(writer, "{}", atom_line)?; } } @@ -139,35 +205,31 @@ impl MolecularFile for BgfFile { writeln!(writer, "{}", line)?; } - let mut bond_map: HashMap> = HashMap::new(); + let mut bond_serials_written: HashMap> = HashMap::new(); for bond in system.bonds() { - bond_map - .entry(bond.atom1_idx) - .or_default() - .push(bond.atom2_idx); - bond_map - .entry(bond.atom2_idx) - .or_default() - .push(bond.atom1_idx); + let atom1 = system.atom(bond.atom1_id).unwrap(); + let atom2 = system.atom(bond.atom2_id).unwrap(); + + if atom1.serial < atom2.serial { + bond_serials_written + .entry(atom1.serial) + .or_default() + .push(atom2.serial); + } else if atom2.serial < atom1.serial { + bond_serials_written + .entry(atom2.serial) + .or_default() + .push(atom1.serial); + } } - let mut sorted_indices: Vec<_> = bond_map.keys().copied().collect(); - sorted_indices.sort_unstable(); - - for &atom_index in &sorted_indices { - let base_serial = system.get_atom(atom_index).unwrap().serial; - let mut neighbor_serials: Vec<_> = bond_map[&atom_index] - .iter() - .filter(|&&neighbor_idx| neighbor_idx > atom_index) - .map(|&idx| system.get_atom(idx).unwrap().serial) - .collect(); + let mut sorted_base_serials: Vec<_> = bond_serials_written.keys().copied().collect(); + sorted_base_serials.sort_unstable(); - if neighbor_serials.is_empty() { - continue; - } - neighbor_serials.sort_unstable(); - - for chunk in neighbor_serials.chunks(14) { + for base_serial in sorted_base_serials { + let neighbors = bond_serials_written.get_mut(&base_serial).unwrap(); + neighbors.sort_unstable(); + for chunk in neighbors.chunks(14) { let mut connect_line = format!("CONECT {:>5}", base_serial); for &neighbor_serial in chunk { connect_line.push_str(&format!("{:>6}", neighbor_serial)); @@ -195,29 +257,25 @@ impl MolecularFile for BgfFile { "FORMAT ATOM (a6,1x,i5,1x,a5,1x,a3,1x,a1,1x,a5,3f10.5,1x,a5,i3,i2,1x,f8.5)" .to_string(), ], - order_lines: Vec::new(), - footer_lines: Vec::new(), + ..Default::default() }; Self::write_to(system, &minimal_metadata, writer) } } -fn parse_atom_line( - line: &str, -) -> Result< - ( - String, - usize, - String, - String, - char, - isize, - Point3, - f64, - String, - ), - String, -> { +struct ParsedAtomInfo { + record_type: String, + serial: usize, + name: String, + res_name: String, + chain_char: char, + res_seq: isize, + pos: Point3, + charge: f64, + ff_type: String, +} + +fn parse_atom_line(line: &str) -> Result { let get_slice = |start: usize, end: usize| -> Result<&str, String> { line.get(start..end) .ok_or_else(|| format!("Line is too short for slice {}-{}", start, end)) @@ -230,8 +288,8 @@ fn parse_atom_line( .map_err(|e| format!("Invalid serial: {}", e))?; let name = get_slice(13, 18)?.trim().to_string(); let res_name = get_slice(19, 22)?.trim().to_string(); - let chain_id = get_slice(23, 24)?.chars().next().unwrap_or(' '); - let res_id = get_slice(25, 30)? + let chain_char = get_slice(23, 24)?.chars().next().unwrap_or(' '); + let res_seq = get_slice(25, 30)? .trim() .parse() .map_err(|e| format!("Invalid residue ID: {}", e))?; @@ -250,34 +308,36 @@ fn parse_atom_line( let ff_type = get_slice(61, 66)?.trim().to_string(); let charge = get_slice(72, 80)?.trim().parse().unwrap_or(0.0); - Ok(( + Ok(ParsedAtomInfo { record_type, serial, name, res_name, - chain_id, - res_id, - Point3::new(x, y, z), + chain_char, + res_seq, + pos: Point3::new(x, y, z), charge, ff_type, - )) + }) } fn parse_conect_line(line: &str) -> Result<(usize, Vec), String> { - let parts: Vec<_> = line.split_whitespace().collect(); - if parts.len() < 2 || parts[0] != "CONECT" { - return Err("Invalid CONECT line format".to_string()); - } - let base_serial = parts[1] - .parse() - .map_err(|e| format!("Invalid base serial in CONECT: {}", e))?; - let connected_serials = parts[2..].iter().filter_map(|s| s.parse().ok()).collect(); + let mut parts = line.split_whitespace().skip(1).map(str::parse); + let base_serial = parts + .next() + .ok_or("Missing base serial in CONECT")? + .map_err(|e| format!("Invalid base serial: {}", e))?; + let connected_serials = parts + .collect::, _>>() + .map_err(|e| format!("Invalid connected serial: {}", e))?; Ok((base_serial, connected_serials)) } fn format_atom_line( record_type: &str, - atom: &crate::core::models::atom::Atom, + atom: &Atom, + residue: &crate::core::models::residue::Residue, + chain: &crate::core::models::chain::Chain, atoms_connected: u8, lone_pair: u8, ) -> String { @@ -286,9 +346,9 @@ fn format_atom_line( record_type, atom.serial, atom.name, - atom.res_name, - atom.chain_id, - atom.res_id, + residue.name, + chain.id, + residue.id, atom.position.x, atom.position.y, atom.position.z, @@ -298,202 +358,3 @@ fn format_atom_line( atom.partial_charge ) } - -#[cfg(test)] -mod tests { - use super::*; - use crate::core::models::atom::{Atom, AtomFlags}; - use std::io::Cursor; - - fn create_test_system() -> MolecularSystem { - let mut builder = MolecularSystemBuilder::new(); - builder - .start_chain('A', ChainType::Protein) - .start_residue(1, "GLY") - .add_atom( - 1, - "N", - "GLY", - Point3::new(-0.416, -0.535, 0.0), - Some(-0.35), - Some("N_3"), - ) - .add_atom( - 2, - "H", - "GLY", - Point3::new(0.564, -0.535, 0.0), - Some(0.27), - Some("H_"), - ) - .add_bond(1, 2, BondOrder::Single); - builder.build() - } - - #[test] - fn read_from_parses_valid_bgf_file() { - let bgf_data = r#"BIOGRF 332 -FORMAT ATOM (a6,1x,i5,1x,a5,1x,a3,1x,a1,1x,a5,3f10.5,1x,a5,i3,i2,1x,f8.5) -ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000 -ATOM 2 H GLY A 1 0.56400 -0.53500 0.00000 H_ 1 1 0.27000 -CONECT 1 2 -END -"#; - let mut reader = Cursor::new(bgf_data); - let result = BgfFile::read_from(&mut reader); - assert!(result.is_ok()); - let (system, metadata) = result.unwrap(); - - assert_eq!(system.atoms().len(), 2); - assert_eq!(system.chains().len(), 1); - assert_eq!(system.chains()[0].residues().len(), 1); - assert_eq!(system.bonds().len(), 1); - assert_eq!(metadata.header_lines.len(), 1); - assert_eq!(metadata.header_lines[0], "BIOGRF 332"); - assert_eq!(system.get_atom_by_serial(1).unwrap().name, "N"); - assert_eq!(system.get_atom_by_serial(2).unwrap().partial_charge, 0.27); - } - - #[test] - fn read_from_handles_hetatm_and_water_chain() { - let bgf_data = "HETATM 1 O HOH A 1 0.00000 0.00000 0.00000 O_3 1 2 -0.83400\nEND\n"; - let mut reader = Cursor::new(bgf_data); - let (system, _) = BgfFile::read_from(&mut reader).unwrap(); - assert_eq!(system.chains().len(), 1); - assert_eq!(system.chains()[0].chain_type, ChainType::Water); - } - - #[test] - fn read_from_handles_multiple_chains_and_residues() { - let bgf_data = r#" -ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000 -ATOM 2 CA ALA A 2 1.00000 1.00000 1.00000 C_3 1 4 0.10000 -ATOM 3 N SER B 1 2.00000 2.00000 2.00000 N_3 1 3 -0.35000 -END -"#; - let mut reader = Cursor::new(bgf_data); - let (system, _) = BgfFile::read_from(&mut reader).unwrap(); - assert_eq!(system.chains().len(), 2); - assert_eq!(system.chains()[0].residues().len(), 2); - assert_eq!(system.chains()[1].residues().len(), 1); - assert_eq!(system.atoms().len(), 3); - } - - #[test] - fn read_from_returns_error_for_malformed_atom_line() { - let bgf_data = "ATOM X N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000\n"; - let mut reader = Cursor::new(bgf_data); - let result = BgfFile::read_from(&mut reader); - assert!(matches!(result, Err(BgfError::Parse { line_num: 1, .. }))); - } - - #[test] - fn read_from_ignores_malformed_conect_line_parts() { - let bgf_data = "ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000\nCONECT 1 X\nEND\n"; - let mut reader = Cursor::new(bgf_data); - let (system, _) = BgfFile::read_from(&mut reader).unwrap(); - assert!(system.bonds().is_empty()); - } - - #[test] - fn write_to_produces_correct_bgf_format() { - let system = create_test_system(); - let metadata = BgfMetadata { - header_lines: vec!["BIOGRF 332".to_string()], - order_lines: vec![], - footer_lines: vec!["END".to_string()], - }; - let mut buffer = Vec::new(); - let result = BgfFile::write_to(&system, &metadata, &mut buffer); - assert!(result.is_ok()); - - let output = String::from_utf8(buffer).unwrap(); - assert!(output.starts_with("BIOGRF 332")); - assert!(output.contains( - "ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 0 0 -0.35000" - )); - assert!(output.contains("CONECT 1 2")); - assert!(output.ends_with("END\n")); - } - - #[test] - fn write_system_to_uses_default_metadata() { - let system = create_test_system(); - let mut buffer = Vec::new(); - let result = BgfFile::write_system_to(&system, &mut buffer); - assert!(result.is_ok()); - - let output = String::from_utf8(buffer).unwrap(); - assert!(output.contains("BIOGRF 332")); - assert!(output.contains("FORCEFIELD DREIDING")); - } - - #[test] - fn parse_atom_line_succeeds_on_valid_input() { - let line = - "ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000"; - let result = parse_atom_line(line); - assert!(result.is_ok()); - let (_, serial, name, _, _, _, _, charge, ff_type) = result.unwrap(); - assert_eq!(serial, 1); - assert_eq!(name, "N"); - assert_eq!(charge, -0.35); - assert_eq!(ff_type, "N_3"); - } - - #[test] - fn parse_atom_line_fails_on_invalid_number() { - let line = "ATOM X N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000"; - let result = parse_atom_line(line); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err(), - "Invalid serial: invalid digit found in string".to_string() - ); - } - - #[test] - fn parse_conect_line_succeeds_on_valid_input() { - let line = "CONECT 1 2 3 4"; - let result = parse_conect_line(line); - assert!(result.is_ok()); - let (base, connected) = result.unwrap(); - assert_eq!(base, 1); - assert_eq!(connected, vec![2, 3, 4]); - } - - #[test] - fn parse_conect_line_fails_on_invalid_format() { - let line = "CONNECT 1 2"; - let result = parse_conect_line(line); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err(), - "Invalid CONECT line format".to_string() - ); - } - - #[test] - fn format_atom_line_produces_correct_string() { - let atom = Atom { - index: 1, - serial: 1, - name: "CA".to_string(), - res_name: "ALA".to_string(), - res_id: 5, - chain_id: 'A', - force_field_type: "C_3".to_string(), - partial_charge: 0.123, - position: Point3::new(1.1, 2.2, 3.3), - flags: AtomFlags::empty(), - delta: 0.0, - vdw_radius: 1.0, - vdw_well_depth: 0.1, - hbond_type_id: 0, - }; - let line = format_atom_line("ATOM ", &atom, 4, 0); - let expected = - "ATOM 1 CA ALA A 5 1.10000 2.20000 3.30000 C_3 4 0 0.12300"; - assert_eq!(line, expected); - } -} From 8eee55a75d0fca4a5f87f0a4a4d63a8e1c885837 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Wed, 2 Jul 2025 23:46:55 -0700 Subject: [PATCH 18/20] test(core): Add unit tests for BgfFile read_from and write_to methods --- crates/scream-core/src/core/io/bgf.rs | 216 ++++++++++++++++++++++++++ 1 file changed, 216 insertions(+) diff --git a/crates/scream-core/src/core/io/bgf.rs b/crates/scream-core/src/core/io/bgf.rs index 400b8229..c7a428d5 100644 --- a/crates/scream-core/src/core/io/bgf.rs +++ b/crates/scream-core/src/core/io/bgf.rs @@ -263,6 +263,7 @@ impl MolecularFile for BgfFile { } } +#[derive(Debug, Clone)] struct ParsedAtomInfo { record_type: String, serial: usize, @@ -358,3 +359,218 @@ fn format_atom_line( atom.partial_charge ) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::models::chain::Chain; + use crate::core::models::ids::{ChainId, ResidueId}; + use crate::core::models::residue::Residue; + use std::io::Cursor; + + fn create_test_system() -> MolecularSystem { + let mut system = MolecularSystem::new(); + let chain_a = system.add_chain('A', ChainType::Protein); + let res1 = system.add_residue(chain_a, 1, "GLY").unwrap(); + let mut atom_n = Atom::new(1, "N", res1, Point3::new(-0.416, -0.535, 0.0)); + atom_n.partial_charge = -0.35; + atom_n.force_field_type = "N_3".to_string(); + let mut atom_h = Atom::new(2, "H", res1, Point3::new(0.564, -0.535, 0.0)); + atom_h.partial_charge = 0.27; + atom_h.force_field_type = "H_".to_string(); + + let id_n = system.add_atom_to_residue(res1, atom_n).unwrap(); + let id_h = system.add_atom_to_residue(res1, atom_h).unwrap(); + system.add_bond(id_n, id_h, BondOrder::Single); + system + } + + const TEST_BGF_DATA: &str = r#"BIOGRF 332 +FORMAT ATOM (a6,1x,i5,1x,a5,1x,a3,1x,a1,1x,a5,3f10.5,1x,a5,i3,i2,1x,f8.5) +ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000 +ATOM 2 H GLY A 1 0.56400 -0.53500 0.00000 H_ 1 1 0.27000 +CONECT 1 2 +END +"#; + + #[test] + fn read_from_parses_valid_bgf_file() { + let mut reader = Cursor::new(TEST_BGF_DATA); + let result = BgfFile::read_from(&mut reader); + assert!(result.is_ok()); + let (system, metadata) = result.unwrap(); + + assert_eq!(system.atoms_iter().count(), 2); + assert_eq!(system.chains_iter().count(), 1); + assert_eq!(system.residues_iter().count(), 1); + assert_eq!(system.bonds().len(), 1); + assert_eq!(metadata.header_lines.len(), 1); + assert_eq!(metadata.header_lines[0], "BIOGRF 332"); + let atom_n_id = system.find_atom_by_serial(1).unwrap(); + let atom_h_id = system.find_atom_by_serial(2).unwrap(); + assert_eq!(system.atom(atom_n_id).unwrap().name, "N"); + assert_eq!(system.atom(atom_h_id).unwrap().partial_charge, 0.27); + } + + #[test] + fn read_from_handles_hetatm_and_water_chain() { + let bgf_data = "HETATM 1 O HOH A 1 0.00000 0.00000 0.00000 O_3 1 2 -0.83400\nEND\n"; + let mut reader = Cursor::new(bgf_data); + let (system, _) = BgfFile::read_from(&mut reader).unwrap(); + let chain_id = system.find_chain_by_id('A').unwrap(); + assert_eq!(system.chain(chain_id).unwrap().chain_type, ChainType::Water); + } + + #[test] + fn read_from_handles_multiple_chains_and_residues() { + let bgf_data = r#" +ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000 +ATOM 2 CA ALA A 2 1.00000 1.00000 1.00000 C_3 1 4 0.10000 +ATOM 3 N SER B 1 2.00000 2.00000 2.00000 N_3 1 3 -0.35000 +END +"#; + let mut reader = Cursor::new(bgf_data); + let (system, _) = BgfFile::read_from(&mut reader).unwrap(); + assert_eq!(system.chains_iter().count(), 2); + assert_eq!(system.atoms_iter().count(), 3); + + let chain_a_id = system.find_chain_by_id('A').unwrap(); + let chain_b_id = system.find_chain_by_id('B').unwrap(); + assert_eq!(system.chain(chain_a_id).unwrap().residues().len(), 2); + assert_eq!(system.chain(chain_b_id).unwrap().residues().len(), 1); + } + + #[test] + fn read_from_returns_error_for_malformed_atom_line() { + let bgf_data = "ATOM X N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000\n"; + let mut reader = Cursor::new(bgf_data); + let result = BgfFile::read_from(&mut reader); + assert!(matches!(result, Err(BgfError::Parse { line_num: 1, .. }))); + } + + #[test] + fn read_from_returns_error_for_malformed_conect_line() { + let bgf_data_1 = "ATOM 1 N GLY A 1 0.0 0.0 0.0 N_3 0 0 0.0\nCONECT 1 X\nEND\n"; + let mut reader_1 = Cursor::new(bgf_data_1); + let result_1 = BgfFile::read_from(&mut reader_1); + assert!(matches!(result_1, Err(BgfError::Parse { line_num: 2, .. }))); + + let bgf_data_2 = "ATOM 1 N GLY A 1 0.0 0.0 0.0 N_3 0 0 0.0\nCONECT 1 99\nEND\n"; + let mut reader_2 = Cursor::new(bgf_data_2); + let result_2 = BgfFile::read_from(&mut reader_2); + assert!( + matches!(result_2, Err(BgfError::Parse { line_num: 2, message: m, .. }) if m.contains("Atom serial 99 not found")) + ); + } + + #[test] + fn write_to_produces_correct_bgf_format() { + let system = create_test_system(); + let metadata = BgfMetadata { + header_lines: vec!["BIOGRF 332".to_string()], + ..Default::default() + }; + let mut buffer = Vec::new(); + let result = BgfFile::write_to(&system, &metadata, &mut buffer); + assert!(result.is_ok()); + + let output = String::from_utf8(buffer).unwrap(); + assert!(output.starts_with("BIOGRF 332")); + assert!(output.contains("ATOM 1 N GLY A 1")); + assert!(output.contains("ATOM 2 H GLY A 1")); + assert!(output.contains("CONECT 1 2")); + assert!(output.trim_end().ends_with("END")); + } + + #[test] + fn write_system_to_uses_default_metadata() { + let system = create_test_system(); + let mut buffer = Vec::new(); + let result = BgfFile::write_system_to(&system, &mut buffer); + assert!(result.is_ok()); + + let output = String::from_utf8(buffer).unwrap(); + assert!(output.contains("BIOGRF 332")); + assert!(output.contains("FORCEFIELD DREIDING")); + } + + #[test] + fn read_from_handles_duplicate_conect_records() { + let bgf_data = r#"ATOM 1 N GLY A 1 0.0 0.0 0.0 N_3 0 0 0.0 +ATOM 2 C GLY A 1 1.0 0.0 0.0 C_3 0 0 0.0 +CONECT 1 2 +CONECT 2 1 +END +"#; + let mut reader = Cursor::new(bgf_data); + let (system, _) = BgfFile::read_from(&mut reader).unwrap(); + assert_eq!( + system.bonds().len(), + 1, + "Should only create one bond for duplicate CONECT records" + ); + } + + #[test] + fn parse_atom_line_succeeds_on_valid_input() { + let line = + "ATOM 1 N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000"; + let result = parse_atom_line(line); + assert!(result.is_ok()); + let info = result.unwrap(); + assert_eq!(info.serial, 1); + assert_eq!(info.name, "N"); + assert_eq!(info.charge, -0.35); + assert_eq!(info.ff_type, "N_3"); + } + + #[test] + fn parse_atom_line_fails_on_invalid_number() { + let line = "ATOM X N GLY A 1 -0.41600 -0.53500 0.00000 N_3 1 3 -0.35000"; + let result = parse_atom_line(line); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err(), + "Invalid serial: invalid digit found in string".to_string() + ); + } + + #[test] + fn parse_conect_line_succeeds_on_valid_input() { + let line = "CONECT 1 2 3 4"; + let result = parse_conect_line(line); + assert!(result.is_ok()); + let (base, connected) = result.unwrap(); + assert_eq!(base, 1); + assert_eq!(connected, vec![2, 3, 4]); + } + + #[test] + fn parse_conect_line_fails_on_invalid_format() { + let line = "CONECT 1 X"; + let result = parse_conect_line(line); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err(), + "Invalid connected serial: invalid digit found in string".to_string() + ); + } + + #[test] + fn format_atom_line_produces_correct_string() { + let chain_id = ChainId::default(); + let res_id = ResidueId::default(); + let atom = Atom::new(1, "CA", res_id, Point3::new(1.1, 2.2, 3.3)); + let residue = Residue::new(5, "ALA", chain_id); + let chain = Chain::new('A', ChainType::Protein); + + let mut atom_to_format = atom; + atom_to_format.force_field_type = "C_3".to_string(); + atom_to_format.partial_charge = 0.123; + + let line = format_atom_line("ATOM ", &atom_to_format, &residue, &chain, 4, 0); + let expected = + "ATOM 1 CA ALA A 5 1.10000 2.20000 3.30000 C_3 4 0 0.12300"; + assert_eq!(line, expected); + } +} From 1bdc944cd1c8920e370f105a425caf9a3d833fd0 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Thu, 3 Jul 2025 00:14:19 -0700 Subject: [PATCH 19/20] feat(core): Ensure consistent atom ordering in Bond struct initialization --- crates/scream-core/src/core/models/topology.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/scream-core/src/core/models/topology.rs b/crates/scream-core/src/core/models/topology.rs index 2828911e..40ac5f21 100644 --- a/crates/scream-core/src/core/models/topology.rs +++ b/crates/scream-core/src/core/models/topology.rs @@ -59,6 +59,11 @@ pub struct Bond { impl Bond { pub fn new(atom1_id: AtomId, atom2_id: AtomId, order: BondOrder) -> Self { + let (atom1_id, atom2_id) = if atom1_id <= atom2_id { + (atom1_id, atom2_id) + } else { + (atom2_id, atom1_id) + }; Self { atom1_id, atom2_id, From 74b04e13be0e1978517a51b319bc4a2f185f1e30 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Thu, 3 Jul 2025 00:17:14 -0700 Subject: [PATCH 20/20] refactor(core): Simplify chain residue access in add_residue method --- crates/scream-core/src/core/models/system.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/scream-core/src/core/models/system.rs b/crates/scream-core/src/core/models/system.rs index 0c002fff..a7c4da5b 100644 --- a/crates/scream-core/src/core/models/system.rs +++ b/crates/scream-core/src/core/models/system.rs @@ -92,7 +92,8 @@ impl MolecularSystem { self.residues.insert(residue) }); - let chain_residues = &mut self.chains.get_mut(chain_id).unwrap().residues; + let chain = self.chains.get_mut(chain_id).unwrap(); + let chain_residues = &mut chain.residues; if !chain_residues.contains(&residue_id) { chain_residues.push(residue_id); }