1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

use nalgebra::Point3;

use super::Atom;

/// Struct of Array style, memory allocation is continuous when modifying the same attribute for all atoms.
#[derive(Debug, PartialEq, PartialOrd, Clone)]
pub struct AtomCollections {
    symbols: Vec<String>,
    atomic_numbers: Vec<u8>,
    cartesian_coords: Vec<Point3<f64>>,
    indexes: Vec<usize>,
    size: usize,
}

impl AtomCollections {
    pub fn new(size: usize) -> Self {
        AtomCollections {
            symbols: Vec::with_capacity(size),
            atomic_numbers: Vec::with_capacity(size),
            cartesian_coords: Vec::with_capacity(size),
            indexes: Vec::with_capacity(size),
            size,
        }
    }
    pub fn symbols(&self) -> &[String] {
        self.symbols.as_ref()
    }

    pub fn atomic_numbers(&self) -> &[u8] {
        self.atomic_numbers.as_ref()
    }

    pub fn cartesian_coords(&self) -> &[Point3<f64>] {
        self.cartesian_coords.as_ref()
    }

    pub fn indexes(&self) -> &[usize] {
        self.indexes.as_ref()
    }
    /// Retrieve an `Atom` at given 0th-based index.
    pub fn get_atom_at(&self, index: usize) -> Option<Atom> {
        if index >= self.size {
            None
        } else {
            let symbol = self.symbols().get(index).unwrap();
            let atomic_num = self.atomic_numbers().get(index).unwrap();
            let cartesian_coord = self.cartesian_coords().get(index).unwrap();
            Some(
                Atom::new_builder()
                    .with_symbol(symbol)
                    .with_atomic_number(*atomic_num)
                    .with_coord(cartesian_coord)
                    .with_index(index)
                    .ready()
                    .build(),
            )
        }
    }

    pub fn size(&self) -> usize {
        self.size
    }
}

/// Conversion from `&[Atom]` to `AtomCollections`
impl From<&[Atom]> for AtomCollections {
    fn from(value: &[Atom]) -> Self {
        value.to_vec().into()
    }
}

/// Conversion from `Vec<Atom>` to `AtomCollections`
impl From<Vec<Atom>> for AtomCollections {
    fn from(value: Vec<Atom>) -> Self {
        let collection_size = value.len();
        let mut collections = AtomCollections::new(collection_size);
        value.iter().for_each(|atom| {
            collections.symbols.push(atom.symbol().into());
            collections.atomic_numbers.push(atom.atomic_number());
            collections.cartesian_coords.push(atom.cartesian_coord());
            collections.indexes.push(atom.index);
        });
        collections
    }
}

/// Conversion from `AtomCollections` to `Vec<Atom>`
impl From<AtomCollections> for Vec<Atom> {
    fn from(value: AtomCollections) -> Self {
        let collection_size = value.indexes().len();
        let mut atom_vec: Vec<Atom> = Vec::with_capacity(collection_size);
        for i in 0..collection_size {
            atom_vec.push(value.get_atom_at(i).unwrap())
        }
        atom_vec
    }
}