haddock_restraints/core/commands/
z.rs

1use crate::*;
2use std::{collections::HashMap, error::Error};
3
4// TODO: Docstring
5pub fn generate_z_restraints(
6    pdb: pdbtbx::PDB,
7    output_file: &str,
8    selections: &[Vec<isize>],
9    grid_size: &usize,
10    grid_spacing: &f64,
11) -> Result<(), Box<dyn Error>> {
12    // // DEVELOPMENT, move the pdb to the origin --------------------------------------------------
13    // let mut debug_pdb = pdb.clone();
14    // move_to_origin(&mut debug_pdb);
15    // let output_path = Path::new("input.pdb");
16    // let file = File::create(output_path)?;
17    // pdbtbx::save_pdb_raw(
18    //     &debug_pdb,
19    //     BufWriter::new(file),
20    //     pdbtbx::StrictnessLevel::Strict,
21    // );
22    // // -----------------------------------------------------------------------------------------
23
24    let atoms1: Vec<pdbtbx::Atom>;
25    let atoms2: Vec<pdbtbx::Atom>;
26
27    let mut restraints: HashMap<usize, Vec<Bead>> = HashMap::new();
28
29    if selections.len() >= 2 {
30        (atoms1, atoms2) = find_furthest_selections(selections, &pdb);
31    } else {
32        atoms1 = get_atoms_from_resnumbers(&pdb, &selections[0]);
33        atoms2 = vec![
34            pdbtbx::Atom::new(false, 1, "CA", 0.0, 0.0, 0.0, 1.0, 0.0, "C", 0)
35                .expect("Failed to create atom"),
36        ];
37    }
38
39    let center1 = calculate_geometric_center(&atoms1);
40    let center2 = calculate_geometric_center(&atoms2);
41
42    // Project endpoints onto global Z-axis and center at origin
43    let min_z = center1.z.min(center2.z);
44    let max_z = center1.z.max(center2.z);
45    let half_length = (max_z - min_z) / 2.0;
46
47    // Generate grids at both ends, perpendicular to global Z-axis
48    let grid_beads1 = generate_grid_beads(-half_length, *grid_size, *grid_spacing);
49    let grid_beads2 = generate_grid_beads(half_length, *grid_size, *grid_spacing);
50
51    restraints.insert(0, grid_beads1.clone());
52    restraints.insert(1, grid_beads2.clone());
53
54    let mut all_beads = Vec::new();
55    all_beads.extend(grid_beads1);
56    all_beads.extend(grid_beads2);
57
58    // It can be that `selections` contains more than 2 selections, if that's the case, we need to place more grids in between
59    if selections.len() > 2 {
60        for (i, selection) in selections.iter().enumerate().skip(2) {
61            let atoms = get_atoms_from_resnumbers(&pdb, selection);
62            let center = calculate_geometric_center(&atoms);
63            let grid_beads = generate_grid_beads(center.z, *grid_size, *grid_spacing);
64            restraints.insert(i, grid_beads.clone());
65            all_beads.extend(grid_beads);
66        }
67    }
68
69    // Write the beads to a PDB file
70    write_beads_pdb(&all_beads, output_file)?;
71
72    let mut interactors: Vec<Interactor> = Vec::new();
73    let mut counter = 0;
74    let restraint_distance = ((grid_spacing / 2.0) - 2.0).max(2.0);
75    selections
76        .iter()
77        .enumerate()
78        .for_each(|(index, selection)| {
79            let beads = restraints.get(&(index)).unwrap();
80            let z = beads[0].position.z;
81
82            let comparison_operator = if z >= 0.0 { "ge" } else { "le" };
83
84            selection.iter().for_each(|resnum| {
85                let mut interactor_i = Interactor::new(counter);
86                counter += 1;
87                let mut interactor_j = Interactor::new(counter);
88                interactor_j.add_target(counter - 1);
89                interactor_i.add_target(counter);
90                counter += 1;
91
92                interactor_i.set_chain("A");
93                interactor_i.set_active(vec![*resnum as i16]);
94                interactor_i.set_active_atoms(vec!["CA".to_string()]);
95                interactor_i.set_passive_atoms(vec!["SHA".to_string()]);
96                interactor_i.set_target_distance(restraint_distance);
97                interactor_i.set_lower_margin(restraint_distance);
98                interactor_i.set_upper_margin(0.0);
99
100                interactor_j.set_chain("S");
101                interactor_j
102                    .set_wildcard(format!("and attr z {} {:.3}", comparison_operator, z).as_str());
103
104                interactors.push(interactor_i);
105                interactors.push(interactor_j);
106            });
107        });
108
109    let air = Air::new(interactors);
110    let tbl = air.gen_tbl().unwrap();
111
112    println!("{}", tbl);
113
114    Ok(())
115}
116
117#[cfg(test)]
118mod tests {}