Update

src/main.rs

@@ -7,7 +7,6 @@ use polars::prelude::*;
 use clap::Parser;
 use colored::Colorize;
 
-const EXPOSED_ASA_THRESHOLD: f64 = 0.0; // Units is angstroms squared
 
 #[derive(Debug)]
 struct ASAData {
@@ -23,11 +22,20 @@ struct Args {
     #[arg(short, long)]
     input_path: String,
 
-    /// Path for the CSV to be created
+    /// Optional: Path for the CSV to be created
     #[arg(short, long)]
-    output_path: String,
+    csv_output_path: Option<String>,
+
+    /// Optional: Path for the PDB to be created - Will not be created if path is not specified
+    #[arg(short, long)]
+    pdb_output_path: Option<String>,
+
+    ///  Exposed ASA Threshold - Defaults to 10 angstroms squared (absolute)
+    #[arg(short, long)]
+    exposed_asa_threshold: Option<f64>
 }
 
+// Calculate euclidean distance between points
 fn distance(p1: &(f64, f64, f64), p2: &(f64, f64, f64)) -> f64 {
     let (x1, y1, z1) = p1;
     let (x2, y2, z2) = p2;
@@ -72,18 +80,30 @@ fn calculate_atomic_info_for_residue(residue: &Residue) -> (f64,(f64,f64,f64)) {
 fn main() {
     let args = Args::parse();
 
+    let mut exposed_asa_threshold = 10.0;
+
+    if args.exposed_asa_threshold.is_some() {
+        exposed_asa_threshold = args.exposed_asa_threshold.unwrap();
+    }
+
     if Path::new(&args.input_path).exists() == false {
         println!("{}",format!("ERROR: Input file ({}) does not exist",&args.input_path).red());
         exit(1);
     }
     println!("{}", "WARNING: This program expects the B-factor field of the input PDB file to be the ASA/SASA of each residue. IF THIS IS NOT THE CASE THIS PROGRAM WILL NOT WORK.".yellow());
     let (mut pdb, _errors) = pdbtbx::open(
         args.input_path,
-        StrictnessLevel::Medium
+        StrictnessLevel::Loose
     ).unwrap();
+    if args.csv_output_path.is_none() && args.pdb_output_path.is_none() {
+        println!("{}","You must specify --pdb-output-path and or --csv-output-path".red());
+        exit(1);
+    }
 
     pdb.remove_atoms_by(|atom| atom.element() == Some(&Element::H)); // Remove all H atoms
 
+    let tree = pdb.create_atom_rtree();
+
     // First loop to build list of ASA values for each residue
     let mut asa_values : Vec<ASAData> = Vec::new();
     for residue in pdb.residues() { // Iterate over all residues in the structure
@@ -98,36 +118,54 @@ fn main() {
     // Second loop to generate DPX values for each residue
     let mut dpx_residue_ids: Vec<i32> = vec![];
     let mut dpx_values: Vec<f64> = vec![];
-    for residue in pdb.residues() {
+    for residue in pdb.residues_mut() {
         let (average_asa,position) = calculate_atomic_info_for_residue(residue);
-        if average_asa <= EXPOSED_ASA_THRESHOLD {
+        if average_asa <= exposed_asa_threshold {
             // Atom is not exposed
             sort_by_distance(&mut asa_values,position);
             for asa_value in &asa_values {
                 if asa_value.residue_id == residue.id().0 as i32 {
                     continue
                 }
-                if asa_value.asa > EXPOSED_ASA_THRESHOLD {
+                if asa_value.asa > exposed_asa_threshold {
                     let dpx = distance(&position, &asa_value.position);
                     dpx_residue_ids.push(residue.id().0 as i32);
                     dpx_values.push(dpx);
+                    if args.pdb_output_path.is_some() {
+                        for atom in residue.atoms_mut() {
+                            atom.set_b_factor(dpx).unwrap();
+                        }
+                    }
                     break
                 }
             }
         } else {
             // Atom is exposed
             dpx_residue_ids.push(residue.id().0 as i32);
             dpx_values.push(0.0);
+            if args.pdb_output_path.is_some() {
+                for atom in residue.atoms_mut() {
+                    atom.set_b_factor(0.0).unwrap();
+                }
+            }
         }
     }
-    let mut df: DataFrame = df!(
-        "Residue ID" => &dpx_residue_ids,
-        "DPX" => &dpx_values
-    ).unwrap();
-    let mut file = File::create(&args.output_path).expect("could not create file");
-    println!("{} {}", "Wrote output CSV to ".green(), &args.output_path);
-    CsvWriter::new(&mut file)
-        .include_header(true)
-        .with_separator(b',')
-        .finish(&mut df).unwrap();
+    if args.pdb_output_path.is_some() {
+        let pdb_path = args.pdb_output_path.unwrap();
+        pdbtbx::save(&pdb, &pdb_path, pdbtbx::StrictnessLevel::Loose).unwrap();
+        println!("{}", format!("Wrote output PDB to {} - DPX Values are saved in B-factor column",pdb_path).green());
+    }
+    if args.csv_output_path.is_some() {
+        let csv_path = args.csv_output_path.unwrap();
+        let mut df: DataFrame = df!(
+            "Residue ID" => &dpx_residue_ids,
+            "DPX" => &dpx_values
+        ).unwrap();
+        let mut file = File::create(&csv_path).expect("could not create file");
+        CsvWriter::new(&mut file)
+            .include_header(true)
+            .with_separator(b',')
+            .finish(&mut df).unwrap();
+        println!("{} {}", "Wrote output CSV to ".green(), &csv_path);
+    }
 }