update main

src/main.rs

@@ -47,8 +47,8 @@ enum Commands {
     Z {
         #[arg(help = "Input file")]
         input: String,
-        #[arg(help = "List of comma separated residue indexes")]
-        residues: String,
+        #[arg(long, required = true, help = "Group of residue indexes (can be specified multiple times)", value_parser = parse_residues, number_of_values = 1)]
+        residues: Vec<Vec<isize>>,
         #[arg(help = "Output file")]
         output: String,
         #[arg(help = "Spacing between two beads")]
@@ -60,6 +60,17 @@ enum Commands {
     },
 }
 
+// Parse a comma-separated list of residues
+fn parse_residues(arg: &str) -> Result<Vec<isize>, String> {
+    arg.split(',')
+        .map(|s| {
+            s.trim()
+                .parse::<isize>()
+                .map_err(|_| format!("Invalid number: {}", s))
+        })
+        .collect()
+}
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let cli = Cli::parse();
 
@@ -84,7 +95,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             x_size,
             y_size,
         } => {
-            let _ = generate_z(input, residues, output, spacing, x_size, y_size);
+            let _ = generate_z_restraints(input, residues, output, spacing, x_size, y_size);
         }
     }
 
@@ -229,13 +240,77 @@ fn list_interface(input_file: &str, cutoff: &f64) -> Result<(), Box<dyn Error>>
     Ok(())
 }
 
-fn generate_z(
+fn generate_z_restraints(
     input_file: &str,
-    residues: &str,
+    residues: &Vec<Vec<isize>>,
     output_file: &str,
     spacing: &f64,
     x_size: &f64,
     y_size: &f64,
 ) -> Result<(), Box<dyn Error>> {
-    todo!("Generate Z matrix from input file");
+    let pdb = match pdbtbx::open_pdb(input_file, pdbtbx::StrictnessLevel::Loose) {
+        Ok((pdb, _warnings)) => pdb,
+        Err(e) => {
+            panic!("Error opening PDB file: {:?}", e);
+        }
+    };
+
+    // DEVELOPMENT ----------------------------------------------------------------------------
+    let sele1: Vec<isize> = vec![19, 83, 145, 167];
+    let sele2: Vec<isize> = vec![98, 101, 126, 129];
+    // ----------------------------------------------------------------------------------------
+
+    let target_residues1 = structure::get_residues(&pdb, sele1);
+    let target_residues2 = structure::get_residues(&pdb, sele2);
+
+    // Transform the Residues into a Vector of pdbtbx::Atom
+    let mut atoms1: Vec<pdbtbx::Atom> = Vec::new();
+    for residue in target_residues1.iter() {
+        for atom in residue.atoms() {
+            atoms1.push(atom.clone());
+        }
+    }
+    let mut atoms2: Vec<pdbtbx::Atom> = Vec::new();
+    for residue in target_residues2.iter() {
+        for atom in residue.atoms() {
+            atoms2.push(atom.clone());
+        }
+    }
+
+    let z_axis = structure::calculate_z_axis(&atoms1, &atoms2);
+    let center1 = structure::calculate_geometric_center(&atoms1);
+    let center2 = structure::calculate_geometric_center(&atoms2);
+
+    // Project endpoints onto global Z-axis and center at origin
+    let min_z = center1.z.min(center2.z);
+    let max_z = center1.z.max(center2.z);
+    let mid_z = (min_z + max_z) / 2.0;
+    let half_length = (max_z - min_z) / 2.0;
+
+    // Project endpoints onto global Z-axis
+    let start_z = center1.z.min(center2.z);
+    let end_z = center1.z.max(center2.z);
+    let start = nalgebra::Vector3::new(0.0, 0.0, start_z);
+    let end = nalgebra::Vector3::new(0.0, 0.0, end_z);
+
+    // Generate beads along the global Z-axis
+    let num_axis_beads = 10; // Number of beads along the axis
+    let axis_beads = structure::generate_axis_beads(-half_length, half_length, num_axis_beads);
+
+    // Generate grids at both ends, perpendicular to global Z-axis
+    // Generate grids at both ends, perpendicular to global Z-axis
+    let grid_size = 5; // 5x5 grid
+    let grid_spacing = 2.0; // 2 Angstroms between grid points
+    let grid_beads1 = structure::generate_grid_beads(-half_length, grid_size, grid_spacing);
+    let grid_beads2 = structure::generate_grid_beads(half_length, grid_size, grid_spacing);
+
+    // Combine all beads
+    // let mut all_beads = axis_beads;
+    let mut all_beads = Vec::new();
+    all_beads.extend(grid_beads1);
+    all_beads.extend(grid_beads2);
+
+    structure::write_beads_pdb(&all_beads, "z_beads.pdb")?;
+
+    Ok(())
 }