Skip to content

Reproducibility release: HF datasets/embeddings, fixed evaluators, reproduce scripts - #1

Merged
ComDec merged 23 commits into
mainfrom
repro/v1
Sep 18, 2026
Merged

Reproducibility release: HF datasets/embeddings, fixed evaluators, reproduce scripts#1
ComDec merged 23 commits into
mainfrom
repro/v1

Conversation

@ComDec

@ComDec ComDec commented Sep 17, 2026

Copy link
Copy Markdown
Owner

This PR collects the reproducibility work for the ICLR 2026 paper on one branch: the released
Hugging Face data, evaluators whose defaults reproduce the published numbers, per-table
reproduction scripts with reference values, and documentation of what is and is not reproducible.

It merges three topic branches (fix/rotation-infra, fix/chirality, fix/energy) plus
integration fixes. Draft: the Table 1 full recompute and the errata wording are still open.

Released data (already on the Hub, nothing to merge here)

Repository What changed
EscheWang/3dcs combined repo with the configs chirality, rotation, traj_frames, traj_energies, so load_dataset("EscheWang/3dcs", name=..., split="train") works
EscheWang/3dcs-embeddings new, public: the 139 baseline embedding files used for the paper (21.2 GB) with a manifest.csv (path, size, SHA-256, keys, shape, provenance), plus the original per-molecule metric outputs and the rMD17 split files
traj_energies (both in EscheWang/3dcs and EscheWang/3dcs-traj-energies) the float32-quantized energies were replaced by the original float64 rMD17 energies (bitwise equal to the source .npz); the previous revision is named in each card's changelog
EscheWang/3dcs-rotation ungated; all five dataset cards now carry cc-by-sa-4.0 and the citation

Code

Energy (Tables 3, 6, 7). Dropped the float32 cast of the rMD17 energies; added a quantization
check (--energy-precision-check, default error), the published window protocol
(--window-scheme legacy, the default), automatic Tanimoto for fingerprint pickles, --molecules
and --n-jobs. The CLI now reproduces the legacy runner bit for bit (301,000 per-window values and
7 × 43 summary cells).

Chirality (Table 2). Added --distance {euclidean,cosine}. The default stays Euclidean,
which is what produced the published Table 2; the paper text says cosine, and both variants are
published in expected.csv. Added --unsup-kmax (default n-1, the published setting),
--n-jobs, strict input validation (row counts, offsets, NaN/Inf, ambiguous keys) and support for
the E3FP pickle via --embedding-key e3fp.

Geometry (Table 1). The rotation offset column is per shard, so embeddings are now sliced
shard by shard (--layout by-shard, --offset-mode auto). The previous flat slicing picked the
wrong conformers for every shard but shard 0. The paper's metric definitions were recovered from
the backed-up per-molecule outputs (Kendall only for molecules with ≥ 11 conformers, the CKA
bandwidth, the isotonic-R² direction, LIE@k with k=10 including the conformer itself, AS as a
median per radian).

Dual-track metrics. Every evaluator takes --metric-version {paper,v2}. paper (the default)
reproduces the published numbers; v2 implements the corrected definitions, each documented with
its rationale in docs/METRICS.md and docs/metrics/{geometry,chirality,energy}.md.

Infrastructure. New python -m three_dbench download {dataset,embeddings} with SHA-256
verification against the manifest; rdkit declared correctly (rdkit-pypi removed),
scikit-learn-extra made optional, joblib declared; no import-time side effects; CITATION.cff;
the official rMD17 split files under splits/rmd17/; rewritten README.md, docs/USAGE.md,
docs/EMBEDDINGS.md, docs/METRICS.md.

Reproduction scripts. reproduce/{table1_geometry,table2_chirality,energy_tables_3_6_7}/ each
have a run.sh (download → evaluate → results.csv) and an expected.csv with the printed value,
the recomputed value at 6 decimals, a tolerance and notes; reproduce/compare.py prints PASS/FAIL.

Findings recorded in the documentation

  • Table 2 distance. Published with Euclidean distance, not the cosine distance of the paper text.
  • Table 2, GemNet SCI. Recomputes as −0.015151; the paper prints +0.015 (sign).
  • Table 1, embedding-cursor drift. In the published full run, two molecules could not be
    processed and the embedding cursor was not advanced, so the molecules after them in shards 1 and 2
    (91,093 of 1,464,493, 6.2 %) were scored with embeddings shifted by 3 and 7 rows. This affects the
    LIE@k and AS columns. --replicate-offset-drift reproduces it; the regular columns use aligned
    embeddings. Estimated effect on LIE@k: about 0.012–0.024 too high; the model ranking is unchanged.
  • Energy tables. Some printed cells are truncated rather than rounded; the MACE and FMG columns
    of Table 6 are only partially reproducible from the backed-up embeddings. Every case is listed in
    the notes column of expected.csv.
  • Extraction recipes. Known for E3FP, partially for GemNet, unknown for UniMol, MolAE,
    MolSpectra, MACE and FMG (docs/EMBEDDINGS.md).

Verification

All checks were run on a 128-core Linux box with the merge commit of this branch.

Tests and lint

Environment Install Result
Python 3.11 (conda, numpy 2.4.6, rdkit 2026.03.6) PYTHONPATH=src pytest 131 passed
Python 3.12 (fresh venv) pip install . (non-editable) pytest 131 passed
Python 3.9 (fresh venv, rdkit 2025.09.2) pip install ".[dev]" pytest 131 passed
all – ruff check and ruff format --check on src/ tests/ examples/ reproduce/ clean

python examples/demo.py all runs end to end against the Hub without a token.

Table 2 (chirality), end to end from the Hub, anonymously

bash reproduce/table2_chirality/run.sh downloaded the chirality config and all 7 embedding
files (SHA-256 verified against models.csv), ran 28 evaluations (7 models x
{euclidean, cosine, v2_euclidean, v2_cosine}, ~55 s each with 22 workers) and
reproduce/compare.py reported 136 PASS, 0 FAIL (largest deviation 5e-7).

Tables 3, 6, 7 (energy)

  • Full 7-model run (both metric versions), compared against expected.csv: 413 PASS, 0 FAIL.
  • Re-run end to end from the Hub anonymously for GemNet and E3FP (all 10 molecules, 1,000 windows
    each, --metric-version paper, 331 s and 302 s with 22 workers): 78 PASS, 0 FAIL. This
    exercises run.sh, the float64 energies from the re-uploaded traj_energies, and the
    traj/<model>/ embeddings from the new repository.

Data access

python -m three_dbench download {dataset,embeddings} was run against the live repositories without
a token, including the SHA-256 verification against manifest.csv.

Not re-run here: the full 16-shard Table 1 GemNet run, and the 5 energy models other than GemNet
and E3FP in the end-to-end script (their values come from the full run above).

Still open

  • Embedding-extraction scripts are not in the release; the per-model status is documented.
  • Fine-tuning (Tables 5, 8, 9): no code or checkpoints here. Checkpoints for the chirality
    fine-tuning are being prepared for a later release; none are planned for the rMD17 tables.
  • Table 1: only GemNet rotation embeddings survive, and the full 16-shard recompute has not been
    run; expected.csv carries the reference values from the backed-up per-molecule outputs, and the
    paper_aligned / v2 cells stay empty until that run.
  • Table 4 (chirality correlation) evaluation code is not in the release; the backed-up summary is
    published.
  • Errata wording is still being discussed and is deliberately not part of this PR.

🤖 Generated with Claude Code

ComDec and others added 23 commits September 18, 2026 01:57
@ComDec @claude
traj: keep rMD17 energies float64 and detect quantized energies
1b37aeb 
Remove the float32 cast in convert_traj_energy_npz_to_hf; the published
energy tables were computed with float64 rMD17 energies. Add
detect_quantized_energies/check_energy_precision to flag float32-cast
absolute energies (few distinct values or float32-representable values
with a coarse step).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
traj: reproduce the published energy protocol from the CLI
cccdbe3 
- --window-scheme {legacy,shared}, default legacy: per-molecule
  default_rng(seed) reseeding with starts in [0, 100000 - window), as in
  run_trajectory_benchmark; bitwise-identical to the legacy runner.
- --metric-version {paper,v2}: paper = published implementation; v2 =
  appendix definitions (RMS EJS sigma, separate CKA bandwidths, isoR2
  dE ~ f(Delta), scale-free KS, TS/Smoothness only for time-ordered
  frames, float32 distances).
- --n-jobs (order-preserving process pool, one BLAS thread per worker),
  --molecules, --energy-precision-check, --time-ordered.
- E3FP pickles (lists of RDKit bit vectors) load from directories and
  use Tanimoto distance; dense bit arrays can use --metric-embed tanimoto.
- Vectorise the pair-index helpers in traj/evaluation.py (values unchanged).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
tests: trajectory window schemes, float64 energies, Tanimoto, metric …
46dc441 
…versions

Includes a bitwise equivalence test between the CLI protocol and the
legacy runner (serial and parallel).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
docs: energy metric definitions for paper and v2
9f7584f 
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
chirality: --distance, --unsup-kmax, --metric-version, --n-jobs and i…
f88e337 
…nput validation

- distance switch for continuous embeddings (euclidean default = published
  Table 2; cosine with an exact-zero diagonal); fingerprints keep Tanimoto
- best-k silhouette scans k up to n-1 by default (the published run); the
  hard-coded min(10, n-1) of release 0.1.0 is available as --unsup-kmax 10
- metric_version paper (default, published code path) and v2 (tie-aware
  NN@1 over points with a same-class partner, centroid DBI, Hopkins on the
  supervised population, best-k silhouette on the selected distance)
- per-molecule evaluation shared by both entry points; optional process
  pool (results independent of n_jobs)
- clear errors for embedding rows != sum(n_conformers), non-finite values,
  inconsistent dataset offsets, duplicate keys and ambiguous pickle keys
- no import-time dependency on utils.paths or sklearn_extra; legacy driver
  paths resolved lazily; config.json with options, coverage and versions

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
tests: chirality distance switch, unsup kmax, metric versions, valida…
3559cfc 
…tion

Fast synthetic tests for --distance (euclidean/cosine, fingerprints use
Tanimoto), the unbounded best-k silhouette default, paper vs v2 definitions
(NN@1 ties and structural zeros, centroid DBI vs sklearn, Hopkins
population, cosine geometry), n_jobs invariance, input validation, the CLI
options and the absence of import-time side effects.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
reproduce: energy Tables 3/6/7 runner, collector and paper values
db26e82 
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
chirality: run molecule batches with joblib/loky instead of a fork pool
3af40ab 
loky avoids forking a process that may already hold an OpenMP runtime and
works the same on Linux and macOS. Each batch carries only its molecules'
rows; workers use one BLAS/OpenMP thread. Per-molecule outputs are
bit-identical to the serial path (checked on 300 molecules, UniMol and E3FP).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
reproduce/table2_chirality: run.sh, models.csv, make_results.py
7482a68 
run.sh downloads the chirality config of EscheWang/3dcs and the seven
Table 2 embedding files of EscheWang/3dcs-embeddings (SHA-256 checked
against models.csv), evaluates all models for the variants euclidean
(published protocol), cosine, v2_euclidean and v2_cosine, writes
results.csv and calls reproduce/compare.py when present.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
rotation: shard-aware embedding slicing, versioned geometry metrics
ad5eab9 
- Slice per-shard embeddings with row['shard'] and the per-shard offset of the
  published HF rotation config (offsets restart per shard; shards appear in
  string order). Flat arrays are aligned by the cumulative conformer count in
  dataset row order; converter datasets with global offsets are detected.
- Add --metric-version presets (rotation/metrics.py): 'paper' reproduces the
  published Table 1 definitions recovered from the backup per-molecule outputs
  (Kendall only for >=11 conformers, CKA bandwidth of the original run, isotonic
  fit D=f(Delta) without guard, LIE k=10 incl. self on (1-cos)/2, AS median of
  (1-cos)/2 per rad incl. wrap-around); 'v2' follows the paper text (LIE k=3
  excluding self, AS median ||dz|| per rad, median-heuristic CKA).
- Selection options: shards, molecule list, per-shard sampler, min conformers.
- Molecules whose RMSD cannot be computed are reported instead of aborting.
- Optional replicate_offset_drift: recompute metrics with the embedding-cursor
  shift of the published full run (shards 1 and 2) in __offset_drift columns.
- Parallel evaluation (fork), parquet per-key output, summary and config.
- Embedding IO: skip string arrays when auto-selecting NPZ keys, accept a
  directory with one file, registry of the published embeddings layout.
- Lazy lmdb imports; no directory creation at import time.
- examples/run_rotation_from_hf.py: by-shard evaluation, flat .npy cache via
  numpy.lib.format.open_memmap with shard-specific cache names.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
infra: rdkit dependency, optional scikit-learn-extra, CWD-based paths…
0400a52 
…, CI lint paths

- Declare rdkit>=2023.9 and huggingface_hub; drop 'pip install rdkit-pypi'
  from Makefile, CONTRIBUTING and CI (no cp312 wheel; breaks with numpy 2).
- scikit-learn-extra becomes the optional 'kmedoids' extra (0.3.0 fails to
  import with numpy 2; the default chirality path does not use it).
- utils/paths: resolve data/results roots from $THREE_DBENCH_HOME or the
  current working directory instead of the package location.
- CI: fail-fast off, lint examples/ and reproduce/, codecov fail_ci_if_error.
- Authors and project URLs in pyproject; add CITATION.cff.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
cli: download subcommand and rotation evaluate options
2f2a4d0 
- python -m three_dbench download dataset --task {chirality,traj,rotation,all}
  saves the EscheWang/3dcs configs to data/hf/... with save_to_disk.
- python -m three_dbench download embeddings --task ... [--models] [--results]
  fetches files listed in EscheWang/3dcs-embeddings manifest.csv and verifies
  their SHA-256 (cached files are not downloaded again).
- evaluate rotation: --layout by-shard, --offset-mode, --shards, --molecule-list,
  --sample-ratio/--sample-seed, --min-conformers, --max-keys, --metric-version,
  --lie-k, --lie-self, --as-variant, --extra-metrics, --replicate-offset-drift,
  --n-jobs.
- Tests: mocked downloads, CLI help, import side effects, path resolution.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
reproduce: compare.py, Table 1 geometry scripts, rMD17 splits
7fefcd5 
- reproduce/compare.py: PASS/FAIL/MISSING/NO_EXPECTED per expected row,
  non-zero exit on failures (contract shared by all reproduce/ directories).
- reproduce/table1_geometry: run.sh (download, GemNet paper and v2 runs,
  collect, compare; QUICK=1 for shard 1), collect.py, expected.csv (paper
  values and means of the original per-molecule outputs for five models),
  expected_quick.csv, and the 146,389 molecule keys of the original 10 % run.
- splits/rmd17: the official rMD17 train/test index files (splits 01-05) with
  checksums; Tables 8/9 inputs correspond to split 01.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
docs: quick start, reproducing the paper, metrics and embeddings
b122c24 
- README: installation without rdkit-pypi, download helper, evaluation of
  published embeddings, per-table reproduction scripts and status table,
  links to datasets, embeddings and metric docs, BibTeX.
- docs/METRICS.md index; docs/metrics/geometry.md (paper/v2/legacy
  definitions, provenance of Table 1, per-molecule agreement, embedding offset
  shift of the original full run, runtime).
- docs/USAGE.md and docs/EMBEDDINGS.md for the new CLI options and the
  EscheWang/3dcs-embeddings layout (with integrator TODOs for the chirality
  and trajectory options).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
traj: task-agnostic --metric-version help, run.sh EXTRA_ARGS, quiet e…
e8260f4 
…xpected warning

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
chirality: expected values, metric documentation, v2 Hopkins populati…
05c74c1 
…on fix

- reproduce/table2_chirality/expected.csv: 136 rows (7 models x 4 variants
  x 5 metrics, E3FP Hopkins omitted) computed on the box with run.sh on the
  published embeddings; euclidean rows carry the printed paper values and
  notes (rounding rule, GemNet SCI sign, comparison with the Sep-2025
  en_sep_results)
- docs/metrics/chirality.md: distance switch, paper vs v2 definitions with
  code pointers, measured rationale, populations, expected values
- make_results.py: pandas writes NaN as an empty field
- v2: Hopkins is also NaN in the opt-in unsupervised_only branch

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
reproduce: expected values for energy Tables 3/6/7 (paper and v2)
9053551 
Recomputed on the A100 box with the new CLI (legacy windows, float64 HF
energies, 7 models x 1000 windows). The paper variant is bitwise equal to
the legacy runner output used for the published tables.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
Merge branch 'fix/rotation-infra' into repro/v1
db66fbe 
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
Merge branch 'fix/chirality' into repro/v1
f6d91cf 
Keep a single --metric-version / --n-jobs on 'evaluate' (both branches added
them); chirality contributes --distance and --unsup-kmax.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
Merge branch 'fix/energy' into repro/v1
60238d6 
Keep one --metric-version / --n-jobs on 'evaluate' (all three branches added
them) and make their help task-agnostic.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
integrate: resolve USAGE TODOs, declare joblib, make examples importa…
6dc3fdf 
…ble in tests

- docs/USAGE.md: chirality and trajectory option tables checked against the
  merged CLI; add --do-unsup-when-single-en, --time-ordered, --legacy-traj-len
  and --block-size.
- pyproject: declare joblib (imported directly by the chirality --n-jobs path;
  previously only a transitive scikit-learn dependency).
- tests/conftest.py: put the checkout root on sys.path so tests/test_demo.py
  can import examples.demo when pytest is started outside the repository
  (non-editable installs).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
integrate: make the dense Tanimoto test independent of the RDKit version
f6e1d72 
RDKit's Tanimoto similarity of two all-zero fingerprints changed from 1.0
(<= 2025.09) to 0.0 (>= 2026.03), which made the dense-vs-RDKit comparison
fail on Python 3.9 with the older wheel. The comparison now skips pairs of
two all-zero fingerprints and the dense convention is asserted separately.
The published E3FP fingerprints contain no all-zero vector (checked for all
999,988 trajectory frames and all 52,391 chirality conformers), so the
reproduced tables do not depend on this; noted in docs/metrics/energy.md.

Also: README status for Table 1 (full 16-shard GemNet run still pending, the
published LIE@k / AS include the embedding-cursor shift) and for the
fine-tuning tables.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ComDec @claude
integrate: ignore reproduce/*/out/ (default output dir of the energy …
510fa53 
…runner)

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@codacy-production

codacy-production Bot commented Sep 17, 2026

Copy link
Copy Markdown

Not up to standards ⛔

🔴 Issues 57 high · 28 medium · 15 minor

Alerts:
⚠ 100 issues (≤ 0 issues of at least minor severity)

Results:
100 new issues

Category Results
BestPractice 2 medium
Security 9 medium
3 minor
57 high
CodeStyle 12 minor
Complexity 17 medium

View in Codacy

🟢 Metrics 796 complexity · -4 duplication

Metric Results
Complexity 796
Duplication -4

View in Codacy

NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer
TIP This summary will be updated as you push new changes.

@ComDec
ComDec marked this pull request as ready for review September 18, 2026 03:14
@ComDec
ComDec merged commit 00d3c32 into main Sep 18, 2026
5 of 6 checks passed
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 18, 2026 •
edited
Loading

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-09-18T03:15:43.885391Z 510fa53 Draft marked ready
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@ComDec ComDec mentioned this pull request Sep 18, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@ComDec