use std::env;
use std::process::ExitCode;
use readcon_db::{
join_drained_roots, open_single_env_for_export, AnnotateTopologyOpts, ConCorpus, FrameKey,
Select, ShardedConCorpus, DEFAULT_N_SHARDS,
};
fn usage() -> ExitCode {
eprintln!(
"Usage:
readcon-db ingest <corpus_dir> [--start-id N] [--units JSON] <file.con>...
readcon-db ingest-dir <corpus_dir> [--units JSON] <dir_with_con_files>
readcon-db select <corpus_dir> [--traj N] [--symbol S] [--natoms-min A] [--natoms-max B]
[--energy-min E] [--energy-max E] [--fmax-min F] [--fmax-max F]
[--elem SYM:COUNT] [--elem-min SYM:COUNT] [--formula Cu:2|H:2]
[--mass-min M] [--mass-max M] [--volume-min V] [--volume-max V]
[--pbc X,Y,Z] [--time-min T] [--time-max T] [--timestep-min DT] [--timestep-max DT]
[--frame-index-min I] [--frame-index-max I]
[--neb-bead-min N] [--neb-bead-max N] [--neb-band-min B] [--neb-band-max B]
[--charge-min C] [--charge-max C] [--magmom-min M] [--magmom-max M]
[--require-forces] [--require-velocities] [--require-energy]
[--topo-key HEX]
[--export out.xyz]
readcon-db dedup-export <corpus_dir> [same filters as select] -o out.xyz
readcon-db reindex <corpus_dir>
readcon-db annotate-topology <corpus_dir> --cutoff A [--graph G] [--hops N] [--seams PATH]
readcon-db find-by-topology <corpus_dir> <file.con>
readcon-db shard-init <root> [--shards N]
readcon-db shard-ingest <root> --shard S --start-id T [--units JSON] <file.con>...
readcon-db shard-select <root> [--symbol S] ...
readcon-db drain <local_root> <pfs_root>
readcon-db join-drained <single_dst> <drained_root>...
readcon-db compact-join <sharded_root> <single_dst>
readcon-db compact-split <single_src> <sharded_dst> [--shards N]
readcon-db compact-export-extxyz <corpus_or_shard_root> <out.xyz> [--sharded] [--symbol S]
readcon-db cook-frame <corpus_dir> --traj T --frame F
readcon-db recook-all <corpus_dir>
readcon-db delete-cooked <corpus_dir> --traj T --frame F
readcon-db positions <corpus_dir> --traj T --frame F
readcon-db hash-file <file.con>
(RCSO cooked tier is opt-in; CON text in frames remains authority — docs/orgmode/cooked-soa.org)
"
);
ExitCode::from(2)
}
fn parse_sym_count(s: &str) -> Result<(String, u32), Box<dyn std::error::Error>> {
let (sym, cnt) = s.split_once(':').ok_or("expected SYM:COUNT")?;
Ok((sym.to_string(), cnt.parse()?))
}
fn main() -> ExitCode {
let mut args: Vec<String> = env::args().skip(1).collect();
if args.is_empty() {
return usage();
}
let cmd = args.remove(0);
let run = (|| -> Result<(), Box<dyn std::error::Error>> {
match cmd.as_str() {
"ingest" => {
let corpus = args.first().ok_or("corpus_dir")?.clone();
let mut start = 1u64;
let mut files = Vec::new();
let mut units = None;
let mut i = 1;
while i < args.len() {
if args[i] == "--start-id" {
start = args.get(i + 1).ok_or("id")?.parse()?;
i += 2;
continue;
}
if args[i] == "--units" {
let raw = args.get(i + 1).ok_or("units json")?;
units = Some(serde_json::from_str(raw)?);
i += 2;
continue;
}
files.push(args[i].clone());
i += 1;
}
let db = ConCorpus::open(&corpus)?;
let mut tid = start;
for f in files {
let n = db.append_trajectory_path_units(tid, &f, units.clone())?;
println!("traj {tid}: {n} frames from {f}");
tid += 1;
}
}
"ingest-dir" => {
let corpus = args.first().ok_or("corpus")?.clone();
let mut units = None;
let mut dir = None;
let mut i = 1;
while i < args.len() {
if args[i] == "--units" {
let raw = args.get(i + 1).ok_or("units json")?;
units = Some(serde_json::from_str(raw)?);
i += 2;
continue;
}
dir = Some(args[i].clone());
i += 1;
}
let dir = dir.ok_or("dir")?;
let db = ConCorpus::open(&corpus)?;
for (tid, n, p) in db.ingest_directory_units(&dir, 1, units)? {
println!("traj {tid}: {n} frames <- {p}");
}
}
"reindex" => {
let corpus = args.first().ok_or("corpus")?.clone();
let db = ConCorpus::open_existing(&corpus)?;
let n = db.reindex()?;
println!("reindexed {n} frames");
}
"annotate-topology" => {
let corpus = args.first().ok_or("corpus_dir")?.clone();
let mut cutoff = None::<f64>;
let mut graph = None::<String>;
let mut hops = None::<u32>;
let mut seams = None::<String>;
let mut i = 1;
while i < args.len() {
match args[i].as_str() {
"--cutoff" => {
cutoff = Some(args.get(i + 1).ok_or("cutoff")?.parse()?);
i += 2;
}
"--graph" => {
graph = Some(args.get(i + 1).ok_or("graph")?.clone());
i += 2;
}
"--hops" => {
hops = Some(args.get(i + 1).ok_or("hops")?.parse()?);
i += 2;
}
"--seams" => {
seams = Some(args.get(i + 1).ok_or("seams")?.clone());
i += 2;
}
_ => i += 1,
}
}
let cutoff = cutoff.ok_or("--cutoff is required")?;
let mut opts = AnnotateTopologyOpts::new(cutoff);
if let Some(g) = graph {
opts = opts.graph(g);
}
if let Some(h) = hops {
opts = opts.hops(h);
}
if let Some(s) = seams {
opts = opts.seams(s);
}
let db = ConCorpus::open_existing(&corpus)?;
let n = db.annotate_topology(opts)?;
println!("annotated {n} frames");
}
"find-by-topology" => {
let corpus = args.first().ok_or("corpus_dir")?.clone();
let file = args.get(1).ok_or("file.con")?.clone();
let db = ConCorpus::open_readonly(&corpus)?;
let keys = db.find_by_topology_path(&file)?;
println!("{} keys", keys.len());
for k in keys.iter().take(20) {
println!(" traj={} frame={}", k.traj_id, k.frame_idx);
}
if keys.len() > 20 {
println!(" ...");
}
}
"select" | "dedup-export" => {
let corpus = args.first().ok_or("corpus")?.clone();
let mut traj = None;
let mut symbol = None;
let mut nmin = 0u32;
let mut nmax = u32::MAX;
let mut emin: Option<f64> = None;
let mut emax: Option<f64> = None;
let mut fmin: Option<f64> = None;
let mut fmax: Option<f64> = None;
let mut mass_min: Option<f64> = None;
let mut mass_max: Option<f64> = None;
let mut vol_min: Option<f64> = None;
let mut vol_max: Option<f64> = None;
let mut pbc: Option<[bool; 3]> = None;
let mut time_min: Option<f64> = None;
let mut time_max: Option<f64> = None;
let mut timestep_min: Option<f64> = None;
let mut timestep_max: Option<f64> = None;
let mut fi_min: Option<f64> = None;
let mut fi_max: Option<f64> = None;
let mut bead_min: Option<f64> = None;
let mut bead_max: Option<f64> = None;
let mut band_min: Option<f64> = None;
let mut band_max: Option<f64> = None;
let mut charge_min: Option<f64> = None;
let mut charge_max: Option<f64> = None;
let mut mag_min: Option<f64> = None;
let mut mag_max: Option<f64> = None;
let mut elem_exact = Vec::new();
let mut elem_min = Vec::new();
let mut formula = None;
let mut topo_key = None;
let mut req_forces = false;
let mut req_vels = false;
let mut req_energy = false;
let mut export = None;
let mut i = 1;
while i < args.len() {
match args[i].as_str() {
"--traj" => {
traj = Some(args.get(i + 1).ok_or("traj")?.parse()?);
i += 2;
}
"--symbol" => {
symbol = Some(args.get(i + 1).ok_or("sym")?.clone());
i += 2;
}
"--natoms-min" => {
nmin = args.get(i + 1).ok_or("n")?.parse()?;
i += 2;
}
"--natoms-max" => {
nmax = args.get(i + 1).ok_or("n")?.parse()?;
i += 2;
}
"--energy-min" => {
emin = Some(args.get(i + 1).ok_or("e")?.parse()?);
i += 2;
}
"--energy-max" => {
emax = Some(args.get(i + 1).ok_or("e")?.parse()?);
i += 2;
}
"--fmax-min" => {
fmin = Some(args.get(i + 1).ok_or("f")?.parse()?);
i += 2;
}
"--fmax-max" => {
fmax = Some(args.get(i + 1).ok_or("f")?.parse()?);
i += 2;
}
"--elem" => {
elem_exact.push(parse_sym_count(args.get(i + 1).ok_or("SYM:COUNT")?)?);
i += 2;
}
"--elem-min" => {
elem_min.push(parse_sym_count(args.get(i + 1).ok_or("SYM:COUNT")?)?);
i += 2;
}
"--formula" => {
formula = Some(args.get(i + 1).ok_or("formula")?.clone());
i += 2;
}
"--topo-key" => {
topo_key = Some(args.get(i + 1).ok_or("hex")?.clone());
i += 2;
}
"--mass-min" => {
mass_min = Some(args.get(i + 1).ok_or("m")?.parse()?);
i += 2;
}
"--mass-max" => {
mass_max = Some(args.get(i + 1).ok_or("m")?.parse()?);
i += 2;
}
"--volume-min" => {
vol_min = Some(args.get(i + 1).ok_or("v")?.parse()?);
i += 2;
}
"--volume-max" => {
vol_max = Some(args.get(i + 1).ok_or("v")?.parse()?);
i += 2;
}
"--pbc" => {
let s = args.get(i + 1).ok_or("x,y,z")?;
let parts: Vec<_> = s.split(',').collect();
if parts.len() != 3 {
return Err("pbc needs X,Y,Z as 0/1 or true/false".into());
}
let parse_b = |x: &str| -> Result<bool, Box<dyn std::error::Error>> {
Ok(matches!(
x.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "t" | "yes"
))
};
pbc =
Some([parse_b(parts[0])?, parse_b(parts[1])?, parse_b(parts[2])?]);
i += 2;
}
"--time-min" => {
time_min = Some(args.get(i + 1).ok_or("t")?.parse()?);
i += 2;
}
"--time-max" => {
time_max = Some(args.get(i + 1).ok_or("t")?.parse()?);
i += 2;
}
"--timestep-min" => {
timestep_min = Some(args.get(i + 1).ok_or("dt")?.parse()?);
i += 2;
}
"--timestep-max" => {
timestep_max = Some(args.get(i + 1).ok_or("dt")?.parse()?);
i += 2;
}
"--frame-index-min" => {
fi_min = Some(args.get(i + 1).ok_or("i")?.parse()?);
i += 2;
}
"--frame-index-max" => {
fi_max = Some(args.get(i + 1).ok_or("i")?.parse()?);
i += 2;
}
"--neb-bead-min" => {
bead_min = Some(args.get(i + 1).ok_or("n")?.parse()?);
i += 2;
}
"--neb-bead-max" => {
bead_max = Some(args.get(i + 1).ok_or("n")?.parse()?);
i += 2;
}
"--neb-band-min" => {
band_min = Some(args.get(i + 1).ok_or("b")?.parse()?);
i += 2;
}
"--neb-band-max" => {
band_max = Some(args.get(i + 1).ok_or("b")?.parse()?);
i += 2;
}
"--charge-min" => {
charge_min = Some(args.get(i + 1).ok_or("c")?.parse()?);
i += 2;
}
"--charge-max" => {
charge_max = Some(args.get(i + 1).ok_or("c")?.parse()?);
i += 2;
}
"--magmom-min" => {
mag_min = Some(args.get(i + 1).ok_or("m")?.parse()?);
i += 2;
}
"--magmom-max" => {
mag_max = Some(args.get(i + 1).ok_or("m")?.parse()?);
i += 2;
}
"--require-forces" => {
req_forces = true;
i += 1;
}
"--require-velocities" => {
req_vels = true;
i += 1;
}
"--require-energy" => {
req_energy = true;
i += 1;
}
"--export" | "-o" => {
export = Some(args.get(i + 1).ok_or("path")?.clone());
i += 2;
}
_ => i += 1,
}
}
let db = ConCorpus::open_readonly(&corpus)?;
let mut sel = Select::new().natoms_range(nmin, nmax);
if let Some(t) = traj {
sel = sel.trajectory(t);
}
if let Some(s) = symbol.clone() {
sel = sel.require_symbol(s);
}
if emin.is_some() || emax.is_some() {
sel = sel.energy_range(
emin.unwrap_or(f64::NEG_INFINITY),
emax.unwrap_or(f64::INFINITY),
);
}
if fmin.is_some() || fmax.is_some() {
sel = sel.fmax_range(fmin.unwrap_or(0.0), fmax.unwrap_or(f64::INFINITY));
}
if mass_min.is_some() || mass_max.is_some() {
sel = sel.mass_range(
mass_min.unwrap_or(f64::NEG_INFINITY),
mass_max.unwrap_or(f64::INFINITY),
);
}
if vol_min.is_some() || vol_max.is_some() {
sel = sel.volume_range(
vol_min.unwrap_or(f64::NEG_INFINITY),
vol_max.unwrap_or(f64::INFINITY),
);
}
if let Some(p) = pbc {
sel = sel.pbc(p);
}
if time_min.is_some() || time_max.is_some() {
sel = sel.time_range(
time_min.unwrap_or(f64::NEG_INFINITY),
time_max.unwrap_or(f64::INFINITY),
);
}
if timestep_min.is_some() || timestep_max.is_some() {
sel = sel.timestep_range(
timestep_min.unwrap_or(f64::NEG_INFINITY),
timestep_max.unwrap_or(f64::INFINITY),
);
}
if fi_min.is_some() || fi_max.is_some() {
sel = sel.frame_index_range(
fi_min.unwrap_or(f64::NEG_INFINITY),
fi_max.unwrap_or(f64::INFINITY),
);
}
if bead_min.is_some() || bead_max.is_some() {
sel = sel.neb_bead_range(
bead_min.unwrap_or(f64::NEG_INFINITY),
bead_max.unwrap_or(f64::INFINITY),
);
}
if band_min.is_some() || band_max.is_some() {
sel = sel.neb_band_range(
band_min.unwrap_or(f64::NEG_INFINITY),
band_max.unwrap_or(f64::INFINITY),
);
}
if charge_min.is_some() || charge_max.is_some() {
sel = sel.charge_range(
charge_min.unwrap_or(f64::NEG_INFINITY),
charge_max.unwrap_or(f64::INFINITY),
);
}
if mag_min.is_some() || mag_max.is_some() {
sel = sel.magmom_range(
mag_min.unwrap_or(f64::NEG_INFINITY),
mag_max.unwrap_or(f64::INFINITY),
);
}
for (sym, c) in elem_exact {
sel = sel.element_exact(sym, c);
}
for (sym, c) in elem_min {
sel = sel.element_min(sym, c);
}
if let Some(f) = formula {
sel = sel.exact_composition(f);
}
if let Some(hex) = topo_key {
sel = sel.topo_key(hex);
}
if req_forces {
sel = sel.require_forces();
}
if req_vels {
sel = sel.require_velocities();
}
if req_energy {
sel = sel.require_energy();
}
let keys = if cmd == "dedup-export" {
db.unique_frame_keys(&sel)?
} else {
db.select(&sel)?
};
println!("{} keys", keys.len());
if let Some(path) = export {
let n = db.export_extxyz(&keys, &path, "energy")?;
println!("wrote {n} frames -> {path}");
} else {
for k in keys.iter().take(20) {
println!(" traj={} frame={}", k.traj_id, k.frame_idx);
}
if keys.len() > 20 {
println!(" ...");
}
}
}
"shard-init" => {
let root = args.first().ok_or("root")?.clone();
let mut ns = DEFAULT_N_SHARDS;
let mut i = 1;
while i < args.len() {
if args[i] == "--shards" {
ns = args.get(i + 1).ok_or("n")?.parse()?;
i += 2;
} else {
i += 1;
}
}
let _ = ShardedConCorpus::open(&root, ns)?;
println!("initialized {root} with {ns} shards");
}
"shard-ingest" => {
let root = args.first().ok_or("root")?.clone();
let mut shard = None::<u32>;
let mut start = 1u64;
let mut files = Vec::new();
let mut units = None;
let mut i = 1;
while i < args.len() {
match args[i].as_str() {
"--shard" => {
shard = Some(args.get(i + 1).ok_or("shard")?.parse()?);
i += 2;
}
"--start-id" => {
start = args.get(i + 1).ok_or("id")?.parse()?;
i += 2;
}
"--units" => {
let raw = args.get(i + 1).ok_or("units json")?;
units = Some(serde_json::from_str(raw)?);
i += 2;
}
_ => {
files.push(args[i].clone());
i += 1;
}
}
}
let sid = shard.ok_or(
"--shard required (one writer per shard id across the job; overlapping shard ids drain to unique dests then join-drained)",
)?;
let man = std::path::Path::new(&root).join("shards.json");
if !man.is_file() {
return Err("shard-ingest: missing shards.json; run shard-init".into());
}
let m: readcon_db::ShardManifest =
serde_json::from_str(&std::fs::read_to_string(&man)?)?;
if files.is_empty() {
return Err("shard-ingest needs at least one file".into());
}
let mut tid = start;
if ShardedConCorpus::shard_for_traj(tid, m.n_shards) != sid {
return Err(format!(
"traj {tid} routes to shard {}, not {sid}; choose start-id ≡ {sid} (mod n_shards)",
ShardedConCorpus::shard_for_traj(tid, m.n_shards)
)
.into());
}
let db = ShardedConCorpus::open_shard(&root, sid)?;
for f in files {
let routed = ShardedConCorpus::shard_for_traj(tid, m.n_shards);
if routed != sid {
return Err(format!(
"traj {tid} routes to shard {routed}, not {sid}; choose start-id ≡ {sid} (mod n_shards)"
)
.into());
}
let n = db.append_trajectory_path_units(tid, &f, units.clone())?;
println!("shard {sid} traj {tid}: {n} frames from {f}");
tid += 1;
while ShardedConCorpus::shard_for_traj(tid, m.n_shards) != sid {
tid += 1;
}
}
}
"shard-select" => {
let root = args.first().ok_or("root")?.clone();
let mut symbol = None;
let mut i = 1;
while i < args.len() {
if args[i] == "--symbol" {
symbol = Some(args.get(i + 1).ok_or("sym")?.clone());
i += 2;
} else {
i += 1;
}
}
let mut db = ShardedConCorpus::open_existing(&root)?;
let mut sel = Select::new();
if let Some(s) = symbol {
sel = sel.require_symbol(s);
}
let keys = db.select(&sel)?;
println!("{} keys across {} shards", keys.len(), db.n_shards());
for k in keys.iter().take(20) {
println!(" traj={} frame={}", k.traj_id, k.frame_idx);
}
}
"drain" => {
let src = args.first().ok_or("local_root")?.clone();
let dst = args.get(1).ok_or("pfs_root")?.clone();
let n = ShardedConCorpus::drain_to(&src, &dst)?;
println!("drained {n} shards {src} -> {dst} (close writers first)");
}
"join-drained" => {
let dst = args.first().ok_or("single_dst")?.clone();
let srcs: Vec<std::path::PathBuf> =
args.iter().skip(1).map(std::path::PathBuf::from).collect();
if srcs.is_empty() {
return Err("join-drained needs at least one drained root".into());
}
let n = join_drained_roots(&srcs, &dst)?;
println!("joined {n} frames from {} roots -> {dst}", srcs.len());
}
"compact-join" => {
let src = args.first().ok_or("sharded_root")?.clone();
let dst = args.get(1).ok_or("single_dst")?.clone();
let mut s = ShardedConCorpus::open_existing(&src)?;
let n = s.join_to_single_env(&dst)?;
println!("joined {n} frames -> {dst} (single-env-lmdb)");
}
"compact-split" => {
let src = args.first().ok_or("single_src")?.clone();
let dst = args.get(1).ok_or("sharded_dst")?.clone();
let mut ns = DEFAULT_N_SHARDS;
let mut i = 2;
while i < args.len() {
if args[i] == "--shards" {
ns = args.get(i + 1).ok_or("n")?.parse()?;
i += 2;
} else {
i += 1;
}
}
let single = ConCorpus::open_readonly(&src)?;
let n = ShardedConCorpus::split_single_to_sharded(&single, &dst, ns)?;
println!("split {n} frames -> {dst} ({ns} shards, sharded-lmdb)");
}
"compact-export-extxyz" => {
let src = args.first().ok_or("src")?.clone();
let out = args.get(1).ok_or("out.xyz")?.clone();
let mut sharded = false;
let mut symbol = None;
let mut i = 2;
while i < args.len() {
match args[i].as_str() {
"--sharded" => {
sharded = true;
i += 1;
}
"--symbol" => {
symbol = Some(args.get(i + 1).ok_or("sym")?.clone());
i += 2;
}
_ => i += 1,
}
}
let mut sel = Select::new();
if let Some(s) = symbol {
sel = sel.require_symbol(s);
}
let n = if sharded {
let mut sc2 = ShardedConCorpus::open_existing(&src)?;
sc2.export_extxyz(&sel, &out, "energy")?
} else {
let db = open_single_env_for_export(&src)?;
let keys = db.select(&sel)?;
db.export_extxyz(&keys, &out, "energy")? as u32
};
println!("wrote {n} frames extxyz -> {out} (analysis export)");
}
"cook-frame" => {
let corpus = args.first().ok_or("corpus")?.clone();
let mut traj = 0u64;
let mut frame = 0u32;
let mut i = 1;
while i < args.len() {
match args[i].as_str() {
"--traj" => {
traj = args.get(i + 1).ok_or("traj")?.parse()?;
i += 2;
}
"--frame" => {
frame = args.get(i + 1).ok_or("frame")?.parse()?;
i += 2;
}
_ => i += 1,
}
}
let db = ConCorpus::open_existing(&corpus)?;
let n = db.cook_frame(FrameKey {
traj_id: traj,
frame_idx: frame,
})?;
println!("cooked traj={traj} frame={frame} ({n} RCSO bytes); CON text unchanged");
}
"recook-all" => {
let corpus = args.first().ok_or("corpus")?.clone();
let db = ConCorpus::open_existing(&corpus)?;
let n = db.recook_all()?;
println!("recooked {n} frames");
}
"delete-cooked" => {
let corpus = args.first().ok_or("corpus")?.clone();
let mut traj = 0u64;
let mut frame = 0u32;
let mut i = 1;
while i < args.len() {
match args[i].as_str() {
"--traj" => {
traj = args.get(i + 1).ok_or("traj")?.parse()?;
i += 2;
}
"--frame" => {
frame = args.get(i + 1).ok_or("frame")?.parse()?;
i += 2;
}
_ => i += 1,
}
}
let db = ConCorpus::open_existing(&corpus)?;
db.delete_cooked_soa(FrameKey {
traj_id: traj,
frame_idx: frame,
})?;
println!("deleted cooked traj={traj} frame={frame}");
}
"positions" => {
let corpus = args.first().ok_or("corpus")?.clone();
let mut traj = 0u64;
let mut frame = 0u32;
let mut i = 1;
while i < args.len() {
match args[i].as_str() {
"--traj" => {
traj = args.get(i + 1).ok_or("traj")?.parse()?;
i += 2;
}
"--frame" => {
frame = args.get(i + 1).ok_or("frame")?.parse()?;
i += 2;
}
_ => i += 1,
}
}
let db = ConCorpus::open_readonly(&corpus)?;
let key = FrameKey {
traj_id: traj,
frame_idx: frame,
};
let cooked = db.has_valid_cooked_soa(key)?;
let pos = db.get_positions(key)?;
println!(
"traj={traj} frame={frame} natoms={} valid_cooked={cooked}",
pos.len()
);
for (i, row) in pos.iter().enumerate().take(5) {
println!(" [{i}] {} {} {}", row[0], row[1], row[2]);
}
if pos.len() > 5 {
println!(" ... ({} more)", pos.len() - 5);
}
}
"hash-file" => {
let f = args.first().ok_or("file")?;
let text = std::fs::read_to_string(f)?;
let h = ConCorpus::hash_con_text(&text)?;
println!("{}", h.to_hex());
}
_ => return Err("unknown command".into()),
}
Ok(())
})();
match run {
Ok(()) => ExitCode::SUCCESS,
Err(e) => {
eprintln!("error: {e}");
usage()
}
}
}