ferritin_plms/featurize/
utilities.rs

1use candle_core::{D, DType, Device, IndexOp, Result, Tensor};
2use candle_nn::encoding::one_hot;
3use candle_nn::{Linear, Module};
4
5// Core biochemistry constants live in ferritin-core; re-export so existing
6// callers (`use crate::featurize::utilities::AAAtom`, etc.) still compile.
7pub use ferritin_core::info::amino_acids::{ALPHABET, aa1to_int, aa3to1, int_to_aa1};
8pub use ferritin_core::info::atom37::AAAtom;
9
10macro_rules! define_residues {
11    ($($name:ident: $code3:expr_2021, $code1:expr_2021, $idx:expr_2021, $features:expr_2021, $atoms14:expr_2021),* $(,)?) => {
12        #[derive(Debug, Copy, Clone)]
13        pub enum Residue {
14            $($name),*
15        }
16
17        impl Residue {
18            pub const fn code3(&self) -> &'static str {
19                match self {
20                    $(Self::$name => $code3),*
21                }
22            }
23            pub const fn code1(&self) -> char {
24                match self {
25                    $(Self::$name => $code1),*
26                }
27            }
28            pub const fn atoms14(&self) -> [AAAtom; 14] {
29                match self {
30                    $(Self::$name => $atoms14),*
31                }
32            }
33            pub fn from_int(value: i32) -> Self {
34                match value {
35                    $($idx => Self::$name,)*
36                    _ => Self::UNK
37                }
38            }
39            pub fn to_int(&self) -> i32 {
40                match self {
41                    $(Self::$name => $idx),*
42                }
43            }
44        }
45    }
46}
47
48define_residues! {
49    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],
50    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],
51    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],
52    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],
53    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],
54    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],
55    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],
56    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],
57    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],
58    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],
59    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],
60    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],
61    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],
62    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],
63    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],
64    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],
65    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],
66    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],
67    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],
68    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],
69    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],
70}
71
72/// Per-residue ligand atom context: the `number_of_ligand_atoms` ligand atoms
73/// nearest each residue's Cβ.
74///
75/// Port of `data_utils.get_nearest_neighbours`. Given `L` residues and `M`
76/// ligand atoms it returns `(Y, Y_t, Y_m, D_closest)` shaped
77/// `(L, N, 3)`, `(L, N)`, `(L, N)` and `(L,)`, where `N` is
78/// `number_of_ligand_atoms`.
79///
80/// Masked-out pairs are pushed to a distance of 1000 so they sort last, and
81/// when a structure has fewer than `N` ligand atoms the remaining slots are
82/// zero-filled — including `Y_m`, so the padding contributes nothing.
83///
84/// # Arguments
85/// * `cb` — virtual Cβ per residue, `(L, 3)` or `(1, L, 3)`.
86/// * `mask` — per-residue validity, `(L,)` or `(1, L)`.
87/// * `y` — ligand atom coordinates, `(M, 3)`.
88/// * `y_t` — ligand atomic numbers, `(M,)`.
89/// * `y_m` — ligand atom validity, `(M,)`.
90pub fn get_nearest_neighbours(
91    cb: &Tensor,
92    mask: &Tensor,
93    y: &Tensor,
94    y_t: &Tensor,
95    y_m: &Tensor,
96    number_of_ligand_atoms: i64,
97) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
98    let cb = if cb.dims().len() == 3 {
99        cb.squeeze(0)?
100    } else {
101        cb.clone()
102    };
103    let mask = if mask.dims().len() == 2 {
104        mask.squeeze(0)?
105    } else {
106        mask.clone()
107    };
108    let device = cb.device();
109    let num_residues = cb.dim(0)?;
110    let num_ligand_atoms = y.dim(0)?;
111    let context = number_of_ligand_atoms as usize;
112
113    let mask = mask.to_dtype(DType::F32)?;
114    let y_m_f32 = y_m.to_dtype(DType::F32)?;
115
116    // mask_CBY = mask[:, None] * Y_m[None, :]
117    let mask_cby = mask.unsqueeze(1)?.broadcast_mul(&y_m_f32.unsqueeze(0)?)?;
118
119    // L2_AB = sum((CB[:, None, :] - Y[None, :, :]) ** 2, -1)
120    let l2_ab = cb
121        .unsqueeze(1)?
122        .broadcast_sub(&y.unsqueeze(0)?)?
123        .powf(2.0)?
124        .sum(D::Minus1)?;
125
126    // L2_AB = L2_AB * mask_CBY + (1 - mask_CBY) * 1000.0
127    let l2_ab = ((l2_ab * &mask_cby)? + ((1.0 - &mask_cby)? * 1000.0)?)?;
128
129    // ASCENDING: `torch.argsort` defaults to ascending, and these are
130    // distances. Sorting descending here selected each residue's FARTHEST
131    // ligand atoms as its context (ferritin-100.11).
132    let take = context.min(num_ligand_atoms);
133    let nn_idx = l2_ab
134        .contiguous()?
135        .arg_sort_last_dim(true)?
136        .narrow(1, 0, take)?
137        .contiguous()?;
138
139    let l2_ab_nn = l2_ab.contiguous()?.gather(&nn_idx, 1)?;
140    let d_ab_closest = l2_ab_nn.i((.., 0))?.sqrt()?;
141
142    let y_sel = y
143        .unsqueeze(0)?
144        .broadcast_as((num_residues, num_ligand_atoms, 3))?
145        .contiguous()?
146        .gather(
147            &nn_idx
148                .unsqueeze(2)?
149                .broadcast_as((num_residues, take, 3))?
150                .contiguous()?,
151            1,
152        )?;
153    let y_t_sel = y_t
154        .unsqueeze(0)?
155        .broadcast_as((num_residues, num_ligand_atoms))?
156        .contiguous()?
157        .gather(&nn_idx, 1)?;
158    let y_m_sel = y_m
159        .unsqueeze(0)?
160        .broadcast_as((num_residues, num_ligand_atoms))?
161        .contiguous()?
162        .gather(&nn_idx, 1)?;
163
164    // Fewer ligand atoms than the model's context width: zero-fill the rest,
165    // Y_m included, so the padding slots contribute nothing.
166    if take == context {
167        return Ok((y_sel, y_t_sel, y_m_sel, d_ab_closest));
168    }
169    let pad = context - take;
170    let pad_y = Tensor::zeros((num_residues, pad, 3), y_sel.dtype(), device)?;
171    let pad_t = Tensor::zeros((num_residues, pad), y_t_sel.dtype(), device)?;
172    let pad_m = Tensor::zeros((num_residues, pad), y_m_sel.dtype(), device)?;
173    Ok((
174        Tensor::cat(&[&y_sel, &pad_y], 1)?,
175        Tensor::cat(&[&y_t_sel, &pad_t], 1)?,
176        Tensor::cat(&[&y_m_sel, &pad_m], 1)?,
177        d_ab_closest,
178    ))
179}
180
181/// Apply a [`Linear`] across the last dimension of a tensor of any rank.
182///
183/// candle's `Linear::forward` matmuls directly, which caps the input at rank 3
184/// — beyond that it fails with `shape mismatch in matmul`. LigandMPNN's ligand
185/// graph is rank 5 (`(B, L, M, M, C)`: one M x M edge matrix per residue), so
186/// `y_edges` and every projection inside `DecLayerJ` needs this
187/// (ferritin-100.11).
188///
189/// Flattens the leading axes, applies the layer, and restores the shape with
190/// the new feature width.
191pub fn linear_last_dim(layer: &Linear, xs: &Tensor) -> Result<Tensor> {
192    let dims = xs.dims();
193    if dims.len() <= 3 {
194        return layer.forward(xs);
195    }
196    let features = dims[dims.len() - 1];
197    let rows: usize = dims[..dims.len() - 1].iter().product();
198    let out = layer.forward(&xs.reshape((rows, features))?.contiguous()?)?;
199    let mut shape = dims[..dims.len() - 1].to_vec();
200    shape.push(out.dim(D::Minus1)?);
201    out.reshape(shape.as_slice())
202}
203
204pub fn cat_neighbors_nodes(
205    h_nodes: &Tensor,
206    h_neighbors: &Tensor,
207    e_idx: &Tensor,
208) -> Result<Tensor> {
209    let h_nodes_gathered = gather_nodes(h_nodes, e_idx)?;
210    let h_neighbors = h_neighbors.expand((
211        h_neighbors.dim(0)?, // 1
212        h_nodes.dim(1)?,     // 93
213        h_neighbors.dim(2)?, // 24
214        h_neighbors.dim(3)?, // 128
215    ))?;
216
217    Tensor::cat(
218        &[h_neighbors, h_nodes_gathered.to_dtype(DType::F32)?],
219        D::Minus1,
220    )
221}
222
223/// Retrieve the nearest Neighbor of a set of coordinates.
224/// Usually used for CA carbon distance.
225pub fn compute_nearest_neighbors(
226    coords: &Tensor,
227    mask: &Tensor,
228    k: usize,
229    eps: f32,
230) -> Result<(Tensor, Tensor)> {
231    let (_batch_size, seq_len, _) = coords.dims3()?;
232    // broadcast_matmul handles broadcasting automatically
233    // [2, 3, 1] × [2, 1, 3] -> [2, 3, 3]
234
235    let mask_2d = mask
236        .unsqueeze(2)?
237        .broadcast_matmul(&mask.unsqueeze(1)?)?
238        .to_dtype(DType::F32)?;
239    // Compute pairwise distances with broadcasting
240
241    let distances = (coords
242        .unsqueeze(2)?
243        .broadcast_sub(&coords.unsqueeze(1)?)?
244        .powf(2.)?
245        .sum(D::Minus1)?
246        + eps as f64)? // also  doesn't have add
247        .sqrt()?
248        .to_dtype(DType::F32)?;
249
250    // Apply mask
251    // Get max values for adjustment
252    let masked_distances = (&distances * &mask_2d.to_dtype(DType::F32)?)?;
253    // println!("after masked_distances");
254    let d_max = masked_distances.max_keepdim(D::Minus1)?;
255    let mask_term = ((&mask_2d.to_dtype(DType::F32)? * -1.0)? + 1.0)?;
256    let d_adjust = (&masked_distances + mask_term.broadcast_mul(&d_max)?)?;
257    let d_adjust = d_adjust.to_dtype(DType::F32)?;
258    topk_last_dim(&d_adjust, k.min(seq_len))
259}
260
261// https://github.com/huggingface/candle/pull/2375/files#diff-e4d52a71060a80ac8c549f2daffcee77f9bf4de8252ad067c47b1c383c3ac828R957
262//
263// NOTE the ascending sort. This is the k SMALLEST values, matching
264// `torch.topk(.., largest=False)` — which is what every caller here wants,
265// since the values being ranked are distances. It used to pass `false`
266// (descending) to `arg_sort_last_dim`, so `compute_nearest_neighbors` returned
267// each residue's k FARTHEST neighbours: on 1BC8 only 2% of the resulting
268// `E_idx` entries matched the reference, and residue 0's neighbour list did not
269// even contain residue 0, whose self-distance is zero (ferritin-100.11).
270pub fn topk_last_dim(xs: &Tensor, topk: usize) -> Result<(Tensor, Tensor)> {
271    let sorted_indices = xs.arg_sort_last_dim(true)?.to_dtype(DType::U32)?;
272    let topk_indices = sorted_indices.narrow(D::Minus1, 0, topk)?.contiguous()?;
273    let gathered = xs.gather(&topk_indices, D::Minus1)?;
274    Ok((gathered, topk_indices))
275}
276
277/// Input coords. Output 1 <batch  x 1 > Tensor
278/// representing whether each residue has all 4 backbone atoms.
279/// note that the internal ordering is different between
280/// backbone only [N/CA/C/O] and all-atom [N/CA/C/CB/O]....
281pub fn create_backbone_mask_37(xyz_37: &Tensor) -> Result<Tensor> {
282    let (b, l, rescount, _) = xyz_37.dims4()?;
283    // Create a vector with 1s at positions 0,1,2,4 and 0s elsewhere
284    let mut values = vec![0f32; rescount];
285    for &idx in &[0, 1, 2, 4] {
286        if idx < rescount {
287            values[idx] = 1.0;
288        }
289    }
290
291    // Create the base mask for one sequence, explicitly specifying the data type
292    let base_mask = Tensor::new(values.as_slice(), xyz_37.device())?.to_dtype(DType::F32)?;
293
294    // Create the full mask by repeating it for each batch and length
295    let mask = base_mask.unsqueeze(0)?.unsqueeze(0)?; // Add batch and length dimensions
296    let mask = mask.broadcast_as((b, l, rescount))?;
297
298    Ok(mask)
299}
300/// Get Pseudo CB
301pub fn calculate_cb(xyz_37: &Tensor) -> Result<Tensor> {
302    // make sure we are dealing with
303    let (_, dim37, dim3) = xyz_37.dims3()?;
304    assert_eq!(dim37, 37);
305    assert_eq!(dim3, 3);
306
307    // Constants for CB calculation
308    let a_coeff = -0.58273431f64;
309    let b_coeff = 0.56802827f64;
310    let c_coeff = -0.54067466f64;
311
312    // Get N, CA, C coordinates
313    let n = xyz_37.i((.., 0, ..))?; // N  at index 0
314    let ca = xyz_37.i((.., 1, ..))?; // CA at index 1
315    let c = xyz_37.i((.., 2, ..))?; // C  at index 2
316
317    // Calculate vectors
318    let b = (&ca - &n)?; // CA - N
319    let c = (&c - &ca)?; // C - CA
320
321    // Manual cross product components
322    // a_x = b_y * c_z - b_z * c_y
323    // a_y = b_z * c_x - b_x * c_z
324    // a_z = b_x * c_y - b_y * c_x
325    let b_x = b.i((.., 0))?;
326    let b_y = b.i((.., 1))?;
327    let b_z = b.i((.., 2))?;
328    let c_x = c.i((.., 0))?;
329    let c_y = c.i((.., 1))?;
330    let c_z = c.i((.., 2))?;
331
332    let a_x = ((&b_y * &c_z)? - (&b_z * &c_y)?)?;
333    let a_y = ((&b_z * &c_x)? - (&b_x * &c_z)?)?;
334    let a_z = ((&b_x * &c_y)? - (&b_y * &c_x)?)?;
335
336    // Stack the cross product components back together
337    let a = Tensor::stack(&[&a_x, &a_y, &a_z], 1)?;
338
339    // Final CB calculation: -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + CA
340
341    (&a * a_coeff)? + (&b * b_coeff)? + (&c * c_coeff)? + &ca
342}
343
344/// Custom Cross-Product Fn.
345pub fn cross_product(a: &Tensor, b: &Tensor) -> Result<Tensor> {
346    let last_dim = a.dims().len() - 1;
347
348    // Extract components
349    let a0 = a.narrow(last_dim, 0, 1)?;
350    let a1 = a.narrow(last_dim, 1, 1)?;
351    let a2 = a.narrow(last_dim, 2, 1)?;
352
353    let b0 = b.narrow(last_dim, 0, 1)?;
354    let b1 = b.narrow(last_dim, 1, 1)?;
355    let b2 = b.narrow(last_dim, 2, 1)?;
356
357    // Compute cross produAAct components
358    let c0 = ((&a1 * &b2)? - (&a2 * &b1)?)?;
359    let c1 = ((&a2 * &b0)? - (&a0 * &b2)?)?;
360    let c2 = ((&a0 * &b1)? - (&a1 * &b0)?)?;
361
362    // Stack the results
363    Tensor::cat(&[&c0, &c1, &c2], last_dim)
364}
365
366/// Gather_edges
367/// Features `[B,N,N,C]` at Neighbor indices `[B,N,K]` => Neighbor features `[B,N,K,C]`
368pub fn gather_edges(edges: &Tensor, neighbor_idx: &Tensor) -> Result<Tensor> {
369    let (d1, d2, d3) = neighbor_idx.dims3()?;
370    let neighbors =
371        neighbor_idx
372            .unsqueeze(D::Minus1)?
373            .expand((d1, d2, d3, edges.dim(D::Minus1)?))?;
374    let edge_gather = edges.gather(&neighbors, 2)?;
375    Ok(edge_gather)
376}
377
378/// Gather Nodes
379///
380/// Features `[B,N,C]` at Neighbor indices `[B,N,K]` => `[B,N,K,C]`
381/// Flatten and expand indices per batch `[B,N,K]` => `[B,NK]` => `[B,NK,C]`
382pub fn gather_nodes(nodes: &Tensor, neighbor_idx: &Tensor) -> Result<Tensor> {
383    let (batch_size, n_nodes, n_features) = nodes.dims3()?;
384    let (_, _, k_neighbors) = neighbor_idx.dims3()?;
385    // Reshape neighbor_idx to [B, N*K]
386    let neighbors_flat = neighbor_idx.reshape((batch_size, n_nodes * k_neighbors))?;
387    // Add feature dimension and expand
388    let neighbors_flat = neighbors_flat
389        .unsqueeze(2)? // Add feature dimension [B, N*K, 1]
390        .expand((batch_size, n_nodes * k_neighbors, n_features))?; // Expand to [B, N*K, C]
391    // make contiguous for the gather.
392    let neighbors_flat = neighbors_flat.contiguous()?;
393    // Gather features
394    let neighbor_features = nodes.gather(&neighbors_flat, 1)?;
395    // Reshape back to [B, N, K, C]
396    neighbor_features.reshape((batch_size, n_nodes, k_neighbors, n_features))
397}
398
399pub fn gather_nodes_t(nodes: &Tensor, neighbor_idx: &Tensor) -> Result<Tensor> {
400    // Features [B,N,C] at Neighbor index [B,K] => Neighbor features[B,K,C]
401    let (d1, d2, d3) = nodes.dims3()?;
402    let idx_flat = neighbor_idx.unsqueeze(D::Minus1)?.expand((d1, d2, d3))?;
403    nodes.gather(&idx_flat, 1)
404}
405
406#[allow(dead_code)]
407fn get_seq_rec(s: &Tensor, s_pred: &Tensor, mask: &Tensor) -> Result<Tensor> {
408    // S: true sequence shape=[batch, length]
409    // S_pred: predicted sequence shape=[batch, length]
410    // mask: mask to compute average over the region shape=[batch, length]
411    // Returns: averaged sequence recovery shape=[batch]
412    //
413    // Compute the match tensor
414    let match_tensor = s.eq(s_pred)?;
415    let match_f32 = match_tensor.to_dtype(DType::F32)?;
416    let numerator = (match_f32 * mask)?.sum_keepdim(1)?;
417    let denominator = mask.sum_keepdim(1)?;
418    let average = numerator.broadcast_div(&denominator)?;
419    // Remove the last dimension to get shape=[batch]
420    average.squeeze(1)
421}
422
423#[allow(dead_code)]
424fn get_score(s: &Tensor, log_probs: &Tensor, mask: &Tensor) -> Result<(Tensor, Tensor)> {
425    //     S : true sequence shape=[batch, length]
426    //     log_probs : predicted sequence shape=[batch, length]
427    //     mask : mask to compute average over the region shape=[batch, length]
428    //     average_loss : averaged categorical cross entropy (CCE) [batch]
429    //     loss_per_resdue : per position CCE [batch, length]
430
431    //     """
432    //     S_one_hot = torch.nn.functional.one_hot(S, 21)
433    //     loss_per_residue = -(S_one_hot * log_probs).sum(-1)  # [B, L]
434    //     average_loss = torch.sum(loss_per_residue * mask, dim=-1) / (
435    //         torch.sum(mask, dim=-1) + 1e-8
436    //     )
437    //     return average_loss, loss_per_residue
438
439    // S: true sequence shape=[batch, length]
440    // log_probs: predicted sequence shape=[batch, length, 21]
441    // mask: mask to compute average over the region shape=[batch, length]
442    // Returns:
443    //   - average_loss: averaged categorical cross entropy (CCE) [batch]
444    //   - loss_per_residue: per position CCE [batch, length]
445
446    // Create one-hot encoding of S.
447    // see https://docs.rs/candle-nn/0.7.2/candle_nn/encoding/fn.one_hot.html
448    // this could be wrong...
449    let s_one_hot = one_hot(s.clone(), 21, 1., 0.)?;
450    let loss_per_residue = s_one_hot.mul(&log_probs.neg()?)?.sum(D::Minus1)?;
451    let average_loss = loss_per_residue
452        .mul(mask)?
453        .sum_keepdim(D::Minus1)?
454        .div(&(mask.sum_keepdim(D::Minus1)? + 1e-8f64)?)?
455        .squeeze(D::Minus1)?;
456
457    Ok((average_loss, loss_per_residue))
458}
459
460pub fn linspace(
461    start: f64,
462    stop: f64,
463    steps: usize,
464    device: &Device,
465    return_type: DType,
466) -> Result<Tensor> {
467    if steps == 0 {
468        Tensor::from_vec(Vec::<f64>::new(), steps, device)
469    } else if steps == 1 {
470        Tensor::from_vec(vec![start], steps, device)
471    } else {
472        let delta = (stop - start) / (steps - 1) as f64;
473        let vs = (0..steps)
474            .map(|step| start + step as f64 * delta)
475            .collect::<Vec<_>>();
476        Tensor::from_vec(vs, steps, device)?.to_dtype(return_type)
477    }
478}
479
480pub fn linspace_f32(start: f32, stop: f32, steps: usize, device: &Device) -> Result<Tensor> {
481    match steps {
482        0 => Tensor::from_vec(Vec::<f32>::new(), steps, device),
483        1 => Tensor::from_vec(vec![start], steps, device),
484        _ => {
485            let delta = (stop - start) / (steps - 1) as f32;
486            let vs = (0..steps)
487                .map(|step| start + step as f32 * delta)
488                .collect::<Vec<_>>();
489            Tensor::from_vec(vs, steps, device)
490        }
491    }
492}
493
494#[cfg(test)]
495mod tests {
496    use super::*;
497    use crate::StructureFeatures;
498    use anyhow::Result;
499    use ferritin_core::info::elements::Element;
500    use ferritin_core::load_structure;
501    use ferritin_test_data::TestFile;
502
503    #[test]
504    fn test_residue_codes() {
505        let ala = Residue::ALA;
506        assert_eq!(ala.code3(), "ALA");
507        assert_eq!(ala.code1(), 'A');
508        assert_eq!(ala.to_int(), 0);
509    }
510
511    #[test]
512    fn test_residue_from_int() {
513        assert!(matches!(Residue::from_int(0), Residue::ALA));
514        assert!(matches!(Residue::from_int(1), Residue::CYS));
515        assert!(matches!(Residue::from_int(999), Residue::UNK));
516    }
517
518    #[test]
519    fn test_residue_atoms() {
520        let trp = Residue::TRP;
521        let atoms = trp.atoms14();
522        assert_eq!(atoms[0], AAAtom::N);
523        assert_eq!(atoms[13], AAAtom::CH2);
524
525        let gly = Residue::GLY;
526        let atoms = gly.atoms14();
527        assert_eq!(atoms[4], AAAtom::Unknown);
528    }
529
530    #[test]
531    fn test_atom_backbone_tensor() -> Result<()> {
532        let device = Device::Cpu;
533        let (pdb_file, _temp) = TestFile::protein_01().create_temp()?;
534        let ac = load_structure(pdb_file)?;
535        let ac_backbone_tensor: Tensor = ac.to_numeric_backbone_atoms(&device).expect("REASON");
536        // batch size of 1;154 residues; N/CA/C/O; positions
537        assert_eq!(ac_backbone_tensor.dims(), &[1, 154, 4, 3]);
538
539        // Check my residue coords in the Tensor
540        // ATOM   1    N  N   . MET A 1 1   ? 24.277 8.374   -9.854  1.00 38.41  ? 0   MET A N   1
541        // ATOM   2    C  CA  . MET A 1 1   ? 24.404 9.859   -9.939  1.00 37.90  ? 0   MET A CA  1
542        // ATOM   3    C  C   . MET A 1 1   ? 25.814 10.249  -10.359 1.00 36.65  ? 0   MET A C   1
543        // ATOM   4    O  O   . MET A 1 1   ? 26.748 9.469   -10.197 1.00 37.13  ? 0   MET A O   1
544        let backbone_coords = [
545            // Methionine - AA00
546            ("N", (0, 0, 0, ..), vec![24.277, 8.374, -9.854]),
547            ("CA", (0, 0, 1, ..), vec![24.404, 9.859, -9.939]),
548            ("C", (0, 0, 2, ..), vec![25.814, 10.249, -10.359]),
549            ("O", (0, 0, 3, ..), vec![26.748, 9.469, -10.197]),
550            // Valine - AA01
551            ("N", (0, 1, 0, ..), vec![25.964, 11.453, -10.903]),
552            ("CA", (0, 1, 1, ..), vec![27.263, 11.924, -11.359]),
553            ("C", (0, 1, 2, ..), vec![27.392, 13.428, -11.115]),
554            ("O", (0, 1, 3, ..), vec![26.443, 14.184, -11.327]),
555            // Glycing - AAlast
556            ("N", (0, 153, 0, ..), vec![23.474, -3.227, 5.994]),
557            ("CA", (0, 153, 1, ..), vec![22.818, -2.798, 7.211]),
558            ("C", (0, 153, 2, ..), vec![22.695, -1.282, 7.219]),
559            ("O", (0, 153, 3, ..), vec![21.870, -0.745, 7.992]),
560        ];
561
562        for (atom_name, (b, i, j, k), expected) in backbone_coords {
563            // assert_eq!(ac_backbone_tensor.dims(), &[1, 154, 4, 3])
564            let actual: Vec<f32> = ac_backbone_tensor.i((b, i, j, k))?.to_vec1()?;
565            assert_eq!(actual, expected, "Mismatch for atom {}", atom_name);
566        }
567        Ok(())
568    }
569
570    #[test]
571    fn test_all_atom37_tensor() -> Result<()> {
572        let device = Device::Cpu;
573        let (pdb_file, _temp) = TestFile::protein_01().create_temp()?;
574        let ac = load_structure(pdb_file)?;
575        let ac_backbone_tensor: Tensor = ac.to_numeric_atom37(&device).expect("REASON");
576        assert_eq!(ac_backbone_tensor.dims(), &[1, 154, 37, 3]);
577
578        // Check my residue coords in the Tensor
579        // ATOM   1    N  N   . MET A 1 1   ? 24.277 8.374   -9.854  1.00 38.41  ? 0   MET A N   1
580        // ATOM   2    C  CA  . MET A 1 1   ? 24.404 9.859   -9.939  1.00 37.90  ? 0   MET A CA  1
581        // ATOM   3    C  C   . MET A 1 1   ? 25.814 10.249  -10.359 1.00 36.65  ? 0   MET A C   1
582        // ATOM   4    O  O   . MET A 1 1   ? 26.748 9.469   -10.197 1.00 37.13  ? 0   MET A O   1
583        // ATOM   5    C  CB  . MET A 1 1   ? 24.070 10.495  -8.596  1.00 39.58  ? 0   MET A CB  1
584        // ATOM   6    C  CG  . MET A 1 1   ? 24.880 9.939   -7.442  1.00 41.49  ? 0   MET A CG  1
585        // ATOM   7    S  SD  . MET A 1 1   ? 24.262 10.555  -5.873  1.00 44.70  ? 0   MET A SD  1
586        // ATOM   8    C  CE  . MET A 1 1   ? 24.822 12.266  -5.967  1.00 41.59  ? 0   MET A CE  1
587        //
588        // pub enum AAAtom {
589        //     N = 0,    CA = 1,   C = 2,    CB = 3,   O = 4,
590        //     CG = 5,   CG1 = 6,  CG2 = 7,  OG = 8,   OG1 = 9,
591        //     SG = 10,  CD = 11,  CD1 = 12, CD2 = 13, ND1 = 14,
592        //     ND2 = 15, OD1 = 16, OD2 = 17, SD = 18,  CE = 19,
593        //     CE1 = 20, CE2 = 21, CE3 = 22, NE = 23,  NE1 = 24,
594        //     NE2 = 25, OE1 = 26, OE2 = 27, CH2 = 28, NH1 = 29,
595        //     NH2 = 30, OH = 31,  CZ = 32,  CZ2 = 33, CZ3 = 34,
596        //     NZ = 35,  OXT = 36,
597        //     Unknown = -1,
598        // }
599        let allatom_coords = [
600            // Methionine - AA00
601            // We iterate through these positions. Not all AA's have each
602            ("N", (0, 0, 0, ..), vec![24.277, 8.374, -9.854]),
603            ("CA", (0, 0, 1, ..), vec![24.404, 9.859, -9.939]),
604            ("C", (0, 0, 2, ..), vec![25.814, 10.249, -10.359]),
605            ("CB", (0, 0, 3, ..), vec![24.070, 10.495, -8.596]),
606            ("O", (0, 0, 4, ..), vec![26.748, 9.469, -10.197]),
607            ("CG", (0, 0, 5, ..), vec![24.880, 9.939, -7.442]),
608            ("CG1", (0, 0, 6, ..), vec![0.0, 0.0, 0.0]),
609            ("CG2", (0, 0, 7, ..), vec![0.0, 0.0, 0.0]),
610            ("OG", (0, 0, 8, ..), vec![0.0, 0.0, 0.0]),
611            ("OG1", (0, 0, 9, ..), vec![0.0, 0.0, 0.0]),
612            ("SG", (0, 0, 10, ..), vec![0.0, 0.0, 0.0]),
613            ("CD", (0, 0, 11, ..), vec![0.0, 0.0, 0.0]),
614            ("CD1", (0, 0, 12, ..), vec![0.0, 0.0, 0.0]),
615            ("CD2", (0, 0, 13, ..), vec![0.0, 0.0, 0.0]),
616            ("ND1", (0, 0, 14, ..), vec![0.0, 0.0, 0.0]),
617            ("ND2", (0, 0, 15, ..), vec![0.0, 0.0, 0.0]),
618            ("OD1", (0, 0, 16, ..), vec![0.0, 0.0, 0.0]),
619            ("OD2", (0, 0, 17, ..), vec![0.0, 0.0, 0.0]),
620            ("SD", (0, 0, 18, ..), vec![24.262, 10.555, -5.873]),
621            ("CE", (0, 0, 19, ..), vec![24.822, 12.266, -5.967]),
622            ("CE1", (0, 0, 20, ..), vec![0.0, 0.0, 0.0]),
623            ("CE2", (0, 0, 21, ..), vec![0.0, 0.0, 0.0]),
624            ("CE3", (0, 0, 22, ..), vec![0.0, 0.0, 0.0]),
625            ("NE", (0, 0, 23, ..), vec![0.0, 0.0, 0.0]),
626            ("NE1", (0, 0, 24, ..), vec![0.0, 0.0, 0.0]),
627            ("NE2", (0, 0, 25, ..), vec![0.0, 0.0, 0.0]),
628            ("OE1", (0, 0, 26, ..), vec![0.0, 0.0, 0.0]),
629            ("OE2", (0, 0, 27, ..), vec![0.0, 0.0, 0.0]),
630            ("CH2", (0, 0, 28, ..), vec![0.0, 0.0, 0.0]),
631            ("NH1", (0, 0, 29, ..), vec![0.0, 0.0, 0.0]),
632            ("NH2", (0, 0, 30, ..), vec![0.0, 0.0, 0.0]),
633            ("OH", (0, 0, 31, ..), vec![0.0, 0.0, 0.0]),
634            ("CZ", (0, 0, 32, ..), vec![0.0, 0.0, 0.0]),
635            ("CZ2", (0, 0, 33, ..), vec![0.0, 0.0, 0.0]),
636            ("CZ3", (0, 0, 34, ..), vec![0.0, 0.0, 0.0]),
637            ("NZ", (0, 0, 35, ..), vec![0.0, 0.0, 0.0]),
638            ("OXT", (0, 0, 36, ..), vec![0.0, 0.0, 0.0]),
639        ];
640        for (atom_name, (b, i, j, k), expected) in allatom_coords {
641            let actual: Vec<f32> = ac_backbone_tensor.i((b, i, j, k))?.to_vec1()?;
642            assert_eq!(actual, expected, "Mismatch for atom {}", atom_name);
643        }
644        Ok(())
645    }
646
647    #[test]
648    fn test_ligand_tensor() -> Result<()> {
649        let device = Device::Cpu;
650        let (pdb_file, _temp) = TestFile::protein_01().create_temp()?;
651        let ac = load_structure(pdb_file)?;
652        let (ligand_coords, ligand_elements, ligand_mask) =
653            ac.to_numeric_ligand_atoms(&device).expect("REASON");
654
655        // 54 ligand heavy atoms, as RAW per-atom tensors. This used to assert
656        // [1, 154, 54, 3] — one copy per residue — because the per-residue
657        // context selection happened here. It now happens in the LigandMPNN
658        // featurizer, which is where `atom_context_num` is known
659        // (ferritin-100.11).
660        assert_eq!(ligand_coords.dims(), &[1, 54, 3]);
661        assert_eq!(ligand_elements.dims(), &[1, 54]);
662        assert_eq!(ligand_mask.dims(), &[1, 54]);
663
664        // Every real atom is valid — one flag per ATOM, so all ones rather
665        // than the 3.0 a (M, 3) mask used to collapse to.
666        let mask: Vec<f32> = ligand_mask.i((0, ..))?.to_vec1()?;
667        assert!(mask.iter().all(|&m| m == 1.0), "every ligand atom is valid");
668
669        // Check the first sulfate's coords against the mmCIF:
670        //
671        // HETATM 1222 S  S   . SO4 B 2 .   ? 30.746 18.706  28.896 ...
672        // HETATM 1223 O  O1  . SO4 B 2 .   ? 30.697 20.077  28.620 ...
673        // HETATM 1224 O  O2  . SO4 B 2 .   ? 31.104 18.021  27.725 ...
674        // HETATM 1225 O  O3  . SO4 B 2 .   ? 29.468 18.179  29.331 ...
675        // HETATM 1226 O  O4  . SO4 B 2 .   ? 31.722 18.578  29.881 ...
676        let expected = [
677            ("S", vec![30.746, 18.706, 28.896]),
678            ("O1", vec![30.697, 20.077, 28.620]),
679            ("O2", vec![31.104, 18.021, 27.725]),
680            ("O3", vec![29.468, 18.179, 29.331]),
681            ("O4", vec![31.722, 18.578, 29.881]),
682        ];
683        for (i, (atom_name, coords)) in expected.iter().enumerate() {
684            let actual: Vec<f32> = ligand_coords.i((0, i, ..))?.to_vec1()?;
685            assert_eq!(&actual, coords, "Mismatch for atom {atom_name}");
686        }
687
688        let elements: Vec<&str> = ligand_elements
689            .i((0, ..))?
690            .to_vec1::<i64>()?
691            .into_iter()
692            .map(|elem| Element::new(elem as usize).unwrap().symbol())
693            .collect();
694        assert_eq!(&elements[..5], &["S", "O", "O", "O", "O"]);
695        Ok(())
696    }
697
698    #[test]
699    fn test_backbone_tensor() {
700        let device = Device::Cpu;
701        let (pdb_file, _temp) = TestFile::protein_01().create_temp().unwrap();
702        let ac = load_structure(pdb_file).unwrap();
703        let xyz_37 = ac
704            .to_numeric_atom37(&device)
705            .expect("XYZ creation for all-atoms");
706        assert_eq!(xyz_37.dims(), [1, 154, 37, 3]);
707
708        // # xyz_37_m = feature_dict["xyz_37_m"] #[B,L,37] - mask for all coords
709        let xyz_m = create_backbone_mask_37(&xyz_37).expect("masking procedure should work");
710        assert_eq!(xyz_m.dims(), &[1, 154, 37]);
711    }
712    #[test]
713    fn test_compute_nearest_neighbors() {
714        let device = Device::Cpu;
715        let test_dtype = DType::F32;
716
717        // Create a simple 2x3x3 tensor representing 2 sequences of 3 points in 3D space
718        let coords = Tensor::new(
719            &[
720                [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]], // First sequence
721                [[0.0, 1.0, 0.0], [1.0, 1.0, 0.0], [2.0, 1.0, 0.0]], // Second sequence
722            ],
723            &device,
724        )
725        .unwrap()
726        .to_dtype(test_dtype)
727        .unwrap();
728
729        // Create mask indicating all points are valid
730        let mask = Tensor::ones((2, 3), test_dtype, &device).unwrap();
731
732        // Get 2 nearest neighbors for each point
733        let eps = 1e-6f32;
734        let (distances, indices) = compute_nearest_neighbors(&coords, &mask, 2, eps).unwrap();
735
736        // Check shapes
737        assert_eq!(distances.dims(), &[2, 3, 2]); // [batch, seq_len, k]
738        assert_eq!(indices.dims(), &[2, 3, 2]); // [batch, seq_len, k]
739
740        // A point's own index comes FIRST: its self-distance is zero, and
741        // `torch.topk(.., largest=False)` puts the smallest first.
742        //
743        // This assertion used to read `vec![0, 2]` — the two points at distance
744        // 1 — which is what you get from the k LARGEST distances. It passed
745        // because `topk_last_dim` sorted descending, so the test pinned the bug
746        // rather than the behaviour: `compute_nearest_neighbors` was returning
747        // each residue's farthest neighbours (ferritin-100.11).
748        let point_neighbors: Vec<u32> = indices.i((0, 1, ..)).unwrap().to_vec1().unwrap();
749        assert_eq!(
750            point_neighbors[0], 1,
751            "a point's nearest neighbour is itself, at distance 0"
752        );
753        assert!(
754            point_neighbors[1] == 0 || point_neighbors[1] == 2,
755            "the second neighbour is one of the two points at distance 1, got {}",
756            point_neighbors[1]
757        );
758
759        // Check distances are correct
760        let point_distances: Vec<f32> = distances.i((0, 1, ..)).unwrap().to_vec1().unwrap();
761        // sqrt(0 + eps), not 0: eps sits inside the square root.
762        assert!(
763            (point_distances[0] - eps.sqrt()).abs() < 1e-5,
764            "self-distance should be sqrt(eps) = {}, got {}",
765            eps.sqrt(),
766            point_distances[0]
767        );
768        assert!((point_distances[1] - 1.0).abs() < 1e-5);
769    }
770}
ferritin_plms/featurize/utilities.rs

ferritin_plms/featurize/
utilities.rs
