electrum_atomic.py

import numpy as np
import pandas as pd
import hashlib

from rdkit import Chem
from rdkit.Chem import rdmolops

metals_dict = {
    'H': 1,
    'He': 2,
    'Li': 3,
    'Be': 4,
    'B': 5,
    'C': 6,
    'N': 7,
    'O': 8,
    'F': 9,
    'Ne': 10,
    'Na': 11,
    'Mg': 12,
    'Al': 13,
    'Si': 14,
    'P': 15,
    'S': 16,
    'Cl': 17,
    'Ar': 18,
    'K': 19,
    'Ca': 20,
    'Sc': 21,
    'Ti': 22,
    'V': 23,
    'Cr': 24,
    'Mn': 25,
    'Fe': 26,
    'Co': 27,
    'Ni': 28,
    'Cu': 29,
    'Zn': 30,
    'Ga': 31,
    'Ge': 32,
    'As': 33,
    'Se': 34,
    'Br': 35,
    'Kr': 36,
    'Rb': 37,
    'Sr': 38,
    'Y': 39,
    'Zr': 40,
    'Nb': 41,
    'Mo': 42,
    'Tc': 43,
    'Ru': 44,
    'Rh': 45,
    'Pd': 46,
    'Ag': 47,
    'Cd': 48,
    'In': 49,
    'Sn': 50,
    'Sb': 51,
    'Te': 52,
    'I': 53,
    'Xe': 54,
    'Cs': 55,
    'Ba': 56,
    'La': 57,
    'Ce': 58,
    'Pr': 59,
    'Nd': 60,
    'Pm': 61,
    'Sm': 62,
    'Eu': 63,
    'Gd': 64,
    'Tb': 65,
    'Dy': 66,
    'Ho': 67,
    'Er': 68,
    'Tm': 69,
    'Yb': 70,
    'Lu': 71,
    'Hf': 72,
    'Ta': 73,
    'W': 74,
    'Re': 75,
    'Os': 76,
    'Ir': 77,
    'Pt': 78,
    'Au': 79,
    'Hg': 80,
    'Tl': 81,
    'Pb': 82,
    'Bi': 83,
}

def get_atom_env(mol, radius:int, atom:int) -> str:
    """
    Extracts the local chemical environment around a specified atom within a molecular structure.

    Parameters:
    ----------
    mol (RDKit Mol): The molecular structure from which the atom environment will be extracted.
    radius: The radius of the desired atom environment. Defines the number of bonds away from the central atom to be considered in the local environment.
    atom: The index of the target atom in the molecular structure. The function will extract the environment around this specific atom.

    Returns:
    ----------
    smiles: The SMILES representation of the identified substructure.

    Example Usage:
    ----------
    ```python
    mol = Chem.MolFromSmiles('C1CC(=O)NC(=O)[C@@H]1N2C(=O)C3=CC=CC=C3C2=O')
    substructure = get_atom_env(mol, radius=2, atom_index=1)
    print(substructure)
    ```
    """
    env = Chem.FindAtomEnvironmentOfRadiusN(mol, radius, atom)
    atom_map = {}
    mol = Chem.PathToSubmol(mol, env, atomMap=atom_map)
    smiles = Chem.MolToSmiles(mol, isomericSmiles=False, canonical=True)
    return smiles

def get_mol_substructs(mol, radius:int) -> list:
    """
    Generates a list of unique substructures within a molecular structure, considering atoms within a specified radius.

    Parameters:
    ----------
    mol (RDKit Mol): The molecular structure from which the substructures will be extracted.
    radius (int): The maximum radius to consider for the extraction of substructures. Defines the number of bonds away from each atom to be considered in the local environment.

    Returns:
    ----------
    list: A list of unique SMILES representations of the identified substructures.

    Example Usage:
    ----------
    ```python
    mol = Chem.MolFromSmiles('C1CC(=O)NC(=O)[C@@H]1N2C(=O)C3=CC=CC=C3C2=O')
    substructures = get_mol_substructs(mol, radius=2)
    print(substructures)
    ```
    """
    substructs = []
    for r in range(1, radius+1):
        for a in range(mol.GetNumAtoms()):
            substructs.append(get_atom_env(mol, r, a))
    return list(set(substructs))

def hash_and_fold(mol, radius:int, n_bits:int, return_mapping=False) -> np.ndarray or tuple:
    """
    Generates a folded fingerprint for a molecular structure by hashing substructures.

    Parameters:
    ----------
    mol (RDKit Mol): The molecular structure for which the fingerprint will be generated.
    radius (int): The maximum radius to consider for the extraction of substructures. Defines the number of bonds away from each atom to be considered in the local environment.
    n_bits (int): The number of bits in the fingerprint.
    return_mapping (bool, optional): Whether to return a mapping of hash values to substructures. Defaults to False.

    Returns:
    ----------
    np.ndarray or tuple: If `return_mapping` is False, returns a binary fingerprint represented as a numpy array. If `return_mapping` is True, returns a tuple containing the fingerprint and a dictionary mapping hash values to substructures.

    Example Usage:
    ----------
    ```python
    mol = Chem.MolFromSmiles('C1CC(=O)NC(=O)[C@@H]1N2C(=O)C3=CC=CC=C3C2=O')
    fingerprint = hash_and_fold(mol, radius=2, n_bits=1024)
    print(fingerprint)
    ```
    """
    substructs = get_mol_substructs(mol, radius)
    unique_hashes = set()
    hash_to_position = {}

    for s in substructs:
        hash_val = int(hashlib.sha1(s.encode('utf-8')).hexdigest(), 16) % n_bits
        unique_hashes.add(hash_val)
        hash_to_position[hash_val] = s

    fp = np.zeros(n_bits, dtype=int)
    fp[list(unique_hashes)] = 1

    if return_mapping:
        return fp, hash_to_position
    else:
        return fp

def process_smiles(smiles:str) -> list:
    """
    Processes a string containing multiple SMILES strings separated by '.' and returns a list of RDKit Mol objects.

    Parameters:
    ----------
    smiles (str): A string containing one or more SMILES strings separated by '.'.

    Returns:
    ----------
    list: A list of RDKit Mol objects corresponding to the input SMILES strings.

    Example Usage:
    ----------
    ```python
    smiles = 'CCO.CN.CC'
    mol_list = process_smiles(smiles)
    for mol in mol_list:
        print(Chem.MolToSmiles(mol))
    ```
    """
    smiles_list = smiles.split('.')
    return [Chem.MolFromSmiles(s, sanitize=False) for s in smiles_list]

def calculate_fingerprint(smiles:str, metal:str, radius:int, n_bits:int, return_mapping=False) -> np.ndarray or tuple:
    """
    Calculates a combined fingerprint for a set of SMILES strings with a specified metal ion.

    Parameters:
    ----------
    smiles (str): A string containing one or more SMILES strings separated by '.'.
    metal (str): The SMILES representation of the metal ion to be included in the fingerprint.
    radius (int): The maximum radius to consider for the extraction of substructures. Defines the number of bonds away from each atom to be considered in the local environment.
    n_bits (int): The number of bits in the fingerprint.
    return_mapping (bool, optional): Whether to return a mapping of hash values to substructures. Defaults to False.

    Returns:
    ----------
    np.ndarray or tuple: If `return_mapping` is False, returns a combined binary fingerprint represented as a numpy array. If `return_mapping` is True, returns a tuple containing the fingerprint and a dictionary mapping hash values to substructures.

    Example Usage:
    ----------
    ```python
    smiles = 'CCO.CN.CC'
    metal = 'Fe'
    fingerprint = calculate_fingerprint(smiles, metal, radius=2, n_bits=1024)
    print(fingerprint)
    ```
    """
    mol_list = process_smiles(smiles)
    fp_list = [hash_and_fold(m, radius, n_bits, return_mapping=return_mapping) for m in mol_list]

    if return_mapping:
        combined_mapping = {}

        for idx, (_, substructure_mapping) in enumerate(fp_list):
            for substruct_idx, substruct_smiles in substructure_mapping.items():
                combined_mapping[substruct_idx + idx * n_bits] = substruct_smiles

        fingerprint = np.append(np.sum([fps for fps, _ in fp_list], axis=0), metals_dict[metal])
        adjusted_mapping = {pos % n_bits: smiles for pos, smiles in combined_mapping.items()}

        return fingerprint, adjusted_mapping
    else:
        fingerprint = np.append(np.sum(fp_list, axis=0), metals_dict[metal])
        return fingerprint

def calculate_fingerprints(smiles_list:list, metals_list:list, radius:int, n_bits:int, return_mapping=False) -> np.ndarray or tuple:
    """
    Calculates fingerprints for a list of SMILES strings with corresponding metal ions.

    Parameters:
    ----------
    smiles_list (list): A list of strings, each containing one or more SMILES strings separated by '.'.
    metals_list (list): A list of strings, each containing the SMILES representation of the metal ion corresponding to the SMILES string in `smiles_list`.
    radius (int): The maximum radius to consider for the extraction of substructures. Defines the number of bonds away from each atom to be considered in the local environment.
    n_bits (int): The number of bits in each fingerprint.
    return_mapping (bool, optional): Whether to return a mapping of hash values to substructures. Defaults to False.

    Returns:
    ----------
    np.ndarray or tuple: If `return_mapping` is False, returns a list of combined binary fingerprints represented as numpy arrays. If `return_mapping` is True, returns a tuple containing the list of fingerprints and a dictionary mapping hash values to substructures.

    Example Usage:
    ----------
    ```python
    smiles_list = ['CCO.CN.CC', 'CCC.O=O.CN']
    metals_list = ['[Fe]', '[Cu]']
    fingerprints = calculate_fingerprints(smiles_list, metals_list, radius=2, n_bits=1024)
    print(fingerprints)
    ```
    """
    fingerprints = []
    combined_mapping = {}

    for smiles, metal in zip(smiles_list, metals_list):
        fingerprint, substruct_mapping = calculate_fingerprint(smiles, metal, radius, n_bits, return_mapping=True)

        if return_mapping:
            combined_mapping.update(substruct_mapping)

        fingerprints.append(fingerprint)

    if return_mapping:
        adjusted_mapping = {pos % n_bits: smiles for pos, smiles in combined_mapping.items()}
        return fingerprints, adjusted_mapping
    else:
        return fingerprints