Baseline embedding-extraction scripts and documentation pass - #2
Conversation
Rewrite the user-facing text of the release so that it states, for every option and every published file, what the code computes and what the artifact contains: - README: replace the per-model status matrix with a coverage table of the published artifacts; point at docs/EMBEDDINGS.md for the per-file details. - docs/METRICS.md, docs/metrics/*: describe `paper` and `v2` as the two sets of definitions the evaluators implement, keep the measured effect of each choice, and document the embedding indexing that --replicate-offset-drift reproduces. - docs/EMBEDDINGS.md: list the format, key, dimension and E3FP parameters of each published embedding file. - reproduce/: expected.csv keeps paper_value as printed and expected_value as the reference value of this release; the notes column now describes the metric version, the variant, the molecule population and the provenance of each reference value. - CLI help and docstrings: `v2` is the alternative definition set. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
One script per baseline model (E3FP, GemNet-Q, Uni-Mol, Mol-AE, MolSpectra, MACE, FMG) that turns the chirality conformers into the embedding matrix the evaluator reads, in dataset row order and under the documented array key. Each script takes the Hugging Face config or a pickle of RDKit molecules, prints the versions and checksums of everything it used, and has a --verify flag that compares the file it wrote with the published embedding of that model in EscheWang/3dcs-embeddings (checksums, elementwise differences, per-row cosine). baselines/common.py holds the input loading, the npz writing and the comparison. No third-party code or weights are redistributed: every model directory carries an ENVIRONMENT.md with the upstream repository and commit, the weight file with its SHA-256 and where to download it, and the exact install commands. baselines/README.md tabulates the output dimension, array key, hydrogen handling, pooling and upstream dependency of each model, and the measured agreement between a full 52,391-conformer run of each script and the published file. tests/test_baselines.py covers --help of every script and the helpers in common.py offline; baselines/ is added to the ruff paths of the Makefile and of CI. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds a "Baseline embeddings" section to the README and a pointer from the embedding documentation to the per-model extraction scripts, and records that recomputing the E3FP fingerprints with the documented parameters reproduces all 52,391 published fingerprints bit for bit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The MACE batch-size figure is from a 128-conformer sample, and the MolSpectra entry is the largest difference over the five metrics of expected.csv. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Not up to standards ⛔🔴 Issues
|
| Category | Results |
|---|---|
| Documentation | 48 minor |
| Security | 10 minor 1 high 3 critical 8 medium |
| CodeStyle | 8 minor |
| Complexity | 22 medium |
🟢 Metrics 1052 complexity · 45 duplication
Metric Results Complexity 1052 Duplication 45
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.
…row subset `baselines/common.py` resolves the input specification for all seven scripts (`parse_dataset_spec`, `load_conformers`, `iter_conformers`): `hf:<repo>[:<config>]`, `hfdisk:<dir>` or a plain `save_to_disk` directory, a bare Hub dataset id, a pickle of RDKit molecules (a list or a dict of lists), and `lmdb:<file>` for a rotation shard. Each script calls it through its own `load_molecules`, so the same `--dataset` value works everywhere; the Mol-AE script reads the `mol_blocks` list of every dataset row, i.e. one row per conformer in ascending `offset`. `--limit` and `--start` run a slice of the conformers in every script. `--verify` compares the rows a run produced. `parse_row_selection` resolves `--verify-rows` into indices into the reference file: `prefix` (the default for a shorter output), `full`, a `<start>:<stop>` slice, `<start>+`, or `@<file>` of 0-based indices; `--start N` implies `N+`. The report names the selection, the number of rows compared and the reference rows used, and the per-row cosine, allclose and difference statistics are computed over them. For fingerprints the same selection applies to the Tanimoto comparison. `tests/test_baselines.py` covers the specification table, reading a `save_to_disk` copy in `offset` order with `--limit`/`--start`, every script accepting every input form and resolving the Hub form, the row selections and their errors, subset comparison of vectors and fingerprints, and that each command in `baselines/README.md` uses flags its script has. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ile belongs to `baselines/fmg/requirements.txt` carries `tqdm>=4.66.3`, the floor `datasets==4.0.0` requires, so `pip install -r` resolves in one pass; verified in a fresh conda env (python 3.10, torch 2.0.0+cu118 first, then the file), which gives torch 2.0.0+cu118, numpy 1.26.4, rdkit 2024.9.6, einops 0.7.0, datasets 4.0.0 and tqdm 4.70.1 — the versions `fmg/ENVIRONMENT.md` lists. Its header names `baselines/fmg/extract_chirality.py`, the path of the script it builds the environment for. The `datasets` / `huggingface_hub` lines of e3fp, mace, molae and molspectra name the `--dataset` forms and `--verify` they are needed for, in the wording the other files use. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`fmg/ENVIRONMENT.md` gives the public Google Drive folder the upstream FMG README links for the pre-trained weights, the file id of `model-120qm9_3rd_run.pt`, the `gdown` command for the folder or the single file, and the sha256 to check the download against. `molae/ENVIRONMENT.md` says that the Mol-AE Drive download is a ZIP archive named `checkpoint_7_1000000.pt` (611,355,678 bytes) holding the 760,699,523-byte checkpoint, gives the unzip command and both checksums, and notes that the file-id form of `gdown` works on every version. Every `ENVIRONMENT.md` now opens its "Input precision" section with the input the published file and the numbers in `baselines/README.md` were computed from — the source RDKit molecules, at full float precision — and with what the four-decimal `mol_blocks` of the public dataset change. Mol-AE and GemNet-Q have that measurement: over the first 2,000 conformers, one machine and one environment, Mol-AE moves by `max|diff| = 4.4e-01`, per-row cosine mean 0.9999948 and minimum 0.9921 with 2.8 % of rows below 1 - 1e-6; GemNet-Q by `max|diff| = 5.1e-04`, per-row cosine at least 0.999999996 with no row below 1 - 1e-6. `baselines/README.md` splits the agreement table into the run from the full-precision molecules over all 52,391 conformers and the run from the published `mol_blocks` over the first 2,000, lists the per-model effect of the rounding with a link to each `ENVIRONMENT.md`, documents the shared `--dataset` table, `--limit`/`--start` and `--verify-rows`, and carries one runnable command per model. The root `README.md` and `docs/EMBEDDINGS.md` state the same input syntax and subset verification. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
# Conflicts: # README.md
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6aabd7e999
ℹ️ 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".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not os.path.exists(lmdb_path): | ||
| mols = load_molecules(args.dataset, args.hf_split, args.hf_config, limit=args.limit, start=args.start) | ||
| n = build_lmdb(mols, lmdb_path) | ||
| print(f"[info] wrote {n} records to {lmdb_path}") | ||
| else: | ||
| print(f"[info] reusing existing LMDB {lmdb_path}") |
There was a problem hiding this comment.
When --work-dir points to a directory used by an earlier run, the script silently reuses molae_input.lmdb without checking the current dataset, --start, or --limit. For example, a full run after a --limit 2000 smoke test will emit only the old 2,000 embeddings while otherwise succeeding, and default verification merely treats them as a valid prefix. Rebuild the LMDB for each invocation or persist and validate its source parameters before reuse.
Useful? React with 👍 / 👎.
| rec = _record_from_mol(mol, align, allowed, bond_order_map) | ||
| if rec is None: | ||
| skipped += 1 | ||
| continue |
There was a problem hiding this comment.
When a supported pickle or LMDB input contains a conformer with no atom in --atom-channels (or lacks a conformer), this branch drops it and shifts every subsequent embedding row. The evaluator assumes one output per input conformer, so the resulting metrics are assigned to the wrong molecules; this also violates the shared baseline contract that unfeaturizable rows cause an error rather than being skipped. Abort with the offending input index instead of continuing.
Useful? React with 👍 / 👎.
| # only for the Hugging Face --dataset forms (hf:<repo>[:<config>], hfdisk:<dir>, a | ||
| # save_to_disk directory or a bare Hub id) | ||
| # datasets | ||
| # only for --verify | ||
| # huggingface_hub |
There was a problem hiding this comment.
After following the Mol-AE installation instructions, the documented --dataset hf:EscheWang/3dcs:chirality ... --verify command fails before extraction because both datasets and huggingface_hub are commented out here, and the explicit install command in ENVIRONMENT.md also omits them. In particular, load_conformers immediately imports datasets; the environment guide even incorrectly says this dependency is included in this requirements file.
Useful? React with 👍 / 👎.
1 participant
This PR adds
baselines/, the embedding-extraction script of each baseline model, and adocumentation pass over the toolkit's text. It builds on the reproducibility release already on
main(#1): the Hugging Face datasets and baseline embeddings, evaluators whose defaults reproducethe published numbers, and per-table reproduction scripts with reference values.
Baseline extraction scripts
baselines/<model>/extract_chirality.pyturns the chirality conformers into the embedding matrixthe evaluator reads, in dataset row order and under the documented array key.
e3fp(pickle)e3fp==1.2.7gemnetRemoveAllHs, heavy atoms onlyTUM-DAML/gemnet_pytorch@a0164f7+ itspretrained/GemNet-Qarr_0[CLS]tokendeepmodeling/Uni-Mol@90f52c4,dptech-corp/Uni-Core@ace6fae,mol_pre_no_h_220816.ptarr_0[CLS]tokencheckpoint_7_1000000.ptarr_0Z != 1)AzureLeon1/MolSpectra@8846530+ a checkpoint supplied by the callerarr_0mace-torch==0.3.15, MACE-MP-0mediumembeddingsDumitrescu-Alexandru/FMG@13a0a7c+model-120qm9_3rd_run.ptCommon to all of them, resolved by
baselines/common.py:--datasettakes the same value in every script:hf:EscheWang/3dcs:chirality(Hub repositoryand config),
hf:<repo>or a bare Hub id (configchirality),hfdisk:<dir>or a plainsave_to_diskdirectory of that config, a pickle of RDKit molecules (a list or a dict oflists), or
lmdb:<file>for a rotation shard. Hugging Face rows are read in ascendingoffsetand each row's
offsetis checked against the running conformer count, somol_blocksconcatenated that way is the row order of the published files.
--limit N/--start Nrun a slice of the conformers, in every script.SHA-256 of the weights it loaded and the SHA-256 of the file it wrote, and refuses to skip a
conformer it cannot featurise rather than shifting the row order.
--verify [REFERENCE]compares the file just written with the published embedding of that modelin
EscheWang/3dcs-embeddings(orhub:<path>, or a local file) and prints both checksums, theelementwise differences and the per-row cosine similarity.
--verify-rowssays which rows of the reference the output covers, so a partial run isverifiable:
prefix(the default when the output is shorter),full, a<start>:<stop>slice,<start>+, or@<file>of 0-based row indices;--start NimpliesN+. The report names theselection, how many rows it compared and which rows of the reference it used.
baselines/README.mdcarries one runnable command per model, with the flags that model needs.No third-party code or weights are redistributed. Every model directory has an
ENVIRONMENT.mdwith the upstream repository and commit, the weight file with its SHA-256 and where to download it,
the exact install commands and the measured cost of a full run, plus a pinned
requirements.txt.Every weight these scripts need is publicly downloadable without an account: the FMG checkpoint
from the Drive folder the upstream README links (
gdowncommand and file id infmg/ENVIRONMENT.md), the Mol-AE checkpoint from the Drive link in the Mol-AE README, whicharrives as a ZIP archive that has to be unpacked before
torch.load(command and both checksums inmolae/ENVIRONMENT.md), and the rest from GitHub releases or checkouts.Agreement with the published embeddings
The published embedding files were computed from the source RDKit molecules, whose coordinates
carry full float precision; the public
EscheWang/3dcsdataset stores the same geometries as V2000MOL blocks, which hold four decimals. The table below is the comparison from the full-precision
molecules;
baselines/README.mdcarries the same comparison from the publishedmol_blocksoverthe first 2,000 conformers, which is what a run from the released dataset gives, and the per-model
effect of the rounding (2.8 % of rows for Mol-AE, 1.2 % for MolSpectra, 0.4 % for FMG, 0.1 % for
Uni-Mol,
max|diff| = 5.1e-04for GemNet-Q, 3.8e-05 for MACE, 2,846 of 3,000 sampled fingerprintsfor E3FP).
Each script was run over all 52,391 conformers and compared with the published file of that model.
"Metrics" is the largest absolute difference over the six chirality metrics (ES-AUC, NN@1-Acc, SCI,
SCI_unsup, Hopkins, DBI) computed by
evaluate chirality --distance euclidean --metric-version paperfrom the published file and from the regenerated file.The neural files are not byte-identical and are not expected to be: the forward passes run in
float32 on the GPU, where the reduction order depends on the batch size, the library build and the
device. Two runs of the same script on the same machine measure the size of that effect (FMG: up to
6e-4 per element; MACE, batch size 1 against 16: 1.4e-7 on a 128-conformer sample; GemNet, TF32 on
against off: 2.0e-3). The MolSpectra entry is the largest difference over the five metrics of
reproduce/table2_chirality/expected.csv; DBI was not recorded for that run.MolSpectra publishes no checkpoint, so one has to be supplied by the caller; the script computes the
documented quantity (the per-atom scalar representation of the MolSpectra equivariant Transformer
summed over the atoms of the molecule), and the row above used the public
denoised-pcqm4mv2.ckptthat MolSpectra initialises from. The published
chirality/molspectra/sampled_mol_feature.npzis the artifact the paper's MolSpectra values werecomputed from.
baselines/README.mdstates this per model.Documentation
README.md: a "Baseline embeddings" section linkingbaselines/and the published embeddings.docs/EMBEDDINGS.md: a pointer from each published file to the script that computes it, and theE3FP fingerprint parameters.
docs/METRICS.md,docs/metrics/{geometry,chirality,energy}.md,reproduce/*/README.mdand thenotescolumn of everyexpected.csv: a wording pass that states what each evaluator computes inpaperand inv2, the molecule populations behind every mean, and the provenance of everyreference value. No numeric value, tolerance or
paper_valuewas changed.Released data (already on the Hub, nothing to merge here)
EscheWang/3dcschirality,rotation,traj_frames,traj_energies, loadable withload_dataset("EscheWang/3dcs", name=..., split="train")EscheWang/3dcs-embeddingsmanifest.csv(path, size, SHA-256, keys, shapes, provenance), the per-molecule metric outputs of the original runs, and the rMD17 split filesVerification
Tests and lint.
pytest217 passed on Python 3.11 withPYTHONPATH=src;ruff checkandruff format --checkonsrc/ tests/ examples/ reproduce/ baselines/clean.tests/test_baselines.py(86 tests) runs offline and without any model: the--datasetspecification table, reading a
save_to_diskcopy inoffsetorder with--limit/--start, eachof the seven scripts accepting each input form and resolving the Hub form, the
--verify-rowsselections and their error messages, subset comparison of vectors and of fingerprints, and that
every command in
baselines/README.mduses flags its script has.E3FP, end to end.
baselines/e3fp/extract_chirality.pyover the 52,391 chirality conformers(214 s,
--jobs 24, e3fp 1.2.7, rdkit 2026.03.6) reproduces the publishedchirality/e3fp/sampled_chi.pklbit for bit (identical 52391,tanimoto_min 1), and evaluatingthe regenerated file gives ES-AUC 0.485935, NN@1-Acc 0.177939, SCI −0.012543, SCI_unsup 0.033825 —
equal at six decimals to
reproduce/table2_chirality/expected.csv.All seven from the released dataset, on one A100-80GB. Each script run from the Hugging Face
input in the environment its
ENVIRONMENT.mdbuilds, with--limit/--startand--verifyagainst the published file: E3FP (
hfdisk:, 300 rows, 292 bit-identical, Tanimoto mean 0.9942),MACE (
hf:, 200 rows, cosine min 0.99999999), Uni-Mol (hfdisk:, 500 rows, cosine min 0.99994),Mol-AE (
hfdisk:andhf:, 2,000 and 64 rows, cosine mean 0.9999948), GemNet-Q (hfdisk:, 2,000rows, cosine min 0.99999966), MolSpectra (
hfdisk:, 256 rows, cosine mean −0.17 with thesubstitute checkpoint, as documented), FMG (
hfdisk:with--start 100 --limit 100, compared withreference rows 100..199, cosine mean 0.9999998).
baselines/fmg/requirements.txtwas installedinto a fresh conda env (python 3.10, torch 2.0.0+cu118 first) and FMG run in it against the
published file (64 rows, cosine mean 0.9999999996).
Mol-AE and GemNet-Q input precision. Measured by running each of them twice over the first
2,000 conformers, once from the published
mol_blocksand once from the full-precision molecules,on the same machine in the same environment; the numbers are in their
ENVIRONMENT.md.Not re-run here: the full 52,391-conformer neural extractions, whose numbers in the table above
come from the runs that produced them.
Scope
not part of it; the rMD17 fine-tuning inputs correspond to the official split 01, included under
splits/rmd17/.per-shard rotation inputs (see
baselines/gemnet/ENVIRONMENT.md); the trajectory embeddings arepublished and are not regenerated here.
🤖 Generated with Claude Code