Extending cuik-molmaker to reactions (cuik-reactmaker)#4
Conversation
Implements C++ CGR (Condensed Graph of Reaction) featurization to accelerate chemprop reaction property prediction (~9x over Python path). New C++ code: - src/reaction_features.cpp: batch_reaction_featurizer supporting all 6 RxnModes (REAC_DIFF, REAC_PROD, PROD_DIFF and BALANCE variants) and all 4 atom featurizer modes (V1, V2, ORGANIC, RIGR). Uses O(bonds) hash-map bond enumeration instead of O(n^2) atom-pair scan. - src/features.h: ReactionMode enum, reaction_mode_names_to_array, and parse_reaction declarations; get_atomic_num_onehot_index helper for num_only representation of unmatched atoms - src/one_hot.cpp/h: get_atomic_num_onehot_index for num_only encoding - src/cuik_molmaker_cpp.cpp: exports reaction_mode_names_to_array and batch_reaction_featurizer to Python - CMakeLists.txt: add reaction_features.cpp to cuik_molmaker_core sources Test fixtures: - tests/data/sample_rxns_100.csv: 100 balanced reactions (50 E2 + 50 SN2) plus 10 hand-crafted unbalanced reactions covering num_only and BALANCE mode edge cases. Verified against chemprop CondensedGraphOfReactionFeaturizer across all 6 modes with max_diff=0.
test_reaction_features.py: parametrized over all 4 atom featurizer versions (V1, V2, ORGANIC, RIGR) × all 6 reaction modes (REAC_DIFF, REAC_PROD, PROD_DIFF and BALANCE variants) = 24 test cases. RIGR uses reduced bond features (["is-null", "in-ring"], bond_fdim=2 per side = 4 total) unlike V1/V2/ORGANIC which use 5 bond features (bond_fdim=14 per side = 28 total). Golden .xz files generated from C++ batch_reaction_featurizer output after verifying agreement with chemprop CondensedGraphOfReactionFeaturizer (max_diff=0 on E2/SN2 data across all 6 modes).
|
Hi @sveccham , just a friendly ping on the PR whenever you have a chance. I'd appreciate any feedback/comments before it is ready to merge. |
|
Hi @akshatzalte, Thank you for your patience. I will have a look at this PR in the next few days. We are currently working on releasing |
sveccham
left a comment
sveccham
left a comment
There was a problem hiding this comment.
Hi @akshatzalte,
Thanks again for the PR. Very neat idea using a hashmap to reduce the problem complexity. This is also a very good use case for moving to C++.
I had a high level look at the PR and left a couple of initial comments. I had one question about the overall design. The CGR bond lookup is constructed once for each reaction in a batch. Would it be more computationally efficient to construct it for the entire batch (with appropriate atom and bond offsets) at once and exploit the O(1) lookup?
|
Hi @sveccham , For the high-level design question, doing a batch level map construction won't be super advantageous. My current understanding is that merging into a single batch-wide table wouldn't change the number of lookups: since bonds only exist within a single reaction, the queries are all intra-reaction and their count is set by the pair scan rather than by how the table is partitioned. So a batch table would be queried the same number of times. It would be just larger and colder in cache, with some added offset bookkeeping. If that's right, the bigger lever for the bond step might be the iteration strategy rather than the table's scope: replacing the O(atoms^2) pair scan with direct bond iteration (O(atoms + bonds)), still reproducing chemprop's exact edge ordering. I'd be very happy to prototype that and benchmark it if you think it's worth it. This would make the code look significantly different from the current python-only handling in chemprop and the real advantage will be visible for super large molecules only. Open to your thoughts on this. |
Add a terse binding docstring and a CGR batch example in docs/USAGE.md covering all 5 returned arrays, mirroring the existing batch_mol_featurizer documentation. Addresses review feedback on the reaction return value.
The reaction regression fixtures in tests/data/sample_rxns_100.csv are sampled from the RDB7 dataset (Spiekermann, K.; Pattanaik, L.; Green, W. H. Scientific Data 2022, 9, 417), distributed under CC BY 4.0. Record the required attribution in LICENSE/third-party.txt alongside the other third-party notices.
Replace the reaction regression fixtures with 100 atom-mapped reactions sampled from RDB7 (CC BY 4.0; see LICENSE/third-party.txt) and regenerate all 24 golden references (V1/V2/ORGANIC/RIGR x 6 RxnModes). Reactions are real published chemistry rather than the previous mixed/hand-crafted set.
Run the repo's pre-commit hooks over the branch's source files. Only the reaction test module required reformatting (black list-wrapping + isort); the C++ sources were already clang-format clean.
There was a problem hiding this comment.
How were these reference files generated? It would be good to document the process in this codebase. Perhaps as a README file with some associated code. Open to your suggestions here.
There was a problem hiding this comment.
Added tests/data/gen_reaction_refs.py, a single self-contained script that downloads the RDB7 source from Zenodo, samples the 100 reactions deterministically, and regenerates all 24 golden files. Added tests/data/README.md documenting the pipeline, provenance, and the cross-verification against Chemprop's CGR featurizer. I previously also included a license under third-party.txt to ensure we comply with the data license.
Happy to adjust the location or format if you prefer.
There was a problem hiding this comment.
So golden values should not be generated with the same codebase that we are testing. It would be great if you could document the process of generating golden values using Chemprop's CGR featurizer. I will try to reproduce the golden values following those instructions.
You can delete the tests/data/gen_reaction_refs.py
There was a problem hiding this comment.
Got it! Sorry I misunderstood your previous comment. I originally had the generation using Chemprop's CGR code but did not want to add that script as it adds new dependency to be able to run everything in this repo. I will document the recipe in a README and provide code snippets from my script so that the process is fully reproducible and unambiguous. Thank you for clarifying this!
There was a problem hiding this comment.
Does it take a significant amount of time for chemprop to generate these? If it doesn't pull too many dependencies in, would seem to make more sense to have it be a test dependency. I know since cuik-molmaker is a dep of chemprop, that would be circular if we put it in the source code, but for a dev test setup it doesn't seem to bad, compared to binary reference files, even if they're regeneratable.
There was a problem hiding this comment.
Thank you @scal444 for adding to this discussion.
It doesn't take that long for chemprop to generate these (just a few seconds). However, I don't think it's worth adding due to the weight of the dependencies. chemprop's base install pulls torch, lightning, pandas, scikit-learn and the CUDA runtime, which comes to 76 packages and roughly 5.5 GB, most of it torch plus CUDA. That is a lot to require in the test environment of an otherwise lean featurization package (today the reaction tests need only pytest, numpy and rdkit) to replace about 370 KB of frozen references.
So for now I am leaning toward keeping the committed goldens plus the documentation of fully reproducible recipe in the README. Happy to discuss further!
- bond-type-float: map bond order to chemically meaningful floats (aromatic -> 1.5, not the raw RDKit enum 12.0), matching the molecule featurizer in float_features.cpp. - reaction mode: replace reaction_mode_names_to_array (array) with a scalar reaction_mode_to_int(str) that raises on an unknown name; guard batch_reaction_featurizer against invalid/UNKNOWN mode ints so a typo no longer silently produces REAC_PROD features. - CGR layout: explicitly require the first one-hot atom feature to be an atomic-number block (atomic-number / -common / -organic) instead of assuming it silently, since that block is stripped from the CGR second half and used by num_only. All input validation runs before the parse loop. Updates docs/USAGE.md and the reaction test to the new scalar API.
- Add tests/data/gen_reaction_refs.py: a single end-to-end reproducible generator that downloads the RDB7 wb97xd3.csv directly from Zenodo (record 6618262), builds the forward+reverse reaction pool, samples 100 reactions deterministically (seed=7) into sample_rxns_100.csv, and writes the 24 golden .xz files (V1/V2/ORGANIC/RIGR x 6 reaction modes). Fully self-contained, no manual data setup. - Add tests/data/README.md documenting the data provenance, the one-command regeneration pipeline, and the chemprop cross-verification result (bond features and edge indices exact; atom features within float32 round-off). - Regenerate sample_rxns_100.csv and all 24 goldens from this pipeline. - Correct the RDB7 attribution in LICENSE/third-party.txt to the record the data is actually sampled from (Zenodo 6618262, Sci. Data 2022, 9, 417).
|
Thank you for all your comments @sveccham! I think the current version resolves all of them. We can discuss about adding more tests so that the molecular and reaction featurizers do not diverge by mistake (e.g. bond-type-float bug) but that can be a separate PR for better long term maintainability. |
| // --------------------------------------------------------------------------- | ||
| // parse_reaction | ||
| // --------------------------------------------------------------------------- | ||
| CompactReaction parse_reaction(const std::string& reac_smi, const std::string& prod_smi, bool keep_h, bool add_h) { |
There was a problem hiding this comment.
Atom-map uniqueness is not validated. Product maps overwrite at src/reaction_features.cpp:116, and reactant/product inverse maps overwrite at src/reaction_features.cpp:130. Duplicate atom-map numbers are accepted and produce arbitrary correspondence.
Leaving this comment here as it is the entry point.
Is the onus of passing the correct mapping on the user? If so, it would be good to document this.
There was a problem hiding this comment.
Yes, providing the correct atom-mapped reaction SMILES in user's responsibility. This is how we do it in Chemprop. I will document it similar to how we do it in Chemprop documentation.
| mode = cuik_molmaker.reaction_mode_names_to_array(["REAC_DIFF"])[0] | ||
|
|
||
| keep_h = True # keep explicit (mapped) hydrogens, e.g. [H:3] | ||
| add_h = False # add implicit hydrogens via RDKit AddHs (per reaction side) |
There was a problem hiding this comment.
Is add_h=True intended for CGRs? An identity [CH4:1] >> [CH4:1] returns 9 nodes with add_h=True because RDKit-added H atoms are unmapped on both sides. If this is not intentional, I recommend disabling it or documenting it very explicitly.
There was a problem hiding this comment.
This is the expected behavior. In my personal opinion, add_h=True is one of the most useless flags in chemprop. The most used is keep_h as it allows you to keep the hydrogen atoms already expressed in the atom-mapped SMILES provided by the user. I will document this in the code (along with your example as it makes the behavior clear). There may be some reactions where add_h may be used but it's definitely rare. I propose we still keep it to mirror the Chemprop functionality.
Thank you! Left a couple of minor comments/requests for fixes. Some of your answers also helps me understand what is going on better. |
- Note that batch_reaction_featurizer operates on atom-mapped reaction SMILES (Daylight SMILES/SMIRKS) and that providing a correct, unique mapping is the caller's responsibility (uniqueness is not validated). - Explain keep_h (keeps explicit hydrogens already in the SMILES) vs add_h (adds new unmapped hydrogens), with a note that add_h hydrogens become phantom atoms in a CGR: the identity [CH4:1] >> [CH4:1] yields 9 nodes, not 5. Documented in docs/USAGE.md and the batch_reaction_featurizer pybind docstring.
Address review feedback that golden values should not be produced by the code under test. The reference .xz files are now documented as, and reproducible from, Chemprop's CondensedGraphOfReactionFeaturizer + BatchMolGraph collation (pinned to chemprop==2.2.4), with cuik-molmaker never imported during golden generation. - Rewrite tests/data/README.md to embed the full, self-contained generation recipe (env setup pinned to chemprop 2.2.4, exact dependency versions used, and the complete script) so the goldens are reproducible without ambiguity. - Delete tests/data/gen_reaction_refs.py, the cuik-based generator. The committed goldens are byte-identical whether produced by cuik-molmaker or by stock Chemprop 2.2.4, so no golden values change and the regression test still passes with an exact positional comparison (24/24).
- USAGE.md: the identity reaction [CH4:1] >> [CH4:1] with add_h=True yields 9 CGR nodes (1 matched carbon + 8 unmatched hydrogens), verified against both cuik-molmaker and Chemprop 2.2.4. The previous "9 nodes, not 5" was misleading: no setting produces 5 (add_h=False yields a single carbon node). - LICENSE/third-party.txt: point the RDB7 sampling note at tests/data/README.md; the referenced tests/data/gen_rdb7_sample.py never existed, and the actual generator has been removed in favor of the documented Chemprop recipe.
Describe the two flags in the general terms Chemprop uses (keep_h keeps hydrogens already written in the SMILES; add_h adds new ones), and note that add_h's hydrogens are unmapped and become phantom atoms in a CGR. Kept to a short inline comment rather than a separate worked example.
The embedded golden-generation script had one print statement exceeding the project's 88-character line limit. Wrap it so the recipe is black-clean if a user runs it through the project's pre-commit hooks.
Hope this version addresses your comments and requests! Your feedback helped me understand a few things better too. I originally started this because I was trying to use Chemprop CGR for one of my own projects and was frustrated by how slow it was. It turned into a fun weekend learning project, and extending cuik-molmaker felt like the obvious direction to take. Hopefully this makes it easier for people to train foundation models for reaction property prediction, things like yield, bond dissociation energies, barrier heights, solvation, etc. |
There was a problem hiding this comment.
Does it take a significant amount of time for chemprop to generate these? If it doesn't pull too many dependencies in, would seem to make more sense to have it be a test dependency. I know since cuik-molmaker is a dep of chemprop, that would be circular if we put it in the source code, but for a dev test setup it doesn't seem to bad, compared to binary reference files, even if they're regeneratable.
| const RDKit::Bond* b = mol.getBondWithIdx(i); | ||
| gd.bonds[i] = CompactBond{uint8_t(b->getBondType()), | ||
| b->getIsConjugated(), | ||
| ringInfo->numBondRings(i) != 0, |
There was a problem hiding this comment.
Is ringInfo guaranteed to be initialized by whatever pathway is calling this? Just checking, if not might want to form an error message.
There was a problem hiding this comment.
Good question. On the current paths it is always initialized: the only caller reaches populate_graph_arrays via SmilesToMol, which sanitizes by default and runs ring perception. I also confirmed it holds after add_h on ring systems (the appended C-H bonds come back as not-in-ring, with no out-of-range access). I added an explicit ringInfo->isInitialized() check that throws a clear error, so the precondition is enforced rather than assumed if a non-sanitizing path is ever added.
| gd.atoms = std::unique_ptr<CompactAtom[]>(new CompactAtom[num_atoms]); | ||
| for (size_t i = 0; i < num_atoms; ++i) { | ||
| const RDKit::Atom* a = mol.getAtomWithIdx(i); | ||
| gd.atoms[i] = CompactAtom{uint8_t(a->getAtomicNum()), |
There was a problem hiding this comment.
Might want to use designated initializers here and elsewhere with CompactAtom/bond instantiation, helps reduce ordering bugs and improve readability.
There was a problem hiding this comment.
Agreed. Switched CompactAtom, CompactBond, and the two GraphData aggregates to designated initializers. These structs have several same-typed fields, so naming each one removes the risk of a silent positional swap. Thank you!
| @@ -0,0 +1,746 @@ | |||
| // SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
| // SPDX-License-Identifier: Apache-2.0 | |||
|
|
|||
There was a problem hiding this comment.
General note - I know the data struct layout requires direct pointer math in some of these cases, but since the codebase is C++ 20, a subset of the raw pointer array stuff here could be more safely represented using std::span
There was a problem hiding this comment.
Good call. I applied std::span<float> to the two clearest spots, build_num_only and fill_single_bond_feats
| // Build r2p / p2r maps; classify reactant atoms as matched or reac-only | ||
| std::unordered_map<uint32_t, uint32_t> r2p_idx_map, p2r_idx_map; | ||
| std::vector<uint32_t> reac_only_idxs; | ||
| r2p_idx_map.reserve(n_reac_atoms); |
There was a problem hiding this comment.
For dense int->int mappings where the size is known and manageable, it'll likely be much more performant to use a vector, with a sentinel value indicating no entry. While std::unordered_map is definitely the clearer semantic choice, it separately allocates each value, so it's more expensive to populate and has worse cache usage on lookups.
Best of both worlds would be a flat map, with map semantics and vector storage - if this ever bumps to C++23, that's included in the standard library, or cuik_molmaker could pull in one of several third party flat map implementations. That would also be able to optimize your bond lookup table, which currently can't trivially be ported to vector like r2p and p2r. For now I'd suggest using a vector for just r2p and p2r.
There was a problem hiding this comment.
Done for r2p and p2r. Both are dense over 0..n-1, so they are now std::vector<uint32_t> initialized to a NO_IDX sentinel, which avoids the per-node allocation and is more cache-friendly on lookup. I left the two bond-lookup maps as hash maps since they are keyed on packed atom pairs and are not densely enumerable. One thing I fixed along the way: p2r is indexed by product atom index, so it now sizes to n_prod_atoms (the previous reserve used n_reac_atoms, which was fine for a hash map but would undersize a vector when the product has more atoms than the reactant).
… parse Addresses two review nits in populate_graph_arrays: - Use C++20 designated initializers for CompactAtom, CompactBond, and the two GraphData aggregates. These structs have several same-typed fields, so positional init could silently reorder them; naming each field guards against that class of bug. - Assert ringInfo->isInitialized() before calling numBondRings(). SmilesToMol sanitizes by default so ring perception has always run on the current paths (verified across add_h=True on ring systems); the guard makes the precondition explicit and fails loudly if a non-sanitizing path is ever added. No behavior change: all 24 reaction goldens and the full python suite pass.
The reactant->product and product->reactant atom-index maps are dense over 0..n-1, so replace the unordered_map<uint32_t,uint32_t> with std::vector<uint32_t> initialized to a NO_IDX sentinel. This avoids per-node heap allocation and gives contiguous, cache-friendly lookups on the hot CGR path. The two bond-lookup maps are left as hash maps (keyed on packed atom-pair values, not densely enumerable). p2r is now correctly sized to n_prod_atoms (it is indexed by product atom index); the previous reserve() used n_reac_atoms, which was harmless for a hash map but would be undersized for a vector on reactions where the product has more atoms. No behavior change: all 24 reaction goldens are byte-identical, and asymmetric reactions (product larger/smaller than reactant, plus add_h phantom Hs) run correctly across all six modes.
Replace the (float* buf, size_t len) pair with a single std::span<float> in build_num_only() and fill_single_bond_feats(). The span carries the length with the pointer, so call sites pass the backing std::vector directly and can no longer pass a mismatched length. Narrow, targeted use of C++20 span; the resizing feature buffers and tight inner diff loops are left as raw pointers. No behavior change: all 24 reaction goldens remain byte-identical.
3 participants
Add
batch_reaction_featurizerfor CGR reaction featurizationWhat this PR adds
A new
batch_reaction_featurizerfunction — the reaction analogue ofbatch_mol_featurizer— using the Condensed Graph of Reaction (CGR) representation (same as Chemprop'sCondensedGraphOfReactionFeaturizer). API is consistent with the existing package:Supported: all 4 atom featurizer modes (V1, V2, ORGANIC, RIGR) and all 6 reaction modes (
REAC_DIFF,REAC_PROD,PROD_DIFF, and their_BALANCEvariants).keep_h/add_hsemantics match Chemprop'smake_molexactly.Implementation notes
All new code is in
src/reaction_features.cpp(701 lines); additions tofeatures.h,one_hot.cpp, andcuik_molmaker_cpp.cpp. No existing code was modified.Two design choices worth noting:
parse_rxn_side_moldoes not clear atom-map numbers. The existingparse_molstrips them for reordering purposes; reaction featurization needs them for reactant↔product correspondence. Fully additive — molecule featurization is unchanged.unordered_map<uint64_t, size_t>keyed by(min_idx << 32) | max_idx), replacing the Python CGR featurizer's O(n²) atom-pair scan.Node ordering matches Chemprop exactly: reactant atoms 0..n_reac−1, then product-only atoms n_reac..n_cgr−1.
Correctness
Verified against Chemprop's Python
CondensedGraphOfReactionFeaturizeron:[H-:2]lone-hydride nucleophiles)num_onlyvs BALANCE divergence)A test fixture CSV (
tests/data/sample_rxns_100.csv, 110 reactions) is included. Golden.xzreference files can be committed once you've had a look at the design.Speedup benchmarks present here