molframe_core/structure/
ensemble.rs

1//! Dense ensembles as one `(models, atoms, 3)` block, and the refusal ragged ones get.
2//!
3//! A dense ensemble keeps one coordinate block per model over a single topology,
4//! so asking for all of them is a gather into one exactly sized buffer. A ragged
5//! ensemble has no shared atom axis: there is nothing to gather into, and the
6//! honest answer is a typed refusal that points at the per-model structures.
7
8use super::Structure;
9use crate::diagnostic::{Code, Diagnostic};
10use crate::index::ModelIndex;
11
12/// The positions of several models of one dense ensemble, model-major.
13#[derive(Clone, Debug, PartialEq)]
14pub struct EnsembleBlock {
15    models: usize,
16    atoms: usize,
17    coordinates: Vec<f32>,
18}
19
20impl EnsembleBlock {
21    /// Models in the block.
22    #[must_use]
23    pub const fn models(&self) -> usize {
24        self.models
25    }
26
27    /// Atoms in every model.
28    #[must_use]
29    pub const fn atoms(&self) -> usize {
30        self.atoms
31    }
32
33    /// `models * atoms * 3` values: model, then atom, then axis.
34    #[must_use]
35    pub fn as_flat(&self) -> &[f32] {
36        &self.coordinates
37    }
38
39    /// Hands the buffer over without copying it.
40    #[must_use]
41    pub fn into_flat(self) -> Vec<f32> {
42        self.coordinates
43    }
44
45    /// One model's positions.
46    #[must_use]
47    pub fn model(&self, index: usize) -> Option<&[[f32; 3]]> {
48        let width = self.atoms.checked_mul(3)?;
49        let start = index.checked_mul(width)?;
50        let flat = self.coordinates.get(start..start.checked_add(width)?)?;
51        Some(flat.as_chunks::<3>().0)
52    }
53}
54
55/// The message every refusal of a ragged ensemble carries.
56const RAGGED_MESSAGE: &str = "ragged ensemble: its models hold different atoms and have no shared atom axis; \
57     iterate `ragged_models` (Python `structure.ragged_models`, Rust `Structure::ragged_models()`) \
58     and analyse each model on its own";
59
60impl Structure {
61    /// The refusal a ragged ensemble gets from every operation that needs one atom axis.
62    pub(crate) fn ragged_refusal(&self) -> Diagnostic {
63        Diagnostic::new(Code::E6003)
64            .with_message(RAGGED_MESSAGE)
65            .with_context("reason", RAGGED_MESSAGE)
66            .with_context("models", self.model_count().to_string())
67    }
68
69    /// Gathers the positions of `models` (every model when `None`) into one buffer.
70    ///
71    /// The buffer is allocated once at its final size, so peak memory is the
72    /// result itself. Models keep the order given.
73    ///
74    /// # Errors
75    ///
76    /// Returns [`Code::E6003`] for a ragged ensemble (naming `ragged_models`) or
77    /// for a model index the structure does not have, and [`Code::E7001`] when
78    /// the buffer size overflows `usize` or cannot be reserved.
79    pub fn ensemble_positions(
80        &self,
81        models: Option<&[ModelIndex]>,
82    ) -> Result<EnsembleBlock, Diagnostic> {
83        self.require_resolved_model()?;
84        let count = self.model_count();
85        let atoms = self.positions().len();
86        let all: Vec<ModelIndex> = if models.is_none() {
87            (0..count)
88                .map(|index| {
89                    ModelIndex::new(match u32::try_from(index) {
90                        Ok(index) => index,
91                        Err(_) => u32::MAX,
92                    })
93                })
94                .collect()
95        } else {
96            Vec::new()
97        };
98        let chosen = match models {
99            Some(models) => models,
100            None => &all,
101        };
102        let too_big = || {
103            Diagnostic::new(Code::E7001)
104                .with_message("ensemble coordinates exceed addressable memory")
105                .with_context("models", chosen.len().to_string())
106                .with_context("atoms", atoms.to_string())
107        };
108        let total = chosen
109            .len()
110            .checked_mul(atoms)
111            .and_then(|rows| rows.checked_mul(3))
112            .ok_or_else(too_big)?;
113        let mut coordinates = Vec::new();
114        coordinates
115            .try_reserve_exact(total)
116            .map_err(|_| too_big())?;
117        for model in chosen {
118            let Some(positions) = self.model_positions(*model) else {
119                return Err(Diagnostic::new(Code::E6003)
120                    .with_message(format!(
121                        "model {} does not exist: the structure has {count} models, addressed by position from 0",
122                        model.get()
123                    ))
124                    .with_context("model", model.get().to_string())
125                    .with_context("models", count.to_string()));
126            };
127            coordinates.extend(positions.iter().flatten());
128        }
129        Ok(EnsembleBlock {
130            models: chosen.len(),
131            atoms,
132            coordinates,
133        })
134    }
135}
molframe_core/structure/ensemble.rs

molframe_core/structure/
ensemble.rs
