ferritin_bevy/
selection.rs

1//! AtomMask selection layer: maps MolViewSpec ComponentSelector → per-atom boolean mask.
2use ferritin_core::AtomCollection;
3use ferritin_molviewspec::molviewspec::nodes::{
4    ComponentExpression, ComponentSelector, ComponentSelectorT,
5};
6
7/// Per-atom boolean mask. Length always equals the atom count of the target model.
8pub struct AtomMask(pub Vec<bool>);
9
10impl AtomMask {
11    pub fn count_true(&self) -> usize {
12        self.0.iter().filter(|&&b| b).count()
13    }
14
15    pub fn all_false(&self) -> bool {
16        self.0.iter().all(|&b| !b)
17    }
18}
19
20/// Evaluate a MolViewSpec `ComponentSelector` against an `AtomCollection`, returning a per-atom mask.
21pub fn evaluate_selector(selector: &ComponentSelector, model: &AtomCollection) -> AtomMask {
22    match selector {
23        ComponentSelector::Selector(s) => evaluate_named(s, model),
24        ComponentSelector::Expression(expr) => evaluate_expression(expr, model),
25        ComponentSelector::ExpressionList(exprs) => {
26            let n = model.get_size();
27            let mut mask = vec![false; n];
28            for expr in exprs {
29                let sub = evaluate_expression(expr, model);
30                for (dst, src) in mask.iter_mut().zip(sub.0.iter()) {
31                    *dst |= src;
32                }
33            }
34            AtomMask(mask)
35        }
36    }
37}
38
39fn evaluate_named(selector: &ComponentSelectorT, model: &AtomCollection) -> AtomMask {
40    let n = model.get_size();
41    match selector {
42        ComponentSelectorT::All => AtomMask(vec![true; n]),
43        ComponentSelectorT::Polymer => {
44            let mut mask = vec![false; n];
45            for residue in model.iter_residues() {
46                let is_polymer = residue.is_amino_acid() || residue.is_nucleotide();
47                for idx in residue.atom_range() {
48                    mask[idx] = is_polymer;
49                }
50            }
51            AtomMask(mask)
52        }
53        ComponentSelectorT::Protein => {
54            let mut mask = vec![false; n];
55            for residue in model.iter_residues() {
56                if residue.is_amino_acid() {
57                    for idx in residue.atom_range() {
58                        mask[idx] = true;
59                    }
60                }
61            }
62            AtomMask(mask)
63        }
64        ComponentSelectorT::Nucleic => {
65            let mut mask = vec![false; n];
66            for residue in model.iter_residues() {
67                if residue.is_nucleotide() {
68                    for idx in residue.atom_range() {
69                        mask[idx] = true;
70                    }
71                }
72            }
73            AtomMask(mask)
74        }
75        ComponentSelectorT::Branched => {
76            let mut mask = vec![false; n];
77            for residue in model.iter_residues() {
78                if residue.is_carbohydrate() {
79                    for idx in residue.atom_range() {
80                        mask[idx] = true;
81                    }
82                }
83            }
84            AtomMask(mask)
85        }
86        ComponentSelectorT::Ligand => {
87            let mut mask = vec![false; n];
88            for residue in model.iter_residues() {
89                if residue.is_hetero() && !residue.is_water() && !residue.is_ion() {
90                    for idx in residue.atom_range() {
91                        mask[idx] = true;
92                    }
93                }
94            }
95            AtomMask(mask)
96        }
97        ComponentSelectorT::Ion => AtomMask(model.select_ion()),
98        ComponentSelectorT::Water => AtomMask(model.select_water()),
99    }
100}
101
102fn evaluate_expression(expr: &ComponentExpression, model: &AtomCollection) -> AtomMask {
103    // An all-None expression matches everything (MVS spec: "match whole structure").
104    if is_all_none(expr) {
105        return AtomMask(vec![true; model.get_size()]);
106    }
107
108    let n = model.get_size();
109    let mut mask = vec![false; n];
110
111    for (res_iter_idx, residue) in model.iter_residues().enumerate() {
112        // Chain filter (label_asym_id matched against auth_asym_id in MVP)
113        if let Some(ref chain) = expr.label_asym_id
114            && residue.chain_id() != chain
115        {
116            continue;
117        }
118        if let Some(ref chain) = expr.auth_asym_id
119            && residue.chain_id() != chain
120        {
121            continue;
122        }
123
124        // residue_index: 0-based residue iteration index
125        if let Some(ri) = expr.residue_index
126            && res_iter_idx as i32 != ri
127        {
128            continue;
129        }
130
131        let res_id = residue.residue_id();
132
133        // Exact residue seq id (label_seq_id matched against auth_seq_id in MVP)
134        if let Some(seq) = expr.label_seq_id
135            && res_id != seq
136        {
137            continue;
138        }
139        if let Some(seq) = expr.auth_seq_id
140            && res_id != seq
141        {
142            continue;
143        }
144
145        // Range filters (beg inclusive, end inclusive)
146        if let Some(beg) = expr.beg_label_seq_id
147            && res_id < beg
148        {
149            continue;
150        }
151        if let Some(end) = expr.end_label_seq_id
152            && res_id > end
153        {
154            continue;
155        }
156        if let Some(beg) = expr.beg_auth_seq_id
157            && res_id < beg
158        {
159            continue;
160        }
161        if let Some(end) = expr.end_auth_seq_id
162            && res_id > end
163        {
164            continue;
165        }
166
167        // residue_index: 0-based index of this residue in iteration order
168        // We skip this check at residue level; handle per-atom below since it's unusual.
169
170        // Residue passed — now check per-atom fields
171        for idx in residue.atom_range() {
172            // atom_index: 0-based global atom index
173            if let Some(ai) = expr.atom_index
174                && idx as i32 != ai
175            {
176                continue;
177            }
178
179            // label_atom_id / auth_atom_id
180            let atom_name = model.get_atom_name(idx);
181            if let Some(ref la) = expr.label_atom_id
182                && atom_name != la
183            {
184                continue;
185            }
186            if let Some(ref aa) = expr.auth_atom_id
187                && atom_name != aa
188            {
189                continue;
190            }
191
192            // type_symbol (element symbol, uppercase)
193            if let Some(ref sym) = expr.type_symbol
194                && model.get_element(idx).symbol() != sym.as_str()
195            {
196                continue;
197            }
198
199            mask[idx] = true;
200        }
201    }
202
203    AtomMask(mask)
204}
205
206fn is_all_none(expr: &ComponentExpression) -> bool {
207    expr.label_entity_id.is_none()
208        && expr.label_asym_id.is_none()
209        && expr.auth_asym_id.is_none()
210        && expr.label_seq_id.is_none()
211        && expr.auth_seq_id.is_none()
212        && expr.pdbx_pdb_ins_code.is_none()
213        && expr.beg_label_seq_id.is_none()
214        && expr.end_label_seq_id.is_none()
215        && expr.beg_auth_seq_id.is_none()
216        && expr.end_auth_seq_id.is_none()
217        && expr.residue_index.is_none()
218        && expr.label_atom_id.is_none()
219        && expr.auth_atom_id.is_none()
220        && expr.type_symbol.is_none()
221        && expr.atom_id.is_none()
222        && expr.atom_index.is_none()
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use ferritin_core::AtomCollection;
229    use ferritin_core::info::elements::Element;
230    use ferritin_molviewspec::molviewspec::nodes::ComponentExpression;
231
232    /// 8-atom collection: ALA(3, chain A, res 1), GLY(2, chain A, res 2),
233    /// HOH(1, chain B, res 3 hetero), ZN(1, chain B, res 4 hetero), LIG(1, chain B, res 5 hetero multi-atom)
234    fn make_test_collection() -> AtomCollection {
235        let coords = vec![[0.0; 3]; 9];
236        let res_ids: Vec<i32> = vec![1, 1, 1, 2, 2, 3, 4, 5, 5];
237        let res_names: Vec<String> =
238            vec!["ALA", "ALA", "ALA", "GLY", "GLY", "HOH", "ZN", "LIG", "LIG"]
239                .into_iter()
240                .map(|s| s.to_string())
241                .collect();
242        let is_hetero = vec![false, false, false, false, false, true, true, true, true];
243        let elements = vec![
244            Element::N,
245            Element::C,
246            Element::C,
247            Element::N,
248            Element::C,
249            Element::O,
250            Element::Zn,
251            Element::C,
252            Element::O,
253        ];
254        let atom_names: Vec<String> = vec!["N", "CA", "C", "N", "CA", "O", "ZN", "C1", "O1"]
255            .into_iter()
256            .map(|s| s.to_string())
257            .collect();
258        let chain_ids: Vec<String> = vec!["A", "A", "A", "A", "A", "B", "B", "B", "B"]
259            .into_iter()
260            .map(|s| s.to_string())
261            .collect();
262        AtomCollection::new(
263            9, coords, res_ids, res_names, is_hetero, elements, atom_names, chain_ids, None,
264        )
265    }
266
267    // T2-01: All selector → all true
268    #[test]
269    fn test_selector_all() {
270        let ac = make_test_collection();
271        let sel = ComponentSelector::Selector(ComponentSelectorT::All);
272        let mask = evaluate_selector(&sel, &ac);
273        assert_eq!(mask.count_true(), 9);
274    }
275
276    // T2-02: Protein selector → only ALA + GLY atoms (5 atoms)
277    #[test]
278    fn test_selector_protein() {
279        let ac = make_test_collection();
280        let sel = ComponentSelector::Selector(ComponentSelectorT::Protein);
281        let mask = evaluate_selector(&sel, &ac);
282        assert_eq!(
283            mask.count_true(),
284            5,
285            "protein: expected 5 atoms (ALA×3 + GLY×2)"
286        );
287    }
288
289    // T2-03: Water selector → only HOH atom (1 atom)
290    #[test]
291    fn test_selector_water() {
292        let ac = make_test_collection();
293        let sel = ComponentSelector::Selector(ComponentSelectorT::Water);
294        let mask = evaluate_selector(&sel, &ac);
295        assert_eq!(mask.count_true(), 1, "water: expected 1 HOH atom");
296        assert!(mask.0[5], "atom 5 (HOH) must be true");
297    }
298
299    // T2-04: Ion selector → only ZN atom (1 atom)
300    #[test]
301    fn test_selector_ion() {
302        let ac = make_test_collection();
303        let sel = ComponentSelector::Selector(ComponentSelectorT::Ion);
304        let mask = evaluate_selector(&sel, &ac);
305        assert_eq!(mask.count_true(), 1, "ion: expected 1 ZN atom");
306        assert!(mask.0[6], "atom 6 (ZN) must be true");
307    }
308
309    // T2-05: Ligand selector → LIG atoms (2 atoms), not water or ion
310    #[test]
311    fn test_selector_ligand() {
312        let ac = make_test_collection();
313        let sel = ComponentSelector::Selector(ComponentSelectorT::Ligand);
314        let mask = evaluate_selector(&sel, &ac);
315        assert_eq!(mask.count_true(), 2, "ligand: expected 2 LIG atoms");
316        assert!(mask.0[7] && mask.0[8], "atoms 7,8 (LIG) must be true");
317    }
318
319    // T2-06: Polymer selector → protein + nucleic (ALA + GLY = 5)
320    #[test]
321    fn test_selector_polymer() {
322        let ac = make_test_collection();
323        let sel = ComponentSelector::Selector(ComponentSelectorT::Polymer);
324        let mask = evaluate_selector(&sel, &ac);
325        assert_eq!(mask.count_true(), 5);
326    }
327
328    // T2-07: Expression chain A → first 5 atoms
329    #[test]
330    fn test_expression_chain_a() {
331        let ac = make_test_collection();
332        let expr = ComponentExpression {
333            auth_asym_id: Some("A".to_string()),
334            ..Default::default()
335        };
336        let sel = ComponentSelector::Expression(expr);
337        let mask = evaluate_selector(&sel, &ac);
338        assert_eq!(mask.count_true(), 5);
339        assert!(mask.0[..5].iter().all(|&b| b));
340        assert!(mask.0[5..].iter().all(|&b| !b));
341    }
342
343    // T2-08: Expression chain B → last 4 atoms
344    #[test]
345    fn test_expression_chain_b() {
346        let ac = make_test_collection();
347        let expr = ComponentExpression {
348            auth_asym_id: Some("B".to_string()),
349            ..Default::default()
350        };
351        let sel = ComponentSelector::Expression(expr);
352        let mask = evaluate_selector(&sel, &ac);
353        assert_eq!(mask.count_true(), 4);
354        assert!(mask.0[5..].iter().all(|&b| b));
355    }
356
357    // T2-09: Expression auth_seq_id = 1 → 3 atoms (ALA residue)
358    #[test]
359    fn test_expression_seq_id() {
360        let ac = make_test_collection();
361        let expr = ComponentExpression {
362            auth_seq_id: Some(1),
363            ..Default::default()
364        };
365        let sel = ComponentSelector::Expression(expr);
366        let mask = evaluate_selector(&sel, &ac);
367        assert_eq!(mask.count_true(), 3);
368    }
369
370    // T2-10: Expression beg/end range 1..2 → 5 atoms (ALA+GLY)
371    #[test]
372    fn test_expression_seq_range() {
373        let ac = make_test_collection();
374        let expr = ComponentExpression {
375            beg_auth_seq_id: Some(1),
376            end_auth_seq_id: Some(2),
377            ..Default::default()
378        };
379        let sel = ComponentSelector::Expression(expr);
380        let mask = evaluate_selector(&sel, &ac);
381        assert_eq!(mask.count_true(), 5);
382    }
383
384    // T2-11: Expression auth_atom_id = "CA" → 2 atoms
385    #[test]
386    fn test_expression_atom_name() {
387        let ac = make_test_collection();
388        let expr = ComponentExpression {
389            auth_atom_id: Some("CA".to_string()),
390            ..Default::default()
391        };
392        let sel = ComponentSelector::Expression(expr);
393        let mask = evaluate_selector(&sel, &ac);
394        assert_eq!(mask.count_true(), 2, "CA atoms in ALA and GLY");
395    }
396
397    // T2-12: Expression type_symbol = "N" → 2 atoms
398    #[test]
399    fn test_expression_type_symbol() {
400        let ac = make_test_collection();
401        let expr = ComponentExpression {
402            type_symbol: Some("N".to_string()),
403            ..Default::default()
404        };
405        let sel = ComponentSelector::Expression(expr);
406        let mask = evaluate_selector(&sel, &ac);
407        assert_eq!(mask.count_true(), 2, "N atoms in ALA and GLY");
408    }
409
410    // T2-13: Expression atom_index = 0 → 1 atom
411    #[test]
412    fn test_expression_atom_index() {
413        let ac = make_test_collection();
414        let expr = ComponentExpression {
415            atom_index: Some(0),
416            ..Default::default()
417        };
418        let sel = ComponentSelector::Expression(expr);
419        let mask = evaluate_selector(&sel, &ac);
420        assert_eq!(mask.count_true(), 1);
421        assert!(mask.0[0]);
422    }
423
424    // T2-14: ExpressionList (chain A res 1) OR (chain B res 4) → 3 + 1 = 4 atoms
425    #[test]
426    fn test_expression_list_union() {
427        let ac = make_test_collection();
428        let expr1 = ComponentExpression {
429            auth_asym_id: Some("A".to_string()),
430            auth_seq_id: Some(1),
431            ..Default::default()
432        };
433        let expr2 = ComponentExpression {
434            auth_asym_id: Some("B".to_string()),
435            auth_seq_id: Some(4),
436            ..Default::default()
437        };
438        let sel = ComponentSelector::ExpressionList(vec![expr1, expr2]);
439        let mask = evaluate_selector(&sel, &ac);
440        assert_eq!(mask.count_true(), 4, "union: ALA(3) + ZN(1)");
441    }
442
443    // T2-21: All-None expression → all-true mask
444    #[test]
445    fn test_expression_all_none_is_all_true() {
446        let ac = make_test_collection();
447        let expr = ComponentExpression::default();
448        let sel = ComponentSelector::Expression(expr);
449        let mask = evaluate_selector(&sel, &ac);
450        assert_eq!(
451            mask.count_true(),
452            9,
453            "all-None expression must match all atoms"
454        );
455    }
456
457    // T2-22: count_true and all_false helpers
458    #[test]
459    fn test_atom_mask_helpers() {
460        let full = AtomMask(vec![true, true, false]);
461        assert_eq!(full.count_true(), 2);
462        assert!(!full.all_false());
463
464        let empty = AtomMask(vec![false, false, false]);
465        assert_eq!(empty.count_true(), 0);
466        assert!(empty.all_false());
467    }
468
469    // T2-23: label_asym_id matched against auth_asym_id (MVP documented behaviour)
470    #[test]
471    fn test_label_asym_id_matched_against_auth() {
472        let ac = make_test_collection();
473        let expr = ComponentExpression {
474            label_asym_id: Some("A".to_string()),
475            ..Default::default()
476        };
477        let sel = ComponentSelector::Expression(expr);
478        let mask = evaluate_selector(&sel, &ac);
479        assert_eq!(
480            mask.count_true(),
481            5,
482            "label_asym_id A must match chain A (auth)"
483        );
484    }
485
486    // T2-24: label_seq_id matched against auth_seq_id (MVP documented behaviour)
487    #[test]
488    fn test_label_seq_id_matched_against_auth() {
489        let ac = make_test_collection();
490        let expr = ComponentExpression {
491            label_seq_id: Some(2),
492            ..Default::default()
493        };
494        let sel = ComponentSelector::Expression(expr);
495        let mask = evaluate_selector(&sel, &ac);
496        assert_eq!(
497            mask.count_true(),
498            2,
499            "label_seq_id 2 must match res_id 2 (GLY×2)"
500        );
501    }
502
503    // T2-25: residue_index (0-based) selects exactly the Nth residue in iteration order
504    #[test]
505    fn test_expression_residue_index() {
506        let ac = make_test_collection();
507        // residue_index=0 → ALA (3 atoms)
508        let expr = ComponentExpression {
509            residue_index: Some(0),
510            ..Default::default()
511        };
512        let sel = ComponentSelector::Expression(expr);
513        let mask = evaluate_selector(&sel, &ac);
514        assert_eq!(
515            mask.count_true(),
516            3,
517            "residue_index 0 must select ALA (3 atoms)"
518        );
519        assert!(
520            mask.0[0] && mask.0[1] && mask.0[2],
521            "atoms 0-2 should be selected"
522        );
523
524        // residue_index=1 → GLY (2 atoms)
525        let expr2 = ComponentExpression {
526            residue_index: Some(1),
527            ..Default::default()
528        };
529        let mask2 = evaluate_selector(&ComponentSelector::Expression(expr2), &ac);
530        assert_eq!(
531            mask2.count_true(),
532            2,
533            "residue_index 1 must select GLY (2 atoms)"
534        );
535
536        // residue_index=3 → ZN (1 atom); residues: 0=ALA, 1=GLY, 2=HOH, 3=ZN, 4=LIG
537        let expr3 = ComponentExpression {
538            residue_index: Some(3),
539            ..Default::default()
540        };
541        let mask3 = evaluate_selector(&ComponentSelector::Expression(expr3), &ac);
542        assert_eq!(
543            mask3.count_true(),
544            1,
545            "residue_index 3 must select ZN (1 atom)"
546        );
547    }
548
549    // T2-26: residue_index combined with atom_name → single atom
550    #[test]
551    fn test_expression_residue_index_with_atom_name() {
552        let ac = make_test_collection();
553        // Select only the CA of the second residue (GLY)
554        let expr = ComponentExpression {
555            residue_index: Some(1),
556            auth_atom_id: Some("CA".to_string()),
557            ..Default::default()
558        };
559        let mask = evaluate_selector(&ComponentSelector::Expression(expr), &ac);
560        assert_eq!(
561            mask.count_true(),
562            1,
563            "residue 1 CA should be exactly 1 atom"
564        );
565        assert!(mask.0[4], "atom index 4 is GLY CA");
566    }
567}
ferritin_bevy/selection.rs

ferritin_bevy/
selection.rs
