Skip to main content

haddock_restraints/core/
air.rs

1use std::collections::HashSet;
2
3use crate::core::{interactor, utils};
4
5/// Represents the Air (Ambiguous Interaction Restraints) structure.
6///
7/// This struct holds a collection of interactors and provides methods
8/// to find interaction partners and generate tables.
9pub struct Air(Vec<interactor::Interactor>);
10
11impl Air {
12    /// Creates a new `Air` instance with the given interactors.
13    ///
14    /// # Arguments
15    ///
16    /// * `interactors` - A vector of `Interactor` objects.
17    ///
18    /// # Returns
19    ///
20    /// A new `Air` instance.
21    pub fn new(interactors: Vec<interactor::Interactor>) -> Self {
22        Air(interactors)
23    }
24
25    /// Finds potential interaction partners for a given interactor.
26    ///
27    /// # Arguments
28    ///
29    /// * `interactor` - The interactor to find partners for.
30    ///
31    /// # Returns
32    ///
33    /// A vector of references to `Interactor` objects that are potential partners.
34    pub fn find_partners(
35        &self,
36        interactor: &interactor::Interactor,
37    ) -> Vec<&interactor::Interactor> {
38        self.0
39            .iter()
40            .filter(|x| x.target().contains(&interactor.id()) && x.id() != interactor.id())
41            .collect()
42    }
43
44    /// Generates a table representation of the AIR data.
45    ///
46    /// This method validates all interactors and creates a string representation
47    /// of the AIR data, including interaction blocks for each valid interactor
48    /// with partners.
49    ///
50    /// # Returns
51    ///
52    /// * `Ok(String)` - A string containing the generated table.
53    /// * `Err(&str)` - An error message if any interactor is invalid.
54    pub fn gen_tbl(&self) -> Result<String, &str> {
55        let mut tbl = String::new();
56
57        for interactor in self.0.iter() {
58            match interactor.is_valid() {
59                Ok(_) => (),
60                Err(err) => {
61                    eprintln!("## Interactor {} is not valid ##", interactor.id());
62                    return Err(err);
63                }
64            }
65
66            let partners = self.find_partners(interactor);
67
68            if partners.is_empty() {
69                continue;
70            }
71
72            // let header = append_header(i);
73            // tbl.push_str(&header);
74
75            let target_res = interactor::collect_residues(partners);
76            let block = interactor.create_block(target_res);
77            tbl.push_str(&block);
78        }
79        Ok(tbl)
80    }
81
82    pub fn gen_pml(&self, output_f: &str, structure: Option<&str>) {
83        let mut pml = String::new();
84
85        if let Some(structure) = structure {
86            pml.push_str(format!("load {}\n", structure).as_str());
87        }
88
89        // General settings
90        pml.push_str("set label_size, 0\n");
91        pml.push_str("set dash_gap, 0\n");
92        pml.push_str("set dash_color, yellow\n");
93
94        let mut active: HashSet<(i16, &str)> = HashSet::new();
95        let mut passive: HashSet<(i16, &str)> = HashSet::new();
96
97        for interactor in self.0.iter() {
98            let partners = self.find_partners(interactor);
99
100            if partners.is_empty() {
101                continue;
102            }
103
104            interactor.active().iter().for_each(|r| {
105                active.insert((*r, interactor.chain()));
106            });
107
108            let target_res = interactor::collect_residues(partners);
109
110            target_res.iter().for_each(|r| {
111                let resnum = r.res_number.unwrap_or(0);
112                passive.insert((resnum, r.chain_id));
113            });
114
115            let block = interactor.make_pml_string(target_res);
116            pml.push_str(&block);
117        }
118
119        pml.push_str("color white\n");
120        passive.iter().for_each(|(resnum, chain)| {
121            pml.push_str(format!("color green, (resi {} and chain {})\n", resnum, chain).as_str())
122        });
123        active.iter().for_each(|(resnum, chain)| {
124            pml.push_str(format!("color red, (resi {} and chain {})\n", resnum, chain).as_str())
125        });
126
127        utils::write_string_to_file(&pml, output_f).expect("Could not write pml")
128    }
129}
130
131/// Generates a header string for AIR restraints.
132///
133/// This function is currently not used in the implementation (prefixed with _).
134///
135/// # Arguments
136///
137/// * `index` - The index of the selection, used in the header.
138///
139/// # Returns
140///
141/// A formatted string representing the header.
142fn _append_header(index: usize) -> String {
143    format!(
144        "!========================================!\n\
145         ! HADDOCK AIR restraints for selection {} !\n\
146         !========================================!\n",
147        index + 1
148    )
149}