diff --git a/docs/USAGE.md b/docs/USAGE.md index d22ba44..036e33e 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -38,6 +38,36 @@ 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. + +#### 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) @@ -97,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 b27ee82..1680ec8 100644 --- a/src/cuik_molmaker_cpp.cpp +++ b/src/cuik_molmaker_cpp.cpp @@ -49,14 +49,42 @@ 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. 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"), + 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, + 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. 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"), + 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, + 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."); @@ -79,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, @@ -88,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.cpp b/src/features.cpp index 81542c9..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) { - std::unique_ptr mol{parse_mol(smiles_string, explicit_H)}; +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)}; @@ -519,7 +523,10 @@ 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, + bool keep_h, + bool ignore_stereo) { return batch_mol_featurizer(std::vector{smiles_string}, atom_property_list_onehot, atom_property_list_float, @@ -527,7 +534,10 @@ std::vector mol_featurizer(const std::string& smiles_string, explicit_H, offset_carbon, duplicate_edges, - add_self_loop); + add_self_loop, + ordered, + keep_h, + ignore_stereo); } std::vector batch_mol_featurizer(const std::vector& smiles_list, @@ -537,7 +547,10 @@ 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, + bool keep_h, + bool ignore_stereo) { const size_t n_smiles = smiles_list.size(); // Create graphs @@ -547,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); + 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)); @@ -680,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; @@ -694,6 +712,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 +722,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 +735,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; @@ -742,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 e3a0fe8..96d855b 100644 --- a/src/features.h +++ b/src/features.h @@ -408,6 +408,13 @@ 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. +//! @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 @@ -423,18 +430,29 @@ 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, + bool keep_h = false, + bool ignore_stereo = false); //! 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. +//! +//! 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`. //! @@ -455,6 +473,12 @@ 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. +//! @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 @@ -471,21 +495,31 @@ 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); @@ -497,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, @@ -506,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_atom_map_order.py b/tests/python/test_atom_map_order.py new file mode 100644 index 0000000..16f270d --- /dev/null +++ b/tests/python/test_atom_map_order.py @@ -0,0 +1,166 @@ +# 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 feature_arrays(): + 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"]) + return atom_onehot_property_list, atom_float_property_list, bond_property_list + + +def atom_features(smiles, **kwargs): + """Atom feature rows for `smiles`, one row per atom, in featurizer order. + + Extra keyword arguments are forwarded to `mol_featurizer`, so passing + nothing exercises the default value of `ordered`. + """ + explicit_H, offset_carbon, duplicate_edges, add_self_loop = ( + False, + False, + False, + False, + ) + atom_feats = cuik_molmaker.mol_featurizer( + smiles, + *feature_arrays(), + explicit_H, + offset_carbon, + duplicate_edges, + add_self_loop, + **kwargs, + )[0] + return atom_feats + + +def batch_atom_features(smiles_list, **kwargs): + """Same as `atom_features`, but through `batch_mol_featurizer`.""" + explicit_H, offset_carbon, duplicate_edges, add_self_loop = ( + False, + False, + False, + False, + ) + atom_feats = cuik_molmaker.batch_mol_featurizer( + smiles_list, + *feature_arrays(), + explicit_H, + offset_carbon, + duplicate_edges, + add_self_loop, + **kwargs, + )[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")) + + +@pytest.mark.parametrize("smiles", ["[OH:1][CH3:0]", "[OH:2][CH3:1]"]) +def test_ordered_false_disables_reordering(smiles): + """With ordered=False the atom maps are ignored entirely, leaving the atoms + in the order the SMILES string lists them (O, then C).""" + np.testing.assert_allclose( + atom_features(smiles, ordered=False), atom_features("OC") + ) + + +@pytest.mark.parametrize("smiles", ["[OH:1][CH3:0]", "[OH:2][CH3:1]"]) +def test_ordered_defaults_to_true(smiles): + np.testing.assert_allclose( + atom_features(smiles), atom_features(smiles, ordered=True) + ) + + +def test_ordered_accepted_positionally(): + """`ordered` is the ninth positional argument, after add_self_loop.""" + np.testing.assert_allclose( + cuik_molmaker.mol_featurizer( + "[OH:2][CH3:1]", *feature_arrays(), False, False, False, False, False + )[0], + atom_features("OC"), + ) + + +@pytest.mark.parametrize("ordered", [True, False]) +def test_batch_featurizer_honors_ordered(ordered): + expected = "CO" if ordered else "OC" + np.testing.assert_allclose( + batch_atom_features(["[OH:2][CH3:1]"], ordered=ordered), + atom_features(expected), + ) diff --git a/tests/python/test_parse_options.py b/tests/python/test_parse_options.py new file mode 100644 index 0000000..821dedc --- /dev/null +++ b/tests/python/test_parse_options.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # noqa: E501 +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the `keep_h` and `ignore_stereo` parsing options on +`mol_featurizer` / `batch_mol_featurizer`, and `ignore_stereo` on +`batch_reaction_featurizer`. +""" + +import numpy as np +import pytest + +import cuik_molmaker + + +def featurize( + smiles, + explicit_H=False, + keep_h=False, + ignore_stereo=False, + ordered=True, + atom_onehot=("atomic-number-common",), + atom_float=(), + bond=(), +): + """mol_featurizer output for `smiles`, with the new options exposed as + keyword arguments and everything else defaulted.""" + oh = cuik_molmaker.atom_onehot_feature_names_to_array(list(atom_onehot)) + fl = cuik_molmaker.atom_float_feature_names_to_array(list(atom_float)) + bp = cuik_molmaker.bond_feature_names_to_array(list(bond)) + return cuik_molmaker.mol_featurizer( + smiles, + oh, + fl, + bp, + explicit_H=explicit_H, + offset_carbon=False, + duplicate_edges=False, + add_self_loop=False, + ordered=ordered, + keep_h=keep_h, + ignore_stereo=ignore_stereo, + ) + + +# --------------------------------------------------------------------------- +# keep_h +# --------------------------------------------------------------------------- + + +def test_keep_h_defaults_to_false(): + """Without keep_h, explicit [H] atoms written in the SMILES are folded + back into implicit valence on the heavy atoms (C, O), leaving 2 atoms.""" + feats = featurize("[H]C([H])([H])O") + assert feats[0].shape[0] == 2 + + +def test_keep_h_true_retains_explicit_atoms(): + """With keep_h, the 3 explicit [H] on carbon are kept as real atoms; the + implicit O-H is not (nothing adds it), giving 5 atoms total.""" + feats = featurize("[H]C([H])([H])O", keep_h=True) + assert feats[0].shape[0] == 5 + + +def test_keep_h_and_explicit_H_converge(): + """Whether the SMILES starts with some explicit Hs (keep_h) or none, once + explicit_H adds every remaining implicit H, methanol always has 6 atoms.""" + default_then_add = featurize("[H]C([H])([H])O", explicit_H=True) + keep_then_add = featurize("[H]C([H])([H])O", explicit_H=True, keep_h=True) + assert default_then_add[0].shape[0] == 6 + assert keep_then_add[0].shape[0] == 6 + + +def test_batch_mol_featurizer_honors_keep_h(): + oh = cuik_molmaker.atom_onehot_feature_names_to_array(["atomic-number-common"]) + fl = cuik_molmaker.atom_float_feature_names_to_array([]) + bp = cuik_molmaker.bond_feature_names_to_array([]) + feats = cuik_molmaker.batch_mol_featurizer( + ["[H]C([H])([H])O"], + oh, + fl, + bp, + explicit_H=False, + offset_carbon=False, + duplicate_edges=False, + add_self_loop=False, + keep_h=True, + ) + assert feats[0].shape[0] == 5 + + +# --------------------------------------------------------------------------- +# ignore_stereo +# --------------------------------------------------------------------------- + + +def chirality_features(smiles, ignore_stereo): + return featurize( + smiles, + atom_onehot=("chirality",), + atom_float=("chirality",), + ignore_stereo=ignore_stereo, + )[0] + + +def test_ignore_stereo_clears_onehot_and_float_chirality(): + """The stereocenter is atom index 1 (N, C@@H, CH3, C(=O), O, O). Clearing + stereo must change BOTH the one-hot chirality block (reads GetChiralTag) + and the float chirality feature (reads the _CIPCode property) -- this is + exactly the gap a manual port of Chemprop's SetChiralTag-only loop would + miss, which is why ignore_stereo is implemented with + RDKit::MolOps::removeStereochemistry instead.""" + with_stereo = chirality_features("N[C@@H](C)C(=O)O", ignore_stereo=False) + without_stereo = chirality_features("N[C@@H](C)C(=O)O", ignore_stereo=True) + plain = chirality_features("NC(C)C(=O)O", ignore_stereo=False) + + assert not np.allclose(with_stereo[1], without_stereo[1]) + np.testing.assert_allclose(without_stereo, plain) + + +def bond_stereo_features(smiles, ignore_stereo): + return featurize(smiles, bond=("stereo",), ignore_stereo=ignore_stereo)[1] + + +def test_ignore_stereo_clears_bond_stereo(): + """cis and trans difluoroethylene must differ in bond features by + default, and become identical (and match the unspecified form) once + ignore_stereo strips E/Z information.""" + cis = bond_stereo_features(r"F/C=C\F", ignore_stereo=False) + trans = bond_stereo_features("F/C=C/F", ignore_stereo=False) + cis_cleared = bond_stereo_features(r"F/C=C\F", ignore_stereo=True) + trans_cleared = bond_stereo_features("F/C=C/F", ignore_stereo=True) + plain = bond_stereo_features("FC=CF", ignore_stereo=False) + + 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])