ferritin_core/columns/
mod.rs

1//! Stable columnar boundary for WASM, FFI, and GPU consumers.
2//!
3//! The [`Model`](crate::Model) layer is already struct-of-arrays, but its
4//! public fields are Rust types (`Vec<String>`, `Option<char>`, `Element`,
5//! `Arc`, borrowed views) that cannot cross a wasm-bindgen or C ABI boundary,
6//! or be uploaded to a GPU, as they are. This module is the contract for doing
7//! so. It adds no dependencies: bindings crates (wasm-bindgen, PyO3, a C
8//! header) and renderers (bevy, wgpu) build on it.
9//!
10//! # What crosses the boundary
11//!
12//! * **Numeric columns** — contiguous `u8`/`i8`/`u32`/`i32`/`f32` slices
13//!   ([`ColumnData`]), each with a stable name ([`Column`]). Bytes are
14//!   little-endian ([`ColumnData::le_bytes`]).
15//! * **String columns** — dictionary-encoded ([`StringDictionary`]): one `u32`
16//!   code per row plus a UTF-8 value buffer and offsets.
17//! * **Optional values** — encoded in place with sentinels
18//!   ([`NO_ALT_LOC`], [`NO_INS_CODE`], [`NO_FORMAL_CHARGE`],
19//!   [`MISSING_SEQ_ID`](crate::model::MISSING_SEQ_ID)), never as a separate
20//!   validity bitmap.
21//! * **Handles** — [`AtomHandle`], [`ResidueHandle`], [`ChainHandle`]: `u32`
22//!   row indices into a topology snapshot, navigated with checked lookups.
23//! * **Errors** — [`BoundaryError`], with a stable numeric
24//!   [`code`](BoundaryError::code) instead of panics.
25//! * **GPU uploads** — [`GpuBuffer`]: bytes plus a [`BufferLayout`]
26//!   (scalar type, components, stride, count).
27//!
28//! # Ownership, borrowing, and invalidation
29//!
30//! | Data | Source | Copy or borrow | Valid until |
31//! |---|---|---|---|
32//! | `x`, `y`, `z`, `occupancy`, `b_iso`, `confidence` | [`FrameColumns`] | borrow of the model's conformation | the model is dropped or its conformation mutated |
33//! | sequence ids, segment offsets, bonds | [`TopologyColumns`] | borrow of the pinned hierarchy | the snapshot is dropped |
34//! | element, alt_loc, charge, ins_code, group, dictionaries | [`TopologyColumns`] | owned, encoded once per hierarchy | the snapshot is dropped |
35//! | interleaved positions, bond line indices | [`GpuBuffer`] | owned copy | never (independent of the model) |
36//!
37//! Inside Rust the borrow checker enforces the "valid until" column. Across
38//! FFI it is the caller's obligation: a raw pointer from
39//! [`ColumnData::as_ptr`] is valid exactly as long as the Rust borrow it came
40//! from.
41//!
42//! A [`TopologyColumns`] snapshot keeps its `Arc<AtomicHierarchy>` alive and
43//! therefore immutable (`Arc::make_mut` will clone rather than mutate). A
44//! topology edit thus produces a *new* hierarchy, and pairing a frame with the
45//! old snapshot fails with [`StaleTopology`](BoundaryErrorKind::StaleTopology)
46//! rather than exporting mismatched columns. All frames of a trajectory share
47//! one hierarchy, so one snapshot serves every frame.
48//!
49//! ## wasm-bindgen
50//!
51//! A zero-copy JS typed-array view over a borrowed column (e.g.
52//! `js_sys::Float32Array::view`) points into wasm linear memory and is
53//! **detached by any allocation that grows that memory**, including
54//! allocations made later by ferritin itself. Use views only for an immediate
55//! copy or upload within the same call; copy (`Float32Array::from`) anything
56//! JS keeps.
57//!
58//! # Example
59//!
60//! ```text
61//! let topology = TopologyColumns::from_model(&model)?;   // once per hierarchy
62//! let frame = topology.frame(&model)?;                   // per frame, zero-copy
63//! upload(frame.positions(PositionLayout::Padded4));      // owned, 16-byte stride
64//! upload(GpuBuffer::from_column(column)?);               // borrowed scalar column
65//! upload(topology.bond_line_indices());                  // owned [a, b] pairs
66//! ```
67
68mod buffer;
69mod dictionary;
70mod error;
71mod frame;
72mod gpu;
73mod handles;
74mod topology;
75
76pub use buffer::{Column, ColumnData, ScalarType};
77pub use dictionary::StringDictionary;
78pub use error::{BoundaryError, BoundaryErrorKind};
79pub use frame::FrameColumns;
80pub use gpu::{BufferLayout, GpuBuffer, PositionLayout};
81pub use handles::{AtomHandle, ChainHandle, ResidueHandle};
82pub use topology::{
83    NO_ALT_LOC, NO_FORMAL_CHARGE, NO_INS_CODE, RESIDUE_GROUP_NON_POLYMER, RESIDUE_GROUP_POLYMER,
84    TopologyColumns,
85};
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use crate::data::Segmentation;
91    use crate::info::elements::Element;
92    use crate::model::{
93        AtomicConformation, AtomicHierarchy, AtomsTable, Bonds, ChainsTable, MISSING_SEQ_ID, Model,
94        ResidueGroup, ResiduesTable,
95    };
96    use std::borrow::Cow;
97    use std::sync::Arc;
98
99    /// Two chains: chain A has ALA (N, CA, C) + GLY (N, CA); chain B has one HOH (O).
100    fn make_model() -> Model {
101        let atoms = AtomsTable {
102            atom_name: ["N", "CA", "C", "N", "CA", "O"].map(String::from).to_vec(),
103            auth_atom_name: ["N", "CA", "C", "N", "CA", "O"].map(String::from).to_vec(),
104            element: vec![
105                Element::N,
106                Element::C,
107                Element::C,
108                Element::N,
109                Element::C,
110                Element::O,
111            ],
112            alt_loc: vec![None, Some('A'), None, None, None, None],
113            formal_charge: vec![None, None, None, Some(1), None, Some(-2)],
114        };
115        let residues = ResiduesTable {
116            comp_id: ["ALA", "GLY", "HOH"].map(String::from).to_vec(),
117            auth_comp_id: ["ALA", "GLY", "HOH"].map(String::from).to_vec(),
118            label_seq_id: vec![1, 2, MISSING_SEQ_ID],
119            auth_seq_id: vec![10, 11, 101],
120            ins_code: vec![None, Some('B'), None],
121            group: vec![
122                ResidueGroup::Polymer,
123                ResidueGroup::Polymer,
124                ResidueGroup::NonPolymer,
125            ],
126        };
127        let chains = ChainsTable {
128            label_asym_id: ["A", "B"].map(String::from).to_vec(),
129            auth_asym_id: ["A", "B"].map(String::from).to_vec(),
130            entity_id: ["1", "2"].map(String::from).to_vec(),
131        };
132        let hierarchy = AtomicHierarchy {
133            atoms,
134            residues,
135            chains,
136            atom_to_residue: Segmentation::from_offsets(vec![0, 3, 5, 6]),
137            residue_to_chain: Segmentation::from_offsets(vec![0, 2, 3]),
138            bonds: Bonds::from_unsorted(vec![1, 0, 2], vec![2, 1, 3], vec![1, 1, 1], 6),
139        };
140        let conformation = AtomicConformation {
141            x: vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
142            y: vec![0.5; 6],
143            z: vec![-1.0; 6],
144            occupancy: Some(vec![1.0; 6]),
145            b_iso: None,
146            confidence: None,
147        };
148        Model::new(Arc::new(hierarchy), conformation)
149    }
150
151    #[test]
152    fn test_topology_encodes_scalars_and_sentinels() {
153        let model = make_model();
154        let topology = TopologyColumns::from_model(&model).unwrap();
155
156        assert_eq!(topology.n_atoms(), 6);
157        assert_eq!(topology.n_residues(), 3);
158        assert_eq!(topology.n_chains(), 2);
159        assert_eq!(topology.n_bonds(), 3);
160        assert_eq!(topology.element(), &[7, 6, 6, 7, 6, 8]);
161        assert_eq!(topology.alt_loc(), &[NO_ALT_LOC, 'A' as u32, 0, 0, 0, 0]);
162        assert_eq!(
163            topology.formal_charge(),
164            &[
165                NO_FORMAL_CHARGE,
166                NO_FORMAL_CHARGE,
167                NO_FORMAL_CHARGE,
168                1,
169                NO_FORMAL_CHARGE,
170                -2
171            ]
172        );
173        assert_eq!(topology.ins_code(), &[NO_INS_CODE, 'B' as u32, NO_INS_CODE]);
174        assert_eq!(
175            topology.residue_group(),
176            &[
177                RESIDUE_GROUP_POLYMER,
178                RESIDUE_GROUP_POLYMER,
179                RESIDUE_GROUP_NON_POLYMER
180            ]
181        );
182        assert_eq!(topology.label_seq_id(), &[1, 2, MISSING_SEQ_ID]);
183        assert_eq!(topology.residue_atom_offsets(), &[0, 3, 5, 6]);
184        assert_eq!(topology.chain_residue_offsets(), &[0, 2, 3]);
185    }
186
187    #[test]
188    fn test_topology_dictionaries_roundtrip() {
189        let model = make_model();
190        let topology = TopologyColumns::from_model(&model).unwrap();
191        let atoms = &model.hierarchy.atoms;
192
193        assert_eq!(topology.atom_name().n_values(), 4); // N, CA, C, O
194        for (i, name) in atoms.atom_name.iter().enumerate() {
195            assert_eq!(topology.atom_name().get(i), Some(name.as_str()));
196        }
197        assert_eq!(topology.comp_id().codes(), &[0, 1, 2]);
198        assert_eq!(topology.entity_id().get(1), Some("2"));
199    }
200
201    #[test]
202    fn test_topology_borrows_hierarchy_columns() {
203        let model = make_model();
204        let topology = TopologyColumns::from_model(&model).unwrap();
205        assert!(std::ptr::eq(
206            topology.auth_seq_id(),
207            model.hierarchy.residues.auth_seq_id.as_slice()
208        ));
209        assert!(std::ptr::eq(
210            topology.bond_atom_b(),
211            model.hierarchy.bonds.atom_b.as_slice()
212        ));
213    }
214
215    #[test]
216    fn test_topology_columns_lengths_and_names_are_unique() {
217        let model = make_model();
218        let topology = TopologyColumns::from_model(&model).unwrap();
219        let columns = topology.columns();
220
221        let mut names: Vec<&str> = columns.iter().map(|c| c.name).collect();
222        names.sort_unstable();
223        names.dedup();
224        assert_eq!(names.len(), columns.len());
225
226        for column in &columns {
227            // CSR offsets carry one sentinel entry beyond their parent table.
228            let expected = match column.name {
229                "residue.atom_offsets" => topology.n_residues() + 1,
230                "chain.residue_offsets" => topology.n_chains() + 1,
231                "atom.bond_offsets" => topology.n_atoms() + 1,
232                name => match name.split('.').next().unwrap() {
233                    "atom" => topology.n_atoms(),
234                    "residue" => topology.n_residues(),
235                    "chain" => topology.n_chains(),
236                    "bond" => topology.n_bonds(),
237                    other => panic!("unexpected column prefix {other}"),
238                },
239            };
240            assert_eq!(column.data.len() as u32, expected, "{}", column.name);
241        }
242    }
243
244    #[test]
245    fn test_handle_navigation_is_checked() {
246        let model = make_model();
247        let topology = TopologyColumns::from_model(&model).unwrap();
248
249        assert_eq!(
250            topology.residue_of(AtomHandle(4)).unwrap(),
251            ResidueHandle(1)
252        );
253        assert_eq!(topology.chain_of(ResidueHandle(2)).unwrap(), ChainHandle(1));
254        assert_eq!(topology.atoms_of(ResidueHandle(1)).unwrap(), 3..5);
255        assert_eq!(topology.residues_of(ChainHandle(0)).unwrap(), 0..2);
256
257        let err = topology.residue_of(AtomHandle(6)).unwrap_err();
258        assert_eq!(err.kind(), BoundaryErrorKind::IndexOutOfRange);
259        assert_eq!(err.code(), 2);
260        assert_eq!(err.message(), "atom 6 is out of range (len=6)");
261        assert!(topology.residues_of(ChainHandle(2)).is_err());
262    }
263
264    #[test]
265    fn test_frame_columns_are_zero_copy() {
266        let model = make_model();
267        let topology = TopologyColumns::from_model(&model).unwrap();
268        let frame = topology.frame(&model).unwrap();
269
270        assert!(std::ptr::eq(frame.x(), model.conformation.x.as_slice()));
271        assert_eq!(frame.occupancy().map(<[f32]>::len), Some(6));
272        assert!(frame.b_iso().is_none());
273
274        let names: Vec<&str> = frame.columns().iter().map(|c| c.name).collect();
275        assert_eq!(names, vec!["x", "y", "z", "occupancy"]);
276
277        let buffer = GpuBuffer::from_column(frame.columns()[0]).unwrap();
278        assert_eq!(buffer.layout, BufferLayout::new(ScalarType::F32, 1, 4, 6));
279        if cfg!(target_endian = "little") {
280            assert!(matches!(buffer.bytes, Cow::Borrowed(_)));
281        }
282    }
283
284    #[test]
285    fn test_frames_sharing_a_hierarchy_reuse_one_snapshot() {
286        let model = make_model();
287        let topology = TopologyColumns::from_model(&model).unwrap();
288        let mut conformation = model.conformation.clone();
289        conformation.x[0] = 42.0;
290        let other_frame = Model::new(Arc::clone(&model.hierarchy), conformation);
291
292        let frame = topology.frame(&other_frame).unwrap();
293        assert_eq!(frame.x()[0], 42.0);
294    }
295
296    #[test]
297    fn test_topology_edit_makes_snapshot_stale() {
298        let mut model = make_model();
299        let topology = TopologyColumns::from_model(&model).unwrap();
300        assert!(topology.is_current_for(&model));
301
302        // The snapshot pins the hierarchy, so a topology edit must copy-on-write.
303        Arc::make_mut(&mut model.hierarchy).atoms.atom_name[0] = "N1".to_string();
304        assert!(!topology.is_current_for(&model));
305        assert_eq!(topology.atom_name().get(0), Some("N"));
306
307        let err = topology.frame(&model).unwrap_err();
308        assert_eq!(err.kind(), BoundaryErrorKind::StaleTopology);
309        assert_eq!(err.code(), 3);
310    }
311
312    #[test]
313    fn test_reserved_sentinel_values_are_rejected() {
314        let model = make_model();
315        let mut hierarchy = (*model.hierarchy).clone();
316        hierarchy.atoms.formal_charge[0] = Some(NO_FORMAL_CHARGE);
317        let err = TopologyColumns::encode(&Arc::new(hierarchy)).unwrap_err();
318        assert_eq!(err.kind(), BoundaryErrorKind::InvalidModel);
319
320        let mut hierarchy = (*model.hierarchy).clone();
321        hierarchy.residues.ins_code[0] = Some('\0');
322        let err = TopologyColumns::encode(&Arc::new(hierarchy)).unwrap_err();
323        assert_eq!(err.kind(), BoundaryErrorKind::InvalidModel);
324    }
325
326    #[test]
327    fn test_invalid_hierarchy_maps_to_invalid_model() {
328        let model = make_model();
329        let mut hierarchy = (*model.hierarchy).clone();
330        hierarchy.atoms.element.pop();
331        let err = TopologyColumns::encode(&Arc::new(hierarchy)).unwrap_err();
332        assert_eq!(err.kind(), BoundaryErrorKind::InvalidModel);
333        assert_eq!(err.code(), 1);
334        assert_eq!(err.message(), "atoms.element has length 5, expected 6");
335    }
336
337    #[test]
338    fn test_positions_layouts() {
339        let model = make_model();
340        let frame = FrameColumns::new(&model).unwrap();
341
342        let packed = frame.positions(PositionLayout::Packed3);
343        assert_eq!(packed.layout.stride_bytes, 12);
344        assert_eq!(packed.bytes.len() as u64, packed.layout.byte_len());
345        assert_eq!(&packed.bytes[12..16], &1.0f32.to_le_bytes()); // atom 1 x
346
347        let padded = frame.positions(PositionLayout::Padded4);
348        assert_eq!(padded.layout, BufferLayout::new(ScalarType::F32, 4, 16, 6));
349        assert_eq!(padded.bytes.len(), 6 * 16);
350        assert_eq!(&padded.bytes[12..16], &1.0f32.to_le_bytes()); // atom 0 w
351        assert_eq!(&padded.bytes[16..20], &1.0f32.to_le_bytes()); // atom 1 x
352        assert_eq!(&padded.bytes[20..24], &0.5f32.to_le_bytes()); // atom 1 y
353    }
354
355    #[test]
356    fn test_bond_line_indices() {
357        let model = make_model();
358        let topology = TopologyColumns::from_model(&model).unwrap();
359        let lines = topology.bond_line_indices();
360
361        assert_eq!(lines.layout, BufferLayout::new(ScalarType::U32, 2, 8, 3));
362        let pairs: Vec<u32> = lines
363            .bytes
364            .as_chunks::<4>()
365            .0
366            .iter()
367            .map(|&b| u32::from_le_bytes(b))
368            .collect();
369        // Bonds are sorted by atom_a.
370        assert_eq!(pairs, vec![0, 1, 1, 2, 2, 3]);
371    }
372}
ferritin_core/columns/mod.rs

ferritin_core/columns/
mod.rs
