diff --git a/src/pymolparsing/parsing.rs b/src/pymolparsing/parsing.rs index 868267f..704f6df 100644 --- a/src/pymolparsing/parsing.rs +++ b/src/pymolparsing/parsing.rs @@ -265,6 +265,16 @@ pub struct CoordSet { pub symmetry: Option>, } +impl CoordSet { + /// coords are stored in a 1D vector of x,y,z,x,y,x,z,x,y,z + ///returns as a Vec<[x,y,z]> + pub fn get_coords_as_vec(&self) -> Vec<[f32; 3]> { + self.coord + .chunks(3) + .map(|chunk| [chunk[0], chunk[1], chunk[2]]) + .collect() + } +} /// Custom Value /// /// needed for the settings triplet. @@ -422,6 +432,7 @@ impl PyObjectMolecule { ); atom.unwrap() } + /// Get unique chain names pub fn get_chains(&self) -> Vec { self.atom diff --git a/src/pymolparsing/psedata.rs b/src/pymolparsing/psedata.rs index 967ad1c..488068a 100644 --- a/src/pymolparsing/psedata.rs +++ b/src/pymolparsing/psedata.rs @@ -113,6 +113,7 @@ impl PSEData { // https://github.com/schrodinger/pymol-open-source/blob/master/layer1/Color.cpp#L415 unimplemented!() } + /// session is where all the action happens pub fn get_session_names(&self) -> Vec { self.names diff --git a/tests/test_pse_parsing.rs b/tests/test_pse_parsing.rs index cf18936..e36def7 100644 --- a/tests/test_pse_parsing.rs +++ b/tests/test_pse_parsing.rs @@ -1,5 +1,6 @@ +use itertools::assert_equal; use pseutils::pymolparsing::colors::Color; -use pseutils::pymolparsing::parsing::{CustomValue, SettingsEnum}; +use pseutils::pymolparsing::parsing::{CoordSet, CustomValue, SettingsEnum}; use pseutils::pymolparsing::representation::RepBitmask; use pseutils::PSEData; const TEST_OUTPUT_DIR: &str = "./test_temporary"; @@ -38,6 +39,17 @@ fn test_pdb_00() { let mols = psedata.get_molecule_data(); assert_eq!(mols.len(), 1); + // check coordinates. theres 1 mon and 1519 atoms + let coord_sets: Vec<&CoordSet> = mols.iter().flat_map(|mol| mol.coord_set.iter()).collect(); + assert_eq!(coord_sets.len(), 1); + + let coords_01 = coord_sets[0].get_coords_as_vec(); + assert_eq!(coords_01.len(), 1519); + + let coords_01_atom_01 = coords_01; + assert_equal(coords_01_atom_01[0], [50.873, 32.978, 2.387]); + assert_equal(coords_01_atom_01[1518], [52.372, 15.397, -15.323]); + let atom01 = mols[0].get_atom(0); assert!(atom01.x() == 50.87300109863281); assert!(atom01.y() == 32.97800064086914);