Skip to content

Commit

Permalink
add support for multiple selections
Browse files Browse the repository at this point in the history
  • Loading branch information
@rvhonorato
rvhonorato committed Aug 12, 2024
1 parent 15b669a commit 5d50aee
Show file tree
Hide file tree
Showing 2 changed files with 174 additions and 77 deletions.
182 changes: 121 additions & 61 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,14 +49,10 @@ enum Commands {
input: 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")]
spacing: f64,
#[arg(help = "Size in x dimension")]
x_size: f64,
#[arg(help = "Size in y dimension")]
y_size: f64,
#[arg(help = "Spacing between two beads in Angstrom", default_value = "2.0")]
grid_spacing: f64,
#[arg(help = "Size in xy dimension", default_value = "10")]
grid_size: usize,
},
}

Expand Down Expand Up @@ -90,12 +86,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
Commands::Z {
input,
residues,
output,
spacing,
x_size,
y_size,
grid_size,
grid_spacing,
} => {
let _ = generate_z_restraints(input, residues, output, spacing, x_size, y_size);
let _ = generate_z_restraints(input, residues, grid_size, grid_spacing);
}
}

Expand Down Expand Up @@ -240,13 +234,86 @@ fn list_interface(input_file: &str, cutoff: &f64) -> Result<(), Box<dyn Error>>
Ok(())
}

// fn generate_z_restraints(
// input_file: &str,
// selections: &[Vec<isize>],
// grid_size: &usize,
// grid_spacing: &f64,
// ) -> Result<(), Box<dyn Error>> {
// let pdb = match pdbtbx::open_pdb(input_file, pdbtbx::StrictnessLevel::Loose) {
// Ok((pdb, _warnings)) => pdb,
// Err(e) => {
// panic!("Error opening PDB file: {:?}", e);
// }
// };

// let atoms1: Vec<pdbtbx::Atom>;
// let atoms2: Vec<pdbtbx::Atom>;

// if selections.len() >= 2 {
// (atoms1, atoms2) = structure::find_furthest_selections(selections, &pdb);
// } else {
// atoms1 = structure::get_atoms_from_resnumbers(&pdb, &selections[0]);
// atoms2 = vec![pdbtbx::Atom::new(
// false, // hetero
// 1, // serial_number
// "CA", // atom_name
// 0.0, // x
// 0.0, // y
// 0.0, // z
// 1.0, // occupancy
// 0.0, // b_factor
// "C", // element
// 0, // charge
// )
// .expect("Failed to create atom")];
// }

// 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
// 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);

// let atoms3 = structure::get_atoms_from_resnumbers(&pdb, &selections[2]);
// let center_atoms3 = structure::calculate_geometric_center(&atoms3);
// let grid_beads3 = structure::generate_grid_beads(center_atoms3.z, *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);
// all_beads.extend(grid_beads3);

// structure::write_beads_pdb(&all_beads, "z_beads.pdb")?;

// Ok(())
// }

fn generate_z_restraints(
input_file: &str,
_residues: &[Vec<isize>],
_output_file: &str,
_spacing: &f64,
_x_size: &f64,
_y_size: &f64,
selections: &[Vec<isize>],
grid_size: &usize,
grid_spacing: &f64,
) -> Result<(), Box<dyn Error>> {
let pdb = match pdbtbx::open_pdb(input_file, pdbtbx::StrictnessLevel::Loose) {
Ok((pdb, _warnings)) => pdb,
Expand All @@ -255,61 +322,54 @@ fn generate_z_restraints(
}
};

// 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 atoms1: Vec<pdbtbx::Atom>;
let atoms2: Vec<pdbtbx::Atom>;

if selections.len() >= 2 {
(atoms1, atoms2) = structure::find_furthest_selections(selections, &pdb);
} else {
atoms1 = structure::get_atoms_from_resnumbers(&pdb, &selections[0]);
atoms2 = vec![pdbtbx::Atom::new(
false, // hetero
1, // serial_number
"CA", // atom_name
0.0, // x
0.0, // y
0.0, // z
1.0, // occupancy
0.0, // b_factor
"C", // element
0, // charge
)
.expect("Failed to create atom")];
}
// ----------------------------------------------------------------------------------------

let _z_axis = structure::calculate_z_axis(&atoms1, &atoms2);
// 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
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);
// Generate grids centered around the geometric centers of each selection
let grid_beads1 = structure::generate_grid_beads(center1.z, *grid_size, *grid_spacing);
let grid_beads2 = structure::generate_grid_beads(center2.z, *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);

// If there's a third selection, generate a grid for it too
if selections.len() > 2 {
let atoms3 = structure::get_atoms_from_resnumbers(&pdb, &selections[2]);
let center3 = structure::calculate_geometric_center(&atoms3);
let grid_beads3 = structure::generate_grid_beads(center3.z, *grid_size, *grid_spacing);
all_beads.extend(grid_beads3);
}

structure::write_beads_pdb(&all_beads, "z_beads.pdb")?;

Ok(())
}

// sele s1, resid 19+83+145+167 and protein
// sele s2, resid 98+101+126+129 and protein
// sele s3, resid 23+62+87+111+116+153+163 and protein
69 changes: 53 additions & 16 deletions src/structure.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
use std::collections::{HashMap, HashSet};

use itertools::Itertools;
use kd_tree::KdTree;
use nalgebra::Vector3;
use pdbtbx::Residue;
use pdbtbx::{Atom, PDB};
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::SeedableRng;
Expand Down Expand Up @@ -221,11 +223,11 @@ pub fn create_iter_body_gaps(
pairs
}

pub fn calculate_z_axis(selection1: &[pdbtbx::Atom], selection2: &[pdbtbx::Atom]) -> Vector3<f64> {
let center1 = calculate_geometric_center(selection1);
let center2 = calculate_geometric_center(selection2);
(center2 - center1).normalize()
}
// pub fn calculate_z_axis(selection1: &[pdbtbx::Atom], selection2: &[pdbtbx::Atom]) -> Vector3<f64> {
// let center1 = calculate_geometric_center(selection1);
// let center2 = calculate_geometric_center(selection2);
// (center2 - center1).normalize()
// }

pub fn calculate_geometric_center(atoms: &[pdbtbx::Atom]) -> Vector3<f64> {
let mut center = Vector3::zeros();
Expand All @@ -235,18 +237,18 @@ pub fn calculate_geometric_center(atoms: &[pdbtbx::Atom]) -> Vector3<f64> {
center / (atoms.len() as f64)
}

pub fn generate_axis_beads(start_z: f64, end_z: f64, num_beads: usize) -> Vec<Bead> {
let mut beads = Vec::with_capacity(num_beads);
let length = end_z - start_z;
// pub fn generate_axis_beads(start_z: f64, end_z: f64, num_beads: usize) -> Vec<Bead> {
// let mut beads = Vec::with_capacity(num_beads);
// let length = end_z - start_z;

for i in 0..num_beads {
let t = (i as f64) / ((num_beads - 1) as f64);
let z = start_z + t * length;
let position = Vector3::new(0.0, 0.0, z);
beads.push(Bead { position });
}
beads
}
// for i in 0..num_beads {
// let t = (i as f64) / ((num_beads - 1) as f64);
// let z = start_z + t * length;
// let position = Vector3::new(0.0, 0.0, z);
// beads.push(Bead { position });
// }
// beads
// }

pub fn generate_grid_beads(center_z: f64, grid_size: usize, spacing: f64) -> Vec<Bead> {
let mut beads = Vec::with_capacity(grid_size * grid_size);
Expand Down Expand Up @@ -287,6 +289,41 @@ pub fn write_beads_pdb(beads: &[Bead], filename: &str) -> std::io::Result<()> {
Ok(())
}

pub fn find_furthest_selections(selections: &[Vec<isize>], pdb: &PDB) -> (Vec<Atom>, Vec<Atom>) {
let sele: HashMap<usize, (Vec<Atom>, Vector3<f64>)> = selections
.iter()
.enumerate()
.map(|(i, sel)| {
let atoms = get_atoms_from_resnumbers(pdb, sel);
let center = calculate_geometric_center(&atoms);
(i, (atoms, center))
})
.collect();

// Find the two selections that are the furthest apart
let mut max_distance = 0.0;
let mut atoms1 = Vec::new();
let mut atoms2 = Vec::new();
for (i, j) in (0..selections.len()).tuple_combinations() {
let distance = (sele[&i].1 - sele[&j].1).norm();
if distance > max_distance {
max_distance = distance;
atoms1 = sele[&i].0.clone();
atoms2 = sele[&j].0.clone();
}
}

(atoms1, atoms2)
}

pub fn get_atoms_from_resnumbers(pdb: &PDB, selection: &[isize]) -> Vec<Atom> {
pdb.residues()
.filter(|res| selection.contains(&res.serial_number()))
.flat_map(|res| res.atoms())
.cloned() // This clones each &Atom into an owned Atom
.collect()
}

#[cfg(test)]
mod tests {

Expand Down

0 comments on commit 5d50aee

Please sign in to comment.