From 6e8590536446337ff64dfcab167b3d50078c7ccb Mon Sep 17 00:00:00 2001 From: T4ras123 Date: Mon, 6 Apr 2026 12:23:11 +0400 Subject: [PATCH 1/2] Add revisited data processing scripts and helpers Support the revisited preprocessing workflow with grouped/counting utilities and shared serialization helpers so the data pipeline can be reviewed independently from training and evaluation changes. Made-with: Cursor --- scripts/analyze_binning_precision.py | 200 +++++++ scripts/analyze_coord_centering.py | 213 +++++++ scripts/center_and_analyze.py | 249 +++++++++ scripts/compute_range_R.py | 96 ++++ scripts/count_geom_revisited_train_pickle.sh | 98 ++++ scripts/count_revisited_train_tokens.sh | 167 ++++++ scripts/fit_bins.py | 132 +++++ scripts/preprocess_geom_revisited.sh | 76 +++ ...reprocess_geom_revisited_paired_grouped.sh | 94 ++++ .../data_processing/data_preprocessing.py | 79 ++- .../data_preprocessing_revisited.py | 467 ++++++++++++++++ .../preprocess_geom_grouped.py | 157 ++---- .../preprocess_geom_grouped_revisited.py | 521 ++++++++++++++++++ .../data_processing/smiles_encoder_decoder.py | 317 ++++++++--- src/molgen3D/data_processing/utils.py | 222 +++++++- .../dataprocessing/count_tokens.py | 6 +- 16 files changed, 2866 insertions(+), 228 deletions(-) create mode 100644 scripts/analyze_binning_precision.py create mode 100644 scripts/analyze_coord_centering.py create mode 100644 scripts/center_and_analyze.py create mode 100644 scripts/compute_range_R.py create mode 100755 scripts/count_geom_revisited_train_pickle.sh create mode 100755 scripts/count_revisited_train_tokens.sh create mode 100644 scripts/fit_bins.py create mode 100755 scripts/preprocess_geom_revisited.sh create mode 100755 scripts/preprocess_geom_revisited_paired_grouped.sh create mode 100644 src/molgen3D/data_processing/data_preprocessing_revisited.py create mode 100644 src/molgen3D/data_processing/preprocess_geom_grouped_revisited.py diff --git a/scripts/analyze_binning_precision.py b/scripts/analyze_binning_precision.py new file mode 100644 index 0000000..b63e759 --- /dev/null +++ b/scripts/analyze_binning_precision.py @@ -0,0 +1,200 @@ +""" +Measure coordinate precision loss from binning on the validation set. + +Compares three encoding methods: + 1. Uniform bins (256 bins, raw/non-centered) + 2. Quantile bins (256 bins, raw/non-centered) + 3. Raw text encoding (truncate to 4 decimal places) + +For each conformer, computes RMSD between original and round-tripped coordinates. + +Usage: + python scripts/analyze_binning_precision.py \ + --data /data/molgen/geom_revisited/val_data.pickle \ + --n_mols 10000 +""" + +import argparse +import math +import os +import pickle +import sys + +import numpy as np + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from molgen3D.data_processing.smiles_encoder_decoder import ( + BinConfig, + _encode_scalar, + _decode_scalar, +) + +BIN_CONFIGS_DIR = os.path.join(os.path.dirname(__file__), "bin_configs") + + +def truncate_coord(x, precision=4): + """Replicate the raw text encoding's truncation.""" + factor = 10 ** precision + truncated = math.trunc(x * factor) / factor + if abs(truncated) < 10 ** (-precision): + truncated = 0.0 + return truncated + + +def roundtrip_bins(coords_flat, config): + """Bin and unbin a flat array of coordinate values.""" + decoded = np.empty_like(coords_flat) + for i, c in enumerate(coords_flat): + idx = _encode_scalar(float(c), config) + decoded[i] = _decode_scalar(idx, config) + return decoded + + +def roundtrip_raw(coords_flat, precision=4): + """Truncate to `precision` decimal places (raw text encoding).""" + return np.array([truncate_coord(float(c), precision) for c in coords_flat]) + + +def conformer_rmsd(original, decoded): + """RMSD between two (N, 3) coordinate arrays.""" + diff = original - decoded.reshape(original.shape) + return np.sqrt(np.mean(diff ** 2)) + + +def conformer_max_error(original, decoded): + """Max absolute error across all coordinates.""" + return np.max(np.abs(original - decoded.reshape(original.shape))) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--data", type=str, + default="/data/molgen/geom_revisited/val_data.pickle") + parser.add_argument("--n_mols", type=int, default=10000) + parser.add_argument("--precision", type=int, default=4, + help="Decimal precision for raw text encoding") + args = parser.parse_args() + + # Load bin configs (raw = fit on non-centered data) + uniform_cfg = BinConfig.load(os.path.join(BIN_CONFIGS_DIR, "uniform_bins.json")) + quantile_cfg = BinConfig.load(os.path.join(BIN_CONFIGS_DIR, "quantile_bins.json")) + + print(f"Uniform bins: L={uniform_cfg.L:.4f}, H={uniform_cfg.H:.4f}, " + f"n_bins={uniform_cfg.n_bins}, " + f"bin_width={(uniform_cfg.H - uniform_cfg.L) / uniform_cfg.n_bins:.6f} A") + print(f"Quantile bins: L={quantile_cfg.L:.4f}, H={quantile_cfg.H:.4f}, " + f"n_bins={quantile_cfg.n_bins}, " + f"median_bin_width={np.median(np.diff(quantile_cfg.edges)):.6f} A") + print(f"Raw text precision: {args.precision} decimal places") + + # Load validation data + print(f"\nLoading {args.data} ...") + with open(args.data, "rb") as f: + data = pickle.load(f) + print(f" {len(data)} molecules in validation set") + + # Collect conformers from first n_mols molecules + n_mols = min(args.n_mols, len(data)) + print(f" Using first {n_mols} molecules\n") + + methods = ["uniform", "quantile", "raw"] + stats = {m: {"rmsd": [], "max_err": [], "overflow": 0, "total_coords": 0} for m in methods} + + n_conformers = 0 + for mol_idx, (smiles, confs) in enumerate(data[:n_mols]): + for mol in confs: + try: + pos = mol.GetConformer().GetPositions() # (n_atoms, 3) + except Exception: + continue + + flat = pos.flatten() + n_conformers += 1 + + # Uniform bins + decoded_u = roundtrip_bins(flat, uniform_cfg) + stats["uniform"]["rmsd"].append(conformer_rmsd(pos, decoded_u)) + stats["uniform"]["max_err"].append(conformer_max_error(pos, decoded_u)) + stats["uniform"]["overflow"] += int(np.sum((flat < uniform_cfg.L) | (flat > uniform_cfg.H))) + stats["uniform"]["total_coords"] += len(flat) + + # Quantile bins + decoded_q = roundtrip_bins(flat, quantile_cfg) + stats["quantile"]["rmsd"].append(conformer_rmsd(pos, decoded_q)) + stats["quantile"]["max_err"].append(conformer_max_error(pos, decoded_q)) + stats["quantile"]["overflow"] += int(np.sum((flat < quantile_cfg.L) | (flat > quantile_cfg.H))) + stats["quantile"]["total_coords"] += len(flat) + + # Raw text truncation + decoded_r = roundtrip_raw(flat, args.precision) + stats["raw"]["rmsd"].append(conformer_rmsd(pos, decoded_r)) + stats["raw"]["max_err"].append(conformer_max_error(pos, decoded_r)) + stats["raw"]["total_coords"] += len(flat) + + if (mol_idx + 1) % 2000 == 0: + print(f" Processed {mol_idx + 1}/{n_mols} molecules " + f"({n_conformers} conformers so far)") + + print(f"\nTotal: {n_mols} molecules, {n_conformers} conformers\n") + + # Print results + print("=" * 72) + print(f"{'Method':<12} {'Mean RMSD':>10} {'Median RMSD':>12} {'p95 RMSD':>10} " + f"{'p99 RMSD':>10} {'Max RMSD':>10} {'Mean MaxErr':>12}") + print("=" * 72) + + for method in methods: + rmsds = np.array(stats[method]["rmsd"]) + max_errs = np.array(stats[method]["max_err"]) + + print(f"{method:<12} " + f"{np.mean(rmsds):>10.6f} " + f"{np.median(rmsds):>12.6f} " + f"{np.percentile(rmsds, 95):>10.6f} " + f"{np.percentile(rmsds, 99):>10.6f} " + f"{np.max(rmsds):>10.6f} " + f"{np.mean(max_errs):>12.6f}") + + print("=" * 72) + + # Overflow stats (only for binned methods) + print(f"\nOverflow statistics (coords outside [L, H]):") + for method in ["uniform", "quantile"]: + s = stats[method] + pct = 100.0 * s["overflow"] / s["total_coords"] if s["total_coords"] else 0 + print(f" {method:<12}: {s['overflow']:>8} / {s['total_coords']:<10} ({pct:.4f}%)") + + # Per-axis error distribution (sample from last batch) + print(f"\nPer-scalar absolute error distribution:") + print(f"{'Method':<12} {'Mean':>10} {'Median':>10} {'p95':>10} {'p99':>10} {'Max':>10}") + print("-" * 64) + + # Recompute on a sample for per-scalar stats + sample_flat = [] + for mol_idx, (smiles, confs) in enumerate(data[:min(1000, n_mols)]): + for mol in confs: + try: + pos = mol.GetConformer().GetPositions() + sample_flat.append(pos.flatten()) + except Exception: + continue + sample_flat = np.concatenate(sample_flat) + + for method, cfg in [("uniform", uniform_cfg), ("quantile", quantile_cfg)]: + decoded = roundtrip_bins(sample_flat, cfg) + errs = np.abs(sample_flat - decoded) + print(f"{method:<12} {np.mean(errs):>10.6f} {np.median(errs):>10.6f} " + f"{np.percentile(errs, 95):>10.6f} {np.percentile(errs, 99):>10.6f} " + f"{np.max(errs):>10.6f}") + + decoded_r = roundtrip_raw(sample_flat, args.precision) + errs_r = np.abs(sample_flat - decoded_r) + print(f"{'raw':<12} {np.mean(errs_r):>10.6f} {np.median(errs_r):>10.6f} " + f"{np.percentile(errs_r, 95):>10.6f} {np.percentile(errs_r, 99):>10.6f} " + f"{np.max(errs_r):>10.6f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/analyze_coord_centering.py b/scripts/analyze_coord_centering.py new file mode 100644 index 0000000..cf228cd --- /dev/null +++ b/scripts/analyze_coord_centering.py @@ -0,0 +1,213 @@ +""" +Analyze GEOM conformer coordinates: +1. Center all conformers (subtract centroid) and save centered pickle files +2. Verify centering (centroid displacement should be ~0) +3. Compute L/H bounds on centered train data (xyz-pooled and x-only) + at quantile levels: 1%, 0.5%, 0.1% +4. Count overflow for all thresholds on centered train and test +""" + +import pickle +import numpy as np +from rdkit.Geometry import Point3D + +DATA_DIR = "/data/molgen/geom_revisited" + + +def get_coords(mol): + """Extract N×3 coordinate array from an RDKit Mol with a conformer.""" + conf = mol.GetConformer() + return conf.GetPositions().astype(np.float64) + + +def load_split(name): + path = f"{DATA_DIR}/{name}_data.pickle" + if 'centered' in name: + name = name.replace("_centered", "") + path = f"{DATA_DIR}/{name}_data_centered.pickle" + print(f"Loading {path} ...", flush=True) + with open(path, "rb") as f: + data = pickle.load(f) + print(f" → {len(data)} molecules", flush=True) + return data + + +# ────────────────────────────────────────────────────────────────────────────── +# 1. Center conformers in-place and save +# ────────────────────────────────────────────────────────────────────────────── + +def center_split(data, split_name): + """ + Center every conformer in-place (subtract centroid from atom positions). + Saves the result to DATA_DIR/{split_name}_data_centered.pickle. + Returns the modified data list. + """ + print(f"\nCentering {split_name} ({len(data)} mols)...", flush=True) + n_confs = 0 + for smiles, confs in data: + for mol in confs: + conf = mol.GetConformer() + pos = conf.GetPositions() + centroid = pos.mean(axis=0) + new_pos = pos - centroid + for i in range(mol.GetNumAtoms()): + conf.SetAtomPosition(i, Point3D(*new_pos[i].tolist())) + n_confs += 1 + print(f" Centered {n_confs} conformers.", flush=True) + + out_path = f"{DATA_DIR}/{split_name}_data_centered.pickle" + print(f" Saving to {out_path} ...", flush=True) + with open(out_path, "wb") as f: + pickle.dump(data, f) + print(f" Saved.", flush=True) + return data + + +# ────────────────────────────────────────────────────────────────────────────── +# 2. Verify centering +# ────────────────────────────────────────────────────────────────────────────── + +def centering_stats(data, label, max_mols=None): + """Compute centroid L2-norm for each conformer (should be ~0 after centering).""" + norms = [] + mols_done = 0 + for smiles, confs in data: + for mol in confs: + X = get_coords(mol) + norms.append(np.linalg.norm(X.mean(axis=0))) + mols_done += 1 + if max_mols and mols_done >= max_mols: + break + norms = np.array(norms) + print(f"\n[{label}] centroid displacement over {len(norms)} conformers" + f" (from {mols_done} mols)") + print(f" mean = {norms.mean():.2e} Å") + print(f" median = {np.median(norms):.2e} Å") + print(f" max = {norms.max():.2e} Å") + print(f" >1e-9Å : {(norms > 1e-9).sum()} ({100*(norms > 1e-9).mean():.4f}%)") + return norms + + +# ────────────────────────────────────────────────────────────────────────────── +# 3. Collect coordinates (data already centered, no re-centering needed) +# ────────────────────────────────────────────────────────────────────────────── + +def collect_coords(data, label, axis=None): + """ + Returns a 1D array of raw scalar coordinate values (no centering). + axis=None → pool all three axes (x, y, z) + axis=0/1/2 → x / y / z only + """ + axis_label = {None: "xyz pooled", 0: "x only", 1: "y only", 2: "z only"} + all_vals = [] + print(f"\nCollecting coords from {label} ({len(data)} mols)" + f" [{axis_label[axis]}]...", flush=True) + for smiles, confs in data: + for mol in confs: + X = get_coords(mol) + vals = X.flatten() if axis is None else X[:, axis] + all_vals.append(vals) + all_vals = np.concatenate(all_vals) + print(f" total scalar values: {len(all_vals):,}", flush=True) + return all_vals + + +# ────────────────────────────────────────────────────────────────────────────── +# 4. Overflow counting (data already centered) +# ────────────────────────────────────────────────────────────────────────────── + +def count_overflow(data, L, H, label): + """Count mols and confs with at least one raw coordinate outside [L, H].""" + n_mols = len(data) + n_confs_total = 0 + n_confs_overflow = 0 + n_mols_overflow = 0 + + for smiles, confs in data: + mol_has_overflow = False + for mol in confs: + flat = get_coords(mol).flatten() # raw, no centering + has_out = bool((flat < L).any() or (flat > H).any()) + n_confs_total += 1 + if has_out: + n_confs_overflow += 1 + mol_has_overflow = True + if mol_has_overflow: + n_mols_overflow += 1 + + print(f"\n[{label}] L={L:.4f}, H={H:.4f}") + print(f" Mols with overflow: {n_mols_overflow:>6} / {n_mols:>6}" + f" ({100*n_mols_overflow/n_mols:.2f}%)") + print(f" Confs with overflow: {n_confs_overflow:>6} / {n_confs_total:>6}" + f" ({100*n_confs_overflow/n_confs_total:.2f}%)") + return n_mols_overflow, n_confs_overflow, n_mols, n_confs_total + + +# ────────────────────────────────────────────────────────────────────────────── +# Main +# ────────────────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + + # ── Load raw (uncentered) data ── + train_data = load_split("train") + test_data = load_split("test") + + # ── Step 1: Collect raw coords ── + print("\n" + "="*60) + print("STEP 1 — COORDINATE COLLECTION (raw, uncentered)") + print("="*60) + train_vals_xyz = collect_coords(train_data, "train", axis=None) + train_vals_x = collect_coords(train_data, "train", axis=0) + + # Tail quantiles: 1%, 0.1%, 0.01%, 0.001%, 0.0001% + q_pairs = { + "1% / 99%": (0.01, 0.99), + "0.1% / 99.9%": (0.001, 0.999), + "0.01% / 99.99%": (0.0001, 0.9999), + "0.001%/ 99.999%": (0.00001, 0.99999), + "0.0001%/99.9999%": (0.000001, 0.999999), + } + + print("\n" + "="*60) + print("STEP 2 — L / H BOUNDS (raw train): xyz-pooled vs x-only") + print("="*60) + + configs = {} + for q_name, (qlo, qhi) in q_pairs.items(): + for src_name, src_vals in [("xyz", train_vals_xyz), ("x-only", train_vals_x)]: + L = float(np.quantile(src_vals, qlo)) + H = float(np.quantile(src_vals, qhi)) + key = f"{q_name} [{src_name}]" + configs[key] = (L, H) + print(f"\n {key}") + print(f" L = {L:.4f} Å, H = {H:.4f} Å") + + # ── Step 2: Count overflow ── + print("\n" + "="*60) + print("STEP 3 — OVERFLOW COUNTS (raw, uncentered)") + print("="*60) + + results = {} + for thresh_name, (L, H) in configs.items(): + for split_label, split_data in [("TRAIN", train_data), ("TEST", test_data)]: + key = f"{thresh_name} | {split_label}" + nm_ov, nc_ov, nm_tot, nc_tot = count_overflow( + split_data, L, H, f"{thresh_name} — {split_label}") + results[key] = (nm_ov, nc_ov, nm_tot, nc_tot) + + # ── Summary table ── + print("\n" + "="*60) + print("SUMMARY TABLE") + print("="*60) + header = (f"{'Threshold + source':<32} {'Split':<7}" + f" {'Mol overflow':>22} {'Conf overflow':>22}") + print(header) + print("-" * len(header)) + for thresh_name in configs: + for split_label in ["TRAIN", "TEST"]: + key = f"{thresh_name} | {split_label}" + nm_ov, nc_ov, nm_tot, nc_tot = results[key] + print(f"{thresh_name:<32} {split_label:<7}" + f" {nm_ov:>7}/{nm_tot:<7} ({100*nm_ov/nm_tot:5.2f}%)" + f" {nc_ov:>8}/{nc_tot:<8} ({100*nc_ov/nc_tot:5.2f}%)") diff --git a/scripts/center_and_analyze.py b/scripts/center_and_analyze.py new file mode 100644 index 0000000..5e14bbf --- /dev/null +++ b/scripts/center_and_analyze.py @@ -0,0 +1,249 @@ +""" +Center all GEOM conformers, report centering stats, and save coordinate +distributions (as numpy arrays and plots) before and after centering. + +Outputs go to: scripts/centering_analysis/ + - centroid_norms_before.npy centroid displacement per conformer (before) + - centroid_norms_after.npy centroid displacement per conformer (after) + - coords_before_{xyz,x,y,z}.npy coordinate values before centering + - coords_after_{xyz,x,y,z}.npy coordinate values after centering + - distributions.png before/after coordinate histograms + - centroid_norms.png before/after centroid displacement histograms + +Centered data saved to DATA_DIR as: + train_data_centered.pickle / val_data_centered.pickle / test_data_centered.pickle +""" + +import os +import pickle +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from rdkit.Geometry import Point3D + +DATA_DIR = "/data/molgen/geom_revisited" +OUT_DIR = os.path.join(os.path.dirname(__file__), "centering_analysis") +os.makedirs(OUT_DIR, exist_ok=True) + +SAMPLE_MOLS = 5000 # mols used for distribution snapshots and centroid stats + + +# ────────────────────────────────────────────────────────────────────────────── +# I/O helpers +# ────────────────────────────────────────────────────────────────────────────── + +def get_coords(mol): + return mol.GetConformer().GetPositions().astype(np.float64) + + +def load_split(split): + path = f"{DATA_DIR}/{split}_data.pickle" + print(f"Loading {path} ...", flush=True) + with open(path, "rb") as f: + data = pickle.load(f) + print(f" → {len(data)} molecules", flush=True) + return data + + +def save_split(data, split): + path = f"{DATA_DIR}/{split}_data_centered.pickle" + print(f" Saving → {path} ...", flush=True) + with open(path, "wb") as f: + pickle.dump(data, f) + print(f" Done.", flush=True) + + +# ────────────────────────────────────────────────────────────────────────────── +# Centering +# ────────────────────────────────────────────────────────────────────────────── + +def center_in_place(data, split): + """Subtract per-conformer centroid from every atom. Mutates data in-place.""" + print(f"\nCentering {split} ({len(data)} mols)...", flush=True) + n_confs = 0 + for smiles, confs in data: + for mol in confs: + conf = mol.GetConformer() + pos = conf.GetPositions() + mu = pos.mean(axis=0) + new = pos - mu + for i in range(mol.GetNumAtoms()): + conf.SetAtomPosition(i, Point3D(*new[i].tolist())) + n_confs += 1 + print(f" {n_confs:,} conformers centered.", flush=True) + + +# ────────────────────────────────────────────────────────────────────────────── +# Stats collection +# ────────────────────────────────────────────────────────────────────────────── + +def collect_centroid_norms(data, max_mols=None): + """Centroid L2-norm for each conformer.""" + norms, done = [], 0 + for smiles, confs in data: + for mol in confs: + X = get_coords(mol) + norms.append(np.linalg.norm(X.mean(axis=0))) + done += 1 + if max_mols and done >= max_mols: + break + return np.array(norms) + + +def print_norm_stats(norms, label): + print(f"\n[{label}] n={len(norms):,} conformers") + for p in [50, 75, 90, 95, 99, 100]: + print(f" p{p:3d} = {np.percentile(norms, p):.6f} Å") + print(f" mean = {norms.mean():.6f} Å") + print(f" >1e-9 : {(norms > 1e-9).sum():,} ({100*(norms > 1e-9).mean():.4f}%)") + print(f" >0.01 : {(norms > 0.01).sum():,} ({100*(norms > 0.01).mean():.2f}%)") + print(f" >0.1 : {(norms > 0.1).sum():,} ({100*(norms > 0.1).mean():.2f}%)") + + +def collect_coord_distributions(data, max_mols=None): + """Pool all coordinate values (optionally limited to max_mols molecules).""" + xyz, x, y, z = [], [], [], [] + done = 0 + for smiles, confs in data: + for mol in confs: + X = get_coords(mol) + xyz.append(X.flatten()) + x.append(X[:, 0]) + y.append(X[:, 1]) + z.append(X[:, 2]) + done += 1 + if max_mols and done >= max_mols: + break + return { + "xyz": np.concatenate(xyz), + "x": np.concatenate(x), + "y": np.concatenate(y), + "z": np.concatenate(z), + } + + +# ────────────────────────────────────────────────────────────────────────────── +# Plotting +# ────────────────────────────────────────────────────────────────────────────── + +def plot_distributions(before, after, out_path): + keys = ["xyz", "x", "y", "z"] + fig, axes = plt.subplots(2, 4, figsize=(20, 8)) + fig.suptitle( + f"Coordinate distributions — before vs after centering " + f"(first {SAMPLE_MOLS:,} train mols)", fontsize=13) + + bins = np.linspace(-30, 30, 300) + for col, key in enumerate(keys): + b = np.clip(before[key], -30, 30) + a = np.clip(after[key], -30, 30) + + for row, (arr, raw, color, tag) in enumerate([ + (b, before[key], "steelblue", "BEFORE"), + (a, after[key], "coral", "AFTER"), + ]): + ax = axes[row, col] + ax.hist(arr, bins=bins, color=color, alpha=0.85) + ax.set_title(f"{tag} — {key}", fontsize=11) + ax.set_xlabel("Å") + ax.set_ylabel("count") + ax.text(0.97, 0.95, + f"μ={raw.mean():.3f}\nσ={raw.std():.3f}", + transform=ax.transAxes, ha="right", va="top", fontsize=9) + + plt.tight_layout() + plt.savefig(out_path, dpi=150) + plt.close() + print(f" Saved → {out_path}") + + +def plot_centroid_norms(before, after, out_path): + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) + fig.suptitle("Centroid displacement |μ| per conformer", fontsize=13) + + ax1.hist(np.clip(before, 0, 5), bins=200, color="steelblue", alpha=0.85) + ax1.set_title("BEFORE centering") + ax1.set_xlabel("|centroid| (Å)") + ax1.set_ylabel("count") + ax1.text(0.97, 0.95, + f"mean={before.mean():.3f}\nmedian={np.median(before):.3f}\n" + f"p99={np.percentile(before,99):.3f}\nmax={before.max():.3f}", + transform=ax1.transAxes, ha="right", va="top", fontsize=9) + + ax2.hist(after, bins=50, color="coral", alpha=0.85) + ax2.set_title("AFTER centering") + ax2.set_xlabel("|centroid| (Å)") + ax2.set_ylabel("count") + ax2.text(0.97, 0.95, + f"mean={after.mean():.2e}\nmax={after.max():.2e}", + transform=ax2.transAxes, ha="right", va="top", fontsize=9) + + plt.tight_layout() + plt.savefig(out_path, dpi=150) + plt.close() + print(f" Saved → {out_path}") + + +# ────────────────────────────────────────────────────────────────────────────── +# Main +# ────────────────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + + train_data = load_split("train") + val_data = load_split("val") + test_data = load_split("test") + + # ── BEFORE ── + print("\n" + "="*60) + print(f"BEFORE CENTERING (sample {SAMPLE_MOLS:,} train mols)") + print("="*60) + + norms_before = collect_centroid_norms(train_data, max_mols=SAMPLE_MOLS) + print_norm_stats(norms_before, "centroid |μ| BEFORE") + np.save(os.path.join(OUT_DIR, "centroid_norms_before.npy"), norms_before) + print(f" Saved centroid_norms_before.npy ({len(norms_before):,} values)") + + print(f"\nCollecting coord distributions...", flush=True) + dists_before = collect_coord_distributions(train_data, max_mols=SAMPLE_MOLS) + for key, arr in dists_before.items(): + np.save(os.path.join(OUT_DIR, f"coords_before_{key}.npy"), arr) + print(f" Saved coords_before_{key}.npy " + f"({len(arr):,} vals, mean={arr.mean():.3f}, std={arr.std():.3f})") + + # ── CENTER + SAVE ── + print("\n" + "="*60) + print("CENTERING AND SAVING ALL SPLITS") + print("="*60) + for data, split in [(train_data, "train"), (val_data, "val"), (test_data, "test")]: + center_in_place(data, split) + save_split(data, split) + + # ── AFTER ── + print("\n" + "="*60) + print(f"AFTER CENTERING (same {SAMPLE_MOLS:,} train mols)") + print("="*60) + + norms_after = collect_centroid_norms(train_data, max_mols=SAMPLE_MOLS) + print_norm_stats(norms_after, "centroid |μ| AFTER") + np.save(os.path.join(OUT_DIR, "centroid_norms_after.npy"), norms_after) + print(f" Saved centroid_norms_after.npy ({len(norms_after):,} values)") + + print(f"\nCollecting coord distributions after centering...", flush=True) + dists_after = collect_coord_distributions(train_data, max_mols=SAMPLE_MOLS) + for key, arr in dists_after.items(): + np.save(os.path.join(OUT_DIR, f"coords_after_{key}.npy"), arr) + print(f" Saved coords_after_{key}.npy " + f"({len(arr):,} vals, mean={arr.mean():.2e}, std={arr.std():.3f})") + + # ── PLOTS ── + print("\n" + "="*60) + print("SAVING PLOTS") + print("="*60) + plot_distributions(dists_before, dists_after, + os.path.join(OUT_DIR, "distributions.png")) + plot_centroid_norms(norms_before, norms_after, + os.path.join(OUT_DIR, "centroid_norms.png")) + + print(f"\nAll outputs → {OUT_DIR}") diff --git a/scripts/compute_range_R.py b/scripts/compute_range_R.py new file mode 100644 index 0000000..01a0075 --- /dev/null +++ b/scripts/compute_range_R.py @@ -0,0 +1,96 @@ +import argparse +import math +import pickle + +import numpy as np + + +DATA_DIR = "/data/molgen/geom_revisited" + + +def load_split(data_dir, name): + path = f"{data_dir}/{name}_data.pickle" + print(f"Loading {path} ...", flush=True) + with open(path, "rb") as f: + data = pickle.load(f) + print(f" -> {len(data)} molecules", flush=True) + return data + + +def conformer_radius_proxies(data): + radii = [] + for _smiles, confs in data: + for mol in confs: + pos = mol.GetConformer().GetPositions() + radii.append(np.abs(pos).max()) + return np.array(radii, dtype=np.float64) + + +def round_up(value, step=0.5): + return math.ceil(value / step) * step + + +def count_overflow(data, R, label): + n_confs = 0 + n_overflow = 0 + for _smiles, confs in data: + for mol in confs: + pos = mol.GetConformer().GetPositions() + n_confs += 1 + if np.abs(pos).max() > R: + n_overflow += 1 + pct = 100 * n_overflow / n_confs if n_confs else 0 + print(f" [{label}] overflow: {n_overflow:>7} / {n_confs:<7} ({pct:.4f}%)") + return n_overflow, n_confs + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--data_dir", type=str, default=DATA_DIR, + help="Directory with {train,val,test}_data_centered.pickle", + ) + parser.add_argument( + "--quantile", type=float, default=0.9999, + help="Quantile level for choosing R (default: 0.9999 = 99.99%%)", + ) + parser.add_argument( + "--round_step", type=float, default=0.5, + help="Round R up to this multiple (default: 0.5 A)", + ) + args = parser.parse_args() + + # -- Load centered data -- + train_data = load_split(args.data_dir, "train") + + # -- Conformer radius proxies on train -- + print("\nComputing conformer radius proxies (m = max|Xc|) on train...", flush=True) + m_train = conformer_radius_proxies(train_data) + print(f" {len(m_train):,} conformers") + + # -- Percentile table -- + print("\nPercentile table (train):") + for p in [50, 75, 90, 95, 99, 99.5, 99.9, 99.95, 99.99, 99.999, 100]: + val = np.percentile(m_train, p) + print(f" p{p:<8} = {val:.4f} A") + + # -- Pick R -- + R_raw = float(np.quantile(m_train, args.quantile)) + R = round_up(R_raw, args.round_step) + print(f"\nR_raw (quantile {args.quantile}) = {R_raw:.4f} A") + print(f"R (rounded up to {args.round_step}) = {R:.1f} A") + + # -- Overflow counts -- + print(f"\nOverflow counts with R = {R:.1f} (coords outside [-{R:.1f}, {R:.1f}]):") + count_overflow(train_data, R, "train") + + # Load val/test if available + for split in ("val", "test"): + try: + split_data = load_split(args.data_dir, split) + count_overflow(split_data, R, split) + except FileNotFoundError: + print(f" [{split}] skipped (file not found)") + + print(f"\n>>> Recommended range: [-{R:.1f}, {R:.1f}]") + print(f'>>> CLI flag: --ranges "[-{R:.1f}, {R:.1f}], [-{R:.1f}, {R:.1f}], [-{R:.1f}, {R:.1f}]"') diff --git a/scripts/count_geom_revisited_train_pickle.sh b/scripts/count_geom_revisited_train_pickle.sh new file mode 100755 index 0000000..4b4f1bf --- /dev/null +++ b/scripts/count_geom_revisited_train_pickle.sh @@ -0,0 +1,98 @@ +#!/bin/bash +# Usage: +# sbatch --partition=research --job-name=count-geom-revisited-train \ +# --cpus-per-task=1 --mem=32G --time=02:00:00 \ +# scripts/count_geom_revisited_train_pickle.sh +# +# Optional overrides: +# REPO_ROOT=/path/to/3DMolGen +# VENV_PATH=/path/to/.venv/bin/activate +# PICKLE_PATH=/data/vtarasov/geom_revisited/train_data.pickle +# +#SBATCH --output=/home/vtarasov/slurm_logs/dataset_counts/geom-revisited-train-%j.out +#SBATCH --error=/home/vtarasov/slurm_logs/dataset_counts/geom-revisited-train-%j.err + +set -euo pipefail + +REPO_ROOT=${REPO_ROOT:-/home/vtarasov/code/3DMolGen} +VENV_PATH=${VENV_PATH:-$REPO_ROOT/.venv/bin/activate} +PICKLE_PATH=${PICKLE_PATH:-/data/vtarasov/geom_revisited/train_data.pickle} + +mkdir -p /home/vtarasov/slurm_logs/dataset_counts + +echo "Job ID: ${SLURM_JOB_ID:-local}" +echo "Node: $(hostname)" +echo "Started: $(date)" +echo "Repo root: $REPO_ROOT" +echo "Pickle path: $PICKLE_PATH" +echo "" + +source "$VENV_PATH" +cd "$REPO_ROOT" + +python -u - <<'PY' +import pickle +from pathlib import Path + + +pickle_path = Path("/data/vtarasov/geom_revisited/train_data.pickle") +override = Path(__import__("os").environ.get("PICKLE_PATH", str(pickle_path))) +pickle_path = override + +if not pickle_path.exists(): + raise FileNotFoundError(f"Pickle file not found: {pickle_path}") + +with pickle_path.open("rb") as fh: + data = pickle.load(fh) + + +def count_from_entry(entry): + if isinstance(entry, tuple) and len(entry) >= 2: + mols = entry[1] + try: + return 1, len(mols) + except TypeError: + return 1, 0 + if isinstance(entry, dict): + if "mols" in entry: + mols = entry["mols"] + try: + return 1, len(mols) + except TypeError: + return 1, 0 + if "conformers" in entry: + conformers = entry["conformers"] + try: + return 1, len(conformers) + except TypeError: + return 1, 0 + return 1, 0 + + +if isinstance(data, dict): + items = list(data.items()) + molecule_count = len(items) + conformer_count = 0 + for _, value in items: + try: + conformer_count += len(value) + except TypeError: + pass +elif isinstance(data, list): + molecule_count = 0 + conformer_count = 0 + for entry in data: + mols, confs = count_from_entry(entry) + molecule_count += mols + conformer_count += confs +else: + raise TypeError(f"Unsupported pickle top-level type: {type(data).__name__}") + +print("GEOM revisited train counts") +print(f"pickle_path: {pickle_path}") +print(f"molecules: {molecule_count:,}") +print(f"conformers: {conformer_count:,}") +PY + +echo "" +echo "Finished: $(date)" diff --git a/scripts/count_revisited_train_tokens.sh b/scripts/count_revisited_train_tokens.sh new file mode 100755 index 0000000..74831a4 --- /dev/null +++ b/scripts/count_revisited_train_tokens.sh @@ -0,0 +1,167 @@ +#!/bin/bash +# Usage: +# sbatch --partition=research --job-name=count-revisited-train-tokens \ +# --cpus-per-task=8 --mem=64G --time=24:00:00 \ +# scripts/count_revisited_train_tokens.sh +# +# Counts tokens on the revisited train splits currently defined in paths.yaml. +# Grouped datasets are counted with serialization_mode=isomer_units and +# non-grouped datasets with serialization_mode=pairs. +# Note: paths.yaml currently contains a duplicate +# `revisited_cartesian_isomeric_grouped_train` entry and does not define +# `revisited_cartesian_isomeric_train`, so this script targets the unique train +# aliases that are available today. +# +# Optional overrides: +# REPO_ROOT=/path/to/3DMolGen +# VENV_PATH=/path/to/.venv/bin/activate +# OUTPUT_ROOT=/path/to/output-dir +# SEQ_LEN=4096 +# SAMPLE_LINES=1000 +# BATCH_SIZE=4 +# UNIT_BATCH_SIZE=64 +# SAMPLE_UNITS=10000 +# SAMPLE_SAMPLES=1000 +# SAMPLE_LINES_FOR_UNITS=1000 +# SHUFFLE=true +# GROUPED_ESTIMATE=exact # exact | fast | sample_only +# +#SBATCH --output=/home/vtarasov/slurm_logs/token_counts/revisited-train-%j.out +#SBATCH --error=/home/vtarasov/slurm_logs/token_counts/revisited-train-%j.err + +set -euo pipefail + +REPO_ROOT=${REPO_ROOT:-/home/vtarasov/code/3DMolGen} +VENV_PATH=${VENV_PATH:-$REPO_ROOT/.venv/bin/activate} +SEQ_LEN=${SEQ_LEN:-4096} +SAMPLE_LINES=${SAMPLE_LINES:-1000} +BATCH_SIZE=${BATCH_SIZE:-4} +UNIT_BATCH_SIZE=${UNIT_BATCH_SIZE:-64} +SAMPLE_UNITS=${SAMPLE_UNITS:-10000} +SAMPLE_SAMPLES=${SAMPLE_SAMPLES:-1000} +SAMPLE_LINES_FOR_UNITS=${SAMPLE_LINES_FOR_UNITS:-1000} +SHUFFLE=${SHUFFLE:-false} +GROUPED_ESTIMATE=${GROUPED_ESTIMATE:-exact} + +TIMESTAMP=$(date +%Y%m%d-%H%M%S) +JOB_TAG=${SLURM_JOB_ID:-local} +OUTPUT_ROOT=${OUTPUT_ROOT:-$REPO_ROOT/outputs/token_counts/revisited_train/${JOB_TAG}-${TIMESTAMP}} + +mkdir -p "$OUTPUT_ROOT" +mkdir -p /home/vtarasov/slurm_logs/token_counts + +echo "Job ID: ${SLURM_JOB_ID:-local}" +echo "Node: $(hostname)" +echo "Started: $(date)" +echo "Repo root: $REPO_ROOT" +echo "Output root: $OUTPUT_ROOT" +echo "Seq len: $SEQ_LEN" +echo "Sample lines: $SAMPLE_LINES" +echo "Batch size: $BATCH_SIZE" +echo "Unit batch size: $UNIT_BATCH_SIZE" +echo "Sample units: $SAMPLE_UNITS" +echo "Sample samples: $SAMPLE_SAMPLES" +echo "Sample unit lines: $SAMPLE_LINES_FOR_UNITS" +echo "Shuffle: $SHUFFLE" +echo "Grouped estimate: $GROUPED_ESTIMATE" +echo "" + +source "$VENV_PATH" +cd "$REPO_ROOT" + +build_grouped_extra_args() { + case "$GROUPED_ESTIMATE" in + exact) + echo "--exact-estimate" + ;; + fast) + echo "--fast-estimate --sample-units $SAMPLE_UNITS" + ;; + sample_only) + echo "--sample-only --sample-samples $SAMPLE_SAMPLES --sample-lines-for-units $SAMPLE_LINES_FOR_UNITS" + ;; + *) + echo "Unsupported GROUPED_ESTIMATE='$GROUPED_ESTIMATE' (expected: exact, fast, sample_only)" >&2 + exit 1 + ;; + esac +} + +run_count() { + local dataset_alias="$1" + local serialization_mode="$2" + local tokenizer_alias="$3" + local label="$4" + local safe_name="${dataset_alias//\//_}" + local log_path="$OUTPUT_ROOT/${safe_name}.log" + local extra_args=() + + if [[ "$serialization_mode" == "isomer_units" ]]; then + read -r -a extra_args <<< "$(build_grouped_extra_args)" + extra_args+=( + --unit-batch-size "$UNIT_BATCH_SIZE" + ) + else + extra_args+=( + --sample-lines "$SAMPLE_LINES" + --batch-size "$BATCH_SIZE" + ) + fi + + if [[ "$SHUFFLE" == "true" ]]; then + extra_args+=(--shuffle) + fi + + echo "============================================================" + echo "Dataset: $dataset_alias" + echo "Label: $label" + echo "Serialization mode: $serialization_mode" + echo "Tokenizer: $tokenizer_alias" + echo "Log path: $log_path" + echo "Started: $(date)" + echo "============================================================" + + python -m molgen3D.training.pretraining.dataprocessing.count_tokens \ + --train-path "$dataset_alias" \ + --skip-validation \ + --seq-len "$SEQ_LEN" \ + --batch-size "$BATCH_SIZE" \ + --serialization-mode "$serialization_mode" \ + --tokenizers "$tokenizer_alias" \ + "${extra_args[@]}" | tee "$log_path" + + echo "" +} + +run_count \ + "revisited_cartesian_isomeric_grouped_train" \ + "isomer_units" \ + "qwen3_0.6b_custom" \ + "grouped cartesian" + +run_count \ + "revisited_quantile_binned_isomeric_grouped_train" \ + "isomer_units" \ + "qwen3_0.6b_binned_258" \ + "grouped quantile_binned" + +run_count \ + "revisited_uniform_binned_isomeric_grouped_train" \ + "isomer_units" \ + "qwen3_0.6b_binned_258" \ + "grouped uniform_binned" + +run_count \ + "revisited_quantile_binned_isomeric_train" \ + "pairs" \ + "qwen3_0.6b_binned_258" \ + "non-grouped quantile_binned" + +run_count \ + "revisited_uniform_binned_isomeric_train" \ + "pairs" \ + "qwen3_0.6b_binned_258" \ + "non-grouped uniform_binned" + +echo "Finished: $(date)" +echo "Logs written to: $OUTPUT_ROOT" diff --git a/scripts/fit_bins.py b/scripts/fit_bins.py new file mode 100644 index 0000000..39f6fe7 --- /dev/null +++ b/scripts/fit_bins.py @@ -0,0 +1,132 @@ +import argparse +import os +import pickle + +import numpy as np + +# Append project root so the import works when running as a script +import sys +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from molgen3D.data_processing.smiles_encoder_decoder import ( + fit_uniform_bins, + fit_quantile_bins, +) + + +DATA_DIR = "/data/molgen/geom_revisited" + + +def load_split(data_dir, name): + path = f"{data_dir}/{name}_data.pickle" + print(f"Loading {path} ...", flush=True) + with open(path, "rb") as f: + data = pickle.load(f) + print(f" -> {len(data)} molecules", flush=True) + return data + + +def pool_coords(data, x_only=True): + all_vals = [] + for _smiles, confs in data: + for mol in confs: + pos = mol.GetConformer().GetPositions() + if x_only: + all_vals.append(pos[:, 0]) + else: + all_vals.append(pos.flatten()) + return np.concatenate(all_vals) + + +def overflow_stats(data, L, H, label): + n_confs = n_overflow = 0 + for _smiles, confs in data: + for mol in confs: + flat = mol.GetConformer().GetPositions().flatten() + n_confs += 1 + if (flat < L).any() or (flat > H).any(): + n_overflow += 1 + pct = 100 * n_overflow / n_confs if n_confs else 0 + print(f" [{label}] overflow: {n_overflow:>7} / {n_confs:<7} ({pct:.4f}%)") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--data_dir", type=str, default=DATA_DIR) + parser.add_argument("--out_dir", type=str, default="../src/molgen3D/config/bin_configs") + parser.add_argument("--n_bins", type=int, default=256) + parser.add_argument("--q_low", type=float, default=0.0001) + parser.add_argument("--q_high", type=float, default=0.9999) + args = parser.parse_args() + + os.makedirs(args.out_dir, exist_ok=True) + + # -- Load and pool coordinates -- + train_data = load_split(args.data_dir, "train") + + # Pool all xyz for uniform bins and distribution summary + print("\nPooling all scalar coordinates from train (xyz)...", flush=True) + V_all = pool_coords(train_data) + print(f" {len(V_all):,} scalar values") + + # Pool X-only for quantile bins (X is representative of all axes + # since molecular orientations are random) + print("Pooling X-only coordinates from train...", flush=True) + V_x = pool_coords(train_data, x_only=True) + print(f" {len(V_x):,} scalar values") + + # -- Distribution summary -- + print("\nCoordinate distribution (train, pooled xyz):") + for p in [0.01, 0.1, 1, 5, 25, 50, 75, 95, 99, 99.9, 99.99]: + print(f" p{p:<6} = {np.percentile(V_all, p):+.4f} A") + print(f" min = {V_all.min():+.4f} A") + print(f" max = {V_all.max():+.4f} A") + + # -- Fit uniform (from pooled xyz) -- + print(f"\n{'='*60}") + print(f"Fitting UNIFORM bins (B={args.n_bins}, q=[{args.q_low}, {args.q_high}])") + print(f"{'='*60}") + uniform_cfg = fit_uniform_bins(V_all, n_bins=args.n_bins, + q_low=args.q_low, q_high=args.q_high) + print(f" L = {uniform_cfg.L:+.4f} A") + print(f" H = {uniform_cfg.H:+.4f} A") + print(f" w = {(uniform_cfg.H - uniform_cfg.L) / uniform_cfg.n_bins:.6f} A") + print(f" digit_width = {uniform_cfg.digit_width}") + + uniform_path = os.path.join(args.out_dir, "uniform_bins.json") + uniform_cfg.save(uniform_path) + print(f" Saved -> {uniform_path}") + + # -- Fit quantile (from X-only) -- + print(f"\n{'='*60}") + print(f"Fitting QUANTILE bins (B={args.n_bins}, q=[{args.q_low}, {args.q_high}], X-only)") + print(f"{'='*60}") + quantile_cfg = fit_quantile_bins(V_x, n_bins=args.n_bins, + q_low=args.q_low, q_high=args.q_high) + print(f" L = {quantile_cfg.L:+.4f} A") + print(f" H = {quantile_cfg.H:+.4f} A") + print(f" edge range: [{quantile_cfg.edges[0]:+.4f}, {quantile_cfg.edges[-1]:+.4f}]") + print(f" median bin width = {np.median(np.diff(quantile_cfg.edges)):.6f} A") + print(f" digit_width = {quantile_cfg.digit_width}") + + quantile_path = os.path.join(args.out_dir, "quantile_bins.json") + quantile_cfg.save(quantile_path) + print(f" Saved -> {quantile_path}") + + # -- Overflow stats -- + print(f"\nOverflow counts (coords outside [L, H]):") + overflow_stats(train_data, uniform_cfg.L, uniform_cfg.H, "train-uniform") + overflow_stats(train_data, quantile_cfg.L, quantile_cfg.H, "train-quantile") + + for split in ("val", "test"): + try: + split_data = load_split(args.data_dir, split) + overflow_stats(split_data, uniform_cfg.L, uniform_cfg.H, f"{split}-uniform") + overflow_stats(split_data, quantile_cfg.L, quantile_cfg.H, f"{split}-quantile") + except FileNotFoundError: + print(f" [{split}] skipped (file not found)") + + print(f"\nConfigs saved to {args.out_dir}/") + print(f" uniform_bins.json — use with encode_cartesian_with_config") + print(f" quantile_bins.json — use with encode_cartesian_with_config") diff --git a/scripts/preprocess_geom_revisited.sh b/scripts/preprocess_geom_revisited.sh new file mode 100755 index 0000000..31c5be0 --- /dev/null +++ b/scripts/preprocess_geom_revisited.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# Usage: +# sbatch --partition=research --job-name=preprocess-revisited \ +# --cpus-per-task=32 --mem=256G --time=48:00:00 \ +# scripts/preprocess_geom_revisited.sh +# +# Memory note: the train split pickle is ~6.5 GB on disk and expands to +# ~50 GB in RAM. multiprocessing.Pool forks the parent, so Python's +# refcount CoW quickly multiplies that by the worker count. Cap workers +# at 32 to stay well within 256 GB; more workers give negligible extra +# throughput for this CPU-light, fork-heavy workload. +#SBATCH --output=/home/vtarasov/slurm_logs/preprocess/revisited-%j.out +#SBATCH --error=/home/vtarasov/slurm_logs/preprocess/revisited-%j.err + +set -euo pipefail + +GEOM_RAW=/data/vtarasov/geom_revisited +DEST=/data/vtarasov +BIN_CONFIGS="/home/vtarasov/code/3DMolGen/src/molgen3D/config/bin_configs" +# Cap at 32: more workers multiply CoW memory without improving throughput. +_RAW_CPUS=${SLURM_CPUS_PER_TASK:-32} +WORKERS=$(( _RAW_CPUS > 32 ? 32 : _RAW_CPUS )) + +echo "Job ID: ${SLURM_JOB_ID:-local}" +echo "Node: $(hostname)" +echo "CPUs: $WORKERS" +echo "Started: $(date)" +echo "" + +source "/home/vtarasov/code/3DMolGen/.venv/bin/activate" + +cd "/home/vtarasov/code/3DMolGen" + +echo "==========================================" +echo " 1/3 cartesian_v2" +echo "==========================================" +python -m molgen3D.data_processing.data_preprocessing_revisited \ + --geom_raw_path "$GEOM_RAW" \ + --dest "$DEST" \ + --run_name geom_revisited_cartesian_isomeric \ + --embedding_type cartesian_v2 \ + --num_workers "$WORKERS" \ + --isomeric + +echo "==========================================" +echo " 2/3 quantile_binned" +echo "==========================================" +python -m molgen3D.data_processing.data_preprocessing_revisited \ + --geom_raw_path "$GEOM_RAW" \ + --dest "$DEST" \ + --run_name geom_revisited_quantile_binned_isomeric \ + --embedding_type quantile_binned \ + --bin_config_path "$BIN_CONFIGS/quantile_bins.json" \ + --num_workers "$WORKERS" \ + --isomeric + +echo "==========================================" +echo " 3/3 uniform_binned" +echo "==========================================" +python -m molgen3D.data_processing.data_preprocessing_revisited \ + --geom_raw_path "$GEOM_RAW" \ + --dest "$DEST" \ + --run_name geom_revisited_uniform_binned_isomeric \ + --embedding_type uniform_binned \ + --bin_config_path "$BIN_CONFIGS/uniform_bins.json" \ + --num_workers "$WORKERS" \ + --isomeric + +echo "==========================================" +echo " All 3 runs complete." +echo " Finished: $(date)" +echo " Output dirs:" +echo " $DEST/geom_revisited_cartesian_isomeric" +echo " $DEST/geom_revisited_quantile_binned_isomeric" +echo " $DEST/geom_revisited_uniform_binned_isomeric" +echo "==========================================" diff --git a/scripts/preprocess_geom_revisited_paired_grouped.sh b/scripts/preprocess_geom_revisited_paired_grouped.sh new file mode 100755 index 0000000..1c0f715 --- /dev/null +++ b/scripts/preprocess_geom_revisited_paired_grouped.sh @@ -0,0 +1,94 @@ +#!/bin/bash +# Usage: +# sbatch --partition=research --job-name=preprocess-revisited-grouped \ +# --cpus-per-task=128 --mem=256G --time=12:00:00 \ +# scripts/preprocess_geom_revisited_paired_grouped.sh +# +# This script runs grouped revisited preprocessing across three encodings: +# cartesian, quantile_binned, and uniform_binned. + +set -euo pipefail + +GEOM_RAW=${GEOM_RAW:-/data/vtarasov/geom_revisited} +DEST=${DEST:-/data/vtarasov} +BIN_CONFIGS=${BIN_CONFIGS:-/home/vtarasov/code/3DMolGen/src/molgen3D/config/bin_configs} +REPO_ROOT=${REPO_ROOT:-/home/vtarasov/code/3DMolGen} +VENV_PATH=${VENV_PATH:-$REPO_ROOT/.venv/bin/activate} + +# Cap at 32: more workers multiply CoW memory without improving throughput. +_RAW_CPUS=${SLURM_CPUS_PER_TASK:-32} +WORKERS=$(( _RAW_CPUS > 32 ? 32 : _RAW_CPUS )) + +echo "Job ID: ${SLURM_JOB_ID:-local}" +echo "Node: $(hostname)" +echo "CPUs: $WORKERS" +echo "Started: $(date)" +echo "GEOM_RAW: $GEOM_RAW" +echo "DEST: $DEST" +echo "" + +source "$VENV_PATH" +cd "$REPO_ROOT" + +run_job() { + local step="$1" + local total="$2" + local label="$3" + local module="$4" + local run_name="$5" + local embedding_type="$6" + local bin_config="${7:-}" + + echo "==========================================" + echo " ${step}/${total} ${label}" + echo "==========================================" + + if [[ -n "$bin_config" ]]; then + python -m "$module" \ + --geom_raw_path "$GEOM_RAW" \ + --dest "$DEST" \ + --run_name "$run_name" \ + --embedding_type "$embedding_type" \ + --bin_config_path "$bin_config" \ + --num_workers "$WORKERS" \ + --isomeric + else + python -m "$module" \ + --geom_raw_path "$GEOM_RAW" \ + --dest "$DEST" \ + --run_name "$run_name" \ + --embedding_type "$embedding_type" \ + --num_workers "$WORKERS" \ + --isomeric + fi +} + + +run_job 1 3 \ + "grouped cartesian" \ + "molgen3D.data_processing.preprocess_geom_grouped_revisited" \ + "geom_revisited_cartesian_isomeric_grouped" \ + "cartesian" + +run_job 2 3 \ + "grouped quantile_binned" \ + "molgen3D.data_processing.preprocess_geom_grouped_revisited" \ + "geom_revisited_quantile_binned_isomeric_grouped" \ + "quantile_binned" \ + "$BIN_CONFIGS/quantile_bins.json" + +run_job 3 3 \ + "grouped uniform_binned" \ + "molgen3D.data_processing.preprocess_geom_grouped_revisited" \ + "geom_revisited_uniform_binned_isomeric_grouped" \ + "uniform_binned" \ + "$BIN_CONFIGS/uniform_bins.json" + +echo "==========================================" +echo " All 3 grouped runs complete." +echo " Finished: $(date)" +echo " Output dirs:" +echo " $DEST/geom_revisited_cartesian_isomeric_grouped" +echo " $DEST/geom_revisited_quantile_binned_isomeric_grouped" +echo " $DEST/geom_revisited_uniform_binned_isomeric_grouped" +echo "==========================================" diff --git a/src/molgen3D/data_processing/data_preprocessing.py b/src/molgen3D/data_processing/data_preprocessing.py index 878167c..8c411e3 100644 --- a/src/molgen3D/data_processing/data_preprocessing.py +++ b/src/molgen3D/data_processing/data_preprocessing.py @@ -1,5 +1,4 @@ import argparse -import ast import glob import json import os @@ -7,7 +6,7 @@ import random from collections import defaultdict from multiprocessing import Pool -from typing import Any, Dict, Optional, Set, Tuple, List +from typing import Any, Dict, List, Optional, Set, Tuple import numpy as np from loguru import logger as log from rdkit import Chem, RDLogger @@ -16,12 +15,15 @@ from molgen3D.data_processing.utils import ( JsonlSplitWriter, filter_mols, + get_embedding_func_and_config, + parse_coordinate_ranges, save_processed_pickle, ) from molgen3D.data_processing.smiles_encoder_decoder import ( - encode_cartesian_v2, + BinConfig, encode_cartesian_binned, encode_cartesian_binned_v2, + encode_cartesian_with_config, ) from molgen3D.utils.utils import load_pkl @@ -29,12 +31,16 @@ def read_mol( - args: Tuple[str, int, int, Any, float, List[Tuple[float, float]], bool, str, str] + args: Tuple, ) -> Optional[Tuple[List[str], Dict[str, Any]]]: - mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir, _geom_root = args + mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir, _geom_root = args[:9] + bin_config = args[9] if len(args) > 9 else None + use_isomeric_smiles = args[10] if len(args) > 10 else False try: return _read_mol_impl( - mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir + mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir, + bin_config=bin_config, + use_isomeric_smiles=use_isomeric_smiles, ) except Exception as exc: log.error("Unhandled exception in read_mol | path={} | error={}", mol_path, exc) @@ -50,6 +56,8 @@ def _read_mol_impl( ranges: List[Tuple[float, float]], do_filter: bool, pickle_dir: str, + bin_config: Optional[BinConfig] = None, + use_isomeric_smiles: bool = False, ) -> Tuple[List[str], Dict[str, Any]]: mol_object = load_pkl(mol_path) geom_smiles = mol_object["smiles"] @@ -75,7 +83,9 @@ def _read_mol_impl( continue try: - if embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): embedded_smile, iso_smile = embedding_func(mol, bin_size=bin_size, ranges=ranges) else: embedded_smile, iso_smile = embedding_func(mol, precision=precision) @@ -85,6 +95,7 @@ def _read_mol_impl( continue # Compute nonisomeric SMILES only for conformers that encoded successfully + noniso = None try: noniso = Chem.MolToSmiles(Chem.RemoveHs(mol, sanitize=False), canonical=True, isomericSmiles=False) nonisomeric_smiles.add(noniso) @@ -93,10 +104,11 @@ def _read_mol_impl( except Exception: pass + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) samples.append( json.dumps( { - "canonical_smiles": iso_smile, + "canonical_smiles": canonical_smiles, "embedded_smiles": embedded_smile, }, separators=(",", ":"), @@ -158,19 +170,16 @@ def preprocess( bin_size: float = 0.104, ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", filter_ranges: str = None, + bin_config_path: Optional[str] = None, + use_isomeric_smiles: bool = False, ) -> None: if dest_path is None: raise ValueError("dest_path must be provided for preprocessing output") - embedding_registry = { - "cartesian_v2": encode_cartesian_v2, - "cartesian": encode_cartesian_v2, - "cartesian_binned": encode_cartesian_binned, - "cartesian_binned_v2": encode_cartesian_binned_v2, - } - if embedding_type not in embedding_registry: - raise ValueError(f"Unsupported embedding_type '{embedding_type}'. Options: {sorted(embedding_registry)}") - embedding_func = embedding_registry[embedding_type] + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) overall_total_input_mols = overall_total_confs = overall_total_mols = 0 overall_multi_distinct_graphs = overall_mol_with_dotted_smiles = overall_total_dotted_smiles = 0 @@ -199,13 +208,7 @@ def preprocess( if pickle_paths.size == 0: raise FileNotFoundError(f"No pickle files found under pattern {pickle_glob}") - # Parse ranges string once - try: - parsed_ranges = ast.literal_eval(f"[{ranges}]") - parsed_ranges = [tuple(r) for r in parsed_ranges] - except Exception as e: - log.error(f"Failed to parse ranges: {ranges}. Error: {e}") - parsed_ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + parsed_ranges = parse_coordinate_ranges(ranges) do_filter = False if filter_ranges is not None: @@ -250,6 +253,8 @@ def preprocess( do_filter, split_pickle_dirs[split_name], geom_raw_path, + bin_config, + use_isomeric_smiles, ) for path in mol_paths ] @@ -379,7 +384,8 @@ def preprocess( "--embedding_type", "-et", type=str, - choices=["cartesian", "cartesian_v2", "cartesian_binned", "cartesian_binned_v2"], + choices=["cartesian", "cartesian_v2", "cartesian_binned", "cartesian_binned_v2", + "uniform_binned", "quantile_binned"], default="cartesian_v2", help="Embedding type to use for enrichment.", ) @@ -446,6 +452,25 @@ def preprocess( default=None, help="Filter ranges for binned embedding.", ) + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + + parser.add_argument( + "--isomeric", + action="store_true", + help="Alias for --use_isomeric_smiles.", + ) + parser.add_argument( + "--sort_by", + type=str, + choices=["energy", "weight", "none"], + default="energy", + help="Sort conformers by energy, weight, or keep original order.", + ) args = parser.parse_args() dest_path = osp.join(args.dest, args.run_name) @@ -458,6 +483,8 @@ def preprocess( diagnose=False, ) + use_isomeric = args.isomeric + preprocess( geom_raw_path=args.geom_raw_path, indices_path=args.indices_path, @@ -471,5 +498,7 @@ def preprocess( bin_size=args.bin_size, ranges=args.ranges, filter_ranges=args.filter_ranges, + bin_config_path=args.bin_config_path, + use_isomeric_smiles=use_isomeric, ) \ No newline at end of file diff --git a/src/molgen3D/data_processing/data_preprocessing_revisited.py b/src/molgen3D/data_processing/data_preprocessing_revisited.py new file mode 100644 index 0000000..e521802 --- /dev/null +++ b/src/molgen3D/data_processing/data_preprocessing_revisited.py @@ -0,0 +1,467 @@ +import argparse +import json +import os +import os.path as osp +import pickle +import random +from collections import defaultdict +from multiprocessing import Pool +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np +from loguru import logger as log +from rdkit import Chem, RDLogger +from tqdm.auto import tqdm + +from molgen3D.data_processing.smiles_encoder_decoder import ( + encode_cartesian_binned, + encode_cartesian_binned_v2, + encode_cartesian_with_config, +) +from molgen3D.data_processing.utils import ( + JsonlSplitWriter, + filter_revisited_mols, + get_embedding_func_and_config, + get_revisited_split_path, + parse_coordinate_ranges, + save_processed_pickle, +) + +RDLogger.DisableLog("rdApp.*") + +_REVISITED_DATA: Optional[List] = None + + +def _process_revisited_mol_impl( + idx: int, + max_confs: int, + precision: int, + embedding_func: Any, + bin_size: float, + parsed_ranges: List[Tuple[float, float]], + do_filter: bool, + bin_config: Any, + use_isomeric_smiles: bool, +) -> Tuple[List[str], Dict[str, Any], List[Chem.Mol]]: + smiles, mols = _REVISITED_DATA[idx] + local_failures: Dict[str, int] = defaultdict(int) + + valid = filter_revisited_mols( + smiles=smiles, + mols=mols, + failures=local_failures, + max_confs=max_confs, + ) + + nonisomeric_smiles, dotted_smiles, isomeric_smiles = set(), set(), set() + samples: List[str] = [] + filtered_mols: List[Chem.Mol] = [] + + for mol in valid: + if do_filter: + pos = mol.GetConformer().GetPositions() + out_of_range = False + for i in range(3): + min_val, max_val = parsed_ranges[i] + if np.any(pos[:, i] < min_val) or np.any(pos[:, i] > max_val): + out_of_range = True + break + if out_of_range: + local_failures["coord_out_of_range"] += 1 + continue + + try: + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + embedded_smile, iso_smile = embedding_func( + mol, + bin_size=bin_size, + ranges=parsed_ranges, + ) + else: + embedded_smile, iso_smile = embedding_func(mol, precision=precision) + except Exception as exc: + log.error("Error encoding conformer | smiles={} | failure={}", smiles, exc) + local_failures["encoding_error"] += 1 + continue + + noniso = None + try: + noniso = Chem.MolToSmiles( + Chem.RemoveHs(mol, sanitize=False), + canonical=True, + isomericSmiles=False, + ) + nonisomeric_smiles.add(noniso) + if "." in noniso: + dotted_smiles.add(noniso) + except Exception: + pass + + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) + samples.append( + json.dumps( + { + "canonical_smiles": canonical_smiles, + "embedded_smiles": embedded_smile, + }, + separators=(",", ":"), + ) + + "\n" + ) + isomeric_smiles.add(iso_smile) + filtered_mols.append(mol) + + if not samples: + local_failures["no_samples_after_filtering"] += 1 + + stats = { + "path": smiles, + "geom_smiles": smiles, + "confs_count_pre_filter": len(mols), + "confs_count_post_filter": len(samples), + "nonisomeric_smiles_post_filter": len(nonisomeric_smiles), + "isomeric_smiles_post_filter": isomeric_smiles, + "num_distinct_smiles_with_dot": len(dotted_smiles), + "has_dotted_smiles": bool(dotted_smiles), + "failures": local_failures, + "processed_pickle_path": None, + } + return samples, stats, filtered_mols + + +def _process_revisited_by_idx(args: Tuple) -> Optional[Tuple]: + idx = args[0] + try: + return _process_revisited_mol_impl(idx, *args[1:]) + except Exception as exc: + log.error("Unhandled exception | idx={} | error={}", idx, exc) + return None + + +def preprocess_revisited( + geom_raw_path: str, + embedding_type: str, + num_workers: int = 20, + precision: int = 4, + splits: Optional[str] = None, + dest_path: Optional[str] = None, + max_confs: int = 30, + bin_size: float = 0.104, + ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + filter_ranges: str = None, + bin_config_path: Optional[str] = None, + isomeric: bool = False, + use_centered: bool = True, +) -> None: + global _REVISITED_DATA + + if dest_path is None: + raise ValueError("dest_path must be provided for preprocessing output") + + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) + parsed_ranges = parse_coordinate_ranges(ranges) + + do_filter = False + if filter_ranges is not None: + if isinstance(filter_ranges, str): + do_filter = filter_ranges.lower() in ("true", "1", "yes", "on") + else: + do_filter = bool(filter_ranges) + + requested_splits = [splits] if splits else ["train", "valid", "test"] + + strings_root = osp.join(dest_path, "processed_strings") + split_writers = { + split: JsonlSplitWriter(osp.join(strings_root, split), split, chunk_size=50_000) + for split in ("train", "valid", "test") + } + split_pickle_dirs = { + split: osp.join(dest_path, "processed_pickles", split) for split in ("train", "valid", "test") + } + for split_dir in split_pickle_dirs.values(): + os.makedirs(split_dir, exist_ok=True) + + overall_total_input = overall_total_confs = overall_total_mols = 0 + overall_failure_counts: Dict[str, int] = defaultdict(int) + + for split_name in requested_splits: + data_path = get_revisited_split_path( + geom_raw_path=geom_raw_path, + split_name=split_name, + use_centered=use_centered, + ) + if not osp.exists(data_path): + raise FileNotFoundError(f"Revisited split file not found: {data_path}") + + log.info("Loading {} ...", data_path) + with open(data_path, "rb") as fh: + _REVISITED_DATA = pickle.load(fh) + log.info("Loaded {:,} molecules for split '{}'", len(_REVISITED_DATA), split_name) + + n = len(_REVISITED_DATA) + overall_total_input += n + + job_args = [ + ( + idx, + max_confs, + precision, + embedding_func, + bin_size, + parsed_ranges, + do_filter, + bin_config, + isomeric, + ) + for idx in range(n) + ] + + conf_count_pre = conf_count_post = mol_count_post = 0 + split_num_mol_with_multi_distinct_graphs = split_num_mol_with_dotted_smiles = total_dotted_smiles = 0 + failure_counts: Dict[str, int] = defaultdict(int) + split_geom_to_iso_map: Dict[str, Set[str]] = defaultdict(set) + + geom_to_iso_path = osp.join(dest_path, f"{split_name}_geom_to_isomeric_smiles.jsonl") + geom_to_iso_fh = open(geom_to_iso_path, "w") + try: + with tqdm(total=n, dynamic_ncols=True, mininterval=0.2) as pbar: + with Pool(processes=num_workers) as pool: + chunk_size = max(1, n // max(num_workers * 2, 1)) + for result in pool.imap_unordered( + _process_revisited_by_idx, + job_args, + chunksize=chunk_size, + ): + if result is None: + failure_counts["unhandled_exception"] += 1 + pbar.update() + continue + + samples, stats, filtered_mols = result + split_writers[split_name].write(samples) + + geom_smiles = stats.get("geom_smiles") + if filtered_mols and geom_smiles: + try: + processed_pickle_path = save_processed_pickle( + split_dir=split_pickle_dirs[split_name], + geom_smiles=geom_smiles, + mols=filtered_mols, + ) + stats["processed_pickle_path"] = processed_pickle_path + except Exception as exc: + log.error( + "Failed to write processed pickle | geom_smiles={} | failure={}", + geom_smiles, + exc, + ) + + conf_count_pre += stats["confs_count_pre_filter"] + conf_count_post += stats["confs_count_post_filter"] + overall_total_confs += stats["confs_count_post_filter"] + total_dotted_smiles += stats.get("num_distinct_smiles_with_dot", 0) + + if stats["nonisomeric_smiles_post_filter"] > 1: + split_num_mol_with_multi_distinct_graphs += 1 + if stats.get("has_dotted_smiles", False): + split_num_mol_with_dotted_smiles += 1 + + for reason, count in stats["failures"].items(): + failure_counts[reason] += int(count) + + if stats["confs_count_post_filter"] > 0: + mol_count_post += 1 + overall_total_mols += 1 + + if geom_smiles: + for iso in stats.get("isomeric_smiles_post_filter", ()): + split_geom_to_iso_map[geom_smiles].add(iso) + geom_to_iso_fh.write( + json.dumps( + { + "geom_smiles": geom_smiles, + "isomeric_smiles": iso, + }, + separators=(",", ":"), + ) + + "\n" + ) + + pbar.update() + finally: + geom_to_iso_fh.close() + _REVISITED_DATA = None + + avg_confs = conf_count_post / mol_count_post if mol_count_post else 0.0 + success_rate = mol_count_post / n if n else 0.0 + split_report = { + "split": split_name, + "num_input_molecules": n, + "num_output_molecules": mol_count_post, + "num_input_conformers": conf_count_pre, + "total_conformers_after": conf_count_post, + "avg_conformers_per_molecule_after": avg_confs, + "success_rate": success_rate, + "failure_counts": dict(failure_counts), + "molecules_with_multiple_distinct_graphs": split_num_mol_with_multi_distinct_graphs, + "molecules_with_dotted_smiles": split_num_mol_with_dotted_smiles, + "total_dotted_smiles": total_dotted_smiles, + } + log.info(json.dumps({"split_summary": split_report}, ensure_ascii=False, separators=(",", ":"))) + + for reason, count in failure_counts.items(): + overall_failure_counts[reason] += count + + for writer in split_writers.values(): + writer.close() + + grand_total = sum(writer.total_samples for writer in split_writers.values()) + overall_success_rate = float(overall_total_mols) / max(1, overall_total_input) + run_summary = { + "grand_total_samples_written": grand_total, + "total_input_molecules": overall_total_input, + "molecules_after_filter": overall_total_mols, + "conformers_after_filter": overall_total_confs, + "avg_confs_per_mol_after": float(overall_total_confs) / max(1, overall_total_mols), + "overall_success_rate": overall_success_rate, + "overall_failure_counts": dict(overall_failure_counts), + } + log.info(json.dumps({"run_summary": run_summary}, ensure_ascii=False, separators=(",", ":"))) + + +if __name__ == "__main__": + random.seed(42) + parser = argparse.ArgumentParser() + parser.add_argument( + "--geom_raw_path", + "-p", + type=str, + required=True, + help="Path to the revisited GEOM split pickle files.", + ) + parser.add_argument( + "--dest", + "-d", + type=str, + default="/data/molgen/", + help="Destination directory for processed outputs.", + ) + parser.add_argument( + "--embedding_type", + "-et", + type=str, + choices=[ + "cartesian", + "cartesian_v2", + "cartesian_binned", + "cartesian_binned_v2", + "uniform_binned", + "quantile_binned", + ], + default="cartesian_v2", + help="Embedding type to use for enrichment.", + ) + parser.add_argument( + "--num_workers", + "-nw", + type=int, + default=max(4, os.cpu_count() or 4), + help="Number of worker processes.", + ) + parser.add_argument( + "--precision", + type=int, + default=4, + help="Numeric precision for encoded coordinates.", + ) + parser.add_argument( + "--splits", + type=str, + choices=["train", "valid", "test"], + default=None, + help="Optional single split to process.", + ) + parser.add_argument( + "--run_name", + type=str, + default="", + help="Run name, appended to destination directory.", + ) + parser.add_argument( + "--max_confs", + type=int, + default=30, + help="Maximum number of conformers per molecule.", + ) + parser.add_argument( + "--bin_size", + type=float, + default=0.104, + help="Bin size for binned embedding.", + ) + parser.add_argument( + "--ranges", + type=str, + default="[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + help="Ranges for binned embedding.", + ) + parser.add_argument( + "--filter_ranges", + type=str, + default=None, + help="Filter ranges for binned embedding.", + ) + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + parser.add_argument( + "--use_isomeric_smiles", + action="store_true", + help="If set, writes isomeric SMILES to canonical_smiles; otherwise writes non-isomeric SMILES.", + ) + parser.add_argument( + "--isomeric", + action="store_true", + help="Alias for --use_isomeric_smiles.", + ) + parser.add_argument( + "--use_centered", + action="store_true", + default=False, + help="Load *_data_centered.pickle instead of *_data.pickle.", + ) + args = parser.parse_args() + + dest_path = osp.join(args.dest, args.run_name) + os.makedirs(dest_path, exist_ok=True) + log.add( + osp.join(dest_path, "preprocessing.log"), + mode="w", + enqueue=True, + backtrace=False, + diagnose=False, + ) + + preprocess_revisited( + geom_raw_path=args.geom_raw_path, + embedding_type=args.embedding_type, + dest_path=dest_path, + max_confs=args.max_confs, + num_workers=args.num_workers, + precision=args.precision, + splits=args.splits, + bin_size=args.bin_size, + ranges=args.ranges, + filter_ranges=args.filter_ranges, + bin_config_path=args.bin_config_path, + isomeric=args.isomeric or args.use_isomeric_smiles, + use_centered=args.use_centered, + ) diff --git a/src/molgen3D/data_processing/preprocess_geom_grouped.py b/src/molgen3D/data_processing/preprocess_geom_grouped.py index d2132f8..d44241b 100644 --- a/src/molgen3D/data_processing/preprocess_geom_grouped.py +++ b/src/molgen3D/data_processing/preprocess_geom_grouped.py @@ -1,10 +1,8 @@ import argparse -import ast import glob import json import os import os.path as osp -import pickle import random from collections import defaultdict from multiprocessing import Pool @@ -15,11 +13,18 @@ from rdkit import Chem, RDLogger from tqdm.auto import tqdm -from molgen3D.data_processing.utils import JsonlSplitWriter +from molgen3D.data_processing.utils import ( + JsonlSplitWriter, + copy_single_conformer_mol, + extract_conf_meta, + get_embedding_func_and_config, + parse_coordinate_ranges, + save_grouped_pickle, +) from molgen3D.data_processing.smiles_encoder_decoder import ( encode_cartesian_binned, - encode_cartesian_v2, encode_cartesian_binned_v2, + encode_cartesian_with_config, ) from molgen3D.utils.utils import load_pkl @@ -35,81 +40,6 @@ def infer_geom_id(mol_path: str) -> str: return osp.splitext(osp.basename(mol_path))[0] -def copy_single_conformer_mol(mol: Chem.Mol) -> Chem.Mol: - copied = Chem.Mol(mol) - if copied.GetNumConformers() > 1: - conf = Chem.Conformer(copied.GetConformer(0)) - copied.RemoveAllConformers() - copied.AddConformer(conf, assignId=True) - return copied - - -def _get_prop_float(props: Dict[str, Any], key: str) -> Optional[float]: - if key in props: - try: - return float(props[key]) - except Exception: - return None - return None - - -def extract_conf_meta( - conf_meta: Optional[Dict[str, Any]], mol: Chem.Mol -) -> Tuple[Optional[float], Optional[float], Optional[int]]: - energy = weight = None - conf_id = None - - if conf_meta: - if "totalenergy" in conf_meta: - energy = conf_meta.get("totalenergy") - elif "relativeenergy" in conf_meta: - energy = conf_meta.get("relativeenergy") - - if "boltzmannweight" in conf_meta: - weight = conf_meta.get("boltzmannweight") - elif "weight" in conf_meta: - weight = conf_meta.get("weight") - - if "geom_id" in conf_meta: - conf_id = conf_meta.get("geom_id") - - if mol is not None: - mol_props = mol.GetPropsAsDict() - if energy is None: - energy = _get_prop_float(mol_props, "totalenergy") - if weight is None: - weight = _get_prop_float(mol_props, "boltzmannweight") - if conf_id is None and "geom_id" in mol_props: - try: - conf_id = int(mol_props["geom_id"]) - except Exception: - conf_id = None - - if mol.GetNumConformers() > 0: - conf_props = mol.GetConformer().GetPropsAsDict() - if energy is None: - energy = _get_prop_float(conf_props, "totalenergy") - if weight is None: - weight = _get_prop_float(conf_props, "boltzmannweight") - if conf_id is None and "geom_id" in conf_props: - try: - conf_id = int(conf_props["geom_id"]) - except Exception: - conf_id = None - - try: - energy = float(energy) if energy is not None else None - except Exception: - energy = None - - try: - weight = float(weight) if weight is not None else None - except Exception: - weight = None - - return energy, weight, conf_id - - def filter_conformers_keep_dotted( mol_object: Dict[str, Any], failures: Dict[str, int], @@ -171,6 +101,8 @@ def _read_mol_impl( Dict[str, List[Dict[str, Any]]], ]: mol_path, max_confs, precision, embedding_func, bin_size, ranges, sort_by = args + bin_config = args[7] if len(args) > 7 else None + use_isomeric_smiles = args[8] if len(args) > 8 else False mol_object = load_pkl(mol_path) geom_key = mol_object.get("smiles") geom_id = infer_geom_id(mol_path) @@ -204,7 +136,9 @@ def _read_mol_impl( for _, _, mol, _conf_meta, energy, weight, conf_id in scored_candidates: try: - if embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): embedded_smile, iso_smile = embedding_func( mol, bin_size=bin_size, @@ -217,6 +151,7 @@ def _read_mol_impl( local_failures["encoding_error"] += 1 continue + noniso = None # Compute nonisomeric SMILES only for conformers that encoded successfully try: noniso = Chem.MolToSmiles( @@ -230,7 +165,8 @@ def _read_mol_impl( except Exception: pass - isomeric_smiles.add(iso_smile) + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) + isomeric_smiles.add(canonical_smiles) sample_entry: Dict[str, Any] = { "embedded_smiles": embedded_smile, @@ -240,7 +176,7 @@ def _read_mol_impl( } if conf_id is not None: sample_entry["conf_id"] = conf_id - sample_isomers[iso_smile].append(sample_entry) + sample_isomers[canonical_smiles].append(sample_entry) pickle_entry: Dict[str, Any] = { "mol": copy_single_conformer_mol(mol), @@ -251,7 +187,7 @@ def _read_mol_impl( } if conf_id is not None: pickle_entry["conf_id"] = conf_id - pickle_isomers[iso_smile].append(pickle_entry) + pickle_isomers[canonical_smiles].append(pickle_entry) if len(nonisomeric_smiles) > 1: log.info( @@ -296,14 +232,6 @@ def _read_mol_impl( return json_line, stats, geom_id, isomeric_smiles, pickle_isomers -def save_grouped_pickle(output_path: str, iso_to_confs: Dict[str, List[Dict[str, Any]]]) -> None: - parent = osp.dirname(output_path) - if parent: - os.makedirs(parent, exist_ok=True) - with open(output_path, "wb") as fh: - pickle.dump(iso_to_confs, fh) - - def preprocess( geom_raw_path: str, indices_path: str, @@ -317,19 +245,16 @@ def preprocess( bin_size: float = 0.104, ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", sort_by: str = "energy", + bin_config_path: Optional[str] = None, + isomeric: bool = False, ) -> None: if dest_path is None: raise ValueError("dest_path must be provided for preprocessing output") - embedding_registry = { - "cartesian_v2": encode_cartesian_v2, - "cartesian": encode_cartesian_v2, - "cartesian_binned": encode_cartesian_binned, - "cartesian_binned_v2": encode_cartesian_binned_v2, - } - if embedding_type not in embedding_registry: - raise ValueError(f"Unsupported embedding_type '{embedding_type}'. Options: {sorted(embedding_registry)}") - embedding_func = embedding_registry[embedding_type] + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) overall_total_input_mols = overall_total_confs = overall_total_mols = 0 overall_multi_distinct_graphs = overall_mol_with_dotted_smiles = overall_total_dotted_smiles = 0 @@ -352,12 +277,7 @@ def preprocess( split_indices_array = np.load(indices_path, allow_pickle=True) - try: - parsed_ranges = ast.literal_eval(f"[{ranges}]") - parsed_ranges = [tuple(r) for r in parsed_ranges] - except Exception as exc: - log.error("Failed to parse ranges: {} | failure={}", ranges, exc) - parsed_ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + parsed_ranges = parse_coordinate_ranges(ranges) pickle_glob = osp.join(geom_raw_path, f"{dataset_type}/*.pickle") pickle_paths = np.array(sorted(glob.glob(pickle_glob))) @@ -403,6 +323,8 @@ def preprocess( bin_size, parsed_ranges, sort_by, + bin_config, + isomeric, ) for path in mol_paths ] @@ -553,7 +475,14 @@ def preprocess( "--embedding_type", "-et", type=str, - choices=["cartesian", "cartesian_v2", "cartesian_binned", "cartesian_binned_v2"], + choices=[ + "cartesian", + "cartesian_v2", + "cartesian_binned", + "cartesian_binned_v2", + "uniform_binned", + "quantile_binned", + ], default="cartesian_v2", help="Embedding type to use for enrichment.", ) @@ -621,7 +550,17 @@ def preprocess( default="energy", help="Sort conformers by energy, weight, or keep original order.", ) - + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + parser.add_argument( + "--isomeric", + action="store_true", + help="If set, use isomeric SMILES keys; otherwise use non-isomeric keys when available.", + ) args = parser.parse_args() random.seed(42) @@ -649,4 +588,6 @@ def preprocess( sort_by=args.sort_by, bin_size=args.bin_size, ranges=args.ranges, + bin_config_path=args.bin_config_path, + isomeric=args.isomeric, ) diff --git a/src/molgen3D/data_processing/preprocess_geom_grouped_revisited.py b/src/molgen3D/data_processing/preprocess_geom_grouped_revisited.py new file mode 100644 index 0000000..e91df4c --- /dev/null +++ b/src/molgen3D/data_processing/preprocess_geom_grouped_revisited.py @@ -0,0 +1,521 @@ +import argparse +import json +import os +import os.path as osp +import pickle +import random +from collections import defaultdict +from multiprocessing import Pool +from typing import Any, Dict, List, Optional, Set, Tuple + +from loguru import logger as log +from rdkit import Chem, RDLogger +from tqdm.auto import tqdm + +from molgen3D.data_processing.smiles_encoder_decoder import ( + encode_cartesian_binned, + encode_cartesian_binned_v2, + encode_cartesian_with_config, +) +from molgen3D.data_processing.utils import ( + JsonlSplitWriter, + copy_single_conformer_mol, + extract_conf_meta, + filter_revisited_conformers_keep_dotted, + get_embedding_func_and_config, + get_revisited_split_path, + parse_coordinate_ranges, + save_grouped_pickle, +) + +RDLogger.DisableLog("rdApp.*") + +_REVISITED_DATA: Optional[List] = None + + +def _process_revisited_mol_impl( + idx: int, + max_confs: int, + precision: int, + embedding_func: Any, + bin_size: float, + parsed_ranges: List[Tuple[float, float]], + sort_by: str, + bin_config: Any, + use_isomeric_smiles: bool, +) -> Tuple[ + Optional[str], + Dict[str, Any], + str, + Set[str], + Dict[str, List[Dict[str, Any]]], +]: + smiles, mols = _REVISITED_DATA[idx] + geom_id = str(idx) + local_failures: Dict[str, int] = defaultdict(int) + + candidates = filter_revisited_conformers_keep_dotted( + smiles=smiles, + mols=mols, + failures=local_failures, + ) + + scored_candidates: List[ + Tuple[float, int, Chem.Mol, Optional[float], Optional[float], Optional[int]] + ] = [] + for mol, conf_idx in candidates: + energy, weight, conf_id = extract_conf_meta(None, mol) + if energy is None: + local_failures["missing_energy"] += 1 + if sort_by == "energy": + scored_value = energy if energy is not None else float("inf") + elif sort_by == "weight": + scored_value = -(weight if weight is not None else float("-inf")) + else: + scored_value = float(conf_idx) + scored_candidates.append((scored_value, conf_idx, mol, energy, weight, conf_id)) + + scored_candidates.sort(key=lambda x: (x[0], x[1])) + scored_candidates = scored_candidates[:max_confs] + + nonisomeric_smiles: Set[str] = set() + dotted_smiles: Set[str] = set() + isomeric_smiles: Set[str] = set() + sample_isomers: Dict[str, List[Dict[str, Any]]] = defaultdict(list) + pickle_isomers: Dict[str, List[Dict[str, Any]]] = defaultdict(list) + + for _, _, mol, energy, weight, conf_id in scored_candidates: + try: + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + embedded_smile, iso_smile = embedding_func( + mol, + bin_size=bin_size, + ranges=parsed_ranges, + ) + else: + embedded_smile, iso_smile = embedding_func(mol, precision=precision) + except Exception as exc: + log.error("Error encoding conformer | smiles={} | failure={}", smiles, exc) + local_failures["encoding_error"] += 1 + continue + + noniso = None + try: + noniso = Chem.MolToSmiles( + Chem.RemoveHs(mol, sanitize=False), + canonical=True, + isomericSmiles=False, + ) + nonisomeric_smiles.add(noniso) + if "." in noniso: + dotted_smiles.add(noniso) + except Exception: + pass + + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) + isomeric_smiles.add(canonical_smiles) + + sample_entry: Dict[str, Any] = { + "embedded_smiles": embedded_smile, + "energy": energy, + "weight": weight, + "geom_id": geom_id, + } + if conf_id is not None: + sample_entry["conf_id"] = conf_id + sample_isomers[canonical_smiles].append(sample_entry) + + pickle_entry: Dict[str, Any] = { + "mol": copy_single_conformer_mol(mol), + "embedded_smiles": embedded_smile, + "energy": energy, + "weight": weight, + "geom_id": geom_id, + } + if conf_id is not None: + pickle_entry["conf_id"] = conf_id + pickle_isomers[canonical_smiles].append(pickle_entry) + + if len(nonisomeric_smiles) > 1: + log.info( + "multiple_distinct_nonisomeric_smiles | path={} | distinct_smiles={}", + smiles, + nonisomeric_smiles, + ) + for dotted in dotted_smiles: + log.info("dot_in_conformer_smiles | path={} | smile={}", smiles, dotted) + + if not sample_isomers: + log.warning("No samples after filtering | path={}", smiles) + local_failures["no_samples_after_filtering"] += 1 + + json_line = None + if sample_isomers: + json_line = ( + json.dumps( + {"geom_key": smiles, "geom_id": geom_id, "isomers": sample_isomers}, + separators=(",", ":"), + ) + + "\n" + ) + + stats = { + "path": smiles, + "geom_smiles": smiles, + "confs_count_pre_filter": len(mols), + "confs_count_post_filter": sum(len(v) for v in sample_isomers.values()), + "nonisomeric_smiles_post_filter": len(nonisomeric_smiles), + "isomeric_smiles_post_filter": isomeric_smiles, + "num_distinct_smiles_with_dot": len(dotted_smiles), + "has_dotted_smiles": bool(dotted_smiles), + "failures": local_failures, + "processed_pickle_path": None, + } + return json_line, stats, geom_id, isomeric_smiles, pickle_isomers + + +def _process_revisited_by_idx(args: Tuple) -> Optional[Tuple]: + idx = args[0] + try: + return _process_revisited_mol_impl(idx, *args[1:]) + except Exception as exc: + log.error("Unhandled exception | idx={} | error={}", idx, exc) + return None + + +def preprocess_revisited( + geom_raw_path: str, + embedding_type: str, + num_workers: int = 20, + precision: int = 4, + splits: Optional[str] = None, + dest_path: Optional[str] = None, + max_confs: int = 30, + bin_size: float = 0.104, + ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + sort_by: str = "energy", + bin_config_path: Optional[str] = None, + isomeric: bool = False, + use_centered: bool = True, +) -> None: + global _REVISITED_DATA + + if dest_path is None: + raise ValueError("dest_path must be provided for preprocessing output") + + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) + parsed_ranges = parse_coordinate_ranges(ranges) + requested_splits = [splits] if splits else ["train", "valid", "test"] + + strings_root = osp.join(dest_path, "processed_strings") + split_writers = { + split: JsonlSplitWriter(osp.join(strings_root, split), split, chunk_size=50_000) + for split in ("train", "valid", "test") + } + split_pickle_dirs = { + split: osp.join(dest_path, "processed_pickles", split) for split in ("train", "valid", "test") + } + for split_dir in split_pickle_dirs.values(): + os.makedirs(split_dir, exist_ok=True) + + overall_total_input = overall_total_confs = overall_total_mols = 0 + overall_failure_counts: Dict[str, int] = defaultdict(int) + + for split_name in requested_splits: + data_path = get_revisited_split_path( + geom_raw_path=geom_raw_path, + split_name=split_name, + use_centered=use_centered, + ) + if not osp.exists(data_path): + raise FileNotFoundError(f"Revisited split file not found: {data_path}") + + log.info("Loading {} ...", data_path) + with open(data_path, "rb") as fh: + _REVISITED_DATA = pickle.load(fh) + log.info("Loaded {:,} molecules for split '{}'", len(_REVISITED_DATA), split_name) + + n = len(_REVISITED_DATA) + overall_total_input += n + + job_args = [ + ( + idx, + max_confs, + precision, + embedding_func, + bin_size, + parsed_ranges, + sort_by, + bin_config, + isomeric, + ) + for idx in range(n) + ] + + conf_count_pre = conf_count_post = mol_count_post = 0 + split_num_mol_with_multi_distinct_graphs = split_num_mol_with_dotted_smiles = total_dotted_smiles = 0 + failure_counts: Dict[str, int] = defaultdict(int) + split_geom_to_iso_map: Dict[str, Set[str]] = defaultdict(set) + + geom_to_iso_path = osp.join(dest_path, f"{split_name}_geom_to_isomeric_smiles.jsonl") + iso_to_geom_path = osp.join(dest_path, f"{split_name}_isomeric_to_geom.jsonl") + geom_to_iso_fh = open(geom_to_iso_path, "w") + iso_to_geom_fh = open(iso_to_geom_path, "w") + try: + with tqdm(total=n, dynamic_ncols=True, mininterval=0.2) as pbar: + with Pool(processes=num_workers) as pool: + chunk_size = max(1, n // max(num_workers * 2, 1)) + for result in pool.imap_unordered( + _process_revisited_by_idx, + job_args, + chunksize=chunk_size, + ): + if result is None: + failure_counts["unhandled_exception"] += 1 + pbar.update() + continue + + json_line, stats, geom_id, iso_set, iso_to_pickle_confs = result + + if json_line: + split_writers[split_name].write([json_line]) + + if iso_to_pickle_confs: + processed_pickle_path = osp.join(split_pickle_dirs[split_name], f"{geom_id}.pkl") + try: + save_grouped_pickle(processed_pickle_path, iso_to_pickle_confs) + except Exception as exc: + log.error( + "Failed to write processed pickle | path={} | failure={}", + processed_pickle_path, + exc, + ) + + conf_count_pre += stats["confs_count_pre_filter"] + conf_count_post += stats["confs_count_post_filter"] + overall_total_confs += stats["confs_count_post_filter"] + total_dotted_smiles += stats.get("num_distinct_smiles_with_dot", 0) + + if stats["nonisomeric_smiles_post_filter"] > 1: + split_num_mol_with_multi_distinct_graphs += 1 + if stats.get("has_dotted_smiles", False): + split_num_mol_with_dotted_smiles += 1 + + for reason, count in stats["failures"].items(): + failure_counts[reason] += int(count) + + if stats["confs_count_post_filter"] > 0: + mol_count_post += 1 + overall_total_mols += 1 + + geom_smiles = stats.get("geom_smiles") + if geom_smiles and iso_set: + for iso in iso_set: + split_geom_to_iso_map[geom_smiles].add(iso) + geom_to_iso_fh.write( + json.dumps( + { + "geom_key": geom_smiles, + "geom_id": str(geom_id), + "isomeric_smiles": iso, + }, + separators=(",", ":"), + ) + + "\n" + ) + iso_to_geom_fh.write( + json.dumps( + { + "isomeric_smiles": iso, + "geom_key": geom_smiles, + "geom_id": str(geom_id), + }, + separators=(",", ":"), + ) + + "\n" + ) + + pbar.update() + finally: + geom_to_iso_fh.close() + iso_to_geom_fh.close() + _REVISITED_DATA = None + + avg_confs = conf_count_post / mol_count_post if mol_count_post else 0.0 + success_rate = mol_count_post / n if n else 0.0 + split_report = { + "split": split_name, + "num_input_molecules": n, + "num_output_molecules": mol_count_post, + "num_input_conformers": conf_count_pre, + "total_conformers_after": conf_count_post, + "avg_conformers_per_molecule_after": avg_confs, + "success_rate": success_rate, + "failure_counts": dict(failure_counts), + "molecules_with_multiple_distinct_graphs": split_num_mol_with_multi_distinct_graphs, + "molecules_with_dotted_smiles": split_num_mol_with_dotted_smiles, + "total_dotted_smiles": total_dotted_smiles, + } + log.info(json.dumps({"split_summary": split_report}, ensure_ascii=False, separators=(",", ":"))) + + for reason, count in failure_counts.items(): + overall_failure_counts[reason] += count + + for writer in split_writers.values(): + writer.close() + + grand_total = sum(writer.total_samples for writer in split_writers.values()) + overall_success_rate = float(overall_total_mols) / max(1, overall_total_input) + run_summary = { + "grand_total_samples_written": grand_total, + "total_input_molecules": overall_total_input, + "molecules_after_filter": overall_total_mols, + "conformers_after_filter": overall_total_confs, + "avg_confs_per_mol_after": float(overall_total_confs) / max(1, overall_total_mols), + "overall_success_rate": overall_success_rate, + "overall_failure_counts": dict(overall_failure_counts), + } + log.info(json.dumps({"run_summary": run_summary}, ensure_ascii=False, separators=(",", ":"))) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--geom_raw_path", + "-p", + type=str, + required=True, + help="Path to the revisited GEOM split pickle files.", + ) + parser.add_argument( + "--dest", + "-d", + type=str, + default="/data/molgen/", + help="Destination directory for processed outputs.", + ) + parser.add_argument( + "--embedding_type", + "-et", + type=str, + choices=[ + "cartesian", + "cartesian_v2", + "cartesian_binned", + "cartesian_binned_v2", + "uniform_binned", + "quantile_binned", + ], + default="cartesian_v2", + help="Embedding type to use for enrichment.", + ) + parser.add_argument( + "--num_workers", + "-nw", + type=int, + default=max(4, os.cpu_count() or 4), + help="Number of worker processes.", + ) + parser.add_argument( + "--precision", + type=int, + default=4, + help="Numeric precision for encoded coordinates.", + ) + parser.add_argument( + "--dataset_type", + "-dt", + type=str, + default="drugs", + help="Unused placeholder kept for CLI parity.", + ) + parser.add_argument( + "--splits", + type=str, + choices=["train", "valid", "test"], + default=None, + help="Optional single split to process.", + ) + parser.add_argument( + "--run_name", + type=str, + default="", + help="Run name, appended to destination directory.", + ) + parser.add_argument( + "--max_confs", + type=int, + default=30, + help="Maximum number of conformers per molecule.", + ) + parser.add_argument( + "--bin_size", + type=float, + default=0.104, + help="Bin size for binned embedding.", + ) + parser.add_argument( + "--ranges", + type=str, + default="[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + help="Ranges for binned embedding.", + ) + parser.add_argument( + "--sort_by", + type=str, + choices=["energy", "weight", "none"], + default="none", + help="Sort conformers by energy, weight, or keep original order.", + ) + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + parser.add_argument( + "--isomeric", + action="store_true", + help="If set, use isomeric SMILES keys; otherwise use non-isomeric keys when available.", + ) + parser.add_argument( + "--use_centered", + action="store_true", + default=False, + help="Load *_data_centered.pickle instead of *_data.pickle.", + ) + args = parser.parse_args() + + random.seed(42) + dest_path = osp.join(args.dest, args.run_name) + os.makedirs(dest_path, exist_ok=True) + enqueue_logs = os.environ.get("LOGURU_ENQUEUE", "1") not in {"0", "false", "False"} + log.add( + osp.join(dest_path, "preprocessing.log"), + mode="w", + enqueue=enqueue_logs, + backtrace=False, + diagnose=False, + ) + + preprocess_revisited( + geom_raw_path=args.geom_raw_path, + embedding_type=args.embedding_type, + dest_path=dest_path, + max_confs=args.max_confs, + num_workers=args.num_workers, + precision=args.precision, + splits=args.splits, + bin_size=args.bin_size, + ranges=args.ranges, + sort_by=args.sort_by, + bin_config_path=args.bin_config_path, + isomeric=args.isomeric, + use_centered=args.use_centered, + ) diff --git a/src/molgen3D/data_processing/smiles_encoder_decoder.py b/src/molgen3D/data_processing/smiles_encoder_decoder.py index bc853ae..3f17218 100644 --- a/src/molgen3D/data_processing/smiles_encoder_decoder.py +++ b/src/molgen3D/data_processing/smiles_encoder_decoder.py @@ -1,6 +1,11 @@ import ast +import json import math import re +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + from rdkit import Chem from rdkit.Chem import AllChem from rdkit.Chem.rdchem import ChiralType @@ -9,7 +14,6 @@ def truncate(x, precision=4): - """Format a float with exactly ``precision`` decimal places (truncation, not rounding).""" if precision < 0: raise ValueError("precision must be non-negative") @@ -62,29 +66,13 @@ def _parse_float_token(token: str) -> float: _ORGANIC_SUBSET = {"B", "C", "N", "O", "P", "S", "F", "Cl", "Br", "I", "b", "c", "n", "o", "p", "s"} def strip_smiles(s: str) -> str: - """ - Normalize enriched SMILES strings into a 'canonical-ish' comparison form. - - Supported inputs: - - Legacy enriched strings: C<...>N<...> (atoms without brackets + coords) - - Current enriched strings: [C]<...>[N]<...> (atoms wrapped in brackets) - - Binned enriched strings: [C]123456789[N]123456789 (atoms wrapped in brackets + digits) - - Plain SMILES: C[NH2+]Cc1... - - Steps: - 1. Remove every <...> coordinate block. - 2. Remove binned coordinate digits: [C]123456789 -> [C]. - 3. Collapse decorative carbon H-counts: [CH3],[CH2],[CH],[cH] -> C/c. - 4. Drop brackets around simple atoms: [C]->C, [c]->c, [N]->N, ... - 5. Keep chemically meaningful brackets: [NH2+], [nH], [H], [Pt+2], etc. - """ - if not s: return "" # Remove any tags that might be present s = s.replace("[CONFORMER]", "").replace("[/CONFORMER]", "") s = s.replace("[SMILES]", "").replace("[/SMILES]", "") + s = re.sub(r"\[SERIALIZATION\].*?\[/SERIALIZATION\]", "", s, flags=re.IGNORECASE | re.DOTALL) s = s.replace(";", "") s = _WHITESPACE_RE.sub('', s) @@ -137,7 +125,6 @@ def _expected_plain_token(atom) -> str: def tokenize_smiles(smiles_str, expected_atom_tokens=None): - """Tokenize a canonical SMILES string into atom/non-atom tokens.""" tokens = [] i = 0 n = len(smiles_str) @@ -236,7 +223,6 @@ def tokenize_smiles(smiles_str, expected_atom_tokens=None): def _format_atom_descriptor(atom, *, allow_chirality: bool = True): - """Return a bracketed atom descriptor that preserves valence information.""" symbol = atom.GetSymbol() aromatic = atom.GetIsAromatic() if aromatic and len(symbol) == 1: @@ -276,9 +262,6 @@ def _format_atom_descriptor(atom, *, allow_chirality: bool = True): def _normalize_atom_descriptor(descriptor: str) -> str: - """ - Collapse decorative hydrogen counts on neutral carbon descriptors. - """ match = _CARBON_DESCRIPTOR_RE.match(descriptor) if not match or match.group("iso"): return descriptor @@ -297,7 +280,6 @@ def _normalize_atom_descriptor(descriptor: str) -> str: def encode_cartesian_v2(mol, precision=4): - """Serialize a 3D RDKit Mol into the enriched text representation.""" mol_no_h = Chem.RemoveHs(mol) if mol_no_h.GetNumConformers() == 0: raise ValueError("Molecule has no conformer / 3D coordinates.") @@ -364,7 +346,6 @@ def encode_cartesian_v2(mol, precision=4): ) def tokenize_enriched(enriched): - """Tokenize the enriched representation back into atoms (with coords) and other tokens.""" tokens = [] pos = 0 for match in _ENRICHED_TOKEN_PATTERN.finditer(enriched): @@ -469,7 +450,6 @@ def tokenize_enriched_v2(enriched, digit_width): def decode_cartesian_v2(enriched_string): - """Reconstruct an RDKit Mol (with conformer) from the enriched string produced by the encoder.""" tokens = tokenize_enriched(enriched_string) smiles_parts = [] @@ -505,7 +485,6 @@ def decode_cartesian_v2(enriched_string): def embed_3d_conformer_from_smiles(smiles, seed=0): - """Generate a 3D conformer for a SMILES, drop implicit hydrogens, and return the resulting Mol.""" mol = Chem.MolFromSmiles(smiles) if mol is None: raise ValueError(f"Could not parse SMILES: {smiles}") @@ -543,7 +522,6 @@ def embed_3d_conformer_from_smiles(smiles, seed=0): def coords_rmsd(mol_a, mol_b): - """Compute RMSD between conformer-0 coordinates assuming identical atom order.""" if mol_a.GetNumAtoms() != mol_b.GetNumAtoms(): raise ValueError("Cannot compare coordinates for molecules with different atom counts.") @@ -566,7 +544,6 @@ def coords_rmsd(mol_a, mol_b): return min(math.sqrt(sse / n), rmsd_rdkit) def get_bins_for_coords(ranges, bin_size=0.104): - """Get bins for coordinates based on the ranges and bin size.""" bins = [] for start, end in ranges: bins.append(np.arange(start, end, bin_size)) @@ -574,60 +551,216 @@ def get_bins_for_coords(ranges, bin_size=0.104): def get_digit_width(bins): - """Determine the zero-padding width based on the maximum number of bins.""" max_bin_len = max(len(b) for b in bins) return max(3, len(str(max_bin_len))) def coords_to_bins(coords, bins): - """Convert coordinates to bins.""" return np.digitize(coords, bins) def bins_to_coords(bin_indices, bins, use_bin_center=False): - """ - Convert bin indices to coordinates by choosing a value within the bin interval for each bin. - - Parameters: - bin_indices (array-like): Indices of the bins. - bins (array-like): The bin edges as used in np.digitize (e.g., output from np.arange). - use_bin_center (bool): If True, use the center of the bin; if False, uniformly sample within the bin. - - Returns: - np.ndarray: Coordinates either as bin centers or randomly sampled within each bin. - """ + step = float(bins[-1] - bins[-2]) if len(bins) > 1 else 1.0 coords = [] for bin_idx in bin_indices: - # Find the left and right edges of the bin if bin_idx <= 0: - left = bins[0] - right = bins[1] if len(bins) > 1 else bins[0] + # Tail low (BIN_L): snap to range start + coords.append(float(bins[0])) elif bin_idx >= len(bins): - left = bins[-1] - right = bins[-1] + (bins[-1] - bins[-2]) if len(bins) > 1 else bins[-1] + 1.0 + # Tail high (BIN_H): snap to range end + coords.append(float(bins[-1] + step)) else: left = bins[bin_idx - 1] right = bins[bin_idx] - # Choose value: random within [left, right) or use center - if use_bin_center: - coord = (left + right) / 2.0 - else: - coord = np.random.uniform(left, right) - coords.append(coord) + if use_bin_center: + coords.append((left + right) / 2.0) + else: + coords.append(np.random.uniform(left, right)) return np.array(coords) +@dataclass +class BinConfig: + mode: str + L: float + H: float + n_bins: int + edges: np.ndarray + digit_width: int = 3 + + def __post_init__(self): + self.digit_width = max(3, len(str(self.n_bins + 1))) + + # -- persistence --------------------------------------------------------- + def save(self, path: str) -> None: + obj = { + "mode": self.mode, + "L": self.L, + "H": self.H, + "n_bins": self.n_bins, + "edges": self.edges.tolist(), + } + with open(path, "w") as f: + json.dump(obj, f, indent=2) + + @classmethod + def load(cls, path: str) -> "BinConfig": + with open(path) as f: + obj = json.load(f) + return cls( + mode=obj["mode"], + L=obj["L"], + H=obj["H"], + n_bins=obj["n_bins"], + edges=np.array(obj["edges"], dtype=np.float64), + ) + + +def fit_uniform_bins( + values: np.ndarray, + n_bins: int = 256, + q_low: float = 0.01, + q_high: float = 0.99, +) -> BinConfig: + L = float(np.quantile(values, q_low)) + H = float(np.quantile(values, q_high)) + edges = np.linspace(L, H, n_bins + 1) + return BinConfig(mode="uniform", L=L, H=H, n_bins=n_bins, edges=edges) + + +def fit_quantile_bins( + values: np.ndarray, + n_bins: int = 256, + q_low: float = 0.01, + q_high: float = 0.99, +) -> BinConfig: + L = float(np.quantile(values, q_low)) + H = float(np.quantile(values, q_high)) + clipped = np.clip(values, L, H) + edges = np.quantile(clipped, np.linspace(0, 1, n_bins + 1)) + return BinConfig(mode="quantile", L=L, H=H, n_bins=n_bins, edges=edges) + + +def _encode_scalar(c: float, config: BinConfig) -> int: + if c < config.L: + return 0 + if c > config.H: + return config.n_bins + 1 + i = int(np.searchsorted(config.edges, c, side="right")) - 1 + i = max(0, min(i, config.n_bins - 1)) + return i + 1 # 1-based interior index + + +def _decode_scalar(idx: int, config: BinConfig) -> float: + if idx <= 0: + return config.L + if idx > config.n_bins: + return config.H + return float((config.edges[idx - 1] + config.edges[idx]) / 2.0) + + +def encode_cartesian_with_config(mol, config: BinConfig): + mol_no_h = Chem.RemoveHs(mol) + if mol_no_h.GetNumConformers() == 0: + raise ValueError("Molecule has no conformer / 3D coordinates.") + + smiles = Chem.MolToSmiles( + mol_no_h, + canonical=True, + isomericSmiles=True, + allHsExplicit=False, + allBondsExplicit=False, + ) + + if not mol_no_h.HasProp("_smilesAtomOutputOrder"): + raise ValueError("Mol is missing _smilesAtomOutputOrder after MolToSmiles.") + + atom_order_raw = mol_no_h.GetProp("_smilesAtomOutputOrder") + atom_order = list(map(int, ast.literal_eval(atom_order_raw))) + + expected_atom_tokens = [ + _expected_plain_token(mol_no_h.GetAtomWithIdx(idx)) for idx in atom_order + ] + + tokens = tokenize_smiles(smiles, expected_atom_tokens=expected_atom_tokens) + dw = config.digit_width + + out_parts = [] + atom_idx_in_smiles = 0 + conformer = mol_no_h.GetConformer() + + for token in tokens: + if token["type"] == "atom": + if atom_idx_in_smiles >= len(atom_order): + raise ValueError("SMILES atom tokens exceed atom order mapping.") + + rd_idx = atom_order[atom_idx_in_smiles] + atom_text = token["text"] + atom_descriptor = atom_text if atom_text.startswith("[") else f"[{atom_text}]" + + pos = conformer.GetAtomPosition(rd_idx) + ix = _encode_scalar(pos.x, config) + iy = _encode_scalar(pos.y, config) + iz = _encode_scalar(pos.z, config) + + out_parts.append( + f"{atom_descriptor}{ix:0{dw}d}{iy:0{dw}d}{iz:0{dw}d};" + ) + atom_idx_in_smiles += 1 + else: + out_parts.append(token["text"]) + + if atom_idx_in_smiles != len(atom_order): + raise ValueError( + f"Atom count mismatch: mapped {atom_idx_in_smiles} atoms " + f"but expected {len(atom_order)}." + ) + + return "".join(out_parts), smiles + + +def decode_cartesian_with_config(enriched_string: str, config: BinConfig): + normalized = enriched_string.replace(";", "") + tokens = tokenize_enriched_v2(normalized, config.digit_width) + + smiles_parts = [] + coords = [] + for token in tokens: + if token["type"] == "atom_with_coords": + desc = token["atom_desc"] + desc_inner = desc[1:-1] + if desc_inner in _ORGANIC_SUBSET: + smiles_parts.append(desc_inner) + else: + smiles_parts.append(desc) + + ix, iy, iz = (int(round(v)) for v in token["coords"]) + x = _decode_scalar(ix, config) + y = _decode_scalar(iy, config) + z = _decode_scalar(iz, config) + coords.append((x, y, z)) + else: + smiles_parts.append(token["text"]) + + smiles = "".join(smiles_parts) + mol = Chem.MolFromSmiles(smiles, sanitize=False) + if mol is None: + raise ValueError(f"Failed to parse rebuilt SMILES: {smiles}") + if mol.GetNumAtoms() != len(coords): + raise ValueError( + f"Atom count mismatch: mol has {mol.GetNumAtoms()} atoms, " + f"coords list has {len(coords)} entries." + ) + + Chem.SanitizeMol(mol) + + conformer = Chem.Conformer(mol.GetNumAtoms()) + for idx, (x, y, z) in enumerate(coords): + conformer.SetAtomPosition(idx, Point3D(x, y, z)) + mol.AddConformer(conformer, assignId=True) + return mol + def encode_cartesian_binned(mol, bin_size, ranges=None): - """ - Serialize a 3D RDKit Mol into an enriched text representation where - the Cartesian coordinates are replaced by bin indices. - - Returns: - enriched_string (str): SMILES-like string with [atom] tokens. - smiles (str): Canonical SMILES of the heavy-atom molecule. - bins (list[np.ndarray]): [bins_x, bins_y, bins_z] used for binning. - ranges (list[tuple[float, float]]): Axis ranges used to construct bins. - """ mol_no_h = Chem.RemoveHs(mol) if mol_no_h.GetNumConformers() == 0: raise ValueError("Molecule has no conformer / 3D coordinates.") @@ -653,7 +786,7 @@ def encode_cartesian_binned(mol, bin_size, ranges=None): tokens = tokenize_smiles(smiles, expected_atom_tokens=expected_atom_tokens) if ranges is None: - ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + ranges = [(-11.0, 11.0), (-11.0, 11.0), (-11.0, 11.0)] if len(ranges) != 3: raise ValueError("ranges must be a sequence of three (start, end) tuples.") bins = get_bins_for_coords(ranges, bin_size=bin_size) @@ -705,16 +838,6 @@ def encode_cartesian_binned(mol, bin_size, ranges=None): def encode_cartesian_binned_v2(mol, bin_size, ranges=None): - """ - Serialize a 3D RDKit Mol into an enriched text representation where - the Cartesian coordinates are replaced by bin indices. - - Returns: - enriched_string (str): SMILES-like string with [atom] tokens. - smiles (str): Canonical SMILES of the heavy-atom molecule. - bins (list[np.ndarray]): [bins_x, bins_y, bins_z] used for binning. - ranges (list[tuple[float, float]]): Axis ranges used to construct bins. - """ mol_no_h = Chem.RemoveHs(mol) if mol_no_h.GetNumConformers() == 0: raise ValueError("Molecule has no conformer / 3D coordinates.") @@ -740,7 +863,7 @@ def encode_cartesian_binned_v2(mol, bin_size, ranges=None): tokens = tokenize_smiles(smiles, expected_atom_tokens=expected_atom_tokens) if ranges is None: - ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + ranges = [(-11.0, 11.0), (-11.0, 11.0), (-11.0, 11.0)] if len(ranges) != 3: raise ValueError("ranges must be a sequence of three (start, end) tuples.") bins = get_bins_for_coords(ranges, bin_size=bin_size) @@ -777,7 +900,7 @@ def encode_cartesian_binned_v2(mol, bin_size, ranges=None): iy_txt = f"{iy:0{digit_width}d}" iz_txt = f"{iz:0{digit_width}d}" - out_parts.append(f"{atom_descriptor}{ix_txt}{iy_txt}{iz_txt}") + out_parts.append(f"{atom_descriptor}{ix_txt}{iy_txt}{iz_txt};") atom_idx_in_smiles += 1 else: out_parts.append(token["text"]) @@ -792,20 +915,6 @@ def encode_cartesian_binned_v2(mol, bin_size, ranges=None): def decode_cartesian_binned_v2(enriched_string, bins, use_bin_center=True): - """ - Reconstruct an RDKit Mol (with conformer) from a v2 binned enriched string. - - Supports both schemas: - - With semicolons: [c]123123123;[n]456456456; - - Without semicolons: [c]123123123[n]456456456 - - The string must have been produced by ``encode_cartesian_binned_v2`` using - the same set of ``bins`` (one array per axis). - - Semicolons (';') are treated as optional delimiters and are stripped - before parsing, so both ``[c]123123123;`` and ``[c]123123123`` decode - identically here, while upstream code can still see the raw schema. - """ if len(bins) != 3: raise ValueError("bins must be a sequence of three bin arrays (x, y, z).") @@ -855,3 +964,31 @@ def decode_cartesian_binned_v2(enriched_string, bins, use_bin_center=True): conformer.SetAtomPosition(idx, Point3D(x, y, z)) mol.AddConformer(conformer, assignId=True) return mol + + +def decode_conformer_by_serialization( + enriched_string: str, + serialization_tag: str, + *, + bins=None, + uniform_config_path: Optional[str] = None, + quantile_config_path: Optional[str] = None, +): + mode = str(serialization_tag) + if mode == "cartesian": + return decode_cartesian_v2(enriched_string) + if mode == "cartesian_binned": + if bins is None: + raise ValueError("`bins` must be provided for cartesian_binned decoding.") + return decode_cartesian_binned_v2(enriched_string, bins) + if mode == "uniform": + cfg_path = uniform_config_path or str( + Path(__file__).resolve().parents[1] / "config" / "bin_configs" / "uniform_bins.json" + ) + return decode_cartesian_with_config(enriched_string, BinConfig.load(cfg_path)) + if mode == "quantile": + cfg_path = quantile_config_path or str( + Path(__file__).resolve().parents[1] / "config" / "bin_configs" / "quantile_bins.json" + ) + return decode_cartesian_with_config(enriched_string, BinConfig.load(cfg_path)) + raise ValueError(f"Unsupported serialization mode: {serialization_tag}") diff --git a/src/molgen3D/data_processing/utils.py b/src/molgen3D/data_processing/utils.py index 978649f..bfd41dc 100644 --- a/src/molgen3D/data_processing/utils.py +++ b/src/molgen3D/data_processing/utils.py @@ -1,3 +1,4 @@ +import ast import os import os.path as osp import re @@ -7,15 +8,22 @@ import numpy as np import cloudpickle from collections import defaultdict -from typing import Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Optional, Tuple from pathlib import Path +from loguru import logger as log from rdkit.Chem.rdchem import HybridizationType from rdkit.Chem.rdchem import BondType as BT from rdkit.Chem.rdchem import ChiralType from rdkit import Chem from rdkit.Geometry import Point3D -from molgen3D.data_processing.smiles_encoder_decoder import encode_cartesian_v2 +from molgen3D.data_processing.smiles_encoder_decoder import ( + BinConfig, + encode_cartesian_binned, + encode_cartesian_binned_v2, + encode_cartesian_v2, + encode_cartesian_with_config, +) dihedral_pattern = Chem.MolFromSmarts('[*]~[*]~[*]~[*]') chirality = {ChiralType.CHI_TETRAHEDRAL_CW: -1., @@ -30,6 +38,17 @@ 'Ga': 21, 'Ge': 22, 'As': 23, 'Se': 24, 'Br': 25, 'Ag': 26, 'In': 27, 'Sb': 28, 'I': 29, 'Gd': 30, 'Pt': 31, 'Au': 32, 'Hg': 33, 'Bi': 34} +EMBEDDING_REGISTRY = { + "cartesian_v2": encode_cartesian_v2, + "cartesian": encode_cartesian_v2, + "cartesian_binned": encode_cartesian_binned, + "cartesian_binned_v2": encode_cartesian_binned_v2, + "uniform_binned": encode_cartesian_with_config, + "quantile_binned": encode_cartesian_with_config, +} + +REVISITED_SPLIT_FILE_MAP = {"train": "train", "valid": "val", "test": "test"} + def encode_cartesian_raw(mol, precision=4): """Legacy compatibility wrapper for the enriched representation.""" @@ -103,6 +122,137 @@ def save_processed_pickle( return output_path + +def save_grouped_pickle(output_path: str, iso_to_confs: Dict[str, List[Dict[str, Any]]]) -> None: + parent = osp.dirname(output_path) + if parent: + os.makedirs(parent, exist_ok=True) + with open(output_path, "wb") as fh: + pickle.dump(iso_to_confs, fh) + + +def get_embedding_func_and_config( + embedding_type: str, + bin_config_path: Optional[str] = None, +): + bin_config = None + if embedding_type in ("uniform_binned", "quantile_binned"): + if bin_config_path is None: + raise ValueError( + f"--bin_config_path is required for embedding_type={embedding_type!r}." + ) + bin_config = BinConfig.load(bin_config_path) + log.info( + "Loaded BinConfig from {} | mode={} L={:.4f} H={:.4f} n_bins={}", + bin_config_path, + bin_config.mode, + bin_config.L, + bin_config.H, + bin_config.n_bins, + ) + if embedding_type not in EMBEDDING_REGISTRY: + raise ValueError( + f"Unsupported embedding_type '{embedding_type}'. Options: {sorted(EMBEDDING_REGISTRY)}" + ) + return EMBEDDING_REGISTRY[embedding_type], bin_config + + +def parse_coordinate_ranges(ranges: str) -> List[Tuple[float, float]]: + try: + parsed_ranges = ast.literal_eval(f"[{ranges}]") + return [tuple(r) for r in parsed_ranges] + except Exception as exc: + log.error("Failed to parse ranges: {} | failure={}", ranges, exc) + return [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + + +def get_revisited_split_path( + geom_raw_path: str, + split_name: str, + use_centered: bool = False, +) -> str: + if split_name not in REVISITED_SPLIT_FILE_MAP: + raise ValueError(f"Unsupported revisited split: {split_name!r}") + suffix = "_data_centered.pickle" if use_centered else "_data.pickle" + file_stem = REVISITED_SPLIT_FILE_MAP[split_name] + return osp.join(geom_raw_path, f"{file_stem}{suffix}") + + +def copy_single_conformer_mol(mol: Chem.Mol) -> Chem.Mol: + copied = Chem.Mol(mol) + if copied.GetNumConformers() > 1: + conf = Chem.Conformer(copied.GetConformer(0)) + copied.RemoveAllConformers() + copied.AddConformer(conf, assignId=True) + return copied + + +def _get_prop_float(props: Dict[str, Any], key: str) -> Optional[float]: + if key in props: + try: + return float(props[key]) + except Exception: + return None + return None + + +def extract_conf_meta( + conf_meta: Optional[Dict[str, Any]], + mol: Chem.Mol, +) -> Tuple[Optional[float], Optional[float], Optional[int]]: + energy = weight = None + conf_id = None + + if conf_meta: + if "totalenergy" in conf_meta: + energy = conf_meta.get("totalenergy") + elif "relativeenergy" in conf_meta: + energy = conf_meta.get("relativeenergy") + + if "boltzmannweight" in conf_meta: + weight = conf_meta.get("boltzmannweight") + elif "weight" in conf_meta: + weight = conf_meta.get("weight") + + if "geom_id" in conf_meta: + conf_id = conf_meta.get("geom_id") + + if mol is not None: + mol_props = mol.GetPropsAsDict() + if energy is None: + energy = _get_prop_float(mol_props, "totalenergy") + if weight is None: + weight = _get_prop_float(mol_props, "boltzmannweight") + if conf_id is None and "geom_id" in mol_props: + try: + conf_id = int(mol_props["geom_id"]) + except Exception: + conf_id = None + + if mol.GetNumConformers() > 0: + conf_props = mol.GetConformer().GetPropsAsDict() + if energy is None: + energy = _get_prop_float(conf_props, "totalenergy") + if weight is None: + weight = _get_prop_float(conf_props, "boltzmannweight") + if conf_id is None and "geom_id" in conf_props: + try: + conf_id = int(conf_props["geom_id"]) + except Exception: + conf_id = None + + try: + energy = float(energy) if energy is not None else None + except Exception: + energy = None + + try: + weight = float(weight) if weight is not None else None + except Exception: + weight = None + + return energy, weight, conf_id + class JsonlSplitWriter: def __init__(self, base_dir: str, split_name: str, chunk_size: int = 1_000_000): self.split_name = split_name @@ -170,6 +320,74 @@ def filter_mols( return selected + +def filter_revisited_mols( + smiles: Optional[str], + mols: List[Chem.Mol], + failures: Dict[str, int] = defaultdict(int), + max_confs: int = 30, +) -> List[Chem.Mol]: + if smiles and "." in smiles: + failures["dot_in_smiles"] += 1 + return [] + + mol_from_smiles = Chem.MolFromSmiles(smiles) if smiles else None + if mol_from_smiles is None: + failures["mol_from_smiles_failed"] += 1 + return [] + + num_confs = len(mols) + if max_confs is not None: + num_confs = max_confs + + selected: List[Chem.Mol] = [] + k = 0 + for mol in mols: + if mol is None: + failures["missing_mol"] += 1 + continue + + num_neighbors = [len(a.GetNeighbors()) for a in mol.GetAtoms()] + if not num_neighbors or np.max(num_neighbors) > 4: + failures["large_degree"] += 1 + continue + + selected.append(mol) + k += 1 + if k == num_confs: + break + + return selected + + +def filter_revisited_conformers_keep_dotted( + smiles: Optional[str], + mols: List[Chem.Mol], + failures: Dict[str, int], +) -> List[Tuple[Chem.Mol, int]]: + if smiles and "." in smiles: + failures["dot_in_smiles"] += 1 + + mol_from_smiles = Chem.MolFromSmiles(smiles) if smiles else None + if mol_from_smiles is None: + failures["mol_from_smiles_failed"] += 1 + return [] + + selected: List[Tuple[Chem.Mol, int]] = [] + for idx, mol in enumerate(mols): + if mol is None: + failures["missing_mol"] += 1 + continue + + num_neighbors = [len(a.GetNeighbors()) for a in mol.GetAtoms()] + if not num_neighbors or np.max(num_neighbors) > 4: + failures["large_degree"] += 1 + continue + + selected.append((mol, idx)) + + return selected + def load_pkl(file_path: Path) -> object: if not os.path.exists(file_path): raise FileNotFoundError(f"File {file_path} does not exist.") diff --git a/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py b/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py index 8506719..975d6e2 100644 --- a/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py +++ b/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py @@ -1750,10 +1750,10 @@ def main() -> None: action="store_true", help="Use grouped binned dataset defaults (binned_conformers_* and isomer_units).", ) - parser.add_argument("--seq-len", type=int, default=2048) + parser.add_argument("--seq-len", type=int, default=4096) parser.add_argument("--sample-lines", type=int, default=1000) parser.add_argument("--batch-size", type=int, default=4) - parser.add_argument("--tokenizers", nargs="+", default=["qwen3_0.6b_origin", "qwen3_0.6b_custom", "qwen3_0.6b_binned"]) + parser.add_argument("--tokenizers", nargs="+", default=["qwen3_0.6b_origin", "qwen3_0.6b_custom", "qwen3_0.6b_binned", "qwen3_0.6b_binned_258"]) parser.add_argument("--skip-validation", action="store_true") parser.add_argument("--shuffle", action="store_true", help="Sample random lines via dataloader shuffle") parser.add_argument("--seed", type=int, default=0) @@ -1857,7 +1857,7 @@ def main() -> None: for alias in args.tokenizers: tok_path = str(get_tokenizer_path(alias)) tokenizer = AutoTokenizer.from_pretrained( - tok_path, use_fast=True, fix_mistral_regex=True + tok_path, use_fast=True, ) tokenizer_map[alias] = (tok_path, tokenizer) tokenizer_info_map[alias] = { From a42045baddc35c65291149bbd6bfb4fca1819ce1 Mon Sep 17 00:00:00 2001 From: T4ras123 Date: Mon, 6 Apr 2026 12:31:19 +0400 Subject: [PATCH 2/2] Trim PR to preprocessing modules and shared helpers Keep the PR focused on the preprocessing implementation and shared helper code, and drop auxiliary scripts and token counting changes from the review. Made-with: Cursor --- scripts/analyze_binning_precision.py | 200 -------------- scripts/analyze_coord_centering.py | 213 --------------- scripts/center_and_analyze.py | 249 ------------------ scripts/compute_range_R.py | 96 ------- scripts/count_geom_revisited_train_pickle.sh | 98 ------- scripts/count_revisited_train_tokens.sh | 167 ------------ scripts/fit_bins.py | 132 ---------- scripts/preprocess_geom_revisited.sh | 76 ------ ...reprocess_geom_revisited_paired_grouped.sh | 94 ------- .../dataprocessing/count_tokens.py | 6 +- 10 files changed, 3 insertions(+), 1328 deletions(-) delete mode 100644 scripts/analyze_binning_precision.py delete mode 100644 scripts/analyze_coord_centering.py delete mode 100644 scripts/center_and_analyze.py delete mode 100644 scripts/compute_range_R.py delete mode 100755 scripts/count_geom_revisited_train_pickle.sh delete mode 100755 scripts/count_revisited_train_tokens.sh delete mode 100644 scripts/fit_bins.py delete mode 100755 scripts/preprocess_geom_revisited.sh delete mode 100755 scripts/preprocess_geom_revisited_paired_grouped.sh diff --git a/scripts/analyze_binning_precision.py b/scripts/analyze_binning_precision.py deleted file mode 100644 index b63e759..0000000 --- a/scripts/analyze_binning_precision.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Measure coordinate precision loss from binning on the validation set. - -Compares three encoding methods: - 1. Uniform bins (256 bins, raw/non-centered) - 2. Quantile bins (256 bins, raw/non-centered) - 3. Raw text encoding (truncate to 4 decimal places) - -For each conformer, computes RMSD between original and round-tripped coordinates. - -Usage: - python scripts/analyze_binning_precision.py \ - --data /data/molgen/geom_revisited/val_data.pickle \ - --n_mols 10000 -""" - -import argparse -import math -import os -import pickle -import sys - -import numpy as np - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - -from molgen3D.data_processing.smiles_encoder_decoder import ( - BinConfig, - _encode_scalar, - _decode_scalar, -) - -BIN_CONFIGS_DIR = os.path.join(os.path.dirname(__file__), "bin_configs") - - -def truncate_coord(x, precision=4): - """Replicate the raw text encoding's truncation.""" - factor = 10 ** precision - truncated = math.trunc(x * factor) / factor - if abs(truncated) < 10 ** (-precision): - truncated = 0.0 - return truncated - - -def roundtrip_bins(coords_flat, config): - """Bin and unbin a flat array of coordinate values.""" - decoded = np.empty_like(coords_flat) - for i, c in enumerate(coords_flat): - idx = _encode_scalar(float(c), config) - decoded[i] = _decode_scalar(idx, config) - return decoded - - -def roundtrip_raw(coords_flat, precision=4): - """Truncate to `precision` decimal places (raw text encoding).""" - return np.array([truncate_coord(float(c), precision) for c in coords_flat]) - - -def conformer_rmsd(original, decoded): - """RMSD between two (N, 3) coordinate arrays.""" - diff = original - decoded.reshape(original.shape) - return np.sqrt(np.mean(diff ** 2)) - - -def conformer_max_error(original, decoded): - """Max absolute error across all coordinates.""" - return np.max(np.abs(original - decoded.reshape(original.shape))) - - -def main(): - parser = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--data", type=str, - default="/data/molgen/geom_revisited/val_data.pickle") - parser.add_argument("--n_mols", type=int, default=10000) - parser.add_argument("--precision", type=int, default=4, - help="Decimal precision for raw text encoding") - args = parser.parse_args() - - # Load bin configs (raw = fit on non-centered data) - uniform_cfg = BinConfig.load(os.path.join(BIN_CONFIGS_DIR, "uniform_bins.json")) - quantile_cfg = BinConfig.load(os.path.join(BIN_CONFIGS_DIR, "quantile_bins.json")) - - print(f"Uniform bins: L={uniform_cfg.L:.4f}, H={uniform_cfg.H:.4f}, " - f"n_bins={uniform_cfg.n_bins}, " - f"bin_width={(uniform_cfg.H - uniform_cfg.L) / uniform_cfg.n_bins:.6f} A") - print(f"Quantile bins: L={quantile_cfg.L:.4f}, H={quantile_cfg.H:.4f}, " - f"n_bins={quantile_cfg.n_bins}, " - f"median_bin_width={np.median(np.diff(quantile_cfg.edges)):.6f} A") - print(f"Raw text precision: {args.precision} decimal places") - - # Load validation data - print(f"\nLoading {args.data} ...") - with open(args.data, "rb") as f: - data = pickle.load(f) - print(f" {len(data)} molecules in validation set") - - # Collect conformers from first n_mols molecules - n_mols = min(args.n_mols, len(data)) - print(f" Using first {n_mols} molecules\n") - - methods = ["uniform", "quantile", "raw"] - stats = {m: {"rmsd": [], "max_err": [], "overflow": 0, "total_coords": 0} for m in methods} - - n_conformers = 0 - for mol_idx, (smiles, confs) in enumerate(data[:n_mols]): - for mol in confs: - try: - pos = mol.GetConformer().GetPositions() # (n_atoms, 3) - except Exception: - continue - - flat = pos.flatten() - n_conformers += 1 - - # Uniform bins - decoded_u = roundtrip_bins(flat, uniform_cfg) - stats["uniform"]["rmsd"].append(conformer_rmsd(pos, decoded_u)) - stats["uniform"]["max_err"].append(conformer_max_error(pos, decoded_u)) - stats["uniform"]["overflow"] += int(np.sum((flat < uniform_cfg.L) | (flat > uniform_cfg.H))) - stats["uniform"]["total_coords"] += len(flat) - - # Quantile bins - decoded_q = roundtrip_bins(flat, quantile_cfg) - stats["quantile"]["rmsd"].append(conformer_rmsd(pos, decoded_q)) - stats["quantile"]["max_err"].append(conformer_max_error(pos, decoded_q)) - stats["quantile"]["overflow"] += int(np.sum((flat < quantile_cfg.L) | (flat > quantile_cfg.H))) - stats["quantile"]["total_coords"] += len(flat) - - # Raw text truncation - decoded_r = roundtrip_raw(flat, args.precision) - stats["raw"]["rmsd"].append(conformer_rmsd(pos, decoded_r)) - stats["raw"]["max_err"].append(conformer_max_error(pos, decoded_r)) - stats["raw"]["total_coords"] += len(flat) - - if (mol_idx + 1) % 2000 == 0: - print(f" Processed {mol_idx + 1}/{n_mols} molecules " - f"({n_conformers} conformers so far)") - - print(f"\nTotal: {n_mols} molecules, {n_conformers} conformers\n") - - # Print results - print("=" * 72) - print(f"{'Method':<12} {'Mean RMSD':>10} {'Median RMSD':>12} {'p95 RMSD':>10} " - f"{'p99 RMSD':>10} {'Max RMSD':>10} {'Mean MaxErr':>12}") - print("=" * 72) - - for method in methods: - rmsds = np.array(stats[method]["rmsd"]) - max_errs = np.array(stats[method]["max_err"]) - - print(f"{method:<12} " - f"{np.mean(rmsds):>10.6f} " - f"{np.median(rmsds):>12.6f} " - f"{np.percentile(rmsds, 95):>10.6f} " - f"{np.percentile(rmsds, 99):>10.6f} " - f"{np.max(rmsds):>10.6f} " - f"{np.mean(max_errs):>12.6f}") - - print("=" * 72) - - # Overflow stats (only for binned methods) - print(f"\nOverflow statistics (coords outside [L, H]):") - for method in ["uniform", "quantile"]: - s = stats[method] - pct = 100.0 * s["overflow"] / s["total_coords"] if s["total_coords"] else 0 - print(f" {method:<12}: {s['overflow']:>8} / {s['total_coords']:<10} ({pct:.4f}%)") - - # Per-axis error distribution (sample from last batch) - print(f"\nPer-scalar absolute error distribution:") - print(f"{'Method':<12} {'Mean':>10} {'Median':>10} {'p95':>10} {'p99':>10} {'Max':>10}") - print("-" * 64) - - # Recompute on a sample for per-scalar stats - sample_flat = [] - for mol_idx, (smiles, confs) in enumerate(data[:min(1000, n_mols)]): - for mol in confs: - try: - pos = mol.GetConformer().GetPositions() - sample_flat.append(pos.flatten()) - except Exception: - continue - sample_flat = np.concatenate(sample_flat) - - for method, cfg in [("uniform", uniform_cfg), ("quantile", quantile_cfg)]: - decoded = roundtrip_bins(sample_flat, cfg) - errs = np.abs(sample_flat - decoded) - print(f"{method:<12} {np.mean(errs):>10.6f} {np.median(errs):>10.6f} " - f"{np.percentile(errs, 95):>10.6f} {np.percentile(errs, 99):>10.6f} " - f"{np.max(errs):>10.6f}") - - decoded_r = roundtrip_raw(sample_flat, args.precision) - errs_r = np.abs(sample_flat - decoded_r) - print(f"{'raw':<12} {np.mean(errs_r):>10.6f} {np.median(errs_r):>10.6f} " - f"{np.percentile(errs_r, 95):>10.6f} {np.percentile(errs_r, 99):>10.6f} " - f"{np.max(errs_r):>10.6f}") - - -if __name__ == "__main__": - main() diff --git a/scripts/analyze_coord_centering.py b/scripts/analyze_coord_centering.py deleted file mode 100644 index cf228cd..0000000 --- a/scripts/analyze_coord_centering.py +++ /dev/null @@ -1,213 +0,0 @@ -""" -Analyze GEOM conformer coordinates: -1. Center all conformers (subtract centroid) and save centered pickle files -2. Verify centering (centroid displacement should be ~0) -3. Compute L/H bounds on centered train data (xyz-pooled and x-only) - at quantile levels: 1%, 0.5%, 0.1% -4. Count overflow for all thresholds on centered train and test -""" - -import pickle -import numpy as np -from rdkit.Geometry import Point3D - -DATA_DIR = "/data/molgen/geom_revisited" - - -def get_coords(mol): - """Extract N×3 coordinate array from an RDKit Mol with a conformer.""" - conf = mol.GetConformer() - return conf.GetPositions().astype(np.float64) - - -def load_split(name): - path = f"{DATA_DIR}/{name}_data.pickle" - if 'centered' in name: - name = name.replace("_centered", "") - path = f"{DATA_DIR}/{name}_data_centered.pickle" - print(f"Loading {path} ...", flush=True) - with open(path, "rb") as f: - data = pickle.load(f) - print(f" → {len(data)} molecules", flush=True) - return data - - -# ────────────────────────────────────────────────────────────────────────────── -# 1. Center conformers in-place and save -# ────────────────────────────────────────────────────────────────────────────── - -def center_split(data, split_name): - """ - Center every conformer in-place (subtract centroid from atom positions). - Saves the result to DATA_DIR/{split_name}_data_centered.pickle. - Returns the modified data list. - """ - print(f"\nCentering {split_name} ({len(data)} mols)...", flush=True) - n_confs = 0 - for smiles, confs in data: - for mol in confs: - conf = mol.GetConformer() - pos = conf.GetPositions() - centroid = pos.mean(axis=0) - new_pos = pos - centroid - for i in range(mol.GetNumAtoms()): - conf.SetAtomPosition(i, Point3D(*new_pos[i].tolist())) - n_confs += 1 - print(f" Centered {n_confs} conformers.", flush=True) - - out_path = f"{DATA_DIR}/{split_name}_data_centered.pickle" - print(f" Saving to {out_path} ...", flush=True) - with open(out_path, "wb") as f: - pickle.dump(data, f) - print(f" Saved.", flush=True) - return data - - -# ────────────────────────────────────────────────────────────────────────────── -# 2. Verify centering -# ────────────────────────────────────────────────────────────────────────────── - -def centering_stats(data, label, max_mols=None): - """Compute centroid L2-norm for each conformer (should be ~0 after centering).""" - norms = [] - mols_done = 0 - for smiles, confs in data: - for mol in confs: - X = get_coords(mol) - norms.append(np.linalg.norm(X.mean(axis=0))) - mols_done += 1 - if max_mols and mols_done >= max_mols: - break - norms = np.array(norms) - print(f"\n[{label}] centroid displacement over {len(norms)} conformers" - f" (from {mols_done} mols)") - print(f" mean = {norms.mean():.2e} Å") - print(f" median = {np.median(norms):.2e} Å") - print(f" max = {norms.max():.2e} Å") - print(f" >1e-9Å : {(norms > 1e-9).sum()} ({100*(norms > 1e-9).mean():.4f}%)") - return norms - - -# ────────────────────────────────────────────────────────────────────────────── -# 3. Collect coordinates (data already centered, no re-centering needed) -# ────────────────────────────────────────────────────────────────────────────── - -def collect_coords(data, label, axis=None): - """ - Returns a 1D array of raw scalar coordinate values (no centering). - axis=None → pool all three axes (x, y, z) - axis=0/1/2 → x / y / z only - """ - axis_label = {None: "xyz pooled", 0: "x only", 1: "y only", 2: "z only"} - all_vals = [] - print(f"\nCollecting coords from {label} ({len(data)} mols)" - f" [{axis_label[axis]}]...", flush=True) - for smiles, confs in data: - for mol in confs: - X = get_coords(mol) - vals = X.flatten() if axis is None else X[:, axis] - all_vals.append(vals) - all_vals = np.concatenate(all_vals) - print(f" total scalar values: {len(all_vals):,}", flush=True) - return all_vals - - -# ────────────────────────────────────────────────────────────────────────────── -# 4. Overflow counting (data already centered) -# ────────────────────────────────────────────────────────────────────────────── - -def count_overflow(data, L, H, label): - """Count mols and confs with at least one raw coordinate outside [L, H].""" - n_mols = len(data) - n_confs_total = 0 - n_confs_overflow = 0 - n_mols_overflow = 0 - - for smiles, confs in data: - mol_has_overflow = False - for mol in confs: - flat = get_coords(mol).flatten() # raw, no centering - has_out = bool((flat < L).any() or (flat > H).any()) - n_confs_total += 1 - if has_out: - n_confs_overflow += 1 - mol_has_overflow = True - if mol_has_overflow: - n_mols_overflow += 1 - - print(f"\n[{label}] L={L:.4f}, H={H:.4f}") - print(f" Mols with overflow: {n_mols_overflow:>6} / {n_mols:>6}" - f" ({100*n_mols_overflow/n_mols:.2f}%)") - print(f" Confs with overflow: {n_confs_overflow:>6} / {n_confs_total:>6}" - f" ({100*n_confs_overflow/n_confs_total:.2f}%)") - return n_mols_overflow, n_confs_overflow, n_mols, n_confs_total - - -# ────────────────────────────────────────────────────────────────────────────── -# Main -# ────────────────────────────────────────────────────────────────────────────── - -if __name__ == "__main__": - - # ── Load raw (uncentered) data ── - train_data = load_split("train") - test_data = load_split("test") - - # ── Step 1: Collect raw coords ── - print("\n" + "="*60) - print("STEP 1 — COORDINATE COLLECTION (raw, uncentered)") - print("="*60) - train_vals_xyz = collect_coords(train_data, "train", axis=None) - train_vals_x = collect_coords(train_data, "train", axis=0) - - # Tail quantiles: 1%, 0.1%, 0.01%, 0.001%, 0.0001% - q_pairs = { - "1% / 99%": (0.01, 0.99), - "0.1% / 99.9%": (0.001, 0.999), - "0.01% / 99.99%": (0.0001, 0.9999), - "0.001%/ 99.999%": (0.00001, 0.99999), - "0.0001%/99.9999%": (0.000001, 0.999999), - } - - print("\n" + "="*60) - print("STEP 2 — L / H BOUNDS (raw train): xyz-pooled vs x-only") - print("="*60) - - configs = {} - for q_name, (qlo, qhi) in q_pairs.items(): - for src_name, src_vals in [("xyz", train_vals_xyz), ("x-only", train_vals_x)]: - L = float(np.quantile(src_vals, qlo)) - H = float(np.quantile(src_vals, qhi)) - key = f"{q_name} [{src_name}]" - configs[key] = (L, H) - print(f"\n {key}") - print(f" L = {L:.4f} Å, H = {H:.4f} Å") - - # ── Step 2: Count overflow ── - print("\n" + "="*60) - print("STEP 3 — OVERFLOW COUNTS (raw, uncentered)") - print("="*60) - - results = {} - for thresh_name, (L, H) in configs.items(): - for split_label, split_data in [("TRAIN", train_data), ("TEST", test_data)]: - key = f"{thresh_name} | {split_label}" - nm_ov, nc_ov, nm_tot, nc_tot = count_overflow( - split_data, L, H, f"{thresh_name} — {split_label}") - results[key] = (nm_ov, nc_ov, nm_tot, nc_tot) - - # ── Summary table ── - print("\n" + "="*60) - print("SUMMARY TABLE") - print("="*60) - header = (f"{'Threshold + source':<32} {'Split':<7}" - f" {'Mol overflow':>22} {'Conf overflow':>22}") - print(header) - print("-" * len(header)) - for thresh_name in configs: - for split_label in ["TRAIN", "TEST"]: - key = f"{thresh_name} | {split_label}" - nm_ov, nc_ov, nm_tot, nc_tot = results[key] - print(f"{thresh_name:<32} {split_label:<7}" - f" {nm_ov:>7}/{nm_tot:<7} ({100*nm_ov/nm_tot:5.2f}%)" - f" {nc_ov:>8}/{nc_tot:<8} ({100*nc_ov/nc_tot:5.2f}%)") diff --git a/scripts/center_and_analyze.py b/scripts/center_and_analyze.py deleted file mode 100644 index 5e14bbf..0000000 --- a/scripts/center_and_analyze.py +++ /dev/null @@ -1,249 +0,0 @@ -""" -Center all GEOM conformers, report centering stats, and save coordinate -distributions (as numpy arrays and plots) before and after centering. - -Outputs go to: scripts/centering_analysis/ - - centroid_norms_before.npy centroid displacement per conformer (before) - - centroid_norms_after.npy centroid displacement per conformer (after) - - coords_before_{xyz,x,y,z}.npy coordinate values before centering - - coords_after_{xyz,x,y,z}.npy coordinate values after centering - - distributions.png before/after coordinate histograms - - centroid_norms.png before/after centroid displacement histograms - -Centered data saved to DATA_DIR as: - train_data_centered.pickle / val_data_centered.pickle / test_data_centered.pickle -""" - -import os -import pickle -import numpy as np -import matplotlib -matplotlib.use("Agg") -import matplotlib.pyplot as plt -from rdkit.Geometry import Point3D - -DATA_DIR = "/data/molgen/geom_revisited" -OUT_DIR = os.path.join(os.path.dirname(__file__), "centering_analysis") -os.makedirs(OUT_DIR, exist_ok=True) - -SAMPLE_MOLS = 5000 # mols used for distribution snapshots and centroid stats - - -# ────────────────────────────────────────────────────────────────────────────── -# I/O helpers -# ────────────────────────────────────────────────────────────────────────────── - -def get_coords(mol): - return mol.GetConformer().GetPositions().astype(np.float64) - - -def load_split(split): - path = f"{DATA_DIR}/{split}_data.pickle" - print(f"Loading {path} ...", flush=True) - with open(path, "rb") as f: - data = pickle.load(f) - print(f" → {len(data)} molecules", flush=True) - return data - - -def save_split(data, split): - path = f"{DATA_DIR}/{split}_data_centered.pickle" - print(f" Saving → {path} ...", flush=True) - with open(path, "wb") as f: - pickle.dump(data, f) - print(f" Done.", flush=True) - - -# ────────────────────────────────────────────────────────────────────────────── -# Centering -# ────────────────────────────────────────────────────────────────────────────── - -def center_in_place(data, split): - """Subtract per-conformer centroid from every atom. Mutates data in-place.""" - print(f"\nCentering {split} ({len(data)} mols)...", flush=True) - n_confs = 0 - for smiles, confs in data: - for mol in confs: - conf = mol.GetConformer() - pos = conf.GetPositions() - mu = pos.mean(axis=0) - new = pos - mu - for i in range(mol.GetNumAtoms()): - conf.SetAtomPosition(i, Point3D(*new[i].tolist())) - n_confs += 1 - print(f" {n_confs:,} conformers centered.", flush=True) - - -# ────────────────────────────────────────────────────────────────────────────── -# Stats collection -# ────────────────────────────────────────────────────────────────────────────── - -def collect_centroid_norms(data, max_mols=None): - """Centroid L2-norm for each conformer.""" - norms, done = [], 0 - for smiles, confs in data: - for mol in confs: - X = get_coords(mol) - norms.append(np.linalg.norm(X.mean(axis=0))) - done += 1 - if max_mols and done >= max_mols: - break - return np.array(norms) - - -def print_norm_stats(norms, label): - print(f"\n[{label}] n={len(norms):,} conformers") - for p in [50, 75, 90, 95, 99, 100]: - print(f" p{p:3d} = {np.percentile(norms, p):.6f} Å") - print(f" mean = {norms.mean():.6f} Å") - print(f" >1e-9 : {(norms > 1e-9).sum():,} ({100*(norms > 1e-9).mean():.4f}%)") - print(f" >0.01 : {(norms > 0.01).sum():,} ({100*(norms > 0.01).mean():.2f}%)") - print(f" >0.1 : {(norms > 0.1).sum():,} ({100*(norms > 0.1).mean():.2f}%)") - - -def collect_coord_distributions(data, max_mols=None): - """Pool all coordinate values (optionally limited to max_mols molecules).""" - xyz, x, y, z = [], [], [], [] - done = 0 - for smiles, confs in data: - for mol in confs: - X = get_coords(mol) - xyz.append(X.flatten()) - x.append(X[:, 0]) - y.append(X[:, 1]) - z.append(X[:, 2]) - done += 1 - if max_mols and done >= max_mols: - break - return { - "xyz": np.concatenate(xyz), - "x": np.concatenate(x), - "y": np.concatenate(y), - "z": np.concatenate(z), - } - - -# ────────────────────────────────────────────────────────────────────────────── -# Plotting -# ────────────────────────────────────────────────────────────────────────────── - -def plot_distributions(before, after, out_path): - keys = ["xyz", "x", "y", "z"] - fig, axes = plt.subplots(2, 4, figsize=(20, 8)) - fig.suptitle( - f"Coordinate distributions — before vs after centering " - f"(first {SAMPLE_MOLS:,} train mols)", fontsize=13) - - bins = np.linspace(-30, 30, 300) - for col, key in enumerate(keys): - b = np.clip(before[key], -30, 30) - a = np.clip(after[key], -30, 30) - - for row, (arr, raw, color, tag) in enumerate([ - (b, before[key], "steelblue", "BEFORE"), - (a, after[key], "coral", "AFTER"), - ]): - ax = axes[row, col] - ax.hist(arr, bins=bins, color=color, alpha=0.85) - ax.set_title(f"{tag} — {key}", fontsize=11) - ax.set_xlabel("Å") - ax.set_ylabel("count") - ax.text(0.97, 0.95, - f"μ={raw.mean():.3f}\nσ={raw.std():.3f}", - transform=ax.transAxes, ha="right", va="top", fontsize=9) - - plt.tight_layout() - plt.savefig(out_path, dpi=150) - plt.close() - print(f" Saved → {out_path}") - - -def plot_centroid_norms(before, after, out_path): - fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) - fig.suptitle("Centroid displacement |μ| per conformer", fontsize=13) - - ax1.hist(np.clip(before, 0, 5), bins=200, color="steelblue", alpha=0.85) - ax1.set_title("BEFORE centering") - ax1.set_xlabel("|centroid| (Å)") - ax1.set_ylabel("count") - ax1.text(0.97, 0.95, - f"mean={before.mean():.3f}\nmedian={np.median(before):.3f}\n" - f"p99={np.percentile(before,99):.3f}\nmax={before.max():.3f}", - transform=ax1.transAxes, ha="right", va="top", fontsize=9) - - ax2.hist(after, bins=50, color="coral", alpha=0.85) - ax2.set_title("AFTER centering") - ax2.set_xlabel("|centroid| (Å)") - ax2.set_ylabel("count") - ax2.text(0.97, 0.95, - f"mean={after.mean():.2e}\nmax={after.max():.2e}", - transform=ax2.transAxes, ha="right", va="top", fontsize=9) - - plt.tight_layout() - plt.savefig(out_path, dpi=150) - plt.close() - print(f" Saved → {out_path}") - - -# ────────────────────────────────────────────────────────────────────────────── -# Main -# ────────────────────────────────────────────────────────────────────────────── - -if __name__ == "__main__": - - train_data = load_split("train") - val_data = load_split("val") - test_data = load_split("test") - - # ── BEFORE ── - print("\n" + "="*60) - print(f"BEFORE CENTERING (sample {SAMPLE_MOLS:,} train mols)") - print("="*60) - - norms_before = collect_centroid_norms(train_data, max_mols=SAMPLE_MOLS) - print_norm_stats(norms_before, "centroid |μ| BEFORE") - np.save(os.path.join(OUT_DIR, "centroid_norms_before.npy"), norms_before) - print(f" Saved centroid_norms_before.npy ({len(norms_before):,} values)") - - print(f"\nCollecting coord distributions...", flush=True) - dists_before = collect_coord_distributions(train_data, max_mols=SAMPLE_MOLS) - for key, arr in dists_before.items(): - np.save(os.path.join(OUT_DIR, f"coords_before_{key}.npy"), arr) - print(f" Saved coords_before_{key}.npy " - f"({len(arr):,} vals, mean={arr.mean():.3f}, std={arr.std():.3f})") - - # ── CENTER + SAVE ── - print("\n" + "="*60) - print("CENTERING AND SAVING ALL SPLITS") - print("="*60) - for data, split in [(train_data, "train"), (val_data, "val"), (test_data, "test")]: - center_in_place(data, split) - save_split(data, split) - - # ── AFTER ── - print("\n" + "="*60) - print(f"AFTER CENTERING (same {SAMPLE_MOLS:,} train mols)") - print("="*60) - - norms_after = collect_centroid_norms(train_data, max_mols=SAMPLE_MOLS) - print_norm_stats(norms_after, "centroid |μ| AFTER") - np.save(os.path.join(OUT_DIR, "centroid_norms_after.npy"), norms_after) - print(f" Saved centroid_norms_after.npy ({len(norms_after):,} values)") - - print(f"\nCollecting coord distributions after centering...", flush=True) - dists_after = collect_coord_distributions(train_data, max_mols=SAMPLE_MOLS) - for key, arr in dists_after.items(): - np.save(os.path.join(OUT_DIR, f"coords_after_{key}.npy"), arr) - print(f" Saved coords_after_{key}.npy " - f"({len(arr):,} vals, mean={arr.mean():.2e}, std={arr.std():.3f})") - - # ── PLOTS ── - print("\n" + "="*60) - print("SAVING PLOTS") - print("="*60) - plot_distributions(dists_before, dists_after, - os.path.join(OUT_DIR, "distributions.png")) - plot_centroid_norms(norms_before, norms_after, - os.path.join(OUT_DIR, "centroid_norms.png")) - - print(f"\nAll outputs → {OUT_DIR}") diff --git a/scripts/compute_range_R.py b/scripts/compute_range_R.py deleted file mode 100644 index 01a0075..0000000 --- a/scripts/compute_range_R.py +++ /dev/null @@ -1,96 +0,0 @@ -import argparse -import math -import pickle - -import numpy as np - - -DATA_DIR = "/data/molgen/geom_revisited" - - -def load_split(data_dir, name): - path = f"{data_dir}/{name}_data.pickle" - print(f"Loading {path} ...", flush=True) - with open(path, "rb") as f: - data = pickle.load(f) - print(f" -> {len(data)} molecules", flush=True) - return data - - -def conformer_radius_proxies(data): - radii = [] - for _smiles, confs in data: - for mol in confs: - pos = mol.GetConformer().GetPositions() - radii.append(np.abs(pos).max()) - return np.array(radii, dtype=np.float64) - - -def round_up(value, step=0.5): - return math.ceil(value / step) * step - - -def count_overflow(data, R, label): - n_confs = 0 - n_overflow = 0 - for _smiles, confs in data: - for mol in confs: - pos = mol.GetConformer().GetPositions() - n_confs += 1 - if np.abs(pos).max() > R: - n_overflow += 1 - pct = 100 * n_overflow / n_confs if n_confs else 0 - print(f" [{label}] overflow: {n_overflow:>7} / {n_confs:<7} ({pct:.4f}%)") - return n_overflow, n_confs - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--data_dir", type=str, default=DATA_DIR, - help="Directory with {train,val,test}_data_centered.pickle", - ) - parser.add_argument( - "--quantile", type=float, default=0.9999, - help="Quantile level for choosing R (default: 0.9999 = 99.99%%)", - ) - parser.add_argument( - "--round_step", type=float, default=0.5, - help="Round R up to this multiple (default: 0.5 A)", - ) - args = parser.parse_args() - - # -- Load centered data -- - train_data = load_split(args.data_dir, "train") - - # -- Conformer radius proxies on train -- - print("\nComputing conformer radius proxies (m = max|Xc|) on train...", flush=True) - m_train = conformer_radius_proxies(train_data) - print(f" {len(m_train):,} conformers") - - # -- Percentile table -- - print("\nPercentile table (train):") - for p in [50, 75, 90, 95, 99, 99.5, 99.9, 99.95, 99.99, 99.999, 100]: - val = np.percentile(m_train, p) - print(f" p{p:<8} = {val:.4f} A") - - # -- Pick R -- - R_raw = float(np.quantile(m_train, args.quantile)) - R = round_up(R_raw, args.round_step) - print(f"\nR_raw (quantile {args.quantile}) = {R_raw:.4f} A") - print(f"R (rounded up to {args.round_step}) = {R:.1f} A") - - # -- Overflow counts -- - print(f"\nOverflow counts with R = {R:.1f} (coords outside [-{R:.1f}, {R:.1f}]):") - count_overflow(train_data, R, "train") - - # Load val/test if available - for split in ("val", "test"): - try: - split_data = load_split(args.data_dir, split) - count_overflow(split_data, R, split) - except FileNotFoundError: - print(f" [{split}] skipped (file not found)") - - print(f"\n>>> Recommended range: [-{R:.1f}, {R:.1f}]") - print(f'>>> CLI flag: --ranges "[-{R:.1f}, {R:.1f}], [-{R:.1f}, {R:.1f}], [-{R:.1f}, {R:.1f}]"') diff --git a/scripts/count_geom_revisited_train_pickle.sh b/scripts/count_geom_revisited_train_pickle.sh deleted file mode 100755 index 4b4f1bf..0000000 --- a/scripts/count_geom_revisited_train_pickle.sh +++ /dev/null @@ -1,98 +0,0 @@ -#!/bin/bash -# Usage: -# sbatch --partition=research --job-name=count-geom-revisited-train \ -# --cpus-per-task=1 --mem=32G --time=02:00:00 \ -# scripts/count_geom_revisited_train_pickle.sh -# -# Optional overrides: -# REPO_ROOT=/path/to/3DMolGen -# VENV_PATH=/path/to/.venv/bin/activate -# PICKLE_PATH=/data/vtarasov/geom_revisited/train_data.pickle -# -#SBATCH --output=/home/vtarasov/slurm_logs/dataset_counts/geom-revisited-train-%j.out -#SBATCH --error=/home/vtarasov/slurm_logs/dataset_counts/geom-revisited-train-%j.err - -set -euo pipefail - -REPO_ROOT=${REPO_ROOT:-/home/vtarasov/code/3DMolGen} -VENV_PATH=${VENV_PATH:-$REPO_ROOT/.venv/bin/activate} -PICKLE_PATH=${PICKLE_PATH:-/data/vtarasov/geom_revisited/train_data.pickle} - -mkdir -p /home/vtarasov/slurm_logs/dataset_counts - -echo "Job ID: ${SLURM_JOB_ID:-local}" -echo "Node: $(hostname)" -echo "Started: $(date)" -echo "Repo root: $REPO_ROOT" -echo "Pickle path: $PICKLE_PATH" -echo "" - -source "$VENV_PATH" -cd "$REPO_ROOT" - -python -u - <<'PY' -import pickle -from pathlib import Path - - -pickle_path = Path("/data/vtarasov/geom_revisited/train_data.pickle") -override = Path(__import__("os").environ.get("PICKLE_PATH", str(pickle_path))) -pickle_path = override - -if not pickle_path.exists(): - raise FileNotFoundError(f"Pickle file not found: {pickle_path}") - -with pickle_path.open("rb") as fh: - data = pickle.load(fh) - - -def count_from_entry(entry): - if isinstance(entry, tuple) and len(entry) >= 2: - mols = entry[1] - try: - return 1, len(mols) - except TypeError: - return 1, 0 - if isinstance(entry, dict): - if "mols" in entry: - mols = entry["mols"] - try: - return 1, len(mols) - except TypeError: - return 1, 0 - if "conformers" in entry: - conformers = entry["conformers"] - try: - return 1, len(conformers) - except TypeError: - return 1, 0 - return 1, 0 - - -if isinstance(data, dict): - items = list(data.items()) - molecule_count = len(items) - conformer_count = 0 - for _, value in items: - try: - conformer_count += len(value) - except TypeError: - pass -elif isinstance(data, list): - molecule_count = 0 - conformer_count = 0 - for entry in data: - mols, confs = count_from_entry(entry) - molecule_count += mols - conformer_count += confs -else: - raise TypeError(f"Unsupported pickle top-level type: {type(data).__name__}") - -print("GEOM revisited train counts") -print(f"pickle_path: {pickle_path}") -print(f"molecules: {molecule_count:,}") -print(f"conformers: {conformer_count:,}") -PY - -echo "" -echo "Finished: $(date)" diff --git a/scripts/count_revisited_train_tokens.sh b/scripts/count_revisited_train_tokens.sh deleted file mode 100755 index 74831a4..0000000 --- a/scripts/count_revisited_train_tokens.sh +++ /dev/null @@ -1,167 +0,0 @@ -#!/bin/bash -# Usage: -# sbatch --partition=research --job-name=count-revisited-train-tokens \ -# --cpus-per-task=8 --mem=64G --time=24:00:00 \ -# scripts/count_revisited_train_tokens.sh -# -# Counts tokens on the revisited train splits currently defined in paths.yaml. -# Grouped datasets are counted with serialization_mode=isomer_units and -# non-grouped datasets with serialization_mode=pairs. -# Note: paths.yaml currently contains a duplicate -# `revisited_cartesian_isomeric_grouped_train` entry and does not define -# `revisited_cartesian_isomeric_train`, so this script targets the unique train -# aliases that are available today. -# -# Optional overrides: -# REPO_ROOT=/path/to/3DMolGen -# VENV_PATH=/path/to/.venv/bin/activate -# OUTPUT_ROOT=/path/to/output-dir -# SEQ_LEN=4096 -# SAMPLE_LINES=1000 -# BATCH_SIZE=4 -# UNIT_BATCH_SIZE=64 -# SAMPLE_UNITS=10000 -# SAMPLE_SAMPLES=1000 -# SAMPLE_LINES_FOR_UNITS=1000 -# SHUFFLE=true -# GROUPED_ESTIMATE=exact # exact | fast | sample_only -# -#SBATCH --output=/home/vtarasov/slurm_logs/token_counts/revisited-train-%j.out -#SBATCH --error=/home/vtarasov/slurm_logs/token_counts/revisited-train-%j.err - -set -euo pipefail - -REPO_ROOT=${REPO_ROOT:-/home/vtarasov/code/3DMolGen} -VENV_PATH=${VENV_PATH:-$REPO_ROOT/.venv/bin/activate} -SEQ_LEN=${SEQ_LEN:-4096} -SAMPLE_LINES=${SAMPLE_LINES:-1000} -BATCH_SIZE=${BATCH_SIZE:-4} -UNIT_BATCH_SIZE=${UNIT_BATCH_SIZE:-64} -SAMPLE_UNITS=${SAMPLE_UNITS:-10000} -SAMPLE_SAMPLES=${SAMPLE_SAMPLES:-1000} -SAMPLE_LINES_FOR_UNITS=${SAMPLE_LINES_FOR_UNITS:-1000} -SHUFFLE=${SHUFFLE:-false} -GROUPED_ESTIMATE=${GROUPED_ESTIMATE:-exact} - -TIMESTAMP=$(date +%Y%m%d-%H%M%S) -JOB_TAG=${SLURM_JOB_ID:-local} -OUTPUT_ROOT=${OUTPUT_ROOT:-$REPO_ROOT/outputs/token_counts/revisited_train/${JOB_TAG}-${TIMESTAMP}} - -mkdir -p "$OUTPUT_ROOT" -mkdir -p /home/vtarasov/slurm_logs/token_counts - -echo "Job ID: ${SLURM_JOB_ID:-local}" -echo "Node: $(hostname)" -echo "Started: $(date)" -echo "Repo root: $REPO_ROOT" -echo "Output root: $OUTPUT_ROOT" -echo "Seq len: $SEQ_LEN" -echo "Sample lines: $SAMPLE_LINES" -echo "Batch size: $BATCH_SIZE" -echo "Unit batch size: $UNIT_BATCH_SIZE" -echo "Sample units: $SAMPLE_UNITS" -echo "Sample samples: $SAMPLE_SAMPLES" -echo "Sample unit lines: $SAMPLE_LINES_FOR_UNITS" -echo "Shuffle: $SHUFFLE" -echo "Grouped estimate: $GROUPED_ESTIMATE" -echo "" - -source "$VENV_PATH" -cd "$REPO_ROOT" - -build_grouped_extra_args() { - case "$GROUPED_ESTIMATE" in - exact) - echo "--exact-estimate" - ;; - fast) - echo "--fast-estimate --sample-units $SAMPLE_UNITS" - ;; - sample_only) - echo "--sample-only --sample-samples $SAMPLE_SAMPLES --sample-lines-for-units $SAMPLE_LINES_FOR_UNITS" - ;; - *) - echo "Unsupported GROUPED_ESTIMATE='$GROUPED_ESTIMATE' (expected: exact, fast, sample_only)" >&2 - exit 1 - ;; - esac -} - -run_count() { - local dataset_alias="$1" - local serialization_mode="$2" - local tokenizer_alias="$3" - local label="$4" - local safe_name="${dataset_alias//\//_}" - local log_path="$OUTPUT_ROOT/${safe_name}.log" - local extra_args=() - - if [[ "$serialization_mode" == "isomer_units" ]]; then - read -r -a extra_args <<< "$(build_grouped_extra_args)" - extra_args+=( - --unit-batch-size "$UNIT_BATCH_SIZE" - ) - else - extra_args+=( - --sample-lines "$SAMPLE_LINES" - --batch-size "$BATCH_SIZE" - ) - fi - - if [[ "$SHUFFLE" == "true" ]]; then - extra_args+=(--shuffle) - fi - - echo "============================================================" - echo "Dataset: $dataset_alias" - echo "Label: $label" - echo "Serialization mode: $serialization_mode" - echo "Tokenizer: $tokenizer_alias" - echo "Log path: $log_path" - echo "Started: $(date)" - echo "============================================================" - - python -m molgen3D.training.pretraining.dataprocessing.count_tokens \ - --train-path "$dataset_alias" \ - --skip-validation \ - --seq-len "$SEQ_LEN" \ - --batch-size "$BATCH_SIZE" \ - --serialization-mode "$serialization_mode" \ - --tokenizers "$tokenizer_alias" \ - "${extra_args[@]}" | tee "$log_path" - - echo "" -} - -run_count \ - "revisited_cartesian_isomeric_grouped_train" \ - "isomer_units" \ - "qwen3_0.6b_custom" \ - "grouped cartesian" - -run_count \ - "revisited_quantile_binned_isomeric_grouped_train" \ - "isomer_units" \ - "qwen3_0.6b_binned_258" \ - "grouped quantile_binned" - -run_count \ - "revisited_uniform_binned_isomeric_grouped_train" \ - "isomer_units" \ - "qwen3_0.6b_binned_258" \ - "grouped uniform_binned" - -run_count \ - "revisited_quantile_binned_isomeric_train" \ - "pairs" \ - "qwen3_0.6b_binned_258" \ - "non-grouped quantile_binned" - -run_count \ - "revisited_uniform_binned_isomeric_train" \ - "pairs" \ - "qwen3_0.6b_binned_258" \ - "non-grouped uniform_binned" - -echo "Finished: $(date)" -echo "Logs written to: $OUTPUT_ROOT" diff --git a/scripts/fit_bins.py b/scripts/fit_bins.py deleted file mode 100644 index 39f6fe7..0000000 --- a/scripts/fit_bins.py +++ /dev/null @@ -1,132 +0,0 @@ -import argparse -import os -import pickle - -import numpy as np - -# Append project root so the import works when running as a script -import sys -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - -from molgen3D.data_processing.smiles_encoder_decoder import ( - fit_uniform_bins, - fit_quantile_bins, -) - - -DATA_DIR = "/data/molgen/geom_revisited" - - -def load_split(data_dir, name): - path = f"{data_dir}/{name}_data.pickle" - print(f"Loading {path} ...", flush=True) - with open(path, "rb") as f: - data = pickle.load(f) - print(f" -> {len(data)} molecules", flush=True) - return data - - -def pool_coords(data, x_only=True): - all_vals = [] - for _smiles, confs in data: - for mol in confs: - pos = mol.GetConformer().GetPositions() - if x_only: - all_vals.append(pos[:, 0]) - else: - all_vals.append(pos.flatten()) - return np.concatenate(all_vals) - - -def overflow_stats(data, L, H, label): - n_confs = n_overflow = 0 - for _smiles, confs in data: - for mol in confs: - flat = mol.GetConformer().GetPositions().flatten() - n_confs += 1 - if (flat < L).any() or (flat > H).any(): - n_overflow += 1 - pct = 100 * n_overflow / n_confs if n_confs else 0 - print(f" [{label}] overflow: {n_overflow:>7} / {n_confs:<7} ({pct:.4f}%)") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--data_dir", type=str, default=DATA_DIR) - parser.add_argument("--out_dir", type=str, default="../src/molgen3D/config/bin_configs") - parser.add_argument("--n_bins", type=int, default=256) - parser.add_argument("--q_low", type=float, default=0.0001) - parser.add_argument("--q_high", type=float, default=0.9999) - args = parser.parse_args() - - os.makedirs(args.out_dir, exist_ok=True) - - # -- Load and pool coordinates -- - train_data = load_split(args.data_dir, "train") - - # Pool all xyz for uniform bins and distribution summary - print("\nPooling all scalar coordinates from train (xyz)...", flush=True) - V_all = pool_coords(train_data) - print(f" {len(V_all):,} scalar values") - - # Pool X-only for quantile bins (X is representative of all axes - # since molecular orientations are random) - print("Pooling X-only coordinates from train...", flush=True) - V_x = pool_coords(train_data, x_only=True) - print(f" {len(V_x):,} scalar values") - - # -- Distribution summary -- - print("\nCoordinate distribution (train, pooled xyz):") - for p in [0.01, 0.1, 1, 5, 25, 50, 75, 95, 99, 99.9, 99.99]: - print(f" p{p:<6} = {np.percentile(V_all, p):+.4f} A") - print(f" min = {V_all.min():+.4f} A") - print(f" max = {V_all.max():+.4f} A") - - # -- Fit uniform (from pooled xyz) -- - print(f"\n{'='*60}") - print(f"Fitting UNIFORM bins (B={args.n_bins}, q=[{args.q_low}, {args.q_high}])") - print(f"{'='*60}") - uniform_cfg = fit_uniform_bins(V_all, n_bins=args.n_bins, - q_low=args.q_low, q_high=args.q_high) - print(f" L = {uniform_cfg.L:+.4f} A") - print(f" H = {uniform_cfg.H:+.4f} A") - print(f" w = {(uniform_cfg.H - uniform_cfg.L) / uniform_cfg.n_bins:.6f} A") - print(f" digit_width = {uniform_cfg.digit_width}") - - uniform_path = os.path.join(args.out_dir, "uniform_bins.json") - uniform_cfg.save(uniform_path) - print(f" Saved -> {uniform_path}") - - # -- Fit quantile (from X-only) -- - print(f"\n{'='*60}") - print(f"Fitting QUANTILE bins (B={args.n_bins}, q=[{args.q_low}, {args.q_high}], X-only)") - print(f"{'='*60}") - quantile_cfg = fit_quantile_bins(V_x, n_bins=args.n_bins, - q_low=args.q_low, q_high=args.q_high) - print(f" L = {quantile_cfg.L:+.4f} A") - print(f" H = {quantile_cfg.H:+.4f} A") - print(f" edge range: [{quantile_cfg.edges[0]:+.4f}, {quantile_cfg.edges[-1]:+.4f}]") - print(f" median bin width = {np.median(np.diff(quantile_cfg.edges)):.6f} A") - print(f" digit_width = {quantile_cfg.digit_width}") - - quantile_path = os.path.join(args.out_dir, "quantile_bins.json") - quantile_cfg.save(quantile_path) - print(f" Saved -> {quantile_path}") - - # -- Overflow stats -- - print(f"\nOverflow counts (coords outside [L, H]):") - overflow_stats(train_data, uniform_cfg.L, uniform_cfg.H, "train-uniform") - overflow_stats(train_data, quantile_cfg.L, quantile_cfg.H, "train-quantile") - - for split in ("val", "test"): - try: - split_data = load_split(args.data_dir, split) - overflow_stats(split_data, uniform_cfg.L, uniform_cfg.H, f"{split}-uniform") - overflow_stats(split_data, quantile_cfg.L, quantile_cfg.H, f"{split}-quantile") - except FileNotFoundError: - print(f" [{split}] skipped (file not found)") - - print(f"\nConfigs saved to {args.out_dir}/") - print(f" uniform_bins.json — use with encode_cartesian_with_config") - print(f" quantile_bins.json — use with encode_cartesian_with_config") diff --git a/scripts/preprocess_geom_revisited.sh b/scripts/preprocess_geom_revisited.sh deleted file mode 100755 index 31c5be0..0000000 --- a/scripts/preprocess_geom_revisited.sh +++ /dev/null @@ -1,76 +0,0 @@ -#!/bin/bash -# Usage: -# sbatch --partition=research --job-name=preprocess-revisited \ -# --cpus-per-task=32 --mem=256G --time=48:00:00 \ -# scripts/preprocess_geom_revisited.sh -# -# Memory note: the train split pickle is ~6.5 GB on disk and expands to -# ~50 GB in RAM. multiprocessing.Pool forks the parent, so Python's -# refcount CoW quickly multiplies that by the worker count. Cap workers -# at 32 to stay well within 256 GB; more workers give negligible extra -# throughput for this CPU-light, fork-heavy workload. -#SBATCH --output=/home/vtarasov/slurm_logs/preprocess/revisited-%j.out -#SBATCH --error=/home/vtarasov/slurm_logs/preprocess/revisited-%j.err - -set -euo pipefail - -GEOM_RAW=/data/vtarasov/geom_revisited -DEST=/data/vtarasov -BIN_CONFIGS="/home/vtarasov/code/3DMolGen/src/molgen3D/config/bin_configs" -# Cap at 32: more workers multiply CoW memory without improving throughput. -_RAW_CPUS=${SLURM_CPUS_PER_TASK:-32} -WORKERS=$(( _RAW_CPUS > 32 ? 32 : _RAW_CPUS )) - -echo "Job ID: ${SLURM_JOB_ID:-local}" -echo "Node: $(hostname)" -echo "CPUs: $WORKERS" -echo "Started: $(date)" -echo "" - -source "/home/vtarasov/code/3DMolGen/.venv/bin/activate" - -cd "/home/vtarasov/code/3DMolGen" - -echo "==========================================" -echo " 1/3 cartesian_v2" -echo "==========================================" -python -m molgen3D.data_processing.data_preprocessing_revisited \ - --geom_raw_path "$GEOM_RAW" \ - --dest "$DEST" \ - --run_name geom_revisited_cartesian_isomeric \ - --embedding_type cartesian_v2 \ - --num_workers "$WORKERS" \ - --isomeric - -echo "==========================================" -echo " 2/3 quantile_binned" -echo "==========================================" -python -m molgen3D.data_processing.data_preprocessing_revisited \ - --geom_raw_path "$GEOM_RAW" \ - --dest "$DEST" \ - --run_name geom_revisited_quantile_binned_isomeric \ - --embedding_type quantile_binned \ - --bin_config_path "$BIN_CONFIGS/quantile_bins.json" \ - --num_workers "$WORKERS" \ - --isomeric - -echo "==========================================" -echo " 3/3 uniform_binned" -echo "==========================================" -python -m molgen3D.data_processing.data_preprocessing_revisited \ - --geom_raw_path "$GEOM_RAW" \ - --dest "$DEST" \ - --run_name geom_revisited_uniform_binned_isomeric \ - --embedding_type uniform_binned \ - --bin_config_path "$BIN_CONFIGS/uniform_bins.json" \ - --num_workers "$WORKERS" \ - --isomeric - -echo "==========================================" -echo " All 3 runs complete." -echo " Finished: $(date)" -echo " Output dirs:" -echo " $DEST/geom_revisited_cartesian_isomeric" -echo " $DEST/geom_revisited_quantile_binned_isomeric" -echo " $DEST/geom_revisited_uniform_binned_isomeric" -echo "==========================================" diff --git a/scripts/preprocess_geom_revisited_paired_grouped.sh b/scripts/preprocess_geom_revisited_paired_grouped.sh deleted file mode 100755 index 1c0f715..0000000 --- a/scripts/preprocess_geom_revisited_paired_grouped.sh +++ /dev/null @@ -1,94 +0,0 @@ -#!/bin/bash -# Usage: -# sbatch --partition=research --job-name=preprocess-revisited-grouped \ -# --cpus-per-task=128 --mem=256G --time=12:00:00 \ -# scripts/preprocess_geom_revisited_paired_grouped.sh -# -# This script runs grouped revisited preprocessing across three encodings: -# cartesian, quantile_binned, and uniform_binned. - -set -euo pipefail - -GEOM_RAW=${GEOM_RAW:-/data/vtarasov/geom_revisited} -DEST=${DEST:-/data/vtarasov} -BIN_CONFIGS=${BIN_CONFIGS:-/home/vtarasov/code/3DMolGen/src/molgen3D/config/bin_configs} -REPO_ROOT=${REPO_ROOT:-/home/vtarasov/code/3DMolGen} -VENV_PATH=${VENV_PATH:-$REPO_ROOT/.venv/bin/activate} - -# Cap at 32: more workers multiply CoW memory without improving throughput. -_RAW_CPUS=${SLURM_CPUS_PER_TASK:-32} -WORKERS=$(( _RAW_CPUS > 32 ? 32 : _RAW_CPUS )) - -echo "Job ID: ${SLURM_JOB_ID:-local}" -echo "Node: $(hostname)" -echo "CPUs: $WORKERS" -echo "Started: $(date)" -echo "GEOM_RAW: $GEOM_RAW" -echo "DEST: $DEST" -echo "" - -source "$VENV_PATH" -cd "$REPO_ROOT" - -run_job() { - local step="$1" - local total="$2" - local label="$3" - local module="$4" - local run_name="$5" - local embedding_type="$6" - local bin_config="${7:-}" - - echo "==========================================" - echo " ${step}/${total} ${label}" - echo "==========================================" - - if [[ -n "$bin_config" ]]; then - python -m "$module" \ - --geom_raw_path "$GEOM_RAW" \ - --dest "$DEST" \ - --run_name "$run_name" \ - --embedding_type "$embedding_type" \ - --bin_config_path "$bin_config" \ - --num_workers "$WORKERS" \ - --isomeric - else - python -m "$module" \ - --geom_raw_path "$GEOM_RAW" \ - --dest "$DEST" \ - --run_name "$run_name" \ - --embedding_type "$embedding_type" \ - --num_workers "$WORKERS" \ - --isomeric - fi -} - - -run_job 1 3 \ - "grouped cartesian" \ - "molgen3D.data_processing.preprocess_geom_grouped_revisited" \ - "geom_revisited_cartesian_isomeric_grouped" \ - "cartesian" - -run_job 2 3 \ - "grouped quantile_binned" \ - "molgen3D.data_processing.preprocess_geom_grouped_revisited" \ - "geom_revisited_quantile_binned_isomeric_grouped" \ - "quantile_binned" \ - "$BIN_CONFIGS/quantile_bins.json" - -run_job 3 3 \ - "grouped uniform_binned" \ - "molgen3D.data_processing.preprocess_geom_grouped_revisited" \ - "geom_revisited_uniform_binned_isomeric_grouped" \ - "uniform_binned" \ - "$BIN_CONFIGS/uniform_bins.json" - -echo "==========================================" -echo " All 3 grouped runs complete." -echo " Finished: $(date)" -echo " Output dirs:" -echo " $DEST/geom_revisited_cartesian_isomeric_grouped" -echo " $DEST/geom_revisited_quantile_binned_isomeric_grouped" -echo " $DEST/geom_revisited_uniform_binned_isomeric_grouped" -echo "==========================================" diff --git a/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py b/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py index 975d6e2..8506719 100644 --- a/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py +++ b/src/molgen3D/training/pretraining/dataprocessing/count_tokens.py @@ -1750,10 +1750,10 @@ def main() -> None: action="store_true", help="Use grouped binned dataset defaults (binned_conformers_* and isomer_units).", ) - parser.add_argument("--seq-len", type=int, default=4096) + parser.add_argument("--seq-len", type=int, default=2048) parser.add_argument("--sample-lines", type=int, default=1000) parser.add_argument("--batch-size", type=int, default=4) - parser.add_argument("--tokenizers", nargs="+", default=["qwen3_0.6b_origin", "qwen3_0.6b_custom", "qwen3_0.6b_binned", "qwen3_0.6b_binned_258"]) + parser.add_argument("--tokenizers", nargs="+", default=["qwen3_0.6b_origin", "qwen3_0.6b_custom", "qwen3_0.6b_binned"]) parser.add_argument("--skip-validation", action="store_true") parser.add_argument("--shuffle", action="store_true", help="Sample random lines via dataloader shuffle") parser.add_argument("--seed", type=int, default=0) @@ -1857,7 +1857,7 @@ def main() -> None: for alias in args.tokenizers: tok_path = str(get_tokenizer_path(alias)) tokenizer = AutoTokenizer.from_pretrained( - tok_path, use_fast=True, + tok_path, use_fast=True, fix_mistral_regex=True ) tokenizer_map[alias] = (tok_path, tokenizer) tokenizer_info_map[alias] = {