1use std::collections::BTreeSet;
2use std::fs;
3use std::io::Cursor;
4use std::path::{Path, PathBuf};
5
6use heed::types::{Bytes, Str, Unit};
7use heed::{CompactionOption, Database, Env, EnvFlags, EnvOpenOptions, RwTxn};
8use readcon_core::iterators::ConFrameIterator;
9use readcon_core::types::ConFrame;
10use readcon_core::writer::{ConFrameWriter, FloatFormat};
11use serde::{Deserialize, Serialize};
12
13use crate::error::{Error, Result};
14use crate::export_xyz::write_frame_extxyz;
15use crate::frame_scalars::{
16 frame_cell_volume, frame_charge, frame_frame_index, frame_magmom, frame_neb_band,
17 frame_neb_bead, frame_pbc_mask, frame_time, frame_timestep, frame_total_mass,
18};
19use crate::keys::{
20 composition_formula, elem_count_key, elem_count_symbol_prefix, energy_bin_key, flag_key,
21 fmax_bin_key, formula_key, formula_prefix, hash_frame_bytes, mass_bin_key, meta_scalar_key,
22 natoms_key, ordered_f64_bits, parse_elem_count_key, pbc_key, pbc_mask_from_bools,
23 symbol_key, symbol_prefix, topo_key, topo_prefix, volume_bin_key,
24 ContentHash, FrameKey, TrajId, FLAG_HAS_ENERGY, FLAG_HAS_FORCES, FLAG_HAS_VELOCITIES,
25 META_CHARGE, META_FRAME_INDEX, META_MAGMOM, META_NEB_BAND, META_NEB_BEAD, META_TIME,
26 META_TIMESTEP,
27};
28use crate::select::Select;
29use crate::topology::{
30 fingerprint_con_bytes, mixed_topo_error, normalize_topo_hex, resolve_seams_binary,
31 run_seams_fingerprint, AnnotateTopologyOpts, TopologyParams,
32};
33
34const MAP_SIZE: usize = 2 * 1024 * 1024 * 1024;
36const MAX_DBS: u32 = 48;
37const MAX_READERS: u32 = 512;
41
42#[derive(Clone, Debug, Default, Serialize, Deserialize)]
43pub struct TrajMeta {
44 pub n_frames: u32,
45 pub source: String,
46 #[serde(default)]
47 pub topo_cutoff: Option<f64>,
48 #[serde(default)]
49 pub topo_graph: Option<String>,
50 #[serde(default)]
51 pub topo_hops: Option<u32>,
52 #[serde(default)]
53 pub topo_method: Option<String>,
54}
55
56impl TrajMeta {
57 fn topo_params(&self) -> TopologyParams {
58 TopologyParams {
59 cutoff: self.topo_cutoff,
60 graph: self.topo_graph.clone(),
61 hops: self.topo_hops,
62 method: self.topo_method.clone(),
63 }
64 }
65
66 fn has_topo(&self) -> bool {
67 !self.topo_params().is_empty()
68 }
69
70 fn stamp_topo(&mut self, cutoff: f64, graph: &str, hops: u32, method: &str) {
71 self.topo_cutoff = Some(cutoff);
72 self.topo_graph = Some(graph.to_owned());
73 self.topo_hops = Some(hops);
74 self.topo_method = Some(method.to_owned());
75 }
76}
77
78struct PreparedIndexPuts {
80 fk: FrameKey,
81 blob: String,
82 hash: [u8; 16],
83 natoms_k: [u8; 16],
84 symbol_keys: Vec<Vec<u8>>,
85 elem_count_keys: Vec<Vec<u8>>,
86 formula_k: Option<Vec<u8>>,
87 energy_k: Option<[u8; 20]>,
88 fmax_k: Option<[u8; 20]>,
89 mass_k: Option<[u8; 20]>,
90 volume_k: Option<[u8; 20]>,
91 pbc_k: Option<[u8; 13]>,
92 flag_keys: Vec<[u8; 13]>,
93 meta_keys: Vec<[u8; 21]>,
94}
95
96pub struct ConCorpus {
97 path: PathBuf,
98 env: Env,
99 frames: Database<Bytes, Str>,
100 traj_meta: Database<Bytes, Str>,
101 idx_natoms: Database<Bytes, Unit>,
102 idx_symbol: Database<Bytes, Unit>,
103 idx_energy: Database<Bytes, Unit>,
104 idx_flags: Database<Bytes, Unit>,
105 idx_elem_count: Database<Bytes, Unit>,
107 idx_formula: Database<Bytes, Unit>,
109 idx_fmax: Database<Bytes, Unit>,
111 idx_mass: Database<Bytes, Unit>,
112 idx_volume: Database<Bytes, Unit>,
113 idx_pbc: Database<Bytes, Unit>,
115 idx_meta: Database<Bytes, Unit>,
117 frame_by_hash: Database<Bytes, Bytes>,
118 hash_by_frame: Database<Bytes, Bytes>,
119 frames_soa: Database<Bytes, Bytes>,
121 idx_topo: Option<Database<Bytes, Unit>>,
123 topo_by_frame: Option<Database<Bytes, Bytes>>,
125}
126
127fn frame_has_forces(frame: &ConFrame) -> bool {
128 readcon_core::index_proj::frame_has_forces(frame)
129}
130
131fn frame_has_velocities(frame: &ConFrame) -> bool {
132 readcon_core::index_proj::frame_has_velocities(frame)
133}
134
135fn frame_energy(frame: &ConFrame) -> Option<f64> {
136 readcon_core::index_proj::finite_energy(frame)
137}
138
139pub fn frame_fmax(frame: &ConFrame) -> Option<f64> {
141 readcon_core::index_proj::frame_fmax(frame)
142}
143
144fn frame_species(frame: &ConFrame) -> Vec<(String, u32)> {
145 readcon_core::index_proj::frame_species_counts(frame)
146}
147
148impl ConCorpus {
149 pub fn open(path: impl AsRef<Path>) -> Result<Self> {
150 Self::open_with(path, false)
151 }
152
153 pub fn open_existing(path: impl AsRef<Path>) -> Result<Self> {
155 let path = path.as_ref();
156 if !path.join("data.mdb").is_file() {
157 return Err(Error::Message(format!(
158 "not a corpus directory (missing data.mdb): {}",
159 path.display()
160 )));
161 }
162 Self::open_with(path, false)
163 }
164
165 pub fn open_readonly(path: impl AsRef<Path>) -> Result<Self> {
171 Self::open_with(path, true)
172 }
173
174 fn open_with(path: impl AsRef<Path>, readonly: bool) -> Result<Self> {
175 let path = path.as_ref().to_path_buf();
176 if path.join("shards.json").is_file() {
177 return Err(Error::Message(format!(
178 "refusing to open a sharded root as a single-env corpus: {}",
179 path.display()
180 )));
181 }
182 if readonly {
183 if !path.join("data.mdb").is_file() {
184 return Err(Error::Message(format!(
185 "readonly open: {} is not a corpus directory",
186 path.display()
187 )));
188 }
189 } else {
190 fs::create_dir_all(&path)?;
191 }
192 let mut opts = EnvOpenOptions::new();
193 opts.map_size(MAP_SIZE);
194 opts.max_dbs(MAX_DBS);
195 opts.max_readers(MAX_READERS);
196 if readonly {
197 unsafe {
199 opts.flags(EnvFlags::READ_ONLY);
200 }
201 }
202 let env = unsafe { opts.open(&path)? };
203
204 if readonly {
205 let rtxn = env.read_txn()?;
206 let missing =
207 |name: &str| Error::Message(format!("readonly open: missing database {name}"));
208 let frames = env
209 .open_database(&rtxn, Some("frames"))?
210 .ok_or_else(|| missing("frames"))?;
211 let traj_meta = env
212 .open_database(&rtxn, Some("traj_meta"))?
213 .ok_or_else(|| missing("traj_meta"))?;
214 let idx_natoms = env
215 .open_database(&rtxn, Some("idx_natoms"))?
216 .ok_or_else(|| missing("idx_natoms"))?;
217 let idx_symbol = env
218 .open_database(&rtxn, Some("idx_symbol"))?
219 .ok_or_else(|| missing("idx_symbol"))?;
220 let idx_energy = env
221 .open_database(&rtxn, Some("idx_energy"))?
222 .ok_or_else(|| missing("idx_energy"))?;
223 let idx_flags = env
224 .open_database(&rtxn, Some("idx_flags"))?
225 .ok_or_else(|| missing("idx_flags"))?;
226 let idx_elem_count = env
227 .open_database(&rtxn, Some("idx_elem_count"))?
228 .ok_or_else(|| missing("idx_elem_count"))?;
229 let idx_formula = env
230 .open_database(&rtxn, Some("idx_formula"))?
231 .ok_or_else(|| missing("idx_formula"))?;
232 let idx_fmax = env
233 .open_database(&rtxn, Some("idx_fmax"))?
234 .ok_or_else(|| missing("idx_fmax"))?;
235 let idx_mass = env
236 .open_database(&rtxn, Some("idx_mass"))?
237 .ok_or_else(|| missing("idx_mass"))?;
238 let idx_volume = env
239 .open_database(&rtxn, Some("idx_volume"))?
240 .ok_or_else(|| missing("idx_volume"))?;
241 let idx_pbc = env
242 .open_database(&rtxn, Some("idx_pbc"))?
243 .ok_or_else(|| missing("idx_pbc"))?;
244 let idx_meta = env
245 .open_database(&rtxn, Some("idx_meta"))?
246 .ok_or_else(|| missing("idx_meta"))?;
247 let frame_by_hash = env
248 .open_database(&rtxn, Some("frame_by_hash"))?
249 .ok_or_else(|| missing("frame_by_hash"))?;
250 let hash_by_frame = env
251 .open_database(&rtxn, Some("hash_by_frame"))?
252 .ok_or_else(|| missing("hash_by_frame"))?;
253 let frames_soa = env
254 .open_database(&rtxn, Some("frames_soa"))?
255 .ok_or_else(|| missing("frames_soa"))?;
256 let idx_topo = env.open_database(&rtxn, Some("idx_topo"))?;
258 let topo_by_frame = env.open_database(&rtxn, Some("topo_by_frame"))?;
259 rtxn.commit()?;
261 return Ok(Self {
262 path,
263 env,
264 frames,
265 traj_meta,
266 idx_natoms,
267 idx_symbol,
268 idx_energy,
269 idx_flags,
270 idx_elem_count,
271 idx_formula,
272 idx_fmax,
273 idx_mass,
274 idx_volume,
275 idx_pbc,
276 idx_meta,
277 frame_by_hash,
278 hash_by_frame,
279 frames_soa,
280 idx_topo,
281 topo_by_frame,
282 });
283 }
284
285 let mut wtxn = env.write_txn()?;
286 let frames = env.create_database(&mut wtxn, Some("frames"))?;
287 let traj_meta = env.create_database(&mut wtxn, Some("traj_meta"))?;
288 let idx_natoms = env.create_database(&mut wtxn, Some("idx_natoms"))?;
289 let idx_symbol = env.create_database(&mut wtxn, Some("idx_symbol"))?;
290 let idx_energy = env.create_database(&mut wtxn, Some("idx_energy"))?;
291 let idx_flags = env.create_database(&mut wtxn, Some("idx_flags"))?;
292 let idx_elem_count = env.create_database(&mut wtxn, Some("idx_elem_count"))?;
293 let idx_formula = env.create_database(&mut wtxn, Some("idx_formula"))?;
294 let idx_fmax = env.create_database(&mut wtxn, Some("idx_fmax"))?;
295 let idx_mass = env.create_database(&mut wtxn, Some("idx_mass"))?;
296 let idx_volume = env.create_database(&mut wtxn, Some("idx_volume"))?;
297 let idx_pbc = env.create_database(&mut wtxn, Some("idx_pbc"))?;
298 let idx_meta = env.create_database(&mut wtxn, Some("idx_meta"))?;
299 let frame_by_hash = env.create_database(&mut wtxn, Some("frame_by_hash"))?;
300 let hash_by_frame = env.create_database(&mut wtxn, Some("hash_by_frame"))?;
301 let frames_soa = env.create_database(&mut wtxn, Some("frames_soa"))?;
302 let idx_topo = env.create_database(&mut wtxn, Some("idx_topo"))?;
303 let topo_by_frame = env.create_database(&mut wtxn, Some("topo_by_frame"))?;
304 wtxn.commit()?;
305
306 Ok(Self {
307 path,
308 env,
309 frames,
310 traj_meta,
311 idx_natoms,
312 idx_symbol,
313 idx_energy,
314 idx_flags,
315 idx_elem_count,
316 idx_formula,
317 idx_fmax,
318 idx_mass,
319 idx_volume,
320 idx_pbc,
321 idx_meta,
322 frame_by_hash,
323 hash_by_frame,
324 frames_soa,
325 idx_topo: Some(idx_topo),
326 topo_by_frame: Some(topo_by_frame),
327 })
328 }
329
330 pub fn pack_frame(&self, key: FrameKey) -> Result<Vec<u8>> {
334 if let Some(b) = self.get_cooked_soa_bytes(key)? {
335 if crate::cooked_soa::CookedSoa::try_decode(&b).is_some() {
336 return Ok(b);
337 }
338 }
339 let fr = self.get_frame(key)?;
340 crate::cooked_soa::CookedSoa::encode_frame(&fr)
341 }
342
343 pub fn pack_frames(&self, keys: &[FrameKey]) -> Result<Vec<u8>> {
345 let mut blobs = Vec::with_capacity(keys.len());
346 for k in keys {
347 blobs.push(self.pack_frame(*k)?);
348 }
349 crate::cooked_soa::encode_batch(&blobs)
350 }
351
352 pub fn path(&self) -> &Path {
353 &self.path
354 }
355
356 pub fn close(self) {
359 self.env.prepare_for_closing().wait();
360 }
361
362 pub fn snapshot_to(&self, dest_dir: impl AsRef<Path>) -> Result<()> {
365 let dest = dest_dir.as_ref();
366 fs::create_dir_all(dest)?;
367 let data = dest.join("data.mdb");
368 if data.exists() {
369 return Err(Error::Message(format!(
370 "snapshot dest exists: {}",
371 data.display()
372 )));
373 }
374 self.env.copy_to_file(&data, CompactionOption::Enabled)?;
375 Ok(())
376 }
377
378 pub fn list_frame_keys(&self) -> Result<Vec<FrameKey>> {
380 let rtxn = self.env.read_txn()?;
381 let mut out = Vec::new();
382 let mut iter = self.frames.iter(&rtxn)?;
383 while let Some(Ok((k, _))) = iter.next() {
384 if let Some(fk) = FrameKey::from_bytes(k) {
385 out.push(fk);
386 }
387 }
388 out.sort();
389 Ok(out)
390 }
391
392 pub fn export_frame_blob(&self, key: FrameKey) -> Result<String> {
395 self.get_frame_text(key)
396 }
397
398 fn build_index_puts(fk: FrameKey, frame: &ConFrame, blob: String) -> PreparedIndexPuts {
401 use readcon_core::index_proj::FrameIndexProjection;
402 let hash = hash_frame_bytes(blob.as_bytes()).to_bytes();
403 let proj = FrameIndexProjection::from_frame(frame);
404 let mut symbol_keys = Vec::new();
405 let mut elem_count_keys = Vec::new();
406 let mut seen = BTreeSet::new();
407 for (sym, cnt) in &proj.species_counts {
408 elem_count_keys.push(elem_count_key(sym, *cnt, fk));
409 if seen.insert(sym.clone()) {
410 symbol_keys.push(symbol_key(sym, fk));
411 }
412 }
413 let formula_k = if proj.formula.is_empty() {
414 None
415 } else {
416 Some(formula_key(&proj.formula, fk))
417 };
418 let mut flag_keys = Vec::new();
419 let energy_k = proj.energy.and_then(|e| {
420 flag_keys.push(flag_key(FLAG_HAS_ENERGY, fk));
421 energy_bin_key(e, fk)
422 });
423 let fmax_k = if proj.has_forces {
424 flag_keys.push(flag_key(FLAG_HAS_FORCES, fk));
425 proj.fmax.and_then(|fm| fmax_bin_key(fm, fk))
426 } else {
427 None
428 };
429 if proj.has_velocities {
430 flag_keys.push(flag_key(FLAG_HAS_VELOCITIES, fk));
431 }
432 let mass_k = proj.total_mass.and_then(|m| mass_bin_key(m, fk));
433 let volume_k = proj.cell_volume.and_then(|v| volume_bin_key(v, fk));
434 let pbc_k = proj
435 .pbc
436 .map(|p| pbc_key(crate::keys::pbc_mask_from_bools(p), fk));
437 let mut meta_keys = Vec::new();
438 for (ch, val) in [
439 (META_TIME, proj.time),
440 (META_TIMESTEP, proj.timestep),
441 (META_FRAME_INDEX, proj.frame_index),
442 (META_NEB_BEAD, proj.neb_bead),
443 (META_NEB_BAND, proj.neb_band),
444 (META_CHARGE, proj.charge),
445 (META_MAGMOM, proj.magmom),
446 ] {
447 if let Some(x) = val {
448 if let Some(k) = meta_scalar_key(ch, x, fk) {
449 meta_keys.push(k);
450 }
451 }
452 }
453 PreparedIndexPuts {
454 fk,
455 blob,
456 hash,
457 natoms_k: natoms_key(proj.n_atoms, fk),
458 symbol_keys,
459 elem_count_keys,
460 formula_k,
461 energy_k,
462 fmax_k,
463 mass_k,
464 volume_k,
465 pbc_k,
466 flag_keys,
467 meta_keys,
468 }
469 }
470
471 fn put_index_puts(&self, wtxn: &mut RwTxn, p: &PreparedIndexPuts) -> Result<()> {
473 let fk_b = p.fk.to_bytes();
474 self.frames.put(wtxn, &fk_b[..], p.blob.as_str())?;
475 self.hash_by_frame.put(wtxn, &fk_b[..], &p.hash[..])?;
476 if self.frame_by_hash.get(wtxn, &p.hash[..])?.is_none() {
477 self.frame_by_hash.put(wtxn, &p.hash[..], &fk_b[..])?;
478 }
479 self.idx_natoms.put(wtxn, &p.natoms_k[..], &())?;
480 for sk in &p.symbol_keys {
481 self.idx_symbol.put(wtxn, &sk[..], &())?;
482 }
483 for ek in &p.elem_count_keys {
484 self.idx_elem_count.put(wtxn, &ek[..], &())?;
485 }
486 if let Some(ref fk_form) = p.formula_k {
487 self.idx_formula.put(wtxn, &fk_form[..], &())?;
488 }
489 if let Some(ek) = p.energy_k {
490 self.idx_energy.put(wtxn, &ek[..], &())?;
491 }
492 if let Some(fk_fm) = p.fmax_k {
493 self.idx_fmax.put(wtxn, &fk_fm[..], &())?;
494 }
495 if let Some(k) = p.mass_k {
496 self.idx_mass.put(wtxn, &k[..], &())?;
497 }
498 if let Some(k) = p.volume_k {
499 self.idx_volume.put(wtxn, &k[..], &())?;
500 }
501 if let Some(pk) = p.pbc_k {
502 self.idx_pbc.put(wtxn, &pk[..], &())?;
503 }
504 for fk_flag in &p.flag_keys {
505 self.idx_flags.put(wtxn, &fk_flag[..], &())?;
506 }
507 for mk in &p.meta_keys {
508 self.idx_meta.put(wtxn, &mk[..], &())?;
509 }
510 Ok(())
511 }
512
513 fn delete_index_puts(&self, wtxn: &mut RwTxn, p: &PreparedIndexPuts) -> Result<()> {
514 let fk_b = p.fk.to_bytes();
515 let _ = self.frames.delete(wtxn, &fk_b[..])?;
516 let _ = self.hash_by_frame.delete(wtxn, &fk_b[..])?;
517 let _ = self.frame_by_hash.delete(wtxn, &p.hash[..])?;
518 let _ = self.frames_soa.delete(wtxn, &fk_b[..])?;
519 let _ = self.idx_natoms.delete(wtxn, &p.natoms_k[..])?;
520 for sk in &p.symbol_keys {
521 let _ = self.idx_symbol.delete(wtxn, &sk[..])?;
522 }
523 for ek in &p.elem_count_keys {
524 let _ = self.idx_elem_count.delete(wtxn, &ek[..])?;
525 }
526 if let Some(ref fk_form) = p.formula_k {
527 let _ = self.idx_formula.delete(wtxn, &fk_form[..])?;
528 }
529 if let Some(ek) = p.energy_k {
530 let _ = self.idx_energy.delete(wtxn, &ek[..])?;
531 }
532 if let Some(fk_fm) = p.fmax_k {
533 let _ = self.idx_fmax.delete(wtxn, &fk_fm[..])?;
534 }
535 if let Some(k) = p.mass_k {
536 let _ = self.idx_mass.delete(wtxn, &k[..])?;
537 }
538 if let Some(k) = p.volume_k {
539 let _ = self.idx_volume.delete(wtxn, &k[..])?;
540 }
541 if let Some(pk) = p.pbc_k {
542 let _ = self.idx_pbc.delete(wtxn, &pk[..])?;
543 }
544 for fk_flag in &p.flag_keys {
545 let _ = self.idx_flags.delete(wtxn, &fk_flag[..])?;
546 }
547 for mk in &p.meta_keys {
548 let _ = self.idx_meta.delete(wtxn, &mk[..])?;
549 }
550 Ok(())
551 }
552
553 fn index_frame(
555 &self,
556 wtxn: &mut RwTxn,
557 fk: FrameKey,
558 frame: &ConFrame,
559 blob: &str,
560 ) -> Result<()> {
561 let puts = Self::build_index_puts(fk, frame, blob.to_owned());
562 self.put_index_puts(wtxn, &puts)
563 }
564
565 fn prepare_trajectory_str(
567 traj_id: TrajId,
568 file_contents: &str,
569 ) -> Result<Vec<PreparedIndexPuts>> {
570 let mut out = Vec::new();
571 let mut frame_idx: u32 = 0;
572 let mut iter = ConFrameIterator::new(file_contents);
573 loop {
574 let (frame, blob) = match iter.next_with_raw_span(file_contents) {
575 None => break,
576 Some(Err(e)) => return Err(Error::Parse(e.to_string())),
577 Some(Ok(x)) => x,
578 };
579 let fk = FrameKey { traj_id, frame_idx };
580 out.push(Self::build_index_puts(fk, &frame, blob.to_owned()));
581 frame_idx += 1;
582 }
583 Ok(out)
584 }
585
586 fn prepare_trajectory_frames(
588 traj_id: TrajId,
589 frames: &[ConFrame],
590 start_idx: u32,
591 ) -> Result<Vec<PreparedIndexPuts>> {
592 Self::prepare_trajectory_frames_formatted(traj_id, frames, start_idx, FloatFormat::default())
593 }
594
595 fn prepare_trajectory_frames_formatted(
596 traj_id: TrajId,
597 frames: &[ConFrame],
598 start_idx: u32,
599 float_format: FloatFormat,
600 ) -> Result<Vec<PreparedIndexPuts>> {
601 let mut out = Vec::with_capacity(frames.len());
602 let mut frame_idx = start_idx;
603 for fr in frames {
604 let mut buf = Cursor::new(Vec::new());
605 {
606 let mut w = ConFrameWriter::with_float_format(&mut buf, float_format);
607 w.write_frame(fr)
608 .map_err(|e| Error::Parse(format!("serialize: {e}")))?;
609 }
610 let blob = String::from_utf8(buf.into_inner())
611 .map_err(|e| Error::Message(format!("utf8: {e}")))?;
612 let fk = FrameKey { traj_id, frame_idx };
613 out.push(Self::build_index_puts(fk, fr, blob));
614 frame_idx += 1;
615 }
616 Ok(out)
617 }
618
619 fn commit_prepared(
621 &self,
622 traj_id: TrajId,
623 prepared: &[PreparedIndexPuts],
624 source: String,
625 replace_meta: bool,
626 ) -> Result<u32> {
627 let mut wtxn = self.env.write_txn()?;
628 let tid_key = traj_id.to_be_bytes();
629 if !replace_meta {
630 if let Some(existing) = self.traj_meta.get(&wtxn, &tid_key[..])? {
631 let meta: TrajMeta = serde_json::from_str(existing)?;
632 return Err(Error::TrajExists(traj_id, meta.n_frames));
633 }
634 }
635 for p in prepared {
636 self.put_index_puts(&mut wtxn, p)?;
637 }
638 let n_frames = prepared.last().map(|p| p.fk.frame_idx + 1).unwrap_or(0);
639 let existing: Option<TrajMeta> = if replace_meta {
640 self.traj_meta
641 .get(&wtxn, &tid_key[..])?
642 .map(serde_json::from_str)
643 .transpose()?
644 } else {
645 None
646 };
647 let mut meta = existing.unwrap_or_default();
648 meta.n_frames = n_frames;
649 meta.source = source;
650 self.traj_meta
651 .put(&mut wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
652 wtxn.commit()?;
653 Ok(n_frames)
654 }
655
656 pub fn append_trajectory_path(&self, traj_id: TrajId, file: impl AsRef<Path>) -> Result<u32> {
657 self.append_trajectory_path_units(traj_id, file, None)
658 }
659
660 pub fn append_trajectory_path_units(
662 &self,
663 traj_id: TrajId,
664 file: impl AsRef<Path>,
665 units: Option<serde_json::Value>,
666 ) -> Result<u32> {
667 let text = fs::read_to_string(file.as_ref())?;
668 let source = file.as_ref().display().to_string();
669 let Some(units) = units else {
670 return self.append_trajectory_str(traj_id, &text, source);
671 };
672 let units = crate::units_canon::canonicalize_units_object(units)?;
673 let mut frames = Vec::new();
674 for item in ConFrameIterator::new(&text) {
675 let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
676 fr.header.set_units(units.clone());
677 frames.push(fr);
678 }
679 self.append_trajectory_frames(traj_id, &frames, source)
680 }
681
682 pub fn set_trajectory_units(&self, traj_id: TrajId, units: serde_json::Value) -> Result<u32> {
685 let keys = self.select(&Select::new().trajectory(traj_id))?;
686 if keys.is_empty() {
687 return Err(Error::Message("no frames for traj".into()));
688 }
689 let mut frames = Vec::with_capacity(keys.len());
690 let mut old_puts = Vec::with_capacity(keys.len());
691 for k in &keys {
692 let blob = self.get_frame_text(*k)?;
693 let mut fr = self.get_frame(*k)?;
694 old_puts.push(Self::build_index_puts(*k, &fr, blob));
695 let merged = crate::units_canon::merge_units(fr.header.units(), units.clone())?;
696 crate::units_canon::rescale_frame_units(&mut fr, &merged)?;
697 fr.header.set_units(merged);
698 frames.push(fr);
699 }
700 let prepared = Self::prepare_trajectory_frames_formatted(
702 traj_id, &frames, 0, FloatFormat::RoundTrip,
703 )?;
704 let mut wtxn = self.env.write_txn()?;
705 for p in &old_puts {
706 self.delete_index_puts(&mut wtxn, p)?;
707 }
708 for p in &prepared {
709 self.put_index_puts(&mut wtxn, p)?;
710 }
711 wtxn.commit()?;
712 Ok(keys.len() as u32)
713 }
714
715 pub fn frame_units(&self, key: FrameKey) -> Result<Option<serde_json::Value>> {
717 let fr = self.get_frame(key)?;
718 Ok(fr.header.metadata.get("units").cloned())
719 }
720
721 pub fn append_trajectory_str(
722 &self,
723 traj_id: TrajId,
724 file_contents: &str,
725 source: impl Into<String>,
726 ) -> Result<u32> {
727 let source = source.into();
728 let prepared = Self::prepare_trajectory_str(traj_id, file_contents)?;
730 self.commit_prepared(traj_id, &prepared, source, false)
731 }
732
733 pub fn append_trajectory_frames(
735 &self,
736 traj_id: TrajId,
737 frames: &[ConFrame],
738 source: impl Into<String>,
739 ) -> Result<u32> {
740 let source = source.into();
741 let prepared = Self::prepare_trajectory_frames(traj_id, frames, 0)?;
742 self.commit_prepared(traj_id, &prepared, source, false)
743 }
744
745 pub fn append_trajectory_frames_with_precision(
749 &self,
750 traj_id: TrajId,
751 frames: &[ConFrame],
752 source: impl Into<String>,
753 precision: usize,
754 ) -> Result<u32> {
755 self.append_trajectory_frames_with_float_format(
756 traj_id, frames, source, FloatFormat::DecimalPlaces(precision),
757 )
758 }
759
760 pub fn append_trajectory_frames_with_float_format(
762 &self,
763 traj_id: TrajId,
764 frames: &[ConFrame],
765 source: impl Into<String>,
766 float_format: FloatFormat,
767 ) -> Result<u32> {
768 let source = source.into();
769 let prepared = Self::prepare_trajectory_frames_formatted(
770 traj_id, frames, 0, float_format,
771 )?;
772 self.commit_prepared(traj_id, &prepared, source, false)
773 }
774
775 pub fn extend_trajectory_frames(
778 &self,
779 traj_id: TrajId,
780 frames: &[ConFrame],
781 source_hint: impl Into<String>,
782 ) -> Result<u32> {
783 let (start_idx, prior) = {
784 let rtxn = self.env.read_txn()?;
785 let tid_key = traj_id.to_be_bytes();
786 if let Some(existing) = self.traj_meta.get(&rtxn, &tid_key[..])? {
787 let meta: TrajMeta = serde_json::from_str(existing)?;
788 (meta.n_frames, Some(meta))
789 } else {
790 (0u32, None)
791 }
792 };
793 let prepared = Self::prepare_trajectory_frames(traj_id, frames, start_idx)?;
795 let mut wtxn = self.env.write_txn()?;
796 let tid_key = traj_id.to_be_bytes();
797 let live: Option<TrajMeta> =
798 if let Some(existing) = self.traj_meta.get(&wtxn, &tid_key[..])? {
799 Some(serde_json::from_str(existing)?)
800 } else {
801 None
802 };
803 let live_n = live.as_ref().map(|m| m.n_frames).unwrap_or(0);
804 if live_n != start_idx {
805 return Err(Error::Message(format!(
806 "concurrent extend race on traj {traj_id}: expected start {start_idx}, live {live_n}"
807 )));
808 }
809 for p in &prepared {
810 self.put_index_puts(&mut wtxn, p)?;
811 }
812 let n_frames = prepared
813 .last()
814 .map(|p| p.fk.frame_idx + 1)
815 .unwrap_or(start_idx);
816 let mut meta = live.or(prior).unwrap_or_else(|| TrajMeta {
817 source: source_hint.into(),
818 ..Default::default()
819 });
820 meta.n_frames = n_frames;
821 self.traj_meta
822 .put(&mut wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
823 wtxn.commit()?;
824 Ok(n_frames)
825 }
826
827 pub fn extend_trajectory_path(&self, traj_id: TrajId, file: impl AsRef<Path>) -> Result<u32> {
830 self.extend_trajectory_path_units(traj_id, file, None)
831 }
832
833 pub fn extend_trajectory_path_units(
835 &self,
836 traj_id: TrajId,
837 file: impl AsRef<Path>,
838 units: Option<serde_json::Value>,
839 ) -> Result<u32> {
840 let text = fs::read_to_string(file.as_ref())?;
841 let mut frames = Vec::new();
842 let units = match units {
843 Some(u) => Some(crate::units_canon::canonicalize_units_object(u)?),
844 None => None,
845 };
846 for item in ConFrameIterator::new(&text) {
847 let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
848 if let Some(ref u) = units {
849 fr.header.set_units(u.clone());
850 }
851 frames.push(fr);
852 }
853 self.extend_trajectory_frames(traj_id, &frames, file.as_ref().display().to_string())
854 }
855
856 fn clear_secondary(&self, wtxn: &mut RwTxn) -> Result<()> {
857 self.idx_natoms.clear(wtxn)?;
858 self.idx_symbol.clear(wtxn)?;
859 self.idx_energy.clear(wtxn)?;
860 self.idx_flags.clear(wtxn)?;
861 self.idx_elem_count.clear(wtxn)?;
862 self.idx_formula.clear(wtxn)?;
863 self.idx_fmax.clear(wtxn)?;
864 self.idx_mass.clear(wtxn)?;
865 self.idx_volume.clear(wtxn)?;
866 self.idx_pbc.clear(wtxn)?;
867 self.idx_meta.clear(wtxn)?;
868 self.frame_by_hash.clear(wtxn)?;
869 self.hash_by_frame.clear(wtxn)?;
870 Ok(())
871 }
872
873 pub fn reindex(&self) -> Result<u32> {
880 let recorded = self.collected_topo_params()?;
881 let mut wtxn = self.env.write_txn()?;
882 self.clear_secondary(&mut wtxn)?;
883 if let Some(idx) = self.idx_topo {
884 idx.clear(&mut wtxn)?;
885 }
886
887 let mut n = 0u32;
888 let mut keys: Vec<Vec<u8>> = Vec::new();
890 {
891 let mut iter = self.frames.iter(&wtxn)?;
892 while let Some(Ok((k, _))) = iter.next() {
893 keys.push(k.to_vec());
894 }
895 }
896 for k in keys {
897 let Some(fk) = FrameKey::from_bytes(&k) else {
898 continue;
899 };
900 let Some(blob) = self.frames.get(&wtxn, &k)? else {
901 continue;
902 };
903 let blob_owned = blob.to_owned();
904 let mut it = ConFrameIterator::new(&blob_owned);
905 let frame = it
906 .next()
907 .ok_or_else(|| Error::Parse("empty blob on reindex".into()))?
908 .map_err(|e| Error::Parse(e.to_string()))?;
909 let hash = hash_frame_bytes(blob_owned.as_bytes());
911 let hb = hash.to_bytes();
912 let fk_b = fk.to_bytes();
913 self.hash_by_frame.put(&mut wtxn, &fk_b[..], &hb[..])?;
914 if self.frame_by_hash.get(&wtxn, &hb[..])?.is_none() {
915 self.frame_by_hash.put(&mut wtxn, &hb[..], &fk_b[..])?;
916 }
917 let n_atoms = frame.atom_data.len() as u32;
918 self.idx_natoms
919 .put(&mut wtxn, &natoms_key(n_atoms, fk)[..], &())?;
920 let counts = frame_species(&frame);
921 for (sym, cnt) in &counts {
922 self.idx_elem_count
923 .put(&mut wtxn, &elem_count_key(sym, *cnt, fk)[..], &())?;
924 self.idx_symbol
925 .put(&mut wtxn, &symbol_key(sym, fk)[..], &())?;
926 }
927 let formula = composition_formula(&counts);
928 if !formula.is_empty() {
929 self.idx_formula
930 .put(&mut wtxn, &formula_key(&formula, fk)[..], &())?;
931 }
932 if let Some(e) = frame_energy(&frame) {
933 if let Some(ek) = energy_bin_key(e, fk) {
934 self.idx_energy.put(&mut wtxn, &ek[..], &())?;
935 }
936 self.idx_flags
937 .put(&mut wtxn, &flag_key(FLAG_HAS_ENERGY, fk)[..], &())?;
938 }
939 if frame_has_forces(&frame) {
940 self.idx_flags
941 .put(&mut wtxn, &flag_key(FLAG_HAS_FORCES, fk)[..], &())?;
942 if let Some(fm) = frame_fmax(&frame) {
943 if let Some(fk_fm) = fmax_bin_key(fm, fk) {
944 self.idx_fmax.put(&mut wtxn, &fk_fm[..], &())?;
945 }
946 }
947 }
948 if frame_has_velocities(&frame) {
949 self.idx_flags
950 .put(&mut wtxn, &flag_key(FLAG_HAS_VELOCITIES, fk)[..], &())?;
951 }
952 if let Some(m) = frame_total_mass(&frame) {
953 if let Some(k) = mass_bin_key(m, fk) {
954 self.idx_mass.put(&mut wtxn, &k[..], &())?;
955 }
956 }
957 if let Some(v) = frame_cell_volume(&frame) {
958 if let Some(k) = volume_bin_key(v, fk) {
959 self.idx_volume.put(&mut wtxn, &k[..], &())?;
960 }
961 }
962 if let Some(mask) = frame_pbc_mask(&frame) {
963 self.idx_pbc.put(&mut wtxn, &pbc_key(mask, fk)[..], &())?;
964 }
965 let meta_puts: [(u8, Option<f64>); 7] = [
966 (META_TIME, frame_time(&frame)),
967 (META_TIMESTEP, frame_timestep(&frame)),
968 (META_FRAME_INDEX, frame_frame_index(&frame)),
969 (META_NEB_BEAD, frame_neb_bead(&frame)),
970 (META_NEB_BAND, frame_neb_band(&frame)),
971 (META_CHARGE, frame_charge(&frame)),
972 (META_MAGMOM, frame_magmom(&frame)),
973 ];
974 for (ch, val) in meta_puts {
975 if let Some(x) = val {
976 if let Some(k) = meta_scalar_key(ch, x, fk) {
977 self.idx_meta.put(&mut wtxn, &k[..], &())?;
978 }
979 }
980 }
981 n += 1;
982 }
983 if recorded.is_some() {
984 self.rebuild_idx_topo_from_stored(&mut wtxn)?;
985 }
986 wtxn.commit()?;
987 Ok(n)
988 }
989
990 pub fn traj_meta(&self, traj_id: TrajId) -> Result<Option<TrajMeta>> {
991 let rtxn = self.env.read_txn()?;
992 let tid_key = traj_id.to_be_bytes();
993 match self.traj_meta.get(&rtxn, &tid_key[..])? {
994 None => Ok(None),
995 Some(s) => Ok(Some(serde_json::from_str(s)?)),
996 }
997 }
998
999 fn list_traj_metas(&self) -> Result<Vec<(TrajId, TrajMeta)>> {
1000 let rtxn = self.env.read_txn()?;
1001 let mut out = Vec::new();
1002 let mut iter = self.traj_meta.iter(&rtxn)?;
1003 while let Some(Ok((k, v))) = iter.next() {
1004 if k.len() != 8 {
1005 continue;
1006 }
1007 let mut tb = [0u8; 8];
1008 tb.copy_from_slice(k);
1009 out.push((u64::from_be_bytes(tb), serde_json::from_str(v)?));
1010 }
1011 Ok(out)
1012 }
1013
1014 fn collected_topo_params(&self) -> Result<Option<TopologyParams>> {
1015 let mut seen: Option<TopologyParams> = None;
1016 for (_, meta) in self.list_traj_metas()? {
1017 if !meta.has_topo() {
1018 continue;
1019 }
1020 let rec = meta.topo_params();
1021 match &seen {
1022 None => seen = Some(rec),
1023 Some(prev) if prev.agrees(&rec) => {}
1024 Some(prev) => return Err(mixed_topo_error(prev, &rec)),
1025 }
1026 }
1027 Ok(seen)
1028 }
1029
1030 fn require_topo_dbs(&self) -> Result<(Database<Bytes, Unit>, Database<Bytes, Bytes>)> {
1031 match (self.idx_topo, self.topo_by_frame) {
1032 (Some(idx), Some(by)) => Ok((idx, by)),
1033 _ => Err(Error::Message(
1034 "topology indexes missing; open the corpus for write to create idx_topo".into(),
1035 )),
1036 }
1037 }
1038
1039 fn rebuild_idx_topo_from_stored(&self, wtxn: &mut RwTxn) -> Result<u32> {
1040 let (idx_topo, topo_by_frame) = self.require_topo_dbs()?;
1041 let mut pairs: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
1042 {
1043 let mut iter = topo_by_frame.iter(wtxn)?;
1044 while let Some(Ok((k, v))) = iter.next() {
1045 pairs.push((k.to_vec(), v.to_vec()));
1046 }
1047 }
1048 let mut n = 0u32;
1049 for (fk_b, hex_b) in pairs {
1050 let Some(fk) = FrameKey::from_bytes(&fk_b) else {
1051 continue;
1052 };
1053 let hex = std::str::from_utf8(&hex_b)?;
1054 let hex = normalize_topo_hex(hex)?;
1055 idx_topo.put(wtxn, &topo_key(&hex, fk)[..], &())?;
1056 n += 1;
1057 }
1058 Ok(n)
1059 }
1060
1061 fn put_topo_entry(
1062 &self,
1063 wtxn: &mut RwTxn,
1064 idx_topo: Database<Bytes, Unit>,
1065 topo_by_frame: Database<Bytes, Bytes>,
1066 fk: FrameKey,
1067 hex: &str,
1068 ) -> Result<()> {
1069 let fk_b = fk.to_bytes();
1070 if let Some(old) = topo_by_frame.get(wtxn, &fk_b[..])? {
1071 let old_hex = std::str::from_utf8(old)?;
1072 let _ = idx_topo.delete(wtxn, &topo_key(old_hex, fk)[..])?;
1073 }
1074 topo_by_frame.put(wtxn, &fk_b[..], hex.as_bytes())?;
1075 idx_topo.put(wtxn, &topo_key(hex, fk)[..], &())?;
1076 Ok(())
1077 }
1078
1079 pub fn frame_topo_key(&self, key: FrameKey) -> Result<Option<String>> {
1081 let Some(by) = self.topo_by_frame else {
1082 return Ok(None);
1083 };
1084 let rtxn = self.env.read_txn()?;
1085 let fk_b = key.to_bytes();
1086 match by.get(&rtxn, &fk_b[..])? {
1087 None => Ok(None),
1088 Some(b) => Ok(Some(std::str::from_utf8(b)?.to_owned())),
1089 }
1090 }
1091
1092 pub fn annotate_topology(&self, opts: AnnotateTopologyOpts) -> Result<u32> {
1099 let seams = resolve_seams_binary(opts.seams.as_deref())?;
1100 let (idx_topo, topo_by_frame) = self.require_topo_dbs()?;
1101 let recorded = self.collected_topo_params()?;
1102 if let Some(ref rec) = recorded {
1103 let cutoff_diff = rec.cutoff.is_some() && rec.cutoff != Some(opts.cutoff);
1104 let graph_diff = rec.graph.as_deref().is_some_and(|g| g != opts.graph);
1105 let hops_diff = rec.hops.is_some_and(|h| h != opts.hops);
1106 if cutoff_diff || graph_diff || hops_diff {
1107 let incoming = TopologyParams {
1108 cutoff: Some(opts.cutoff),
1109 graph: Some(opts.graph.clone()),
1110 hops: Some(opts.hops),
1111 method: rec.method.clone(),
1112 };
1113 return Err(mixed_topo_error(rec, &incoming));
1114 }
1115 }
1116 let metas = self.list_traj_metas()?;
1117 let mut pending: Vec<(FrameKey, String)> = Vec::new();
1118 let mut run_method: Option<String> = None;
1119 for (tid, meta) in &metas {
1120 if meta.n_frames == 0 {
1121 continue;
1122 }
1123 let mut blob = String::new();
1124 for frame_idx in 0..meta.n_frames {
1125 blob.push_str(&self.get_frame_text(FrameKey {
1126 traj_id: *tid,
1127 frame_idx,
1128 })?);
1129 }
1130 let recs = fingerprint_con_bytes(
1131 &seams,
1132 blob.as_bytes(),
1133 opts.cutoff,
1134 &opts.graph,
1135 opts.hops,
1136 )?;
1137 if recs.len() != meta.n_frames as usize {
1138 return Err(Error::Message(format!(
1139 "seams fingerprint returned {} objects for traj {tid} with {} frames",
1140 recs.len(),
1141 meta.n_frames
1142 )));
1143 }
1144 for (i, rec) in recs.into_iter().enumerate() {
1145 match &run_method {
1146 None => run_method = Some(rec.method.clone()),
1147 Some(m) if m == &rec.method => {}
1148 Some(m) => {
1149 return Err(Error::Message(format!(
1150 "seams fingerprint mixed methods in one run ({m} vs {})",
1151 rec.method
1152 )));
1153 }
1154 }
1155 pending.push((
1156 FrameKey {
1157 traj_id: *tid,
1158 frame_idx: i as u32,
1159 },
1160 rec.key,
1161 ));
1162 }
1163 }
1164 if let (Some(rec), Some(method)) = (recorded.as_ref(), run_method.as_deref()) {
1165 if rec.method.as_deref().is_some_and(|m| m != method) {
1166 let incoming = TopologyParams {
1167 cutoff: Some(opts.cutoff),
1168 graph: Some(opts.graph.clone()),
1169 hops: Some(opts.hops),
1170 method: Some(method.to_owned()),
1171 };
1172 return Err(mixed_topo_error(rec, &incoming));
1173 }
1174 }
1175 let Some(method) = run_method else {
1176 return Ok(0);
1177 };
1178 let mut wtxn = self.env.write_txn()?;
1179 for (fk, hex) in &pending {
1180 self.put_topo_entry(&mut wtxn, idx_topo, topo_by_frame, *fk, hex)?;
1181 }
1182 for (tid, mut meta) in metas {
1183 meta.stamp_topo(opts.cutoff, &opts.graph, opts.hops, &method);
1184 let tid_key = tid.to_be_bytes();
1185 self.traj_meta
1186 .put(&mut wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
1187 }
1188 wtxn.commit()?;
1189 Ok(pending.len() as u32)
1190 }
1191
1192 pub fn find_by_topology_path(&self, file: impl AsRef<Path>) -> Result<Vec<FrameKey>> {
1195 let rec = self.collected_topo_params()?.ok_or_else(|| {
1196 Error::Message(
1197 "no recorded topology parameters on this corpus; run annotate-topology first"
1198 .into(),
1199 )
1200 })?;
1201 let cutoff = rec.cutoff_or_err()?;
1202 let graph = rec.graph_or_default();
1203 let hops = rec.hops_or_default();
1204 let seams = resolve_seams_binary(None)?;
1205 let recs = run_seams_fingerprint(&seams, file.as_ref(), cutoff, &graph, hops)?;
1206 self.lookup_topo_keys(recs.into_iter().map(|r| r.key))
1207 }
1208
1209 pub fn find_by_topology_text(&self, text: &str) -> Result<Vec<FrameKey>> {
1211 let rec = self.collected_topo_params()?.ok_or_else(|| {
1212 Error::Message(
1213 "no recorded topology parameters on this corpus; run annotate-topology first"
1214 .into(),
1215 )
1216 })?;
1217 let cutoff = rec.cutoff_or_err()?;
1218 let graph = rec.graph_or_default();
1219 let hops = rec.hops_or_default();
1220 let seams = resolve_seams_binary(None)?;
1221 let recs = fingerprint_con_bytes(&seams, text.as_bytes(), cutoff, &graph, hops)?;
1222 self.lookup_topo_keys(recs.into_iter().map(|r| r.key))
1223 }
1224
1225 fn lookup_topo_keys<I>(&self, hexes: I) -> Result<Vec<FrameKey>>
1226 where
1227 I: IntoIterator<Item = String>,
1228 {
1229 let mut out = BTreeSet::new();
1230 let Some(idx) = self.idx_topo else {
1231 return Ok(Vec::new());
1232 };
1233 let rtxn = self.env.read_txn()?;
1234 for hex in hexes {
1235 let pref = topo_prefix(&normalize_topo_hex(&hex)?);
1236 let mut iter = idx.prefix_iter(&rtxn, &pref)?;
1237 while let Some(Ok((k, _))) = iter.next() {
1238 if k.len() < 12 {
1239 continue;
1240 }
1241 if let Some(fk) = FrameKey::from_bytes(&k[k.len() - 12..]) {
1242 out.insert(fk);
1243 }
1244 }
1245 }
1246 Ok(out.into_iter().collect())
1247 }
1248
1249 pub fn get_frame_text(&self, key: FrameKey) -> Result<String> {
1250 let rtxn = self.env.read_txn()?;
1251 self.get_frame_text_txn(&rtxn, key)
1252 }
1253
1254 fn get_frame_text_txn(&self, rtxn: &heed::RoTxn<'_>, key: FrameKey) -> Result<String> {
1255 let fk_b = key.to_bytes();
1256 match self.frames.get(rtxn, &fk_b[..])? {
1257 Some(s) => Ok(s.to_owned()),
1258 None => Err(Error::MissingFrame(key)),
1259 }
1260 }
1261
1262 pub fn get_frame_texts(&self, keys: &[FrameKey]) -> Result<Vec<String>> {
1264 let rtxn = self.env.read_txn()?;
1265 let mut out = Vec::with_capacity(keys.len());
1266 for &key in keys {
1267 out.push(self.get_frame_text_txn(&rtxn, key)?);
1268 }
1269 Ok(out)
1270 }
1271
1272 pub fn touch_trajectory_blobs(&self, traj_id: TrajId, n_frames: u32) -> Result<(u64, u64)> {
1278 let rtxn = self.env.read_txn()?;
1279 let mut total = 0u64;
1280 let mut checksum = 0u64;
1281 for frame_idx in 0..n_frames {
1282 let key = FrameKey { traj_id, frame_idx };
1283 let owned = self.get_frame_text_txn(&rtxn, key)?;
1285 total += owned.len() as u64;
1286 for (i, b) in owned.as_bytes().iter().enumerate() {
1287 checksum =
1288 checksum.wrapping_add((*b as u64).wrapping_mul((i as u64).wrapping_add(1)));
1289 }
1290 }
1291 std::hint::black_box(checksum);
1293 Ok((total, checksum))
1294 }
1295
1296 pub fn get_frame(&self, key: FrameKey) -> Result<ConFrame> {
1297 let text = self.get_frame_text(key)?;
1298 let mut it = ConFrameIterator::new(&text);
1299 let fr = it
1300 .next()
1301 .ok_or_else(|| Error::Parse("empty blob".into()))?
1302 .map_err(|e| Error::Parse(e.to_string()))?;
1303 Ok(fr)
1304 }
1305
1306 pub fn get_cooked_soa_bytes(&self, key: FrameKey) -> Result<Option<Vec<u8>>> {
1309 let rtxn = self.env.read_txn()?;
1310 let fk_b = key.to_bytes();
1311 Ok(self.frames_soa.get(&rtxn, &fk_b[..])?.map(|b| b.to_vec()))
1312 }
1313
1314 pub fn has_valid_cooked_soa(&self, key: FrameKey) -> Result<bool> {
1317 Ok(self.get_cooked_soa(key)?.is_some())
1318 }
1319
1320 pub fn get_cooked_soa(&self, key: FrameKey) -> Result<Option<crate::cooked_soa::CookedSoa>> {
1323 Ok(self
1324 .get_cooked_soa_bytes(key)?
1325 .and_then(|b| crate::cooked_soa::CookedSoa::try_decode(&b)))
1326 }
1327
1328 pub fn get_positions(&self, key: FrameKey) -> Result<Vec<[f64; 3]>> {
1331 if let Some(c) = self.get_cooked_soa(key)? {
1333 return Ok(c.positions);
1334 }
1335 let fr = self.get_frame(key)?;
1337 Ok(fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect())
1338 }
1339
1340 pub fn get_forces(&self, key: FrameKey) -> Result<Option<Vec<[f64; 3]>>> {
1343 if let Some(c) = self.get_cooked_soa(key)? {
1344 if c.forces.is_some() {
1345 return Ok(c.forces);
1346 }
1347 }
1349 let fr = self.get_frame(key)?;
1350 if !fr.atom_data.iter().any(|a| a.force.is_some()) {
1351 return Ok(None);
1352 }
1353 Ok(Some(
1354 fr.atom_data
1355 .iter()
1356 .map(|a| a.force.unwrap_or([0.0; 3]))
1357 .collect(),
1358 ))
1359 }
1360
1361 pub fn get_velocities(&self, key: FrameKey) -> Result<Option<Vec<[f64; 3]>>> {
1363 if let Some(c) = self.get_cooked_soa(key)? {
1364 if c.velocities.is_some() {
1365 return Ok(c.velocities);
1366 }
1367 }
1368 let fr = self.get_frame(key)?;
1369 if !fr.atom_data.iter().any(|a| a.velocity.is_some()) {
1370 return Ok(None);
1371 }
1372 Ok(Some(
1373 fr.atom_data
1374 .iter()
1375 .map(|a| a.velocity.unwrap_or([0.0; 3]))
1376 .collect(),
1377 ))
1378 }
1379
1380 pub fn cook_frame(&self, key: FrameKey) -> Result<usize> {
1382 let fr = self.get_frame(key)?;
1383 let bytes = crate::cooked_soa::CookedSoa::encode_frame(&fr)?;
1384 let mut wtxn = self.env.write_txn()?;
1385 let fk_b = key.to_bytes();
1386 self.frames_soa.put(&mut wtxn, &fk_b[..], &bytes[..])?;
1387 wtxn.commit()?;
1388 Ok(bytes.len())
1389 }
1390
1391 pub fn recook_all(&self) -> Result<u32> {
1393 let keys = self.list_frame_keys()?;
1394 let mut n = 0u32;
1395 for k in keys {
1396 self.cook_frame(k)?;
1397 n += 1;
1398 }
1399 Ok(n)
1400 }
1401
1402 pub fn delete_cooked_soa(&self, key: FrameKey) -> Result<()> {
1404 let mut wtxn = self.env.write_txn()?;
1405 let fk_b = key.to_bytes();
1406 let _ = self.frames_soa.delete(&mut wtxn, &fk_b[..])?;
1407 wtxn.commit()?;
1408 Ok(())
1409 }
1410
1411 pub fn append_trajectory_path_cook(
1413 &self,
1414 traj_id: TrajId,
1415 file: impl AsRef<Path>,
1416 cook: bool,
1417 ) -> Result<u32> {
1418 let n = self.append_trajectory_path(traj_id, file)?;
1419 if cook {
1420 for i in 0..n {
1421 self.cook_frame(FrameKey {
1422 traj_id,
1423 frame_idx: i,
1424 })?;
1425 }
1426 }
1427 Ok(n)
1428 }
1429
1430 pub fn frame_hash(&self, key: FrameKey) -> Result<ContentHash> {
1431 let rtxn = self.env.read_txn()?;
1432 let fk_b = key.to_bytes();
1433 match self.hash_by_frame.get(&rtxn, &fk_b[..])? {
1434 Some(b) => ContentHash::from_bytes(b).ok_or_else(|| Error::Message("bad hash".into())),
1435 None => {
1436 let text = self.get_frame_text(key)?;
1437 Ok(hash_frame_bytes(text.as_bytes()))
1438 }
1439 }
1440 }
1441
1442 pub fn find_by_hash(&self, hash: ContentHash) -> Result<Option<FrameKey>> {
1443 let rtxn = self.env.read_txn()?;
1444 let hb = hash.to_bytes();
1445 match self.frame_by_hash.get(&rtxn, &hb[..])? {
1446 Some(b) => Ok(FrameKey::from_bytes(b)),
1447 None => Ok(None),
1448 }
1449 }
1450
1451 pub fn hash_con_text(text: &str) -> Result<ContentHash> {
1452 let mut it = ConFrameIterator::new(text);
1453 let frame = it
1454 .next()
1455 .ok_or_else(|| Error::Parse("no frame".into()))?
1456 .map_err(|e| Error::Parse(e.to_string()))?;
1457 let mut buf = Cursor::new(Vec::new());
1458 {
1459 let mut w = ConFrameWriter::new(&mut buf);
1460 w.write_frame(&frame)
1461 .map_err(|e| Error::Parse(format!("re-serialize: {e}")))?;
1462 }
1463 Ok(hash_frame_bytes(&buf.into_inner()))
1464 }
1465
1466 pub fn frame_formula(&self, key: FrameKey) -> Result<String> {
1468 let fr = self.get_frame(key)?;
1469 Ok(composition_formula(&frame_species(&fr)))
1470 }
1471
1472 pub fn select(&self, sel: &Select) -> Result<Vec<FrameKey>> {
1473 let rtxn = self.env.read_txn()?;
1474 let mut sets: Vec<BTreeSet<FrameKey>> = Vec::new();
1475
1476 if let Some(h) = sel.content_hash {
1477 let mut s = BTreeSet::new();
1478 if let Some(b) = self.frame_by_hash.get(&rtxn, &h[..])? {
1479 if let Some(fk) = FrameKey::from_bytes(b) {
1480 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1481 s.insert(fk);
1482 }
1483 }
1484 }
1485 sets.push(s);
1486 }
1487
1488 if sel.natoms_min.is_some() || sel.natoms_max.is_some() {
1489 let lo = sel.natoms_min.unwrap_or(0);
1490 let hi = sel.natoms_max.unwrap_or(u32::MAX);
1491 let mut s = BTreeSet::new();
1492 let mut iter = self.idx_natoms.iter(&rtxn)?;
1493 while let Some(Ok((k, _))) = iter.next() {
1494 if k.len() < 16 {
1495 continue;
1496 }
1497 let mut nb = [0u8; 4];
1498 nb.copy_from_slice(&k[..4]);
1499 let n = u32::from_be_bytes(nb);
1500 if n > hi {
1501 break;
1502 }
1503 if n >= lo {
1504 if let Some(fk) = FrameKey::from_bytes(&k[4..16]) {
1505 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1506 s.insert(fk);
1507 }
1508 }
1509 }
1510 }
1511 sets.push(s);
1512 }
1513
1514 for sym in &sel.symbols_all {
1515 let mut s = BTreeSet::new();
1516 let pref = symbol_prefix(sym);
1517 let mut iter = self.idx_symbol.prefix_iter(&rtxn, &pref)?;
1518 while let Some(Ok((k, _))) = iter.next() {
1519 if k.len() < 12 {
1520 continue;
1521 }
1522 let fk_bytes = &k[k.len() - 12..];
1523 if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
1524 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1525 s.insert(fk);
1526 }
1527 }
1528 }
1529 sets.push(s);
1530 }
1531
1532 if let Some(ref formula) = sel.exact_formula {
1533 let mut s = BTreeSet::new();
1534 let pref = formula_prefix(formula);
1535 let mut iter = self.idx_formula.prefix_iter(&rtxn, &pref)?;
1536 while let Some(Ok((k, _))) = iter.next() {
1537 if k.len() < 12 {
1538 continue;
1539 }
1540 let fk_bytes = &k[k.len() - 12..];
1541 if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
1542 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1543 s.insert(fk);
1544 }
1545 }
1546 }
1547 sets.push(s);
1548 }
1549
1550 if let Some(ref hex) = sel.topo_key {
1551 let mut s = BTreeSet::new();
1552 if let Some(idx) = self.idx_topo {
1553 let pref = topo_prefix(&normalize_topo_hex(hex)?);
1554 let mut iter = idx.prefix_iter(&rtxn, &pref)?;
1555 while let Some(Ok((k, _))) = iter.next() {
1556 if k.len() < 12 {
1557 continue;
1558 }
1559 let fk_bytes = &k[k.len() - 12..];
1560 if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
1561 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1562 s.insert(fk);
1563 }
1564 }
1565 }
1566 }
1567 sets.push(s);
1568 }
1569
1570 for (sym, exact_c) in &sel.element_count_exact {
1571 let mut s = BTreeSet::new();
1572 let pref = elem_count_symbol_prefix(sym);
1573 let mut iter = self.idx_elem_count.prefix_iter(&rtxn, &pref)?;
1574 while let Some(Ok((k, _))) = iter.next() {
1575 if let Some((c, fk)) = parse_elem_count_key(k, sym) {
1576 if c == *exact_c && sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1577 s.insert(fk);
1578 }
1579 }
1580 }
1581 sets.push(s);
1582 }
1583
1584 for (sym, min_c) in &sel.element_count_min {
1585 let mut s = BTreeSet::new();
1586 let pref = elem_count_symbol_prefix(sym);
1587 let mut iter = self.idx_elem_count.prefix_iter(&rtxn, &pref)?;
1588 while let Some(Ok((k, _))) = iter.next() {
1589 if let Some((c, fk)) = parse_elem_count_key(k, sym) {
1590 if c >= *min_c && sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1591 s.insert(fk);
1592 }
1593 }
1594 }
1595 sets.push(s);
1596 }
1597
1598 if sel.energy_min.is_some() || sel.energy_max.is_some() {
1599 let lo_e = sel.energy_min.unwrap_or(f64::NEG_INFINITY);
1600 let hi_e = sel.energy_max.unwrap_or(f64::INFINITY);
1601 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
1602 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
1603 let mut s = BTreeSet::new();
1604 let mut iter = self.idx_energy.iter(&rtxn)?;
1605 while let Some(Ok((k, _))) = iter.next() {
1606 if k.len() < 20 {
1607 continue;
1608 }
1609 let mut eb = [0u8; 8];
1610 eb.copy_from_slice(&k[..8]);
1611 let bits = u64::from_be_bytes(eb);
1612 if bits > hi_bits {
1613 break;
1614 }
1615 if bits >= lo_bits {
1616 if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
1617 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1618 s.insert(fk);
1619 }
1620 }
1621 }
1622 }
1623 sets.push(s);
1624 }
1625
1626 if sel.fmax_min.is_some() || sel.fmax_max.is_some() {
1627 let lo_e = sel.fmax_min.unwrap_or(0.0);
1628 let hi_e = sel.fmax_max.unwrap_or(f64::INFINITY);
1629 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
1630 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
1631 let mut s = BTreeSet::new();
1632 let mut iter = self.idx_fmax.iter(&rtxn)?;
1633 while let Some(Ok((k, _))) = iter.next() {
1634 if k.len() < 20 {
1635 continue;
1636 }
1637 let mut eb = [0u8; 8];
1638 eb.copy_from_slice(&k[..8]);
1639 let bits = u64::from_be_bytes(eb);
1640 if bits > hi_bits {
1641 break;
1642 }
1643 if bits >= lo_bits {
1644 if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
1645 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1646 s.insert(fk);
1647 }
1648 }
1649 }
1650 }
1651 sets.push(s);
1652 }
1653
1654 let scan_ord = |db: Database<Bytes, Unit>,
1655 lo: Option<f64>,
1656 hi: Option<f64>|
1657 -> Result<BTreeSet<FrameKey>> {
1658 let lo_e = lo.unwrap_or(f64::NEG_INFINITY);
1659 let hi_e = hi.unwrap_or(f64::INFINITY);
1660 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
1661 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
1662 let mut s = BTreeSet::new();
1663 let mut iter = db.iter(&rtxn)?;
1664 while let Some(Ok((k, _))) = iter.next() {
1665 if k.len() < 20 {
1666 continue;
1667 }
1668 let mut eb = [0u8; 8];
1669 eb.copy_from_slice(&k[..8]);
1670 let bits = u64::from_be_bytes(eb);
1671 if bits > hi_bits {
1672 break;
1673 }
1674 if bits >= lo_bits {
1675 if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
1676 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1677 s.insert(fk);
1678 }
1679 }
1680 }
1681 }
1682 Ok(s)
1683 };
1684
1685 if sel.mass_min.is_some() || sel.mass_max.is_some() {
1686 sets.push(scan_ord(self.idx_mass, sel.mass_min, sel.mass_max)?);
1687 }
1688 if sel.volume_min.is_some() || sel.volume_max.is_some() {
1689 sets.push(scan_ord(self.idx_volume, sel.volume_min, sel.volume_max)?);
1690 }
1691
1692 if let Some(pbc) = sel.pbc_exact {
1693 let mask = pbc_mask_from_bools(pbc);
1694 let mut s = BTreeSet::new();
1695 let pref = [mask];
1696 let mut iter = self.idx_pbc.prefix_iter(&rtxn, &pref)?;
1697 while let Some(Ok((k, _))) = iter.next() {
1698 if k.len() < 13 {
1699 continue;
1700 }
1701 if let Some(fk) = FrameKey::from_bytes(&k[1..13]) {
1702 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1703 s.insert(fk);
1704 }
1705 }
1706 }
1707 sets.push(s);
1708 }
1709
1710 let scan_meta = |ch: u8, lo: Option<f64>, hi: Option<f64>| -> Result<BTreeSet<FrameKey>> {
1711 let lo_e = lo.unwrap_or(f64::NEG_INFINITY);
1712 let hi_e = hi.unwrap_or(f64::INFINITY);
1713 let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
1714 let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
1715 let mut s = BTreeSet::new();
1716 let pref = [ch];
1717 let mut iter = self.idx_meta.prefix_iter(&rtxn, &pref)?;
1718 while let Some(Ok((k, _))) = iter.next() {
1719 if k.len() < 21 {
1720 continue;
1721 }
1722 let mut eb = [0u8; 8];
1723 eb.copy_from_slice(&k[1..9]);
1724 let bits = u64::from_be_bytes(eb);
1725 if bits > hi_bits {
1726 continue; }
1728 if bits >= lo_bits {
1729 if let Some(fk) = FrameKey::from_bytes(&k[9..21]) {
1730 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1731 s.insert(fk);
1732 }
1733 }
1734 }
1735 }
1736 Ok(s)
1737 };
1738
1739 if sel.time_min.is_some() || sel.time_max.is_some() {
1740 sets.push(scan_meta(META_TIME, sel.time_min, sel.time_max)?);
1741 }
1742 if sel.timestep_min.is_some() || sel.timestep_max.is_some() {
1743 sets.push(scan_meta(
1744 META_TIMESTEP,
1745 sel.timestep_min,
1746 sel.timestep_max,
1747 )?);
1748 }
1749 if sel.frame_index_min.is_some() || sel.frame_index_max.is_some() {
1750 sets.push(scan_meta(
1751 META_FRAME_INDEX,
1752 sel.frame_index_min,
1753 sel.frame_index_max,
1754 )?);
1755 }
1756 if sel.neb_bead_min.is_some() || sel.neb_bead_max.is_some() {
1757 sets.push(scan_meta(
1758 META_NEB_BEAD,
1759 sel.neb_bead_min,
1760 sel.neb_bead_max,
1761 )?);
1762 }
1763 if sel.neb_band_min.is_some() || sel.neb_band_max.is_some() {
1764 sets.push(scan_meta(
1765 META_NEB_BAND,
1766 sel.neb_band_min,
1767 sel.neb_band_max,
1768 )?);
1769 }
1770 if sel.charge_min.is_some() || sel.charge_max.is_some() {
1771 sets.push(scan_meta(META_CHARGE, sel.charge_min, sel.charge_max)?);
1772 }
1773 if sel.magmom_min.is_some() || sel.magmom_max.is_some() {
1774 sets.push(scan_meta(META_MAGMOM, sel.magmom_min, sel.magmom_max)?);
1775 }
1776
1777 let mut push_flag = |flag: u8| -> Result<()> {
1778 let mut s = BTreeSet::new();
1779 let pref = [flag];
1780 let mut iter = self.idx_flags.prefix_iter(&rtxn, &pref)?;
1781 while let Some(Ok((k, _))) = iter.next() {
1782 if k.len() < 13 {
1783 continue;
1784 }
1785 if let Some(fk) = FrameKey::from_bytes(&k[1..13]) {
1786 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1787 s.insert(fk);
1788 }
1789 }
1790 }
1791 sets.push(s);
1792 Ok(())
1793 };
1794 if sel.require_forces {
1795 push_flag(FLAG_HAS_FORCES)?;
1796 }
1797 if sel.require_velocities {
1798 push_flag(FLAG_HAS_VELOCITIES)?;
1799 }
1800 if sel.require_energy {
1801 push_flag(FLAG_HAS_ENERGY)?;
1802 }
1803
1804 if sets.is_empty() {
1805 let mut s = BTreeSet::new();
1806 let mut iter = self.frames.iter(&rtxn)?;
1807 while let Some(Ok((k, _))) = iter.next() {
1808 if let Some(fk) = FrameKey::from_bytes(k) {
1809 if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
1810 s.insert(fk);
1811 }
1812 }
1813 }
1814 sets.push(s);
1815 } else if let Some(tid) = sel.traj_id {
1816 for set in &mut sets {
1817 set.retain(|fk| fk.traj_id == tid);
1818 }
1819 }
1820
1821 sets.sort_by_key(|s| s.len());
1823 let mut acc = sets.remove(0);
1824 for s in sets {
1825 if acc.is_empty() {
1826 break;
1827 }
1828 if s.len() < acc.len() {
1829 acc = s.intersection(&acc).copied().collect();
1830 } else {
1831 acc = acc.intersection(&s).copied().collect();
1832 }
1833 }
1834
1835 let mut out: Vec<FrameKey> = acc.into_iter().collect();
1836 out.sort_unstable();
1837 if let Some(lim) = sel.limit {
1838 out.truncate(lim);
1839 }
1840 Ok(out)
1841 }
1842
1843 pub fn export_extxyz(
1844 &self,
1845 keys: &[FrameKey],
1846 path: impl AsRef<Path>,
1847 energy_key: &str,
1848 ) -> Result<usize> {
1849 use std::fs::OpenOptions;
1850 use std::io::BufWriter;
1851 let path = path.as_ref();
1852 if path.exists() {
1853 return Err(Error::Message(format!(
1854 "export_extxyz dest exists: {}",
1855 path.display()
1856 )));
1857 }
1858 let file = OpenOptions::new().write(true).create_new(true).open(path)?;
1859 let written = (|| -> Result<usize> {
1860 use std::io::Write;
1861 let mut w = BufWriter::new(file);
1862 let mut n = 0usize;
1863 for k in keys {
1864 let fr = self.get_frame(*k)?;
1865 write_frame_extxyz(&mut w, &fr, energy_key)?;
1866 n += 1;
1867 }
1868 w.flush()?;
1869 Ok(n)
1870 })();
1871 if written.is_err() {
1872 let _ = std::fs::remove_file(path);
1873 }
1874 written
1875 }
1876
1877 pub fn ingest_directory(
1878 &self,
1879 dir: impl AsRef<Path>,
1880 start_traj_id: TrajId,
1881 ) -> Result<Vec<(TrajId, u32, String)>> {
1882 self.ingest_directory_units(dir, start_traj_id, None)
1883 }
1884
1885 pub fn ingest_directory_units(
1887 &self,
1888 dir: impl AsRef<Path>,
1889 start_traj_id: TrajId,
1890 units: Option<serde_json::Value>,
1891 ) -> Result<Vec<(TrajId, u32, String)>> {
1892 let mut out = Vec::new();
1893 let mut tid = start_traj_id;
1894 let mut paths: Vec<_> = fs::read_dir(dir)?
1895 .filter_map(|e| e.ok())
1896 .map(|e| e.path())
1897 .filter(|p| {
1898 let name = p.file_name().and_then(|s| s.to_str()).unwrap_or("");
1899 name.ends_with(".con") || name.ends_with(".convel")
1900 })
1901 .collect();
1902 paths.sort();
1903 for p in paths {
1904 let n = self.append_trajectory_path_units(tid, &p, units.clone())?;
1905 out.push((tid, n, p.display().to_string()));
1906 tid += 1;
1907 }
1908 Ok(out)
1909 }
1910
1911 pub fn unique_frame_keys(&self, sel: &Select) -> Result<Vec<FrameKey>> {
1912 let keys = self.select(sel)?;
1913 let mut uniq = Vec::new();
1914 for k in keys {
1915 let h = self.frame_hash(k)?;
1916 if self.find_by_hash(h)? == Some(k) {
1917 uniq.push(k);
1918 }
1919 }
1920 Ok(uniq)
1921 }
1922}
1923
1924#[cfg(test)]
1925mod tests {
1926 use super::*;
1927 use std::path::PathBuf;
1928
1929 fn fixture(name: &str) -> PathBuf {
1930 PathBuf::from(env!("CARGO_MANIFEST_DIR"))
1931 .join("resources/test")
1932 .join(name)
1933 }
1934
1935 fn fixtures_dir() -> PathBuf {
1936 PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test")
1937 }
1938
1939 #[test]
1940 fn ingest_hash_dedup_and_select() {
1941 let dir = tempfile::tempdir().unwrap();
1942 let db = ConCorpus::open(dir.path()).unwrap();
1943 let n = db
1944 .append_trajectory_path(1, fixture("tiny_cuh2.con"))
1945 .unwrap();
1946 assert!(n >= 1);
1947 let k0 = FrameKey {
1948 traj_id: 1,
1949 frame_idx: 0,
1950 };
1951 let h = db.frame_hash(k0).unwrap();
1952 assert_eq!(db.find_by_hash(h).unwrap(), Some(k0));
1953
1954 let by_hash = db.select(&Select::new().exact_hash(h.to_bytes())).unwrap();
1955 assert_eq!(by_hash, vec![k0]);
1956
1957 let cu = db
1958 .select(&Select::new().require_symbol("Cu").trajectory(1))
1959 .unwrap();
1960 assert!(!cu.is_empty());
1961
1962 let n2 = db
1963 .append_trajectory_path(2, fixture("tiny_multi_cuh2.con"))
1964 .unwrap();
1965 assert!(n2 >= 2);
1966 }
1967
1968 #[test]
1969 fn get_velocities_returns_native_convel() {
1970 let dir = tempfile::tempdir().unwrap();
1971 let db = ConCorpus::open(dir.path()).unwrap();
1972 db.append_trajectory_path(1, fixture("tiny_cuh2.convel"))
1973 .unwrap();
1974 let v = db
1975 .get_velocities(FrameKey {
1976 traj_id: 1,
1977 frame_idx: 0,
1978 })
1979 .unwrap()
1980 .expect("velocities");
1981 assert!((v[0][0] - 0.001234).abs() < 1e-12, "got {}", v[0][0]);
1982 assert!(db
1983 .get_velocities(FrameKey {
1984 traj_id: 1,
1985 frame_idx: 0,
1986 })
1987 .unwrap()
1988 .is_some());
1989 db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
1990 .unwrap();
1991 assert!(db
1992 .get_velocities(FrameKey {
1993 traj_id: 2,
1994 frame_idx: 0,
1995 })
1996 .unwrap()
1997 .is_none());
1998 }
1999
2000 #[test]
2001 fn caller_units_canonical_in_con_metadata() {
2002 let dir = tempfile::tempdir().unwrap();
2003 let db = ConCorpus::open(dir.path()).unwrap();
2004 db.append_trajectory_path_units(
2005 1,
2006 fixture("tiny_cuh2.con"),
2007 Some(serde_json::json!({
2008 "length": "A",
2009 "energy": "ev",
2010 "time": "femtosecond"
2011 })),
2012 )
2013 .unwrap();
2014 let k = FrameKey {
2015 traj_id: 1,
2016 frame_idx: 0,
2017 };
2018 let u = db.frame_units(k).unwrap().expect("units");
2019 assert_eq!(u["length"], "angstrom");
2020 assert_eq!(u["energy"], "eV");
2021 assert_eq!(u["time"], "fs");
2022 let text = db.get_frame_text(k).unwrap();
2023 assert!(text.contains("angstrom"), "{text}");
2024 assert!(text.contains("fs"), "{text}");
2025 db.set_trajectory_units(
2026 1,
2027 serde_json::json!({"length": "nm", "energy": "hartree", "time": "ps"}),
2028 )
2029 .unwrap();
2030 let u2 = db.frame_units(k).unwrap().expect("units");
2031 assert_eq!(u2["length"], "nm");
2032 assert_eq!(u2["energy"], "hartree");
2033 assert_eq!(u2["time"], "ps");
2034 let fr0 = db.get_frame(k).unwrap();
2035 let vol_nm = crate::frame_scalars::frame_cell_volume(&fr0).expect("vol");
2036 let old_ang_vol = vol_nm * 1000.0;
2037 let stale = db
2038 .select(&Select::new().volume_range(old_ang_vol * 0.5, old_ang_vol * 1.5))
2039 .unwrap();
2040 assert!(
2041 stale.is_empty(),
2042 "old angstrom volume key must be gone after set_units"
2043 );
2044 let fresh = db
2045 .select(&Select::new().volume_range(vol_nm * 0.5, vol_nm * 1.5))
2046 .unwrap();
2047 assert_eq!(fresh, vec![k]);
2048 }
2049
2050 #[test]
2051 fn ingest_directory_units_stamps_canonical() {
2052 let dir = tempfile::tempdir().unwrap();
2053 let src = dir.path().join("cons");
2054 std::fs::create_dir(&src).unwrap();
2055 std::fs::copy(fixture("tiny_cuh2.con"), src.join("a.con")).unwrap();
2056 let db = ConCorpus::open(dir.path().join("db")).unwrap();
2057 let n = db
2058 .ingest_directory_units(
2059 &src,
2060 1,
2061 Some(serde_json::json!({"length":"A","energy":"ev"})),
2062 )
2063 .unwrap();
2064 assert_eq!(n.len(), 1);
2065 let u = db
2066 .frame_units(FrameKey {
2067 traj_id: 1,
2068 frame_idx: 0,
2069 })
2070 .unwrap()
2071 .unwrap();
2072 assert_eq!(u["length"], "angstrom");
2073 assert_eq!(u["energy"], "eV");
2074 }
2075
2076 #[test]
2077 fn readonly_open_and_pack_roundtrip() {
2078 let dir = tempfile::tempdir().unwrap();
2079 let key = FrameKey {
2080 traj_id: 1,
2081 frame_idx: 0,
2082 };
2083 let key1 = FrameKey {
2084 traj_id: 1,
2085 frame_idx: 1,
2086 };
2087 let packed = {
2088 let db = ConCorpus::open(dir.path()).unwrap();
2089 db.append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
2090 .unwrap();
2091 let packed = db.pack_frame(key).unwrap();
2092 db.close();
2093 packed
2094 };
2095 let decoded = crate::cooked_soa::CookedSoa::decode(&packed).unwrap();
2096 assert!(decoded.natoms >= 1);
2097
2098 let ro = ConCorpus::open_readonly(dir.path()).unwrap();
2099 let packed2 = ro.pack_frame(key).unwrap();
2100 assert_eq!(packed, packed2);
2101 let batch = ro.pack_frames(&[key, key1]).unwrap();
2102 let parts = crate::cooked_soa::decode_batch(&batch).unwrap();
2103 assert_eq!(parts.len(), 2);
2104 assert_eq!(parts[0], packed);
2105 let d1 = crate::cooked_soa::CookedSoa::decode(&parts[1]).unwrap();
2106 assert!((d1.positions[2][0] - 8.8549).abs() < 1e-4);
2107 assert!(ConCorpus::open_readonly("/no/such/corpus-dir-readonly").is_err());
2108 assert!(ro
2109 .append_trajectory_path(9, fixture("tiny_cuh2.con"))
2110 .is_err());
2111 }
2112
2113 #[test]
2114 fn workflow_metatrain_extxyz_export() {
2115 let dir = tempfile::tempdir().unwrap();
2116 let db = ConCorpus::open(dir.path().join("corpus")).unwrap();
2117 let ingested = db.ingest_directory(fixtures_dir(), 1).unwrap();
2118 assert!(ingested.len() >= 3);
2119
2120 let keys = db
2121 .select(
2122 &Select::new()
2123 .require_symbol("Cu")
2124 .natoms_range(1, 10_000)
2125 .limit(50),
2126 )
2127 .unwrap();
2128 assert!(!keys.is_empty());
2129
2130 let uniq = db
2131 .unique_frame_keys(&Select::new().require_symbol("Cu"))
2132 .unwrap();
2133 assert!(!uniq.is_empty());
2134
2135 let xyz = dir.path().join("train_subset.xyz");
2136 let n = db.export_extxyz(&uniq, &xyz, "energy").unwrap();
2137 assert_eq!(n, uniq.len());
2138 let text = std::fs::read_to_string(&xyz).unwrap();
2139 assert!(text.contains("Lattice="));
2140 assert!(text.contains("Properties="));
2141 assert!(text.lines().any(|l| l.trim_start().starts_with("Cu ")));
2142 let err = db.export_extxyz(&uniq, &xyz, "energy").unwrap_err();
2143 assert!(err.to_string().contains("dest exists"), "{err}");
2144 let missing = dir.path().join("missing_key.xyz");
2145 let bad = db.export_extxyz(
2146 &[FrameKey {
2147 traj_id: 99,
2148 frame_idx: 0,
2149 }],
2150 &missing,
2151 "energy",
2152 );
2153 assert!(bad.is_err());
2154 assert!(!missing.exists());
2155 }
2156
2157 #[test]
2158 fn open_refuses_sharded_root() {
2159 let dir = tempfile::tempdir().unwrap();
2160 let root = dir.path().join("sharded");
2161 crate::ShardedConCorpus::open(&root, 2).unwrap();
2162 assert!(root.join("shards.json").is_file());
2163 match ConCorpus::open(&root) {
2164 Ok(_) => panic!("expected sharded-root refuse"),
2165 Err(err) => assert!(err.to_string().contains("sharded root"), "{err}"),
2166 }
2167 assert!(!root.join("data.mdb").exists());
2168 }
2169
2170 #[test]
2171 fn workflow_dedup_identical_ingest() {
2172 let dir = tempfile::tempdir().unwrap();
2173 let db = ConCorpus::open(dir.path()).unwrap();
2174 let f = fixture("tiny_cuh2.con");
2175 db.append_trajectory_path(1, &f).unwrap();
2176 db.append_trajectory_path(2, &f).unwrap();
2177 let k1 = FrameKey {
2178 traj_id: 1,
2179 frame_idx: 0,
2180 };
2181 let k2 = FrameKey {
2182 traj_id: 2,
2183 frame_idx: 0,
2184 };
2185 let h1 = db.frame_hash(k1).unwrap();
2186 let h2 = db.frame_hash(k2).unwrap();
2187 assert_eq!(h1, h2);
2188 assert_eq!(db.find_by_hash(h1).unwrap(), Some(k1));
2189 let uniq = db.unique_frame_keys(&Select::new()).unwrap();
2190 assert!(uniq.contains(&k1));
2191 assert!(!uniq.contains(&k2));
2192 }
2193
2194 #[test]
2195 fn metadata_indexes_forces_velocities_energy() {
2196 let dir = tempfile::tempdir().unwrap();
2197 let db = ConCorpus::open(dir.path()).unwrap();
2198 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2199 .unwrap();
2200 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
2201 .unwrap();
2202 db.append_trajectory_path(3, fixture("tiny_cuh2.convel"))
2203 .unwrap();
2204 db.append_trajectory_path(4, fixture("tiny_cuh2_vel_forces.con"))
2205 .unwrap();
2206
2207 let forces = db.select(&Select::new().require_forces()).unwrap();
2208 assert!(forces.iter().any(|k| k.traj_id == 2));
2209 assert!(forces.iter().any(|k| k.traj_id == 4));
2210 assert!(!forces.iter().any(|k| k.traj_id == 1));
2211
2212 let vels = db.select(&Select::new().require_velocities()).unwrap();
2213 assert!(vels.iter().any(|k| k.traj_id == 3));
2214 assert!(vels.iter().any(|k| k.traj_id == 4));
2215
2216 let both = db
2217 .select(&Select::new().require_forces().require_velocities())
2218 .unwrap();
2219 assert!(both.iter().any(|k| k.traj_id == 4));
2220 assert!(!both.iter().any(|k| k.traj_id == 2));
2221
2222 let with_e = db.select(&Select::new().require_energy()).unwrap();
2223 assert!(with_e.iter().any(|k| k.traj_id == 2));
2224 assert!(with_e.iter().any(|k| k.traj_id == 4));
2225 assert!(!with_e.iter().any(|k| k.traj_id == 1));
2226
2227 let in_range = db
2228 .select(&Select::new().energy_range(-43.0, -42.0).require_forces())
2229 .unwrap();
2230 assert!(in_range.iter().any(|k| k.traj_id == 2));
2231 let miss = db.select(&Select::new().energy_range(0.0, 1.0)).unwrap();
2232 assert!(miss.is_empty());
2233 }
2234
2235 #[test]
2236 fn composition_and_fmax_indexes() {
2237 let dir = tempfile::tempdir().unwrap();
2238 let db = ConCorpus::open(dir.path()).unwrap();
2239 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2240 .unwrap();
2241 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
2242 .unwrap();
2243 db.append_trajectory_path(3, fixture("sulfolene.con"))
2244 .unwrap();
2245
2246 let formula = db
2247 .frame_formula(FrameKey {
2248 traj_id: 1,
2249 frame_idx: 0,
2250 })
2251 .unwrap();
2252 assert_eq!(formula, "Cu:2|H:2");
2253
2254 let exact = db
2255 .select(&Select::new().exact_composition("Cu:2|H:2"))
2256 .unwrap();
2257 assert!(exact.iter().any(|k| k.traj_id == 1));
2258 assert!(exact.iter().any(|k| k.traj_id == 2));
2259 assert!(!exact.iter().any(|k| k.traj_id == 3));
2260
2261 let cu2 = db
2262 .select(&Select::new().element_exact("Cu", 2).element_exact("H", 2))
2263 .unwrap();
2264 assert_eq!(cu2.len(), exact.len());
2265
2266 let cu_min2 = db.select(&Select::new().element_min("Cu", 2)).unwrap();
2267 assert!(cu_min2.iter().any(|k| k.traj_id == 1));
2268
2269 let wrong = db.select(&Select::new().exact_composition("Fe:1")).unwrap();
2270 assert!(wrong.is_empty());
2271
2272 let fr_f = db
2273 .get_frame(FrameKey {
2274 traj_id: 2,
2275 frame_idx: 0,
2276 })
2277 .unwrap();
2278 let expected_fmax = frame_fmax(&fr_f).expect("forces fixture has fmax");
2279 let in_fmax = db
2280 .select(
2281 &Select::new()
2282 .require_forces()
2283 .fmax_range(0.0, expected_fmax + 1e-6),
2284 )
2285 .unwrap();
2286 assert!(in_fmax.iter().any(|k| k.traj_id == 2));
2287 assert!(!in_fmax.iter().any(|k| k.traj_id == 1));
2288
2289 let too_small = db
2290 .select(&Select::new().fmax_range(0.0, expected_fmax * 0.0 - 1.0))
2291 .unwrap();
2292 let impossible = db
2294 .select(&Select::new().fmax_range(1e9, 1e9 + 1.0))
2295 .unwrap();
2296 assert!(impossible.is_empty());
2297 let _ = too_small;
2298 }
2299
2300 #[test]
2301 fn reindex_and_append_frames() {
2302 let dir = tempfile::tempdir().unwrap();
2303 let db = ConCorpus::open(dir.path()).unwrap();
2304 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2305 .unwrap();
2306 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
2307 .unwrap();
2308
2309 let before = db
2310 .select(&Select::new().exact_composition("Cu:2|H:2"))
2311 .unwrap();
2312 assert!(!before.is_empty());
2313
2314 let n = db.reindex().unwrap();
2315 assert!(n >= 2);
2316 let after = db
2317 .select(&Select::new().exact_composition("Cu:2|H:2"))
2318 .unwrap();
2319 assert_eq!(before, after);
2320
2321 let fr = db
2322 .get_frame(FrameKey {
2323 traj_id: 1,
2324 frame_idx: 0,
2325 })
2326 .unwrap();
2327 let n2 = db
2328 .append_trajectory_frames(10, &[fr.clone()], "mem")
2329 .unwrap();
2330 assert_eq!(n2, 1);
2331 let mem = db
2332 .select(&Select::new().trajectory(10).exact_composition("Cu:2|H:2"))
2333 .unwrap();
2334 assert_eq!(mem.len(), 1);
2335
2336 let n3 = db.extend_trajectory_frames(10, &[fr], "mem-ext").unwrap();
2337 assert_eq!(n3, 2);
2338 }
2339
2340 #[test]
2341 fn extend_trajectory_path_creates_then_appends() {
2342 let dir = tempfile::tempdir().unwrap();
2343 let db = ConCorpus::open(dir.path()).unwrap();
2344 let path = fixture("tiny_cuh2.con");
2345 let n1 = db.extend_trajectory_path(3, &path).unwrap();
2346 assert!(n1 >= 1);
2347 let n2 = db.extend_trajectory_path(3, &path).unwrap();
2348 assert_eq!(n2, n1 * 2);
2349 }
2350
2351 #[test]
2352 fn reindex_twice_stable() {
2353 let dir = tempfile::tempdir().unwrap();
2354 let db = ConCorpus::open(dir.path()).unwrap();
2355 db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
2356 .unwrap();
2357 let a = db.reindex().unwrap();
2358 let keys1 = db.select(&Select::new().require_forces()).unwrap();
2359 let b = db.reindex().unwrap();
2360 let keys2 = db.select(&Select::new().require_forces()).unwrap();
2361 assert_eq!(a, b);
2362 assert_eq!(keys1, keys2);
2363 }
2364
2365 #[test]
2366 fn batch_touch_and_texts() {
2367 let dir = tempfile::tempdir().unwrap();
2368 let db = ConCorpus::open(dir.path()).unwrap();
2369 let n = db
2370 .append_trajectory_path(1, fixture("tiny_cuh2.con"))
2371 .unwrap();
2372 assert!(n >= 1);
2373 let (total, checksum) = db.touch_trajectory_blobs(1, n).unwrap();
2374 assert!(total > 0);
2375 assert!(checksum > 0);
2376 let texts = db
2377 .get_frame_texts(&[FrameKey {
2378 traj_id: 1,
2379 frame_idx: 0,
2380 }])
2381 .unwrap();
2382 assert_eq!(texts.len(), 1);
2383 assert_eq!(texts[0].len() as u64, total);
2384 let mut expect = 0u64;
2386 for (i, b) in texts[0].as_bytes().iter().enumerate() {
2387 expect = expect.wrapping_add((*b as u64).wrapping_mul((i as u64).wrapping_add(1)));
2388 }
2389 assert_eq!(checksum, expect);
2390 }
2391
2392 #[test]
2394 fn concurrent_writers_distinct_traj() {
2395 use std::sync::Arc;
2396 use std::thread;
2397 let dir = tempfile::tempdir().unwrap();
2398 let db = Arc::new(ConCorpus::open(dir.path()).unwrap());
2399 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
2400 let mut handles = Vec::new();
2401 for tid in 1u64..=8 {
2402 let db = Arc::clone(&db);
2403 let text = text.clone();
2404 handles.push(thread::spawn(move || {
2405 db.append_trajectory_str(tid, &text, format!("t{tid}"))
2406 .unwrap()
2407 }));
2408 }
2409 let mut counts = Vec::new();
2410 for h in handles {
2411 counts.push(h.join().unwrap());
2412 }
2413 assert!(counts.iter().all(|&c| c >= 1));
2414 for tid in 1u64..=8 {
2415 let meta = db.traj_meta(tid).unwrap().expect("meta");
2416 assert_eq!(meta.n_frames, counts[(tid - 1) as usize]);
2417 let _ = db
2418 .get_frame_text(FrameKey {
2419 traj_id: tid,
2420 frame_idx: 0,
2421 })
2422 .unwrap();
2423 }
2424 let cu = db.select(&Select::new().require_symbol("Cu")).unwrap();
2425 assert_eq!(cu.len(), 8);
2426 }
2427
2428 #[test]
2430 fn multiproc_cli_concurrent_select() {
2431 use std::process::Command;
2432 let dir = tempfile::tempdir().unwrap();
2433 let corpus = dir.path().join("mp_corpus");
2434 {
2435 let db = ConCorpus::open(&corpus).unwrap();
2436 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2437 .unwrap();
2438 }
2439 let Some(bin) = option_env!("CARGO_BIN_EXE_readcon-db") else {
2440 return;
2441 };
2442 let corpus_s = corpus.to_str().unwrap();
2443 let mut kids = Vec::new();
2444 for _ in 0..4 {
2445 kids.push(
2446 Command::new(bin)
2447 .args(["select", corpus_s, "--symbol", "Cu"])
2448 .stdout(std::process::Stdio::null())
2449 .spawn()
2450 .expect("spawn"),
2451 );
2452 }
2453 for mut c in kids {
2454 assert!(c.wait().unwrap().success());
2455 }
2456 }
2457
2458 #[test]
2459 fn ase_competitive_mass_volume_pbc_meta_charge() {
2460 use readcon_core::types::ConFrameBuilder;
2461 let dir = tempfile::tempdir().unwrap();
2462 let db = ConCorpus::open(dir.path()).unwrap();
2463
2464 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2466 .unwrap();
2467 let fr1 = db
2468 .get_frame(FrameKey {
2469 traj_id: 1,
2470 frame_idx: 0,
2471 })
2472 .unwrap();
2473 let mass1 = frame_total_mass(&fr1).expect("mass");
2474 let vol1 = frame_cell_volume(&fr1).expect("vol");
2475
2476 let mut fr2 = fr1.clone();
2478 fr2.header.set_frame_index(7);
2479 fr2.header.set_time(1.5);
2480 fr2.header.set_timestep(0.25);
2481 fr2.header.set_neb_bead(2);
2482 fr2.header.set_neb_band(1);
2483 fr2.header.set_pbc([true, true, false]);
2484 fr2.header
2485 .metadata
2486 .insert("charge".into(), serde_json::json!(-1.0));
2487 fr2.header
2488 .metadata
2489 .insert("magmom".into(), serde_json::json!(2.0));
2490 fr2.header.boxl = [20.0, 20.0, 20.0];
2492 db.append_trajectory_frames(2, &[fr2], "meta").unwrap();
2493
2494 let by_mass = db
2495 .select(&Select::new().mass_range(mass1 - 1e-6, mass1 + 1e-6))
2496 .unwrap();
2497 assert!(by_mass.iter().any(|k| k.traj_id == 1));
2498
2499 let by_vol_small = db
2500 .select(&Select::new().volume_range(vol1 - 1.0, vol1 + 1.0))
2501 .unwrap();
2502 assert!(by_vol_small.iter().any(|k| k.traj_id == 1));
2503 assert!(!by_vol_small.iter().any(|k| k.traj_id == 2));
2504
2505 let pbc_match = db.select(&Select::new().pbc([true, true, false])).unwrap();
2506 assert_eq!(pbc_match.len(), 1);
2507 assert_eq!(pbc_match[0].traj_id, 2);
2508 let pbc_miss = db.select(&Select::new().pbc([true, true, true])).unwrap();
2509 assert!(pbc_miss.is_empty());
2510
2511 let fi = db
2512 .select(&Select::new().frame_index_range(7.0, 7.0))
2513 .unwrap();
2514 assert_eq!(
2515 fi,
2516 vec![FrameKey {
2517 traj_id: 2,
2518 frame_idx: 0
2519 }]
2520 );
2521 let fi_miss = db
2522 .select(&Select::new().frame_index_range(99.0, 100.0))
2523 .unwrap();
2524 assert!(fi_miss.is_empty());
2525
2526 let bead = db.select(&Select::new().neb_bead_range(2.0, 2.0)).unwrap();
2527 assert_eq!(bead.len(), 1);
2528 let t = db.select(&Select::new().time_range(1.4, 1.6)).unwrap();
2529 assert_eq!(t.len(), 1);
2530
2531 let dt = db
2532 .select(&Select::new().timestep_range(0.25, 0.25))
2533 .unwrap();
2534 assert_eq!(
2535 dt,
2536 vec![FrameKey {
2537 traj_id: 2,
2538 frame_idx: 0
2539 }]
2540 );
2541 let dt_miss = db.select(&Select::new().timestep_range(9.0, 10.0)).unwrap();
2542 assert!(dt_miss.is_empty());
2543 assert!(!dt.iter().any(|k| k.traj_id == 1));
2545
2546 let band = db.select(&Select::new().neb_band_range(1.0, 1.0)).unwrap();
2547 assert_eq!(
2548 band,
2549 vec![FrameKey {
2550 traj_id: 2,
2551 frame_idx: 0
2552 }]
2553 );
2554 let band_miss = db
2555 .select(&Select::new().neb_band_range(99.0, 100.0))
2556 .unwrap();
2557 assert!(band_miss.is_empty());
2558
2559 let ch = db.select(&Select::new().charge_range(-1.0, -1.0)).unwrap();
2560 assert_eq!(ch.len(), 1);
2561 let mm = db.select(&Select::new().magmom_range(2.0, 2.0)).unwrap();
2562 assert_eq!(mm.len(), 1);
2563 let ch_miss = db.select(&Select::new().charge_range(0.0, 0.0)).unwrap();
2564 assert!(ch_miss.is_empty());
2565
2566 let n = db.reindex().unwrap();
2567 assert!(n >= 2);
2568 let after = db
2569 .select(&Select::new().frame_index_range(7.0, 7.0))
2570 .unwrap();
2571 assert_eq!(after.len(), 1);
2572 let after_dt = db
2573 .select(&Select::new().timestep_range(0.25, 0.25))
2574 .unwrap();
2575 assert_eq!(after_dt.len(), 1);
2576 let after_band = db.select(&Select::new().neb_band_range(1.0, 1.0)).unwrap();
2577 assert_eq!(after_band.len(), 1);
2578
2579 let form = db
2581 .select(&Select::new().exact_composition("Cu:2|H:2"))
2582 .unwrap();
2583 assert!(form.len() >= 2);
2584 }
2585
2586 #[test]
2587 fn cooked_soa_roundtrip_and_authority() {
2588 let dir = tempfile::tempdir().unwrap();
2589 let db = ConCorpus::open(dir.path()).unwrap();
2590 let key = FrameKey {
2591 traj_id: 1,
2592 frame_idx: 0,
2593 };
2594 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2595 .unwrap();
2596 assert!(db.get_cooked_soa_bytes(key).unwrap().is_none());
2598 let pos_via_parse = db.get_positions(key).unwrap();
2599 let fr = db.get_frame(key).unwrap();
2600 assert_eq!(pos_via_parse.len(), fr.atom_data.len());
2601 for (i, a) in fr.atom_data.iter().enumerate() {
2602 assert_eq!(pos_via_parse[i], [a.x, a.y, a.z]);
2603 }
2604
2605 let text_before = db.get_frame_text(key).unwrap();
2606 let h_before = db.frame_hash(key).unwrap();
2607 let nbytes = db.cook_frame(key).unwrap();
2608 assert!(nbytes > 0);
2609 let cooked = db.get_cooked_soa(key).unwrap().expect("cooked present");
2610 assert_eq!(cooked.natoms as usize, fr.atom_data.len());
2611 for (i, a) in fr.atom_data.iter().enumerate() {
2612 assert_eq!(cooked.positions[i], [a.x, a.y, a.z]);
2613 }
2614 assert_eq!(db.get_frame_text(key).unwrap(), text_before);
2616 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h_before.to_bytes());
2617 assert_eq!(db.find_by_hash(h_before).unwrap(), Some(key));
2618 let pos_cooked = db.get_positions(key).unwrap();
2619 assert_eq!(pos_cooked, cooked.positions);
2620
2621 let keys = db.select(&Select::new().require_symbol("Cu")).unwrap();
2623 assert!(keys.contains(&key));
2624
2625 db.reindex().unwrap();
2627 assert!(db.get_cooked_soa(key).unwrap().is_some());
2628 assert!(db
2629 .select(&Select::new().require_symbol("Cu"))
2630 .unwrap()
2631 .contains(&key));
2632
2633 db.delete_cooked_soa(key).unwrap();
2635 assert!(db.get_cooked_soa_bytes(key).unwrap().is_none());
2636 let pos_again = db.get_positions(key).unwrap();
2637 assert_eq!(pos_again, pos_via_parse);
2638 assert_eq!(db.get_frame_text(key).unwrap(), text_before);
2639 }
2640
2641 #[test]
2642 fn cooked_soa_forces_and_cook_on_ingest() {
2643 let dir = tempfile::tempdir().unwrap();
2644 let db = ConCorpus::open(dir.path()).unwrap();
2645 db.append_trajectory_path_cook(3, fixture("tiny_cuh2_forces.con"), true)
2646 .unwrap();
2647 let key = FrameKey {
2648 traj_id: 3,
2649 frame_idx: 0,
2650 };
2651 assert!(db.get_cooked_soa_bytes(key).unwrap().is_some());
2652 let fr = db.get_frame(key).unwrap();
2653 let forces = db.get_forces(key).unwrap().expect("forces");
2654 for (i, a) in fr.atom_data.iter().enumerate() {
2655 if let Some(f) = a.force {
2656 assert_eq!(forces[i], f);
2657 }
2658 }
2659 let mut wtxn = db.env.write_txn().unwrap();
2661 let fk_b = key.to_bytes();
2662 db.frames_soa
2663 .put(&mut wtxn, &fk_b[..], b"XXXX_not_valid")
2664 .unwrap();
2665 wtxn.commit().unwrap();
2666 assert!(db.get_cooked_soa(key).unwrap().is_none());
2667 let pos = db.get_positions(key).unwrap();
2668 assert_eq!(pos.len(), fr.atom_data.len());
2669 assert_eq!(
2670 pos[0],
2671 [fr.atom_data[0].x, fr.atom_data[0].y, fr.atom_data[0].z]
2672 );
2673 }
2674
2675 #[test]
2676 fn cooked_soa_recook_all() {
2677 let dir = tempfile::tempdir().unwrap();
2678 let db = ConCorpus::open(dir.path()).unwrap();
2679 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2680 .unwrap();
2681 db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
2682 .unwrap();
2683 let n = db.recook_all().unwrap();
2684 assert_eq!(n, 2);
2685 assert!(db
2686 .get_cooked_soa(FrameKey {
2687 traj_id: 1,
2688 frame_idx: 0
2689 })
2690 .unwrap()
2691 .is_some());
2692 assert!(db
2693 .get_cooked_soa(FrameKey {
2694 traj_id: 2,
2695 frame_idx: 0
2696 })
2697 .unwrap()
2698 .is_some());
2699 }
2700
2701 #[test]
2703 fn cooked_soa_cannot_replace_missing_frames_text() {
2704 let dir = tempfile::tempdir().unwrap();
2705 let db = ConCorpus::open(dir.path()).unwrap();
2706 let key = FrameKey {
2707 traj_id: 9,
2708 frame_idx: 0,
2709 };
2710 db.append_trajectory_path_cook(9, fixture("tiny_cuh2.con"), true)
2711 .unwrap();
2712 assert!(db.has_valid_cooked_soa(key).unwrap());
2713 let pos_ok = db.get_positions(key).unwrap();
2714 assert!(!pos_ok.is_empty());
2715 {
2717 let mut wtxn = db.env.write_txn().unwrap();
2718 let fk_b = key.to_bytes();
2719 db.frames.delete(&mut wtxn, &fk_b[..]).unwrap();
2720 wtxn.commit().unwrap();
2721 }
2722 assert!(db.get_cooked_soa(key).unwrap().is_some());
2724 let pos_cooked_only = db.get_positions(key).unwrap();
2725 assert_eq!(pos_cooked_only, pos_ok);
2726 assert!(db.get_frame_text(key).is_err());
2728 assert!(db.get_frame(key).is_err());
2729 assert!(db.frame_formula(key).is_err());
2732 }
2733
2734 #[test]
2735 fn numeric_fast_path_matches_con_twice() {
2736 let dir = tempfile::tempdir().unwrap();
2737 let db = ConCorpus::open(dir.path()).unwrap();
2738 let key = FrameKey {
2739 traj_id: 1,
2740 frame_idx: 0,
2741 };
2742 db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
2743 .unwrap();
2744 let fr = db.get_frame(key).unwrap();
2745 let expect_pos: Vec<_> = fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect();
2746 let expect_f: Vec<_> = fr
2747 .atom_data
2748 .iter()
2749 .map(|a| a.force.unwrap_or([0.0; 3]))
2750 .collect();
2751 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
2753 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
2754 db.cook_frame(key).unwrap();
2755 assert!(db.has_valid_cooked_soa(key).unwrap());
2756 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
2758 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
2759 let text = db.get_frame_text(key).unwrap();
2760 let h = db.frame_hash(key).unwrap();
2761 db.delete_cooked_soa(key).unwrap();
2762 assert_eq!(db.get_frame_text(key).unwrap(), text);
2763 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
2764 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
2765 }
2766
2767 #[test]
2769 fn numeric_extract_velocities_match_con() {
2770 let dir = tempfile::tempdir().unwrap();
2771 let db = ConCorpus::open(dir.path()).unwrap();
2772 let key = FrameKey {
2773 traj_id: 1,
2774 frame_idx: 0,
2775 };
2776 db.append_trajectory_path(1, fixture("tiny_cuh2_vel_forces.con"))
2777 .unwrap();
2778 assert!(!db.has_valid_cooked_soa(key).unwrap());
2779 let fr = db.get_frame(key).unwrap();
2780 let expect_pos: Vec<_> = fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect();
2781 let expect_f: Vec<_> = fr
2782 .atom_data
2783 .iter()
2784 .map(|a| a.force.unwrap_or([0.0; 3]))
2785 .collect();
2786 let expect_v: Vec<_> = fr
2787 .atom_data
2788 .iter()
2789 .map(|a| a.velocity.unwrap_or([0.0; 3]))
2790 .collect();
2791 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
2792 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
2793 assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
2794 assert!((expect_v[0][0] - 0.001234).abs() < 1e-12);
2795 let text = db.get_frame_text(key).unwrap();
2796 let h = db.frame_hash(key).unwrap();
2797 db.cook_frame(key).unwrap();
2798 assert!(db.has_valid_cooked_soa(key).unwrap());
2799 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
2800 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
2801 assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
2802 assert_eq!(db.get_frame_text(key).unwrap(), text);
2803 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
2804 assert_eq!(db.find_by_hash(h).unwrap(), Some(key));
2805 db.delete_cooked_soa(key).unwrap();
2806 assert!(!db.has_valid_cooked_soa(key).unwrap());
2807 assert_eq!(db.get_positions(key).unwrap(), expect_pos);
2808 assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
2809 assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
2810 assert_eq!(db.get_frame_text(key).unwrap(), text);
2811 }
2812
2813 #[test]
2815 fn numeric_extract_valid_rcso_skips_con_parse() {
2816 let dir = tempfile::tempdir().unwrap();
2817 let db = ConCorpus::open(dir.path()).unwrap();
2818 let key = FrameKey {
2819 traj_id: 1,
2820 frame_idx: 0,
2821 };
2822 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2823 .unwrap();
2824 let text = db.get_frame_text(key).unwrap();
2825 let h = db.frame_hash(key).unwrap();
2826 let mut fr = db.get_frame(key).unwrap();
2827 let con_x = fr.atom_data[0].x;
2828 fr.atom_data[0].x = 99.0;
2829 let poison = crate::cooked_soa::CookedSoa::encode_frame(&fr).unwrap();
2830 let mut wtxn = db.env.write_txn().unwrap();
2831 let fk_b = key.to_bytes();
2832 db.frames_soa
2833 .put(&mut wtxn, &fk_b[..], &poison[..])
2834 .unwrap();
2835 wtxn.commit().unwrap();
2836 assert!(db.has_valid_cooked_soa(key).unwrap());
2837 let pos = db.get_positions(key).unwrap();
2838 assert!(
2839 (pos[0][0] - 99.0).abs() < 1e-12,
2840 "valid RCSO must skip CON parse, got {}",
2841 pos[0][0]
2842 );
2843 assert!((con_x - 0.6394).abs() < 1e-4);
2844 assert_eq!(db.get_frame_text(key).unwrap(), text);
2845 assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
2846 }
2847
2848 #[test]
2849 fn traj_meta_old_json_loads_without_topo_fields() {
2850 let m: TrajMeta = serde_json::from_str(r#"{"n_frames":3,"source":"a.con"}"#).unwrap();
2851 assert_eq!(m.n_frames, 3);
2852 assert_eq!(m.source, "a.con");
2853 assert!(m.topo_cutoff.is_none());
2854 assert!(m.topo_graph.is_none());
2855 assert!(m.topo_hops.is_none());
2856 assert!(m.topo_method.is_none());
2857 }
2858
2859 fn seed_topo(db: &ConCorpus, fk: FrameKey, hex: &str) {
2860 let (idx, by) = db.require_topo_dbs().unwrap();
2861 let mut wtxn = db.env.write_txn().unwrap();
2862 db.put_topo_entry(&mut wtxn, idx, by, fk, hex).unwrap();
2863 wtxn.commit().unwrap();
2864 }
2865
2866 fn stamp_traj_topo(db: &ConCorpus, tid: TrajId, method: &str, cutoff: f64) {
2867 let mut meta = db.traj_meta(tid).unwrap().expect("meta");
2868 meta.stamp_topo(cutoff, "cutoff", 2, method);
2869 let mut wtxn = db.env.write_txn().unwrap();
2870 let tid_key = tid.to_be_bytes();
2871 db.traj_meta
2872 .put(
2873 &mut wtxn,
2874 &tid_key[..],
2875 &serde_json::to_string(&meta).unwrap(),
2876 )
2877 .unwrap();
2878 wtxn.commit().unwrap();
2879 }
2880
2881 #[test]
2882 fn select_topo_key_prefix_without_seams() {
2883 let dir = tempfile::tempdir().unwrap();
2884 let db = ConCorpus::open(dir.path()).unwrap();
2885 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2886 .unwrap();
2887 let k0 = FrameKey {
2888 traj_id: 1,
2889 frame_idx: 0,
2890 };
2891 seed_topo(&db, k0, "abc123");
2892 let hits = db.select(&Select::new().topo_key("ABC123")).unwrap();
2893 assert_eq!(hits, vec![k0]);
2894 let miss = db.select(&Select::new().topo_key("00")).unwrap();
2895 assert!(miss.is_empty());
2896 }
2897
2898 #[test]
2899 fn reindex_without_annotation_leaves_idx_topo_empty() {
2900 let dir = tempfile::tempdir().unwrap();
2901 let db = ConCorpus::open(dir.path()).unwrap();
2902 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2903 .unwrap();
2904 let k0 = FrameKey {
2905 traj_id: 1,
2906 frame_idx: 0,
2907 };
2908 seed_topo(&db, k0, "abc123");
2909 assert_eq!(
2910 db.select(&Select::new().topo_key("abc123")).unwrap(),
2911 vec![k0]
2912 );
2913 db.reindex().unwrap();
2914 assert!(db
2915 .select(&Select::new().topo_key("abc123"))
2916 .unwrap()
2917 .is_empty());
2918 assert_eq!(db.frame_topo_key(k0).unwrap().as_deref(), Some("abc123"));
2919 }
2920
2921 #[test]
2922 fn reindex_rebuilds_idx_topo_when_params_agree() {
2923 let dir = tempfile::tempdir().unwrap();
2924 let db = ConCorpus::open(dir.path()).unwrap();
2925 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2926 .unwrap();
2927 db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
2928 .unwrap();
2929 let k0 = FrameKey {
2930 traj_id: 1,
2931 frame_idx: 0,
2932 };
2933 let k1 = FrameKey {
2934 traj_id: 2,
2935 frame_idx: 0,
2936 };
2937 seed_topo(&db, k0, "deadbeef");
2938 seed_topo(&db, k1, "deadbeef");
2939 stamp_traj_topo(&db, 1, "nauty", 3.0);
2940 stamp_traj_topo(&db, 2, "nauty", 3.0);
2941 let (idx, _) = db.require_topo_dbs().unwrap();
2942 {
2943 let mut wtxn = db.env.write_txn().unwrap();
2944 idx.clear(&mut wtxn).unwrap();
2945 wtxn.commit().unwrap();
2946 }
2947 assert!(db
2948 .select(&Select::new().topo_key("deadbeef"))
2949 .unwrap()
2950 .is_empty());
2951 db.reindex().unwrap();
2952 let hits = db.select(&Select::new().topo_key("deadbeef")).unwrap();
2953 assert_eq!(hits.len(), 2);
2954 }
2955
2956 #[test]
2957 fn reindex_refuses_mixed_methods() {
2958 let dir = tempfile::tempdir().unwrap();
2959 let db = ConCorpus::open(dir.path()).unwrap();
2960 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2961 .unwrap();
2962 db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
2963 .unwrap();
2964 stamp_traj_topo(&db, 1, "nauty", 3.0);
2965 stamp_traj_topo(&db, 2, "wl", 3.0);
2966 let err = db.reindex().unwrap_err().to_string();
2967 assert!(err.contains("mixed topology"), "{err}");
2968 let cu = db.select(&Select::new().require_symbol("Cu")).unwrap();
2969 assert_eq!(cu.len(), 2);
2970 }
2971
2972 #[test]
2973 fn extend_preserves_topology_fields() {
2974 let dir = tempfile::tempdir().unwrap();
2975 let db = ConCorpus::open(dir.path()).unwrap();
2976 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2977 .unwrap();
2978 stamp_traj_topo(&db, 1, "nauty", 3.0);
2979 let fr = db
2980 .get_frame(FrameKey {
2981 traj_id: 1,
2982 frame_idx: 0,
2983 })
2984 .unwrap();
2985 db.extend_trajectory_frames(1, &[fr], "extra").unwrap();
2986 let meta = db.traj_meta(1).unwrap().unwrap();
2987 assert_eq!(meta.n_frames, 2);
2988 assert_eq!(meta.topo_method.as_deref(), Some("nauty"));
2989 assert_eq!(meta.topo_cutoff, Some(3.0));
2990 assert_eq!(meta.topo_graph.as_deref(), Some("cutoff"));
2991 assert_eq!(meta.topo_hops, Some(2));
2992 }
2993
2994 #[test]
2995 fn annotate_without_seams_names_command() {
2996 let dir = tempfile::tempdir().unwrap();
2997 let db = ConCorpus::open(dir.path()).unwrap();
2998 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
2999 .unwrap();
3000 let err = db
3001 .annotate_topology(AnnotateTopologyOpts::new(3.0).seams("/no/such/seams-binary"))
3002 .unwrap_err()
3003 .to_string();
3004 assert!(err.contains("seams fingerprint"), "{err}");
3005 }
3006
3007 #[test]
3008 fn find_by_topology_without_params_errors() {
3009 let dir = tempfile::tempdir().unwrap();
3010 let db = ConCorpus::open(dir.path()).unwrap();
3011 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3012 .unwrap();
3013 let err = db
3014 .find_by_topology_path(fixture("tiny_cuh2.con"))
3015 .unwrap_err()
3016 .to_string();
3017 assert!(err.contains("annotate-topology"), "{err}");
3018 }
3019
3020 #[test]
3021 fn readonly_missing_topo_select_is_empty() {
3022 let dir = tempfile::tempdir().unwrap();
3023 {
3024 let db = ConCorpus::open(dir.path()).unwrap();
3025 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3026 .unwrap();
3027 db.close();
3028 }
3029 let db = ConCorpus::open_readonly(dir.path()).unwrap();
3030 let hits = db.select(&Select::new().topo_key("abc")).unwrap();
3031 assert!(hits.is_empty());
3032 }
3033
3034 fn seams_or_skip() -> Option<std::path::PathBuf> {
3035 match crate::topology::resolve_seams_binary(None) {
3036 Ok(p) => Some(p),
3037 Err(e) => {
3038 eprintln!("skipping topology seams test: {e}");
3039 None
3040 }
3041 }
3042 }
3043
3044 fn permute_cu_rows(text: &str) -> String {
3045 let mut lines: Vec<&str> = text.lines().collect();
3046 let mut cu = Vec::new();
3047 for (i, line) in lines.iter().enumerate() {
3048 let t = line.trim_start();
3049 if t.starts_with("0.6394") || t.starts_with("3.1969") {
3050 cu.push(i);
3051 }
3052 }
3053 assert_eq!(cu.len(), 2, "expected two Cu coordinate rows");
3054 lines.swap(cu[0], cu[1]);
3055 let mut out = lines.join("\n");
3056 if text.ends_with('\n') {
3057 out.push('\n');
3058 }
3059 out
3060 }
3061
3062 fn break_one_bond(text: &str) -> String {
3063 let mut fr = {
3064 let mut it = ConFrameIterator::new(text);
3065 it.next().unwrap().unwrap()
3066 };
3067 let cu = fr
3070 .atom_data
3071 .iter_mut()
3072 .find(|a| a.symbol.as_ref() == "Cu")
3073 .expect("Cu");
3074 cu.z += 10.0;
3075 let mut buf = Cursor::new(Vec::new());
3076 {
3077 let mut w = ConFrameWriter::new(&mut buf);
3078 w.write_frame(&fr).unwrap();
3079 }
3080 String::from_utf8(buf.into_inner()).unwrap()
3081 }
3082
3083 #[test]
3084 fn annotate_permutation_shares_hex_broken_does_not() {
3085 let Some(seams) = seams_or_skip() else {
3086 return;
3087 };
3088 let dir = tempfile::tempdir().unwrap();
3089 let db = ConCorpus::open(dir.path()).unwrap();
3090 db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3091 .unwrap();
3092 let n = db
3093 .annotate_topology(AnnotateTopologyOpts::new(3.0).seams(&seams))
3094 .unwrap();
3095 assert_eq!(n, 1);
3096 let k0 = FrameKey {
3097 traj_id: 1,
3098 frame_idx: 0,
3099 };
3100 let hex = db.frame_topo_key(k0).unwrap().expect("hex");
3101 assert!(!hex.is_empty());
3102 let meta = db.traj_meta(1).unwrap().unwrap();
3103 assert_eq!(meta.topo_cutoff, Some(3.0));
3104 assert_eq!(meta.topo_graph.as_deref(), Some("cutoff"));
3105 assert_eq!(meta.topo_hops, Some(2));
3106 assert!(meta.topo_method.is_some());
3107
3108 let orig = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
3109 let perm = permute_cu_rows(&orig);
3110 let broken = break_one_bond(&orig);
3111 let perm_hits = db.find_by_topology_text(&perm).unwrap();
3112 assert_eq!(perm_hits, vec![k0]);
3113 let broken_hits = db.find_by_topology_text(&broken).unwrap();
3114 assert!(
3115 broken_hits.is_empty(),
3116 "broken Cu-Cu bond must not share topology key, got {broken_hits:?}"
3117 );
3118 let sel = db.select(&Select::new().topo_key(&hex)).unwrap();
3119 assert_eq!(sel, vec![k0]);
3120 }
3121}