Valence correction by adding formal charge in RDKit

Question

I am modeling molecules as graphs, with nodes as atom types and edges as bond types (single, double, triple). I do not model formal charges or hydrogens explicitly (only heavy atoms). When trying to recover valid RDKit molecules/smiles from the graphs, I fail at molecules with charged atoms. My chemical knowledge is very rusty (pun intended), but I feel like it should be possible to add formal charge where needed using posthoc valence checks when translating graphs to molecules.

Following the recommendations on RDKit's FAQ, I am able to handle cases like C[N+](C)(C)C (detect nitrogen has a higher valence, 4, than it should, 3, and add a positive charge). Right now I struggle with negative charges. When translating this molecule COC1=CC=C(CCNC(=O)C2=CC=C(OC)C([N+](=O)[O-])=C2)C=C1 to a graph then back to smiles using only positive charge correction, I get this output COC1=CC=C(CCNC(=O)C2=CC=C(OC)C([N+](=O)O)=C2)C=C1 (notice the missing negative charge in the last oxygen. RDKit automatically replaces it with an implicit hydrogen).

Are there settings/methods within RDKit that I can use to ensure the correct valences? I'm hoping there is a way to turn off setting implicit hydrogens and/or automatically assign formal charges.

Any pointers would be really helpful.

Here is pseudo-code to explain what I am looking for:

# input smiles
smi = 'NC1=CC(Cl)=CC=C1[N+](=O)[O-]'

# my own modeling of molecules 
graph = get_atom_bond_graph(smi) 

# returns an rwmol object from the graph
mol = get_mol_from_graph(graph)

# smi2 is missing all formal charges => cannot be sanitized because the valence of N is off
smi2 = Chem.MolToSmiles(mol) # NC1=CC(Cl)=CC=C1N(=O)O  

# add missing formal charges by checking valence rules 
correct_smiles(mol)

smi3 = Chem.MolToSmiles(mol) # NC1=CC(Cl)=CC=C1[N+](=O)O'

# ... how to get [0-] with posthoc checks?

Here is my code for generating graphs in case someone needs it (I am using PyTorch Geometric for creating graphs):

atom_types =  ['Si', 'P', 'N', 'Mg', 'Se', 'Cu', 'S', 'Br', 'B', 'O', 'C', 'Zn', 'Sn', 'F', 'I', 'Cl']

atom_type_offset = 1 # where to start atom type indexing. 1 if modeling no nodes, 0 otherwise

# starting from 1 because considering 0 an edge type (= no edge)
bond_types = [BT.SINGLE, BT.DOUBLE, BT.TRIPLE]
bond_type_offset = 1

def get_mol_graph(mol, offset=0):
    if type(mol)==str: 
        mol = Chem.MolFromSmiles(mol)
    Chem.Kekulize(mol, clearAromaticFlags=True)
    Chem.RemoveStereochemistry(mol)
    
    m_nodes = get_mol_nodes(mol=mol)
    m_edge_index, m_edge_attr = get_mol_edges(mol=mol, offset=offset)

    return m_nodes, m_edge_index, m_edge_attr

def get_mol_nodes(mol):
    atoms = mol.GetAtoms()

    for i, atom in enumerate(atoms):
        s = atom.GetSymbol()
        atom_type = torch.tensor([atom_types.index(s)+atom_type_offset], 
                                  dtype=torch.long) # needs to be int for one hot
        atom_types_ = torch.cat((atom_types_, atom_type), dim=0) if i > 0 else atom_type
    
    atom_feats = F.one_hot(atom_types_, num_classes=len(atom_types)+atom_type_offset).float()

    return atom_feats

def get_mol_edges(mol, offset=1):
    '''
        Input:
            offset (optional): default: 1. To account for 'no bond' type.
    '''
    for i, b in enumerate(mol.GetBonds()):
        beg_atom_idx = b.GetBeginAtom().GetIdx()
        end_atom_idx = b.GetEndAtom().GetIdx()

        e_beg = torch.tensor([beg_atom_idx+offset, end_atom_idx+offset], dtype=torch.long).unsqueeze(-1)
        e_end = torch.tensor([end_atom_idx+offset, beg_atom_idx+offset], dtype=torch.long).unsqueeze(-1)
        e_type = torch.tensor([bond_types.index(b.GetBondType())+bond_type_offset, 
                               bond_types.index(b.GetBondType())+bond_type_offset], dtype=torch.long) # needs to be int for one hot

        begs = torch.cat((begs, e_beg), dim=0) if i > 0 else e_beg
        ends = torch.cat((ends, e_end), dim=0) if i > 0 else e_end
        edge_type = torch.cat((edge_type, e_type), dim=0) if i > 0 else e_type

    if len(mol.GetBonds())>0:
        edge_index = torch.cat((begs, ends), dim=1).mT.contiguous()
        edge_attr = F.one_hot(edge_type, num_classes=len(bond_types)+bond_type_offset).float() # add 1 to len of bonds to account for no edge
    else: # handle molecules of single atoms
        edge_index = torch.tensor([]).long().reshape(2,0)
        edge_attr = torch.tensor([]).float()

    return edge_index, edge_attr

if __name__ == '__main__':
    smi = 'c1c(Cl)c(Cl)ccc1'
    nodes, edge_index, edge_attr = get_mol_graph(smi.strip(), offset=0)

    from torch_geometric.data import Data

    data = Data(x=nodes, edge_index=edge_index, edge_attr=edge_attr)

Answer 1

After some investigation, I don't think it's possible to add formal charge posthoc without modeling hydrogens explicitly. Here is an example to illustrate:

Consider both molecules C[Mg+] and C[MgH]. My graph only knows that both molecules contain one carbon and one magnesium. When turning the graph to a molecule, RDKit notices that the valence of both atoms is lower than it should be, and adds explicit hydrogens to both. This means that whether my input is C[Mg+] or C[MgH] I would always recover C[MgH] from my graph.

Other examples illustrating the issue include: [NH4+] vs N, and COC[P+](C1=CC=CC=C1)(C1=CC=CC=C1)C1=CC=CC=C1 vs COC[PH](C1=CC=CC=C1)(C1=CC=CC=C1)C1=CC=CC=C1.

I am still trying to understand how implicit hydrogens are added to molecules in general to see if there is away to differentiate between cases needing a formal charge vs implicit hydrogen. For now distinguishing the two does not seem to be possible.