Skip to content

Accept 1-based atom maps and add featurizer keyword options #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
craabreu wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Accept 1-based atom maps and add featurizer keyword options #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8e962ad
Accept 1-based atom maps for explicit atom ordering
craabreu Jul 21, 2026
607e430
Expose the `ordered` flag in the Python API
craabreu Jul 21, 2026
1bfc66f
Add keep_h and ignore_stereo to mol_featurizer / batch_mol_featurizer
craabreu Jul 21, 2026
44792c0
Add ignore_stereo to batch_reaction_featurizer
craabreu Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion docs/USAGE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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).
Expand Down
64 changes: 53 additions & 11 deletions src/cuik_molmaker_cpp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
Expand All @@ -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,
Expand All @@ -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);
}
61 changes: 47 additions & 14 deletions src/features.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<RDKit::RWMol> 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<RDKit::RWMol> mol{parse_mol(smiles_string, explicit_H, ordered, keep_h, ignore_stereo)};

if (!mol) {
return GraphData{0, std::unique_ptr<CompactAtom[]>(), 0, std::unique_ptr<CompactBond[]>(), std::move(mol)};
Expand Down Expand Up @@ -519,15 +523,21 @@ std::vector<py::array> 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,
bond_property_list,
explicit_H,
offset_carbon,
duplicate_edges,
add_self_loop);
add_self_loop,
ordered,
keep_h,
ignore_stereo);
}

std::vector<py::array> batch_mol_featurizer(const std::vector<std::string>& smiles_list,
Expand All @@ -537,7 +547,10 @@ std::vector<py::array> batch_mol_featurizer(const std::vector<std::string>& 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
Expand All @@ -547,7 +560,7 @@ std::vector<py::array> batch_mol_featurizer(const std::vector<std::string>& 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));
Expand Down Expand Up @@ -680,9 +693,14 @@ std::vector<py::array> batch_mol_featurizer(const std::vector<std::string>& smil

// Creates an RWMol from a SMILES string.
// See the declaration in features.h for more details.
std::unique_ptr<RDKit::RWMol> parse_mol(const std::string& smiles_string, bool explicit_H, bool ordered) {
// Parse SMILES string with default options
RDKit::SmilesParserParams params;
std::unique_ptr<RDKit::RWMol> 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<RDKit::RWMol> mol{RDKit::SmilesToMol(smiles_string, params)};
if (!mol) {
return mol;
Expand All @@ -694,6 +712,8 @@ std::unique_ptr<RDKit::RWMol> parse_mol(const std::string& smiles_string, bool e
// if they indicate a valid order.
const unsigned int num_atoms = mol->getNumAtoms();
std::vector<unsigned int> atom_order(num_atoms);
unsigned int map_num_min = std::numeric_limits<unsigned int>::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)) {
Expand All @@ -702,10 +722,11 @@ std::unique_ptr<RDKit::RWMol> 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
Expand All @@ -714,13 +735,18 @@ std::unique_ptr<RDKit::RWMol> 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<unsigned int>::max();
std::vector<unsigned int> 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;
Expand All @@ -742,5 +768,12 @@ std::unique_ptr<RDKit::RWMol> 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;
}
Loading