Export .pml script to visualize restraints in pymol (#47)

* add logic to export `pml` [wip]

* add `write_string_to_file`

* write the `.pml` to a file instead of printing it

* add `--pml <output>` option

* cleaning

* cleaning

src/air.rs

@@ -1,4 +1,6 @@
-use crate::interactor;
+use std::collections::HashSet;
+
+use crate::{interactor, utils};
 
 /// Represents the Air (Ambiguous Interaction Restraints) structure.
 ///
@@ -76,6 +78,50 @@ impl Air {
         }
         Ok(tbl)
     }
+
+    pub fn gen_pml(&self, output_f: &str) {
+        let mut pml = String::new();
+
+        // General settings
+        pml.push_str("set label_size, 0\n");
+        pml.push_str("set dash_gap, 0\n");
+        pml.push_str("set dash_color, yellow\n");
+
+        let mut active: HashSet<(i16, &str)> = HashSet::new();
+        let mut passive: HashSet<(i16, &str)> = HashSet::new();
+
+        for interactor in self.0.iter() {
+            let partners = self.find_partners(interactor);
+
+            if partners.is_empty() {
+                continue;
+            }
+
+            interactor.active().iter().for_each(|r| {
+                active.insert((*r, interactor.chain()));
+            });
+
+            let target_res = interactor::collect_residues(partners);
+
+            target_res.iter().for_each(|r| {
+                let resnum = r.res_number.unwrap_or(0);
+                passive.insert((resnum, r.chain_id));
+            });
+
+            let block = interactor.make_pml_string(target_res);
+            pml.push_str(&block);
+        }
+
+        pml.push_str("color white\n");
+        passive.iter().for_each(|(resnum, chain)| {
+            pml.push_str(format!("color green, (resi {} and chain {})\n", resnum, chain).as_str())
+        });
+        active.iter().for_each(|(resnum, chain)| {
+            pml.push_str(format!("color red, (resi {} and chain {})\n", resnum, chain).as_str())
+        });
+
+        utils::write_string_to_file(&pml, output_f).expect("Could not write pml")
+    }
 }
 
 /// Generates a header string for AIR restraints.

src/interactor.rs

@@ -592,12 +592,44 @@ impl Interactor {
         }
         block
     }
+
+    pub fn make_pml_string(&self, passive_res: Vec<PassiveResidues>) -> String {
+        let mut pml = String::new();
+        let mut _active: Vec<i16> = self.active().iter().cloned().collect();
+        _active.sort();
+
+        let mut passive_res: Vec<PassiveResidues> = passive_res.clone();
+        passive_res.sort_by(|a, b| a.res_number.cmp(&b.res_number));
+
+        for resnum in _active {
+            let identifier = format!("{}-{}", resnum, self.chain);
+            let active_sel = format!("resi {} and name CA and chain {}", resnum, self.chain);
+
+            for passive_resnum in &passive_res {
+                let passive_sel = format!(
+                    "resi {} and name CA and chain {}",
+                    passive_resnum.res_number.unwrap(),
+                    passive_resnum.chain_id
+                );
+
+                pml.push_str(
+                    format!(
+                        "distance {}, ({}), ({})\n",
+                        identifier, active_sel, passive_sel
+                    )
+                    .as_str(),
+                )
+            }
+        }
+
+        pml
+    }
 }
 
 #[derive(Debug, Clone)]
 pub struct PassiveResidues<'a> {
-    chain_id: &'a str,
-    res_number: Option<i16>,
+    pub chain_id: &'a str,
+    pub res_number: Option<i16>,
     wildcard: &'a str,
     // TODO: ADD THE ATOM ATOM NAMES HERE, THEY SHOULD BE USED WHEN GENERATING THE BLOCK
     atom_str: &'a Option<Vec<String>>,

src/main.rs

@@ -3,6 +3,7 @@ mod input;
 mod interactor;
 mod sasa;
 mod structure;
+mod utils;
 use air::Air;
 use core::panic;
 use interactor::Interactor;
@@ -25,25 +26,50 @@ enum Commands {
     Tbl {
         #[arg(help = "Input file")]
         input: String,
+
+        #[arg(
+            long,
+            help = "PyMol Script (.pml) output file",
+            value_name = "output.pml"
+        )]
+        pml: Option<String>,
     },
     #[command(about = "Generate true-interface restraints from a PDB file")]
     Ti {
         #[arg(help = "PDB file")]
         input: String,
         #[arg(help = "Cutoff distance for interface residues")]
         cutoff: f64,
+        #[arg(
+            long,
+            help = "PyMol Script (.pml) output file",
+            value_name = "output.pml"
+        )]
+        pml: Option<String>,
     },
     #[command(about = "Generate unambiguous true-interface restraints from a PDB file")]
     UnambigTi {
         #[arg(help = "PDB file")]
         input: String,
         #[arg(help = "Cutoff distance for interface residues")]
         cutoff: f64,
+        #[arg(
+            long,
+            help = "PyMol Script (.pml) output file",
+            value_name = "output.pml"
+        )]
+        pml: Option<String>,
     },
     #[command(about = "Generate unambiguous restraints to keep molecules together during docking")]
     Restraint {
         #[arg(help = "PDB file")]
         input: String,
+        #[arg(
+            long,
+            help = "PyMol Script (.pml) output file",
+            value_name = "output.pml"
+        )]
+        pml: Option<String>,
     },
     #[command(about = "List residues in the interface")]
     Interface {
@@ -89,17 +115,17 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let cli = Cli::parse();
 
     match &cli.command {
-        Commands::Tbl { input } => {
-            gen_tbl(input);
+        Commands::Tbl { input, pml } => {
+            gen_tbl(input, pml);
         }
-        Commands::Ti { input, cutoff } => {
-            let _ = true_interface(input, cutoff);
+        Commands::Ti { input, cutoff, pml } => {
+            let _ = true_interface(input, cutoff, pml);
         }
-        Commands::UnambigTi { input, cutoff } => {
-            let _ = unambig_ti(input, cutoff);
+        Commands::UnambigTi { input, cutoff, pml } => {
+            let _ = unambig_ti(input, cutoff, pml);
         }
-        Commands::Restraint { input } => {
-            let _ = restraint_bodies(input);
+        Commands::Restraint { input, pml } => {
+            let _ = restraint_bodies(input, pml);
         }
         Commands::Interface { input, cutoff } => {
             let _ = list_interface(input, cutoff);
@@ -156,7 +182,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 /// - `input::read_json_file` for reading the JSON input.
 /// - `Air` struct and its methods for processing the interactors and generating the table.
 /// - Various methods on the `Interactor` struct for processing individual interactors.
-fn gen_tbl(input_file: &str) {
+fn gen_tbl(input_file: &str, pml: &Option<String>) {
     let mut interactors = input::read_json_file(input_file).unwrap();
 
     interactors.iter_mut().for_each(|interactor| {
@@ -178,7 +204,12 @@ fn gen_tbl(input_file: &str) {
     let air = Air::new(interactors);
 
     let tbl = air.gen_tbl().unwrap();
+
     println!("{}", tbl);
+
+    if let Some(output_f) = pml {
+        air.gen_pml(output_f)
+    };
 }
 
 /// Generates Unambiguous Topological Interactions (TIs) from a protein structure.
@@ -208,7 +239,11 @@ fn gen_tbl(input_file: &str) {
 ///
 /// This function will panic if:
 /// - The PDB file cannot be opened or parsed.
-fn unambig_ti(input_file: &str, cutoff: &f64) -> Result<String, Box<dyn Error>> {
+fn unambig_ti(
+    input_file: &str,
+    cutoff: &f64,
+    pml: &Option<String>,
+) -> Result<String, Box<dyn Error>> {
     let pdb = match structure::load_pdb(input_file) {
         Ok(pdb) => pdb,
         Err(e) => {
@@ -246,6 +281,10 @@ fn unambig_ti(input_file: &str, cutoff: &f64) -> Result<String, Box<dyn Error>>
 
     println!("{}", tbl);
 
+    if let Some(output_f) = pml {
+        air.gen_pml(output_f)
+    };
+
     Ok(tbl)
 }
 
@@ -291,7 +330,11 @@ fn unambig_ti(input_file: &str, cutoff: &f64) -> Result<String, Box<dyn Error>>
 /// - `structure::get_true_interface` and `structure::get_chains_in_contact` for interface analysis.
 /// - `Interactor` struct for representing protein chains and their interactions.
 /// - `Air` struct for generating the AIR table.
-fn true_interface(input_file: &str, cutoff: &f64) -> Result<String, Box<dyn Error>> {
+fn true_interface(
+    input_file: &str,
+    cutoff: &f64,
+    pml: &Option<String>,
+) -> Result<String, Box<dyn Error>> {
     // Read PDB file
     let pdb = match structure::load_pdb(input_file) {
         Ok(pdb) => pdb,
@@ -342,6 +385,10 @@ fn true_interface(input_file: &str, cutoff: &f64) -> Result<String, Box<dyn Erro
 
     println!("{}", tbl);
 
+    if let Some(output_f) = pml {
+        air.gen_pml(output_f)
+    };
+
     Ok(tbl)
 }
 
@@ -392,7 +439,7 @@ fn true_interface(input_file: &str, cutoff: &f64) -> Result<String, Box<dyn Erro
 /// - `structure::create_iter_body_gaps` for calculating gaps between bodies.
 /// - `Interactor` struct for representing parts of the protein involved in restraints.
 /// - `Air` struct for generating the AIR table.
-fn restraint_bodies(input_file: &str) -> Result<String, Box<dyn Error>> {
+fn restraint_bodies(input_file: &str, pml: &Option<String>) -> Result<String, Box<dyn Error>> {
     // Read PDB file
     let pdb = match structure::load_pdb(input_file) {
         Ok(pdb) => pdb,
@@ -459,6 +506,9 @@ fn restraint_bodies(input_file: &str) -> Result<String, Box<dyn Error>> {
 
     println!("{}", tbl);
 
+    if let Some(output_f) = pml {
+        air.gen_pml(output_f)
+    };
     Ok(tbl)
 }
 
@@ -688,7 +738,8 @@ assign ( resid 47 and segid B )
 
 "#;
 
-        match true_interface("tests/data/complex.pdb", &3.0) {
+        let opt: Option<String> = None;
+        match true_interface("tests/data/complex.pdb", &3.0, &opt) {
             Ok(tbl) => assert_eq!(tbl, expected_tbl),
             Err(_e) => (),
         };
@@ -725,7 +776,8 @@ assign ( resid 47 and segid B )
 
 "#;
 
-        match true_interface("tests/data/complex_BA.pdb", &3.0) {
+        let opt: Option<String> = None;
+        match true_interface("tests/data/complex_BA.pdb", &3.0, &opt) {
             Ok(tbl) => assert_eq!(tbl, expected_tbl),
             Err(_e) => (),
         };
@@ -739,16 +791,18 @@ assign ( resid 2 and segid A and name CA ) ( resid 8 and segid A and name CA ) 1
 
 ";
 
-        match restraint_bodies("tests/data/gaps.pdb") {
+        let opt: Option<String> = None;
+        match restraint_bodies("tests/data/gaps.pdb", &opt) {
             Ok(tbl) => assert_eq!(tbl, expected_tbl),
             Err(_e) => (),
         }
     }
     #[test]
     fn test_unambigti() {
+        let opt: Option<String> = None;
         let expected_tbl = "assign ( resid 2 and segid A and name CA ) ( resid 10 and segid B and name CA ) 9.1 2.0 0.0\n\n";
 
-        match unambig_ti("tests/data/two_res.pdb", &5.0) {
+        match unambig_ti("tests/data/two_res.pdb", &5.0, &opt) {
             Ok(tbl) => assert_eq!(tbl, expected_tbl),
             Err(_e) => (),
         }

src/utils.rs

@@ -0,0 +1,55 @@
+use std::fs::File;
+use std::io::Write;
+use std::path::Path;
+
+pub fn write_string_to_file(content: &str, file_path: &str) -> std::io::Result<()> {
+    let path = Path::new(file_path);
+
+    let mut file =
+        File::create(path).unwrap_or_else(|e| panic!("Could not create {}: {}", file_path, e));
+
+    file.write_all(content.as_bytes())
+        .unwrap_or_else(|e| panic!("Could not write file {}: {}", file_path, e));
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+    use std::path::Path;
+
+    #[test]
+    fn test_write_string_to_file() {
+        let temp_file = "test_file.txt";
+        let test_content = "test";
+
+        // Clean up any existing test file
+        if Path::new(temp_file).exists() {
+            fs::remove_file(temp_file).expect("Failed to remove existing test file");
+        }
+
+        // Test the function
+        let result = write_string_to_file(test_content, temp_file);
+        assert!(result.is_ok());
+
+        // Verify content was written correctly
+        let read_content = fs::read_to_string(temp_file).expect("Failed to read test file");
+        assert_eq!(read_content, test_content);
+
+        // Clean up
+        fs::remove_file(temp_file).expect("Failed to clean up test file");
+    }
+
+    #[test]
+    fn test_write_string_to_nonexistent_directory() {
+        let invalid_path = "nonexistent_dir/test_file.txt";
+        let test_content = "Hello, world!";
+
+        // This should panic with our custom error message
+        let result = std::panic::catch_unwind(|| write_string_to_file(test_content, invalid_path));
+
+        assert!(result.is_err());
+    }
+}