1use std::collections::BTreeSet;
2use std::fs;
3use std::io::Cursor;
4use std::path::{Path, PathBuf};
5use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
6use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
7
8use heed::types::{Bytes, Str, Unit};
9use heed::{CompactionOption, Database, Env, EnvFlags, EnvOpenOptions, RwTxn};
10use readcon_core::iterators::ConFrameIterator;
11use readcon_core::types::ConFrame;
12use readcon_core::writer::{ConFrameWriter, FloatFormat};
13use serde::{Deserialize, Serialize};
14
15use crate::error::{Error, Result};
16use crate::export_xyz::write_frame_extxyz;
17use crate::frame_scalars::{
18 frame_cell_volume, frame_charge, frame_frame_index, frame_magmom, frame_neb_band,
19 frame_neb_bead, frame_pbc_mask, frame_time, frame_timestep, frame_total_mass,
20};
21use crate::keys::{
22 composition_formula, elem_count_key, elem_count_symbol_prefix, energy_bin_key, flag_key,
23 fmax_bin_key, formula_key, formula_prefix, hash_frame_bytes, mass_bin_key, meta_scalar_key,
24 natoms_key, ordered_f64_bits, parse_elem_count_key, pbc_key, pbc_mask_from_bools, symbol_key,
25 symbol_prefix, topo_key, topo_prefix, volume_bin_key, ContentHash, FrameKey, TrajId,
26 FLAG_HAS_ENERGY, FLAG_HAS_FORCES, FLAG_HAS_VELOCITIES, META_CHARGE, META_FRAME_INDEX,
27 META_MAGMOM, META_NEB_BAND, META_NEB_BEAD, META_TIME, META_TIMESTEP,
28};
29use crate::select::Select;
30use crate::topology::{
31 fingerprint_con_bytes, mixed_topo_error, normalize_topo_hex, resolve_seams_binary,
32 run_seams_fingerprint, AnnotateTopologyOpts, TopologyParams,
33};
34
35pub const DEFAULT_MAP_SIZE: usize = 2 * 1024 * 1024 * 1024;
39pub const MAP_SIZE_ENV: &str = "READCON_DB_MAP_SIZE";
42const MAP_ALIGN: usize = 1 << 20;
45const MAP_GROW_RETRIES: u32 = 8;
48
49pub fn parse_map_size(s: &str) -> Option<usize> {
51 let s = s.trim();
52 let (num, mul) = match s.as_bytes().last()? {
53 b'k' | b'K' => (&s[..s.len() - 1], 1usize << 10),
54 b'm' | b'M' => (&s[..s.len() - 1], 1 << 20),
55 b'g' | b'G' => (&s[..s.len() - 1], 1 << 30),
56 b't' | b'T' => (&s[..s.len() - 1], 1 << 40),
57 _ => (s, 1),
58 };
59 num.trim()
60 .parse::<usize>()
61 .ok()?
62 .checked_mul(mul)
63 .filter(|&n| n > 0)
64}
65
66fn align_map(n: usize) -> usize {
67 n.div_ceil(MAP_ALIGN).saturating_mul(MAP_ALIGN)
68}
69
70fn initial_map_size(path: &Path, requested: Option<usize>) -> Result<usize> {
74 let base = match requested {
75 Some(n) => n,
76 None => match std::env::var(MAP_SIZE_ENV) {
77 Ok(v) => parse_map_size(&v).ok_or_else(|| {
78 Error::Message(format!("{MAP_SIZE_ENV}={v:?} is not a byte count"))
79 })?,
80 Err(_) => DEFAULT_MAP_SIZE,
81 },
82 };
83 let used = fs::metadata(path.join("data.mdb"))
84 .map(|m| m.len() as usize)
85 .unwrap_or(0);
86 Ok(align_map(base.max(used.saturating_mul(2))))
87}
88const MAX_DBS: u32 = 48;
89const MAX_READERS: u32 = 512;
93
94#[derive(Clone, Debug, Default, Serialize, Deserialize)]
95pub struct TrajMeta {
96 pub n_frames: u32,
97 pub source: String,
98 #[serde(default)]
99 pub topo_cutoff: Option<f64>,
100 #[serde(default)]
101 pub topo_graph: Option<String>,
102 #[serde(default)]
103 pub topo_hops: Option<u32>,
104 #[serde(default)]
105 pub topo_method: Option<String>,
106}
107
108impl TrajMeta {
109 fn topo_params(&self) -> TopologyParams {
110 TopologyParams {
111 cutoff: self.topo_cutoff,
112 graph: self.topo_graph.clone(),
113 hops: self.topo_hops,
114 method: self.topo_method.clone(),
115 }
116 }
117
118 fn has_topo(&self) -> bool {
119 !self.topo_params().is_empty()
120 }
121
122 fn stamp_topo(&mut self, cutoff: f64, graph: &str, hops: u32, method: &str) {
123 self.topo_cutoff = Some(cutoff);
124 self.topo_graph = Some(graph.to_owned());
125 self.topo_hops = Some(hops);
126 self.topo_method = Some(method.to_owned());
127 }
128}
129
130struct PreparedIndexPuts {
132 fk: FrameKey,
133 stored: Vec<u8>,
135 hash: [u8; 16],
137 natoms_k: [u8; 16],
138 symbol_keys: Vec<Vec<u8>>,
139 elem_count_keys: Vec<Vec<u8>>,
140 formula_k: Option<Vec<u8>>,
141 energy_k: Option<[u8; 20]>,
142 fmax_k: Option<[u8; 20]>,
143 mass_k: Option<[u8; 20]>,
144 volume_k: Option<[u8; 20]>,
145 pbc_k: Option<[u8; 13]>,
146 flag_keys: Vec<[u8; 13]>,
147 meta_keys: Vec<[u8; 21]>,
148}
149
150pub struct ConCorpus {
151 path: PathBuf,
152 env: Env,
153 txn_gate: RwLock<()>,
158 frames: Database<Bytes, Bytes>,
159 compress_frames: AtomicBool,
161 traj_meta: Database<Bytes, Str>,
162 idx_natoms: Database<Bytes, Unit>,
163 idx_symbol: Database<Bytes, Unit>,
164 idx_energy: Database<Bytes, Unit>,
165 idx_flags: Database<Bytes, Unit>,
166 idx_elem_count: Database<Bytes, Unit>,
168 idx_formula: Database<Bytes, Unit>,
170 idx_fmax: Database<Bytes, Unit>,
172 idx_mass: Database<Bytes, Unit>,
173 idx_volume: Database<Bytes, Unit>,
174 idx_pbc: Database<Bytes, Unit>,
176 idx_meta: Database<Bytes, Unit>,
178 frame_by_hash: Database<Bytes, Bytes>,
179 hash_by_frame: Database<Bytes, Bytes>,
180 frames_soa: Database<Bytes, Bytes>,
182 idx_topo: Option<Database<Bytes, Unit>>,
184 topo_by_frame: Option<Database<Bytes, Bytes>>,
186}
187
188fn frame_has_forces(frame: &ConFrame) -> bool {
189 readcon_core::index_proj::frame_has_forces(frame)
190}
191
192fn frame_has_velocities(frame: &ConFrame) -> bool {
193 readcon_core::index_proj::frame_has_velocities(frame)
194}
195
196fn frame_energy(frame: &ConFrame) -> Option<f64> {
197 readcon_core::index_proj::finite_energy(frame)
198}
199
200pub fn frame_fmax(frame: &ConFrame) -> Option<f64> {
202 readcon_core::index_proj::frame_fmax(frame)
203}
204
205fn frame_species(frame: &ConFrame) -> Vec<(String, u32)> {
206 readcon_core::index_proj::frame_species_counts(frame)
207}
208
209fn ingest_chunk_bounds(n: usize) -> Vec<(usize, usize)> {
211 if n == 0 {
212 return Vec::new();
213 }
214 let chunk = (n / 8).max(1);
215 let mut bounds = Vec::new();
216 let mut i = 0;
217 while i < n {
218 let j = (i + chunk).min(n);
219 bounds.push((i, j));
220 i = j;
221 }
222 bounds
223}
224
225fn ingest_thread_count() -> usize {
226 match std::env::var("RAYON_NUM_THREADS") {
227 Ok(raw) => raw.parse::<usize>().unwrap_or(1).max(1),
228 Err(_) => std::thread::available_parallelism()
229 .map(|n| n.get())
230 .unwrap_or(1),
231 }
232}
233
234struct HeldRead<'a> {
237 txn: heed::RoTxn<'a>,
238 _guard: RwLockReadGuard<'a, ()>,
239}
240
241impl<'a> std::ops::Deref for HeldRead<'a> {
242 type Target = heed::RoTxn<'a>;
243
244 fn deref(&self) -> &Self::Target {
245 &self.txn
246 }
247}
248
249fn heed_map_full(err: &heed::Error) -> bool {
250 matches!(err, heed::Error::Mdb(heed::MdbError::MapFull))
251}
252
253fn is_map_full(err: &Error) -> bool {
254 matches!(err, Error::Heed(e) if heed_map_full(e))
255}
256
257fn prepared_bytes(prepared: &[PreparedIndexPuts]) -> usize {
258 prepared
259 .iter()
260 .fold(0usize, |n, p| n.saturating_add(p.stored.len()))
261}
262
263impl ConCorpus {
264 pub fn open(path: impl AsRef<Path>) -> Result<Self> {
265 Self::open_with(path, false, None)
266 }
267
268 pub fn open_existing(path: impl AsRef<Path>) -> Result<Self> {
270 let path = path.as_ref();
271 if !path.join("data.mdb").is_file() {
272 return Err(Error::Message(format!(
273 "not a corpus directory (missing data.mdb): {}",
274 path.display()
275 )));
276 }
277 Self::open_with(path, false, None)
278 }
279
280 pub fn open_with_map_size(path: impl AsRef<Path>, map_size: usize) -> Result<Self> {
284 Self::open_with(path, false, Some(map_size))
285 }
286
287 pub fn map_size(&self) -> usize {
289 self.env.info().map_size
290 }
291
292 fn read_guard(&self) -> RwLockReadGuard<'_, ()> {
293 self.txn_gate.read().unwrap_or_else(|e| e.into_inner())
294 }
295
296 fn write_guard(&self) -> RwLockWriteGuard<'_, ()> {
297 self.txn_gate.write().unwrap_or_else(|e| e.into_inner())
298 }
299
300 fn read_txn(&self) -> Result<HeldRead<'_>> {
303 let guard = self.read_guard();
304 let txn = self.env.read_txn()?;
305 Ok(HeldRead { txn, _guard: guard })
306 }
307
308 fn used_bytes(&self) -> usize {
309 fs::metadata(self.path.join("data.mdb"))
310 .map(|m| m.len() as usize)
311 .unwrap_or(0)
312 }
313
314 fn grow_after_map_full(&self, input_bytes: usize) -> Result<()> {
317 let cur = self.map_size();
318 let used_plus = self.used_bytes().saturating_add(input_bytes);
319 let mut want = align_map(cur.saturating_mul(2).max(used_plus));
320 if want <= cur {
321 want = align_map(cur.saturating_add(MAP_ALIGN));
322 }
323 unsafe { self.env.resize(want)? };
326 Ok(())
327 }
328
329 fn with_write<T>(
333 &self,
334 input_bytes: usize,
335 mut body: impl FnMut(&mut RwTxn<'_>) -> Result<T>,
336 ) -> Result<T> {
337 let _guard = self.write_guard();
338 let need = self
342 .used_bytes()
343 .saturating_add(input_bytes.saturating_mul(2));
344 let cur = self.map_size();
345 if need > cur {
346 let want = align_map(need.max(cur.saturating_mul(2)));
347 unsafe { self.env.resize(want)? };
349 }
350 let mut last = Error::Message("LMDB map stayed full after growing it".into());
351 for _ in 0..MAP_GROW_RETRIES {
352 let mut wtxn = match self.env.write_txn() {
353 Ok(txn) => txn,
354 Err(e) if heed_map_full(&e) => {
355 last = Error::Heed(e);
356 self.grow_after_map_full(input_bytes)?;
357 continue;
358 }
359 Err(e) => return Err(e.into()),
360 };
361 match body(&mut wtxn) {
362 Ok(value) => match wtxn.commit() {
363 Ok(()) => return Ok(value),
364 Err(e) if heed_map_full(&e) => {
367 last = Error::Heed(e);
368 self.grow_after_map_full(input_bytes)?;
369 }
370 Err(e) => return Err(e.into()),
371 },
372 Err(e) if is_map_full(&e) => {
373 drop(wtxn);
374 last = e;
375 self.grow_after_map_full(input_bytes)?;
376 }
377 Err(e) => return Err(e),
378 }
379 }
380 Err(last)
381 }
382
383 pub fn reserve_map(&self, extra_bytes: usize) -> Result<usize> {
390 let _guard = self.write_guard();
391 let want = align_map(self.used_bytes().saturating_add(extra_bytes));
392 if want > self.map_size() {
393 unsafe { self.env.resize(want)? };
395 }
396 Ok(self.map_size())
397 }
398
399 pub fn open_readonly(path: impl AsRef<Path>) -> Result<Self> {
405 Self::open_with(path, true, None)
406 }
407
408 fn open_with(path: impl AsRef<Path>, readonly: bool, map_size: Option<usize>) -> Result<Self> {
409 let path = path.as_ref().to_path_buf();
410 if path.join("shards.json").is_file() {
411 return Err(Error::Message(format!(
412 "refusing to open a sharded root as a single-env corpus: {}",
413 path.display()
414 )));
415 }
416 if readonly {
417 if !path.join("data.mdb").is_file() {
418 return Err(Error::Message(format!(
419 "readonly open: {} is not a corpus directory",
420 path.display()
421 )));
422 }
423 } else {
424 fs::create_dir_all(&path)?;
425 }
426 let mut opts = EnvOpenOptions::new();
427 opts.map_size(initial_map_size(&path, map_size)?);
428 opts.max_dbs(MAX_DBS);
429 opts.max_readers(MAX_READERS);
430 if readonly {
431 unsafe {
433 opts.flags(EnvFlags::READ_ONLY);
434 }
435 }
436 let env = unsafe { opts.open(&path)? };
437
438 if readonly {
439 let rtxn = env.read_txn()?;
440 let missing =
441 |name: &str| Error::Message(format!("readonly open: missing database {name}"));
442 let frames = env
443 .open_database(&rtxn, Some("frames"))?
444 .ok_or_else(|| missing("frames"))?;
445 let traj_meta = env
446 .open_database(&rtxn, Some("traj_meta"))?
447 .ok_or_else(|| missing("traj_meta"))?;
448 let idx_natoms = env
449 .open_database(&rtxn, Some("idx_natoms"))?
450 .ok_or_else(|| missing("idx_natoms"))?;
451 let idx_symbol = env
452 .open_database(&rtxn, Some("idx_symbol"))?
453 .ok_or_else(|| missing("idx_symbol"))?;
454 let idx_energy = env
455 .open_database(&rtxn, Some("idx_energy"))?
456 .ok_or_else(|| missing("idx_energy"))?;
457 let idx_flags = env
458 .open_database(&rtxn, Some("idx_flags"))?
459 .ok_or_else(|| missing("idx_flags"))?;
460 let idx_elem_count = env
461 .open_database(&rtxn, Some("idx_elem_count"))?
462 .ok_or_else(|| missing("idx_elem_count"))?;
463 let idx_formula = env
464 .open_database(&rtxn, Some("idx_formula"))?
465 .ok_or_else(|| missing("idx_formula"))?;
466 let idx_fmax = env
467 .open_database(&rtxn, Some("idx_fmax"))?
468 .ok_or_else(|| missing("idx_fmax"))?;
469 let idx_mass = env
470 .open_database(&rtxn, Some("idx_mass"))?
471 .ok_or_else(|| missing("idx_mass"))?;
472 let idx_volume = env
473 .open_database(&rtxn, Some("idx_volume"))?
474 .ok_or_else(|| missing("idx_volume"))?;
475 let idx_pbc = env
476 .open_database(&rtxn, Some("idx_pbc"))?
477 .ok_or_else(|| missing("idx_pbc"))?;
478 let idx_meta = env
479 .open_database(&rtxn, Some("idx_meta"))?
480 .ok_or_else(|| missing("idx_meta"))?;
481 let frame_by_hash = env
482 .open_database(&rtxn, Some("frame_by_hash"))?
483 .ok_or_else(|| missing("frame_by_hash"))?;
484 let hash_by_frame = env
485 .open_database(&rtxn, Some("hash_by_frame"))?
486 .ok_or_else(|| missing("hash_by_frame"))?;
487 let frames_soa = env
488 .open_database(&rtxn, Some("frames_soa"))?
489 .ok_or_else(|| missing("frames_soa"))?;
490 let idx_topo = env.open_database(&rtxn, Some("idx_topo"))?;
492 let topo_by_frame = env.open_database(&rtxn, Some("topo_by_frame"))?;
493 rtxn.commit()?;
495 return Ok(Self {
496 path,
497 env,
498 txn_gate: RwLock::new(()),
499 frames,
500 compress_frames: AtomicBool::new(true),
501 traj_meta,
502 idx_natoms,
503 idx_symbol,
504 idx_energy,
505 idx_flags,
506 idx_elem_count,
507 idx_formula,
508 idx_fmax,
509 idx_mass,
510 idx_volume,
511 idx_pbc,
512 idx_meta,
513 frame_by_hash,
514 hash_by_frame,
515 frames_soa,
516 idx_topo,
517 topo_by_frame,
518 });
519 }
520
521 let mut wtxn = env.write_txn()?;
522 let frames = env.create_database(&mut wtxn, Some("frames"))?;
523 let traj_meta = env.create_database(&mut wtxn, Some("traj_meta"))?;
524 let idx_natoms = env.create_database(&mut wtxn, Some("idx_natoms"))?;
525 let idx_symbol = env.create_database(&mut wtxn, Some("idx_symbol"))?;
526 let idx_energy = env.create_database(&mut wtxn, Some("idx_energy"))?;
527 let idx_flags = env.create_database(&mut wtxn, Some("idx_flags"))?;
528 let idx_elem_count = env.create_database(&mut wtxn, Some("idx_elem_count"))?;
529 let idx_formula = env.create_database(&mut wtxn, Some("idx_formula"))?;
530 let idx_fmax = env.create_database(&mut wtxn, Some("idx_fmax"))?;
531 let idx_mass = env.create_database(&mut wtxn, Some("idx_mass"))?;
532 let idx_volume = env.create_database(&mut wtxn, Some("idx_volume"))?;
533 let idx_pbc = env.create_database(&mut wtxn, Some("idx_pbc"))?;
534 let idx_meta = env.create_database(&mut wtxn, Some("idx_meta"))?;
535 let frame_by_hash = env.create_database(&mut wtxn, Some("frame_by_hash"))?;
536 let hash_by_frame = env.create_database(&mut wtxn, Some("hash_by_frame"))?;
537 let frames_soa = env.create_database(&mut wtxn, Some("frames_soa"))?;
538 let idx_topo = env.create_database(&mut wtxn, Some("idx_topo"))?;
539 let topo_by_frame = env.create_database(&mut wtxn, Some("topo_by_frame"))?;
540 wtxn.commit()?;
541
542 Ok(Self {
543 path,
544 env,
545 txn_gate: RwLock::new(()),
546 frames,
547 compress_frames: AtomicBool::new(true),
548 traj_meta,
549 idx_natoms,
550 idx_symbol,
551 idx_energy,
552 idx_flags,
553 idx_elem_count,
554 idx_formula,
555 idx_fmax,
556 idx_mass,
557 idx_volume,
558 idx_pbc,
559 idx_meta,
560 frame_by_hash,
561 hash_by_frame,
562 frames_soa,
563 idx_topo: Some(idx_topo),
564 topo_by_frame: Some(topo_by_frame),
565 })
566 }
567
568 pub fn pack_frame(&self, key: FrameKey) -> Result<Vec<u8>> {
572 if let Some(b) = self.get_cooked_soa_bytes(key)? {
573 if crate::cooked_soa::CookedSoa::try_decode(&b).is_some() {
574 return Ok(b);
575 }
576 }
577 let fr = self.get_frame(key)?;
578 crate::cooked_soa::CookedSoa::encode_frame(&fr)
579 }
580
581 pub fn pack_frames(&self, keys: &[FrameKey]) -> Result<Vec<u8>> {
583 let mut blobs = Vec::with_capacity(keys.len());
584 for k in keys {
585 blobs.push(self.pack_frame(*k)?);
586 }
587 crate::cooked_soa::encode_batch(&blobs)
588 }
589
590 pub fn path(&self) -> &Path {
591 &self.path
592 }
593
594 pub fn close(self) {
597 self.env.prepare_for_closing().wait();
598 }
599
600 pub fn snapshot_to(&self, dest_dir: impl AsRef<Path>) -> Result<()> {
603 let dest = dest_dir.as_ref();
604 fs::create_dir_all(dest)?;
605 let data = dest.join("data.mdb");
606 if data.exists() {
607 return Err(Error::Message(format!(
608 "snapshot dest exists: {}",
609 data.display()
610 )));
611 }
612 self.env.copy_to_file(&data, CompactionOption::Enabled)?;
613 Ok(())
614 }
615
616 pub fn list_frame_keys(&self) -> Result<Vec<FrameKey>> {
618 let rtxn = self.read_txn()?;
619 let mut out = Vec::new();
620 let mut iter = self.frames.iter(&rtxn)?;
621 while let Some(Ok((k, _))) = iter.next() {
622 if let Some(fk) = FrameKey::from_bytes(k) {
623 out.push(fk);
624 }
625 }
626 out.sort();
627 Ok(out)
628 }
629
630 pub fn export_frame_blob(&self, key: FrameKey) -> Result<String> {
633 self.get_frame_text(key)
634 }
635
636 pub fn set_compress_frames(&self, compress: bool) {
638 self.compress_frames.store(compress, Ordering::Relaxed);
639 }
640
641 fn compress_on(&self) -> bool {
642 self.compress_frames.load(Ordering::Relaxed)
643 }
644
645 fn build_index_puts(
650 fk: FrameKey,
651 frame: &ConFrame,
652 text: String,
653 compress: bool,
654 ) -> Result<PreparedIndexPuts> {
655 use readcon_core::index_proj::FrameIndexProjection;
656 let hash = hash_frame_bytes(text.as_bytes()).to_bytes();
657 let stored = crate::frame_codec::store_frame(&text, compress)?;
658 let proj = FrameIndexProjection::from_frame(frame);
659 let mut symbol_keys = Vec::new();
660 let mut elem_count_keys = Vec::new();
661 let mut seen = BTreeSet::new();
662 for (sym, cnt) in &proj.species_counts {
663 elem_count_keys.push(elem_count_key(sym, *cnt, fk));
664 if seen.insert(sym.clone()) {
665 symbol_keys.push(symbol_key(sym, fk));
666 }
667 }
668 let formula_k = if proj.formula.is_empty() {
669 None
670 } else {
671 Some(formula_key(&proj.formula, fk))
672 };
673 let mut flag_keys = Vec::new();
674 let energy_k = proj.energy.and_then(|e| {
675 flag_keys.push(flag_key(FLAG_HAS_ENERGY, fk));
676 energy_bin_key(e, fk)
677 });
678 let fmax_k = if proj.has_forces {
679 flag_keys.push(flag_key(FLAG_HAS_FORCES, fk));
680 proj.fmax.and_then(|fm| fmax_bin_key(fm, fk))
681 } else {
682 None
683 };
684 if proj.has_velocities {
685 flag_keys.push(flag_key(FLAG_HAS_VELOCITIES, fk));
686 }
687 let mass_k = proj.total_mass.and_then(|m| mass_bin_key(m, fk));
688 let volume_k = proj.cell_volume.and_then(|v| volume_bin_key(v, fk));
689 let pbc_k = proj
690 .pbc
691 .map(|p| pbc_key(crate::keys::pbc_mask_from_bools(p), fk));
692 let mut meta_keys = Vec::new();
693 for (ch, val) in [
694 (META_TIME, proj.time),
695 (META_TIMESTEP, proj.timestep),
696 (META_FRAME_INDEX, proj.frame_index),
697 (META_NEB_BEAD, proj.neb_bead),
698 (META_NEB_BAND, proj.neb_band),
699 (META_CHARGE, proj.charge),
700 (META_MAGMOM, proj.magmom),
701 ] {
702 if let Some(x) = val {
703 if let Some(k) = meta_scalar_key(ch, x, fk) {
704 meta_keys.push(k);
705 }
706 }
707 }
708 Ok(PreparedIndexPuts {
709 fk,
710 stored,
711 hash,
712 natoms_k: natoms_key(proj.n_atoms, fk),
713 symbol_keys,
714 elem_count_keys,
715 formula_k,
716 energy_k,
717 fmax_k,
718 mass_k,
719 volume_k,
720 pbc_k,
721 flag_keys,
722 meta_keys,
723 })
724 }
725
726 fn put_index_puts(&self, wtxn: &mut RwTxn, p: &PreparedIndexPuts) -> Result<()> {
728 let fk_b = p.fk.to_bytes();
729 self.frames.put(wtxn, &fk_b[..], &p.stored[..])?;
730 self.hash_by_frame.put(wtxn, &fk_b[..], &p.hash[..])?;
731 if self.frame_by_hash.get(wtxn, &p.hash[..])?.is_none() {
732 self.frame_by_hash.put(wtxn, &p.hash[..], &fk_b[..])?;
733 }
734 self.idx_natoms.put(wtxn, &p.natoms_k[..], &())?;
735 for sk in &p.symbol_keys {
736 self.idx_symbol.put(wtxn, &sk[..], &())?;
737 }
738 for ek in &p.elem_count_keys {
739 self.idx_elem_count.put(wtxn, &ek[..], &())?;
740 }
741 if let Some(ref fk_form) = p.formula_k {
742 self.idx_formula.put(wtxn, &fk_form[..], &())?;
743 }
744 if let Some(ek) = p.energy_k {
745 self.idx_energy.put(wtxn, &ek[..], &())?;
746 }
747 if let Some(fk_fm) = p.fmax_k {
748 self.idx_fmax.put(wtxn, &fk_fm[..], &())?;
749 }
750 if let Some(k) = p.mass_k {
751 self.idx_mass.put(wtxn, &k[..], &())?;
752 }
753 if let Some(k) = p.volume_k {
754 self.idx_volume.put(wtxn, &k[..], &())?;
755 }
756 if let Some(pk) = p.pbc_k {
757 self.idx_pbc.put(wtxn, &pk[..], &())?;
758 }
759 for fk_flag in &p.flag_keys {
760 self.idx_flags.put(wtxn, &fk_flag[..], &())?;
761 }
762 for mk in &p.meta_keys {
763 self.idx_meta.put(wtxn, &mk[..], &())?;
764 }
765 Ok(())
766 }
767
768 fn delete_index_puts(&self, wtxn: &mut RwTxn, p: &PreparedIndexPuts) -> Result<()> {
769 let fk_b = p.fk.to_bytes();
770 let _ = self.frames.delete(wtxn, &fk_b[..])?;
771 let _ = self.hash_by_frame.delete(wtxn, &fk_b[..])?;
772 let _ = self.frame_by_hash.delete(wtxn, &p.hash[..])?;
773 let _ = self.frames_soa.delete(wtxn, &fk_b[..])?;
774 let _ = self.idx_natoms.delete(wtxn, &p.natoms_k[..])?;
775 for sk in &p.symbol_keys {
776 let _ = self.idx_symbol.delete(wtxn, &sk[..])?;
777 }
778 for ek in &p.elem_count_keys {
779 let _ = self.idx_elem_count.delete(wtxn, &ek[..])?;
780 }
781 if let Some(ref fk_form) = p.formula_k {
782 let _ = self.idx_formula.delete(wtxn, &fk_form[..])?;
783 }
784 if let Some(ek) = p.energy_k {
785 let _ = self.idx_energy.delete(wtxn, &ek[..])?;
786 }
787 if let Some(fk_fm) = p.fmax_k {
788 let _ = self.idx_fmax.delete(wtxn, &fk_fm[..])?;
789 }
790 if let Some(k) = p.mass_k {
791 let _ = self.idx_mass.delete(wtxn, &k[..])?;
792 }
793 if let Some(k) = p.volume_k {
794 let _ = self.idx_volume.delete(wtxn, &k[..])?;
795 }
796 if let Some(pk) = p.pbc_k {
797 let _ = self.idx_pbc.delete(wtxn, &pk[..])?;
798 }
799 for fk_flag in &p.flag_keys {
800 let _ = self.idx_flags.delete(wtxn, &fk_flag[..])?;
801 }
802 for mk in &p.meta_keys {
803 let _ = self.idx_meta.delete(wtxn, &mk[..])?;
804 }
805 Ok(())
806 }
807
808 fn index_frame(
810 &self,
811 wtxn: &mut RwTxn,
812 fk: FrameKey,
813 frame: &ConFrame,
814 blob: &str,
815 ) -> Result<()> {
816 let puts = Self::build_index_puts(fk, frame, blob.to_owned(), self.compress_on())?;
817 self.put_index_puts(wtxn, &puts)
818 }
819
820 fn prepare_trajectory_str(
822 traj_id: TrajId,
823 file_contents: &str,
824 compress: bool,
825 ) -> Result<Vec<PreparedIndexPuts>> {
826 let mut out = Vec::new();
827 let mut frame_idx: u32 = 0;
828 let mut iter = ConFrameIterator::new(file_contents);
829 loop {
830 let (frame, blob) = match iter.next_with_raw_span(file_contents) {
831 None => break,
832 Some(Err(e)) => return Err(Error::Parse(e.to_string())),
833 Some(Ok(x)) => x,
834 };
835 let fk = FrameKey { traj_id, frame_idx };
836 out.push(Self::build_index_puts(
837 fk,
838 &frame,
839 blob.to_owned(),
840 compress,
841 )?);
842 frame_idx += 1;
843 }
844 Ok(out)
845 }
846
847 fn prepare_trajectory_frames(
849 traj_id: TrajId,
850 frames: &[ConFrame],
851 start_idx: u32,
852 compress: bool,
853 ) -> Result<Vec<PreparedIndexPuts>> {
854 Self::prepare_trajectory_frames_formatted(
855 traj_id,
856 frames,
857 start_idx,
858 FloatFormat::default(),
859 compress,
860 )
861 }
862
863 fn prepare_trajectory_frames_formatted(
864 traj_id: TrajId,
865 frames: &[ConFrame],
866 start_idx: u32,
867 float_format: FloatFormat,
868 compress: bool,
869 ) -> Result<Vec<PreparedIndexPuts>> {
870 let mut out = Vec::with_capacity(frames.len());
871 let mut frame_idx = start_idx;
872 for fr in frames {
873 let mut buf = Cursor::new(Vec::new());
874 {
875 let mut w = ConFrameWriter::with_float_format(&mut buf, float_format);
876 w.write_frame(fr)
877 .map_err(|e| Error::Parse(format!("serialize: {e}")))?;
878 }
879 let blob = String::from_utf8(buf.into_inner())
880 .map_err(|e| Error::Message(format!("utf8: {e}")))?;
881 let fk = FrameKey { traj_id, frame_idx };
882 out.push(Self::build_index_puts(fk, fr, blob, compress)?);
883 frame_idx += 1;
884 }
885 Ok(out)
886 }
887
888 #[cfg(feature = "mpi")]
890 pub(crate) fn check_traj_free(&self, traj_id: TrajId) -> Result<()> {
891 let rtxn = self.read_txn()?;
892 let tid_key = traj_id.to_be_bytes();
893 if let Some(existing) = self.traj_meta.get(&rtxn, &tid_key[..])? {
894 let meta: TrajMeta = serde_json::from_str(existing)?;
895 return Err(Error::TrajExists(traj_id, meta.n_frames));
896 }
897 Ok(())
898 }
899
900 #[cfg(feature = "mpi")]
903 pub(crate) fn encode_frame_keys(
904 fk: FrameKey,
905 frame: &ConFrame,
906 blob: &str,
907 ) -> Result<([u8; 16], Vec<u8>, Vec<u8>)> {
908 let p = Self::build_index_puts(fk, frame, blob.to_owned(), true)?;
909 let mut v = Vec::with_capacity(160);
910 fn put(v: &mut Vec<u8>, b: &[u8]) {
911 v.extend_from_slice(&(b.len() as u32).to_le_bytes());
912 v.extend_from_slice(b);
913 }
914 fn list<T: AsRef<[u8]>>(v: &mut Vec<u8>, xs: &[T]) {
915 v.extend_from_slice(&(xs.len() as u32).to_le_bytes());
916 for x in xs {
917 put(v, x.as_ref());
918 }
919 }
920 v.extend_from_slice(&p.hash);
921 v.extend_from_slice(&p.natoms_k);
922 list(&mut v, &p.symbol_keys);
923 list(&mut v, &p.elem_count_keys);
924 list(&mut v, p.formula_k.as_slice());
925 list(&mut v, p.energy_k.as_slice());
926 list(&mut v, p.fmax_k.as_slice());
927 list(&mut v, p.mass_k.as_slice());
928 list(&mut v, p.volume_k.as_slice());
929 list(&mut v, p.pbc_k.as_slice());
930 list(&mut v, &p.flag_keys);
931 list(&mut v, &p.meta_keys);
932 Ok((p.hash, v, p.stored))
933 }
934
935 #[cfg(feature = "mpi")]
938 pub(crate) fn put_encoded(&self, rows: Vec<(FrameKey, Vec<u8>, Vec<u8>)>) -> Result<()> {
939 struct Rd<'a>(&'a [u8]);
940 impl<'a> Rd<'a> {
941 fn take(&mut self, n: usize) -> Result<&'a [u8]> {
942 if self.0.len() < n {
943 return Err(Error::Message("truncated MPI frame record".into()));
944 }
945 let (a, b) = self.0.split_at(n);
946 self.0 = b;
947 Ok(a)
948 }
949 fn u32(&mut self) -> Result<usize> {
950 Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()) as usize)
951 }
952 fn list(&mut self) -> Result<Vec<Vec<u8>>> {
953 let n = self.u32()?;
954 (0..n)
955 .map(|_| {
956 let l = self.u32()?;
957 Ok(self.take(l)?.to_vec())
958 })
959 .collect()
960 }
961 fn fixed<const N: usize>(&mut self) -> Result<Vec<[u8; N]>> {
962 self.list()?
963 .into_iter()
964 .map(|x| {
965 x.try_into()
966 .map_err(|_| Error::Message("bad MPI key width".into()))
967 })
968 .collect()
969 }
970 }
971 let prepared = rows
974 .into_iter()
975 .map(|(fk, enc, stored)| {
976 let mut r = Rd(&enc);
977 let hash: [u8; 16] = r.take(16)?.try_into().unwrap();
978 let natoms_k: [u8; 16] = r.take(16)?.try_into().unwrap();
979 Ok(PreparedIndexPuts {
980 fk,
981 stored,
982 hash,
983 natoms_k,
984 symbol_keys: r.list()?,
985 elem_count_keys: r.list()?,
986 formula_k: r.list()?.pop(),
987 energy_k: r.fixed::<20>()?.pop(),
988 fmax_k: r.fixed::<20>()?.pop(),
989 mass_k: r.fixed::<20>()?.pop(),
990 volume_k: r.fixed::<20>()?.pop(),
991 pbc_k: r.fixed::<13>()?.pop(),
992 flag_keys: r.fixed::<13>()?,
993 meta_keys: r.fixed::<21>()?,
994 })
995 })
996 .collect::<Result<Vec<_>>>()?;
997 self.with_write(prepared_bytes(&prepared), |wtxn| {
998 for p in &prepared {
999 self.put_index_puts(wtxn, p)?;
1000 }
1001 Ok(())
1002 })
1003 }
1004
1005 #[cfg(feature = "mpi")]
1007 pub(crate) fn put_traj_meta(
1008 &self,
1009 traj_id: TrajId,
1010 n_frames: u32,
1011 source: String,
1012 ) -> Result<()> {
1013 let input_bytes = source.len().saturating_add(64);
1014 self.with_write(input_bytes, |wtxn| {
1015 let meta = TrajMeta {
1016 n_frames,
1017 source: source.clone(),
1018 ..Default::default()
1019 };
1020 self.traj_meta.put(
1021 wtxn,
1022 &traj_id.to_be_bytes()[..],
1023 &serde_json::to_string(&meta)?,
1024 )?;
1025 Ok(())
1026 })
1027 }
1028
1029 pub fn canonical_dump(&self) -> Result<Vec<u8>> {
1033 let rtxn = self.read_txn()?;
1034 let mut out = Vec::new();
1035 let mut tables: Vec<(&str, Database<Bytes, Bytes>)> = vec![
1036 ("frames", self.frames.remap_types()),
1037 ("traj_meta", self.traj_meta.remap_types()),
1038 ("idx_natoms", self.idx_natoms.remap_types()),
1039 ("idx_symbol", self.idx_symbol.remap_types()),
1040 ("idx_energy", self.idx_energy.remap_types()),
1041 ("idx_flags", self.idx_flags.remap_types()),
1042 ("idx_elem_count", self.idx_elem_count.remap_types()),
1043 ("idx_formula", self.idx_formula.remap_types()),
1044 ("idx_fmax", self.idx_fmax.remap_types()),
1045 ("idx_mass", self.idx_mass.remap_types()),
1046 ("idx_volume", self.idx_volume.remap_types()),
1047 ("idx_pbc", self.idx_pbc.remap_types()),
1048 ("idx_meta", self.idx_meta.remap_types()),
1049 ("frame_by_hash", self.frame_by_hash),
1050 ("hash_by_frame", self.hash_by_frame),
1051 ("frames_soa", self.frames_soa),
1052 ];
1053 if let Some(d) = self.idx_topo {
1054 tables.push(("idx_topo", d.remap_types()));
1055 }
1056 if let Some(d) = self.topo_by_frame {
1057 tables.push(("topo_by_frame", d));
1058 }
1059 for (name, db) in tables {
1060 for part in [name.as_bytes()] {
1061 out.extend_from_slice(&(part.len() as u64).to_le_bytes());
1062 out.extend_from_slice(part);
1063 }
1064 for kv in db.iter(&rtxn)? {
1065 let (k, v) = kv?;
1066 for part in [k, v] {
1067 out.extend_from_slice(&(part.len() as u64).to_le_bytes());
1068 out.extend_from_slice(part);
1069 }
1070 }
1071 }
1072 Ok(out)
1073 }
1074
1075 fn group_prepared(prepared: &mut [PreparedIndexPuts]) -> Result<()> {
1079 if prepared.len() < 2 {
1080 return Ok(());
1081 }
1082 let texts: Vec<(u32, String)> = prepared
1083 .iter()
1084 .map(|p| Ok((p.fk.frame_idx, crate::frame_codec::load_frame(&p.stored)?)))
1085 .collect::<Result<_>>()?;
1086 let view: Vec<(u32, &str)> = texts.iter().map(|(i, s)| (*i, s.as_str())).collect();
1087 let packed = crate::frame_codec::pack_aligned(&view)?;
1088 for (p, stored) in prepared.iter_mut().zip(packed) {
1089 p.stored = stored;
1090 }
1091 Ok(())
1092 }
1093
1094 fn commit_prepared(
1096 &self,
1097 traj_id: TrajId,
1098 prepared: &mut [PreparedIndexPuts],
1099 source: String,
1100 replace_meta: bool,
1101 compress: bool,
1102 ) -> Result<u32> {
1103 if compress {
1104 Self::group_prepared(prepared)?;
1105 }
1106 let input_bytes = prepared_bytes(prepared).saturating_add(source.len());
1107 self.with_write(input_bytes, |wtxn| {
1108 let tid_key = traj_id.to_be_bytes();
1109 if !replace_meta {
1110 if let Some(existing) = self.traj_meta.get(wtxn, &tid_key[..])? {
1111 let meta: TrajMeta = serde_json::from_str(existing)?;
1112 return Err(Error::TrajExists(traj_id, meta.n_frames));
1113 }
1114 }
1115 for p in prepared.iter() {
1116 self.put_index_puts(wtxn, p)?;
1117 }
1118 let n_frames = prepared.last().map(|p| p.fk.frame_idx + 1).unwrap_or(0);
1119 let existing: Option<TrajMeta> = if replace_meta {
1120 self.traj_meta
1121 .get(wtxn, &tid_key[..])?
1122 .map(serde_json::from_str)
1123 .transpose()?
1124 } else {
1125 None
1126 };
1127 let mut meta = existing.unwrap_or_default();
1128 meta.n_frames = n_frames;
1129 meta.source = source.clone();
1130 self.traj_meta
1131 .put(wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
1132 Ok(n_frames)
1133 })
1134 }
1135
1136 pub fn append_trajectory_path(&self, traj_id: TrajId, file: impl AsRef<Path>) -> Result<u32> {
1137 self.append_trajectory_path_units(traj_id, file, None)
1138 }
1139
1140 pub fn append_trajectory_path_units(
1142 &self,
1143 traj_id: TrajId,
1144 file: impl AsRef<Path>,
1145 units: Option<serde_json::Value>,
1146 ) -> Result<u32> {
1147 let text = fs::read_to_string(file.as_ref())?;
1148 let source = file.as_ref().display().to_string();
1149 let Some(units) = units else {
1150 return self.append_trajectory_str(traj_id, &text, source);
1151 };
1152 let units = crate::units_canon::canonicalize_units_object(units)?;
1153 let mut frames = Vec::new();
1154 for item in ConFrameIterator::new(&text) {
1155 let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
1156 fr.header.set_units(units.clone());
1157 frames.push(fr);
1158 }
1159 self.append_trajectory_frames(traj_id, &frames, source)
1160 }
1161
1162 pub fn set_trajectory_units(&self, traj_id: TrajId, units: serde_json::Value) -> Result<u32> {
1165 let keys = self.select(&Select::new().trajectory(traj_id))?;
1166 if keys.is_empty() {
1167 return Err(Error::Message("no frames for traj".into()));
1168 }
1169 let mut frames = Vec::with_capacity(keys.len());
1170 let mut old_puts = Vec::with_capacity(keys.len());
1171 for k in &keys {
1172 let blob = self.get_frame_text(*k)?;
1173 let mut fr = self.get_frame(*k)?;
1174 old_puts.push(Self::build_index_puts(*k, &fr, blob, self.compress_on())?);
1175 let merged = crate::units_canon::merge_units(fr.header.units(), units.clone())?;
1176 crate::units_canon::rescale_frame_units(&mut fr, &merged)?;
1177 fr.header.set_units(merged);
1178 frames.push(fr);
1179 }
1180 let mut prepared = Self::prepare_trajectory_frames_formatted(
1182 traj_id,
1183 &frames,
1184 0,
1185 FloatFormat::RoundTrip,
1186 self.compress_on(),
1187 )?;
1188 if self.compress_on() {
1189 Self::group_prepared(&mut prepared)?;
1190 }
1191 let input_bytes = prepared_bytes(&prepared);
1192 self.with_write(input_bytes, |wtxn| {
1193 for p in &old_puts {
1194 self.delete_index_puts(wtxn, p)?;
1195 }
1196 for p in &prepared {
1197 self.put_index_puts(wtxn, p)?;
1198 }
1199 Ok(keys.len() as u32)
1200 })
1201 }
1202
1203 pub fn frame_units(&self, key: FrameKey) -> Result<Option<serde_json::Value>> {
1205 let fr = self.get_frame(key)?;
1206 Ok(fr.header.metadata.get("units").cloned())
1207 }
1208
1209 pub fn append_trajectory_str(
1210 &self,
1211 traj_id: TrajId,
1212 file_contents: &str,
1213 source: impl Into<String>,
1214 ) -> Result<u32> {
1215 let source = source.into();
1216 let mut prepared =
1218 Self::prepare_trajectory_str(traj_id, file_contents, self.compress_on())?;
1219 self.commit_prepared(traj_id, &mut prepared, source, false, self.compress_on())
1220 }
1221
1222 pub fn append_trajectory_frames(
1224 &self,
1225 traj_id: TrajId,
1226 frames: &[ConFrame],
1227 source: impl Into<String>,
1228 ) -> Result<u32> {
1229 let source = source.into();
1230 let mut prepared = Self::prepare_trajectory_frames(traj_id, frames, 0, self.compress_on())?;
1231 self.commit_prepared(traj_id, &mut prepared, source, false, self.compress_on())
1232 }
1233
1234 pub fn append_trajectory_frames_with_precision(
1238 &self,
1239 traj_id: TrajId,
1240 frames: &[ConFrame],
1241 source: impl Into<String>,
1242 precision: usize,
1243 ) -> Result<u32> {
1244 self.append_trajectory_frames_with_float_format(
1245 traj_id,
1246 frames,
1247 source,
1248 FloatFormat::DecimalPlaces(precision),
1249 )
1250 }
1251
1252 pub fn append_trajectory_frames_with_float_format(
1254 &self,
1255 traj_id: TrajId,
1256 frames: &[ConFrame],
1257 source: impl Into<String>,
1258 float_format: FloatFormat,
1259 ) -> Result<u32> {
1260 let source = source.into();
1261 let mut prepared = Self::prepare_trajectory_frames_formatted(
1262 traj_id,
1263 frames,
1264 0,
1265 float_format,
1266 self.compress_on(),
1267 )?;
1268 self.commit_prepared(traj_id, &mut prepared, source, false, self.compress_on())
1269 }
1270
1271 pub fn extend_trajectory_frames(
1274 &self,
1275 traj_id: TrajId,
1276 frames: &[ConFrame],
1277 source_hint: impl Into<String>,
1278 ) -> Result<u32> {
1279 let (start_idx, prior) = {
1280 let rtxn = self.read_txn()?;
1281 let tid_key = traj_id.to_be_bytes();
1282 if let Some(existing) = self.traj_meta.get(&rtxn, &tid_key[..])? {
1283 let meta: TrajMeta = serde_json::from_str(existing)?;
1284 (meta.n_frames, Some(meta))
1285 } else {
1286 (0u32, None)
1287 }
1288 };
1289 let mut prepared =
1291 Self::prepare_trajectory_frames(traj_id, frames, start_idx, self.compress_on())?;
1292 if self.compress_on() {
1293 Self::group_prepared(&mut prepared)?;
1294 }
1295 let hint = source_hint.into();
1296 let input_bytes = prepared_bytes(&prepared).saturating_add(hint.len());
1297 self.with_write(input_bytes, |wtxn| {
1298 let tid_key = traj_id.to_be_bytes();
1299 let live: Option<TrajMeta> =
1300 if let Some(existing) = self.traj_meta.get(wtxn, &tid_key[..])? {
1301 Some(serde_json::from_str(existing)?)
1302 } else {
1303 None
1304 };
1305 let live_n = live.as_ref().map(|m| m.n_frames).unwrap_or(0);
1306 if live_n != start_idx {
1307 return Err(Error::Message(format!(
1308 "concurrent extend race on traj {traj_id}: expected start {start_idx}, live {live_n}"
1309 )));
1310 }
1311 for p in &prepared {
1312 self.put_index_puts(wtxn, p)?;
1313 }
1314 let n_frames = prepared
1315 .last()
1316 .map(|p| p.fk.frame_idx + 1)
1317 .unwrap_or(start_idx);
1318 let mut meta = live.or_else(|| prior.clone()).unwrap_or_else(|| TrajMeta {
1319 source: hint.clone(),
1320 ..Default::default()
1321 });
1322 meta.n_frames = n_frames;
1323 self.traj_meta
1324 .put(wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
1325 Ok(n_frames)
1326 })
1327 }
1328
1329 pub fn extend_trajectory_path(&self, traj_id: TrajId, file: impl AsRef<Path>) -> Result<u32> {
1332 self.extend_trajectory_path_units(traj_id, file, None)
1333 }
1334
1335 pub fn extend_trajectory_path_units(
1337 &self,
1338 traj_id: TrajId,
1339 file: impl AsRef<Path>,
1340 units: Option<serde_json::Value>,
1341 ) -> Result<u32> {
1342 let text = fs::read_to_string(file.as_ref())?;
1343 let mut frames = Vec::new();
1344 let units = match units {
1345 Some(u) => Some(crate::units_canon::canonicalize_units_object(u)?),
1346 None => None,
1347 };
1348 for item in ConFrameIterator::new(&text) {
1349 let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
1350 if let Some(ref u) = units {
1351 fr.header.set_units(u.clone());
1352 }
1353 frames.push(fr);
1354 }
1355 self.extend_trajectory_frames(traj_id, &frames, file.as_ref().display().to_string())
1356 }
1357
1358 fn clear_secondary(&self, wtxn: &mut RwTxn) -> Result<()> {
1359 self.idx_natoms.clear(wtxn)?;
1360 self.idx_symbol.clear(wtxn)?;
1361 self.idx_energy.clear(wtxn)?;
1362 self.idx_flags.clear(wtxn)?;
1363 self.idx_elem_count.clear(wtxn)?;
1364 self.idx_formula.clear(wtxn)?;
1365 self.idx_fmax.clear(wtxn)?;
1366 self.idx_mass.clear(wtxn)?;
1367 self.idx_volume.clear(wtxn)?;
1368 self.idx_pbc.clear(wtxn)?;
1369 self.idx_meta.clear(wtxn)?;
1370 self.frame_by_hash.clear(wtxn)?;
1371 self.hash_by_frame.clear(wtxn)?;
1372 Ok(())
1373 }
1374
1375 pub fn reindex(&self) -> Result<u32> {
1382 let recorded = self.collected_topo_params()?;
1383 self.with_write(0, |wtxn| {
1384 self.clear_secondary(wtxn)?;
1385 if let Some(idx) = self.idx_topo {
1386 idx.clear(wtxn)?;
1387 }
1388
1389 let mut n = 0u32;
1390 let mut keys: Vec<Vec<u8>> = Vec::new();
1392 {
1393 let mut iter = self.frames.iter(wtxn)?;
1394 while let Some(Ok((k, _))) = iter.next() {
1395 keys.push(k.to_vec());
1396 }
1397 }
1398 for k in keys {
1399 let Some(fk) = FrameKey::from_bytes(&k) else {
1400 continue;
1401 };
1402 let Some(stored) = self.frames.get(wtxn, &k)? else {
1403 continue;
1404 };
1405 let stored_owned = stored.to_vec();
1406 let blob_owned = crate::frame_codec::decode_frame(&stored_owned, &mut |idx| {
1407 let hk = FrameKey {
1408 traj_id: fk.traj_id,
1409 frame_idx: idx,
1410 }
1411 .to_bytes();
1412 match self.frames.get(&*wtxn, &hk[..])? {
1413 Some(b) => Ok(b.to_vec()),
1414 None => Err(Error::MissingFrame(FrameKey {
1415 traj_id: fk.traj_id,
1416 frame_idx: idx,
1417 })),
1418 }
1419 })?;
1420 let mut it = ConFrameIterator::new(&blob_owned);
1421 let frame = it
1422 .next()
1423 .ok_or_else(|| Error::Parse("empty blob on reindex".into()))?
1424 .map_err(|e| Error::Parse(e.to_string()))?;
1425 let hash = hash_frame_bytes(blob_owned.as_bytes());
1427 let hb = hash.to_bytes();
1428 let fk_b = fk.to_bytes();
1429 self.hash_by_frame.put(wtxn, &fk_b[..], &hb[..])?;
1430 if self.frame_by_hash.get(wtxn, &hb[..])?.is_none() {
1431 self.frame_by_hash.put(wtxn, &hb[..], &fk_b[..])?;
1432 }
1433 let n_atoms = frame.atom_data.len() as u32;
1434 self.idx_natoms
1435 .put(wtxn, &natoms_key(n_atoms, fk)[..], &())?;
1436 let counts = frame_species(&frame);
1437 for (sym, cnt) in &counts {
1438 self.idx_elem_count
1439 .put(wtxn, &elem_count_key(sym, *cnt, fk)[..], &())?;
1440 self.idx_symbol.put(wtxn, &symbol_key(sym, fk)[..], &())?;
1441 }
1442 let formula = composition_formula(&counts);
1443 if !formula.is_empty() {
1444 self.idx_formula
1445 .put(wtxn, &formula_key(&formula, fk)[..], &())?;
1446 }
1447 if let Some(e) = frame_energy(&frame) {
1448 if let Some(ek) = energy_bin_key(e, fk) {
1449 self.idx_energy.put(wtxn, &ek[..], &())?;
1450 }
1451 self.idx_flags
1452 .put(wtxn, &flag_key(FLAG_HAS_ENERGY, fk)[..], &())?;
1453 }
1454 if frame_has_forces(&frame) {
1455 self.idx_flags
1456 .put(wtxn, &flag_key(FLAG_HAS_FORCES, fk)[..], &())?;
1457 if let Some(fm) = frame_fmax(&frame) {
1458 if let Some(fk_fm) = fmax_bin_key(fm, fk) {
1459 self.idx_fmax.put(wtxn, &fk_fm[..], &())?;
1460 }
1461 }
1462 }
1463 if frame_has_velocities(&frame) {
1464 self.idx_flags
1465 .put(wtxn, &flag_key(FLAG_HAS_VELOCITIES, fk)[..], &())?;
1466 }
1467 if let Some(m) = frame_total_mass(&frame) {
1468 if let Some(k) = mass_bin_key(m, fk) {
1469 self.idx_mass.put(wtxn, &k[..], &())?;
1470 }
1471 }
1472 if let Some(v) = frame_cell_volume(&frame) {
1473 if let Some(k) = volume_bin_key(v, fk) {
1474 self.idx_volume.put(wtxn, &k[..], &())?;
1475 }
1476 }
1477 if let Some(mask) = frame_pbc_mask(&frame) {
1478 self.idx_pbc.put(wtxn, &pbc_key(mask, fk)[..], &())?;
1479 }
1480 let meta_puts: [(u8, Option<f64>); 7] = [
1481 (META_TIME, frame_time(&frame)),
1482 (META_TIMESTEP, frame_timestep(&frame)),
1483 (META_FRAME_INDEX, frame_frame_index(&frame)),
1484 (META_NEB_BEAD, frame_neb_bead(&frame)),
1485 (META_NEB_BAND, frame_neb_band(&frame)),
1486 (META_CHARGE, frame_charge(&frame)),
1487 (META_MAGMOM, frame_magmom(&frame)),
1488 ];
1489 for (ch, val) in meta_puts {
1490 if let Some(x) = val {
1491 if let Some(k) = meta_scalar_key(ch, x, fk) {
1492 self.idx_meta.put(wtxn, &k[..], &())?;
1493 }
1494 }
1495 }
1496 n += 1;
1497 }
1498 if recorded.is_some() {
1499 self.rebuild_idx_topo_from_stored(wtxn)?;
1500 }
1501 Ok(n)
1502 })
1503 }
1504
1505 pub fn traj_meta(&self, traj_id: TrajId) -> Result<Option<TrajMeta>> {
1506 let rtxn = self.read_txn()?;
1507 let tid_key = traj_id.to_be_bytes();
1508 match self.traj_meta.get(&rtxn, &tid_key[..])? {
1509 None => Ok(None),
1510 Some(s) => Ok(Some(serde_json::from_str(s)?)),
1511 }
1512 }
1513
1514 fn list_traj_metas(&self) -> Result<Vec<(TrajId, TrajMeta)>> {
1515 let rtxn = self.read_txn()?;
1516 let mut out = Vec::new();
1517 let mut iter = self.traj_meta.iter(&rtxn)?;
1518 while let Some(Ok((k, v))) = iter.next() {
1519 if k.len() != 8 {
1520 continue;
1521 }
1522 let mut tb = [0u8; 8];
1523 tb.copy_from_slice(k);
1524 out.push((u64::from_be_bytes(tb), serde_json::from_str(v)?));
1525 }
1526 Ok(out)
1527 }
1528
1529 fn collected_topo_params(&self) -> Result<Option<TopologyParams>> {
1530 let mut seen: Option<TopologyParams> = None;
1531 for (_, meta) in self.list_traj_metas()? {
1532 if !meta.has_topo() {
1533 continue;
1534 }
1535 let rec = meta.topo_params();
1536 match &seen {
1537 None => seen = Some(rec),
1538 Some(prev) if prev.agrees(&rec) => {}
1539 Some(prev) => return Err(mixed_topo_error(prev, &rec)),
1540 }
1541 }
1542 Ok(seen)
1543 }
1544
1545 fn require_topo_dbs(&self) -> Result<(Database<Bytes, Unit>, Database<Bytes, Bytes>)> {
1546 match (self.idx_topo, self.topo_by_frame) {
1547 (Some(idx), Some(by)) => Ok((idx, by)),
1548 _ => Err(Error::Message(
1549 "topology indexes missing; open the corpus for write to create idx_topo".into(),
1550 )),
1551 }
1552 }
1553
1554 fn rebuild_idx_topo_from_stored(&self, wtxn: &mut RwTxn) -> Result<u32> {
1555 let (idx_topo, topo_by_frame) = self.require_topo_dbs()?;
1556 let mut pairs: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
1557 {
1558 let mut iter = topo_by_frame.iter(wtxn)?;
1559 while let Some(Ok((k, v))) = iter.next() {
1560 pairs.push((k.to_vec(), v.to_vec()));
1561 }
1562 }
1563 let mut n = 0u32;
1564 for (fk_b, hex_b) in pairs {
1565 let Some(fk) = FrameKey::from_bytes(&fk_b) else {
1566 continue;
1567 };
1568 let hex = std::str::from_utf8(&hex_b)?;
1569 let hex = normalize_topo_hex(hex)?;
1570 idx_topo.put(wtxn, &topo_key(&hex, fk)[..], &())?;
1571 n += 1;
1572 }
1573 Ok(n)
1574 }
1575
1576 fn put_topo_entry(
1577 &self,
1578 wtxn: &mut RwTxn,
1579 idx_topo: Database<Bytes, Unit>,
1580 topo_by_frame: Database<Bytes, Bytes>,
1581 fk: FrameKey,
1582 hex: &str,
1583 ) -> Result<()> {
1584 let fk_b = fk.to_bytes();
1585 if let Some(old) = topo_by_frame.get(wtxn, &fk_b[..])? {
1586 let old_hex = std::str::from_utf8(old)?;
1587 let _ = idx_topo.delete(wtxn, &topo_key(old_hex, fk)[..])?;
1588 }
1589 topo_by_frame.put(wtxn, &fk_b[..], hex.as_bytes())?;
1590 idx_topo.put(wtxn, &topo_key(hex, fk)[..], &())?;
1591 Ok(())
1592 }
1593
1594 pub fn frame_topo_key(&self, key: FrameKey) -> Result<Option<String>> {
1596 let Some(by) = self.topo_by_frame else {
1597 return Ok(None);
1598 };
1599 let rtxn = self.read_txn()?;
1600 let fk_b = key.to_bytes();
1601 match by.get(&rtxn, &fk_b[..])? {
1602 None => Ok(None),
1603 Some(b) => Ok(Some(std::str::from_utf8(b)?.to_owned())),
1604 }
1605 }
1606
1607 pub fn annotate_topology(&self, opts: AnnotateTopologyOpts) -> Result<u32> {
1614 let seams = resolve_seams_binary(opts.seams.as_deref())?;
1615 let (idx_topo, topo_by_frame) = self.require_topo_dbs()?;
1616 let recorded = self.collected_topo_params()?;
1617 if let Some(ref rec) = recorded {
1618 let cutoff_diff = rec.cutoff.is_some() && rec.cutoff != Some(opts.cutoff);
1619 let graph_diff = rec.graph.as_deref().is_some_and(|g| g != opts.graph);
1620 let hops_diff = rec.hops.is_some_and(|h| h != opts.hops);
1621 if cutoff_diff || graph_diff || hops_diff {
1622 let incoming = TopologyParams {
1623 cutoff: Some(opts.cutoff),
1624 graph: Some(opts.graph.clone()),
1625 hops: Some(opts.hops),
1626 method: rec.method.clone(),
1627 };
1628 return Err(mixed_topo_error(rec, &incoming));
1629 }
1630 }
1631 let metas = self.list_traj_metas()?;
1632 let mut pending: Vec<(FrameKey, String)> = Vec::new();
1633 let mut run_method: Option<String> = None;
1634 for (tid, meta) in &metas {
1635 if meta.n_frames == 0 {
1636 continue;
1637 }
1638 let mut blob = String::new();
1639 for frame_idx in 0..meta.n_frames {
1640 blob.push_str(&self.get_frame_text(FrameKey {
1641 traj_id: *tid,
1642 frame_idx,
1643 })?);
1644 }
1645 let recs = fingerprint_con_bytes(
1646 &seams,
1647 blob.as_bytes(),
1648 opts.cutoff,
1649 &opts.graph,
1650 opts.hops,
1651 )?;
1652 if recs.len() != meta.n_frames as usize {
1653 return Err(Error::Message(format!(
1654 "seams fingerprint returned {} objects for traj {tid} with {} frames",
1655 recs.len(),
1656 meta.n_frames
1657 )));
1658 }
1659 for (i, rec) in recs.into_iter().enumerate() {
1660 match &run_method {
1661 None => run_method = Some(rec.method.clone()),
1662 Some(m) if m == &rec.method => {}
1663 Some(m) => {
1664 return Err(Error::Message(format!(
1665 "seams fingerprint mixed methods in one run ({m} vs {})",
1666 rec.method
1667 )));
1668 }
1669 }
1670 pending.push((
1671 FrameKey {
1672 traj_id: *tid,
1673 frame_idx: i as u32,
1674 },
1675 rec.key,
1676 ));
1677 }
1678 }
1679 if let (Some(rec), Some(method)) = (recorded.as_ref(), run_method.as_deref()) {
1680 if rec.method.as_deref().is_some_and(|m| m != method) {
1681 let incoming = TopologyParams {
1682 cutoff: Some(opts.cutoff),
1683 graph: Some(opts.graph.clone()),
1684 hops: Some(opts.hops),
1685 method: Some(method.to_owned()),
1686 };
1687 return Err(mixed_topo_error(rec, &incoming));
1688 }
1689 }
1690 let Some(method) = run_method else {
1691 return Ok(0);
1692 };
1693 let input_bytes = pending
1694 .iter()
1695 .fold(0usize, |n, (_, hex)| n.saturating_add(hex.len()));
1696 self.with_write(input_bytes, |wtxn| {
1697 for (fk, hex) in &pending {
1698 self.put_topo_entry(wtxn, idx_topo, topo_by_frame, *fk, hex)?;
1699 }
1700 for (tid, meta) in &metas {
1701 let mut meta = meta.clone();
1702 meta.stamp_topo(opts.cutoff, &opts.graph, opts.hops, &method);
1703 let tid_key = tid.to_be_bytes();
1704 self.traj_meta
1705 .put(wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
1706 }
1707 Ok(pending.len() as u32)
1708 })
1709 }
1710
1711 pub fn find_by_topology_path(&self, file: impl AsRef<Path>) -> Result<Vec<FrameKey>> {
1714 let rec = self.collected_topo_params()?.ok_or_else(|| {
1715 Error::Message(
1716 "no recorded topology parameters on this corpus; run annotate-topology first"
1717 .into(),
1718 )
1719 })?;
1720 let cutoff = rec.cutoff_or_err()?;
1721 let graph = rec.graph_or_default();
1722 let hops = rec.hops_or_default();
1723 let seams = resolve_seams_binary(None)?;
1724 let recs = run_seams_fingerprint(&seams, file.as_ref(), cutoff, &graph, hops)?;
1725 self.lookup_topo_keys(recs.into_iter().map(|r| r.key))
1726 }
1727
1728 pub fn find_by_topology_text(&self, text: &str) -> Result<Vec<FrameKey>> {
1730 let rec = self.collected_topo_params()?.ok_or_else(|| {
1731 Error::Message(
1732 "no recorded topology parameters on this corpus; run annotate-topology first"
1733 .into(),
1734 )
1735 })?;
1736 let cutoff = rec.cutoff_or_err()?;
1737 let graph = rec.graph_or_default();
1738 let hops = rec.hops_or_default();
1739 let seams = resolve_seams_binary(None)?;
1740 let recs = fingerprint_con_bytes(&seams, text.as_bytes(), cutoff, &graph, hops)?;
1741 self.lookup_topo_keys(recs.into_iter().map(|r| r.key))
1742 }
1743
1744 fn lookup_topo_keys<I>(&self, hexes: I) -> Result<Vec<FrameKey>>
1745 where
1746 I: IntoIterator<Item = String>,
1747 {
1748 let mut out = BTreeSet::new();
1749 let Some(idx) = self.idx_topo else {
1750 return Ok(Vec::new());
1751 };
1752 let rtxn = self.read_txn()?;
1753 for hex in hexes {
1754 let pref = topo_prefix(&normalize_topo_hex(&hex)?);
1755 let mut iter = idx.prefix_iter(&rtxn, &pref)?;
1756 while let Some(Ok((k, _))) = iter.next() {
1757 if k.len() < 12 {
1758 continue;
1759 }
1760 if let Some(fk) = FrameKey::from_bytes(&k[k.len() - 12..]) {
1761 out.insert(fk);
1762 }
1763 }
1764 }
1765 Ok(out.into_iter().collect())
1766 }
1767
1768 pub fn get_frame_text(&self, key: FrameKey) -> Result<String> {
1769 let rtxn = self.read_txn()?;
1770 self.get_frame_text_txn(&rtxn, key)
1771 }
1772
1773 fn stored_in(&self, rtxn: &heed::RoTxn<'_>, key: FrameKey) -> Result<Vec<u8>> {
1774 let fk_b = key.to_bytes();
1775 match self.frames.get(rtxn, &fk_b[..])? {
1776 Some(b) => Ok(b.to_vec()),
1777 None => Err(Error::MissingFrame(key)),
1778 }
1779 }
1780
1781 fn get_frame_text_txn(&self, rtxn: &heed::RoTxn<'_>, key: FrameKey) -> Result<String> {
1782 let stored = self.stored_in(rtxn, key)?;
1783 crate::frame_codec::decode_frame(&stored, &mut |idx| {
1784 self.stored_in(
1785 rtxn,
1786 FrameKey {
1787 traj_id: key.traj_id,
1788 frame_idx: idx,
1789 },
1790 )
1791 })
1792 }
1793
1794 pub fn get_frame_texts(&self, keys: &[FrameKey]) -> Result<Vec<String>> {
1796 let rtxn = self.read_txn()?;
1797 let mut out = Vec::with_capacity(keys.len());
1798 for &key in keys {
1799 out.push(self.get_frame_text_txn(&rtxn, key)?);
1800 }
1801 Ok(out)
1802 }
1803
1804 pub fn touch_trajectory_blobs(&self, traj_id: TrajId, n_frames: u32) -> Result<(u64, u64)> {
1810 let rtxn = self.read_txn()?;
1811 let mut total = 0u64;
1812 let mut checksum = 0u64;
1813 for frame_idx in 0..n_frames {
1814 let key = FrameKey { traj_id, frame_idx };
1815 let owned = self.get_frame_text_txn(&rtxn, key)?;
1817 total += owned.len() as u64;
1818 for (i, b) in owned.as_bytes().iter().enumerate() {
1819 checksum =
1820 checksum.wrapping_add((*b as u64).wrapping_mul((i as u64).wrapping_add(1)));
1821 }
1822 }
1823 std::hint::black_box(checksum);
1825 Ok((total, checksum))
1826 }
1827
1828 pub fn get_frame(&self, key: FrameKey) -> Result<ConFrame> {
1829 let text = self.get_frame_text(key)?;
1830 let mut it = ConFrameIterator::new(&text);
1831 let fr = it
1832 .next()
1833 .ok_or_else(|| Error::Parse("empty blob".into()))?
1834 .map_err(|e| Error::Parse(e.to_string()))?;
1835 Ok(fr)
1836 }
1837
1838 pub fn get_cooked_soa_bytes(&self, key: FrameKey) -> Result<Option<Vec<u8>>> {
1841 let rtxn = self.read_txn()?;
1842 let fk_b = key.to_bytes();
1843 Ok(self.frames_soa.get(&rtxn, &fk_b[..])?.map(|b| b.to_vec()))
1844 }
1845
1846 pub fn has_valid_cooked_soa(&self, key: FrameKey) -> Result<bool> {
1849 Ok(self.get_cooked_soa(key)?.is_some())
1850 }
1851
1852 pub fn get_cooked_soa(&self, key: FrameKey) -> Result<Option<crate::cooked_soa::CookedSoa>> {
1855 Ok(self
1856 .get_cooked_soa_bytes(key)?
1857 .and_then(|b| crate::cooked_soa::CookedSoa::try_decode(&b)))
1858 }
1859
1860 pub fn get_positions(&self, key: FrameKey) -> Result<Vec<[f64; 3]>> {
1863 if let Some(c) = self.get_cooked_soa(key)? {
1865 return Ok(c.positions);
1866 }
1867 let fr = self.get_frame(key)?;
1869 Ok(fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect())
1870 }
1871
1872 pub fn get_forces(&self, key: FrameKey) -> Result<Option<Vec<[f64; 3]>>> {
1875 if let Some(c) = self.get_cooked_soa(key)? {
1876 if c.forces.is_some() {
1877 return Ok(c.forces);
1878 }
1879 }
1881 let fr = self.get_frame(key)?;
1882 if !fr.atom_data.iter().any(|a| a.force().is_some()) {
1883 return Ok(None);
1884 }
1885 Ok(Some(
1886 fr.atom_data
1887 .iter()
1888 .map(|a| a.force().unwrap_or([0.0; 3]))
1889 .collect(),
1890 ))
1891 }
1892
1893 pub fn get_velocities(&self, key: FrameKey) -> Result<Option<Vec<[f64; 3]>>> {
1895 if let Some(c) = self.get_cooked_soa(key)? {
1896 if c.velocities.is_some() {
1897 return Ok(c.velocities);
1898 }
1899 }
1900 let fr = self.get_frame(key)?;
1901 if !fr.atom_data.iter().any(|a| a.velocity().is_some()) {
1902 return Ok(None);
1903 }
1904 Ok(Some(
1905 fr.atom_data
1906 .iter()
1907 .map(|a| a.velocity().unwrap_or([0.0; 3]))
1908 .collect(),
1909 ))
1910 }
1911
1912 pub fn cook_frame(&self, key: FrameKey) -> Result<usize> {
1914 let fr = self.get_frame(key)?;
1915 let bytes = crate::cooked_soa::CookedSoa::encode_frame(&fr)?;
1916 let n = bytes.len();
1917 self.with_write(n, |wtxn| {
1918 let fk_b = key.to_bytes();
1919 self.frames_soa.put(wtxn, &fk_b[..], &bytes[..])?;
1920 Ok(n)
1921 })
1922 }
1923
1924 fn stored_frame(&self, key: FrameKey) -> Result<Vec<u8>> {
1926 let rtxn = self.read_txn()?;
1927 let fk_b = key.to_bytes();
1928 match self.frames.get(&rtxn, &fk_b[..])? {
1929 Some(b) => Ok(b.to_vec()),
1930 None => Err(Error::MissingFrame(key)),
1931 }
1932 }
1933
1934 pub fn migrate_frames(&self) -> Result<u32> {
1941 let keys = self.list_frame_keys()?;
1942 let mut trajs: Vec<TrajId> = keys.iter().map(|k| k.traj_id).collect();
1943 trajs.sort_unstable();
1944 trajs.dedup();
1945 let mut n = 0u32;
1946 for traj in trajs {
1947 n += self.repack_traj(traj)?;
1948 }
1949 Ok(n)
1950 }
1951
1952 fn repack_traj(&self, traj_id: TrajId) -> Result<u32> {
1953 let keys: Vec<FrameKey> = self
1954 .list_frame_keys()?
1955 .into_iter()
1956 .filter(|k| k.traj_id == traj_id)
1957 .collect();
1958 if keys.is_empty() {
1959 return Ok(0);
1960 }
1961 let mut texts = Vec::with_capacity(keys.len());
1962 let mut old = Vec::with_capacity(keys.len());
1963 for key in &keys {
1964 texts.push((key.frame_idx, self.get_frame_text(*key)?));
1965 old.push(self.stored_frame(*key)?);
1966 }
1967 let view: Vec<(u32, &str)> = texts.iter().map(|(i, s)| (*i, s.as_str())).collect();
1968 let packed = crate::frame_codec::pack_aligned(&view)?;
1969 let changed: Vec<(FrameKey, Vec<u8>)> = keys
1970 .into_iter()
1971 .zip(packed)
1972 .enumerate()
1973 .filter_map(|(i, (key, next))| (next != old[i]).then_some((key, next)))
1974 .collect();
1975 let n = changed.len() as u32;
1976 if changed.is_empty() {
1977 return Ok(0);
1978 }
1979 let bytes = changed.iter().map(|(_, b)| b.len()).sum();
1980 self.with_write(bytes, |wtxn| {
1981 for (key, next) in &changed {
1982 let fk_b = key.to_bytes();
1983 self.frames.put(wtxn, &fk_b[..], next)?;
1984 }
1985 Ok(())
1986 })?;
1987 Ok(n)
1988 }
1989
1990 pub fn recook_all(&self) -> Result<u32> {
1993 let keys = self.list_frame_keys()?;
1994 let mut n = 0u32;
1995 for k in keys {
1996 self.cook_frame(k)?;
1997 n += 1;
1998 }
1999 Ok(n)
2000 }
2001
2002 pub fn delete_cooked_soa(&self, key: FrameKey) -> Result<()> {
2004 self.with_write(0, |wtxn| {
2005 let fk_b = key.to_bytes();
2006 let _ = self.frames_soa.delete(wtxn, &fk_b[..])?;
2007 Ok(())
2008 })
2009 }
2010
2011 pub fn append_trajectory_path_cook(
2013 &self,
2014 traj_id: TrajId,
2015 file: impl AsRef<Path>,
2016 cook: bool,
2017 ) -> Result<u32> {
2018 let n = self.append_trajectory_path(traj_id, file)?;
2019 if cook {
2020 for i in 0..n {
2021 self.cook_frame(FrameKey {
2022 traj_id,
2023 frame_idx: i,
2024 })?;
2025 }
2026 }
2027 Ok(n)
2028 }
2029
2030 pub fn frame_hash(&self, key: FrameKey) -> Result<ContentHash> {
2031 let rtxn = self.read_txn()?;
2032 let fk_b = key.to_bytes();
2033 match self.hash_by_frame.get(&rtxn, &fk_b[..])? {
2034 Some(b) => ContentHash::from_bytes(b).ok_or_else(|| Error::Message("bad hash".into())),
2035 None => {
2036 let text = self.get_frame_text_txn(&rtxn, key)?;
2037 Ok(hash_frame_bytes(text.as_bytes()))
2038 }
2039 }
2040 }
2041
2042 pub fn find_by_hash(&self, hash: ContentHash) -> Result<Option<FrameKey>> {
2043 let rtxn = self.read_txn()?;
2044 let hb = hash.to_bytes();
2045 match self.frame_by_hash.get(&rtxn, &hb[..])? {
2046 Some(b) => Ok(FrameKey::from_bytes(b)),
2047 None => Ok(None),
2048 }
2049 }
2050
2051 pub fn hash_con_text(text: &str) -> Result<ContentHash> {
2052 let mut it = ConFrameIterator::new(text);
2053 let frame = it
2054 .next()
2055 .ok_or_else(|| Error::Parse("no frame".into()))?
2056 .map_err(|e| Error::Parse(e.to_string()))?;
2057 let mut buf = Cursor::new(Vec::new());
2058 {
2059 let mut w = ConFrameWriter::new(&mut buf);
2060 w.write_frame(&frame)
2061 .map_err(|e| Error::Parse(format!("re-serialize: {e}")))?;
2062 }
2063 Ok(hash_frame_bytes(&buf.into_inner()))
2064 }
2065
2066 pub fn frame_formula(&self, key: FrameKey) -> Result<String> {
2068 let fr = self.get_frame(key)?;
2069 Ok(composition_formula(&frame_species(&fr)))
2070 }
2071
2072 pub fn select(&self, sel: &Select) -> Result<Vec<FrameKey>> {
2073 let rtxn = self.read_txn()?;
2074 let mut sets: Vec<BTreeSet<FrameKey>> = Vec::new();
2075
2076 if let Some(h) = sel.content_hash {
2077 let mut s = BTreeSet::new();
2078 if let Some(b) = self.frame_by_hash.get(&rtxn, &h[..])? {
2079 if let Some(fk) = FrameKey::from_bytes(b) {
2080 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2081 s.insert(fk);
2082 }
2083 }
2084 }
2085 sets.push(s);
2086 }
2087
2088 if sel.natoms_min.is_some() || sel.natoms_max.is_some() {
2089 let lo = sel.natoms_min.unwrap_or(0);
2090 let hi = sel.natoms_max.unwrap_or(u32::MAX);
2091 let mut s = BTreeSet::new();
2092 let mut iter = self.idx_natoms.iter(&rtxn)?;
2093 while let Some(Ok((k, _))) = iter.next() {
2094 if k.len() < 16 {
2095 continue;
2096 }
2097 let mut nb = [0u8; 4];
2098 nb.copy_from_slice(&k[..4]);
2099 let n = u32::from_be_bytes(nb);
2100 if n > hi {
2101 break;
2102 }
2103 if n >= lo {
2104 if let Some(fk) = FrameKey::from_bytes(&k[4..16]) {
2105 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2106 s.insert(fk);
2107 }
2108 }
2109 }
2110 }
2111 sets.push(s);
2112 }
2113
2114 for sym in &sel.symbols_all {
2115 let mut s = BTreeSet::new();
2116 let pref = symbol_prefix(sym);
2117 let mut iter = self.idx_symbol.prefix_iter(&rtxn, &pref)?;
2118 while let Some(Ok((k, _))) = iter.next() {
2119 if k.len() < 12 {
2120 continue;
2121 }
2122 let fk_bytes = &k[k.len() - 12..];
2123 if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
2124 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2125 s.insert(fk);
2126 }
2127 }
2128 }
2129 sets.push(s);
2130 }
2131
2132 if let Some(ref formula) = sel.exact_formula {
2133 let mut s = BTreeSet::new();
2134 let pref = formula_prefix(formula);
2135 let mut iter = self.idx_formula.prefix_iter(&rtxn, &pref)?;
2136 while let Some(Ok((k, _))) = iter.next() {
2137 if k.len() < 12 {
2138 continue;
2139 }
2140 let fk_bytes = &k[k.len() - 12..];
2141 if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
2142 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2143 s.insert(fk);
2144 }
2145 }
2146 }
2147 sets.push(s);
2148 }
2149
2150 if let Some(ref hex) = sel.topo_key {
2151 let mut s = BTreeSet::new();
2152 if let Some(idx) = self.idx_topo {
2153 let pref = topo_prefix(&normalize_topo_hex(hex)?);
2154 let mut iter = idx.prefix_iter(&rtxn, &pref)?;
2155 while let Some(Ok((k, _))) = iter.next() {
2156 if k.len() < 12 {
2157 continue;
2158 }
2159 let fk_bytes = &k[k.len() - 12..];
2160 if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
2161 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2162 s.insert(fk);
2163 }
2164 }
2165 }
2166 }
2167 sets.push(s);
2168 }
2169
2170 for (sym, exact_c) in &sel.element_count_exact {
2171 let mut s = BTreeSet::new();
2172 let pref = elem_count_symbol_prefix(sym);
2173 let mut iter = self.idx_elem_count.prefix_iter(&rtxn, &pref)?;
2174 while let Some(Ok((k, _))) = iter.next() {
2175 if let Some((c, fk)) = parse_elem_count_key(k, sym) {
2176 if c == *exact_c && sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2177 s.insert(fk);
2178 }
2179 }
2180 }
2181 sets.push(s);
2182 }
2183
2184 for (sym, min_c) in &sel.element_count_min {
2185 let mut s = BTreeSet::new();
2186 let pref = elem_count_symbol_prefix(sym);
2187 let mut iter = self.idx_elem_count.prefix_iter(&rtxn, &pref)?;
2188 while let Some(Ok((k, _))) = iter.next() {
2189 if let Some((c, fk)) = parse_elem_count_key(k, sym) {
2190 if c >= *min_c && sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2191 s.insert(fk);
2192 }
2193 }
2194 }
2195 sets.push(s);
2196 }
2197
2198 if sel.energy_min.is_some() || sel.energy_max.is_some() {
2199 let lo_e = sel.energy_min.unwrap_or(f64::NEG_INFINITY);
2200 let hi_e = sel.energy_max.unwrap_or(f64::INFINITY);
2201 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2202 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2203 let mut s = BTreeSet::new();
2204 let mut iter = self.idx_energy.iter(&rtxn)?;
2205 while let Some(Ok((k, _))) = iter.next() {
2206 if k.len() < 20 {
2207 continue;
2208 }
2209 let mut eb = [0u8; 8];
2210 eb.copy_from_slice(&k[..8]);
2211 let bits = u64::from_be_bytes(eb);
2212 if bits > hi_bits {
2213 break;
2214 }
2215 if bits >= lo_bits {
2216 if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
2217 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2218 s.insert(fk);
2219 }
2220 }
2221 }
2222 }
2223 sets.push(s);
2224 }
2225
2226 if sel.fmax_min.is_some() || sel.fmax_max.is_some() {
2227 let lo_e = sel.fmax_min.unwrap_or(0.0);
2228 let hi_e = sel.fmax_max.unwrap_or(f64::INFINITY);
2229 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2230 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2231 let mut s = BTreeSet::new();
2232 let mut iter = self.idx_fmax.iter(&rtxn)?;
2233 while let Some(Ok((k, _))) = iter.next() {
2234 if k.len() < 20 {
2235 continue;
2236 }
2237 let mut eb = [0u8; 8];
2238 eb.copy_from_slice(&k[..8]);
2239 let bits = u64::from_be_bytes(eb);
2240 if bits > hi_bits {
2241 break;
2242 }
2243 if bits >= lo_bits {
2244 if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
2245 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2246 s.insert(fk);
2247 }
2248 }
2249 }
2250 }
2251 sets.push(s);
2252 }
2253
2254 let scan_ord = |db: Database<Bytes, Unit>,
2255 lo: Option<f64>,
2256 hi: Option<f64>|
2257 -> Result<BTreeSet<FrameKey>> {
2258 let lo_e = lo.unwrap_or(f64::NEG_INFINITY);
2259 let hi_e = hi.unwrap_or(f64::INFINITY);
2260 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2261 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2262 let mut s = BTreeSet::new();
2263 let mut iter = db.iter(&rtxn)?;
2264 while let Some(Ok((k, _))) = iter.next() {
2265 if k.len() < 20 {
2266 continue;
2267 }
2268 let mut eb = [0u8; 8];
2269 eb.copy_from_slice(&k[..8]);
2270 let bits = u64::from_be_bytes(eb);
2271 if bits > hi_bits {
2272 break;
2273 }
2274 if bits >= lo_bits {
2275 if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
2276 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2277 s.insert(fk);
2278 }
2279 }
2280 }
2281 }
2282 Ok(s)
2283 };
2284
2285 if sel.mass_min.is_some() || sel.mass_max.is_some() {
2286 sets.push(scan_ord(self.idx_mass, sel.mass_min, sel.mass_max)?);
2287 }
2288 if sel.volume_min.is_some() || sel.volume_max.is_some() {
2289 sets.push(scan_ord(self.idx_volume, sel.volume_min, sel.volume_max)?);
2290 }
2291
2292 if let Some(pbc) = sel.pbc_exact {
2293 let mask = pbc_mask_from_bools(pbc);
2294 let mut s = BTreeSet::new();
2295 let pref = [mask];
2296 let mut iter = self.idx_pbc.prefix_iter(&rtxn, &pref)?;
2297 while let Some(Ok((k, _))) = iter.next() {
2298 if k.len() < 13 {
2299 continue;
2300 }
2301 if let Some(fk) = FrameKey::from_bytes(&k[1..13]) {
2302 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2303 s.insert(fk);
2304 }
2305 }
2306 }
2307 sets.push(s);
2308 }
2309
2310 let scan_meta = |ch: u8, lo: Option<f64>, hi: Option<f64>| -> Result<BTreeSet<FrameKey>> {
2311 let lo_e = lo.unwrap_or(f64::NEG_INFINITY);
2312 let hi_e = hi.unwrap_or(f64::INFINITY);
2313 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2314 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2315 let mut s = BTreeSet::new();
2316 let pref = [ch];
2317 let mut iter = self.idx_meta.prefix_iter(&rtxn, &pref)?;
2318 while let Some(Ok((k, _))) = iter.next() {
2319 if k.len() < 21 {
2320 continue;
2321 }
2322 let mut eb = [0u8; 8];
2323 eb.copy_from_slice(&k[1..9]);
2324 let bits = u64::from_be_bytes(eb);
2325 if bits > hi_bits {
2326 continue; }
2328 if bits >= lo_bits {
2329 if let Some(fk) = FrameKey::from_bytes(&k[9..21]) {
2330 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2331 s.insert(fk);
2332 }
2333 }
2334 }
2335 }
2336 Ok(s)
2337 };
2338
2339 if sel.time_min.is_some() || sel.time_max.is_some() {
2340 sets.push(scan_meta(META_TIME, sel.time_min, sel.time_max)?);
2341 }
2342 if sel.timestep_min.is_some() || sel.timestep_max.is_some() {
2343 sets.push(scan_meta(
2344 META_TIMESTEP,
2345 sel.timestep_min,
2346 sel.timestep_max,
2347 )?);
2348 }
2349 if sel.frame_index_min.is_some() || sel.frame_index_max.is_some() {
2350 sets.push(scan_meta(
2351 META_FRAME_INDEX,
2352 sel.frame_index_min,
2353 sel.frame_index_max,
2354 )?);
2355 }
2356 if sel.neb_bead_min.is_some() || sel.neb_bead_max.is_some() {
2357 sets.push(scan_meta(
2358 META_NEB_BEAD,
2359 sel.neb_bead_min,
2360 sel.neb_bead_max,
2361 )?);
2362 }
2363 if sel.neb_band_min.is_some() || sel.neb_band_max.is_some() {
2364 sets.push(scan_meta(
2365 META_NEB_BAND,
2366 sel.neb_band_min,
2367 sel.neb_band_max,
2368 )?);
2369 }
2370 if sel.charge_min.is_some() || sel.charge_max.is_some() {
2371 sets.push(scan_meta(META_CHARGE, sel.charge_min, sel.charge_max)?);
2372 }
2373 if sel.magmom_min.is_some() || sel.magmom_max.is_some() {
2374 sets.push(scan_meta(META_MAGMOM, sel.magmom_min, sel.magmom_max)?);
2375 }
2376
2377 let mut push_flag = |flag: u8| -> Result<()> {
2378 let mut s = BTreeSet::new();
2379 let pref = [flag];
2380 let mut iter = self.idx_flags.prefix_iter(&rtxn, &pref)?;
2381 while let Some(Ok((k, _))) = iter.next() {
2382 if k.len() < 13 {
2383 continue;
2384 }
2385 if let Some(fk) = FrameKey::from_bytes(&k[1..13]) {
2386 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2387 s.insert(fk);
2388 }
2389 }
2390 }
2391 sets.push(s);
2392 Ok(())
2393 };
2394 if sel.require_forces {
2395 push_flag(FLAG_HAS_FORCES)?;
2396 }
2397 if sel.require_velocities {
2398 push_flag(FLAG_HAS_VELOCITIES)?;
2399 }
2400 if sel.require_energy {
2401 push_flag(FLAG_HAS_ENERGY)?;
2402 }
2403
2404 if sets.is_empty() {
2405 let mut s = BTreeSet::new();
2406 let mut iter = self.frames.iter(&rtxn)?;
2407 while let Some(Ok((k, _))) = iter.next() {
2408 if let Some(fk) = FrameKey::from_bytes(k) {
2409 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2410 s.insert(fk);
2411 }
2412 }
2413 }
2414 sets.push(s);
2415 } else if let Some(tid) = sel.traj_id {
2416 for set in &mut sets {
2417 set.retain(|fk| fk.traj_id == tid);
2418 }
2419 }
2420
2421 sets.sort_by_key(|s| s.len());
2423 let mut acc = sets.remove(0);
2424 for s in sets {
2425 if acc.is_empty() {
2426 break;
2427 }
2428 if s.len() < acc.len() {
2429 acc = s.intersection(&acc).copied().collect();
2430 } else {
2431 acc = acc.intersection(&s).copied().collect();
2432 }
2433 }
2434
2435 let mut out: Vec<FrameKey> = acc.into_iter().collect();
2436 out.sort_unstable();
2437 if let Some(lim) = sel.limit {
2438 out.truncate(lim);
2439 }
2440 Ok(out)
2441 }
2442
2443 pub fn export_extxyz(
2444 &self,
2445 keys: &[FrameKey],
2446 path: impl AsRef<Path>,
2447 energy_key: &str,
2448 ) -> Result<usize> {
2449 use std::fs::OpenOptions;
2450 use std::io::BufWriter;
2451 let path = path.as_ref();
2452 if path.exists() {
2453 return Err(Error::Message(format!(
2454 "export_extxyz dest exists: {}",
2455 path.display()
2456 )));
2457 }
2458 let file = OpenOptions::new().write(true).create_new(true).open(path)?;
2459 let written = (|| -> Result<usize> {
2460 use std::io::Write;
2461 let mut w = BufWriter::new(file);
2462 let mut n = 0usize;
2463 for k in keys {
2464 let fr = self.get_frame(*k)?;
2465 write_frame_extxyz(&mut w, &fr, energy_key)?;
2466 n += 1;
2467 }
2468 w.flush()?;
2469 Ok(n)
2470 })();
2471 if written.is_err() {
2472 let _ = std::fs::remove_file(path);
2473 }
2474 written
2475 }
2476
2477 pub fn ingest_directory(
2478 &self,
2479 dir: impl AsRef<Path>,
2480 start_traj_id: TrajId,
2481 ) -> Result<Vec<(TrajId, u32, String)>> {
2482 self.ingest_directory_units(dir, start_traj_id, None)
2483 }
2484
2485 pub fn ingest_directory_units(
2491 &self,
2492 dir: impl AsRef<Path>,
2493 start_traj_id: TrajId,
2494 units: Option<serde_json::Value>,
2495 ) -> Result<Vec<(TrajId, u32, String)>> {
2496 self.ingest_directory_units_nthreads(dir, start_traj_id, units, ingest_thread_count())
2497 }
2498
2499 pub fn ingest_directory_units_nthreads(
2501 &self,
2502 dir: impl AsRef<Path>,
2503 start_traj_id: TrajId,
2504 units: Option<serde_json::Value>,
2505 threads: usize,
2506 ) -> Result<Vec<(TrajId, u32, String)>> {
2507 let mut paths: Vec<_> = fs::read_dir(dir)?
2508 .filter_map(|e| e.ok())
2509 .map(|e| e.path())
2510 .filter(|p| {
2511 let name = p.file_name().and_then(|s| s.to_str()).unwrap_or("");
2512 name.ends_with(".con") || name.ends_with(".convel")
2513 })
2514 .collect();
2515 paths.sort();
2516 let n = paths.len();
2517 let bounds = ingest_chunk_bounds(n);
2518 let threads = threads.max(1);
2519 let compress = self.compress_on();
2520 let mut prepared: Vec<Result<Vec<PreparedIndexPuts>>> =
2521 (0..n).map(|_| Ok(Vec::new())).collect();
2522 if threads == 1 || bounds.len() <= 1 {
2523 for (lo, hi) in bounds {
2524 for i in lo..hi {
2525 let tid = start_traj_id + i as u64;
2526 prepared[i] = Self::prepare_path(tid, &paths[i], units.as_ref(), compress);
2527 }
2528 }
2529 } else {
2530 let next = AtomicUsize::new(0);
2531 let prepared_slots = std::sync::Mutex::new(prepared);
2532 std::thread::scope(|scope| {
2533 for _ in 0..threads.min(bounds.len()) {
2534 let next = &next;
2535 let bounds = &bounds;
2536 let paths = &paths;
2537 let units = &units;
2538 let prepared_slots = &prepared_slots;
2539 scope.spawn(move || loop {
2540 let chunk = next.fetch_add(1, Ordering::Relaxed);
2541 if chunk >= bounds.len() {
2542 break;
2543 }
2544 let (lo, hi) = bounds[chunk];
2545 for (i, path) in paths.iter().enumerate().take(hi).skip(lo) {
2546 let tid = start_traj_id + i as u64;
2547 let row = Self::prepare_path(tid, path, units.as_ref(), compress);
2548 prepared_slots.lock().unwrap()[i] = row;
2549 }
2550 });
2551 }
2552 });
2553 prepared = prepared_slots.into_inner().unwrap();
2554 }
2555 let mut out = Vec::with_capacity(n);
2556 for (i, row) in prepared.into_iter().enumerate() {
2557 let mut puts = row?;
2558 let tid = start_traj_id + i as u64;
2559 let source = paths[i].display().to_string();
2560 let n_frames = self.commit_prepared(tid, &mut puts, source.clone(), false, compress)?;
2561 out.push((tid, n_frames, source));
2562 }
2563 Ok(out)
2564 }
2565
2566 fn prepare_path(
2567 traj_id: TrajId,
2568 path: &Path,
2569 units: Option<&serde_json::Value>,
2570 compress: bool,
2571 ) -> Result<Vec<PreparedIndexPuts>> {
2572 let text = fs::read_to_string(path)?;
2573 let Some(units) = units else {
2574 return Self::prepare_trajectory_str(traj_id, &text, compress);
2575 };
2576 let units = crate::units_canon::canonicalize_units_object(units.clone())?;
2577 let mut frames = Vec::new();
2578 for item in ConFrameIterator::new(&text) {
2579 let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
2580 fr.header.set_units(units.clone());
2581 frames.push(fr);
2582 }
2583 Self::prepare_trajectory_frames(traj_id, &frames, 0, compress)
2584 }
2585
2586 pub fn sorted_logical_dump(&self) -> Result<Vec<String>> {
2589 let rtxn = self.read_txn()?;
2590 let mut lines = Vec::new();
2591 fn hex(bytes: &[u8]) -> String {
2592 const DIGITS: &[u8] = b"0123456789abcdef";
2593 let mut s = String::with_capacity(bytes.len() * 2);
2594 for b in bytes {
2595 s.push(DIGITS[(b >> 4) as usize] as char);
2596 s.push(DIGITS[(b & 0xf) as usize] as char);
2597 }
2598 s
2599 }
2600 {
2601 let mut iter = self.frames.iter(&rtxn)?;
2602 while let Some(Ok((k, v))) = iter.next() {
2603 lines.push(format!("frames\t{}\t{}", hex(k), hex(v)));
2604 }
2605 }
2606 {
2607 let mut iter = self.traj_meta.iter(&rtxn)?;
2608 while let Some(Ok((k, v))) = iter.next() {
2609 lines.push(format!("traj_meta\t{}\t{}", hex(k), hex(v.as_bytes())));
2610 }
2611 }
2612 {
2613 let mut iter = self.frame_by_hash.iter(&rtxn)?;
2614 while let Some(Ok((k, v))) = iter.next() {
2615 lines.push(format!("frame_by_hash\t{}\t{}", hex(k), hex(v)));
2616 }
2617 }
2618 {
2619 let mut iter = self.hash_by_frame.iter(&rtxn)?;
2620 while let Some(Ok((k, v))) = iter.next() {
2621 lines.push(format!("hash_by_frame\t{}\t{}", hex(k), hex(v)));
2622 }
2623 }
2624 {
2625 let mut iter = self.frames_soa.iter(&rtxn)?;
2626 while let Some(Ok((k, v))) = iter.next() {
2627 lines.push(format!("frames_soa\t{}\t{}", hex(k), hex(v)));
2628 }
2629 }
2630 let unit_dbs: [(&str, Database<Bytes, Unit>); 11] = [
2631 ("idx_natoms", self.idx_natoms),
2632 ("idx_symbol", self.idx_symbol),
2633 ("idx_energy", self.idx_energy),
2634 ("idx_flags", self.idx_flags),
2635 ("idx_elem_count", self.idx_elem_count),
2636 ("idx_formula", self.idx_formula),
2637 ("idx_fmax", self.idx_fmax),
2638 ("idx_mass", self.idx_mass),
2639 ("idx_volume", self.idx_volume),
2640 ("idx_pbc", self.idx_pbc),
2641 ("idx_meta", self.idx_meta),
2642 ];
2643 for (name, db) in unit_dbs {
2644 let mut iter = db.iter(&rtxn)?;
2645 while let Some(Ok((k, ()))) = iter.next() {
2646 lines.push(format!("{name}\t{}", hex(k)));
2647 }
2648 }
2649 if let Some(db) = self.idx_topo {
2650 let mut iter = db.iter(&rtxn)?;
2651 while let Some(Ok((k, ()))) = iter.next() {
2652 lines.push(format!("idx_topo\t{}", hex(k)));
2653 }
2654 }
2655 if let Some(db) = self.topo_by_frame {
2656 let mut iter = db.iter(&rtxn)?;
2657 while let Some(Ok((k, v))) = iter.next() {
2658 lines.push(format!("topo_by_frame\t{}\t{}", hex(k), hex(v)));
2659 }
2660 }
2661 lines.sort();
2662 Ok(lines)
2663 }
2664
2665 pub fn unique_frame_keys(&self, sel: &Select) -> Result<Vec<FrameKey>> {
2666 let keys = self.select(sel)?;
2667 let mut uniq = Vec::new();
2668 for k in keys {
2669 let h = self.frame_hash(k)?;
2670 if self.find_by_hash(h)? == Some(k) {
2671 uniq.push(k);
2672 }
2673 }
2674 Ok(uniq)
2675 }
2676}
2677
2678#[cfg(test)]
2679mod tests {
2680 use super::*;
2681 use std::path::PathBuf;
2682
2683 fn fixture(name: &str) -> PathBuf {
2684 PathBuf::from(env!("CARGO_MANIFEST_DIR"))
2685 .join("resources/test")
2686 .join(name)
2687 }
2688
2689 fn fixtures_dir() -> PathBuf {
2690 PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test")
2691 }
2692
2693 #[test]
2694 fn parse_map_size_suffixes() {
2695 assert_eq!(parse_map_size("1048576"), Some(1 << 20));
2696 assert_eq!(parse_map_size("512M"), Some(512 << 20));
2697 assert_eq!(parse_map_size(" 8g "), Some(8usize << 30));
2698 assert_eq!(parse_map_size("2T"), Some(2usize << 40));
2699 assert_eq!(parse_map_size("0"), None);
2700 assert_eq!(parse_map_size("big"), None);
2701 assert_eq!(parse_map_size(""), None);
2702 }
2703
2704 #[test]
2705 fn map_size_env_overrides_default_and_request_overrides_env() {
2706 let dir = tempfile::tempdir().unwrap();
2707 std::env::set_var(MAP_SIZE_ENV, "3G");
2709 let from_env = initial_map_size(dir.path(), None).unwrap();
2710 let requested = initial_map_size(dir.path(), Some(5 << 20)).unwrap();
2711 std::env::remove_var(MAP_SIZE_ENV);
2712 assert_eq!(from_env, 3usize << 30);
2713 assert_eq!(requested, 5 << 20);
2714 assert_eq!(
2715 initial_map_size(dir.path(), None).unwrap(),
2716 DEFAULT_MAP_SIZE
2717 );
2718 }
2719
2720 #[test]
2721 fn small_map_fills_then_reserve_map_grows_it() {
2722 let dir = tempfile::tempdir().unwrap();
2723 let db = ConCorpus::open_with_map_size(dir.path(), 1 << 20).unwrap();
2724 assert_eq!(db.map_size(), 1 << 20);
2725 let grown = db.reserve_map(3 << 20).unwrap();
2726 assert!(grown >= 3 << 20, "{grown}");
2727 assert!(grown >= db.map_size());
2728 }
2729
2730 #[test]
2733 fn serial_path_grows_1mib_map_for_4mib_input() {
2734 let one = std::fs::read_to_string(fixture("sulfolene.con")).unwrap();
2735 let text = one.repeat((4 << 20) / one.len() + 1);
2736 assert!(text.len() >= 4 << 20, "{}", text.len());
2737 let src = tempfile::tempdir().unwrap();
2738 let file = src.path().join("big.con");
2739 std::fs::write(&file, &text).unwrap();
2740 let dir = tempfile::tempdir().unwrap();
2741 let db = ConCorpus::open_with_map_size(dir.path(), 1 << 20).unwrap();
2742 assert_eq!(db.map_size(), 1 << 20);
2743 let n = db.append_trajectory_path(1, &file).unwrap();
2744 assert!(n > 1, "{n}");
2745 assert!(db.map_size() > 1 << 20, "{}", db.map_size());
2746 drop(db);
2747 let used = std::fs::metadata(dir.path().join("data.mdb"))
2748 .unwrap()
2749 .len() as usize;
2750 let db = ConCorpus::open_with_map_size(dir.path(), 1 << 20).unwrap();
2751 assert!(db.map_size() >= 2 * used, "{} vs {used}", db.map_size());
2752 assert_eq!(db.traj_meta(1).unwrap().unwrap().n_frames, n);
2753 }
2754
2755 #[test]
2756 fn ingest_dump_identical_at_1_2_3_8_threads() {
2757 let src = tempfile::tempdir().unwrap();
2758 for name in ["a.con", "b.convel", "c.con", "d.con"] {
2759 let from = if name.ends_with(".convel") {
2760 fixture("tiny_cuh2.convel")
2761 } else if name.starts_with('c') {
2762 fixture("tiny_cuh2_forces.con")
2763 } else if name.starts_with('d') {
2764 fixture("tiny_multi_cuh2.con")
2765 } else {
2766 fixture("tiny_cuh2.con")
2767 };
2768 std::fs::copy(from, src.path().join(name)).unwrap();
2769 }
2770 let mut dumps = Vec::new();
2771 for threads in [1usize, 2, 3, 8] {
2772 let dir = tempfile::tempdir().unwrap();
2773 let db = ConCorpus::open(dir.path()).unwrap();
2774 let rows = db
2775 .ingest_directory_units_nthreads(src.path(), 1, None, threads)
2776 .unwrap();
2777 assert_eq!(rows.len(), 4);
2778 let ids: Vec<u64> = rows.iter().map(|r| r.0).collect();
2779 assert_eq!(ids, vec![1, 2, 3, 4]);
2780 let dump = db.sorted_logical_dump().unwrap();
2781 assert!(dump.iter().any(|l| l.starts_with("frames\t")));
2782 assert!(dump.iter().any(|l| l.starts_with("traj_meta\t")));
2783 dumps.push(dump);
2784 }
2785 for dump in &dumps[1..] {
2786 assert_eq!(dump, &dumps[0]);
2787 }
2788 }
2789
2790 #[test]
2791 fn ingest_first_error_independent_of_thread_count() {
2792 let src = tempfile::tempdir().unwrap();
2793 std::fs::copy(fixture("tiny_cuh2.con"), src.path().join("a.con")).unwrap();
2794 std::fs::write(src.path().join("b.con"), "not a con frame\n").unwrap();
2795 std::fs::write(src.path().join("c.con"), "also not a frame\n").unwrap();
2796 let mut messages = Vec::new();
2797 for threads in [1usize, 2, 3, 8] {
2798 let dir = tempfile::tempdir().unwrap();
2799 let db = ConCorpus::open(dir.path()).unwrap();
2800 let err = db
2801 .ingest_directory_units_nthreads(src.path(), 1, None, threads)
2802 .unwrap_err();
2803 messages.push(err.to_string());
2804 }
2805 for msg in &messages[1..] {
2806 assert_eq!(msg, &messages[0], "{messages:?}");
2807 }
2808 assert!(!messages[0].is_empty(), "{messages:?}");
2809 let lower = messages[0].to_lowercase();
2810 assert!(
2811 lower.contains("parse")
2812 || lower.contains("invalid")
2813 || lower.contains("incomplete")
2814 || lower.contains("header"),
2815 "{messages:?}"
2816 );
2817 }
2818
2819 #[test]
2820 fn ingest_chunk_bounds_depend_on_file_count_only() {
2821 assert!(ingest_chunk_bounds(0).is_empty());
2822 assert_eq!(
2823 ingest_chunk_bounds(8),
2824 (0..8).map(|i| (i, i + 1)).collect::<Vec<_>>()
2825 );
2826 let sixteen = ingest_chunk_bounds(16);
2827 assert_eq!(sixteen.len(), 8);
2828 assert_eq!(sixteen[0], (0, 2));
2829 assert_eq!(sixteen[7], (14, 16));
2830 assert_eq!(ingest_chunk_bounds(10), ingest_chunk_bounds(10));
2831 }
2832
2833 #[test]
2834 fn ingest_hash_dedup_and_select() {
2835 let dir = tempfile::tempdir().unwrap();
2836 let db = ConCorpus::open(dir.path()).unwrap();
2837 let n = db
2838 .append_trajectory_path(1, fixture("tiny_cuh2.con"))
2839 .unwrap();
2840 assert!(n >= 1);
2841 let k0 = FrameKey {
2842 traj_id: 1,
2843 frame_idx: 0,
2844 };
2845 let h = db.frame_hash(k0).unwrap();
2846 assert_eq!(db.find_by_hash(h).unwrap(), Some(k0));
2847
2848 let by_hash = db.select(&Select::new().exact_hash(h.to_bytes())).unwrap();
2849 assert_eq!(by_hash, vec![k0]);
2850
2851 let cu = db
2852 .select(&Select::new().require_symbol("Cu").trajectory(1))
2853 .unwrap();
2854 assert!(!cu.is_empty());
2855
2856 let n2 = db
2857 .append_trajectory_path(2, fixture("tiny_multi_cuh2.con"))
2858 .unwrap();
2859 assert!(n2 >= 2);
2860 }
2861
2862 #[test]
2863 fn get_velocities_returns_native_convel() {
2864 let dir = tempfile::tempdir().unwrap();
2865 let db = ConCorpus::open(dir.path()).unwrap();
2866 db.append_trajectory_path(1, fixture("tiny_cuh2.convel"))
2867 .unwrap();
2868 let v = db
2869 .get_velocities(FrameKey {
2870 traj_id: 1,
2871 frame_idx: 0,
2872 })
2873 .unwrap()
2874 .expect("velocities");
2875 assert!((v[0][0] - 0.001234).abs() < 1e-12, "got {}", v[0][0]);
2876 assert!(db
2877 .get_velocities(FrameKey {
2878 traj_id: 1,
2879 frame_idx: 0,
2880 })
2881 .unwrap()
2882 .is_some());
2883 db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
2884 .unwrap();
2885 assert!(db
2886 .get_velocities(FrameKey {
2887 traj_id: 2,
2888 frame_idx: 0,
2889 })
2890 .unwrap()
2891 .is_none());
2892 }
2893
2894 #[test]
2895 fn caller_units_canonical_in_con_metadata() {
2896 let dir = tempfile::tempdir().unwrap();
2897 let db = ConCorpus::open(dir.path()).unwrap();
2898 db.append_trajectory_path_units(
2899 1,
2900 fixture("tiny_cuh2.con"),
2901 Some(serde_json::json!({
2902 "length": "A",
2903 "energy": "ev",
2904 "time": "femtosecond"
2905 })),
2906 )
2907 .unwrap();
2908 let k = FrameKey {
2909 traj_id: 1,
2910 frame_idx: 0,
2911 };
2912 let u = db.frame_units(k).unwrap().expect("units");
2913 assert_eq!(u["length"], "angstrom");
2914 assert_eq!(u["energy"], "eV");
2915 assert_eq!(u["time"], "fs");
2916 let text = db.get_frame_text(k).unwrap();
2917 assert!(text.contains("angstrom"), "{text}");
2918 assert!(text.contains("fs"), "{text}");
2919 db.set_trajectory_units(
2920 1,
2921 serde_json::json!({"length": "nm", "energy": "hartree", "time": "ps"}),
2922 )
2923 .unwrap();
2924 let u2 = db.frame_units(k).unwrap().expect("units");
2925 assert_eq!(u2["length"], "nm");
2926 assert_eq!(u2["energy"], "hartree");
2927 assert_eq!(u2["time"], "ps");
2928 let fr0 = db.get_frame(k).unwrap();
2929 let vol_nm = crate::frame_scalars::frame_cell_volume(&fr0).expect("vol");
2930 let old_ang_vol = vol_nm * 1000.0;
2931 let stale = db
2932 .select(&Select::new().volume_range(old_ang_vol * 0.5, old_ang_vol * 1.5))
2933 .unwrap();
2934 assert!(
2935 stale.is_empty(),
2936 "old angstrom volume key must be gone after set_units"
2937 );
2938 let fresh = db
2939 .select(&Select::new().volume_range(vol_nm * 0.5, vol_nm * 1.5))
2940 .unwrap();
2941 assert_eq!(fresh, vec![k]);
2942 }
2943
2944 #[test]
2945 fn ingest_directory_units_stamps_canonical() {
2946 let dir = tempfile::tempdir().unwrap();
2947 let src = dir.path().join("cons");
2948 std::fs::create_dir(&src).unwrap();
2949 std::fs::copy(fixture("tiny_cuh2.con"), src.join("a.con")).unwrap();
2950 let db = ConCorpus::open(dir.path().join("db")).unwrap();
2951 let n = db
2952 .ingest_directory_units(
2953 &src,
2954 1,
2955 Some(serde_json::json!({"length":"A","energy":"ev"})),
2956 )
2957 .unwrap();
2958 assert_eq!(n.len(), 1);
2959 let u = db
2960 .frame_units(FrameKey {
2961 traj_id: 1,
2962 frame_idx: 0,
2963 })
2964 .unwrap()
2965 .unwrap();
2966 assert_eq!(u["length"], "angstrom");
2967 assert_eq!(u["energy"], "eV");
2968 }
2969
2970 #[test]
2971 fn readonly_open_and_pack_roundtrip() {
2972 let dir = tempfile::tempdir().unwrap();
2973 let key = FrameKey {
2974 traj_id: 1,
2975 frame_idx: 0,
2976 };
2977 let key1 = FrameKey {
2978 traj_id: 1,
2979 frame_idx: 1,
2980 };
2981 let packed = {
2982 let db = ConCorpus::open(dir.path()).unwrap();
2983 db.append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
2984 .unwrap();
2985 let packed = db.pack_frame(key).unwrap();
2986 db.close();
2987 packed
2988 };
2989 let decoded = crate::cooked_soa::CookedSoa::decode(&packed).unwrap();
2990 assert!(decoded.natoms >= 1);
2991
2992 let ro = ConCorpus::open_readonly(dir.path()).unwrap();
2993 let packed2 = ro.pack_frame(key).unwrap();
2994 assert_eq!(packed, packed2);
2995 let batch = ro.pack_frames(&[key, key1]).unwrap();
2996 let parts = crate::cooked_soa::decode_batch(&batch).unwrap();
2997 assert_eq!(parts.len(), 2);
2998 assert_eq!(parts[0], packed);
2999 let d1 = crate::cooked_soa::CookedSoa::decode(&parts[1]).unwrap();
3000 assert!((d1.positions[2][0] - 8.8549).abs() < 1e-4);
3001 assert!(ConCorpus::open_readonly("/no/such/corpus-dir-readonly").is_err());
3002 assert!(ro
3003 .append_trajectory_path(9, fixture("tiny_cuh2.con"))
3004 .is_err());
3005 }
3006
3007 #[test]
3008 fn workflow_metatrain_extxyz_export() {
3009 let dir = tempfile::tempdir().unwrap();
3010 let db = ConCorpus::open(dir.path().join("corpus")).unwrap();
3011 let ingested = db.ingest_directory(fixtures_dir(), 1).unwrap();
3012 assert!(ingested.len() >= 3);
3013
3014 let keys = db
3015 .select(
3016 &Select::new()
3017 .require_symbol("Cu")
3018 .natoms_range(1, 10_000)
3019 .limit(50),
3020 )
3021 .unwrap();
3022 assert!(!keys.is_empty());
3023
3024 let uniq = db
3025 .unique_frame_keys(&Select::new().require_symbol("Cu"))
3026 .unwrap();
3027 assert!(!uniq.is_empty());
3028
3029 let xyz = dir.path().join("train_subset.xyz");
3030 let n = db.export_extxyz(&uniq, &xyz, "energy").unwrap();
3031 assert_eq!(n, uniq.len());
3032 let text = std::fs::read_to_string(&xyz).unwrap();
3033 assert!(text.contains("Lattice="));
3034 assert!(text.contains("Properties="));
3035 assert!(text.lines().any(|l| l.trim_start().starts_with("Cu ")));
3036 let err = db.export_extxyz(&uniq, &xyz, "energy").unwrap_err();
3037 assert!(err.to_string().contains("dest exists"), "{err}");
3038 let missing = dir.path().join("missing_key.xyz");
3039 let bad = db.export_extxyz(
3040 &[FrameKey {
3041 traj_id: 99,
3042 frame_idx: 0,
3043 }],
3044 &missing,
3045 "energy",
3046 );
3047 assert!(bad.is_err());
3048 assert!(!missing.exists());
3049 }
3050
3051 #[test]
3052 fn open_refuses_sharded_root() {
3053 let dir = tempfile::tempdir().unwrap();
3054 let root = dir.path().join("sharded");
3055 crate::ShardedConCorpus::open(&root, 2).unwrap();
3056 assert!(root.join("shards.json").is_file());
3057 match ConCorpus::open(&root) {
3058 Ok(_) => panic!("expected sharded-root refuse"),
3059 Err(err) => assert!(err.to_string().contains("sharded root"), "{err}"),
3060 }
3061 assert!(!root.join("data.mdb").exists());
3062 }
3063
3064 #[test]
3065 fn workflow_dedup_identical_ingest() {
3066 let dir = tempfile::tempdir().unwrap();
3067 let db = ConCorpus::open(dir.path()).unwrap();
3068 let f = fixture("tiny_cuh2.con");
3069 db.append_trajectory_path(1, &f).unwrap();
3070 db.append_trajectory_path(2, &f).unwrap();
3071 let k1 = FrameKey {
3072 traj_id: 1,
3073 frame_idx: 0,
3074 };
3075 let k2 = FrameKey {
3076 traj_id: 2,
3077 frame_idx: 0,
3078 };
3079 let h1 = db.frame_hash(k1).unwrap();
3080 let h2 = db.frame_hash(k2).unwrap();
3081 assert_eq!(h1, h2);
3082 assert_eq!(db.find_by_hash(h1).unwrap(), Some(k1));
3083 let uniq = db.unique_frame_keys(&Select::new()).unwrap();
3084 assert!(uniq.contains(&k1));
3085 assert!(!uniq.contains(&k2));
3086 }
3087
3088 #[test]
3089 fn metadata_indexes_forces_velocities_energy() {
3090 let dir = tempfile::tempdir().unwrap();
3091 let db = ConCorpus::open(dir.path()).unwrap();
3092 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3093 .unwrap();
3094 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3095 .unwrap();
3096 db.append_trajectory_path(3, fixture("tiny_cuh2.convel"))
3097 .unwrap();
3098 db.append_trajectory_path(4, fixture("tiny_cuh2_vel_forces.con"))
3099 .unwrap();
3100
3101 let forces = db.select(&Select::new().require_forces()).unwrap();
3102 assert!(forces.iter().any(|k| k.traj_id == 2));
3103 assert!(forces.iter().any(|k| k.traj_id == 4));
3104 assert!(!forces.iter().any(|k| k.traj_id == 1));
3105
3106 let vels = db.select(&Select::new().require_velocities()).unwrap();
3107 assert!(vels.iter().any(|k| k.traj_id == 3));
3108 assert!(vels.iter().any(|k| k.traj_id == 4));
3109
3110 let both = db
3111 .select(&Select::new().require_forces().require_velocities())
3112 .unwrap();
3113 assert!(both.iter().any(|k| k.traj_id == 4));
3114 assert!(!both.iter().any(|k| k.traj_id == 2));
3115
3116 let with_e = db.select(&Select::new().require_energy()).unwrap();
3117 assert!(with_e.iter().any(|k| k.traj_id == 2));
3118 assert!(with_e.iter().any(|k| k.traj_id == 4));
3119 assert!(!with_e.iter().any(|k| k.traj_id == 1));
3120
3121 let in_range = db
3122 .select(&Select::new().energy_range(-43.0, -42.0).require_forces())
3123 .unwrap();
3124 assert!(in_range.iter().any(|k| k.traj_id == 2));
3125 let miss = db.select(&Select::new().energy_range(0.0, 1.0)).unwrap();
3126 assert!(miss.is_empty());
3127 }
3128
3129 #[test]
3130 fn composition_and_fmax_indexes() {
3131 let dir = tempfile::tempdir().unwrap();
3132 let db = ConCorpus::open(dir.path()).unwrap();
3133 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3134 .unwrap();
3135 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3136 .unwrap();
3137 db.append_trajectory_path(3, fixture("sulfolene.con"))
3138 .unwrap();
3139
3140 let formula = db
3141 .frame_formula(FrameKey {
3142 traj_id: 1,
3143 frame_idx: 0,
3144 })
3145 .unwrap();
3146 assert_eq!(formula, "Cu:2|H:2");
3147
3148 let exact = db
3149 .select(&Select::new().exact_composition("Cu:2|H:2"))
3150 .unwrap();
3151 assert!(exact.iter().any(|k| k.traj_id == 1));
3152 assert!(exact.iter().any(|k| k.traj_id == 2));
3153 assert!(!exact.iter().any(|k| k.traj_id == 3));
3154
3155 let cu2 = db
3156 .select(&Select::new().element_exact("Cu", 2).element_exact("H", 2))
3157 .unwrap();
3158 assert_eq!(cu2.len(), exact.len());
3159
3160 let cu_min2 = db.select(&Select::new().element_min("Cu", 2)).unwrap();
3161 assert!(cu_min2.iter().any(|k| k.traj_id == 1));
3162
3163 let wrong = db.select(&Select::new().exact_composition("Fe:1")).unwrap();
3164 assert!(wrong.is_empty());
3165
3166 let fr_f = db
3167 .get_frame(FrameKey {
3168 traj_id: 2,
3169 frame_idx: 0,
3170 })
3171 .unwrap();
3172 let expected_fmax = frame_fmax(&fr_f).expect("forces fixture has fmax");
3173 let in_fmax = db
3174 .select(
3175 &Select::new()
3176 .require_forces()
3177 .fmax_range(0.0, expected_fmax + 1e-6),
3178 )
3179 .unwrap();
3180 assert!(in_fmax.iter().any(|k| k.traj_id == 2));
3181 assert!(!in_fmax.iter().any(|k| k.traj_id == 1));
3182
3183 let too_small = db
3184 .select(&Select::new().fmax_range(0.0, expected_fmax * 0.0 - 1.0))
3185 .unwrap();
3186 let impossible = db
3188 .select(&Select::new().fmax_range(1e9, 1e9 + 1.0))
3189 .unwrap();
3190 assert!(impossible.is_empty());
3191 let _ = too_small;
3192 }
3193
3194 #[test]
3195 fn reindex_and_append_frames() {
3196 let dir = tempfile::tempdir().unwrap();
3197 let db = ConCorpus::open(dir.path()).unwrap();
3198 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3199 .unwrap();
3200 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3201 .unwrap();
3202
3203 let before = db
3204 .select(&Select::new().exact_composition("Cu:2|H:2"))
3205 .unwrap();
3206 assert!(!before.is_empty());
3207
3208 let n = db.reindex().unwrap();
3209 assert!(n >= 2);
3210 let after = db
3211 .select(&Select::new().exact_composition("Cu:2|H:2"))
3212 .unwrap();
3213 assert_eq!(before, after);
3214
3215 let fr = db
3216 .get_frame(FrameKey {
3217 traj_id: 1,
3218 frame_idx: 0,
3219 })
3220 .unwrap();
3221 let n2 = db
3222 .append_trajectory_frames(10, &[fr.clone()], "mem")
3223 .unwrap();
3224 assert_eq!(n2, 1);
3225 let mem = db
3226 .select(&Select::new().trajectory(10).exact_composition("Cu:2|H:2"))
3227 .unwrap();
3228 assert_eq!(mem.len(), 1);
3229
3230 let n3 = db.extend_trajectory_frames(10, &[fr], "mem-ext").unwrap();
3231 assert_eq!(n3, 2);
3232 }
3233
3234 #[test]
3235 fn extend_trajectory_path_creates_then_appends() {
3236 let dir = tempfile::tempdir().unwrap();
3237 let db = ConCorpus::open(dir.path()).unwrap();
3238 let path = fixture("tiny_cuh2.con");
3239 let n1 = db.extend_trajectory_path(3, &path).unwrap();
3240 assert!(n1 >= 1);
3241 let n2 = db.extend_trajectory_path(3, &path).unwrap();
3242 assert_eq!(n2, n1 * 2);
3243 }
3244
3245 #[test]
3246 fn reindex_twice_stable() {
3247 let dir = tempfile::tempdir().unwrap();
3248 let db = ConCorpus::open(dir.path()).unwrap();
3249 db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
3250 .unwrap();
3251 let a = db.reindex().unwrap();
3252 let keys1 = db.select(&Select::new().require_forces()).unwrap();
3253 let b = db.reindex().unwrap();
3254 let keys2 = db.select(&Select::new().require_forces()).unwrap();
3255 assert_eq!(a, b);
3256 assert_eq!(keys1, keys2);
3257 }
3258
3259 #[test]
3260 fn batch_touch_and_texts() {
3261 let dir = tempfile::tempdir().unwrap();
3262 let db = ConCorpus::open(dir.path()).unwrap();
3263 let n = db
3264 .append_trajectory_path(1, fixture("tiny_cuh2.con"))
3265 .unwrap();
3266 assert!(n >= 1);
3267 let (total, checksum) = db.touch_trajectory_blobs(1, n).unwrap();
3268 assert!(total > 0);
3269 assert!(checksum > 0);
3270 let texts = db
3271 .get_frame_texts(&[FrameKey {
3272 traj_id: 1,
3273 frame_idx: 0,
3274 }])
3275 .unwrap();
3276 assert_eq!(texts.len(), 1);
3277 assert_eq!(texts[0].len() as u64, total);
3278 let mut expect = 0u64;
3280 for (i, b) in texts[0].as_bytes().iter().enumerate() {
3281 expect = expect.wrapping_add((*b as u64).wrapping_mul((i as u64).wrapping_add(1)));
3282 }
3283 assert_eq!(checksum, expect);
3284 }
3285
3286 #[test]
3288 fn concurrent_writers_distinct_traj() {
3289 use std::sync::Arc;
3290 use std::thread;
3291 let dir = tempfile::tempdir().unwrap();
3292 let db = Arc::new(ConCorpus::open(dir.path()).unwrap());
3293 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
3294 let mut handles = Vec::new();
3295 for tid in 1u64..=8 {
3296 let db = Arc::clone(&db);
3297 let text = text.clone();
3298 handles.push(thread::spawn(move || {
3299 db.append_trajectory_str(tid, &text, format!("t{tid}"))
3300 .unwrap()
3301 }));
3302 }
3303 let mut counts = Vec::new();
3304 for h in handles {
3305 counts.push(h.join().unwrap());
3306 }
3307 assert!(counts.iter().all(|&c| c >= 1));
3308 for tid in 1u64..=8 {
3309 let meta = db.traj_meta(tid).unwrap().expect("meta");
3310 assert_eq!(meta.n_frames, counts[(tid - 1) as usize]);
3311 let _ = db
3312 .get_frame_text(FrameKey {
3313 traj_id: tid,
3314 frame_idx: 0,
3315 })
3316 .unwrap();
3317 }
3318 let cu = db.select(&Select::new().require_symbol("Cu")).unwrap();
3319 assert_eq!(cu.len(), 8);
3320 }
3321
3322 #[test]
3324 fn multiproc_cli_concurrent_select() {
3325 use std::process::Command;
3326 let dir = tempfile::tempdir().unwrap();
3327 let corpus = dir.path().join("mp_corpus");
3328 {
3329 let db = ConCorpus::open(&corpus).unwrap();
3330 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3331 .unwrap();
3332 }
3333 let Some(bin) = option_env!("CARGO_BIN_EXE_readcon-db") else {
3334 return;
3335 };
3336 let corpus_s = corpus.to_str().unwrap();
3337 let mut kids = Vec::new();
3338 for _ in 0..4 {
3339 kids.push(
3340 Command::new(bin)
3341 .args(["select", corpus_s, "--symbol", "Cu"])
3342 .stdout(std::process::Stdio::null())
3343 .spawn()
3344 .expect("spawn"),
3345 );
3346 }
3347 for mut c in kids {
3348 assert!(c.wait().unwrap().success());
3349 }
3350 }
3351
3352 #[test]
3353 fn ase_competitive_mass_volume_pbc_meta_charge() {
3354 use readcon_core::types::ConFrameBuilder;
3355 let dir = tempfile::tempdir().unwrap();
3356 let db = ConCorpus::open(dir.path()).unwrap();
3357
3358 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3360 .unwrap();
3361 let fr1 = db
3362 .get_frame(FrameKey {
3363 traj_id: 1,
3364 frame_idx: 0,
3365 })
3366 .unwrap();
3367 let mass1 = frame_total_mass(&fr1).expect("mass");
3368 let vol1 = frame_cell_volume(&fr1).expect("vol");
3369
3370 let mut fr2 = fr1.clone();
3372 fr2.header.set_frame_index(7);
3373 fr2.header.set_time(1.5);
3374 fr2.header.set_timestep(0.25);
3375 fr2.header.set_neb_bead(2);
3376 fr2.header.set_neb_band(1);
3377 fr2.header.set_pbc([true, true, false]);
3378 fr2.header
3379 .metadata
3380 .insert("charge".into(), serde_json::json!(-1.0));
3381 fr2.header
3382 .metadata
3383 .insert("magmom".into(), serde_json::json!(2.0));
3384 fr2.header.boxl = [20.0, 20.0, 20.0];
3386 db.append_trajectory_frames(2, &[fr2], "meta").unwrap();
3387
3388 let by_mass = db
3389 .select(&Select::new().mass_range(mass1 - 1e-6, mass1 + 1e-6))
3390 .unwrap();
3391 assert!(by_mass.iter().any(|k| k.traj_id == 1));
3392
3393 let by_vol_small = db
3394 .select(&Select::new().volume_range(vol1 - 1.0, vol1 + 1.0))
3395 .unwrap();
3396 assert!(by_vol_small.iter().any(|k| k.traj_id == 1));
3397 assert!(!by_vol_small.iter().any(|k| k.traj_id == 2));
3398
3399 let pbc_match = db.select(&Select::new().pbc([true, true, false])).unwrap();
3400 assert_eq!(pbc_match.len(), 1);
3401 assert_eq!(pbc_match[0].traj_id, 2);
3402 let pbc_miss = db.select(&Select::new().pbc([true, true, true])).unwrap();
3403 assert!(pbc_miss.is_empty());
3404
3405 let fi = db
3406 .select(&Select::new().frame_index_range(7.0, 7.0))
3407 .unwrap();
3408 assert_eq!(
3409 fi,
3410 vec![FrameKey {
3411 traj_id: 2,
3412 frame_idx: 0
3413 }]
3414 );
3415 let fi_miss = db
3416 .select(&Select::new().frame_index_range(99.0, 100.0))
3417 .unwrap();
3418 assert!(fi_miss.is_empty());
3419
3420 let bead = db.select(&Select::new().neb_bead_range(2.0, 2.0)).unwrap();
3421 assert_eq!(bead.len(), 1);
3422 let t = db.select(&Select::new().time_range(1.4, 1.6)).unwrap();
3423 assert_eq!(t.len(), 1);
3424
3425 let dt = db
3426 .select(&Select::new().timestep_range(0.25, 0.25))
3427 .unwrap();
3428 assert_eq!(
3429 dt,
3430 vec![FrameKey {
3431 traj_id: 2,
3432 frame_idx: 0
3433 }]
3434 );
3435 let dt_miss = db.select(&Select::new().timestep_range(9.0, 10.0)).unwrap();
3436 assert!(dt_miss.is_empty());
3437 assert!(!dt.iter().any(|k| k.traj_id == 1));
3439
3440 let band = db.select(&Select::new().neb_band_range(1.0, 1.0)).unwrap();
3441 assert_eq!(
3442 band,
3443 vec![FrameKey {
3444 traj_id: 2,
3445 frame_idx: 0
3446 }]
3447 );
3448 let band_miss = db
3449 .select(&Select::new().neb_band_range(99.0, 100.0))
3450 .unwrap();
3451 assert!(band_miss.is_empty());
3452
3453 let ch = db.select(&Select::new().charge_range(-1.0, -1.0)).unwrap();
3454 assert_eq!(ch.len(), 1);
3455 let mm = db.select(&Select::new().magmom_range(2.0, 2.0)).unwrap();
3456 assert_eq!(mm.len(), 1);
3457 let ch_miss = db.select(&Select::new().charge_range(0.0, 0.0)).unwrap();
3458 assert!(ch_miss.is_empty());
3459
3460 let n = db.reindex().unwrap();
3461 assert!(n >= 2);
3462 let after = db
3463 .select(&Select::new().frame_index_range(7.0, 7.0))
3464 .unwrap();
3465 assert_eq!(after.len(), 1);
3466 let after_dt = db
3467 .select(&Select::new().timestep_range(0.25, 0.25))
3468 .unwrap();
3469 assert_eq!(after_dt.len(), 1);
3470 let after_band = db.select(&Select::new().neb_band_range(1.0, 1.0)).unwrap();
3471 assert_eq!(after_band.len(), 1);
3472
3473 let form = db
3475 .select(&Select::new().exact_composition("Cu:2|H:2"))
3476 .unwrap();
3477 assert!(form.len() >= 2);
3478 }
3479
3480 #[test]
3481 fn cooked_soa_roundtrip_and_authority() {
3482 let dir = tempfile::tempdir().unwrap();
3483 let db = ConCorpus::open(dir.path()).unwrap();
3484 let key = FrameKey {
3485 traj_id: 1,
3486 frame_idx: 0,
3487 };
3488 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3489 .unwrap();
3490 assert!(db.get_cooked_soa_bytes(key).unwrap().is_none());
3492 let pos_via_parse = db.get_positions(key).unwrap();
3493 let fr = db.get_frame(key).unwrap();
3494 assert_eq!(pos_via_parse.len(), fr.atom_data.len());
3495 for (i, a) in fr.atom_data.iter().enumerate() {
3496 assert_eq!(pos_via_parse[i], [a.x, a.y, a.z]);
3497 }
3498
3499 let text_before = db.get_frame_text(key).unwrap();
3500 let h_before = db.frame_hash(key).unwrap();
3501 let nbytes = db.cook_frame(key).unwrap();
3502 assert!(nbytes > 0);
3503 let cooked = db.get_cooked_soa(key).unwrap().expect("cooked present");
3504 assert_eq!(cooked.natoms as usize, fr.atom_data.len());
3505 for (i, a) in fr.atom_data.iter().enumerate() {
3506 assert_eq!(cooked.positions[i], [a.x, a.y, a.z]);
3507 }
3508 assert_eq!(db.get_frame_text(key).unwrap(), text_before);
3510 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h_before.to_bytes());
3511 assert_eq!(db.find_by_hash(h_before).unwrap(), Some(key));
3512 let pos_cooked = db.get_positions(key).unwrap();
3513 assert_eq!(pos_cooked, cooked.positions);
3514
3515 let keys = db.select(&Select::new().require_symbol("Cu")).unwrap();
3517 assert!(keys.contains(&key));
3518
3519 db.reindex().unwrap();
3521 assert!(db.get_cooked_soa(key).unwrap().is_some());
3522 assert!(db
3523 .select(&Select::new().require_symbol("Cu"))
3524 .unwrap()
3525 .contains(&key));
3526
3527 db.delete_cooked_soa(key).unwrap();
3529 assert!(db.get_cooked_soa_bytes(key).unwrap().is_none());
3530 let pos_again = db.get_positions(key).unwrap();
3531 assert_eq!(pos_again, pos_via_parse);
3532 assert_eq!(db.get_frame_text(key).unwrap(), text_before);
3533 }
3534
3535 #[test]
3536 fn cooked_soa_forces_and_cook_on_ingest() {
3537 let dir = tempfile::tempdir().unwrap();
3538 let db = ConCorpus::open(dir.path()).unwrap();
3539 db.append_trajectory_path_cook(3, fixture("tiny_cuh2_forces.con"), true)
3540 .unwrap();
3541 let key = FrameKey {
3542 traj_id: 3,
3543 frame_idx: 0,
3544 };
3545 assert!(db.get_cooked_soa_bytes(key).unwrap().is_some());
3546 let fr = db.get_frame(key).unwrap();
3547 let forces = db.get_forces(key).unwrap().expect("forces");
3548 for (i, a) in fr.atom_data.iter().enumerate() {
3549 if let Some(f) = a.force() {
3550 assert_eq!(forces[i], f);
3551 }
3552 }
3553 {
3555 let _gate = db.write_guard();
3556 let mut wtxn = db.env.write_txn().unwrap();
3557 let fk_b = key.to_bytes();
3558 db.frames_soa
3559 .put(&mut wtxn, &fk_b[..], b"XXXX_not_valid")
3560 .unwrap();
3561 wtxn.commit().unwrap();
3562 }
3563 assert!(db.get_cooked_soa(key).unwrap().is_none());
3564 let pos = db.get_positions(key).unwrap();
3565 assert_eq!(pos.len(), fr.atom_data.len());
3566 assert_eq!(
3567 pos[0],
3568 [fr.atom_data[0].x, fr.atom_data[0].y, fr.atom_data[0].z]
3569 );
3570 }
3571
3572 #[test]
3573 fn cooked_soa_recook_all() {
3574 let dir = tempfile::tempdir().unwrap();
3575 let db = ConCorpus::open(dir.path()).unwrap();
3576 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3577 .unwrap();
3578 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3579 .unwrap();
3580 let n = db.recook_all().unwrap();
3581 assert_eq!(n, 2);
3582 assert!(db
3583 .get_cooked_soa(FrameKey {
3584 traj_id: 1,
3585 frame_idx: 0
3586 })
3587 .unwrap()
3588 .is_some());
3589 assert!(db
3590 .get_cooked_soa(FrameKey {
3591 traj_id: 2,
3592 frame_idx: 0
3593 })
3594 .unwrap()
3595 .is_some());
3596 }
3597
3598 #[test]
3600 fn cooked_soa_cannot_replace_missing_frames_text() {
3601 let dir = tempfile::tempdir().unwrap();
3602 let db = ConCorpus::open(dir.path()).unwrap();
3603 let key = FrameKey {
3604 traj_id: 9,
3605 frame_idx: 0,
3606 };
3607 db.append_trajectory_path_cook(9, fixture("tiny_cuh2.con"), true)
3608 .unwrap();
3609 assert!(db.has_valid_cooked_soa(key).unwrap());
3610 let pos_ok = db.get_positions(key).unwrap();
3611 assert!(!pos_ok.is_empty());
3612 {
3614 let _gate = db.write_guard();
3615 let mut wtxn = db.env.write_txn().unwrap();
3616 let fk_b = key.to_bytes();
3617 db.frames.delete(&mut wtxn, &fk_b[..]).unwrap();
3618 wtxn.commit().unwrap();
3619 }
3620 assert!(db.get_cooked_soa(key).unwrap().is_some());
3622 let pos_cooked_only = db.get_positions(key).unwrap();
3623 assert_eq!(pos_cooked_only, pos_ok);
3624 assert!(db.get_frame_text(key).is_err());
3626 assert!(db.get_frame(key).is_err());
3627 assert!(db.frame_formula(key).is_err());
3630 }
3631
3632 #[test]
3633 fn numeric_fast_path_matches_con_twice() {
3634 let dir = tempfile::tempdir().unwrap();
3635 let db = ConCorpus::open(dir.path()).unwrap();
3636 let key = FrameKey {
3637 traj_id: 1,
3638 frame_idx: 0,
3639 };
3640 db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
3641 .unwrap();
3642 let fr = db.get_frame(key).unwrap();
3643 let expect_pos: Vec<_> = fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect();
3644 let expect_f: Vec<_> = fr
3645 .atom_data
3646 .iter()
3647 .map(|a| a.force().unwrap_or([0.0; 3]))
3648 .collect();
3649 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3651 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3652 db.cook_frame(key).unwrap();
3653 assert!(db.has_valid_cooked_soa(key).unwrap());
3654 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3656 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3657 let text = db.get_frame_text(key).unwrap();
3658 let h = db.frame_hash(key).unwrap();
3659 db.delete_cooked_soa(key).unwrap();
3660 assert_eq!(db.get_frame_text(key).unwrap(), text);
3661 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
3662 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3663 }
3664
3665 #[test]
3667 fn numeric_extract_velocities_match_con() {
3668 let dir = tempfile::tempdir().unwrap();
3669 let db = ConCorpus::open(dir.path()).unwrap();
3670 let key = FrameKey {
3671 traj_id: 1,
3672 frame_idx: 0,
3673 };
3674 db.append_trajectory_path(1, fixture("tiny_cuh2_vel_forces.con"))
3675 .unwrap();
3676 assert!(!db.has_valid_cooked_soa(key).unwrap());
3677 let fr = db.get_frame(key).unwrap();
3678 let expect_pos: Vec<_> = fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect();
3679 let expect_f: Vec<_> = fr
3680 .atom_data
3681 .iter()
3682 .map(|a| a.force().unwrap_or([0.0; 3]))
3683 .collect();
3684 let expect_v: Vec<_> = fr
3685 .atom_data
3686 .iter()
3687 .map(|a| a.velocity().unwrap_or([0.0; 3]))
3688 .collect();
3689 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3690 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3691 assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
3692 assert!((expect_v[0][0] - 0.001234).abs() < 1e-12);
3693 let text = db.get_frame_text(key).unwrap();
3694 let h = db.frame_hash(key).unwrap();
3695 db.cook_frame(key).unwrap();
3696 assert!(db.has_valid_cooked_soa(key).unwrap());
3697 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3698 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3699 assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
3700 assert_eq!(db.get_frame_text(key).unwrap(), text);
3701 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
3702 assert_eq!(db.find_by_hash(h).unwrap(), Some(key));
3703 db.delete_cooked_soa(key).unwrap();
3704 assert!(!db.has_valid_cooked_soa(key).unwrap());
3705 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3706 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3707 assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
3708 assert_eq!(db.get_frame_text(key).unwrap(), text);
3709 }
3710
3711 #[test]
3713 fn numeric_extract_valid_rcso_skips_con_parse() {
3714 let dir = tempfile::tempdir().unwrap();
3715 let db = ConCorpus::open(dir.path()).unwrap();
3716 let key = FrameKey {
3717 traj_id: 1,
3718 frame_idx: 0,
3719 };
3720 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3721 .unwrap();
3722 let text = db.get_frame_text(key).unwrap();
3723 let h = db.frame_hash(key).unwrap();
3724 let mut fr = db.get_frame(key).unwrap();
3725 let con_x = fr.atom_data[0].x;
3726 fr.atom_data[0].x = 99.0;
3727 let poison = crate::cooked_soa::CookedSoa::encode_frame(&fr).unwrap();
3728 {
3729 let _gate = db.write_guard();
3730 let mut wtxn = db.env.write_txn().unwrap();
3731 let fk_b = key.to_bytes();
3732 db.frames_soa
3733 .put(&mut wtxn, &fk_b[..], &poison[..])
3734 .unwrap();
3735 wtxn.commit().unwrap();
3736 }
3737 assert!(db.has_valid_cooked_soa(key).unwrap());
3738 let pos = db.get_positions(key).unwrap();
3739 assert!(
3740 (pos[0][0] - 99.0).abs() < 1e-12,
3741 "valid RCSO must skip CON parse, got {}",
3742 pos[0][0]
3743 );
3744 assert!((con_x - 0.6394).abs() < 1e-4);
3745 assert_eq!(db.get_frame_text(key).unwrap(), text);
3746 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
3747 }
3748
3749 #[test]
3750 fn traj_meta_old_json_loads_without_topo_fields() {
3751 let m: TrajMeta = serde_json::from_str(r#"{"n_frames":3,"source":"a.con"}"#).unwrap();
3752 assert_eq!(m.n_frames, 3);
3753 assert_eq!(m.source, "a.con");
3754 assert!(m.topo_cutoff.is_none());
3755 assert!(m.topo_graph.is_none());
3756 assert!(m.topo_hops.is_none());
3757 assert!(m.topo_method.is_none());
3758 }
3759
3760 fn seed_topo(db: &ConCorpus, fk: FrameKey, hex: &str) {
3761 let (idx, by) = db.require_topo_dbs().unwrap();
3762 let _gate = db.write_guard();
3763 let mut wtxn = db.env.write_txn().unwrap();
3764 db.put_topo_entry(&mut wtxn, idx, by, fk, hex).unwrap();
3765 wtxn.commit().unwrap();
3766 }
3767
3768 fn stamp_traj_topo(db: &ConCorpus, tid: TrajId, method: &str, cutoff: f64) {
3769 let mut meta = db.traj_meta(tid).unwrap().expect("meta");
3770 meta.stamp_topo(cutoff, "cutoff", 2, method);
3771 let _gate = db.write_guard();
3772 let mut wtxn = db.env.write_txn().unwrap();
3773 let tid_key = tid.to_be_bytes();
3774 db.traj_meta
3775 .put(
3776 &mut wtxn,
3777 &tid_key[..],
3778 &serde_json::to_string(&meta).unwrap(),
3779 )
3780 .unwrap();
3781 wtxn.commit().unwrap();
3782 }
3783
3784 #[test]
3785 fn select_topo_key_prefix_without_seams() {
3786 let dir = tempfile::tempdir().unwrap();
3787 let db = ConCorpus::open(dir.path()).unwrap();
3788 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3789 .unwrap();
3790 let k0 = FrameKey {
3791 traj_id: 1,
3792 frame_idx: 0,
3793 };
3794 seed_topo(&db, k0, "abc123");
3795 let hits = db.select(&Select::new().topo_key("ABC123")).unwrap();
3796 assert_eq!(hits, vec![k0]);
3797 let miss = db.select(&Select::new().topo_key("00")).unwrap();
3798 assert!(miss.is_empty());
3799 }
3800
3801 #[test]
3802 fn reindex_without_annotation_leaves_idx_topo_empty() {
3803 let dir = tempfile::tempdir().unwrap();
3804 let db = ConCorpus::open(dir.path()).unwrap();
3805 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3806 .unwrap();
3807 let k0 = FrameKey {
3808 traj_id: 1,
3809 frame_idx: 0,
3810 };
3811 seed_topo(&db, k0, "abc123");
3812 assert_eq!(
3813 db.select(&Select::new().topo_key("abc123")).unwrap(),
3814 vec![k0]
3815 );
3816 db.reindex().unwrap();
3817 assert!(db
3818 .select(&Select::new().topo_key("abc123"))
3819 .unwrap()
3820 .is_empty());
3821 assert_eq!(db.frame_topo_key(k0).unwrap().as_deref(), Some("abc123"));
3822 }
3823
3824 #[test]
3825 fn reindex_rebuilds_idx_topo_when_params_agree() {
3826 let dir = tempfile::tempdir().unwrap();
3827 let db = ConCorpus::open(dir.path()).unwrap();
3828 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3829 .unwrap();
3830 db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
3831 .unwrap();
3832 let k0 = FrameKey {
3833 traj_id: 1,
3834 frame_idx: 0,
3835 };
3836 let k1 = FrameKey {
3837 traj_id: 2,
3838 frame_idx: 0,
3839 };
3840 seed_topo(&db, k0, "deadbeef");
3841 seed_topo(&db, k1, "deadbeef");
3842 stamp_traj_topo(&db, 1, "nauty", 3.0);
3843 stamp_traj_topo(&db, 2, "nauty", 3.0);
3844 let (idx, _) = db.require_topo_dbs().unwrap();
3845 {
3846 let _gate = db.write_guard();
3847 let mut wtxn = db.env.write_txn().unwrap();
3848 idx.clear(&mut wtxn).unwrap();
3849 wtxn.commit().unwrap();
3850 }
3851 assert!(db
3852 .select(&Select::new().topo_key("deadbeef"))
3853 .unwrap()
3854 .is_empty());
3855 db.reindex().unwrap();
3856 let hits = db.select(&Select::new().topo_key("deadbeef")).unwrap();
3857 assert_eq!(hits.len(), 2);
3858 }
3859
3860 #[test]
3861 fn reindex_refuses_mixed_methods() {
3862 let dir = tempfile::tempdir().unwrap();
3863 let db = ConCorpus::open(dir.path()).unwrap();
3864 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3865 .unwrap();
3866 db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
3867 .unwrap();
3868 stamp_traj_topo(&db, 1, "nauty", 3.0);
3869 stamp_traj_topo(&db, 2, "wl", 3.0);
3870 let err = db.reindex().unwrap_err().to_string();
3871 assert!(err.contains("mixed topology"), "{err}");
3872 let cu = db.select(&Select::new().require_symbol("Cu")).unwrap();
3873 assert_eq!(cu.len(), 2);
3874 }
3875
3876 #[test]
3877 fn extend_preserves_topology_fields() {
3878 let dir = tempfile::tempdir().unwrap();
3879 let db = ConCorpus::open(dir.path()).unwrap();
3880 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3881 .unwrap();
3882 stamp_traj_topo(&db, 1, "nauty", 3.0);
3883 let fr = db
3884 .get_frame(FrameKey {
3885 traj_id: 1,
3886 frame_idx: 0,
3887 })
3888 .unwrap();
3889 db.extend_trajectory_frames(1, &[fr], "extra").unwrap();
3890 let meta = db.traj_meta(1).unwrap().unwrap();
3891 assert_eq!(meta.n_frames, 2);
3892 assert_eq!(meta.topo_method.as_deref(), Some("nauty"));
3893 assert_eq!(meta.topo_cutoff, Some(3.0));
3894 assert_eq!(meta.topo_graph.as_deref(), Some("cutoff"));
3895 assert_eq!(meta.topo_hops, Some(2));
3896 }
3897
3898 #[test]
3899 fn annotate_without_seams_names_command() {
3900 let dir = tempfile::tempdir().unwrap();
3901 let db = ConCorpus::open(dir.path()).unwrap();
3902 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3903 .unwrap();
3904 let err = db
3905 .annotate_topology(AnnotateTopologyOpts::new(3.0).seams("/no/such/seams-binary"))
3906 .unwrap_err()
3907 .to_string();
3908 assert!(err.contains("seams fingerprint"), "{err}");
3909 }
3910
3911 #[test]
3912 fn find_by_topology_without_params_errors() {
3913 let dir = tempfile::tempdir().unwrap();
3914 let db = ConCorpus::open(dir.path()).unwrap();
3915 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3916 .unwrap();
3917 let err = db
3918 .find_by_topology_path(fixture("tiny_cuh2.con"))
3919 .unwrap_err()
3920 .to_string();
3921 assert!(err.contains("annotate-topology"), "{err}");
3922 }
3923
3924 #[test]
3925 fn readonly_missing_topo_select_is_empty() {
3926 let dir = tempfile::tempdir().unwrap();
3927 {
3928 let db = ConCorpus::open(dir.path()).unwrap();
3929 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3930 .unwrap();
3931 db.close();
3932 }
3933 let db = ConCorpus::open_readonly(dir.path()).unwrap();
3934 let hits = db.select(&Select::new().topo_key("abc")).unwrap();
3935 assert!(hits.is_empty());
3936 }
3937
3938 fn seams_or_skip() -> Option<std::path::PathBuf> {
3939 match crate::topology::resolve_seams_binary(None) {
3940 Ok(p) => Some(p),
3941 Err(e) => {
3942 eprintln!("skipping topology seams test: {e}");
3943 None
3944 }
3945 }
3946 }
3947
3948 fn permute_cu_rows(text: &str) -> String {
3949 let mut lines: Vec<&str> = text.lines().collect();
3950 let mut cu = Vec::new();
3951 for (i, line) in lines.iter().enumerate() {
3952 let t = line.trim_start();
3953 if t.starts_with("0.6394") || t.starts_with("3.1969") {
3954 cu.push(i);
3955 }
3956 }
3957 assert_eq!(cu.len(), 2, "expected two Cu coordinate rows");
3958 lines.swap(cu[0], cu[1]);
3959 let mut out = lines.join("\n");
3960 if text.ends_with('\n') {
3961 out.push('\n');
3962 }
3963 out
3964 }
3965
3966 fn break_one_bond(text: &str) -> String {
3967 let mut fr = {
3968 let mut it = ConFrameIterator::new(text);
3969 it.next().unwrap().unwrap()
3970 };
3971 let cu = fr
3974 .atom_data
3975 .iter_mut()
3976 .find(|a| a.symbol.as_ref() == "Cu")
3977 .expect("Cu");
3978 cu.z += 10.0;
3979 let mut buf = Cursor::new(Vec::new());
3980 {
3981 let mut w = ConFrameWriter::new(&mut buf);
3982 w.write_frame(&fr).unwrap();
3983 }
3984 String::from_utf8(buf.into_inner()).unwrap()
3985 }
3986
3987 #[test]
3988 fn annotate_permutation_shares_hex_broken_does_not() {
3989 let Some(seams) = seams_or_skip() else {
3990 return;
3991 };
3992 let dir = tempfile::tempdir().unwrap();
3993 let db = ConCorpus::open(dir.path()).unwrap();
3994 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3995 .unwrap();
3996 let n = db
3997 .annotate_topology(AnnotateTopologyOpts::new(3.0).seams(&seams))
3998 .unwrap();
3999 assert_eq!(n, 1);
4000 let k0 = FrameKey {
4001 traj_id: 1,
4002 frame_idx: 0,
4003 };
4004 let hex = db.frame_topo_key(k0).unwrap().expect("hex");
4005 assert!(!hex.is_empty());
4006 let meta = db.traj_meta(1).unwrap().unwrap();
4007 assert_eq!(meta.topo_cutoff, Some(3.0));
4008 assert_eq!(meta.topo_graph.as_deref(), Some("cutoff"));
4009 assert_eq!(meta.topo_hops, Some(2));
4010 assert!(meta.topo_method.is_some());
4011
4012 let orig = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
4013 let perm = permute_cu_rows(&orig);
4014 let broken = break_one_bond(&orig);
4015 let perm_hits = db.find_by_topology_text(&perm).unwrap();
4016 assert_eq!(perm_hits, vec![k0]);
4017 let broken_hits = db.find_by_topology_text(&broken).unwrap();
4018 assert!(
4019 broken_hits.is_empty(),
4020 "broken Cu-Cu bond must not share topology key, got {broken_hits:?}"
4021 );
4022 let sel = db.select(&Select::new().topo_key(&hex)).unwrap();
4023 assert_eq!(sel, vec![k0]);
4024 }
4025
4026 fn raw_spans(path: &std::path::Path) -> Vec<String> {
4027 let raw = std::fs::read_to_string(path).unwrap();
4028 let mut it = ConFrameIterator::new(&raw);
4029 let mut out = Vec::new();
4030 loop {
4031 match it.next_with_raw_span(&raw) {
4032 None => break,
4033 Some(Ok((_, blob))) => out.push(blob.to_owned()),
4034 Some(Err(e)) => panic!("{e}"),
4035 }
4036 }
4037 out
4038 }
4039
4040 #[test]
4041 fn zstd_round_trip_is_byte_exact_and_does_not_cook() {
4042 let path = fixture("sulfolene.con");
4043 let spans = raw_spans(&path);
4044 assert!(!spans.is_empty());
4045 let dir = tempfile::tempdir().unwrap();
4046 let db = ConCorpus::open(dir.path()).unwrap();
4047 let n = db.append_trajectory_path(1, &path).unwrap();
4048 assert_eq!(n as usize, spans.len());
4049 for (i, span) in spans.iter().enumerate() {
4050 let key = FrameKey {
4051 traj_id: 1,
4052 frame_idx: i as u32,
4053 };
4054 assert_eq!(db.get_frame_text(key).unwrap(), *span);
4055 let stored = db.stored_frame(key).unwrap();
4056 assert!(
4057 crate::frame_codec::is_single_zstd_frame(&stored),
4058 "frame {i} was not one tagged zstd frame"
4059 );
4060 assert_eq!(crate::frame_codec::load_frame(&stored).unwrap(), *span);
4061 assert_eq!(
4062 db.frame_hash(key).unwrap(),
4063 hash_frame_bytes(span.as_bytes())
4064 );
4065 assert!(!db.has_valid_cooked_soa(key).unwrap());
4066 }
4067 }
4068
4069 #[test]
4070 fn grouped_frames_round_trip_and_migrate_does_not_wrap_again() {
4071 let one = std::fs::read_to_string(fixture("sulfolene.con")).unwrap();
4072 let mut text = String::new();
4073 for i in 0..80u32 {
4074 let shifted = one.replacen(
4075 "10.477713",
4076 &format!("{:.6}", 10.477713 + f64::from(i) * 1e-4),
4077 1,
4078 );
4079 text.push_str(&shifted);
4080 if !shifted.ends_with('\n') {
4081 text.push('\n');
4082 }
4083 }
4084 let dir = tempfile::tempdir().unwrap();
4085 let path = dir.path().join("in.con");
4086 std::fs::write(&path, &text).unwrap();
4087 let spans = raw_spans(&path);
4088 assert_eq!(spans.len(), 80);
4089 let db = ConCorpus::open(dir.path().join("db")).unwrap();
4090 let n = db.append_trajectory_str(1, &text, "mem").unwrap();
4091 assert_eq!(n, 80);
4092 let mut stored_sum = 0usize;
4093 let mut single_sum = 0usize;
4094 for (i, span) in spans.iter().enumerate() {
4095 let key = FrameKey {
4096 traj_id: 1,
4097 frame_idx: i as u32,
4098 };
4099 assert_eq!(db.get_frame_text(key).unwrap(), *span);
4100 assert_eq!(
4101 db.frame_hash(key).unwrap(),
4102 hash_frame_bytes(span.as_bytes())
4103 );
4104 let stored = db.stored_frame(key).unwrap();
4105 stored_sum += stored.len();
4106 single_sum += crate::frame_codec::store_frame(span, true).unwrap().len();
4107 if i == 0 {
4108 assert!(crate::frame_codec::is_group(&stored));
4109 } else if i < 64 {
4110 assert!(crate::frame_codec::is_redirect(&stored));
4111 }
4112 assert!(!db.has_valid_cooked_soa(key).unwrap());
4113 }
4114 assert!(stored_sum < single_sum, "{stored_sum} vs {single_sum}");
4115 assert_eq!(db.migrate_frames().unwrap(), 0);
4116 assert_eq!(
4117 db.get_frame_text(FrameKey {
4118 traj_id: 1,
4119 frame_idx: 79
4120 })
4121 .unwrap(),
4122 spans[79]
4123 );
4124 }
4125
4126 #[test]
4127 fn legacy_utf8_reads_and_migrate_never_double_encodes() {
4128 let path = fixture("tiny_cuh2.con");
4129 let spans = raw_spans(&path);
4130 let dir = tempfile::tempdir().unwrap();
4131 let db = ConCorpus::open(dir.path()).unwrap();
4132 db.set_compress_frames(false);
4133 let n = db.append_trajectory_path(1, &path).unwrap();
4134 assert_eq!(n as usize, spans.len());
4135 for (i, span) in spans.iter().enumerate() {
4136 let key = FrameKey {
4137 traj_id: 1,
4138 frame_idx: i as u32,
4139 };
4140 assert_eq!(db.stored_frame(key).unwrap(), span.as_bytes());
4141 assert_eq!(db.get_frame_text(key).unwrap(), *span);
4142 assert_eq!(
4143 db.frame_hash(key).unwrap(),
4144 hash_frame_bytes(span.as_bytes())
4145 );
4146 }
4147 drop(db);
4148 let db = ConCorpus::open(dir.path()).unwrap();
4149 let key0 = FrameKey {
4150 traj_id: 1,
4151 frame_idx: 0,
4152 };
4153 let hash0 = db.frame_hash(key0).unwrap();
4154 assert_eq!(db.migrate_frames().unwrap() as usize, spans.len());
4155 assert_eq!(db.get_frame_text(key0).unwrap(), spans[0]);
4156 assert_eq!(db.frame_hash(key0).unwrap(), hash0);
4157 let tagged = db.stored_frame(key0).unwrap();
4158 assert!(crate::frame_codec::is_single_zstd_frame(&tagged));
4159 assert_eq!(crate::frame_codec::load_frame(&tagged).unwrap(), spans[0]);
4160 assert!(std::str::from_utf8(&tagged).is_err());
4161 db.reindex().unwrap();
4162 assert_eq!(db.frame_hash(key0).unwrap(), hash0);
4163 assert_eq!(db.get_frame_text(key0).unwrap(), spans[0]);
4164 assert_eq!(db.migrate_frames().unwrap(), 0);
4165 assert_eq!(db.stored_frame(key0).unwrap(), tagged);
4166 assert!(!db.has_valid_cooked_soa(key0).unwrap());
4167 }
4168
4169 #[test]
4170 fn repeated_ingest_stores_identical_frame_bytes() {
4171 let path = fixture("tiny_multi_cuh2.con");
4172 let mut dumps = Vec::new();
4173 for _ in 0..2 {
4174 let dir = tempfile::tempdir().unwrap();
4175 let db = ConCorpus::open(dir.path()).unwrap();
4176 db.append_trajectory_path(1, &path).unwrap();
4177 dumps.push(db.canonical_dump().unwrap());
4178 }
4179 assert_eq!(dumps[0], dumps[1]);
4180 }
4181}