ferritin_plms/esmfold2/
mmcif.rs

1//! mmCIF output generation for ESMFold2 predicted structures.
2//!
3//! Converts all-atom coordinate tensors and confidence scores to the mmCIF
4//! format used by the Protein Data Bank. Output is compatible with
5//! ferritin-core's mmCIF reader.
6//!
7//! ## Coordinate convention
8//!
9//! ESMFold2 outputs all-atom coordinates using the atom14 convention
10//! (up to 14 heavy atoms per standard amino acid residue, in a fixed order).
11//! This writer handles the backbone atoms (N, CA, C, O, CB) which are present
12//! for all standard residues, plus any additional side-chain atoms at
13//! non-zero positions.
14
15use candle_core::Tensor;
16use std::fmt::Write;
17
18/// Standard 14-atom per residue heavy-atom names (atom14 convention).
19/// Index i corresponds to the i-th atom slot.
20pub const ATOM14_NAMES: [&str; 14] = [
21    "N", "CA", "C", "O", "CB", "CG", "CG1", "CG2", "CD", "CD1", "CD2", "NE", "CE", "CZ",
22];
23
24/// Element symbol for each atom14 slot (used for `type_symbol`).
25pub const ATOM14_ELEMENTS: [&str; 14] = [
26    "N", "C", "C", "O", "C", "C", "C", "C", "C", "C", "C", "N", "C", "C",
27];
28
29/// Three-letter residue name lookup from one-letter code.
30pub fn one_to_three(aa: char) -> &'static str {
31    match aa {
32        'A' => "ALA",
33        'R' => "ARG",
34        'N' => "ASN",
35        'D' => "ASP",
36        'C' => "CYS",
37        'Q' => "GLN",
38        'E' => "GLU",
39        'G' => "GLY",
40        'H' => "HIS",
41        'I' => "ILE",
42        'L' => "LEU",
43        'K' => "LYS",
44        'M' => "MET",
45        'F' => "PHE",
46        'P' => "PRO",
47        'S' => "SER",
48        'T' => "THR",
49        'W' => "TRP",
50        'Y' => "TYR",
51        'V' => "VAL",
52        _ => "UNK",
53    }
54}
55
56/// Convert ESMFold2 output to mmCIF string (single model, single chain).
57///
58/// # Arguments
59/// * `coords` — All-atom coordinates tensor. Accepts shapes:
60///   - `(L, 14, 3)` (atom14, one residue per row)
61///   - `(1, L, 14, 3)` (batched atom14; batch dim is squeezed)
62///   - `(L, 3)` (Cα-only)
63///   - `(1, L, 3)` (batched Cα-only; batch dim is squeezed)
64/// * `sequence` — Protein sequence (one-letter codes), length `L`.
65/// * `plddt` — Per-residue pLDDT in [0, 1], shape `(L,)` or `(1, L)`.
66///   Written as B-factor × 100.
67/// * `chain_id` — Chain identifier, e.g. `"A"`.
68/// * `entry_id` — Data-block name, e.g. `"ESMFold2_pred"`.
69///
70/// # Returns
71/// mmCIF text string compatible with ferritin-core's reader.
72pub fn coords_to_mmcif(
73    coords: &Tensor,
74    sequence: &str,
75    plddt: &Tensor,
76    chain_id: &str,
77    entry_id: &str,
78) -> anyhow::Result<String> {
79    // --- 1. Squeeze optional batch dimension ---
80    let coords = if coords.rank() >= 3 && coords.dim(0)? == 1 {
81        coords.squeeze(0)?
82    } else {
83        coords.clone()
84    };
85
86    let plddt = if plddt.rank() >= 2 && plddt.dim(0)? == 1 {
87        plddt.squeeze(0)?
88    } else {
89        plddt.clone()
90    };
91
92    // --- 2. pLDDT → B-factor (×100) ---
93    let plddt_f32 = plddt.to_dtype(candle_core::DType::F32)?;
94    let plddt_vals = plddt_f32.to_vec1::<f32>()?;
95
96    // --- 3. mmCIF header + loop_ block ---
97    let mut out = String::new();
98
99    writeln!(out, "data_{}", entry_id)?;
100    writeln!(out, "#")?;
101    writeln!(out, "loop_")?;
102    writeln!(out, "_atom_site.group_PDB")?;
103    writeln!(out, "_atom_site.id")?;
104    writeln!(out, "_atom_site.type_symbol")?;
105    writeln!(out, "_atom_site.label_atom_id")?;
106    writeln!(out, "_atom_site.label_alt_id")?;
107    writeln!(out, "_atom_site.label_comp_id")?;
108    writeln!(out, "_atom_site.label_asym_id")?;
109    writeln!(out, "_atom_site.label_entity_id")?;
110    writeln!(out, "_atom_site.label_seq_id")?;
111    writeln!(out, "_atom_site.pdbx_PDB_ins_code")?;
112    writeln!(out, "_atom_site.Cartn_x")?;
113    writeln!(out, "_atom_site.Cartn_y")?;
114    writeln!(out, "_atom_site.Cartn_z")?;
115    writeln!(out, "_atom_site.occupancy")?;
116    writeln!(out, "_atom_site.B_iso_or_equiv")?;
117    writeln!(out, "_atom_site.pdbx_formal_charge")?;
118    writeln!(out, "_atom_site.auth_seq_id")?;
119    writeln!(out, "_atom_site.auth_comp_id")?;
120    writeln!(out, "_atom_site.auth_asym_id")?;
121    writeln!(out, "_atom_site.auth_atom_id")?;
122    writeln!(out, "_atom_site.pdbx_PDB_model_num")?;
123
124    // --- 4. Atom rows ---
125    let mut atom_id: u32 = 1;
126    let l = sequence.len();
127
128    let coords_f32 = coords.to_dtype(candle_core::DType::F32)?;
129
130    match coords_f32.rank() {
131        3 => {
132            // Atom14 format: (L, 14, 3)
133            let data = coords_f32.to_vec3::<f32>()?;
134            for (i, aa) in sequence.chars().enumerate() {
135                if i >= l {
136                    break;
137                }
138                let res_name = one_to_three(aa);
139                let seq_id = i + 1;
140                let b_factor = plddt_vals.get(i).copied().unwrap_or(0.0) * 100.0;
141                let natoms = data[i].len().min(ATOM14_NAMES.len());
142
143                for j in 0..natoms {
144                    let xyz = &data[i][j];
145                    let (x, y, z) = (xyz[0], xyz[1], xyz[2]);
146
147                    // Skip unoccupied atom14 slots (all near-zero)
148                    if x.abs() + y.abs() + z.abs() < 0.001 {
149                        continue;
150                    }
151
152                    let atom_name = ATOM14_NAMES[j];
153                    let element = ATOM14_ELEMENTS[j];
154
155                    writeln!(
156                        out,
157                        "ATOM {} {} {} . {} {} 1 {} ? {:.3} {:.3} {:.3} 1.00 {:.2} ? {} {} {} {} 1",
158                        atom_id,
159                        element,
160                        atom_name,
161                        res_name,
162                        chain_id,
163                        seq_id,
164                        x,
165                        y,
166                        z,
167                        b_factor,
168                        seq_id,
169                        res_name,
170                        chain_id,
171                        atom_name
172                    )?;
173                    atom_id += 1;
174                }
175            }
176        }
177        2 => {
178            // Cα-only format: (L, 3)
179            let data = coords_f32.to_vec2::<f32>()?;
180            for (i, aa) in sequence.chars().enumerate() {
181                if i >= l {
182                    break;
183                }
184                let xyz = &data[i];
185                let (x, y, z) = (xyz[0], xyz[1], xyz[2]);
186
187                if x.abs() + y.abs() + z.abs() < 0.001 {
188                    continue;
189                }
190
191                let res_name = one_to_three(aa);
192                let seq_id = i + 1;
193                let b_factor = plddt_vals.get(i).copied().unwrap_or(0.0) * 100.0;
194
195                writeln!(
196                    out,
197                    "ATOM {} C CA . {} {} 1 {} ? {:.3} {:.3} {:.3} 1.00 {:.2} ? {} {} {} CA 1",
198                    atom_id,
199                    res_name,
200                    chain_id,
201                    seq_id,
202                    x,
203                    y,
204                    z,
205                    b_factor,
206                    seq_id,
207                    res_name,
208                    chain_id
209                )?;
210                atom_id += 1;
211            }
212        }
213        r => {
214            return Err(anyhow::anyhow!(
215                "Unexpected coords rank {r}. Expected 2 (Cα-only) or 3 (atom14).",
216            ));
217        }
218    }
219
220    writeln!(out, "#")?;
221    Ok(out)
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227    use candle_core::{Device, Tensor};
228
229    /// Build a small (2, 5, 3) coordinate tensor with recognizable values.
230    fn make_coords(l: usize, natoms: usize) -> Tensor {
231        let mut vals = Vec::with_capacity(l * natoms * 3);
232        for i in 0..l {
233            for j in 0..natoms {
234                vals.push((i * 10 + j) as f32 + 1.0); // x
235                vals.push((i * 10 + j) as f32 + 2.0); // y
236                vals.push((i * 10 + j) as f32 + 3.0); // z
237            }
238        }
239        Tensor::from_vec(vals, &[l, natoms, 3], &Device::Cpu).unwrap()
240    }
241
242    fn make_plddt(l: usize) -> Tensor {
243        let vals: Vec<f32> = (0..l).map(|i| 0.7 + 0.01 * i as f32).collect();
244        Tensor::from_vec(vals, &[l], &Device::Cpu).unwrap()
245    }
246
247    #[test]
248    fn test_coords_to_mmcif_basic() {
249        let sequence = "MA";
250        let coords = make_coords(2, 5);
251        let plddt = make_plddt(2);
252
253        let result = coords_to_mmcif(&coords, sequence, &plddt, "A", "test_pred").unwrap();
254
255        assert!(result.contains("data_test_pred"), "missing data_ header");
256        assert!(result.contains("loop_"), "missing loop_");
257        assert!(
258            result.contains("_atom_site.Cartn_x"),
259            "missing Cartn_x header"
260        );
261        assert!(result.contains("ATOM"), "missing ATOM records");
262        // First residue is M → MET
263        assert!(result.contains("MET"), "missing MET residue name");
264    }
265
266    #[test]
267    fn test_coords_to_mmcif_ca_only() {
268        let sequence = "GS";
269        let vals: Vec<f32> = vec![
270            1.0, 2.0, 3.0, // Gly CA
271            4.0, 5.0, 6.0, // Ser CA
272        ];
273        let coords = Tensor::from_vec(vals, &[2, 3], &Device::Cpu).unwrap();
274        let plddt = make_plddt(2);
275
276        let result = coords_to_mmcif(&coords, sequence, &plddt, "B", "ca_only").unwrap();
277
278        assert!(result.contains("data_ca_only"));
279        assert!(result.contains("GLY"));
280        assert!(result.contains("SER"));
281        // Only CA atoms should appear
282        let atom_lines: Vec<&str> = result.lines().filter(|l| l.starts_with("ATOM")).collect();
283        assert_eq!(atom_lines.len(), 2, "expected 2 CA atoms");
284        for line in &atom_lines {
285            assert!(line.contains(" CA "), "Cα-only: each atom should be CA");
286        }
287    }
288
289    #[test]
290    fn test_coords_to_mmcif_batched_squeeze() {
291        // Shape (1, 2, 5, 3) — batch dim should be squeezed automatically
292        let inner = make_coords(2, 5);
293        let batched = inner.unsqueeze(0).unwrap(); // (1, 2, 5, 3)
294        let plddt = make_plddt(2).unsqueeze(0).unwrap(); // (1, 2)
295
296        let result = coords_to_mmcif(&batched, "MA", &plddt, "A", "batched").unwrap();
297
298        assert!(result.contains("data_batched"));
299        assert!(result.contains("ATOM"));
300    }
301
302    #[test]
303    fn test_mmcif_output_scannable_for_atom_site() {
304        // Verify the output can be line-scanned for _atom_site — mimicking what
305        // ferritin-core's reader does before parsing.
306        let sequence = "ACDE";
307        let coords = make_coords(4, 5);
308        let plddt = make_plddt(4);
309
310        let result = coords_to_mmcif(&coords, sequence, &plddt, "A", "scan_test").unwrap();
311
312        let has_loop = result.lines().any(|l| l == "loop_");
313        let has_atom_site_col = result.lines().any(|l| l.starts_with("_atom_site."));
314        let has_atom_record = result.lines().any(|l| l.starts_with("ATOM"));
315
316        assert!(has_loop, "output must contain loop_");
317        assert!(
318            has_atom_site_col,
319            "output must contain _atom_site.* headers"
320        );
321        assert!(has_atom_record, "output must contain ATOM records");
322
323        // All ATOM lines should have 21 space-separated fields
324        for line in result.lines().filter(|l| l.starts_with("ATOM")) {
325            let fields: Vec<&str> = line.split_whitespace().collect();
326            assert_eq!(
327                fields.len(),
328                21,
329                "each ATOM line must have 21 fields, got {}: {:?}",
330                fields.len(),
331                line
332            );
333        }
334    }
335
336    #[test]
337    fn test_one_to_three() {
338        assert_eq!(one_to_three('A'), "ALA");
339        assert_eq!(one_to_three('G'), "GLY");
340        assert_eq!(one_to_three('W'), "TRP");
341        assert_eq!(one_to_three('X'), "UNK");
342    }
343
344    #[test]
345    fn test_skip_near_zero_atoms() {
346        // Residue 0 has atoms at zero (should be skipped), residue 1 has real coords
347        let vals: Vec<f32> = vec![
348            // residue 0: 3 atoms all zero
349            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
350            // residue 1: 3 atoms with real coords
351            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0,
352        ];
353        let coords = Tensor::from_vec(vals, &[2, 3, 3], &Device::Cpu).unwrap();
354        let plddt = make_plddt(2);
355
356        let result = coords_to_mmcif(&coords, "AG", &plddt, "A", "zero_test").unwrap();
357
358        let atom_lines: Vec<&str> = result.lines().filter(|l| l.starts_with("ATOM")).collect();
359        // Residue 0's 3 atoms are all zero → skipped. Residue 1 has 3 real atoms.
360        assert_eq!(atom_lines.len(), 3, "only non-zero atoms should appear");
361    }
362}
ferritin_plms/esmfold2/mmcif.rs

ferritin_plms/esmfold2/
mmcif.rs
