datamol-io / datamol Goto Github PK

def calc_ap(mol: dm.Mol):
    """Calculate aromatic proportion #aromatic atoms/#atoms total."""
    aromatic_query = dm.from_smarts("a")
    matches = mol.GetSubstructMatches(aromatic_query)
    return len(matches) / mol.GetNumAtoms()


def compute_esol(data: pd.DataFrame):
    """Compute esol from
    https://github.com/PatWalters/solubility/blob/d1536c58afe5e0e7ac4c96e2ffef496d5b98664b/esol.py

    Dataframe must have the following columns: clogp, mw, n_rotatable_bonds
    """

    intercept = 0.26121066137801696
    coef = {
        "mw": -0.0066138847738667125,
        "clogp": -0.7416739523408995,
        "n_rotatable_bonds": 0.003451545565957996,
        "ap": -0.42624840441316975,
    }

    # Compute aromatic proportion
    aromatic_proportion = data["mol"].apply(calc_ap)

    # Compute ESOL
    esol = (
        intercept
        + coef["clogp"] * data["clogp"]
        + coef["mw"] * data["mw"]
        + coef["n_rotatable_bonds"] * data["n_rotatable_bonds"]
        + coef["ap"] * aromatic_proportion
    )

    return esol

from rdkit import Chem
from rdkit.Chem import rdMolDescriptors
from rdkit.Chem import DataStructs

import numpy as np

mol = Chem.MolFromSmiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")

fp_vec = rdMolDescriptors.GetMorganFingerprintAsBitVect(
    mol,
    radius=2,
    nBits=2048,
    useFeatures=True,
)

%%timeit -n 1000 -r 1000
# 10.1 µs ± 3.62 µs per loop (mean ± std. dev. of 1000 runs, 1000 loops each)
np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0')

%%timeit -n 1000 -r 100
fp = np.zeros((0,), dtype=int)
39.6 µs ± 1.45 µs per loop (mean ± std. dev. of 100 runs, 1000 loops each)
DataStructs.ConvertToNumpyArray(fp_vec, fp)

rows = data.iterrows()
rows = list(rows)   # memory and CPU intensive
data = dm.parallelized(process_fn, rows, n_jobs=-1, progress=True, arg_type="args")

data = dm.parallelized(process_fn, data.iterrows(), n_jobs=-1, progress=True, arg_type="args", total=len(data))

smiles = "CCCC"
mol = dm.to_mol(smiles)
dm.conformers.generate(mol, n_confs=1, minimize_energy=True, rms_cutoff=None)

import datamol as dm
from rdkit import Chem


def _in_range(x, min_val: float = -float("inf"), max_val: float = float("inf")):
    """Check if a value is in a range
    Args:
        x: value to check
        min_val: minimum value
        max_val: maximum value
    """
    return min_val <= x <= max_val


def _compute_ring_system(mol: dm.Mol, include_spiro: bool = True):
    """
    Compute the list of ring system in a molecule. This is based on RDKit's cookbook:
    https://www.rdkit.org/docs/Cookbook.html#rings-aromaticity-and-kekulization

    # EN: move to datamol

    Args:
        mol: input molecule
        include_spiro: whether to include spiro rings. Defaults to False.

    Returns:
        ring_system: list of ring system
    """
    ri = mol.GetRingInfo()
    systems = []
    for ring in ri.AtomRings():
        ringAts = set(ring)
        nSystems = []
        for system in systems:
            nInCommon = len(ringAts.intersection(system))
            if nInCommon and (include_spiro or nInCommon > 1):
                ringAts = ringAts.union(system)
            else:
                nSystems.append(system)
        nSystems.append(ringAts)
        systems = nSystems
    return systems


def _compute_charge(mol: dm.Mol):
    """
    Compute the charge of a molecule.

    Args:
        mol: input molecule

    Returns:
        charge: charge of the molecule
    """
    return Chem.rdmolops.GetFormalCharge(mol)


def _compute_refractivity(mol: dm.Mol):
    """
    Compute the refractivity of a molecule.

    Args:
        mol: input molecule

    Returns:
        mr: molecular refractivity of the molecule
    """
    return Chem.Crippen.MolMR(mol)


def _compute_rigid_bonds(mol: dm.Mol):
    """
    Compute the number of rigid bonds in a molecule.

    Args:
        mol: input molecule

    Returns:
        n_rigid_bonds: number of rigid bonds in the molecule
    """
    # rigid bonds are bonds that are not (single and not in rings)
    # I don't think this is the same as the number of non rotatable bonds ?
    # EN: move this to datamol ?
    non_rigid_bonds_count = dm.from_smarts("*-&!@*")
    n_rigid_bonds = mol.GetNumBonds() - len(
        mol.GetSubstructMatches(non_rigid_bonds_count)
    )
    return n_rigid_bonds


def _compute_n_stereo_center(mol: dm.Mol):
    """
    Compute the number of stereocenters in a molecule.

    Args:
        mol: input molecule

    Returns:
        n_stero_center: number of stereocenters in the molecule
    """
    n_stereo_center = 0
    try:
        Chem.FindPotentialStereo(mol, cleanIt=False)  # type: ignore
        n_stereo_center = Chem.rdMolDescriptors.CalcNumAtomStereoCenters(mol)
    except:
        pass
    return n_stereo_center


def _compute_n_charged_atoms(mol: dm.Mol):
    """
    Compute the number of charged atoms in a molecule.

    Args:
        mol: input molecule

    Returns:
        n_charged_atoms: number of charged atoms in the molecule
    """
    return sum([at.GetFormalCharge() != 0 for at in mol.GetAtoms()])

import datamol as dm

from rdkit import Chem

# SMILES: "N=C(N)O"
inchi3 = "InChI=1/CH4N2O/c2-1(3)4/h(H4,2,3,4)/f/h2,4H,3H2/b2-1?"

# SMILES
inchi4 = "InChI=1/CH4N2O/c2-1(3)4/h(H4,2,3,4)/f/h2-3H2"

mol3 = dm.from_inchi(inchi3)
mol4 = dm.from_inchi(inchi4)

print(dm.to_smiles(mol3))  # N=C(N)O
print(dm.to_smiles(mol4))  # NC(N)=O

# generated inchi and inchikey remain the same since the hydrogen layer has been removed

print(dm.to_inchi(mol3))  # 'InChI=1S/CH4N2O/c2-1(3)4/h(H4,2,3,4)'
print(dm.to_inchi(mol4))  # 'InChI=1S/CH4N2O/c2-1(3)4/h(H4,2,3,4)'

print(dm.to_inchikey(mol3))  # 'XSQUKJJJFZCRTK-UHFFFAOYSA-N'
print(dm.to_inchikey(mol4))  # 'XSQUKJJJFZCRTK-UHFFFAOYSA-N'


# when computing the inchikey directly from the non-standard inchi with the hydrogen layer
# those are different

print(Chem.inchi.InchiToInchiKey(inchi3))  # XSQUKJJJFZCRTK-ZIALIONUNA-N
print(Chem.inchi.InchiToInchiKey(inchi4))  # XSQUKJJJFZCRTK-UBUOBULFNA-N

def sim_search(
    mols: list[dm.Mol],
    qmol: list[dm.Mol],
    threshold: float = .7,
    fp_type: str = "ffp2"
 ) -> list[Tuple[dm.Mol, float]]:
    pass

def substruct_search(
    mols: list[dm.Mol],
    qmol: list[dm.Mol],
    atoms_and_bonds: bool = False,
 ) -> Union[list[dm.Mol], list[Tuple[dm.Mol, list[int], list[int]]]:
    pass

def exact_search(
    mols: list[dm.Mol],
    qmol: list[dm.Mol],
    unique: bool = True
 ) -> Union[dm.Mol, list[dm.Mol]]:
    pass

from typing import Callable
from typing import Sequence
from typing import Any
from typing import Optional

from concurrent import futures
from loguru import logger


def parallelized(
    fn: Callable,
    inputs_list: Sequence[Any],
    scheduler: str = "processes",
    n_jobs: Optional[int] = -1,
    worker_id_key: Any = None,
    progress: bool = True,
    progress_auto: bool = True,
    log_when_done: bool = False,
):

    if not isinstance(inputs_list[0], dict):
        raise ValueError(f"inputs_list elements must be a dict not {type(inputs_list[0])}")

    if progress_auto:
        from tqdm.auto import tqdm
    else:
        from tqdm import tqdm

    if scheduler == "processes":
        pool_executor_cls = futures.ProcessPoolExecutor
        executor_kwargs = dict(max_workers=n_jobs)
        as_completed_fn = futures.as_completed

    elif scheduler == "threads":
        pool_executor_cls = futures.ThreadPoolExecutor
        executor_kwargs = dict(max_workers=n_jobs)
        as_completed_fn = futures.as_completed

    else:
        raise ValueError(f"Wrong scheduler: {scheduler}")

    if worker_id_key:
        worker_id_fn = lambda x: x[worker_id_key]
    else:
        worker_id_fn = lambda x: x

    with pool_executor_cls(**executor_kwargs) as executor:

        futures_list = {executor.submit(fn, **kwargs): kwargs for kwargs in inputs_list}

        results = []
        tqdm_args = {}
        tqdm_args["total"] = len(inputs_list)
        tqdm_args["disable"] = not progress

        for i, future in enumerate(tqdm(as_completed_fn(futures_list), **tqdm_args)):
            kwargs = futures_list[future]
            try:
                result = future.result()
                results.append(result)
            except Exception as ex:
                logger.info(f"Error running the following worker '{worker_id_fn(kwargs)}':", exc_info=1)  # type: ignore
                logger.info("Shutting down the execution...")
                executor.shutdown(True)
                raise ex
            else:
                if log_when_done:
                    logger.info(f"Execution done for '{worker_id_fn(kwargs)}' ({i+1}/{len(inputs_list)})")

    return results