ferritin_plms/featurize/
utilities.rs

1use candle_core::{D, DType, Device, IndexOp, Result, Tensor};
2use candle_nn::encoding::one_hot;
3use strum::{Display, EnumIter, EnumString};
4
5#[rustfmt::skip]
6// todo: better utility library
7pub fn aa3to1(aa: &str) -> char {
8    match aa {
9        "ALA" => 'A', "CYS" => 'C', "ASP" => 'D',
10        "GLU" => 'E', "PHE" => 'F', "GLY" => 'G',
11        "HIS" => 'H', "ILE" => 'I', "LYS" => 'K',
12        "LEU" => 'L', "MET" => 'M', "ASN" => 'N',
13        "PRO" => 'P', "GLN" => 'Q', "ARG" => 'R',
14        "SER" => 'S', "THR" => 'T', "VAL" => 'V',
15        "TRP" => 'W', "TYR" => 'Y', _     => 'X',
16    }
17}
18
19#[rustfmt::skip]
20// todo: better utility library
21pub fn aa1to_int(aa: char) -> u32 {
22    match aa {
23        'A' => 0, 'C' => 1, 'D' => 2,
24        'E' => 3, 'F' => 4, 'G' => 5,
25        'H' => 6, 'I' => 7, 'K' => 8,
26        'L' => 9, 'M' => 10, 'N' => 11,
27        'P' => 12, 'Q' => 13, 'R' => 14,
28        'S' => 15, 'T' => 16, 'V' => 17,
29        'W' => 18, 'Y' => 19, _   => 20,
30    }
31}
32
33#[rustfmt::skip]
34pub fn int_to_aa1(aa_int: u32) -> char {
35    match aa_int {
36        0 => 'A', 1 => 'C', 2 => 'D',
37        3 => 'E', 4 => 'F', 5 => 'G',
38        6 => 'H', 7 => 'I', 8 => 'K',
39        9 => 'L', 10 => 'M', 11 => 'N',
40        12 => 'P', 13 => 'Q', 14 => 'R',
41        15 => 'S', 16 => 'T', 17 => 'V',
42        18 => 'W', 19 => 'Y', 20 => 'X',
43        _ => 'X'
44
45    }
46}
47
48#[allow(dead_code)]
49const ALPHABET: [char; 21] = [
50    'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W',
51    'Y', 'X',
52];
53
54#[allow(dead_code)]
55const ELEMENT_LIST: [&str; 118] = [
56    "H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al", "Si", "P", "S", "Cl",
57    "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn", "Ga", "Ge", "As",
58    "Se", "Br", "Kr", "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", "In",
59    "Sn", "Sb", "Te", "I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb",
60    "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl",
61    "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra", "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk",
62    "Cf", "Es", "Fm", "Md", "No", "Lr", "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds", "Rg", "Cn", "Nh",
63    "Fl", "Mc", "Lv", "Ts", "Og",
64];
65
66#[rustfmt::skip]
67#[derive(Debug, Clone, Copy, PartialEq, Display, EnumString, EnumIter)]
68pub enum AAAtom {
69    N = 0,    CA = 1,   C = 2,    CB = 3,   O = 4,
70    CG = 5,   CG1 = 6,  CG2 = 7,  OG = 8,   OG1 = 9,
71    SG = 10,  CD = 11,  CD1 = 12, CD2 = 13, ND1 = 14,
72    ND2 = 15, OD1 = 16, OD2 = 17, SD = 18,  CE = 19,
73    CE1 = 20, CE2 = 21, CE3 = 22, NE = 23,  NE1 = 24,
74    NE2 = 25, OE1 = 26, OE2 = 27, CH2 = 28, NH1 = 29,
75    NH2 = 30, OH = 31,  CZ = 32,  CZ2 = 33, CZ3 = 34,
76    NZ = 35,  OXT = 36,
77    Unknown = -1,
78}
79impl AAAtom {
80    // Get numeric value (might still be useful in some contexts)
81    pub fn to_index(&self) -> usize {
82        *self as usize
83    }
84}
85
86macro_rules! define_residues {
87    ($($name:ident: $code3:expr_2021, $code1:expr_2021, $idx:expr_2021, $features:expr_2021, $atoms14:expr_2021),* $(,)?) => {
88        #[derive(Debug, Copy, Clone)]
89        pub enum Residue {
90            $($name),*
91        }
92
93        impl Residue {
94            pub const fn code3(&self) -> &'static str {
95                match self {
96                    $(Self::$name => $code3),*
97                }
98            }
99            pub const fn code1(&self) -> char {
100                match self {
101                    $(Self::$name => $code1),*
102                }
103            }
104            pub const fn atoms14(&self) -> [AAAtom; 14] {
105                match self {
106                    $(Self::$name => $atoms14),*
107                }
108            }
109            pub fn from_int(value: i32) -> Self {
110                match value {
111                    $($idx => Self::$name,)*
112                    _ => Self::UNK
113                }
114            }
115            pub fn to_int(&self) -> i32 {
116                match self {
117                    $(Self::$name => $idx),*
118                }
119            }
120        }
121    }
122}
123
124define_residues! {
125    ALA: "ALA", 'A', 0,  [1.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
126    CYS: "CYS", 'C', 1,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::SG, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
127    ASP: "ASP", 'D', 2,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::OD1, AAAtom::OD2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
128    GLU: "GLU", 'E', 3,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD, AAAtom::OE1, AAAtom::OE2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
129    PHE: "PHE", 'F', 4,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD1, AAAtom::CD2, AAAtom::CE1, AAAtom::CE2, AAAtom::CZ, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
130    GLY: "GLY", 'G', 5,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
131    HIS: "HIS", 'H', 6,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::ND1, AAAtom::CD2, AAAtom::CE1, AAAtom::NE2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
132    ILE: "ILE", 'I', 7,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG1, AAAtom::CG2, AAAtom::CD1, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
133    LYS: "LYS", 'K', 8,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD, AAAtom::CE, AAAtom::NZ, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
134    LEU: "LEU", 'L', 9,  [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD1, AAAtom::CD2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
135    MET: "MET", 'M', 10, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::SD, AAAtom::CE, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
136    ASN: "ASN", 'N', 11, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::OD1, AAAtom::ND2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
137    PRO: "PRO", 'P', 12, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
138    GLN: "GLN", 'Q', 13, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD, AAAtom::OE1, AAAtom::NE2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
139    ARG: "ARG", 'R', 14, [0.0, 1.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD, AAAtom::NE, AAAtom::CZ, AAAtom::NH1, AAAtom::NH2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
140    SER: "SER", 'S', 15, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::OG, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
141    THR: "THR", 'T', 16, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::OG1, AAAtom::CG2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
142    VAL: "VAL", 'V', 17, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG1, AAAtom::CG2, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
143    TRP: "TRP", 'W', 18, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD1, AAAtom::CD2, AAAtom::CE2, AAAtom::CE3, AAAtom::NE1, AAAtom::CZ2, AAAtom::CZ3, AAAtom::CH2],
144    TYR: "TYR", 'Y', 19, [0.0, 0.0], [AAAtom::N, AAAtom::CA, AAAtom::C, AAAtom::O, AAAtom::CB, AAAtom::CG, AAAtom::CD1, AAAtom::CD2, AAAtom::CE1, AAAtom::CE2, AAAtom::CZ, AAAtom::OH, AAAtom::Unknown, AAAtom::Unknown],
145    UNK: "UNK", 'X', 20, [0.0, 0.0], [AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown, AAAtom::Unknown],
146}
147
148pub fn get_nearest_neighbours(
149    cb: &Tensor,
150    mask: &Tensor,
151    y: &Tensor,
152    y_t: &Tensor,
153    y_m: &Tensor,
154    number_of_ligand_atoms: i64,
155) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
156    // First, remove batch dimension if present using squeeze(0)
157    let cb = cb.squeeze(0)?;
158    let mask = mask.squeeze(0)?;
159    let y_m = if y_m.dims().len() > 1 {
160        y_m.sum_keepdim(1)?.squeeze(1)? // or .any(1)? depending on your needs
161    } else {
162        y_m.clone()
163    };
164    let num_residues = cb.dim(0)?;
165    let mask_cby = mask.unsqueeze(1)?.matmul(&y_m.unsqueeze(0)?)?;
166    let cb_flat = cb.reshape((cb.dim(0)?, 1, 3))?; // [154, 1, 3]
167    let y_flat = y.reshape((1, y.dim(0)?, 3))?; // [1, 54, 3]
168    // Try broadcasting manually if needed
169    let cb_broadcast = cb_flat.broadcast_as((cb.dim(0)?, y.dim(0)?, 3))?; // [154, 54, 3]
170    let y_broadcast = y_flat.broadcast_as((cb.dim(0)?, y.dim(0)?, 3))?; // [154, 54, 3]
171    let diff = cb_broadcast.sub(&y_broadcast)?;
172    let l2_ab = diff.powf(2.0)?.sum(D::Minus1)?;
173    let complement_mask = (mask_cby.neg()? + 1.0)?;
174    let padding_value = Tensor::full(1000.0_f32, mask_cby.dims(), cb.device())?;
175    let masked_distances = l2_ab.mul(&mask_cby)?;
176    let padding_contribution = complement_mask.mul(&padding_value)?;
177    let l2_ab = masked_distances.add(&padding_contribution)?;
178
179    // Get nearest neighbors
180    let nn_idx = l2_ab
181        .arg_sort_last_dim(false)?
182        .narrow(1, 0, number_of_ligand_atoms as usize)?
183        .contiguous()?;
184    let l2_ab_nn = l2_ab.contiguous()?.gather(&nn_idx, 1)?;
185    let d_ab_closest = l2_ab_nn.i((.., 0))?.sqrt()?;
186    let y_new = y
187        .unsqueeze(0)?
188        .expand((num_residues, y.dim(0)?, 3))?
189        .contiguous()?
190        .gather(
191            &nn_idx
192                .unsqueeze(2)?
193                .expand((num_residues, number_of_ligand_atoms as usize, 3))?
194                .contiguous()?,
195            1,
196        )?;
197
198    let y_t_new = y_t
199        .unsqueeze(0)?
200        .expand((num_residues, y_t.dim(0)?))?
201        .contiguous()?
202        .gather(&nn_idx, 1)?;
203
204    let y_m_new = y_m
205        .unsqueeze(0)?
206        .expand((num_residues, y_m.dim(0)?))?
207        .contiguous()?
208        .gather(&nn_idx, 1)?;
209
210    Ok((y_new, y_t_new, y_m_new, d_ab_closest))
211}
212
213pub fn cat_neighbors_nodes(
214    h_nodes: &Tensor,
215    h_neighbors: &Tensor,
216    e_idx: &Tensor,
217) -> Result<Tensor> {
218    let h_nodes_gathered = gather_nodes(h_nodes, e_idx)?;
219    let h_neighbors = h_neighbors.expand((
220        h_neighbors.dim(0)?, // 1
221        h_nodes.dim(1)?,     // 93
222        h_neighbors.dim(2)?, // 24
223        h_neighbors.dim(3)?, // 128
224    ))?;
225
226    Tensor::cat(
227        &[h_neighbors, h_nodes_gathered.to_dtype(DType::F32)?],
228        D::Minus1,
229    )
230}
231
232/// Retrieve the nearest Neighbor of a set of coordinates.
233/// Usually used for CA carbon distance.
234pub fn compute_nearest_neighbors(
235    coords: &Tensor,
236    mask: &Tensor,
237    k: usize,
238    eps: f32,
239) -> Result<(Tensor, Tensor)> {
240    let (_batch_size, seq_len, _) = coords.dims3()?;
241    // broadcast_matmul handles broadcasting automatically
242    // [2, 3, 1] × [2, 1, 3] -> [2, 3, 3]
243
244    let mask_2d = mask
245        .unsqueeze(2)?
246        .broadcast_matmul(&mask.unsqueeze(1)?)?
247        .to_dtype(DType::F32)?;
248    // Compute pairwise distances with broadcasting
249
250    let distances = (coords
251        .unsqueeze(2)?
252        .broadcast_sub(&coords.unsqueeze(1)?)?
253        .powf(2.)?
254        .sum(D::Minus1)?
255        + eps as f64)? // also  doesn't have add
256        .sqrt()?
257        .to_dtype(DType::F32)?;
258
259    // Apply mask
260    // Get max values for adjustment
261    let masked_distances = (&distances * &mask_2d.to_dtype(DType::F32)?)?;
262    // println!("after masked_distances");
263    let d_max = masked_distances.max_keepdim(D::Minus1)?;
264    let mask_term = ((&mask_2d.to_dtype(DType::F32)? * -1.0)? + 1.0)?;
265    let d_adjust = (&masked_distances + mask_term.broadcast_mul(&d_max)?)?;
266    let d_adjust = d_adjust.to_dtype(DType::F32)?;
267    Ok(topk_last_dim(&d_adjust, k.min(seq_len))?)
268}
269
270// https://github.com/huggingface/candle/pull/2375/files#diff-e4d52a71060a80ac8c549f2daffcee77f9bf4de8252ad067c47b1c383c3ac828R957
271pub fn topk_last_dim(xs: &Tensor, topk: usize) -> Result<(Tensor, Tensor)> {
272    let sorted_indices = xs.arg_sort_last_dim(false)?.to_dtype(DType::U32)?;
273    let topk_indices = sorted_indices.narrow(D::Minus1, 0, topk)?.contiguous()?;
274    let gathered = xs.gather(&topk_indices, D::Minus1)?;
275    Ok((gathered, topk_indices))
276}
277
278/// Input coords. Output 1 <batch  x 1 > Tensor
279/// representing whether each residue has all 4 backbone atoms.
280/// note that the internal ordering is different between
281/// backbone only [N/CA/C/O] and all-atom [N/CA/C/CB/O]....
282pub fn create_backbone_mask_37(xyz_37: &Tensor) -> Result<Tensor> {
283    let (b, l, rescount, _) = xyz_37.dims4()?;
284    // Create a vector with 1s at positions 0,1,2,4 and 0s elsewhere
285    let mut values = vec![0f32; rescount];
286    for &idx in &[0, 1, 2, 4] {
287        if idx < rescount {
288            values[idx] = 1.0;
289        }
290    }
291
292    // Create the base mask for one sequence, explicitly specifying the data type
293    let base_mask = Tensor::new(values.as_slice(), xyz_37.device())?.to_dtype(DType::F32)?;
294
295    // Create the full mask by repeating it for each batch and length
296    let mask = base_mask.unsqueeze(0)?.unsqueeze(0)?; // Add batch and length dimensions
297    let mask = mask.broadcast_as((b, l, rescount))?;
298
299    Ok(mask)
300}
301/// Get Pseudo CB
302pub fn calculate_cb(xyz_37: &Tensor) -> Result<Tensor> {
303    // make sure we are dealing with
304    let (_, dim37, dim3) = xyz_37.dims3()?;
305    assert_eq!(dim37, 37);
306    assert_eq!(dim3, 3);
307
308    // Constants for CB calculation
309    let a_coeff = -0.58273431f64;
310    let b_coeff = 0.56802827f64;
311    let c_coeff = -0.54067466f64;
312
313    // Get N, CA, C coordinates
314    let n = xyz_37.i((.., 0, ..))?; // N  at index 0
315    let ca = xyz_37.i((.., 1, ..))?; // CA at index 1
316    let c = xyz_37.i((.., 2, ..))?; // C  at index 2
317
318    // Calculate vectors
319    let b = (&ca - &n)?; // CA - N
320    let c = (&c - &ca)?; // C - CA
321
322    // Manual cross product components
323    // a_x = b_y * c_z - b_z * c_y
324    // a_y = b_z * c_x - b_x * c_z
325    // a_z = b_x * c_y - b_y * c_x
326    let b_x = b.i((.., 0))?;
327    let b_y = b.i((.., 1))?;
328    let b_z = b.i((.., 2))?;
329    let c_x = c.i((.., 0))?;
330    let c_y = c.i((.., 1))?;
331    let c_z = c.i((.., 2))?;
332
333    let a_x = ((&b_y * &c_z)? - (&b_z * &c_y)?)?;
334    let a_y = ((&b_z * &c_x)? - (&b_x * &c_z)?)?;
335    let a_z = ((&b_x * &c_y)? - (&b_y * &c_x)?)?;
336
337    // Stack the cross product components back together
338    let a = Tensor::stack(&[&a_x, &a_y, &a_z], 1)?;
339
340    // Final CB calculation: -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + CA
341    let cb = (&a * a_coeff)? + (&b * b_coeff)? + (&c * c_coeff)? + &ca;
342
343    Ok(cb?)
344}
345
346/// Custom Cross-Product Fn.
347pub fn cross_product(a: &Tensor, b: &Tensor) -> Result<Tensor> {
348    let last_dim = a.dims().len() - 1;
349
350    // Extract components
351    let a0 = a.narrow(last_dim, 0, 1)?;
352    let a1 = a.narrow(last_dim, 1, 1)?;
353    let a2 = a.narrow(last_dim, 2, 1)?;
354
355    let b0 = b.narrow(last_dim, 0, 1)?;
356    let b1 = b.narrow(last_dim, 1, 1)?;
357    let b2 = b.narrow(last_dim, 2, 1)?;
358
359    // Compute cross produAAct components
360    let c0 = ((&a1 * &b2)? - (&a2 * &b1)?)?;
361    let c1 = ((&a2 * &b0)? - (&a0 * &b2)?)?;
362    let c2 = ((&a0 * &b1)? - (&a1 * &b0)?)?;
363
364    // Stack the results
365    Tensor::cat(&[&c0, &c1, &c2], last_dim)
366}
367
368/// Gather_edges
369/// Features [B,N,N,C] at Neighbor indices [B,N,K] => Neighbor features [B,N,K,C]
370pub fn gather_edges(edges: &Tensor, neighbor_idx: &Tensor) -> Result<Tensor> {
371    let (d1, d2, d3) = neighbor_idx.dims3()?;
372    let neighbors =
373        neighbor_idx
374            .unsqueeze(D::Minus1)?
375            .expand((d1, d2, d3, edges.dim(D::Minus1)?))?;
376    let edge_gather = edges.gather(&neighbors, 2)?;
377    Ok(edge_gather)
378}
379
380/// Gather Nodes
381///
382/// Features [B,N,C] at Neighbor indices [B,N,K] => [B,N,K,C]
383/// Flatten and expand indices per batch [B,N,K] => [B,NK] => [B,NK,C]
384pub fn gather_nodes(nodes: &Tensor, neighbor_idx: &Tensor) -> Result<Tensor> {
385    let (batch_size, n_nodes, n_features) = nodes.dims3()?;
386    let (_, _, k_neighbors) = neighbor_idx.dims3()?;
387    // Reshape neighbor_idx to [B, N*K]
388    let neighbors_flat = neighbor_idx.reshape((batch_size, n_nodes * k_neighbors))?;
389    // Add feature dimension and expand
390    let neighbors_flat = neighbors_flat
391        .unsqueeze(2)? // Add feature dimension [B, N*K, 1]
392        .expand((batch_size, n_nodes * k_neighbors, n_features))?; // Expand to [B, N*K, C]
393    // make contiguous for the gather.
394    let neighbors_flat = neighbors_flat.contiguous()?;
395    // Gather features
396    let neighbor_features = nodes.gather(&neighbors_flat, 1)?;
397    // Reshape back to [B, N, K, C]
398    neighbor_features.reshape((batch_size, n_nodes, k_neighbors, n_features))
399}
400
401pub fn gather_nodes_t(nodes: &Tensor, neighbor_idx: &Tensor) -> Result<Tensor> {
402    // Features [B,N,C] at Neighbor index [B,K] => Neighbor features[B,K,C]
403    let (d1, d2, d3) = nodes.dims3()?;
404    let idx_flat = neighbor_idx.unsqueeze(D::Minus1)?.expand((d1, d2, d3))?;
405    nodes.gather(&idx_flat, 1)
406}
407
408#[allow(dead_code)]
409fn get_seq_rec(s: &Tensor, s_pred: &Tensor, mask: &Tensor) -> Result<Tensor> {
410    // S: true sequence shape=[batch, length]
411    // S_pred: predicted sequence shape=[batch, length]
412    // mask: mask to compute average over the region shape=[batch, length]
413    // Returns: averaged sequence recovery shape=[batch]
414    //
415    // Compute the match tensor
416    let match_tensor = s.eq(s_pred)?;
417    let match_f32 = match_tensor.to_dtype(DType::F32)?;
418    let numerator = (match_f32 * mask)?.sum_keepdim(1)?;
419    let denominator = mask.sum_keepdim(1)?;
420    let average = numerator.broadcast_div(&denominator)?;
421    // Remove the last dimension to get shape=[batch]
422    average.squeeze(1)
423}
424
425#[allow(dead_code)]
426fn get_score(s: &Tensor, log_probs: &Tensor, mask: &Tensor) -> Result<(Tensor, Tensor)> {
427    //     S : true sequence shape=[batch, length]
428    //     log_probs : predicted sequence shape=[batch, length]
429    //     mask : mask to compute average over the region shape=[batch, length]
430    //     average_loss : averaged categorical cross entropy (CCE) [batch]
431    //     loss_per_resdue : per position CCE [batch, length]
432
433    //     """
434    //     S_one_hot = torch.nn.functional.one_hot(S, 21)
435    //     loss_per_residue = -(S_one_hot * log_probs).sum(-1)  # [B, L]
436    //     average_loss = torch.sum(loss_per_residue * mask, dim=-1) / (
437    //         torch.sum(mask, dim=-1) + 1e-8
438    //     )
439    //     return average_loss, loss_per_residue
440
441    // S: true sequence shape=[batch, length]
442    // log_probs: predicted sequence shape=[batch, length, 21]
443    // mask: mask to compute average over the region shape=[batch, length]
444    // Returns:
445    //   - average_loss: averaged categorical cross entropy (CCE) [batch]
446    //   - loss_per_residue: per position CCE [batch, length]
447
448    // Create one-hot encoding of S.
449    // see https://docs.rs/candle-nn/0.7.2/candle_nn/encoding/fn.one_hot.html
450    // this could be wrong...
451    let s_one_hot = one_hot(s.clone(), 21, 1., 0.)?;
452    let loss_per_residue = s_one_hot.mul(&log_probs.neg()?)?.sum(D::Minus1)?;
453    let average_loss = loss_per_residue
454        .mul(&mask)?
455        .sum_keepdim(D::Minus1)?
456        .div(&(mask.sum_keepdim(D::Minus1)? + 1e-8f64)?)?
457        .squeeze(D::Minus1)?;
458
459    Ok((average_loss, loss_per_residue))
460}
461
462pub fn linspace(
463    start: f64,
464    stop: f64,
465    steps: usize,
466    device: &Device,
467    return_type: DType,
468) -> Result<Tensor> {
469    if steps == 0 {
470        Tensor::from_vec(Vec::<f64>::new(), steps, device)
471    } else if steps == 1 {
472        Tensor::from_vec(vec![start], steps, device)
473    } else {
474        let delta = (stop - start) / (steps - 1) as f64;
475        let vs = (0..steps)
476            .map(|step| start + step as f64 * delta)
477            .collect::<Vec<_>>();
478        Tensor::from_vec(vs, steps, device)?.to_dtype(return_type)
479    }
480}
481
482pub fn linspace_f32(start: f32, stop: f32, steps: usize, device: &Device) -> Result<Tensor> {
483    match steps {
484        0 => Tensor::from_vec(Vec::<f32>::new(), steps, device),
485        1 => Tensor::from_vec(vec![start], steps, device),
486        _ => {
487            let delta = (stop - start) / (steps - 1) as f32;
488            let vs = (0..steps)
489                .map(|step| start + step as f32 * delta)
490                .collect::<Vec<_>>();
491            Tensor::from_vec(vs, steps, device)
492        }
493    }
494}
495
496#[cfg(test)]
497mod tests {
498    use super::*;
499    use crate::StructureFeatures;
500    use anyhow::Result;
501    use ferritin_core::info::elements::Element;
502    use ferritin_core::load_structure;
503    use ferritin_test_data::TestFile;
504
505    #[test]
506    fn test_residue_codes() {
507        let ala = Residue::ALA;
508        assert_eq!(ala.code3(), "ALA");
509        assert_eq!(ala.code1(), 'A');
510        assert_eq!(ala.to_int(), 0);
511    }
512
513    #[test]
514    fn test_residue_from_int() {
515        assert!(matches!(Residue::from_int(0), Residue::ALA));
516        assert!(matches!(Residue::from_int(1), Residue::CYS));
517        assert!(matches!(Residue::from_int(999), Residue::UNK));
518    }
519
520    #[test]
521    fn test_residue_atoms() {
522        let trp = Residue::TRP;
523        let atoms = trp.atoms14();
524        assert_eq!(atoms[0], AAAtom::N);
525        assert_eq!(atoms[13], AAAtom::CH2);
526
527        let gly = Residue::GLY;
528        let atoms = gly.atoms14();
529        assert_eq!(atoms[4], AAAtom::Unknown);
530    }
531
532    #[test]
533    fn test_atom_backbone_tensor() -> Result<()> {
534        let device = Device::Cpu;
535        let (pdb_file, _temp) = TestFile::protein_01().create_temp()?;
536        let ac = load_structure(pdb_file)?;
537        let ac_backbone_tensor: Tensor = ac.to_numeric_backbone_atoms(&device).expect("REASON");
538        // batch size of 1;154 residues; N/CA/C/O; positions
539        assert_eq!(ac_backbone_tensor.dims(), &[1, 154, 4, 3]);
540
541        // Check my residue coords in the Tensor
542        // ATOM   1    N  N   . MET A 1 1   ? 24.277 8.374   -9.854  1.00 38.41  ? 0   MET A N   1
543        // ATOM   2    C  CA  . MET A 1 1   ? 24.404 9.859   -9.939  1.00 37.90  ? 0   MET A CA  1
544        // ATOM   3    C  C   . MET A 1 1   ? 25.814 10.249  -10.359 1.00 36.65  ? 0   MET A C   1
545        // ATOM   4    O  O   . MET A 1 1   ? 26.748 9.469   -10.197 1.00 37.13  ? 0   MET A O   1
546        let backbone_coords = [
547            // Methionine - AA00
548            ("N", (0, 0, 0, ..), vec![24.277, 8.374, -9.854]),
549            ("CA", (0, 0, 1, ..), vec![24.404, 9.859, -9.939]),
550            ("C", (0, 0, 2, ..), vec![25.814, 10.249, -10.359]),
551            ("O", (0, 0, 3, ..), vec![26.748, 9.469, -10.197]),
552            // Valine - AA01
553            ("N", (0, 1, 0, ..), vec![25.964, 11.453, -10.903]),
554            ("CA", (0, 1, 1, ..), vec![27.263, 11.924, -11.359]),
555            ("C", (0, 1, 2, ..), vec![27.392, 13.428, -11.115]),
556            ("O", (0, 1, 3, ..), vec![26.443, 14.184, -11.327]),
557            // Glycing - AAlast
558            ("N", (0, 153, 0, ..), vec![23.474, -3.227, 5.994]),
559            ("CA", (0, 153, 1, ..), vec![22.818, -2.798, 7.211]),
560            ("C", (0, 153, 2, ..), vec![22.695, -1.282, 7.219]),
561            ("O", (0, 153, 3, ..), vec![21.870, -0.745, 7.992]),
562        ];
563
564        for (atom_name, (b, i, j, k), expected) in backbone_coords {
565            // assert_eq!(ac_backbone_tensor.dims(), &[1, 154, 4, 3])
566            let actual: Vec<f32> = ac_backbone_tensor.i((b, i, j, k))?.to_vec1()?;
567            assert_eq!(actual, expected, "Mismatch for atom {}", atom_name);
568        }
569        Ok(())
570    }
571
572    #[test]
573    fn test_all_atom37_tensor() -> Result<()> {
574        let device = Device::Cpu;
575        let (pdb_file, _temp) = TestFile::protein_01().create_temp()?;
576        let ac = load_structure(pdb_file)?;
577        let ac_backbone_tensor: Tensor = ac.to_numeric_atom37(&device).expect("REASON");
578        assert_eq!(ac_backbone_tensor.dims(), &[1, 154, 37, 3]);
579
580        // Check my residue coords in the Tensor
581        // ATOM   1    N  N   . MET A 1 1   ? 24.277 8.374   -9.854  1.00 38.41  ? 0   MET A N   1
582        // ATOM   2    C  CA  . MET A 1 1   ? 24.404 9.859   -9.939  1.00 37.90  ? 0   MET A CA  1
583        // ATOM   3    C  C   . MET A 1 1   ? 25.814 10.249  -10.359 1.00 36.65  ? 0   MET A C   1
584        // ATOM   4    O  O   . MET A 1 1   ? 26.748 9.469   -10.197 1.00 37.13  ? 0   MET A O   1
585        // ATOM   5    C  CB  . MET A 1 1   ? 24.070 10.495  -8.596  1.00 39.58  ? 0   MET A CB  1
586        // ATOM   6    C  CG  . MET A 1 1   ? 24.880 9.939   -7.442  1.00 41.49  ? 0   MET A CG  1
587        // ATOM   7    S  SD  . MET A 1 1   ? 24.262 10.555  -5.873  1.00 44.70  ? 0   MET A SD  1
588        // ATOM   8    C  CE  . MET A 1 1   ? 24.822 12.266  -5.967  1.00 41.59  ? 0   MET A CE  1
589        //
590        // pub enum AAAtom {
591        //     N = 0,    CA = 1,   C = 2,    CB = 3,   O = 4,
592        //     CG = 5,   CG1 = 6,  CG2 = 7,  OG = 8,   OG1 = 9,
593        //     SG = 10,  CD = 11,  CD1 = 12, CD2 = 13, ND1 = 14,
594        //     ND2 = 15, OD1 = 16, OD2 = 17, SD = 18,  CE = 19,
595        //     CE1 = 20, CE2 = 21, CE3 = 22, NE = 23,  NE1 = 24,
596        //     NE2 = 25, OE1 = 26, OE2 = 27, CH2 = 28, NH1 = 29,
597        //     NH2 = 30, OH = 31,  CZ = 32,  CZ2 = 33, CZ3 = 34,
598        //     NZ = 35,  OXT = 36,
599        //     Unknown = -1,
600        // }
601        let allatom_coords = [
602            // Methionine - AA00
603            // We iterate through these positions. Not all AA's have each
604            ("N", (0, 0, 0, ..), vec![24.277, 8.374, -9.854]),
605            ("CA", (0, 0, 1, ..), vec![24.404, 9.859, -9.939]),
606            ("C", (0, 0, 2, ..), vec![25.814, 10.249, -10.359]),
607            ("CB", (0, 0, 3, ..), vec![24.070, 10.495, -8.596]),
608            ("O", (0, 0, 4, ..), vec![26.748, 9.469, -10.197]),
609            ("CG", (0, 0, 5, ..), vec![24.880, 9.939, -7.442]),
610            ("CG1", (0, 0, 6, ..), vec![0.0, 0.0, 0.0]),
611            ("CG2", (0, 0, 7, ..), vec![0.0, 0.0, 0.0]),
612            ("OG", (0, 0, 8, ..), vec![0.0, 0.0, 0.0]),
613            ("OG1", (0, 0, 9, ..), vec![0.0, 0.0, 0.0]),
614            ("SG", (0, 0, 10, ..), vec![0.0, 0.0, 0.0]),
615            ("CD", (0, 0, 11, ..), vec![0.0, 0.0, 0.0]),
616            ("CD1", (0, 0, 12, ..), vec![0.0, 0.0, 0.0]),
617            ("CD2", (0, 0, 13, ..), vec![0.0, 0.0, 0.0]),
618            ("ND1", (0, 0, 14, ..), vec![0.0, 0.0, 0.0]),
619            ("ND2", (0, 0, 15, ..), vec![0.0, 0.0, 0.0]),
620            ("OD1", (0, 0, 16, ..), vec![0.0, 0.0, 0.0]),
621            ("OD2", (0, 0, 17, ..), vec![0.0, 0.0, 0.0]),
622            ("SD", (0, 0, 18, ..), vec![24.262, 10.555, -5.873]),
623            ("CE", (0, 0, 19, ..), vec![24.822, 12.266, -5.967]),
624            ("CE1", (0, 0, 20, ..), vec![0.0, 0.0, 0.0]),
625            ("CE2", (0, 0, 21, ..), vec![0.0, 0.0, 0.0]),
626            ("CE3", (0, 0, 22, ..), vec![0.0, 0.0, 0.0]),
627            ("NE", (0, 0, 23, ..), vec![0.0, 0.0, 0.0]),
628            ("NE1", (0, 0, 24, ..), vec![0.0, 0.0, 0.0]),
629            ("NE2", (0, 0, 25, ..), vec![0.0, 0.0, 0.0]),
630            ("OE1", (0, 0, 26, ..), vec![0.0, 0.0, 0.0]),
631            ("OE2", (0, 0, 27, ..), vec![0.0, 0.0, 0.0]),
632            ("CH2", (0, 0, 28, ..), vec![0.0, 0.0, 0.0]),
633            ("NH1", (0, 0, 29, ..), vec![0.0, 0.0, 0.0]),
634            ("NH2", (0, 0, 30, ..), vec![0.0, 0.0, 0.0]),
635            ("OH", (0, 0, 31, ..), vec![0.0, 0.0, 0.0]),
636            ("CZ", (0, 0, 32, ..), vec![0.0, 0.0, 0.0]),
637            ("CZ2", (0, 0, 33, ..), vec![0.0, 0.0, 0.0]),
638            ("CZ3", (0, 0, 34, ..), vec![0.0, 0.0, 0.0]),
639            ("NZ", (0, 0, 35, ..), vec![0.0, 0.0, 0.0]),
640            ("OXT", (0, 0, 36, ..), vec![0.0, 0.0, 0.0]),
641        ];
642        for (atom_name, (b, i, j, k), expected) in allatom_coords {
643            let actual: Vec<f32> = ac_backbone_tensor.i((b, i, j, k))?.to_vec1()?;
644            assert_eq!(actual, expected, "Mismatch for atom {}", atom_name);
645        }
646        Ok(())
647    }
648
649    #[test]
650    fn test_ligand_tensor() -> Result<()> {
651        let device = Device::Cpu;
652        let (pdb_file, _temp) = TestFile::protein_01().create_temp()?;
653        let ac = load_structure(pdb_file)?;
654        let (ligand_coords, ligand_elements, _) =
655            ac.to_numeric_ligand_atoms(&device).expect("REASON");
656        // 154 residues; 54 other atoms.
657        assert_eq!(ligand_coords.dims(), &[1, 154, 54, 3]);
658        // Check my residue coords in the Tensor
659        //
660        // HETATM 1222 S  S   . SO4 B 2 .   ? 30.746 18.706  28.896  1.00 47.98  ? 157 SO4 A S   1
661        // HETATM 1223 O  O1  . SO4 B 2 .   ? 30.697 20.077  28.620  1.00 48.06  ? 157 SO4 A O1  1
662        // HETATM 1224 O  O2  . SO4 B 2 .   ? 31.104 18.021  27.725  1.00 47.52  ? 157 SO4 A O2  1
663        // HETATM 1225 O  O3  . SO4 B 2 .   ? 29.468 18.179  29.331  1.00 47.79  ? 157 SO4 A O3  1
664        // HETATM 1226 O  O4  . SO4 B 2 .   ? 31.722 18.578  29.881  1.00 47.85  ? 157 SO4 A O4  1
665        let allatom_coords = [
666            ("S", (0, 0, 0, ..), vec![30.746, 18.706, 28.896]),
667            ("O1", (0, 0, 1, ..), vec![30.697, 20.077, 28.620]),
668            ("O2", (0, 0, 2, ..), vec![31.104, 18.021, 27.725]),
669            ("O3", (0, 0, 3, ..), vec![29.468, 18.179, 29.331]),
670            ("O4", (0, 0, 4, ..), vec![31.722, 18.578, 29.881]),
671        ];
672        for (atom_name, (b, l, i, j), expected) in allatom_coords {
673            let actual: Vec<f32> = ligand_coords.i((b, l, i, ..))?.to_vec1()?;
674            assert_eq!(actual, expected, "Mismatch for atom {}", atom_name);
675        }
676
677        // Now check the elements
678        let elements: Vec<&str> = ligand_elements
679            .i((0, 0, ..))?
680            .to_vec1::<i64>()?
681            .into_iter()
682            .map(|elem| Element::new(elem as usize).unwrap().symbol())
683            .collect();
684
685        assert_eq!(elements[0], "S");
686        assert_eq!(elements[1], "O");
687        assert_eq!(elements[2], "O");
688        assert_eq!(elements[3], "O");
689
690        Ok(())
691    }
692
693    #[test]
694    fn test_backbone_tensor() {
695        let device = Device::Cpu;
696        let (pdb_file, _temp) = TestFile::protein_01().create_temp().unwrap();
697        let ac = load_structure(pdb_file).unwrap();
698        let xyz_37 = ac
699            .to_numeric_atom37(&device)
700            .expect("XYZ creation for all-atoms");
701        assert_eq!(xyz_37.dims(), [1, 154, 37, 3]);
702
703        // # xyz_37_m = feature_dict["xyz_37_m"] #[B,L,37] - mask for all coords
704        let xyz_m = create_backbone_mask_37(&xyz_37).expect("masking procedure should work");
705        assert_eq!(xyz_m.dims(), &[1, 154, 37]);
706    }
707    #[test]
708    fn test_compute_nearest_neighbors() {
709        let device = Device::Cpu;
710        let test_dtype = DType::F32;
711
712        // Create a simple 2x3x3 tensor representing 2 sequences of 3 points in 3D space
713        let coords = Tensor::new(
714            &[
715                [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]], // First sequence
716                [[0.0, 1.0, 0.0], [1.0, 1.0, 0.0], [2.0, 1.0, 0.0]], // Second sequence
717            ],
718            &device,
719        )
720        .unwrap()
721        .to_dtype(test_dtype)
722        .unwrap();
723
724        // Create mask indicating all points are valid
725        let mask = Tensor::ones((2, 3), test_dtype, &device).unwrap();
726
727        // Get 2 nearest neighbors for each point
728        let (distances, indices) = compute_nearest_neighbors(&coords, &mask, 2, 1e-6).unwrap();
729
730        // Check shapes
731        assert_eq!(distances.dims(), &[2, 3, 2]); // [batch, seq_len, k]
732        assert_eq!(indices.dims(), &[2, 3, 2]); // [batch, seq_len, k]
733
734        // For first sequence, point [1,0,0] should have [0,0,0] and [2,0,0] as nearest neighbors
735        let point_neighbors: Vec<u32> = indices.i((0, 1, ..)).unwrap().to_vec1().unwrap();
736        assert_eq!(point_neighbors, vec![0, 2]);
737
738        // Check distances are correct
739        let point_distances: Vec<f32> = distances.i((0, 1, ..)).unwrap().to_vec1().unwrap();
740        assert!((point_distances[0] - 1.0).abs() < 1e-5);
741        assert!((point_distances[1] - 1.0).abs() < 1e-5);
742    }
743}
ferritin_plms/featurize/utilities.rs

ferritin_plms/featurize/
utilities.rs
