ferritin_plms/esmfold2/
output.rs

1//! ESMFold2 output types.
2//!
3//! Defines [`ESMFold2Output`], the result of a full ESMFold2 forward pass.
4
5use candle_core::Tensor;
6
7/// Output of the ESMFold2 forward pass.
8#[derive(Debug)]
9pub struct ESMFold2Output {
10    /// All-atom coordinates, shape `(Bm, N_atom, 3)` in Ångströms.
11    pub sample_atom_coords: Tensor,
12    /// Per-token pLDDT confidence, shape `(Bm, N_tok)`, range [0, 1].
13    pub plddt: Tensor,
14    /// Predicted TM-score, shape `(Bm,)`.
15    pub ptm: Tensor,
16    /// Interface pTM (for multi-chain), shape `(Bm,)`.
17    pub iptm: Tensor,
18    /// Predicted Aligned Error, shape `(Bm, N_tok, N_tok)` in Å. Optional.
19    pub pae: Option<Tensor>,
20    /// Distogram logits, shape `(Bm, N_tok, N_tok, distogram_bins)`. Optional.
21    pub distogram_logits: Option<Tensor>,
22}
23
24impl ESMFold2Output {
25    /// Convert this output to mmCIF text for a single protein chain.
26    ///
27    /// pLDDT (range [0, 1]) is written as B-factor × 100.
28    pub fn to_mmcif(
29        &self,
30        sequence: &str,
31        chain_id: &str,
32        entry_id: &str,
33    ) -> anyhow::Result<String> {
34        super::mmcif::coords_to_mmcif(
35            &self.sample_atom_coords,
36            sequence,
37            &self.plddt,
38            chain_id,
39            entry_id,
40        )
41    }
42}
ferritin_plms/esmfold2/output.rs

ferritin_plms/esmfold2/
output.rs
