molframe_core/chunk/
builder.rs

1//! Building chunks from a stream of atoms in file order.
2//!
3//! A chunk is closed once it is at or past the target size *and* the next atom
4//! starts a new residue. That is what keeps a chunk from straddling a residue
5//! without splitting a large one, which in turn is what bounds the number of
6//! chunks any residue-level operation has to touch.
7
8use super::atom::{AtomChunk, AtomRecord};
9use super::parent::ParentMapping;
10use super::stats::AtomChunkStats;
11use crate::column::{EncodedColumn, Presence, ValidityMask};
12use crate::coords::CoordinateBlock;
13use crate::symbol::SymbolId;
14
15/// Atoms per chunk before a new one is started.
16///
17/// Chosen so the hot columns come to roughly a quarter of a megabyte, which is
18/// what stays resident in second-level cache on the hardware this targets. It is
19/// a tuned constant, not a promise, and nothing outside this crate depends on it.
20pub const TARGET_CHUNK_ATOMS: u32 = 4096;
21
22/// Accumulates atoms and closes chunks at residue boundaries.
23///
24/// Rows are appended in file order. A chunk is closed once it is at or past the
25/// target size *and* the next atom starts a new residue, which is what keeps the
26/// no-straddling guarantee without splitting a large residue.
27///
28/// # Examples
29///
30/// ```
31/// use molframe_core::chunk::{AtomRecord, ChunkBuilder};
32/// use molframe_core::{AltId, Element, OptionalSymbol, Presence, ResidueIndex, SymbolId};
33///
34/// let mut builder = ChunkBuilder::new();
35/// builder.push(AtomRecord {
36///     position: Some([1.0, 2.0, 3.0]),
37///     element: Element::CARBON,
38///     atom_name: SymbolId::from_raw(0),
39///     auth_atom_name: OptionalSymbol::NONE,
40///     alternate_component_id: OptionalSymbol::NONE,
41///     alt_id: AltId::BLANK,
42///     residue: ResidueIndex::new(0),
43///     occupancy: (1.0, Presence::Present),
44///     b_factor: (12.5, Presence::Present),
45///     formal_charge: (0, Presence::Inapplicable),
46///     atom_site_id: 1,
47/// });
48///
49/// let (chunks, coords) = builder.finish();
50/// assert_eq!(chunks.len(), 1);
51/// assert_eq!(coords.len(), 1);
52/// assert!(chunks[0].stats().elements.contains(Element::CARBON));
53/// ```
54#[derive(Debug, Default)]
55pub struct ChunkBuilder {
56    chunks: Vec<AtomChunk>,
57    coords: CoordinateBlock,
58    target: u32,
59    model: u32,
60    pending: Pending,
61}
62
63/// Column buffers for the chunk currently being filled.
64#[derive(Debug, Default)]
65struct Pending {
66    first_atom: u32,
67    element: Vec<u8>,
68    atom_name: Vec<SymbolId>,
69    auth_atom_name: Vec<u32>,
70    has_auth_names: bool,
71    alternate_component_id: Vec<u32>,
72    has_alternate_components: bool,
73    alt_id: Vec<u32>,
74    residue: Vec<u32>,
75    occupancy: Vec<f32>,
76    occupancy_presence: Vec<Presence>,
77    b_factor: Vec<f32>,
78    b_factor_presence: Vec<Presence>,
79    formal_charge: Vec<i8>,
80    formal_charge_presence: Vec<Presence>,
81    atom_site_id: Vec<u32>,
82    coord_presence: Vec<Presence>,
83    stats: AtomChunkStats,
84}
85
86impl Pending {
87    fn len(&self) -> usize {
88        self.element.len()
89    }
90
91    fn is_empty(&self) -> bool {
92        self.element.is_empty()
93    }
94
95    fn last_residue(&self) -> Option<u32> {
96        self.residue.last().copied()
97    }
98
99    fn starts_new_residue(&self, residue: u32) -> bool {
100        self.last_residue() != Some(residue)
101    }
102
103    fn should_close_before(&self, target: u32, residue: u32) -> bool {
104        self.len() as u64 >= u64::from(target) && self.starts_new_residue(residue)
105    }
106
107    fn begin(&mut self, first_atom: u32, model: u32) {
108        if !self.is_empty() {
109            return;
110        }
111
112        self.first_atom = first_atom;
113        self.stats.model = model;
114    }
115
116    fn push_record(&mut self, record: &AtomRecord) {
117        self.element.push(record.element.atomic_number());
118        self.atom_name.push(record.atom_name);
119        self.push_auth_atom_name(record);
120        self.push_alternate_component(record);
121        self.alt_id.push(record.alt_id.get());
122        self.residue.push(record.residue.get());
123
124        push_with_presence(
125            &mut self.occupancy,
126            &mut self.occupancy_presence,
127            record.occupancy,
128        );
129
130        push_with_presence(
131            &mut self.b_factor,
132            &mut self.b_factor_presence,
133            record.b_factor,
134        );
135
136        push_with_presence(
137            &mut self.formal_charge,
138            &mut self.formal_charge_presence,
139            record.formal_charge,
140        );
141
142        self.atom_site_id.push(record.atom_site_id);
143        self.observe(record);
144    }
145
146    fn push_auth_atom_name(&mut self, record: &AtomRecord) {
147        match record.auth_atom_name.get() {
148            Some(name) => {
149                self.auth_atom_name.push(name.get());
150                self.has_auth_names = true;
151            }
152            None => self.auth_atom_name.push(u32::MAX),
153        }
154    }
155
156    fn push_alternate_component(&mut self, record: &AtomRecord) {
157        match record.alternate_component_id.get() {
158            Some(component) => {
159                self.alternate_component_id.push(component.get());
160                self.has_alternate_components = true;
161            }
162            None => self.alternate_component_id.push(u32::MAX),
163        }
164    }
165
166    fn observe(&mut self, record: &AtomRecord) {
167        self.stats.elements.insert(record.element);
168        self.stats.has_hydrogen |= record.element.is_hydrogen();
169        self.stats.has_altloc |= !record.alt_id.is_blank();
170
171        if record.occupancy.1.is_present() {
172            self.stats.occupancy.observe(record.occupancy.0);
173        }
174
175        if record.b_factor.1.is_present() {
176            self.stats.b_factor.observe(record.b_factor.0);
177        }
178
179        self.observe_residue(record.residue.get());
180    }
181
182    fn observe_residue(&mut self, residue: u32) {
183        if self.residue.len() == 1 {
184            self.stats.residue_min = residue;
185            self.stats.residue_max = residue;
186            return;
187        }
188
189        self.stats.residue_min = self.stats.residue_min.min(residue);
190        self.stats.residue_max = self.stats.residue_max.max(residue);
191    }
192
193    fn reset(&mut self) {
194        self.first_atom = 0;
195        self.has_auth_names = false;
196        self.has_alternate_components = false;
197
198        self.element.clear();
199        self.atom_name.clear();
200        self.auth_atom_name.clear();
201        self.alternate_component_id.clear();
202        self.alt_id.clear();
203        self.residue.clear();
204        self.occupancy.clear();
205        self.occupancy_presence.clear();
206        self.b_factor.clear();
207        self.b_factor_presence.clear();
208        self.formal_charge.clear();
209        self.formal_charge_presence.clear();
210        self.atom_site_id.clear();
211        self.coord_presence.clear();
212
213        self.stats = AtomChunkStats::default();
214    }
215}
216
217impl ChunkBuilder {
218    /// Creates a builder using the default chunk size.
219    #[must_use]
220    pub fn new() -> Self {
221        Self {
222            target: TARGET_CHUNK_ATOMS,
223            ..Self::default()
224        }
225    }
226
227    /// Creates a builder that closes chunks at `target` atoms.
228    ///
229    /// Exists for tests and for callers that know their access pattern; the
230    /// default is what a file read uses.
231    #[must_use]
232    pub fn with_target(target: u32) -> Self {
233        Self {
234            target: target.max(1),
235            ..Self::default()
236        }
237    }
238
239    /// Reserves room for `atoms` positions.
240    pub fn reserve(&mut self, atoms: usize) {
241        self.coords = CoordinateBlock::with_capacity(atoms);
242    }
243
244    /// Declares which model the atoms that follow belong to.
245    ///
246    /// Closes the current chunk, because a chunk's model is constant by
247    /// construction and that is what lets the column vanish entirely.
248    pub fn start_model(&mut self, model: u32) {
249        self.close();
250        self.model = model;
251    }
252
253    /// Appends an atom.
254    pub fn push(&mut self, record: AtomRecord) {
255        self.close_before(&record);
256        self.pending.begin(self.coords.len(), self.model);
257        self.push_position(record.position);
258        self.pending.push_record(&record);
259    }
260
261    fn close_before(&mut self, record: &AtomRecord) {
262        if self
263            .pending
264            .should_close_before(self.target, record.residue.get())
265        {
266            self.close();
267        }
268    }
269
270    fn push_position(&mut self, position: Option<[f32; 3]>) {
271        if let Some(position) = position {
272            self.coords.push(position);
273            self.pending.stats.bounds.extend(position);
274            self.pending.coord_presence.push(Presence::Present);
275        } else {
276            // The row still exists — an atom whose position was not recorded
277            // is not an atom that was never modelled — so the buffer keeps a
278            // slot and the validity mask carries the distinction.
279            self.coords.push([f32::NAN; 3]);
280            self.pending.coord_presence.push(Presence::Unknown);
281            self.pending.stats.has_missing_coords = true;
282        }
283    }
284
285    /// Closes the chunk being filled, if it holds anything.
286    fn close(&mut self) {
287        if self.pending.is_empty() {
288            return;
289        }
290
291        let chunk = self.build_chunk();
292        self.chunks.push(chunk);
293        self.pending.reset();
294    }
295
296    fn build_chunk(&mut self) -> AtomChunk {
297        let end = self.coords.len();
298        let pending = &mut self.pending;
299        let first = pending.first_atom;
300        let len = end - first;
301
302        let occupancy_validity = compact_validity(&pending.occupancy_presence, len);
303        let b_factor_validity = compact_validity(&pending.b_factor_presence, len);
304        let formal_charge_validity = compact_validity(&pending.formal_charge_presence, len);
305        let coord_validity = compact_validity(&pending.coord_presence, len);
306
307        let element = EncodedColumn::encode(&pending.element);
308
309        // Names are read by most selections, so they stay literal.
310        let atom_name = EncodedColumn::plain(&pending.atom_name);
311
312        let auth_atom_name = pending
313            .has_auth_names
314            .then(|| EncodedColumn::plain(&pending.auth_atom_name));
315        let alternate_component_id = pending
316            .has_alternate_components
317            .then(|| EncodedColumn::encode(&pending.alternate_component_id));
318
319        let alt_id = EncodedColumn::encode(&pending.alt_id);
320        let residue = ParentMapping::block_indexed(&pending.residue);
321        let occupancy = EncodedColumn::plain(&pending.occupancy);
322        let b_factor = EncodedColumn::plain(&pending.b_factor);
323        let formal_charge = EncodedColumn::encode(&pending.formal_charge);
324        let atom_site_id = EncodedColumn::encode(&pending.atom_site_id);
325        let stats = std::mem::take(&mut pending.stats);
326
327        AtomChunk {
328            atoms: first..end,
329            model: self.model,
330            element,
331            atom_name,
332            auth_atom_name,
333            alternate_component_id,
334            alt_id,
335            residue,
336            occupancy,
337            occupancy_validity,
338            b_factor,
339            b_factor_validity,
340            formal_charge,
341            formal_charge_validity,
342            atom_site_id,
343            coord_validity,
344            stats,
345        }
346    }
347
348    /// Closes the last chunk and returns everything built.
349    #[must_use]
350    pub fn finish(mut self) -> (Vec<AtomChunk>, CoordinateBlock) {
351        self.close();
352        (self.chunks, self.coords)
353    }
354}
355
356fn push_with_presence<T>(values: &mut Vec<T>, presences: &mut Vec<Presence>, entry: (T, Presence)) {
357    let (value, presence) = entry;
358    values.push(value);
359    presences.push(presence);
360}
361
362fn compact_validity(presences: &[Presence], len: u32) -> ValidityMask {
363    ValidityMask::from_bounded_iter(presences.iter().copied(), len)
364}
365
366#[cfg(test)]
367#[path = "builder_tests.rs"]
368mod tests;
molframe_core/chunk/builder.rs

molframe_core/chunk/
builder.rs
