def bookkeep_reactant()

in python/dgllife/data/uspto.py [0:0]


def bookkeep_reactant(mol, candidate_pairs):
    """Bookkeep reaction-related information of reactants.

    Parameters
    ----------
    mol : rdkit.Chem.rdchem.Mol
        RDKit molecule instance for reactants.
    candidate_pairs : list of 2-tuples
        Pairs of atoms that ranked high by a model for reaction center prediction.
        By assumption, the two atoms are different and the first atom has a smaller
        index than the second.

    Returns
    -------
    info : dict
        Reaction-related information of reactants
    """
    num_atoms = mol.GetNumAtoms()
    info = {
        # free valence of atoms
        'free_val': [0 for _ in range(num_atoms)],
        # Whether it is a carbon atom
        'is_c': [False for _ in range(num_atoms)],
        # Whether it is a carbon atom connected to a nitrogen atom in pyridine
        'is_c2_of_pyridine': [False for _ in range(num_atoms)],
        # Whether it is a phosphorous atom
        'is_p': [False for _ in range(num_atoms)],
        # Whether it is a sulfur atom
        'is_s': [False for _ in range(num_atoms)],
        # Whether it is an oxygen atom
        'is_o': [False for _ in range(num_atoms)],
        # Whether it is a nitrogen atom
        'is_n': [False for _ in range(num_atoms)],
        'pair_to_bond_val': dict(),
        'ring_bonds': set()
    }

    # bookkeep atoms
    for j, atom in enumerate(mol.GetAtoms()):
        info['free_val'][j] += atom.GetTotalNumHs() + abs(atom.GetFormalCharge())
        # An aromatic carbon atom next to an aromatic nitrogen atom can get a
        # carbonyl b/c of bookkeeping of hydroxypyridines
        if atom.GetSymbol() == 'C':
            info['is_c'][j] = True
            if atom.GetIsAromatic():
                for nbr in atom.GetNeighbors():
                    if nbr.GetSymbol() == 'N' and nbr.GetDegree() == 2:
                        info['is_c2_of_pyridine'][j] = True
                        break
        # A nitrogen atom should be allowed to become positively charged
        elif atom.GetSymbol() == 'N':
            info['free_val'][j] += 1 - atom.GetFormalCharge()
            info['is_n'][j] = True
        # Phosphorous atoms can form a phosphonium
        elif atom.GetSymbol() == 'P':
            info['free_val'][j] += 1 - atom.GetFormalCharge()
            info['is_p'][j] = True
        elif atom.GetSymbol() == 'O':
            info['is_o'][j] = True
        elif atom.GetSymbol() == 'S':
            info['is_s'][j] = True

    # bookkeep bonds
    for bond in mol.GetBonds():
        atom1, atom2 = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        atom1, atom2 = min(atom1, atom2), max(atom1, atom2)
        type_val = bond.GetBondTypeAsDouble()
        info['pair_to_bond_val'][(atom1, atom2)] = type_val
        if (atom1, atom2) in candidate_pairs:
            info['free_val'][atom1] += type_val
            info['free_val'][atom2] += type_val
        if bond.IsInRing():
            info['ring_bonds'].add((atom1, atom2))

    return info