rust_sasa/utils/
io.rs

1use pdbtbx::PDB;
2use quick_xml::SeError as XmlError;
3use serde_json::Error as JsonError;
4
5use crate::structures::atomic::SASAResult;
6
7pub fn sasa_result_to_json(result: &SASAResult) -> Result<String, JsonError> {
8    serde_json::to_string(result)
9}
10
11pub fn sasa_result_to_xml(result: &SASAResult) -> Result<String, XmlError> {
12    quick_xml::se::to_string(result)
13}
14
15pub fn sasa_result_to_protein_object(
16    original_pdb: &mut PDB,
17    result: &SASAResult,
18) -> Result<(), String> {
19    match result {
20        SASAResult::Atom(v) => {
21            for (i, atom) in original_pdb.atoms_mut().enumerate() {
22                let item = v[i];
23                atom.set_b_factor(item as f64)?;
24            }
25        }
26        SASAResult::Residue(v) => {
27            for (i, residue) in original_pdb.residues_mut().enumerate() {
28                let item = &v[i];
29                assert!(residue.serial_number() == item.serial_number);
30                for atom in residue.atoms_mut() {
31                    atom.set_b_factor(item.value as f64)?;
32                }
33            }
34        }
35        SASAResult::Chain(v) => {
36            for (i, chain) in original_pdb.chains_mut().enumerate() {
37                let id = chain.id().to_string();
38                for residue in chain.residues_mut() {
39                    for atom in residue.atoms_mut() {
40                        assert!(v[i].name == id);
41                        atom.set_b_factor(v[i].value as f64)?;
42                    }
43                }
44            }
45        }
46        SASAResult::Protein(v) => {
47            for residue in original_pdb.residues_mut() {
48                for atom in residue.atoms_mut() {
49                    atom.set_b_factor(v.global_total as f64)?;
50                }
51            }
52        }
53    }
54    Ok(())
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use crate::structures::atomic::{ChainResult, ProteinResult, ResidueResult};
61    use pdbtbx::{Format, ReadOptions};
62    use std::io::{BufReader, Cursor};
63
64    #[test]
65    fn test_sasa_result_to_protein_object_atom() {
66        // Create a simple PDB with 3 atoms
67        let pdb_content = r#"ATOM      1  N   ALA A   1      20.154  16.967  25.000  1.00 10.00           N
68ATOM      2  CA  ALA A   1      19.030  16.155  25.000  1.00 15.00           C
69ATOM      3  C   ALA A   1      17.948  16.712  25.000  1.00 20.00           C
70END
71"#;
72        let mut pdb = ReadOptions::default()
73            .set_format(Format::Pdb)
74            .read_raw(BufReader::new(Cursor::new(pdb_content.as_bytes())))
75            .unwrap()
76            .0;
77
78        // Create atom-level SASA result
79        let sasa_result = SASAResult::Atom(vec![5.0, 10.0, 15.0]);
80
81        // Apply the result to the PDB
82        sasa_result_to_protein_object(&mut pdb, &sasa_result).unwrap();
83
84        // Check that b-factors were set correctly
85        let atoms: Vec<_> = pdb.atoms().collect();
86        assert_eq!(atoms.len(), 3);
87        assert!((atoms[0].b_factor() - 5.0).abs() < 0.001);
88        assert!((atoms[1].b_factor() - 10.0).abs() < 0.001);
89        assert!((atoms[2].b_factor() - 15.0).abs() < 0.001);
90    }
91
92    #[test]
93    fn test_sasa_result_to_protein_object_residue() {
94        // Create a simple PDB with 2 residues
95        let pdb_content = r#"ATOM      1  N   ALA A   1      20.154  16.967  25.000  1.00 10.00           N
96ATOM      2  CA  ALA A   1      19.030  16.155  25.000  1.00 15.00           C
97ATOM      3  N   GLY A   2      17.948  16.712  25.000  1.00 20.00           N
98ATOM      4  CA  GLY A   2      16.500  17.000  25.000  1.00 25.00           C
99END
100"#;
101        let mut pdb = ReadOptions::default()
102            .set_format(Format::Pdb)
103            .read_raw(BufReader::new(Cursor::new(pdb_content.as_bytes())))
104            .unwrap()
105            .0;
106
107        // Create residue-level SASA result
108        let sasa_result = SASAResult::Residue(vec![
109            ResidueResult {
110                serial_number: 1,
111                value: 100.0,
112                name: "ALA".to_string(),
113                is_polar: false,
114                chain_id: "A".to_string(),
115            },
116            ResidueResult {
117                serial_number: 2,
118                value: 200.0,
119                name: "GLY".to_string(),
120                is_polar: false,
121                chain_id: "A".to_string(),
122            },
123        ]);
124
125        // Apply the result to the PDB
126        sasa_result_to_protein_object(&mut pdb, &sasa_result).unwrap();
127
128        // Check that b-factors were set correctly for each residue
129        let residues: Vec<_> = pdb.residues().collect();
130        assert_eq!(residues.len(), 2);
131
132        // First residue atoms should have b-factor 100.0
133        for atom in residues[0].atoms() {
134            assert!((atom.b_factor() - 100.0).abs() < 0.001);
135        }
136
137        // Second residue atoms should have b-factor 200.0
138        for atom in residues[1].atoms() {
139            assert!((atom.b_factor() - 200.0).abs() < 0.001);
140        }
141    }
142
143    #[test]
144    fn test_sasa_result_to_protein_object_chain() {
145        // Create a simple PDB with 2 chains
146        let pdb_content = r#"ATOM      1  N   ALA A   1      20.154  16.967  25.000  1.00 10.00           N
147ATOM      2  CA  ALA A   1      19.030  16.155  25.000  1.00 15.00           C
148ATOM      3  N   GLY B   1      17.948  16.712  25.000  1.00 20.00           N
149ATOM      4  CA  GLY B   1      16.500  17.000  25.000  1.00 25.00           C
150END
151"#;
152        let mut pdb = ReadOptions::default()
153            .set_format(Format::Pdb)
154            .read_raw(BufReader::new(Cursor::new(pdb_content.as_bytes())))
155            .unwrap()
156            .0;
157
158        // Create chain-level SASA result
159        let sasa_result = SASAResult::Chain(vec![
160            ChainResult {
161                name: "A".to_string(),
162                value: 300.0,
163            },
164            ChainResult {
165                name: "B".to_string(),
166                value: 400.0,
167            },
168        ]);
169
170        // Apply the result to the PDB
171        sasa_result_to_protein_object(&mut pdb, &sasa_result).unwrap();
172
173        // Check that b-factors were set correctly for each chain
174        let chains: Vec<_> = pdb.chains().collect();
175        assert_eq!(chains.len(), 2);
176
177        // Chain A atoms should have b-factor 300.0
178        for atom in chains[0].atoms() {
179            assert!((atom.b_factor() - 300.0).abs() < 0.001);
180        }
181
182        // Chain B atoms should have b-factor 400.0
183        for atom in chains[1].atoms() {
184            assert!((atom.b_factor() - 400.0).abs() < 0.001);
185        }
186    }
187
188    #[test]
189    fn test_sasa_result_to_protein_object_protein() {
190        // Create a simple PDB with multiple atoms
191        let pdb_content = r#"ATOM      1  N   ALA A   1      20.154  16.967  25.000  1.00 10.00           N
192ATOM      2  CA  ALA A   1      19.030  16.155  25.000  1.00 15.00           C
193ATOM      3  N   GLY A   2      17.948  16.712  25.000  1.00 20.00           N
194END
195"#;
196        let mut pdb = ReadOptions::default()
197            .set_format(Format::Pdb)
198            .read_raw(BufReader::new(Cursor::new(pdb_content.as_bytes())))
199            .unwrap()
200            .0;
201
202        // Create protein-level SASA result
203        let sasa_result = SASAResult::Protein(ProteinResult {
204            global_total: 500.0,
205            polar_total: 200.0,
206            non_polar_total: 300.0,
207        });
208
209        // Apply the result to the PDB
210        sasa_result_to_protein_object(&mut pdb, &sasa_result).unwrap();
211
212        // Check that all atoms have the same b-factor (global_total)
213        for atom in pdb.atoms() {
214            assert!((atom.b_factor() - 500.0).abs() < 0.001);
215        }
216    }
217}