spdkit_surface/
layers.rs

1// [[file:../spdkit-surface.note::135ef16f][135ef16f]]
2use crate::common::*;
3
4use gchemol::Lattice;
5use vecfx::*;
6// 135ef16f ends here
7
8// [[file:../spdkit-surface.note::2378149a][2378149a]]
9fn reorder_atoms_by_zfrac_coords(mol: &Molecule) -> Option<Vec<usize>> {
10    let mut zfrac_coords_numbers: Vec<_> = mol.get_scaled_positions()?.map(|[_, _, fz]| fz).zip(mol.numbers()).collect();
11    zfrac_coords_numbers.sort_by_cached_key(|k| k.0.as_ordered_float());
12    zfrac_coords_numbers.into_iter().map(|(_, i)| i).collect_vec().into()
13}
14
15// Check if mol contains atom `i` by atom label set previously
16fn contains_atom(mol: &Molecule, i: usize) -> bool {
17    mol.atoms().find(|(_, a)| a.label() == format!("{i}")).is_some()
18}
19
20#[track_caller]
21/// Return atoms in bottom layer which contains `bottom_atom`
22fn process_bottom_layer(layers: &mut Vec<Molecule>, bottom_atom: usize) -> Vec<usize> {
23    let i = layers
24        .iter()
25        .position(|mol| contains_atom(mol, bottom_atom))
26        .expect("no layer for bottom atom");
27    let mol = layers.remove(i);
28    mol.atoms().map(|(_, a)| a.label().parse().unwrap()).collect()
29}
30// 2378149a ends here
31
32// [[file:../spdkit-surface.note::4299ffbb][4299ffbb]]
33#[track_caller]
34/// Fragment molecule into connected parts by layer for periodic slab model.
35pub fn fragment_atoms_by_layer(mol: &Molecule) -> Result<impl Iterator<Item = Vec<usize>>> {
36    let mut mol = mol.clone();
37    ensure!(mol.is_periodic(), "only works for periodic structure!");
38
39    // reorder atoms by their zfrac coords
40    let reorder_atoms = reorder_atoms_by_zfrac_coords(&mol).unwrap();
41
42    // label each atom with its serial number so we can find it when
43    // fragmented with this label
44    let numbers = mol.numbers().collect_vec();
45    for i in numbers {
46        let a = mol.get_atom_unchecked_mut(i);
47        a.set_label(format!("{i}"));
48    }
49
50    // scale the lattice along c direction
51    let frac_coords: Vec<_> = mol.get_scaled_positions().unwrap().collect();
52    mol.get_lattice_mut().map(|lat| lat.scale_by_c(2.0));
53    mol.set_scaled_positions(frac_coords);
54    // the recalculate connectivity without considiering lattice
55    std::env::set_var("GCHEMOL_REBOND_IGNORE_PBC", format!("true"));
56    mol.rebond();
57    let mut layers = mol.fragmented().collect_vec();
58
59    let mut remained = reorder_atoms.clone();
60    let iter = std::iter::from_fn(move || {
61        // find bottom layer
62        if remained.is_empty() {
63            None
64        } else {
65            let bottom_atom = remained[0];
66            let bottom_atoms = process_bottom_layer(&mut layers, bottom_atom);
67            // remove atoms in bottom layer
68            remained.retain(|sn| !bottom_atoms.contains(sn));
69            Some(bottom_atoms)
70        }
71    });
72
73    Ok(iter)
74}
75// 4299ffbb ends here
76
77// [[file:../spdkit-surface.note::9d23e311][9d23e311]]
78#[test]
79fn test_atom_layer() -> Result<()> {
80    let f = "tests/files/R.xsd";
81    let m = Molecule::from_file(f)?;
82    let layers = fragment_atoms_by_layer(&m)?.collect_vec();
83
84    assert!(layers.len() > 4);
85    assert_eq!(layers[0].len(), 24);
86    assert_eq!(layers[1].len(), 27);
87    assert_eq!(layers[2].len(), 24);
88    assert_eq!(layers[3].len(), 12);
89
90    Ok(())
91}
92// 9d23e311 ends here
spdkit_surface/layers.rs

spdkit_surface/
layers.rs
