From 8e962adf9e5d8815105a800430642b5f9a7413ab Mon Sep 17 00:00:00 2001 From: Charlles Abreu Date: Tue, 21 Jul 2026 16:56:18 -0400 Subject: [PATCH 1/4] Accept 1-based atom maps for explicit atom ordering parse_mol only honored atom map numbers that formed a complete 0-based permutation. Track the min and max map number instead, and accept the ordering when the numbers span a contiguous range starting at either 0 or 1, subtracting the minimum when inverting the order. Anything else (partial mapping, duplicates, non-contiguous numbers, or a range starting above 1) is still ignored, leaving the atoms in the order the SMILES parser produced them. Add tests/python/test_atom_map_order.py covering both numbering conventions and the rejection cases. --- src/features.cpp | 18 +++-- src/features.h | 6 +- tests/python/test_atom_map_order.py | 104 ++++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 8 deletions(-) create mode 100644 tests/python/test_atom_map_order.py diff --git a/src/features.cpp b/src/features.cpp index 81542c9..bf3f03d 100644 --- a/src/features.cpp +++ b/src/features.cpp @@ -694,6 +694,8 @@ std::unique_ptr parse_mol(const std::string& smiles_string, bool e // if they indicate a valid order. const unsigned int num_atoms = mol->getNumAtoms(); std::vector atom_order(num_atoms); + unsigned int map_num_min = std::numeric_limits::max(); + unsigned int map_num_max = 0; for (unsigned int i = 0; i < num_atoms; ++i) { RDKit::Atom* atom = mol->getAtomWithIdx(i); if (!atom->hasProp(RDKit::common_properties::molAtomMapNumber)) { @@ -702,10 +704,11 @@ std::unique_ptr parse_mol(const std::string& smiles_string, bool e // from any following atoms that might have it. } else { atom_order[i] = (unsigned int)atom->getAtomMapNum(); - - // 0-based, and must be in range - if (atom_order[i] >= num_atoms) { - ordered = false; + if (atom_order[i] < map_num_min) { + map_num_min = atom_order[i]; + } + if (atom_order[i] > map_num_max) { + map_num_max = atom_order[i]; } // Clear the property, so that any equivalent molecules will @@ -714,13 +717,18 @@ std::unique_ptr parse_mol(const std::string& smiles_string, bool e } } + // Must be 0-based or 1-based, and must be a complete permutation + if (ordered && num_atoms != 0 && (map_num_min > 1 || map_num_max - map_num_min + 1 != num_atoms)) { + ordered = false; + } + if (ordered) { // Invert the order // Use max value as a "not found yet" value constexpr unsigned int not_found_value = std::numeric_limits::max(); std::vector inverse_order(num_atoms, not_found_value); for (unsigned int i = 0; i < num_atoms; ++i) { - unsigned int index = atom_order[i]; + unsigned int index = atom_order[i] - map_num_min; // Can't have the same index twice if (inverse_order[index] != not_found_value) { ordered = false; diff --git a/src/features.h b/src/features.h index e3a0fe8..1970509 100644 --- a/src/features.h +++ b/src/features.h @@ -428,9 +428,9 @@ CUIK_EXPORT std::vector mol_featurizer(const std::string& sm //! Creates an RWMol from a SMILES string. //! //! If `ordered` is true, and the string contains atom classes, called "bookmarks" in RDKit, -//! that form a complete (0-based) ordering of the atoms, the atoms will be reordered according -//! to this explicit order, and the bookmarks will be removed, so that canonical orders -//! can be correctly compared later. +//! that form a complete (either 0-based or 1-based) ordering of the atoms, the atoms will be +//! reordered according to this explicit order, and the bookmarks will be removed, so that +//! canonical orders can be correctly compared later. //! //! This is implemented in cuik_molmaker_cpp.cpp, but is declared in this header so //! that both labels.cpp and features.cpp can call it. diff --git a/tests/python/test_atom_map_order.py b/tests/python/test_atom_map_order.py new file mode 100644 index 0000000..98371c5 --- /dev/null +++ b/tests/python/test_atom_map_order.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # noqa: E501 +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the atom-map based reordering performed by `parse_mol`. + +Atom maps ("bookmarks" in RDKit) are honored only when they form a complete +permutation of the atoms, numbered either 0-based or 1-based. Anything else is +ignored, leaving the atoms in the order the SMILES parser produced them. +""" + +import numpy as np +import pytest + +import cuik_molmaker + + +def atom_features(smiles): + """Atom feature rows for `smiles`, one row per atom, in featurizer order.""" + atom_onehot_property_list = cuik_molmaker.atom_onehot_feature_names_to_array( + ["atomic-number-common"] + ) + atom_float_property_list = cuik_molmaker.atom_float_feature_names_to_array( + ["atomic-number"] + ) + bond_property_list = cuik_molmaker.bond_feature_names_to_array(["bond-type-float"]) + explicit_H, offset_carbon, duplicate_edges, add_self_loop = ( + False, + False, + False, + False, + ) + atom_feats = cuik_molmaker.mol_featurizer( + smiles, + atom_onehot_property_list, + atom_float_property_list, + bond_property_list, + explicit_H, + offset_carbon, + duplicate_edges, + add_self_loop, + )[0] + return atom_feats + + +def test_atom_order_is_observable(): + """Guard for the other tests: writing the atoms in a different order in the + SMILES really does produce differently ordered features.""" + assert not np.array_equal(atom_features("CO"), atom_features("OC")) + + +@pytest.mark.parametrize( + "smiles", + [ + "[CH3:0][OH:1]", # 0-based, already in order + "[OH:1][CH3:0]", # 0-based, reversed + "[CH3:1][OH:2]", # 1-based, already in order + "[OH:2][CH3:1]", # 1-based, reversed + ], +) +def test_complete_map_reorders_atoms(smiles): + """A complete 0- or 1-based atom map puts the atoms in the mapped order, + regardless of the order they appear in the SMILES.""" + np.testing.assert_allclose(atom_features(smiles), atom_features("CO")) + + +def test_zero_and_one_based_maps_agree(): + np.testing.assert_allclose( + atom_features("[OH:1][CH3:0]"), atom_features("[OH:2][CH3:1]") + ) + + +@pytest.mark.parametrize( + "smiles", + [ + "[CH3:3][NH:2][OH:1]", # 1-based, fully reversed + "[CH3:2][NH:1][OH:0]", # 0-based, fully reversed + ], +) +def test_complete_map_applies_arbitrary_permutation(smiles): + """Three atoms, mapped so the featurizer order is O, N, C.""" + np.testing.assert_allclose(atom_features(smiles), atom_features("ONC")) + + +@pytest.mark.parametrize( + "smiles", + [ + "OC", # no atom maps at all + "[OH:1]C", # only some atoms mapped + "[OH:4][CH3:3]", # complete permutation, but neither 0- nor 1-based + "[OH:1][CH3:1]", # duplicated map number + "[OH:1][CH3:5]", # 1-based but not contiguous + "[OH:0][CH3:5]", # 0-based but not contiguous + ], +) +def test_invalid_map_leaves_atoms_in_parsed_order(smiles): + """Maps that are not a complete 0- or 1-based permutation are ignored, so + the atoms stay in the order the SMILES string lists them (O, then C).""" + np.testing.assert_allclose(atom_features(smiles), atom_features("OC")) + + +def test_map_numbers_do_not_leak_into_features(): + """The map numbers are stripped, so a mapped molecule and the equivalent + unmapped one featurize identically.""" + np.testing.assert_allclose(atom_features("[CH3:1][OH:2]"), atom_features("CO")) From 607e430d026eadf39917ecde3ea7f99d31a422bf Mon Sep 17 00:00:00 2001 From: Charlles Abreu Date: Tue, 21 Jul 2026 17:01:00 -0400 Subject: [PATCH 2/4] Expose the `ordered` flag in the Python API `parse_mol` could already reorder atoms according to their atom map numbers, but the featurizers always requested it and Python callers had no way to turn it off. Thread an `ordered` parameter through `mol_featurizer` and `batch_mol_featurizer` down to `read_graph`/`parse_mol`, and bind both functions with named arguments so `ordered` can be passed by keyword. It defaults to true, so existing callers are unaffected. Document the option in docs/USAGE.md and cover it in the atom-map ordering tests. --- docs/USAGE.md | 16 +++++++ src/cuik_molmaker_cpp.cpp | 36 +++++++++++---- src/features.cpp | 15 +++--- src/features.h | 10 +++- tests/python/test_atom_map_order.py | 72 +++++++++++++++++++++++++++-- 5 files changed, 128 insertions(+), 21 deletions(-) diff --git a/docs/USAGE.md b/docs/USAGE.md index d22ba44..a943d85 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -38,6 +38,22 @@ duplicate_edges = True add_self_loop = False ``` +#### Controlling atom order (optional) +Both `mol_featurizer` and `batch_mol_featurizer` accept an optional `ordered` argument, which +defaults to `True`. When it is true and the SMILES string carries atom map numbers that form a +complete ordering of the atoms — numbered either `0..n-1` or `1..n` — the atoms are reordered +accordingly, so that features are emitted in the order you asked for: + +```python +# Atoms come out in the order O, C rather than the order written in the SMILES. +cuik_molmaker.mol_featurizer("[CH3:2][OH:1]", ..., ordered=True) +``` + +Atom maps that do not form such an ordering (partial mapping, duplicates, gaps, or a range not +starting at 0 or 1) are ignored, leaving the atoms in the order the SMILES parser produced them. +Pass `ordered=False` to ignore atom map numbers entirely. Either way the map numbers are stripped +before featurization and never appear in the features. + #### Generate atom and bond features ```python all_features =cuik_molmaker.mol_featurizer(smiles, atom_onehot_feature_array, atom_float_feature_array, bond_feature_array, explicit_h, offset_carbon, duplicate_edges, add_self_loop) diff --git a/src/cuik_molmaker_cpp.cpp b/src/cuik_molmaker_cpp.cpp index b27ee82..7d851ae 100644 --- a/src/cuik_molmaker_cpp.cpp +++ b/src/cuik_molmaker_cpp.cpp @@ -49,14 +49,34 @@ PYBIND11_MODULE(cuik_molmaker_cpp, m) { m.def("bond_feature_names_to_array", &bond_feature_names_to_array, "Accepts feature names and returns a NumPy array representing them as integers"); - m.def( - "mol_featurizer", - &mol_featurizer, - "Accepts a SMILES string and returns a list of NumPy arrays representing atom and bond features of the molecule."); - m.def( - "batch_mol_featurizer", - &batch_mol_featurizer, - "Accepts a list of SMILES strings and returns a list of NumPy arrays representing atom and bond features of the molecules."); + m.def("mol_featurizer", + &mol_featurizer, + "Accepts a SMILES string and returns a list of NumPy arrays representing atom and bond features of the " + "molecule. If ordered is true (the default), atom map numbers forming a complete 0-based or 1-based ordering " + "of the atoms are used to reorder them; maps that do not form such an ordering are ignored.", + py::arg("smiles_string"), + py::arg("atom_property_list_onehot"), + py::arg("atom_property_list_float"), + py::arg("bond_property_list"), + py::arg("explicit_H"), + py::arg("offset_carbon"), + py::arg("duplicate_edges"), + py::arg("add_self_loop"), + py::arg("ordered") = true); + m.def("batch_mol_featurizer", + &batch_mol_featurizer, + "Accepts a list of SMILES strings and returns a list of NumPy arrays representing atom and bond features of " + "the molecules. If ordered is true (the default), atom map numbers forming a complete 0-based or 1-based " + "ordering of the atoms are used to reorder them; maps that do not form such an ordering are ignored.", + py::arg("smiles_list"), + py::arg("atom_property_list_onehot"), + py::arg("atom_property_list_float"), + py::arg("bond_property_list"), + py::arg("explicit_H"), + py::arg("offset_carbon"), + py::arg("duplicate_edges"), + py::arg("add_self_loop"), + py::arg("ordered") = true); m.def("list_all_atom_onehot_features", &list_all_atom_onehot_features, "Returns a list of all atom one-hot features."); diff --git a/src/features.cpp b/src/features.cpp index bf3f03d..1d1646a 100644 --- a/src/features.cpp +++ b/src/features.cpp @@ -25,8 +25,8 @@ // This is called by `mol_featurizer` and `batch_mol_featurizer` to parse the SMILES string into an RWMol and // cache some data about the atoms and bonds. -static GraphData read_graph(const std::string& smiles_string, bool explicit_H) { - std::unique_ptr mol{parse_mol(smiles_string, explicit_H)}; +static GraphData read_graph(const std::string& smiles_string, bool explicit_H, bool ordered) { + std::unique_ptr mol{parse_mol(smiles_string, explicit_H, ordered)}; if (!mol) { return GraphData{0, std::unique_ptr(), 0, std::unique_ptr(), std::move(mol)}; @@ -519,7 +519,8 @@ std::vector mol_featurizer(const std::string& smiles_string, bool explicit_H, bool offset_carbon, bool duplicate_edges, - bool add_self_loop) { + bool add_self_loop, + bool ordered) { return batch_mol_featurizer(std::vector{smiles_string}, atom_property_list_onehot, atom_property_list_float, @@ -527,7 +528,8 @@ std::vector mol_featurizer(const std::string& smiles_string, explicit_H, offset_carbon, duplicate_edges, - add_self_loop); + add_self_loop, + ordered); } std::vector batch_mol_featurizer(const std::vector& smiles_list, @@ -537,7 +539,8 @@ std::vector batch_mol_featurizer(const std::vector& smil bool explicit_H, bool offset_carbon, bool duplicate_edges, - bool add_self_loop) { + bool add_self_loop, + bool ordered) { const size_t n_smiles = smiles_list.size(); // Create graphs @@ -547,7 +550,7 @@ std::vector batch_mol_featurizer(const std::vector& smil size_t total_num_atoms = 0, total_num_bonds = 0; for (const auto& smiles : smiles_list) { - GraphData igraph = read_graph(smiles, explicit_H); + GraphData igraph = read_graph(smiles, explicit_H, ordered); total_num_atoms += igraph.num_atoms; total_num_bonds += igraph.num_bonds; graph_list.push_back(std::move(igraph)); diff --git a/src/features.h b/src/features.h index 1970509..5862a97 100644 --- a/src/features.h +++ b/src/features.h @@ -408,6 +408,8 @@ CUIK_EXPORT std::vector list_all_bond_features(); //! self-edges. //! @param offset_carbon If true, some atom float features will subtract a //! value representing carbon, so that carbon atoms would have value zero. +//! @param ordered If true, atom map numbers in the SMILES string that form a complete ordering +//! of the atoms will be used to reorder them. See `parse_mol` for details. //! @return A vector of torch NumPy/pybind arrays for the features. The first array is the atom features //! array, `num_atoms` by the number of values required for all one-hot and float atom //! features. The second array is the bond features array, `num_edges` (or @@ -423,7 +425,8 @@ CUIK_EXPORT std::vector mol_featurizer(const std::string& sm bool explicit_H, bool offset_carbon, bool duplicate_edges, - bool add_self_loop); + bool add_self_loop, + bool ordered = true); //! Creates an RWMol from a SMILES string. //! @@ -455,6 +458,8 @@ std::unique_ptr parse_mol(const std::string& smiles_string, bool e //! value representing carbon, so that carbon atoms would have value zero. //! @param add_self_loop If true, bond features will have values stored for //! self-edges. +//! @param ordered If true, atom map numbers in each SMILES string that form a complete ordering +//! of the atoms will be used to reorder them. See `parse_mol` for details. //! @return A vector of NumPy/pybind arrays for the features. The first array is the atom features //! array, total number of atoms by the number of values required for all one-hot and //! float atom features. The second array is the bond features array, total number of @@ -471,7 +476,8 @@ CUIK_EXPORT std::vector batch_mol_featurizer(const std::vector Date: Tue, 21 Jul 2026 17:34:01 -0400 Subject: [PATCH 3/4] Add keep_h and ignore_stereo to mol_featurizer / batch_mol_featurizer --- docs/USAGE.md | 14 +++ src/cuik_molmaker_cpp.cpp | 16 +++- src/features.cpp | 40 +++++++-- src/features.h | 27 +++++- tests/python/test_parse_options.py | 136 +++++++++++++++++++++++++++++ 5 files changed, 217 insertions(+), 16 deletions(-) create mode 100644 tests/python/test_parse_options.py diff --git a/docs/USAGE.md b/docs/USAGE.md index a943d85..d49813e 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -54,6 +54,20 @@ starting at 0 or 1) are ignored, leaving the atoms in the order the SMILES parse Pass `ordered=False` to ignore atom map numbers entirely. Either way the map numbers are stripped before featurization and never appear in the features. +#### Keeping explicit hydrogens and ignoring stereochemistry (optional) +`mol_featurizer` and `batch_mol_featurizer` also accept `keep_h` and `ignore_stereo`, +both defaulting to `False`: + +```python +# keep_h retains explicit hydrogens already written in the SMILES; it does not add +# any that aren't there (use explicit_H for that). +cuik_molmaker.mol_featurizer("[H]C([H])([H])O", ..., keep_h=True) # 5 atoms, not 2 + +# ignore_stereo clears R/S and cis/trans stereochemistry from the molecule +# (RDKit::MolOps::removeStereochemistry) before featurizing. +cuik_molmaker.mol_featurizer("N[C@@H](C)C(=O)O", ..., ignore_stereo=True) +``` + #### Generate atom and bond features ```python all_features =cuik_molmaker.mol_featurizer(smiles, atom_onehot_feature_array, atom_float_feature_array, bond_feature_array, explicit_h, offset_carbon, duplicate_edges, add_self_loop) diff --git a/src/cuik_molmaker_cpp.cpp b/src/cuik_molmaker_cpp.cpp index 7d851ae..db6299b 100644 --- a/src/cuik_molmaker_cpp.cpp +++ b/src/cuik_molmaker_cpp.cpp @@ -53,7 +53,9 @@ PYBIND11_MODULE(cuik_molmaker_cpp, m) { &mol_featurizer, "Accepts a SMILES string and returns a list of NumPy arrays representing atom and bond features of the " "molecule. If ordered is true (the default), atom map numbers forming a complete 0-based or 1-based ordering " - "of the atoms are used to reorder them; maps that do not form such an ordering are ignored.", + "of the atoms are used to reorder them; maps that do not form such an ordering are ignored. keep_h retains " + "explicit hydrogens already written in the SMILES (it does not add any). ignore_stereo clears R/S and " + "cis/trans stereochemistry after parsing.", py::arg("smiles_string"), py::arg("atom_property_list_onehot"), py::arg("atom_property_list_float"), @@ -62,12 +64,16 @@ PYBIND11_MODULE(cuik_molmaker_cpp, m) { py::arg("offset_carbon"), py::arg("duplicate_edges"), py::arg("add_self_loop"), - py::arg("ordered") = true); + py::arg("ordered") = true, + py::arg("keep_h") = false, + py::arg("ignore_stereo") = false); m.def("batch_mol_featurizer", &batch_mol_featurizer, "Accepts a list of SMILES strings and returns a list of NumPy arrays representing atom and bond features of " "the molecules. If ordered is true (the default), atom map numbers forming a complete 0-based or 1-based " - "ordering of the atoms are used to reorder them; maps that do not form such an ordering are ignored.", + "ordering of the atoms are used to reorder them; maps that do not form such an ordering are ignored. keep_h " + "retains explicit hydrogens already written in each SMILES (it does not add any). ignore_stereo clears R/S " + "and cis/trans stereochemistry after parsing.", py::arg("smiles_list"), py::arg("atom_property_list_onehot"), py::arg("atom_property_list_float"), @@ -76,7 +82,9 @@ PYBIND11_MODULE(cuik_molmaker_cpp, m) { py::arg("offset_carbon"), py::arg("duplicate_edges"), py::arg("add_self_loop"), - py::arg("ordered") = true); + py::arg("ordered") = true, + py::arg("keep_h") = false, + py::arg("ignore_stereo") = false); m.def("list_all_atom_onehot_features", &list_all_atom_onehot_features, "Returns a list of all atom one-hot features."); diff --git a/src/features.cpp b/src/features.cpp index 1d1646a..9ef5deb 100644 --- a/src/features.cpp +++ b/src/features.cpp @@ -25,8 +25,12 @@ // This is called by `mol_featurizer` and `batch_mol_featurizer` to parse the SMILES string into an RWMol and // cache some data about the atoms and bonds. -static GraphData read_graph(const std::string& smiles_string, bool explicit_H, bool ordered) { - std::unique_ptr mol{parse_mol(smiles_string, explicit_H, ordered)}; +static GraphData read_graph(const std::string& smiles_string, + bool explicit_H, + bool ordered, + bool keep_h, + bool ignore_stereo) { + std::unique_ptr mol{parse_mol(smiles_string, explicit_H, ordered, keep_h, ignore_stereo)}; if (!mol) { return GraphData{0, std::unique_ptr(), 0, std::unique_ptr(), std::move(mol)}; @@ -520,7 +524,9 @@ std::vector mol_featurizer(const std::string& smiles_string, bool offset_carbon, bool duplicate_edges, bool add_self_loop, - bool ordered) { + bool ordered, + bool keep_h, + bool ignore_stereo) { return batch_mol_featurizer(std::vector{smiles_string}, atom_property_list_onehot, atom_property_list_float, @@ -529,7 +535,9 @@ std::vector mol_featurizer(const std::string& smiles_string, offset_carbon, duplicate_edges, add_self_loop, - ordered); + ordered, + keep_h, + ignore_stereo); } std::vector batch_mol_featurizer(const std::vector& smiles_list, @@ -540,7 +548,9 @@ std::vector batch_mol_featurizer(const std::vector& smil bool offset_carbon, bool duplicate_edges, bool add_self_loop, - bool ordered) { + bool ordered, + bool keep_h, + bool ignore_stereo) { const size_t n_smiles = smiles_list.size(); // Create graphs @@ -550,7 +560,7 @@ std::vector batch_mol_featurizer(const std::vector& smil size_t total_num_atoms = 0, total_num_bonds = 0; for (const auto& smiles : smiles_list) { - GraphData igraph = read_graph(smiles, explicit_H, ordered); + GraphData igraph = read_graph(smiles, explicit_H, ordered, keep_h, ignore_stereo); total_num_atoms += igraph.num_atoms; total_num_bonds += igraph.num_bonds; graph_list.push_back(std::move(igraph)); @@ -683,9 +693,14 @@ std::vector batch_mol_featurizer(const std::vector& smil // Creates an RWMol from a SMILES string. // See the declaration in features.h for more details. -std::unique_ptr parse_mol(const std::string& smiles_string, bool explicit_H, bool ordered) { - // Parse SMILES string with default options - RDKit::SmilesParserParams params; +std::unique_ptr parse_mol(const std::string& smiles_string, + bool explicit_H, + bool ordered, + bool keep_h, + bool ignore_stereo) { + // Parse SMILES string, optionally keeping explicit hydrogens (keep_h) + RDKit::SmilesParserParams params; + params.removeHs = !keep_h; // keep_h=true keeps explicit [H:n] atoms std::unique_ptr mol{RDKit::SmilesToMol(smiles_string, params)}; if (!mol) { return mol; @@ -753,5 +768,12 @@ std::unique_ptr parse_mol(const std::string& smiles_string, bool e // and calling it again shouldn't have any net effect. // RDKit::MolOps::removeHs(*mol); } + if (ignore_stereo) { + // Clears chiral tags, the cached _CIPCode property, bond stereo, and bond + // directions in one call. A manual port of clearing only the chiral tag + // and bond stereo (as some other tools do) would miss _CIPCode, which the + // float `chirality` feature reads directly. + RDKit::MolOps::removeStereochemistry(*mol); + } return mol; } \ No newline at end of file diff --git a/src/features.h b/src/features.h index 5862a97..ff28621 100644 --- a/src/features.h +++ b/src/features.h @@ -410,6 +410,11 @@ CUIK_EXPORT std::vector list_all_bond_features(); //! value representing carbon, so that carbon atoms would have value zero. //! @param ordered If true, atom map numbers in the SMILES string that form a complete ordering //! of the atoms will be used to reorder them. See `parse_mol` for details. +//! @param keep_h If true, explicit hydrogen atoms already written in the SMILES string are +//! kept (SmilesParserParams::removeHs = false). This does not add any hydrogens; +//! see `explicit_H` for that. +//! @param ignore_stereo If true, R/S and cis/trans stereochemistry is cleared from the +//! molecule (RDKit::MolOps::removeStereochemistry) after parsing. //! @return A vector of torch NumPy/pybind arrays for the features. The first array is the atom features //! array, `num_atoms` by the number of values required for all one-hot and float atom //! features. The second array is the bond features array, `num_edges` (or @@ -426,7 +431,9 @@ CUIK_EXPORT std::vector mol_featurizer(const std::string& sm bool offset_carbon, bool duplicate_edges, bool add_self_loop, - bool ordered = true); + bool ordered = true, + bool keep_h = false, + bool ignore_stereo = false); //! Creates an RWMol from a SMILES string. //! @@ -435,9 +442,17 @@ CUIK_EXPORT std::vector mol_featurizer(const std::string& sm //! reordered according to this explicit order, and the bookmarks will be removed, so that //! canonical orders can be correctly compared later. //! +//! If `keep_h` is true, explicit hydrogen atoms already written in the SMILES string are kept +//! (this does not add any hydrogens). If `ignore_stereo` is true, R/S and cis/trans +//! stereochemistry is cleared from the molecule after parsing. +//! //! This is implemented in cuik_molmaker_cpp.cpp, but is declared in this header so //! that both labels.cpp and features.cpp can call it. -std::unique_ptr parse_mol(const std::string& smiles_string, bool explicit_H, bool ordered = true); +std::unique_ptr parse_mol(const std::string& smiles_string, + bool explicit_H, + bool ordered = true, + bool keep_h = false, + bool ignore_stereo = false); //! `batch_mol_featurizer` is called from Python to get feature arrays for `smiles_list`. //! @@ -460,6 +475,10 @@ std::unique_ptr parse_mol(const std::string& smiles_string, bool e //! self-edges. //! @param ordered If true, atom map numbers in each SMILES string that form a complete ordering //! of the atoms will be used to reorder them. See `parse_mol` for details. +//! @param keep_h If true, explicit hydrogen atoms already written in each SMILES string are +//! kept. See `parse_mol` for details. +//! @param ignore_stereo If true, R/S and cis/trans stereochemistry is cleared from each +//! molecule after parsing. See `parse_mol` for details. //! @return A vector of NumPy/pybind arrays for the features. The first array is the atom features //! array, total number of atoms by the number of values required for all one-hot and //! float atom features. The second array is the bond features array, total number of @@ -477,7 +496,9 @@ CUIK_EXPORT std::vector batch_mol_featurizer(const std::vector Date: Tue, 21 Jul 2026 17:39:02 -0400 Subject: [PATCH 4/4] Add ignore_stereo to batch_reaction_featurizer --- docs/USAGE.md | 6 +++++- src/cuik_molmaker_cpp.cpp | 20 +++++++++++++++++--- src/features.h | 16 +++++++++++++--- src/reaction_features.cpp | 22 ++++++++++++++++------ tests/python/test_parse_options.py | 29 +++++++++++++++++++++++++++++ 5 files changed, 80 insertions(+), 13 deletions(-) diff --git a/docs/USAGE.md b/docs/USAGE.md index d49813e..036e33e 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -127,10 +127,14 @@ mode = cuik_molmaker.reaction_mode_to_int("REAC_DIFF") keep_h = True add_h = False +# ignore_stereo clears R/S and cis/trans stereochemistry from both the reactant and +# product before featurizing (RDKit::MolOps::removeStereochemistry). Defaults to False. +ignore_stereo = False + rxn_features = cuik_molmaker.batch_reaction_featurizer( reac_smiles_list, prod_smiles_list, atom_onehot_feature_array, atom_float_feature_array, bond_feature_array, - keep_h, add_h, offset_carbon, mode) + keep_h, add_h, offset_carbon, mode, ignore_stereo=ignore_stereo) # CGR atom features from all reactions, concatenated along dimension 0. # The feature dimension is doubled relative to a single molecule (reactant + product). diff --git a/src/cuik_molmaker_cpp.cpp b/src/cuik_molmaker_cpp.cpp index db6299b..1680ec8 100644 --- a/src/cuik_molmaker_cpp.cpp +++ b/src/cuik_molmaker_cpp.cpp @@ -107,7 +107,8 @@ PYBIND11_MODULE(cuik_molmaker_cpp, m) { bool keep_h, bool add_h, bool offset_carbon, - int64_t mode_int) { + int64_t mode_int, + bool ignore_stereo) { return batch_reaction_featurizer(reac_smiles_list, prod_smiles_list, atom_property_list_onehot, @@ -116,11 +117,24 @@ PYBIND11_MODULE(cuik_molmaker_cpp, m) { keep_h, add_h, offset_carbon, - ReactionMode(mode_int)); + ReactionMode(mode_int), + ignore_stereo); }, "Accepts lists of reactant and product SMILES strings and returns a list of NumPy arrays " "representing the Condensed Graph of Reaction (CGR) atom and bond features of the reactions. " "SMILES must be atom-mapped and providing a correct, unique mapping is the caller's " "responsibility (uniqueness is not validated). keep_h keeps explicit hydrogens already in the " - "SMILES; add_h adds new unmapped hydrogens that become phantom atoms in the CGR."); + "SMILES; add_h adds new unmapped hydrogens that become phantom atoms in the CGR. ignore_stereo " + "clears R/S and cis/trans stereochemistry from both the reactant and product before " + "featurizing.", + py::arg("reac_smiles_list"), + py::arg("prod_smiles_list"), + py::arg("atom_property_list_onehot"), + py::arg("atom_property_list_float"), + py::arg("bond_property_list"), + py::arg("keep_h"), + py::arg("add_h"), + py::arg("offset_carbon"), + py::arg("mode_int"), + py::arg("ignore_stereo") = false); } diff --git a/src/features.h b/src/features.h index ff28621..96d855b 100644 --- a/src/features.h +++ b/src/features.h @@ -504,15 +504,22 @@ CUIK_EXPORT std::vector batch_mol_featurizer(const std::vector parse_rxn_side_mol(const std::string& smiles, bool keep_h, bool add_h); +//! @param ignore_stereo If true, RDKit::MolOps::removeStereochemistry is called after +//! add_h, clearing R/S and cis/trans stereochemistry. +std::unique_ptr parse_rxn_side_mol(const std::string& smiles, + bool keep_h, + bool add_h, + bool ignore_stereo = false); //! Parses a reaction SMILES pair into a CompactReaction (atom correspondence + both GraphData). //! Both reac_smi and prod_smi must contain atom-map numbers. //! keep_h / add_h semantics match chemprop's _ReactionDatapointMixin.from_smi exactly. +//! ignore_stereo is applied identically to both the reactant and product molecules. CUIK_EXPORT CompactReaction parse_reaction(const std::string& reac_smi, const std::string& prod_smi, bool keep_h, - bool add_h); + bool add_h, + bool ignore_stereo = false); //! Converts a reaction mode name to its integer enum value; throws std::runtime_error on an unknown name. CUIK_EXPORT int64_t reaction_mode_to_int(const std::string& mode); @@ -524,6 +531,8 @@ CUIK_EXPORT int64_t reaction_mode_to_int(const std::string& mode); //! @param keep_h If true, retain explicit mapped [H:n] atoms (required for E2/SN2) //! @param add_h If true, add unmapped Hs via RDKit::MolOps::addHs (after parsing) //! @param mode CGR featurization mode (which combination of reac/prod/diff) +//! @param ignore_stereo If true, clear R/S and cis/trans stereochemistry from both the +//! reactant and product before featurizing. //! @return 5 arrays: [atom_feats, bond_feats, edge_index, rev_edge_index, batch] CUIK_EXPORT std::vector batch_reaction_featurizer(const std::vector& reac_smiles_list, const std::vector& prod_smiles_list, @@ -533,4 +542,5 @@ CUIK_EXPORT std::vector batch_reaction_featurizer(const std::vector build_bond_lookup(const GraphData& // --------------------------------------------------------------------------- // parse_rxn_side_mol // --------------------------------------------------------------------------- -std::unique_ptr parse_rxn_side_mol(const std::string& smiles, bool keep_h, bool add_h) { +std::unique_ptr parse_rxn_side_mol(const std::string& smiles, + bool keep_h, + bool add_h, + bool ignore_stereo) { RDKit::SmilesParserParams params; params.removeHs = !keep_h; // keep_h=true keeps explicit [H:n] atoms std::unique_ptr mol{RDKit::SmilesToMol(smiles, params)}; @@ -101,15 +104,21 @@ std::unique_ptr parse_rxn_side_mol(const std::string& smiles, bool // Do NOT reorder atoms — reactions preserve SMILES parse order. if (add_h) RDKit::MolOps::addHs(*mol); + if (ignore_stereo) + RDKit::MolOps::removeStereochemistry(*mol); return mol; } // --------------------------------------------------------------------------- // parse_reaction // --------------------------------------------------------------------------- -CompactReaction parse_reaction(const std::string& reac_smi, const std::string& prod_smi, bool keep_h, bool add_h) { - std::unique_ptr reac_mol = parse_rxn_side_mol(reac_smi, keep_h, add_h); - std::unique_ptr prod_mol = parse_rxn_side_mol(prod_smi, keep_h, add_h); +CompactReaction parse_reaction(const std::string& reac_smi, + const std::string& prod_smi, + bool keep_h, + bool add_h, + bool ignore_stereo) { + std::unique_ptr reac_mol = parse_rxn_side_mol(reac_smi, keep_h, add_h, ignore_stereo); + std::unique_ptr prod_mol = parse_rxn_side_mol(prod_smi, keep_h, add_h, ignore_stereo); if (!reac_mol) throw std::runtime_error("Failed to parse reactant SMILES: " + reac_smi); if (!prod_mol) @@ -628,7 +637,8 @@ std::vector batch_reaction_featurizer(const std::vector& bool keep_h, bool add_h, bool offset_carbon, - ReactionMode mode) { + ReactionMode mode, + bool ignore_stereo) { // Validate inputs up front (before the expensive per-reaction parse loop). if (reac_smiles_list.size() != prod_smiles_list.size()) throw std::runtime_error("reac_smiles_list and prod_smiles_list must have the same length"); @@ -659,7 +669,7 @@ std::vector batch_reaction_featurizer(const std::vector& std::vector reactions; reactions.reserve(n_rxns); for (size_t i = 0; i < n_rxns; ++i) - reactions.push_back(parse_reaction(reac_smiles_list[i], prod_smiles_list[i], keep_h, add_h)); + reactions.push_back(parse_reaction(reac_smiles_list[i], prod_smiles_list[i], keep_h, add_h, ignore_stereo)); // Feature dimensions const size_t single_atom_fdim = compute_atom_dim(atom_property_list_onehot, atom_property_list_float); diff --git a/tests/python/test_parse_options.py b/tests/python/test_parse_options.py index ecbb505..821dedc 100644 --- a/tests/python/test_parse_options.py +++ b/tests/python/test_parse_options.py @@ -134,3 +134,32 @@ def test_ignore_stereo_clears_bond_stereo(): assert not np.allclose(cis, trans) np.testing.assert_allclose(cis_cleared, trans_cleared) np.testing.assert_allclose(cis_cleared, plain) + + +# --------------------------------------------------------------------------- +# ignore_stereo on the reaction (CGR) path +# --------------------------------------------------------------------------- + + +def test_batch_reaction_featurizer_ignore_stereo_clears_chirality(): + """Identity reaction (product == reactant) on an atom-mapped molecule + with one stereocenter (atom map 2). ignore_stereo must change the CGR + atom features, since it strips the stereocenter on both reactant and + product sides.""" + reac = ["[NH2:1][C@@H:2]([CH3:3])[C:4](=[O:5])[OH:6]"] + prod = ["[NH2:1][C@@H:2]([CH3:3])[C:4](=[O:5])[OH:6]"] + oh = cuik_molmaker.atom_onehot_feature_names_to_array( + ["atomic-number-common", "chirality"] + ) + fl = cuik_molmaker.atom_float_feature_names_to_array([]) + bp = cuik_molmaker.bond_feature_names_to_array(["is-null"]) + mode = cuik_molmaker.reaction_mode_to_int("REAC_DIFF") + + with_stereo = cuik_molmaker.batch_reaction_featurizer( + reac, prod, oh, fl, bp, False, False, False, mode, ignore_stereo=False + ) + without_stereo = cuik_molmaker.batch_reaction_featurizer( + reac, prod, oh, fl, bp, False, False, False, mode, ignore_stereo=True + ) + + assert not np.allclose(with_stereo[0], without_stereo[0])