use std::fs::File;
use std::io::prelude::*;
use std::io::BufReader;
fn parse_line(line: String) -> Vec<f64> {
let line = line.split_whitespace();
line.map(|x| x.parse::<f64>().unwrap())
.collect::<Vec<f64>>()
}
#[allow(dead_code)]
#[derive(Debug)]
pub struct Atom {
pub an: usize,
pub charge: f32,
pub pos: Vec<f64>,
}
#[derive(Debug)]
pub struct CubeData {
pub n_atoms: usize,
pub origin: Vec<f64>,
pub shape: Vec<usize>,
pub data: Vec<f64>,
pub cell: Vec<Vec<f64>>,
pub atoms: Vec<Atom>,
}
impl CubeData {
fn new() -> Self {
CubeData {
n_atoms: 0,
origin: vec![0f64; 3],
shape: vec![0usize; 3],
cell: vec![vec![0f64; 3]; 3],
atoms: Vec::new(),
data: Vec::new(),
}
}
}
pub fn read_cube(file_name: &str) -> std::io::Result<CubeData> {
let f = File::open(file_name)?;
let reader = BufReader::new(f);
let mut cube_data = CubeData::new();
let mut lines = reader.lines();
lines.next();
lines.next();
let line3 = parse_line(lines.next().unwrap()?);
cube_data.n_atoms = line3[0] as usize;
cube_data.origin = Vec::from(&line3[1..]);
let line4 = parse_line(lines.next().unwrap()?);
let line5 = parse_line(lines.next().unwrap()?);
let line6 = parse_line(lines.next().unwrap()?);
cube_data.shape[0] = line4[0] as usize;
cube_data.cell[0] = Vec::from(&line4[1..])
.iter()
.map(|x| x * cube_data.shape[0] as f64)
.collect();
cube_data.shape[1] = line5[0] as usize;
cube_data.cell[1] = Vec::from(&line5[1..])
.iter()
.map(|x| x * cube_data.shape[1] as f64)
.collect();
cube_data.shape[2] = line6[0] as usize;
cube_data.cell[2] = Vec::from(&line6[1..])
.iter()
.map(|x| x * cube_data.shape[2] as f64)
.collect();
let mut idx = 0;
while idx < cube_data.n_atoms {
let line_atom = parse_line(lines.next().unwrap()?);
cube_data.atoms.push(Atom {
an: line_atom[0] as usize,
charge: line_atom[1] as f32,
pos: Vec::from(&line_atom[2..]),
});
idx += 1;
}
for line in lines {
cube_data.data.append(&mut parse_line(line?));
}
Ok(cube_data)
}
pub fn write_cube(file_name: &str, cube_data: &CubeData) -> std::io::Result<()> {
let mut file = File::create(file_name)?;
let delimiter = " ";
file.write_all(b"Cube file generated by cubeio.rs\n")?;
file.write_all(b"OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")?;
file.write_all(cube_data.n_atoms.to_string().as_bytes())?;
file.write_all(delimiter.as_bytes())?;
file.write_all(
cube_data
.origin
.iter()
.map(|x| {
let mut x = x.to_string();
x.push_str(delimiter);
return x;
})
.collect::<String>()
.as_bytes(),
)?;
file.write_all(b"\n")?;
for i in 0..3 {
file.write_all(cube_data.shape[i].to_string().as_bytes())?;
file.write_all(delimiter.as_bytes())?;
file.write_all(
cube_data.cell[i]
.iter()
.map(|x| {
let mut x = (x / cube_data.shape[i] as f64).to_string();
x.push_str(delimiter);
return x;
})
.collect::<String>()
.as_bytes(),
)?;
file.write_all(b"\n")?;
}
for i in 0..(cube_data.n_atoms as usize) {
file.write_all(cube_data.atoms[i].an.to_string().as_bytes())?;
file.write_all(delimiter.as_bytes())?;
file.write_all(cube_data.atoms[i].charge.to_string().as_bytes())?;
file.write_all(delimiter.as_bytes())?;
file.write_all(
cube_data.atoms[i]
.pos
.iter()
.map(|x| {
let mut x = x.to_string();
x.push_str(delimiter);
return x;
})
.collect::<String>()
.as_bytes(),
)?;
file.write_all(b"\n")?;
}
for dat in cube_data.data.iter() {
file.write_all(dat.to_string().as_bytes())?;
file.write_all(b"\n")?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use float_cmp::approx_eq;
#[test]
fn read_file() {
let cube_data = super::read_cube("tdc.cube").unwrap();
assert_eq!(cube_data.n_atoms, 30);
assert_eq!(cube_data.origin, vec![0f64; 3]);
assert_eq!(cube_data.shape, vec![79usize, 52usize, 85usize]);
assert!(approx_eq!(f64, cube_data.cell[0][0], 14.831539));
assert!(approx_eq!(f64, cube_data.cell[0][1], 0.0));
assert!(approx_eq!(f64, cube_data.cell[0][2], 0.0));
assert!(approx_eq!(f64, cube_data.cell[1][0], 0.0));
assert!(approx_eq!(f64, cube_data.cell[1][1], 9.67538));
assert!(approx_eq!(f64, cube_data.cell[1][2], 0.0));
assert!(approx_eq!(f64, cube_data.cell[2][0], -1.263185));
assert!(approx_eq!(f64, cube_data.cell[2][1], 0.0));
assert!(approx_eq!(f64, cube_data.cell[2][2], 15.872644999999999));
assert_eq!(cube_data.atoms[0].an, 3usize);
assert_eq!(cube_data.atoms[cube_data.atoms.len() - 1].an, 16usize);
assert!(approx_eq!(f32, cube_data.atoms[0].charge, 0f32));
assert_eq!(
cube_data.data.len(),
(cube_data.shape[0] * cube_data.shape[1] * cube_data.shape[2]) as usize
);
}
#[test]
fn write_file() {
let cube_data = super::read_cube("tdc.cube").unwrap();
super::write_cube("tdc_out.cube", &cube_data).unwrap();
let cube_data_read = super::read_cube("tdc_out.cube").unwrap();
assert_eq!(cube_data.n_atoms, cube_data_read.n_atoms);
assert_eq!(cube_data.origin, cube_data_read.origin);
assert_eq!(cube_data.shape, cube_data_read.shape);
assert_eq!(cube_data.atoms[0].an, cube_data_read.atoms[0].an);
assert_eq!(
cube_data.atoms[cube_data.atoms.len() - 1].an,
cube_data_read.atoms[cube_data_read.atoms.len() - 1].an
);
approx_eq!(
f32,
cube_data.atoms[0].charge,
cube_data_read.atoms[0].charge
);
assert_eq!(
cube_data_read.data.len(),
(cube_data_read.shape[0] * cube_data_read.shape[1] * cube_data_read.shape[2]) as usize
);
}
}