1use crate::corpus::ConCorpus;
6use crate::error::Result;
7use crate::select::Select;
8
9#[derive(Clone, Debug)]
11pub struct H5mdArrays {
12 pub n_frames: usize,
13 pub natoms: usize,
14 pub positions: Vec<f64>,
16 pub edges: Vec<f64>,
18 pub species_z: Vec<i32>,
20 pub forces: Option<Vec<f64>>,
22 pub velocities: Option<Vec<f64>>,
24 pub velocity_unit: String,
25 pub boundary: [String; 3],
27 pub times: Vec<f64>,
29 pub time_unit: String,
30 pub length_unit: String,
31 pub force_unit: String,
32}
33
34pub const H5MD_LENGTH_CORE: &str = "angstrom";
36pub const H5MD_LENGTH_ATTR: &str = "Angstrom";
37pub const H5MD_TIME_CORE: &str = "ps";
38pub const H5MD_TIME_ATTR: &str = "ps";
39pub const CON_TIME_DEFAULT: &str = "fs";
41pub const H5MD_FORCE_ATTR: &str = "kJ mol-1 Angstrom-1";
42pub const H5MD_VELOCITY_ATTR: &str = "Angstrom ps-1";
43const KJ_MOL_ANGSTROM_SI: f64 = (1000.0 / 6.022_140_76e23) / 1e-10;
45
46fn boundary_from_pbc(pbc: Option<[bool; 3]>) -> [String; 3] {
47 let p = pbc.unwrap_or([true, true, true]);
48 std::array::from_fn(|i| {
49 if p[i] {
50 "periodic".into()
51 } else {
52 "none".into()
53 }
54 })
55}
56
57fn boxl_to_edges33(boxl: &[f64; 3]) -> [f64; 9] {
58 [boxl[0], 0.0, 0.0, 0.0, boxl[1], 0.0, 0.0, 0.0, boxl[2]]
59}
60
61fn uc(from: &str, to: &str) -> Result<f64> {
62 readcon_core::units::unit_conversion_factor(from, to)
63 .map_err(|e| crate::error::Error::Message(e.to_string()))
64}
65
66fn header_unit(h: &readcon_core::types::FrameHeader, dim: &str, default: &str) -> String {
67 h.unit_for(dim)
68 .map(str::to_string)
69 .unwrap_or_else(|| default.to_string())
70}
71
72fn force_scale_to_engine(energy_u: &str, length_u: &str) -> Result<f64> {
74 let e_j = uc(energy_u, "J")?;
75 let l_m = uc(length_u, "m")?;
76 Ok((e_j / l_m) / KJ_MOL_ANGSTROM_SI)
77}
78
79fn time_scale_to_ps(from: &str) -> Result<f64> {
80 match uc(from, H5MD_TIME_CORE) {
81 Ok(f) => Ok(f),
82 Err(_) if from.eq_ignore_ascii_case("ns") || from.eq_ignore_ascii_case("nanosecond") => {
83 Ok(1e3)
84 }
85 Err(e) => Err(e),
86 }
87}
88
89fn frame_time_ps(h: &readcon_core::types::FrameHeader, frame_idx: u32) -> Result<f64> {
90 let from = header_unit(h, "time", CON_TIME_DEFAULT);
91 if let Some(t) = h.time() {
92 return Ok(t * time_scale_to_ps(&from)?);
93 }
94 if let Some(dt) = h.timestep().filter(|x| x.is_finite() && *x > 0.0) {
95 return Ok(f64::from(frame_idx) * dt * time_scale_to_ps(&from)?);
96 }
97 Ok(f64::from(frame_idx))
98}
99
100pub(crate) fn edges33_from_header(h: &readcon_core::types::FrameHeader) -> [f64; 9] {
101 if let Some(arr) = h.metadata.get("lattice_vectors").and_then(|v| v.as_array()) {
102 if arr.len() == 3 {
103 let mut out = [0.0f64; 9];
104 let mut ok = true;
105 for (i, row) in arr.iter().enumerate() {
106 let Some(r) = row.as_array() else {
107 ok = false;
108 break;
109 };
110 if r.len() != 3 {
111 ok = false;
112 break;
113 }
114 for (j, c) in r.iter().enumerate() {
115 let Some(x) = c.as_f64() else {
116 ok = false;
117 break;
118 };
119 out[i * 3 + j] = x;
120 }
121 }
122 if ok {
123 return out;
124 }
125 }
126 }
127 boxl_angles_to_edges33(&h.boxl, &h.angles)
128}
129
130fn boxl_angles_to_edges33(boxl: &[f64; 3], angles: &[f64; 3]) -> [f64; 9] {
131 let ortho = angles.iter().all(|&a| a == 0.0 || (a - 90.0).abs() < 1e-9);
132 if ortho {
133 return boxl_to_edges33(boxl);
134 }
135 let deg = |a: f64| a * std::f64::consts::PI / 180.0;
136 let (lx, ly, lz) = (boxl[0], boxl[1], boxl[2]);
137 let (alpha, beta, gamma) = (deg(angles[0]), deg(angles[1]), deg(angles[2]));
138 let ax = lx;
139 let bx = ly * gamma.cos();
140 let by = ly * gamma.sin();
141 let cx = lz * beta.cos();
142 let cy = lz * (alpha.cos() - beta.cos() * gamma.cos()) / gamma.sin();
143 let cz = (lz * lz - cx * cx - cy * cy).max(0.0).sqrt();
144 [ax, 0.0, 0.0, bx, by, 0.0, cx, cy, cz]
145}
146
147impl ConCorpus {
148 pub fn collect_h5md(&self, traj_id: u64) -> Result<H5mdArrays> {
152 let keys = self.select(&Select::new().trajectory(traj_id))?;
153 if keys.is_empty() {
154 return Err(crate::error::Error::Message("no frames for traj".into()));
155 }
156 let first = self.get_frame(keys[0])?;
157 let natoms = first.atom_data.len();
158 let species_z: Vec<i32> = first
159 .atom_data
160 .iter()
161 .map(|a| readcon_core::helpers::symbol_to_atomic_number(a.symbol.as_ref()) as i32)
162 .collect();
163 let n_frames = keys.len();
164 let boundary = boundary_from_pbc(first.header.pbc());
165 let mut positions = Vec::with_capacity(n_frames * natoms * 3);
166 let mut edges = Vec::with_capacity(n_frames * 9);
167 let mut times = Vec::with_capacity(n_frames);
168 let mut force_rows: Vec<Option<Vec<[f64; 3]>>> = Vec::with_capacity(n_frames);
169 let mut vel_rows: Vec<Option<Vec<[f64; 3]>>> = Vec::with_capacity(n_frames);
170 for (i, k) in keys.iter().enumerate() {
171 let fr = if i == 0 {
172 first.clone()
173 } else {
174 self.get_frame(*k)?
175 };
176 let length_u = header_unit(&fr.header, "length", H5MD_LENGTH_CORE);
177 let energy_u = header_unit(&fr.header, "energy", "eV");
178 let time_u = header_unit(&fr.header, "time", CON_TIME_DEFAULT);
179 let len_scale = uc(&length_u, H5MD_LENGTH_CORE)?;
180 let vel_scale = len_scale / time_scale_to_ps(&time_u)?;
181 times.push(frame_time_ps(&fr.header, k.frame_idx)?);
182 let e33 = edges33_from_header(&fr.header);
183 for x in e33 {
184 edges.push(x * len_scale);
185 }
186 if fr.atom_data.len() != natoms {
187 return Err(crate::error::Error::Message(
188 "H5MD export needs fixed natoms in the trajectory".into(),
189 ));
190 }
191 if boundary_from_pbc(fr.header.pbc()) != boundary {
192 return Err(crate::error::Error::Message(
193 "H5MD export needs fixed box/boundary in the trajectory".into(),
194 ));
195 }
196 let z_here: Vec<i32> = fr
197 .atom_data
198 .iter()
199 .map(|a| readcon_core::helpers::symbol_to_atomic_number(a.symbol.as_ref()) as i32)
200 .collect();
201 if z_here != species_z {
202 return Err(crate::error::Error::Message(
203 "H5MD export needs fixed species Z in the trajectory".into(),
204 ));
205 }
206 let (pos_src, vel_src, frc_src) = if let Some(c) = self.get_cooked_soa(*k)? {
207 (c.positions, c.velocities, c.forces)
208 } else {
209 let pos: Vec<[f64; 3]> = fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect();
210 let vel = if fr.atom_data.iter().any(|a| a.velocity.is_some()) {
211 Some(
212 fr.atom_data
213 .iter()
214 .map(|a| a.velocity.unwrap_or([0.0; 3]))
215 .collect(),
216 )
217 } else {
218 None
219 };
220 let frc = if fr.atom_data.iter().any(|a| a.force.is_some()) {
221 Some(
222 fr.atom_data
223 .iter()
224 .map(|a| a.force.unwrap_or([0.0; 3]))
225 .collect(),
226 )
227 } else {
228 None
229 };
230 (pos, vel, frc)
231 };
232 for p in &pos_src {
233 positions.extend_from_slice(&[
234 p[0] * len_scale,
235 p[1] * len_scale,
236 p[2] * len_scale,
237 ]);
238 }
239 vel_rows.push(match vel_src {
240 Some(rows) => Some(
241 rows.into_iter()
242 .map(|r| [r[0] * vel_scale, r[1] * vel_scale, r[2] * vel_scale])
243 .collect(),
244 ),
245 None => None,
246 });
247 force_rows.push(match frc_src {
248 Some(rows) => {
249 let fscale = force_scale_to_engine(&energy_u, &length_u)?;
250 Some(
251 rows.into_iter()
252 .map(|r| [r[0] * fscale, r[1] * fscale, r[2] * fscale])
253 .collect(),
254 )
255 }
256 None => None,
257 });
258 }
259 let forces = if force_rows.iter().any(|f| f.is_some()) {
260 let mut fbuf = vec![0.0f64; n_frames * natoms * 3];
261 for (ti, fo) in force_rows.iter().enumerate() {
262 if let Some(rows) = fo {
263 let off = ti * natoms * 3;
264 for (i, row) in rows.iter().enumerate() {
265 fbuf[off + i * 3] = row[0];
266 fbuf[off + i * 3 + 1] = row[1];
267 fbuf[off + i * 3 + 2] = row[2];
268 }
269 }
270 }
271 Some(fbuf)
272 } else {
273 None
274 };
275 let velocities = if vel_rows.iter().any(|v| v.is_some()) {
276 let mut vbuf = vec![0.0f64; n_frames * natoms * 3];
277 for (ti, vo) in vel_rows.iter().enumerate() {
278 if let Some(rows) = vo {
279 let off = ti * natoms * 3;
280 for (i, row) in rows.iter().enumerate() {
281 vbuf[off + i * 3] = row[0];
282 vbuf[off + i * 3 + 1] = row[1];
283 vbuf[off + i * 3 + 2] = row[2];
284 }
285 }
286 }
287 Some(vbuf)
288 } else {
289 None
290 };
291 Ok(H5mdArrays {
292 n_frames,
293 natoms,
294 positions,
295 edges,
296 species_z,
297 forces,
298 velocities,
299 velocity_unit: H5MD_VELOCITY_ATTR.into(),
300 boundary,
301 times,
302 time_unit: H5MD_TIME_ATTR.into(),
303 length_unit: H5MD_LENGTH_ATTR.into(),
304 force_unit: H5MD_FORCE_ATTR.into(),
305 })
306 }
307}
308
309#[cfg(test)]
310mod tests {
311 use super::*;
312 use std::path::PathBuf;
313
314 fn fixture(name: &str) -> PathBuf {
315 PathBuf::from(env!("CARGO_MANIFEST_DIR"))
316 .join("resources/test")
317 .join(name)
318 }
319
320 #[test]
321 fn collect_h5md_tn3_from_con() {
322 let dir = tempfile::tempdir().unwrap();
323 let db = ConCorpus::open(dir.path()).unwrap();
324 let n = db
325 .append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
326 .unwrap();
327 assert!(n >= 2);
328 let a = db.collect_h5md(1).unwrap();
329 assert_eq!(a.n_frames, n as usize);
330 assert!(a.natoms >= 1);
331 assert_eq!(a.positions.len(), a.n_frames * a.natoms * 3);
332 assert_eq!(a.edges.len(), a.n_frames * 9);
333 assert_eq!(a.species_z.len(), a.natoms);
334 assert!(a.species_z.iter().all(|&z| z > 0));
335 assert_eq!(a.boundary.len(), 3);
336 assert_eq!(
337 a.boundary,
338 [
339 "periodic".to_string(),
340 "periodic".to_string(),
341 "periodic".to_string()
342 ]
343 );
344 assert_eq!(a.times.len(), a.n_frames);
345 for (i, t) in a.times.iter().enumerate() {
346 assert!(
347 (*t - i as f64).abs() < 1e-12,
348 "frame-index dest ps: times[{i}]={t}"
349 );
350 }
351 assert!((a.positions[0] - 0.6394).abs() < 1e-4);
352 let i_h1 = a.natoms * 3 + 2 * 3;
353 assert!(
354 (a.positions[i_h1] - 8.8549).abs() < 1e-4,
355 "tn3 frame1 H dest Å x={}",
356 a.positions[i_h1]
357 );
358 }
359
360 #[test]
361 fn export_extxyz_writes_pbc_and_triclinic_lattice() {
362 let dir = tempfile::tempdir().unwrap();
363 let db = ConCorpus::open(dir.path()).unwrap();
364 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
365 let mut frames = Vec::new();
366 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
367 frames.push(item.unwrap());
368 }
369 frames[0].header.angles = [60.0, 90.0, 70.0];
370 frames[0]
371 .header
372 .metadata
373 .insert("pbc".into(), serde_json::json!([false, true, false]));
374 db.append_trajectory_frames(1, &frames, "t").unwrap();
375 let keys = db
376 .select(&crate::select::Select::new().trajectory(1))
377 .unwrap();
378 let xyz = dir.path().join("t.xyz");
379 db.export_extxyz(&keys, &xyz, "energy").unwrap();
380 let out = std::fs::read_to_string(&xyz).unwrap();
381 assert!(out.contains("pbc=\"F T F\""), "{out}");
382 assert!(!out.contains("pbc=\"T T T\""));
383 let e = crate::export_h5md::edges33_from_header(&frames[0].header);
384 assert!(e[3].abs() > 1e-9, "triclinic b_x");
385 assert!(out.contains(&format!("{:.10}", e[3])), "{out}");
386 }
387
388 #[test]
389 fn collect_h5md_rejects_changing_natoms() {
390 let dir = tempfile::tempdir().unwrap();
391 let db = ConCorpus::open(dir.path()).unwrap();
392 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
393 .unwrap();
394 db.extend_trajectory_path(1, fixture("sulfolene.con"))
395 .unwrap();
396 let err = db.collect_h5md(1).unwrap_err();
397 assert!(err.to_string().contains("fixed natoms"), "{err}");
398 }
399
400 #[test]
401 fn collect_h5md_rejects_changing_pbc() {
402 let dir = tempfile::tempdir().unwrap();
403 let db = ConCorpus::open(dir.path()).unwrap();
404 let text = std::fs::read_to_string(fixture("tiny_multi_cuh2.con")).unwrap();
405 let mut frames = Vec::new();
406 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
407 frames.push(item.unwrap());
408 }
409 assert!(frames.len() >= 2);
410 frames[0]
411 .header
412 .metadata
413 .insert("pbc".into(), serde_json::json!([true, true, true]));
414 frames[1]
415 .header
416 .metadata
417 .insert("pbc".into(), serde_json::json!([false, false, false]));
418 db.append_trajectory_frames(1, &frames, "t").unwrap();
419 let err = db.collect_h5md(1).unwrap_err();
420 assert!(err.to_string().contains("fixed box/boundary"), "{err}");
421 }
422
423 #[test]
424 fn collect_h5md_open_readonly() {
425 let dir = tempfile::tempdir().unwrap();
426 {
427 let db = ConCorpus::open(dir.path()).unwrap();
428 db.append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
429 .unwrap();
430 db.close();
431 }
432 let ro = ConCorpus::open_readonly(dir.path()).unwrap();
433 let a = ro.collect_h5md(1).unwrap();
434 assert!(a.n_frames >= 2);
435 assert_eq!(a.positions.len(), a.n_frames * a.natoms * 3);
436 assert!((a.positions[0] - 0.6394).abs() < 1e-4);
437 let i_h1 = a.natoms * 3 + 2 * 3;
438 assert!(
439 (a.positions[i_h1] - 8.8549).abs() < 1e-4,
440 "readonly frame1 H dest Å x={}",
441 a.positions[i_h1]
442 );
443 }
444
445 #[test]
446 fn collect_h5md_two_frame_distinct_boxl() {
447 let dir = tempfile::tempdir().unwrap();
448 let db = ConCorpus::open(dir.path()).unwrap();
449 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
450 let mut frames = Vec::new();
451 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
452 frames.push(item.unwrap());
453 }
454 let mut f1 = frames[0].clone();
455 f1.header.boxl = [20.0, 21.702, 100.0];
456 frames.push(f1);
457 db.append_trajectory_frames(1, &frames, "t").unwrap();
458 let a = db.collect_h5md(1).unwrap();
459 assert!(a.n_frames >= 2);
460 assert!((a.edges[0] - 15.3456).abs() < 1e-4, "e0={}", a.edges[0]);
461 assert!((a.edges[9] - 20.0).abs() < 1e-9, "e9={}", a.edges[9]);
462 }
463
464 #[test]
465 fn collect_h5md_mixed_pbc_f_t_f() {
466 let dir = tempfile::tempdir().unwrap();
467 let db = ConCorpus::open(dir.path()).unwrap();
468 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
469 let mut frames = Vec::new();
470 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
471 frames.push(item.unwrap());
472 }
473 frames[0]
474 .header
475 .metadata
476 .insert("pbc".into(), serde_json::json!([false, true, false]));
477 db.append_trajectory_frames(1, &frames, "t").unwrap();
478 let a = db.collect_h5md(1).unwrap();
479 assert_eq!(
480 a.boundary,
481 [
482 "none".to_string(),
483 "periodic".to_string(),
484 "none".to_string()
485 ]
486 );
487 }
488
489 #[test]
490 fn collect_h5md_boundary_none_from_pbc_false() {
491 let dir = tempfile::tempdir().unwrap();
492 let db = ConCorpus::open(dir.path()).unwrap();
493 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
494 let mut frames = Vec::new();
495 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
496 frames.push(item.unwrap());
497 }
498 frames[0]
499 .header
500 .metadata
501 .insert("pbc".into(), serde_json::json!([false, false, false]));
502 db.append_trajectory_frames(1, &frames, "t").unwrap();
503 let a = db.collect_h5md(1).unwrap();
504 assert_eq!(
505 a.boundary,
506 ["none".to_string(), "none".to_string(), "none".to_string()]
507 );
508 }
509
510 #[test]
511 fn collect_h5md_pads_mixed_forces() {
512 let dir = tempfile::tempdir().unwrap();
513 let db = ConCorpus::open(dir.path()).unwrap();
514 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
515 .unwrap();
516 db.extend_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
517 .unwrap();
518 let a = db.collect_h5md(1).unwrap();
519 let f = a.forces.expect("second frame carries forces");
520 assert_eq!(f.len(), a.n_frames * a.natoms * 3);
521 assert!(a.n_frames >= 2);
522 assert!(f[..a.natoms * 3].iter().all(|&x| x == 0.0));
523 assert!(f[a.natoms * 3..].iter().any(|&x| x != 0.0));
524 }
525
526 #[test]
527 fn collect_h5md_con_fallback_matches_rcso() {
528 let dir = tempfile::tempdir().unwrap();
529 let db = ConCorpus::open(dir.path()).unwrap();
530 db.append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
531 .unwrap();
532 let key = crate::keys::FrameKey {
533 traj_id: 1,
534 frame_idx: 0,
535 };
536 assert!(!db.has_valid_cooked_soa(key).unwrap());
537 let from_con = db.collect_h5md(1).unwrap();
538 db.recook_all().unwrap();
539 assert!(db.has_valid_cooked_soa(key).unwrap());
540 let from_rcso = db.collect_h5md(1).unwrap();
541 assert_eq!(from_con.n_frames, from_rcso.n_frames);
542 assert_eq!(from_con.positions, from_rcso.positions);
543 assert_eq!(from_con.edges, from_rcso.edges);
544 assert_eq!(from_con.species_z, from_rcso.species_z);
545 let blob = db.get_cooked_soa_bytes(key).unwrap().expect("stored RCSO");
546 let decoded = crate::cooked_soa::CookedSoa::decode(&blob).unwrap();
547 assert_eq!(decoded.positions.len(), from_rcso.natoms);
548 assert!((decoded.positions[0][0] - 0.6394).abs() < 1e-4);
549 assert!((from_rcso.positions[0] - 0.6394).abs() < 1e-4);
550 assert_eq!(decoded.positions[0][0], from_con.positions[0]);
551 let i_h1 = from_rcso.natoms * 3 + 2 * 3;
553 assert!(
554 (from_rcso.positions[i_h1] - 8.8549).abs() < 1e-4,
555 "frame1 H x={}",
556 from_rcso.positions[i_h1]
557 );
558 assert!((from_con.positions[i_h1] - 8.8549).abs() < 1e-4);
559 }
560
561 #[test]
562 fn collect_h5md_uses_con_time_and_fs_unit() {
563 let dir = tempfile::tempdir().unwrap();
564 let db = ConCorpus::open(dir.path()).unwrap();
565 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
566 let mut frames = Vec::new();
567 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
568 frames.push(item.unwrap());
569 }
570 frames[0].header.set_time(12.5);
571 frames[0].header.metadata.insert(
572 "units".into(),
573 serde_json::json!({"length":"angstrom","energy":"eV","mass":"amu","time":"fs"}),
574 );
575 db.append_trajectory_frames(1, &frames, "t").unwrap();
576 let a = db.collect_h5md(1).unwrap();
577 assert!(
578 (a.times[0] - 0.0125).abs() < 1e-12,
579 "12.5 fs -> ps, got {}",
580 a.times[0]
581 );
582 assert_eq!(a.time_unit, "ps");
583 }
584
585 #[test]
586 fn collect_h5md_scales_nm_length_to_angstrom() {
587 let dir = tempfile::tempdir().unwrap();
588 let db = ConCorpus::open(dir.path()).unwrap();
589 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
590 let mut frames = Vec::new();
591 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
592 frames.push(item.unwrap());
593 }
594 let box0 = frames[0].header.boxl[0];
595 frames[0].header.metadata.insert(
596 "units".into(),
597 serde_json::json!({"length":"nm","energy":"eV","mass":"amu","time":"fs"}),
598 );
599 db.append_trajectory_frames(1, &frames, "t").unwrap();
600 let a = db.collect_h5md(1).unwrap();
601 let scale = uc("nm", "angstrom").unwrap();
602 assert!((scale - 10.0).abs() < 1e-12);
603 assert!((a.edges[0] - box0 * scale).abs() < 1e-9);
604 assert!(
605 (a.positions[0] - 0.6394 * scale).abs() < 1e-4,
606 "dest Å x0 from nm-labelled CON, got {}",
607 a.positions[0]
608 );
609 assert_eq!(a.length_unit, "Angstrom");
610 }
611
612 #[test]
613 fn collect_h5md_ns_time_to_ps() {
614 let dir = tempfile::tempdir().unwrap();
615 let db = ConCorpus::open(dir.path()).unwrap();
616 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
617 let mut frames = Vec::new();
618 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
619 frames.push(item.unwrap());
620 }
621 frames[0].header.set_time(2.0);
622 frames[0].header.metadata.insert(
623 "units".into(),
624 serde_json::json!({"length":"angstrom","energy":"eV","mass":"amu","time":"ns"}),
625 );
626 db.append_trajectory_frames(1, &frames, "t").unwrap();
627 let a = db.collect_h5md(1).unwrap();
628 assert!(
629 (a.times[0] - 2000.0).abs() < 1e-9,
630 "2 ns -> ps, got {}",
631 a.times[0]
632 );
633 assert_eq!(a.time_unit, "ps");
634 }
635
636 #[test]
637 fn collect_h5md_converts_force_via_core_units() {
638 let factor = force_scale_to_engine("eV", "angstrom").unwrap();
639 assert!((factor - 96.485_332).abs() < 1e-3, "got {factor}");
640 let dir = tempfile::tempdir().unwrap();
641 let db = ConCorpus::open(dir.path()).unwrap();
642 db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
643 .unwrap();
644 let a = db.collect_h5md(1).unwrap();
645 let f = a.forces.expect("forces");
646 let cooked = crate::cooked_soa::CookedSoa::decode(
647 &db.pack_frame(crate::keys::FrameKey {
648 traj_id: 1,
649 frame_idx: 0,
650 })
651 .unwrap(),
652 )
653 .unwrap();
654 let native = cooked.forces.expect("native");
655 assert!((f[0] - native[0][0] * factor).abs() < 1e-8);
656 assert_eq!(a.force_unit, "kJ mol-1 Angstrom-1");
657 assert_eq!(a.length_unit, "Angstrom");
658 }
659
660 #[test]
661 fn collect_h5md_triclinic_edges_from_angles() {
662 let dir = tempfile::tempdir().unwrap();
663 let db = ConCorpus::open(dir.path()).unwrap();
664 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
665 let mut frames = Vec::new();
666 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
667 frames.push(item.unwrap());
668 }
669 frames[0].header.angles = [60.0, 90.0, 70.0];
670 db.append_trajectory_frames(1, &frames, "t").unwrap();
671 let a = db.collect_h5md(1).unwrap();
672 assert_eq!(a.edges.len(), 9);
673 assert!(a.edges[3].abs() > 1e-9, "b_x from gamma != 90");
674 let ortho = boxl_to_edges33(&frames[0].header.boxl);
675 assert_ne!(a.edges, ortho.to_vec());
676 }
677
678 #[test]
679 fn collect_h5md_undeclared_time_is_con_fs() {
680 let dir = tempfile::tempdir().unwrap();
681 let db = ConCorpus::open(dir.path()).unwrap();
682 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
683 let mut frames = Vec::new();
684 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
685 frames.push(item.unwrap());
686 }
687 frames[0].header.set_time(12.5);
688 frames[0].header.metadata.remove("units");
689 db.append_trajectory_frames(1, &frames, "t").unwrap();
690 let a = db.collect_h5md(1).unwrap();
691 assert!(
692 (a.times[0] - 0.0125).abs() < 1e-12,
693 "12.5 with no units.time is CON fs -> 0.0125 ps, got {}",
694 a.times[0]
695 );
696 }
697
698 #[test]
699 fn collect_h5md_uses_i_times_timestep() {
700 let dir = tempfile::tempdir().unwrap();
701 let db = ConCorpus::open(dir.path()).unwrap();
702 let text = std::fs::read_to_string(fixture("tiny_multi_cuh2.con")).unwrap();
703 let mut frames = Vec::new();
704 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
705 frames.push(item.unwrap());
706 }
707 assert!(frames.len() >= 2);
708 for fr in &mut frames {
709 fr.header.metadata.remove("time");
710 fr.header.set_timestep(10.0);
711 fr.header.metadata.remove("units");
712 }
713 db.append_trajectory_frames(1, &frames, "t").unwrap();
714 let a = db.collect_h5md(1).unwrap();
715 assert!(a.times.len() >= 2);
716 assert!((a.times[0] - 0.0).abs() < 1e-12, "got {}", a.times[0]);
717 assert!(
718 (a.times[1] - 0.01).abs() < 1e-12,
719 "i=1 * 10 fs -> 0.01 ps, got {}",
720 a.times[1]
721 );
722 assert!((a.positions[0] - 0.6394).abs() < 1e-4);
723 let i_h1 = a.natoms * 3 + 2 * 3;
724 assert!(
725 (a.positions[i_h1] - 8.8549).abs() < 1e-4,
726 "frame1 H dest Å x={}",
727 a.positions[i_h1]
728 );
729 }
730
731 #[test]
732 fn collect_h5md_set_units_converts_then_export_matches() {
733 let dir = tempfile::tempdir().unwrap();
734 let db = ConCorpus::open(dir.path()).unwrap();
735 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
736 let mut frames = Vec::new();
737 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
738 frames.push(item.unwrap());
739 }
740 let box0 = frames[0].header.boxl[0];
741 frames[0].header.metadata.insert(
742 "units".into(),
743 serde_json::json!({"length":"angstrom","energy":"eV"}),
744 );
745 db.append_trajectory_frames(1, &frames, "t").unwrap();
746 let before = db.collect_h5md(1).unwrap();
747 db.set_trajectory_units(1, serde_json::json!({"length":"nm","energy":"eV"}))
748 .unwrap();
749 let after = db.collect_h5md(1).unwrap();
750 assert!((before.edges[0] - box0).abs() < 1e-9);
751 assert!((after.edges[0] - box0).abs() < 1e-9);
752 let u = db
753 .frame_units(crate::keys::FrameKey {
754 traj_id: 1,
755 frame_idx: 0,
756 })
757 .unwrap()
758 .unwrap();
759 assert_eq!(u["length"], "nm");
760 assert_eq!(before.times.len(), after.times.len());
761 for (bt, at) in before.times.iter().zip(&after.times) {
762 assert!(
763 (bt - at).abs() < 1e-12,
764 "dest time after set_units {bt} vs {at}"
765 );
766 }
767 }
768
769 #[test]
770 fn collect_h5md_rejects_changing_species_z() {
771 let dir = tempfile::tempdir().unwrap();
772 let db = ConCorpus::open(dir.path()).unwrap();
773 let text = std::fs::read_to_string(fixture("tiny_multi_cuh2.con")).unwrap();
774 let mut frames = Vec::new();
775 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
776 frames.push(item.unwrap());
777 }
778 assert!(frames.len() >= 2);
779 assert!(!frames[1].atom_data.is_empty());
780 frames[1].atom_data[0].symbol = std::sync::Arc::from("Au");
781 db.append_trajectory_frames(1, &frames, "t").unwrap();
782 let err = db.collect_h5md(1).unwrap_err();
783 assert!(err.to_string().contains("fixed species Z"), "{err}");
784 }
785
786 #[test]
787 fn collect_h5md_after_cook_set_units_keeps_dest_vel() {
788 let dir = tempfile::tempdir().unwrap();
789 let db = ConCorpus::open(dir.path()).unwrap();
790 db.append_trajectory_path(1, fixture("tiny_cuh2.convel"))
791 .unwrap();
792 db.cook_frame(crate::keys::FrameKey {
793 traj_id: 1,
794 frame_idx: 0,
795 })
796 .unwrap();
797 let before = db.collect_h5md(1).unwrap();
798 db.set_trajectory_units(1, serde_json::json!({"length": "nm", "energy": "eV"}))
799 .unwrap();
800 assert!(!db
801 .has_valid_cooked_soa(crate::keys::FrameKey {
802 traj_id: 1,
803 frame_idx: 0,
804 })
805 .unwrap());
806 let native_xyz = db
807 .get_positions(crate::keys::FrameKey {
808 traj_id: 1,
809 frame_idx: 0,
810 })
811 .unwrap();
812 assert!(
813 (native_xyz[0][0] - 0.06394).abs() < 1e-5,
814 "native nm x0={}",
815 native_xyz[0][0]
816 );
817 let after = db.collect_h5md(1).unwrap();
818 let bv = before.velocities.expect("vel");
819 let av = after.velocities.expect("vel");
820 assert_eq!(bv.len(), av.len());
821 let native = db
822 .get_velocities(crate::keys::FrameKey {
823 traj_id: 1,
824 frame_idx: 0,
825 })
826 .unwrap()
827 .expect("native vel");
828 let u = db
829 .frame_units(crate::keys::FrameKey {
830 traj_id: 1,
831 frame_idx: 0,
832 })
833 .unwrap();
834 for (i, (b, a)) in bv.iter().zip(av.iter()).enumerate() {
835 assert!(
836 (b - a).abs() < 1e-9,
837 "dest vel[{i}]: before={b} after={a} native0={} units={u:?}",
838 native[0][0]
839 );
840 }
841 db.set_trajectory_units(
842 1,
843 serde_json::json!({"length": "nm", "energy": "eV", "time": "ps"}),
844 )
845 .unwrap();
846 let after_t = db.collect_h5md(1).unwrap();
847 let avt = after_t.velocities.expect("vel after time relabel");
848 for (i, (b, a)) in bv.iter().zip(avt.iter()).enumerate() {
849 assert!(
850 (b - a).abs() < 1e-9,
851 "dest vel after set_units time[{i}]: {b} vs {a}"
852 );
853 }
854 assert_eq!(before.times.len(), after_t.times.len());
855 for (bt, at) in before.times.iter().zip(&after_t.times) {
856 assert!((bt - at).abs() < 1e-12, "dest time {bt} vs {at}");
857 }
858 }
859
860 #[test]
861 fn collect_h5md_after_cook_and_set_units() {
862 let dir = tempfile::tempdir().unwrap();
863 let db = ConCorpus::open(dir.path()).unwrap();
864 db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
865 .unwrap();
866 db.cook_frame(crate::keys::FrameKey {
867 traj_id: 1,
868 frame_idx: 0,
869 })
870 .unwrap();
871 assert!(db
872 .has_valid_cooked_soa(crate::keys::FrameKey {
873 traj_id: 1,
874 frame_idx: 0,
875 })
876 .unwrap());
877 let before = db.collect_h5md(1).unwrap();
878 db.set_trajectory_units(1, serde_json::json!({"length": "nm", "energy": "eV"}))
879 .unwrap();
880 let after = db.collect_h5md(1).unwrap();
881 assert!((before.edges[0] - after.edges[0]).abs() < 1e-9);
882 assert_eq!(before.positions.len(), after.positions.len());
883 for (i, (b, a)) in before
884 .positions
885 .iter()
886 .zip(after.positions.iter())
887 .enumerate()
888 {
889 assert!(
890 (b - a).abs() < 1e-9,
891 "dest Å positions after cook+set_units[{i}]: {b} vs {a}"
892 );
893 }
894 let bf = before.forces.expect("dest forces");
895 let af = after.forces.expect("dest forces after set_units");
896 assert_eq!(bf.len(), af.len());
897 let factor = force_scale_to_engine("eV", "angstrom").unwrap();
898 assert!(
899 (bf[0] - 0.123456 * factor).abs() < 1e-6,
900 "dest force0={}",
901 bf[0]
902 );
903 for (i, (b, a)) in bf.iter().zip(af.iter()).enumerate() {
904 assert!(
905 (b - a).abs() < 1e-8,
906 "dest force after cook+set_units[{i}]: {b} vs {a}"
907 );
908 }
909 db.set_trajectory_units(1, serde_json::json!({"length": "nm", "energy": "hartree"}))
910 .unwrap();
911 let after_e = db.collect_h5md(1).unwrap();
912 let after_energy = after_e
913 .forces
914 .expect("dest forces after set_units energy");
915 assert_eq!(bf.len(), after_energy.len());
916 for (i, (b, a)) in bf.iter().zip(after_energy.iter()).enumerate() {
917 assert!(
918 (b - a).abs() < 1e-8,
919 "dest force after set_units energy[{i}]: {b} vs {a}"
920 );
921 }
922 assert_eq!(
923 db.frame_units(crate::keys::FrameKey {
924 traj_id: 1,
925 frame_idx: 0,
926 })
927 .unwrap()
928 .unwrap()["length"],
929 "nm"
930 );
931 assert_eq!(
932 db.frame_units(crate::keys::FrameKey {
933 traj_id: 1,
934 frame_idx: 0,
935 })
936 .unwrap()
937 .unwrap()["energy"],
938 "hartree"
939 );
940 }
941
942 #[test]
943 fn collect_h5md_species_z_cu_h() {
944 let dir = tempfile::tempdir().unwrap();
945 let db = ConCorpus::open(dir.path()).unwrap();
946 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
947 .unwrap();
948 let a = db.collect_h5md(1).unwrap();
949 assert!(a.species_z.iter().any(|&z| z == 29), "{:?}", a.species_z);
950 assert!(a.species_z.iter().any(|&z| z == 1), "{:?}", a.species_z);
951 }
952
953 #[test]
954 fn collect_h5md_lattice_vectors_win() {
955 let dir = tempfile::tempdir().unwrap();
956 let db = ConCorpus::open(dir.path()).unwrap();
957 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
958 let mut frames = Vec::new();
959 for item in readcon_core::iterators::ConFrameIterator::new(&text) {
960 frames.push(item.unwrap());
961 }
962 frames[0].header.metadata.insert(
963 "lattice_vectors".into(),
964 serde_json::json!([[2.0, 0.0, 0.0], [0.5, 2.0, 0.0], [0.0, 0.0, 3.0]]),
965 );
966 db.append_trajectory_frames(1, &frames, "t").unwrap();
967 let a = db.collect_h5md(1).unwrap();
968 assert!((a.edges[0] - 2.0).abs() < 1e-12);
969 assert!((a.edges[3] - 0.5).abs() < 1e-12);
970 assert!((a.edges[8] - 3.0).abs() < 1e-12);
971 }
972
973 #[test]
974 fn collect_h5md_writes_velocities() {
975 let dir = tempfile::tempdir().unwrap();
976 let db = ConCorpus::open(dir.path()).unwrap();
977 db.append_trajectory_path(1, fixture("tiny_cuh2.convel"))
978 .unwrap();
979 let a = db.collect_h5md(1).unwrap();
980 let v = a.velocities.expect("velocities");
981 assert_eq!(v.len(), a.n_frames * a.natoms * 3);
982 assert!(v.iter().any(|&x| x != 0.0));
983 assert_eq!(a.velocity_unit, "Angstrom ps-1");
984 assert!(
985 (v[0] - 1.234).abs() < 1e-9,
986 "0.001234 A/fs -> 1.234 A/ps, got {}",
987 v[0]
988 );
989 }
990
991 #[test]
992 fn collect_h5md_pads_mixed_velocities() {
993 let dir = tempfile::tempdir().unwrap();
994 let db = ConCorpus::open(dir.path()).unwrap();
995 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
996 .unwrap();
997 db.extend_trajectory_path(1, fixture("tiny_cuh2.convel"))
998 .unwrap();
999 let a = db.collect_h5md(1).unwrap();
1000 let v = a.velocities.expect("velocities");
1001 assert!(a.n_frames >= 2);
1002 assert!(v[..a.natoms * 3].iter().all(|&x| x == 0.0));
1003 assert!(v[a.natoms * 3..].iter().any(|&x| x != 0.0));
1004 }
1005
1006 #[test]
1007 fn extend_trajectory_path_units_stamps() {
1008 let dir = tempfile::tempdir().unwrap();
1009 let db = ConCorpus::open(dir.path()).unwrap();
1010 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
1011 .unwrap();
1012 db.extend_trajectory_path_units(
1013 1,
1014 fixture("tiny_cuh2_forces.con"),
1015 Some(serde_json::json!({"length":"A","energy":"ev"})),
1016 )
1017 .unwrap();
1018 let u = db
1019 .frame_units(crate::keys::FrameKey {
1020 traj_id: 1,
1021 frame_idx: 1,
1022 })
1023 .unwrap()
1024 .unwrap();
1025 assert_eq!(u["length"], "angstrom");
1026 assert_eq!(u["energy"], "eV");
1027 }
1028}