1use std::fs::File;
2use std::io::prelude::*;
3use std::io::BufReader;
4
5fn parse_line(line: String) -> Vec<f64> {
6 let line = line.split_whitespace();
7 line.map(|x| x.parse::<f64>().unwrap())
8 .collect::<Vec<f64>>()
9}
10
11#[allow(dead_code)]
12#[derive(Debug)]
13pub struct Atom {
14 pub an: usize,
15 pub charge: f32,
16 pub pos: Vec<f64>,
17}
18
19#[derive(Debug)]
20pub struct CubeData {
21 pub n_atoms: usize,
22 pub origin: Vec<f64>,
23 pub shape: Vec<usize>,
24 pub data: Vec<f64>,
25 pub cell: Vec<Vec<f64>>,
26 pub atoms: Vec<Atom>,
27}
28
29impl CubeData {
30 fn new() -> Self {
31 CubeData {
32 n_atoms: 0,
33 origin: vec![0f64; 3],
34 shape: vec![0usize; 3],
35 cell: vec![vec![0f64; 3]; 3],
36 atoms: Vec::new(),
37 data: Vec::new(),
38 }
39 }
40}
41
42pub fn read_cube(file_name: &str) -> std::io::Result<CubeData> {
43 let f = File::open(file_name)?;
44 let reader = BufReader::new(f);
45 let mut cube_data = CubeData::new();
46
47 let mut lines = reader.lines();
48 lines.next();
49 lines.next();
50 let line3 = parse_line(lines.next().unwrap()?);
51 cube_data.n_atoms = line3[0] as usize;
52 cube_data.origin = Vec::from(&line3[1..]);
53
54 let line4 = parse_line(lines.next().unwrap()?);
55 let line5 = parse_line(lines.next().unwrap()?);
56 let line6 = parse_line(lines.next().unwrap()?);
57
58 cube_data.shape[0] = line4[0] as usize;
59 cube_data.cell[0] = Vec::from(&line4[1..])
60 .iter()
61 .map(|x| x * cube_data.shape[0] as f64)
62 .collect();
63 cube_data.shape[1] = line5[0] as usize;
64 cube_data.cell[1] = Vec::from(&line5[1..])
65 .iter()
66 .map(|x| x * cube_data.shape[1] as f64)
67 .collect();
68 cube_data.shape[2] = line6[0] as usize;
69 cube_data.cell[2] = Vec::from(&line6[1..])
70 .iter()
71 .map(|x| x * cube_data.shape[2] as f64)
72 .collect();
73
74 let mut idx = 0;
75 while idx < cube_data.n_atoms {
76 let line_atom = parse_line(lines.next().unwrap()?);
77 cube_data.atoms.push(Atom {
78 an: line_atom[0] as usize,
79 charge: line_atom[1] as f32,
80 pos: Vec::from(&line_atom[2..]),
81 });
82 idx += 1;
83 }
84
85 for line in lines {
86 cube_data.data.append(&mut parse_line(line?));
87 }
88
89 Ok(cube_data)
90}
91
92pub fn write_cube(file_name: &str, cube_data: &CubeData) -> std::io::Result<()> {
93 let mut file = File::create(file_name)?;
94 let delimiter = " ";
95 file.write_all(b"Cube file generated by cubeio.rs\n")?;
96 file.write_all(b"OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")?;
97
98 file.write_all(cube_data.n_atoms.to_string().as_bytes())?;
99 file.write_all(delimiter.as_bytes())?;
100 file.write_all(
101 cube_data
102 .origin
103 .iter()
104 .map(|x| {
105 let mut x = x.to_string();
106 x.push_str(delimiter);
107 return x;
108 })
109 .collect::<String>()
110 .as_bytes(),
111 )?;
112 file.write_all(b"\n")?;
113
114 for i in 0..3 {
115 file.write_all(cube_data.shape[i].to_string().as_bytes())?;
116 file.write_all(delimiter.as_bytes())?;
117 file.write_all(
118 cube_data.cell[i]
119 .iter()
120 .map(|x| {
121 let mut x = (x / cube_data.shape[i] as f64).to_string();
122 x.push_str(delimiter);
123 return x;
124 })
125 .collect::<String>()
126 .as_bytes(),
127 )?;
128 file.write_all(b"\n")?;
129 }
130
131 for i in 0..(cube_data.n_atoms as usize) {
132 file.write_all(cube_data.atoms[i].an.to_string().as_bytes())?;
133 file.write_all(delimiter.as_bytes())?;
134 file.write_all(cube_data.atoms[i].charge.to_string().as_bytes())?;
135 file.write_all(delimiter.as_bytes())?;
136 file.write_all(
137 cube_data.atoms[i]
138 .pos
139 .iter()
140 .map(|x| {
141 let mut x = x.to_string();
142 x.push_str(delimiter);
143 return x;
144 })
145 .collect::<String>()
146 .as_bytes(),
147 )?;
148 file.write_all(b"\n")?;
149 }
150 for dat in cube_data.data.iter() {
151 file.write_all(dat.to_string().as_bytes())?;
152 file.write_all(b"\n")?;
153 }
154 Ok(())
155}
156
157#[cfg(test)]
158mod tests {
159 use float_cmp::approx_eq;
160 #[test]
161 fn read_file() {
162 let cube_data = super::read_cube("tdc.cube").unwrap();
163 assert_eq!(cube_data.n_atoms, 30);
164 assert_eq!(cube_data.origin, vec![0f64; 3]);
165 assert_eq!(cube_data.shape, vec![79usize, 52usize, 85usize]);
166
167 assert!(approx_eq!(f64, cube_data.cell[0][0], 14.831539));
168 assert!(approx_eq!(f64, cube_data.cell[0][1], 0.0));
169 assert!(approx_eq!(f64, cube_data.cell[0][2], 0.0));
170
171 assert!(approx_eq!(f64, cube_data.cell[1][0], 0.0));
172 assert!(approx_eq!(f64, cube_data.cell[1][1], 9.67538));
173 assert!(approx_eq!(f64, cube_data.cell[1][2], 0.0));
174
175 assert!(approx_eq!(f64, cube_data.cell[2][0], -1.263185));
176 assert!(approx_eq!(f64, cube_data.cell[2][1], 0.0));
177 assert!(approx_eq!(f64, cube_data.cell[2][2], 15.872644999999999));
178
179 assert_eq!(cube_data.atoms[0].an, 3usize);
180 assert_eq!(cube_data.atoms[cube_data.atoms.len() - 1].an, 16usize);
181 assert!(approx_eq!(f32, cube_data.atoms[0].charge, 0f32));
182 assert_eq!(
183 cube_data.data.len(),
184 (cube_data.shape[0] * cube_data.shape[1] * cube_data.shape[2]) as usize
185 );
186 }
187 #[test]
188 fn write_file() {
189 let cube_data = super::read_cube("tdc.cube").unwrap();
190 super::write_cube("tdc_out.cube", &cube_data).unwrap();
191 let cube_data_read = super::read_cube("tdc_out.cube").unwrap();
192 assert_eq!(cube_data.n_atoms, cube_data_read.n_atoms);
193 assert_eq!(cube_data.origin, cube_data_read.origin);
194 assert_eq!(cube_data.shape, cube_data_read.shape);
195 assert_eq!(cube_data.atoms[0].an, cube_data_read.atoms[0].an);
196 assert_eq!(
197 cube_data.atoms[cube_data.atoms.len() - 1].an,
198 cube_data_read.atoms[cube_data_read.atoms.len() - 1].an
199 );
200 approx_eq!(
201 f32,
202 cube_data.atoms[0].charge,
203 cube_data_read.atoms[0].charge
204 );
205 assert_eq!(
206 cube_data_read.data.len(),
207 (cube_data_read.shape[0] * cube_data_read.shape[1] * cube_data_read.shape[2]) as usize
208 );
209 }
210}