gosh_model/
model_properties.rs1use std::collections::HashMap;
3use std::fmt;
4use std::str::FromStr;
5
6use super::*;
7
8use gchemol::prelude::*;
9use gchemol::Atom;
10use gchemol::Lattice;
11use gchemol::Molecule;
12const MODEL_PROPERTIES_FORMAT_VERSION: &str = "0.1";
16
17#[derive(Debug, Clone, Default, Serialize, Deserialize)]
19pub struct Computed {
20 energy: Option<f64>,
21 forces: Option<Vec<[f64; 3]>>,
22 stress: Option<[[f64; 3]; 3]>,
23 dipole: Option<[f64; 3]>,
24 #[serde(skip_deserializing, skip_serializing)]
25 molecule: Option<Molecule>,
26 #[serde(skip_deserializing, skip_serializing)]
27 force_constants: Option<Vec<[f64; 3]>>,
28}
29#[derive(Debug, Clone)]
33struct Header {
34 name: String,
35 unit_factor: f64,
36}
37
38impl FromStr for Header {
39 type Err = gut::prelude::Error;
40
41 fn from_str(s: &str) -> Result<Self> {
42 if s.starts_with("@") {
43 let mut unit_factor = 1.0;
44 let parts = &s[1..].split_whitespace().collect_vec();
45 let name = parts[0].into();
46 if parts.len() > 1 {
47 for p in &parts[1..] {
48 if let Some((k, v)) = p.split_once('=') {
49 if k == "unit_factor" {
50 unit_factor = v.parse::<f64>()?;
51 }
52 }
53 }
54 }
55 Ok(Self { name, unit_factor })
56 } else {
57 bail!("invalid model properties section header: {}", s);
58 }
59 }
60}
61
62#[test]
63fn test_header() {
64 let s = "@forces ";
65 let h: Header = s.parse().unwrap();
66 assert_eq!(h.name, "forces");
67 assert_eq!(h.unit_factor, 1.0);
68
69 let s = "@forces unit_factor=1";
70 let h: Header = s.parse().unwrap();
71 assert_eq!(h.unit_factor, 1.0);
72
73 let s = "@forces unit_factor=-1 test=2";
74 let h: Header = s.parse().unwrap();
75 assert_eq!(h.unit_factor, -1.0);
76}
77impl Computed {
81 pub fn parse_all(output: &str) -> Result<Vec<Computed>> {
83 parse_model_results(output)
84 }
85
86 pub fn is_empty(&self) -> bool {
88 self.energy.is_none() && self.forces.is_none()
90 }
91}
92
93impl fmt::Display for Computed {
94 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
95 let mut txt = format!("@model_properties_format_version {}\n", MODEL_PROPERTIES_FORMAT_VERSION);
96
97 if let Some(mol) = &self.molecule {
99 txt.push_str("@structure\n");
100 let coords = mol.format_as("text/pxyz").expect("formatted molecule");
101 txt.push_str(&coords);
102 }
103 if let Some(energy) = &self.energy {
105 txt.push_str("@energy\n");
106 txt.push_str(&format!("{:-20.12E}\n", energy));
107 }
108 if let Some(forces) = &self.forces {
110 txt.push_str("@forces\n");
111 for [fx, fy, fz] in forces {
112 let line = format!("{:-20.12E} {:-20.12E} {:-20.12E}\n", fx, fy, fz);
113 txt.push_str(&line);
114 }
115 }
116 if let Some(stress) = &self.stress {
117 txt.push_str("@stress\n");
118 for [fx, fy, fz] in stress {
119 let line = format!("{:-20.12E} {:-20.12E} {:-20.12E}\n", fx, fy, fz);
120 txt.push_str(&line);
121 }
122 }
123
124 if let Some(d) = &self.dipole {
126 txt.push_str("@dipole\n");
127 let line = format!("{:-20.12E} {:-20.12E} {:-20.12E}\n", d[0], d[1], d[2]);
128 txt.push_str(&line);
129 }
130
131 write!(f, "{}", txt)
132 }
133}
134
135impl FromStr for Computed {
136 type Err = gut::prelude::Error;
137
138 fn from_str(s: &str) -> Result<Self> {
139 let all = parse_model_results(s)?;
140
141 let n = all.len();
142 if n == 0 {
143 bail!("no valid results found from:\n {s:?}!");
144 }
145
146 Ok(all[n - 1].clone())
147 }
148}
149
150fn parse_model_results_single(part: &[&str]) -> Result<Computed> {
152 let mut records: HashMap<&str, Vec<&str>> = HashMap::new();
155 let mut header = None;
156 for line in part {
157 let line = line.trim();
158 if line.starts_with("@") {
159 header = line.trim_end().into();
161 } else {
162 if let Some(k) = header {
163 records.entry(k).or_insert(vec![]).push(line);
164 }
165 }
166 }
167
168 if records.len() < 1 {
170 warn!("Collected no results. Please check if the stream is clean!");
171 warn!("suspicious part: {:?}", part);
172 }
173
174 let mut results = Computed::default();
175 for (k, lines) in records {
176 let header: Header = k.parse()?;
177 let unit_factor = header.unit_factor;
178 match header.name.as_str() {
179 "energy" => {
180 assert_eq!(1, lines.len(), "expect one line containing energy");
181 let energy = lines[0].trim().parse::<f64>()? * unit_factor;
182 results.energy = Some(energy);
183 }
184 "forces" => {
185 let mut forces: Vec<[f64; 3]> = vec![];
186 for line in lines {
187 let parts: Vec<_> = line.split_whitespace().collect();
188 if parts.len() != 3 {
189 bail!("expect xyz forces: {}", line);
190 }
191 let fx = parts[0].parse::<f64>()? * unit_factor;
192 let fy = parts[1].parse::<f64>()? * unit_factor;
193 let fz = parts[2].parse::<f64>()? * unit_factor;
194 forces.push([fx, fy, fz]);
195 }
196
197 results.forces = Some(forces);
198 }
199 "structure" => {
200 let mut s = lines.join("\n");
201 s.push_str("\n\n");
202 let mol = Molecule::from_str(&s, "text/pxyz")?;
203 results.molecule = Some(mol);
204 }
205 "dipole" => {
206 assert_eq!(1, lines.len(), "expect one line containing dipole moment");
207 let parts: Vec<_> = lines[0].split_whitespace().collect();
208 let fx = parts[0].parse::<f64>()? * unit_factor;
209 let fy = parts[1].parse::<f64>()? * unit_factor;
210 let fz = parts[2].parse::<f64>()? * unit_factor;
211 results.dipole = Some([fx, fy, fz]);
212 }
213 _ => {
214 warn!("ignored record: {:?}", k);
215 }
216 }
217 }
218
219 Ok(results)
220}
221
222fn parse_model_results(stream: &str) -> Result<Vec<Computed>> {
223 if stream.trim().is_empty() {
224 bail!("Attemp to parse empty string!");
225 }
226
227 let lines: Vec<_> = stream
229 .lines()
230 .filter(|l| {
231 let l = l.trim();
232 !l.starts_with("#") && !l.is_empty()
233 })
234 .collect();
235
236 let parts = lines[1..].split(|l| l.starts_with("@model_properties_format_version"));
237
238 let mut all_results = vec![];
239 for part in parts {
240 if part.is_empty() {
244 continue;
245 } else {
246 let mp = parse_model_results_single(part)?;
247 all_results.push(mp);
248 }
249 }
250
251 Ok(all_results)
252}
253impl Computed {
256 pub fn set_energy(&mut self, e: f64) {
258 self.energy = Some(e);
259 }
260
261 pub fn set_forces(&mut self, f: Vec<[f64; 3]>) {
263 self.forces = Some(f);
264 }
265
266 pub fn set_stress(&mut self, s: [[f64; 3]; 3]) {
268 self.stress = s.into();
269 }
270
271 pub fn set_dipole(&mut self, d: [f64; 3]) {
273 self.dipole = Some(d);
274 }
275
276 pub fn set_molecule(&mut self, m: Molecule) {
278 self.molecule = Some(m);
279 }
280
281 pub fn set_force_constants(&mut self, fc: Vec<[f64; 3]>) {
283 self.force_constants = Some(fc);
284 }
285
286 pub fn get_energy(&self) -> Option<f64> {
288 self.energy
289 }
290
291 pub fn get_dipole(&self) -> Option<[f64; 3]> {
293 self.dipole
294 }
295
296 pub fn get_forces(&self) -> Option<&Vec<[f64; 3]>> {
298 self.forces.as_ref()
299 }
300
301 pub fn get_stress(&self) -> Option<[[f64; 3]; 3]> {
303 self.stress
304 }
305
306 pub fn get_molecule(&self) -> Option<&Molecule> {
308 self.molecule.as_ref()
309 }
310
311 pub fn get_force_constants(&self) -> Option<&Vec<[f64; 3]>> {
313 self.force_constants.as_ref()
314 }
315
316 pub fn set_structure<A, C>(&mut self, atoms: A, cell: Option<C>, scaled: bool)
326 where
327 A: IntoIterator,
328 A::Item: Into<Atom>,
329 C: Into<[[f64; 3]; 3]>,
330 {
331 let mut mol = Molecule::from_atoms(atoms);
332
333 if let Some(lat) = cell.map(|x| Lattice::new(x.into())) {
334 mol.set_lattice(lat);
335 if scaled {
336 let positions: Vec<_> = mol.positions().collect();
337 mol.set_scaled_positions(positions);
338 }
339 }
340
341 self.molecule = Some(mol);
342 }
343}
344
345#[test]
347fn test_model_parse_results() {
348 use vecfx::approx::*;
349
350 use serde_json;
351
352 let txt = gchemol::io::read_file("tests/files/sample.txt").unwrap();
353 let r: Computed = txt.parse().expect("model results");
354
355 let serialized = serde_json::to_string(&r).unwrap();
357 let _: Computed = serde_json::from_str(&serialized).unwrap();
359
360 let txt = format!("{}", r);
362
363 let r: Computed = txt.parse().expect("model results");
365
366 assert!(&r.molecule.is_some());
367 let ref mol = r.molecule.unwrap();
368 assert_eq!(3, mol.natoms());
369 let e = r.energy.expect("model result: energy");
370 assert_relative_eq!(-0.329336, e, epsilon = 1e-4);
371}
372
373#[test]
374fn test_model_parse_results_special() -> Result<()> {
375
376 let txt = gchemol::io::read_file("./tests/files/sample_special.txt")?;
377 let r: Computed = txt.parse()?;
378 assert_eq!(r.energy.unwrap(), -0.32933619218901);
379 assert_eq!(r.forces.unwrap()[0][0], 0.10525500903260E-03);
380
381 Ok(())
382}
383