ferritin_core/columns/
topology.rs

1//! Owned, boundary-ready snapshot of an [`AtomicHierarchy`].
2
3use super::buffer::{Column, ColumnData};
4use super::dictionary::StringDictionary;
5use super::error::{BoundaryError, BoundaryErrorKind, to_u32};
6use super::frame::FrameColumns;
7use super::gpu::{BufferLayout, GpuBuffer};
8use super::handles::{AtomHandle, ChainHandle, ResidueHandle};
9use crate::model::{AtomicHierarchy, Model, ResidueGroup};
10use std::borrow::Cow;
11use std::ops::Range;
12use std::sync::Arc;
13
14/// `atom.alt_loc` value for atoms without an alternate-location indicator.
15pub const NO_ALT_LOC: u32 = 0;
16/// `residue.ins_code` value for residues without an insertion code.
17pub const NO_INS_CODE: u32 = 0;
18/// `atom.formal_charge` value for atoms without a recorded formal charge.
19pub const NO_FORMAL_CHARGE: i8 = i8::MIN;
20/// `residue.group` value for polymer residues.
21pub const RESIDUE_GROUP_POLYMER: u8 = 0;
22/// `residue.group` value for non-polymer residues.
23pub const RESIDUE_GROUP_NON_POLYMER: u8 = 1;
24
25/// Boundary-ready, frame-independent topology columns.
26///
27/// Built once per `Arc<AtomicHierarchy>` and reused for every frame that shares
28/// it (see [`Self::frame`]). Columns whose in-memory representation is already
29/// a plain numeric array (sequence ids, segment offsets, bonds) are borrowed
30/// from the hierarchy; everything else is encoded into owned buffers:
31///
32/// | Column | Type | Encoding |
33/// |---|---|---|
34/// | `atom.element` | `u8` | atomic number |
35/// | `atom.alt_loc` | `u32` | Unicode scalar, [`NO_ALT_LOC`] if absent |
36/// | `atom.formal_charge` | `i8` | charge, [`NO_FORMAL_CHARGE`] if absent |
37/// | `residue.ins_code` | `u32` | Unicode scalar, [`NO_INS_CODE`] if absent |
38/// | `residue.group` | `u8` | [`RESIDUE_GROUP_POLYMER`] / [`RESIDUE_GROUP_NON_POLYMER`] |
39/// | `residue.label_seq_id` | `i32` | [`MISSING_SEQ_ID`](crate::model::MISSING_SEQ_ID) if absent |
40/// | string fields | `u32` codes | [`StringDictionary`] |
41///
42/// # Ownership and invalidation
43///
44/// The snapshot holds a clone of the `Arc<AtomicHierarchy>` it was built from.
45/// While it is alive, that hierarchy cannot be mutated in place: `Arc::get_mut`
46/// fails and `Arc::make_mut` clones. So the snapshot never silently goes out of
47/// date — a model whose hierarchy was replaced or copied-on-write simply stops
48/// matching, which [`Self::is_current_for`] and [`Self::frame`] detect by
49/// pointer identity.
50#[derive(Clone, Debug)]
51pub struct TopologyColumns {
52    hierarchy: Arc<AtomicHierarchy>,
53    n_atoms: u32,
54    n_residues: u32,
55    n_chains: u32,
56    n_bonds: u32,
57    atom_name: StringDictionary,
58    auth_atom_name: StringDictionary,
59    element: Vec<u8>,
60    alt_loc: Vec<u32>,
61    formal_charge: Vec<i8>,
62    comp_id: StringDictionary,
63    auth_comp_id: StringDictionary,
64    ins_code: Vec<u32>,
65    residue_group: Vec<u8>,
66    label_asym_id: StringDictionary,
67    auth_asym_id: StringDictionary,
68    entity_id: StringDictionary,
69}
70
71fn encode_char(field: &str, row: usize, value: Option<char>) -> Result<u32, BoundaryError> {
72    match value {
73        None => Ok(0),
74        Some('\0') => Err(BoundaryError::new(
75            BoundaryErrorKind::InvalidModel,
76            format!("{field}[{row}] is NUL, which is reserved for 'absent'"),
77        )),
78        Some(c) => Ok(c as u32),
79    }
80}
81
82fn dictionary(values: &[String]) -> Result<StringDictionary, BoundaryError> {
83    StringDictionary::encode(values.iter().map(String::as_str))
84}
85
86impl TopologyColumns {
87    /// Validate `hierarchy` and encode its columns.
88    pub fn encode(hierarchy: &Arc<AtomicHierarchy>) -> Result<Self, BoundaryError> {
89        hierarchy.validate()?;
90        let atoms = &hierarchy.atoms;
91        let residues = &hierarchy.residues;
92        let chains = &hierarchy.chains;
93
94        let alt_loc = atoms
95            .alt_loc
96            .iter()
97            .enumerate()
98            .map(|(i, &c)| encode_char("atom.alt_loc", i, c))
99            .collect::<Result<_, _>>()?;
100        let formal_charge = atoms
101            .formal_charge
102            .iter()
103            .enumerate()
104            .map(|(i, &charge)| match charge {
105                None => Ok(NO_FORMAL_CHARGE),
106                Some(NO_FORMAL_CHARGE) => Err(BoundaryError::new(
107                    BoundaryErrorKind::InvalidModel,
108                    format!("atom.formal_charge[{i}] is {NO_FORMAL_CHARGE}, which is reserved for 'absent'"),
109                )),
110                Some(charge) => Ok(charge),
111            })
112            .collect::<Result<_, _>>()?;
113        let ins_code = residues
114            .ins_code
115            .iter()
116            .enumerate()
117            .map(|(i, &c)| encode_char("residue.ins_code", i, c))
118            .collect::<Result<_, _>>()?;
119        let residue_group = residues
120            .group
121            .iter()
122            .map(|group| match group {
123                ResidueGroup::Polymer => RESIDUE_GROUP_POLYMER,
124                ResidueGroup::NonPolymer => RESIDUE_GROUP_NON_POLYMER,
125            })
126            .collect();
127
128        Ok(Self {
129            hierarchy: Arc::clone(hierarchy),
130            // `atom_bond_starts` holds n_atoms + 1 entries, so it bounds the atom count too.
131            n_atoms: to_u32("atom count", atoms.len())?,
132            n_residues: to_u32("residue count", residues.len())?,
133            n_chains: to_u32("chain count", chains.len())?,
134            n_bonds: to_u32("bond count", hierarchy.bonds.len())?,
135            atom_name: dictionary(&atoms.atom_name)?,
136            auth_atom_name: dictionary(&atoms.auth_atom_name)?,
137            element: atoms
138                .element
139                .iter()
140                .map(|e| e.atomic_number() as u8)
141                .collect(),
142            alt_loc,
143            formal_charge,
144            comp_id: dictionary(&residues.comp_id)?,
145            auth_comp_id: dictionary(&residues.auth_comp_id)?,
146            ins_code,
147            residue_group,
148            label_asym_id: dictionary(&chains.label_asym_id)?,
149            auth_asym_id: dictionary(&chains.auth_asym_id)?,
150            entity_id: dictionary(&chains.entity_id)?,
151        })
152    }
153
154    /// Encode the topology of `model`.
155    pub fn from_model(model: &Model) -> Result<Self, BoundaryError> {
156        Self::encode(&model.hierarchy)
157    }
158
159    /// The hierarchy this snapshot was built from (and keeps alive).
160    pub fn hierarchy(&self) -> &Arc<AtomicHierarchy> {
161        &self.hierarchy
162    }
163
164    /// Whether `model` still uses the exact hierarchy this snapshot describes.
165    pub fn is_current_for(&self, model: &Model) -> bool {
166        Arc::ptr_eq(&self.hierarchy, &model.hierarchy)
167    }
168
169    /// Borrow the per-frame columns of `model`, checking it shares this topology.
170    ///
171    /// Fails with [`StaleTopology`](BoundaryErrorKind::StaleTopology) if it does not.
172    pub fn frame<'m>(&self, model: &'m Model) -> Result<FrameColumns<'m>, BoundaryError> {
173        if !self.is_current_for(model) {
174            return Err(BoundaryError::new(
175                BoundaryErrorKind::StaleTopology,
176                "model does not share the hierarchy this topology snapshot was built from",
177            ));
178        }
179        FrameColumns::new(model)
180    }
181
182    /// Number of atoms.
183    pub fn n_atoms(&self) -> u32 {
184        self.n_atoms
185    }
186
187    /// Number of residues.
188    pub fn n_residues(&self) -> u32 {
189        self.n_residues
190    }
191
192    /// Number of chains.
193    pub fn n_chains(&self) -> u32 {
194        self.n_chains
195    }
196
197    /// Number of bonds.
198    pub fn n_bonds(&self) -> u32 {
199        self.n_bonds
200    }
201
202    // --- atoms -------------------------------------------------------------
203
204    /// `atom.name` (`label_atom_id`).
205    pub fn atom_name(&self) -> &StringDictionary {
206        &self.atom_name
207    }
208
209    /// `atom.auth_name` (`auth_atom_id`).
210    pub fn auth_atom_name(&self) -> &StringDictionary {
211        &self.auth_atom_name
212    }
213
214    /// `atom.element` as atomic numbers.
215    pub fn element(&self) -> &[u8] {
216        &self.element
217    }
218
219    /// `atom.alt_loc` as Unicode scalars ([`NO_ALT_LOC`] if absent).
220    pub fn alt_loc(&self) -> &[u32] {
221        &self.alt_loc
222    }
223
224    /// `atom.formal_charge` ([`NO_FORMAL_CHARGE`] if absent).
225    pub fn formal_charge(&self) -> &[i8] {
226        &self.formal_charge
227    }
228
229    // --- residues ----------------------------------------------------------
230
231    /// `residue.comp_id` (`label_comp_id`).
232    pub fn comp_id(&self) -> &StringDictionary {
233        &self.comp_id
234    }
235
236    /// `residue.auth_comp_id`.
237    pub fn auth_comp_id(&self) -> &StringDictionary {
238        &self.auth_comp_id
239    }
240
241    /// `residue.label_seq_id` (borrowed from the hierarchy).
242    pub fn label_seq_id(&self) -> &[i32] {
243        &self.hierarchy.residues.label_seq_id
244    }
245
246    /// `residue.auth_seq_id` (borrowed from the hierarchy).
247    pub fn auth_seq_id(&self) -> &[i32] {
248        &self.hierarchy.residues.auth_seq_id
249    }
250
251    /// `residue.ins_code` as Unicode scalars ([`NO_INS_CODE`] if absent).
252    pub fn ins_code(&self) -> &[u32] {
253        &self.ins_code
254    }
255
256    /// `residue.group` codes.
257    pub fn residue_group(&self) -> &[u8] {
258        &self.residue_group
259    }
260
261    // --- chains ------------------------------------------------------------
262
263    /// `chain.label_asym_id`.
264    pub fn label_asym_id(&self) -> &StringDictionary {
265        &self.label_asym_id
266    }
267
268    /// `chain.auth_asym_id`.
269    pub fn auth_asym_id(&self) -> &StringDictionary {
270        &self.auth_asym_id
271    }
272
273    /// `chain.entity_id`.
274    pub fn entity_id(&self) -> &StringDictionary {
275        &self.entity_id
276    }
277
278    // --- hierarchy segmentation and bonds (borrowed) -------------------------
279
280    /// CSR offsets grouping atoms into residues (`n_residues + 1` entries).
281    pub fn residue_atom_offsets(&self) -> &[u32] {
282        self.hierarchy.atom_to_residue.offsets()
283    }
284
285    /// CSR offsets grouping residues into chains (`n_chains + 1` entries).
286    pub fn chain_residue_offsets(&self) -> &[u32] {
287        self.hierarchy.residue_to_chain.offsets()
288    }
289
290    /// First atom of each bond, sorted ascending.
291    pub fn bond_atom_a(&self) -> &[u32] {
292        &self.hierarchy.bonds.atom_a
293    }
294
295    /// Second atom of each bond.
296    pub fn bond_atom_b(&self) -> &[u32] {
297        &self.hierarchy.bonds.atom_b
298    }
299
300    /// Bond order of each bond.
301    pub fn bond_order(&self) -> &[u8] {
302        &self.hierarchy.bonds.order
303    }
304
305    /// CSR offsets grouping bonds by `atom_a` (`n_atoms + 1` entries).
306    pub fn atom_bond_offsets(&self) -> &[u32] {
307        &self.hierarchy.bonds.atom_bond_starts
308    }
309
310    /// Every numeric topology column, including dictionary codes, by stable name.
311    ///
312    /// Dictionary value tables are not listed; fetch them from the matching
313    /// [`StringDictionary`] accessor.
314    pub fn columns(&self) -> Vec<Column<'_>> {
315        use ColumnData::{I8, I32, U8, U32};
316        [
317            ("atom.name", U32(self.atom_name.codes())),
318            ("atom.auth_name", U32(self.auth_atom_name.codes())),
319            ("atom.element", U8(&self.element)),
320            ("atom.alt_loc", U32(&self.alt_loc)),
321            ("atom.formal_charge", I8(&self.formal_charge)),
322            ("residue.comp_id", U32(self.comp_id.codes())),
323            ("residue.auth_comp_id", U32(self.auth_comp_id.codes())),
324            ("residue.label_seq_id", I32(self.label_seq_id())),
325            ("residue.auth_seq_id", I32(self.auth_seq_id())),
326            ("residue.ins_code", U32(&self.ins_code)),
327            ("residue.group", U8(&self.residue_group)),
328            ("residue.atom_offsets", U32(self.residue_atom_offsets())),
329            ("chain.label_asym_id", U32(self.label_asym_id.codes())),
330            ("chain.auth_asym_id", U32(self.auth_asym_id.codes())),
331            ("chain.entity_id", U32(self.entity_id.codes())),
332            ("chain.residue_offsets", U32(self.chain_residue_offsets())),
333            ("bond.atom_a", U32(self.bond_atom_a())),
334            ("bond.atom_b", U32(self.bond_atom_b())),
335            ("bond.order", U8(self.bond_order())),
336            ("atom.bond_offsets", U32(self.atom_bond_offsets())),
337        ]
338        .into_iter()
339        .map(|(name, data)| Column { name, data })
340        .collect()
341    }
342
343    /// Bonds as an interleaved `[atom_a, atom_b]` `u32` index buffer, ready to
344    /// use as a GPU line-list index buffer. Always an owned copy.
345    pub fn bond_line_indices(&self) -> GpuBuffer<'static> {
346        let bonds = &self.hierarchy.bonds;
347        let mut bytes = Vec::with_capacity(bonds.len() * 8);
348        for (&a, &b) in bonds.atom_a.iter().zip(&bonds.atom_b) {
349            bytes.extend_from_slice(&a.to_le_bytes());
350            bytes.extend_from_slice(&b.to_le_bytes());
351        }
352        GpuBuffer {
353            label: "bond.line_indices",
354            layout: BufferLayout::new(super::ScalarType::U32, 2, 8, self.n_bonds),
355            bytes: Cow::Owned(bytes),
356        }
357    }
358
359    // --- checked handle navigation ------------------------------------------
360
361    /// Residue containing `atom`.
362    pub fn residue_of(&self, atom: AtomHandle) -> Result<ResidueHandle, BoundaryError> {
363        if atom.0 >= self.n_atoms {
364            return Err(BoundaryError::out_of_range("atom", atom.0, self.n_atoms));
365        }
366        Ok(ResidueHandle(
367            self.hierarchy.residue_of_atom(atom.index()) as u32
368        ))
369    }
370
371    /// Chain containing `residue`.
372    pub fn chain_of(&self, residue: ResidueHandle) -> Result<ChainHandle, BoundaryError> {
373        if residue.0 >= self.n_residues {
374            return Err(BoundaryError::out_of_range(
375                "residue",
376                residue.0,
377                self.n_residues,
378            ));
379        }
380        Ok(ChainHandle(
381            self.hierarchy.chain_of_residue(residue.index()) as u32,
382        ))
383    }
384
385    /// Atom rows belonging to `residue`.
386    pub fn atoms_of(&self, residue: ResidueHandle) -> Result<Range<u32>, BoundaryError> {
387        let offsets = self.residue_atom_offsets();
388        if residue.0 >= self.n_residues {
389            return Err(BoundaryError::out_of_range(
390                "residue",
391                residue.0,
392                self.n_residues,
393            ));
394        }
395        Ok(offsets[residue.index()]..offsets[residue.index() + 1])
396    }
397
398    /// Residue rows belonging to `chain`.
399    pub fn residues_of(&self, chain: ChainHandle) -> Result<Range<u32>, BoundaryError> {
400        let offsets = self.chain_residue_offsets();
401        if chain.0 >= self.n_chains {
402            return Err(BoundaryError::out_of_range("chain", chain.0, self.n_chains));
403        }
404        Ok(offsets[chain.index()]..offsets[chain.index() + 1])
405    }
406}
ferritin_core/columns/topology.rs

ferritin_core/columns/
topology.rs
