diff --git a/src/air.rs b/src/air.rs index 4fa12f3..e321020 100644 --- a/src/air.rs +++ b/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. diff --git a/src/interactor.rs b/src/interactor.rs index 6102eea..a929f8d 100644 --- a/src/interactor.rs +++ b/src/interactor.rs @@ -592,12 +592,44 @@ impl Interactor { } block } + + pub fn make_pml_string(&self, passive_res: Vec) -> String { + let mut pml = String::new(); + let mut _active: Vec = self.active().iter().cloned().collect(); + _active.sort(); + + let mut passive_res: Vec = 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, + pub chain_id: &'a str, + pub res_number: Option, wildcard: &'a str, // TODO: ADD THE ATOM ATOM NAMES HERE, THEY SHOULD BE USED WHEN GENERATING THE BLOCK atom_str: &'a Option>, diff --git a/src/main.rs b/src/main.rs index 3d5169c..214957b 100644 --- a/src/main.rs +++ b/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,6 +26,13 @@ enum Commands { Tbl { #[arg(help = "Input file")] input: String, + + #[arg( + long, + help = "PyMol Script (.pml) output file", + value_name = "output.pml" + )] + pml: Option, }, #[command(about = "Generate true-interface restraints from a PDB file")] Ti { @@ -32,6 +40,12 @@ enum Commands { 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, }, #[command(about = "Generate unambiguous true-interface restraints from a PDB file")] UnambigTi { @@ -39,11 +53,23 @@ enum Commands { 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, }, #[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, }, #[command(about = "List residues in the interface")] Interface { @@ -89,17 +115,17 @@ fn main() -> Result<(), Box> { 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> { /// - `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) { 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> { +fn unambig_ti( + input_file: &str, + cutoff: &f64, + pml: &Option, +) -> Result> { 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> 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> /// - `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> { +fn true_interface( + input_file: &str, + cutoff: &f64, + pml: &Option, +) -> Result> { // 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 Result Result> { +fn restraint_bodies(input_file: &str, pml: &Option) -> Result> { // 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> { 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 = 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 = 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 = 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 = 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) => (), } diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..c8d3e04 --- /dev/null +++ b/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()); + } +}