ak_core/io/xyz/
xyz_main.rs

1use crate::{
2    Structure,
3    io::xyz::{
4        errors::{CellError, PBCError, XYZReaderError},
5        parse_atoms::parse_atoms,
6        parse_cell::parse_cell,
7        parse_pbc::parse_pbc,
8        xyz_main::XYZReaderState::{FindNumberOfAtoms, FindPositionLines, FindProperties},
9    },
10};
11
12#[derive(Debug)]
13enum XYZReaderState {
14    FindNumberOfAtoms,
15    FindProperties,
16    FindPositionLines(usize),
17}
18
19struct StructureBuffer {
20    n_atoms: usize,
21    properties: String,
22    positions: Vec<String>,
23}
24
25impl StructureBuffer {
26    fn new() -> StructureBuffer {
27        StructureBuffer {
28            n_atoms: 0,
29            properties: String::new(),
30            positions: Vec::new(),
31        }
32    }
33
34    fn set_n_atoms(&mut self, n_atoms: usize) {
35        self.n_atoms = n_atoms
36    }
37
38    fn get_n_atoms(&self) -> usize {
39        self.n_atoms
40    }
41
42    fn set_properties(&mut self, properties: String) {
43        self.properties = properties
44    }
45
46    fn push_position(&mut self, position: String) {
47        self.positions.push(position)
48    }
49
50    fn positions(&self) -> &[String] {
51        &self.positions
52    }
53
54    fn properties(&self) -> &str {
55        &self.properties
56    }
57}
58
59pub fn read_xyz<R: std::io::BufRead>(r: R) -> Result<Vec<Structure>, XYZReaderError> {
60    let mut structures: Vec<Structure> = Vec::new();
61    let mut state = XYZReaderState::FindNumberOfAtoms;
62    let mut structure_buffer = StructureBuffer::new();
63
64    for line_result in r.lines() {
65        let line = line_result.map_err(|_| XYZReaderError::LineIO)?;
66
67        match state {
68            FindNumberOfAtoms => {
69                let n_atoms: usize = line.trim().parse().map_err(|_| XYZReaderError::IntParse)?;
70                structure_buffer.set_n_atoms(n_atoms);
71                state = FindProperties;
72            }
73            FindProperties => {
74                structure_buffer.set_properties(line);
75                state = FindPositionLines(structure_buffer.n_atoms);
76            }
77            FindPositionLines(lines_left) => {
78                state = if lines_left - 1 == 0 || lines_left == 0 {
79                    structure_buffer.push_position(line);
80                    let structure = buffer_to_structure(structure_buffer)?;
81                    structures.push(structure);
82                    structure_buffer = StructureBuffer::new();
83                    FindNumberOfAtoms
84                } else {
85                    structure_buffer.push_position(line);
86                    FindPositionLines(lines_left - 1)
87                };
88            }
89        }
90    }
91
92    match state {
93        FindNumberOfAtoms => Ok(structures),
94        _ => Err(XYZReaderError::IncompleteInput),
95    }
96}
97
98fn buffer_to_structure(buffer: StructureBuffer) -> Result<Structure, XYZReaderError> {
99    if !(buffer.get_n_atoms() == buffer.positions().len()) {
100        return Err(XYZReaderError::IncorrectNumberOfAtoms);
101    }
102
103    let (positions, numbers) = parse_atoms(buffer.positions(), buffer.get_n_atoms())?;
104    let cell = parse_cell(buffer.properties());
105
106    let cell = match cell {
107        Ok(cell_result) => cell_result,
108        Err(cell_error) => match cell_error {
109            CellError::NoCellSpecified => [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
110            _ => return Err(XYZReaderError::CellError(cell_error)),
111        },
112    };
113
114    let pbc_result = parse_pbc(buffer.properties());
115    let pbc = match pbc_result {
116        Ok(pbc) => pbc,
117        Err(pbc_error) => match pbc_error {
118            PBCError::NoPBCSpecified => [false, false, false],
119            _ => return Err(XYZReaderError::PBCError(pbc_error)),
120        },
121    };
122
123    Ok(Structure::new(positions, numbers, cell, pbc))
124}
125
126pub fn read_xyz_single<R: std::io::BufRead>(r: R) -> Result<Structure, XYZReaderError> {
127    let structures = read_xyz(r)?;
128    Ok(structures.get(0).ok_or(XYZReaderError::EmptyFile)?.clone())
129}
130
131#[cfg(test)]
132mod tests {
133    use std::io::{BufReader, Cursor};
134
135    use super::read_xyz;
136
137    #[test]
138    fn simple_test() {
139        let xyz = r#"4
140Lattice="10.0 0.0 0.0 0.0 10.0 0.0 0.0 0.0 10.0" Properties=species:S:1:pos:R:3 pbc="F F F"
141P        5.00000000       4.69983800       5.37385700
142H        5.00000000       5.90048500       4.62614300
143H        6.03979100       4.09951500       4.62614300
144H        3.96020900       4.09951500       4.62614300
1452
146Lattice="10.0 0.0 0.0 0.0 10.0 0.0 0.0 0.0 10.0" Properties=species:S:1:pos:R:3 pbc="F F F"
147P        5.00000000       5.00000000       5.96614400
148P        5.00000000       5.00000000       4.03385600
1497
150Lattice="10.0 0.0 0.0 0.0 10.0 0.0 0.0 0.0 10.0" Properties=species:S:1:pos:R:3 pbc="F F F"
151O        6.40718800       5.44571050       5.00000000
152C        5.18913300       5.54860350       5.00000000
153H        4.71189200       6.54976550       5.00000000
154C        4.24103100       4.38433250       5.00000000
155H        4.80318700       3.45023450       5.00000000
156H        3.59281200       4.43199550       5.88094600
157H        3.59281200       4.43199550       4.11905400
158"#;
159        let mut reader = BufReader::new(Cursor::new(xyz));
160        let structures = read_xyz(&mut reader).unwrap();
161
162        assert_eq!(structures.len(), 3);
163
164        let n_atoms = structures
165            .iter()
166            .map(|x| x.numbers.len())
167            .collect::<Vec<usize>>();
168
169        assert_eq!(n_atoms, vec![4, 2, 7])
170    }
171}
ak_core/io/xyz/xyz_main.rs

ak_core/io/xyz/
xyz_main.rs
