Yn patch0411 (#339)

* [feat] add params `save_sdf`: save conformers to sdf

* [style] black --skip-string-normalization

* [docs] add the annotations

* [feat] save_sdf : optional [if_not_exists, always, never]

* [docs] update the annotations

* [feat] update save_sdf judge order

* [feat] remove `save_sdf`, add `conf_cache_level` for saveing conformers to sdf files. `conf_cache_level` optional [0, 1, 2]

* Apply suggestions from code review

Revise several grammatical error by copilot.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: emotionor

unimol_tools/setup.py

@@ -6,7 +6,9 @@
 setup(
     name="unimol_tools",
     version="0.1.3.post1",
-    description=("unimol_tools is a Python package for property prediction with Uni-Mol in molecule, materials and protein."),
+    description=(
+        "unimol_tools is a Python package for property prediction with Uni-Mol in molecule, materials and protein."
+    ),
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',
     author="DP Technology",
@@ -20,16 +22,18 @@
             "dist",
         ],
     ),
-    install_requires=["numpy<2.0.0,>=1.22.4",
-                      "pandas<2.0.0",
-                      "torch",
-                      "joblib",
-                      "rdkit",
-                      "pyyaml",
-                      "addict",
-                      "scikit-learn",
-                      "numba",
-                      "tqdm"],
+    install_requires=[
+        "numpy<2.0.0,>=1.22.4",
+        "pandas<2.0.0",
+        "torch",
+        "joblib",
+        "rdkit",
+        "pyyaml",
+        "addict",
+        "scikit-learn",
+        "numba",
+        "tqdm",
+    ],
     python_requires=">=3.6",
     include_package_data=True,
     classifiers=[
@@ -43,4 +47,4 @@
         "Programming Language :: Python :: 3.10",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
     ],
-)
+)

unimol_tools/unimol_tools/data/conformer.py

@@ -116,16 +116,17 @@ def single_process(self, smiles):
         :raises ValueError: If the conformer generation method is unrecognized.
         """
         if self.method == 'rdkit_random':
-            atoms, coordinates = inner_smi2coords(
+            atoms, coordinates, mol = inner_smi2coords(
                 smiles, seed=self.seed, mode=self.mode, remove_hs=self.remove_hs
             )
-            return coords2unimol(
+            feat = coords2unimol(
                 atoms,
                 coordinates,
                 self.dictionary,
                 self.max_atoms,
                 remove_hs=self.remove_hs,
             )
+            return feat, mol
         else:
             raise ValueError(
                 'Unknown conformer generation method: {}'.format(self.method)
@@ -146,16 +147,36 @@ def transform_raw(self, atoms_list, coordinates_list):
             )
         return inputs
 
+    def transform_mols(self, mols_list):
+        inputs = []
+        for mol in mols_list:
+            atoms = np.array([atom.GetSymbol() for atom in mol.GetAtoms()])
+            coordinates = mol.GetConformer().GetPositions().astype(np.float32)
+            inputs.append(
+                coords2unimol(
+                    atoms,
+                    coordinates,
+                    self.dictionary,
+                    self.max_atoms,
+                    remove_hs=self.remove_hs,
+                )
+            )
+        return inputs
+
     def transform(self, smiles_list):
         logger.info('Start generating conformers...')
         if self.multi_process:
             pool = Pool(processes=min(8, os.cpu_count()))
-            inputs = [
+            results = [
                 item for item in tqdm(pool.imap(self.single_process, smiles_list))
             ]
             pool.close()
         else:
-            inputs = [self.single_process(smiles) for smiles in tqdm(smiles_list)]
+            results = [self.single_process(smiles) for smiles in tqdm(smiles_list)]
+
+        inputs, mols = zip(*results)
+        inputs = list(inputs)
+        mols = list(mols)
 
         failed_conf = [(item['src_coord'] == 0.0).all() for item in inputs]
         logger.info(
@@ -192,7 +213,7 @@ def transform(self, smiles_list):
                     [smiles_list[index] for index in failed_conf_3d_indices]
                 )
             )
-        return inputs
+        return inputs, mols
 
 
 def inner_smi2coords(smi, seed=42, mode='fast', remove_hs=True, return_mol=False):
@@ -253,9 +274,9 @@ def inner_smi2coords(smi, seed=42, mode='fast', remove_hs=True, return_mol=False
         assert len(atoms_no_h) == len(
             coordinates_no_h
         ), "coordinates shape is not align with {}".format(smi)
-        return atoms_no_h, coordinates_no_h
+        return atoms_no_h, coordinates_no_h, mol
     else:
-        return atoms, coordinates
+        return atoms, coordinates, mol
 
 
 def inner_coords(atoms, coordinates, remove_hs=True):
@@ -391,7 +412,8 @@ def single_process(self, smiles):
                 remove_hs=self.remove_hs,
                 return_mol=True,
             )
-            return mol2unimolv2(mol, self.max_atoms, remove_hs=self.remove_hs)
+            feat = mol2unimolv2(mol, self.max_atoms, remove_hs=self.remove_hs)
+            return feat, mol
         else:
             raise ValueError(
                 'Unknown conformer generation method: {}'.format(self.method)
@@ -405,16 +427,26 @@ def transform_raw(self, atoms_list, coordinates_list):
             inputs.append(mol2unimolv2(mol, self.max_atoms, remove_hs=self.remove_hs))
         return inputs
 
+    def transform_mols(self, mols_list):
+        inputs = []
+        for mol in mols_list:
+            inputs.append(mol2unimolv2(mol, self.max_atoms, remove_hs=self.remove_hs))
+        return inputs
+
     def transform(self, smiles_list):
         logger.info('Start generating conformers...')
         if self.multi_process:
             pool = Pool(processes=min(8, os.cpu_count()))
-            inputs = [
+            results = [
                 item for item in tqdm(pool.imap(self.single_process, smiles_list))
             ]
             pool.close()
         else:
-            inputs = [self.single_process(smiles) for smiles in tqdm(smiles_list)]
+            results = [self.single_process(smiles) for smiles in tqdm(smiles_list)]
+
+        inputs, mols = zip(*results)
+        inputs = list(inputs)
+        mols = list(mols)
 
         failed_conf = [(item['src_coord'] == 0.0).all() for item in inputs]
         logger.info(
@@ -452,7 +484,7 @@ def transform(self, smiles_list):
                 )
             )
 
-        return inputs
+        return inputs, mols
 
 
 def create_mol_from_atoms_and_coords(atoms, coordinates):

unimol_tools/unimol_tools/data/datahub.py

@@ -4,7 +4,9 @@
 
 from __future__ import absolute_import, division, print_function
 
+import os
 import numpy as np
+from rdkit.Chem import PandasTools
 
 from ..utils import logger
 from .conformer import ConformerGen, UniMolV2Feature
@@ -29,13 +31,14 @@ def __init__(self, data=None, is_train=True, save_path=None, **params):
         :param save_path: (str) Path to save any necessary files, like scalers.
         :param params: Additional parameters for data preprocessing and model configuration.
         """
-        self.data = data
+        self.raw_data = data
         self.is_train = is_train
         self.save_path = save_path
         self.task = params.get('task', None)
         self.target_cols = params.get('target_cols', None)
         self.multiclass_cnt = params.get('multiclass_cnt', None)
         self.ss_method = params.get('target_normalize', 'none')
+        self.conf_cache_level = params.get('conf_cache_level', 1)
         self._init_data(**params)
         self._init_split(**params)
 
@@ -50,7 +53,7 @@ def _init_data(self, **params):
         :param params: Additional parameters for data processing.
         :raises ValueError: If the task type is unknown.
         """
-        self.data = MolDataReader().read_data(self.data, self.is_train, **params)
+        self.data = MolDataReader().read_data(self.raw_data, self.is_train, **params)
         self.data['target_scaler'] = TargetScaler(
             self.ss_method, self.task, self.save_path
         )
@@ -93,24 +96,35 @@ def _init_data(self, **params):
             raise ValueError('Unknown task: {}'.format(self.task))
 
         if params.get('model_name', None) == 'unimolv1':
-            if 'atoms' in self.data and 'coordinates' in self.data:
+            if 'mols' in self.data:
+                no_h_list = ConformerGen(**params).transform_mols(self.data['mols'])
+                mols = None
+            elif 'atoms' in self.data and 'coordinates' in self.data:
                 no_h_list = ConformerGen(**params).transform_raw(
                     self.data['atoms'], self.data['coordinates']
                 )
+                mols = None
             else:
                 smiles_list = self.data["smiles"]
-                no_h_list = ConformerGen(**params).transform(smiles_list)
+                no_h_list, mols = ConformerGen(**params).transform(smiles_list)
         elif params.get('model_name', None) == 'unimolv2':
-            if 'atoms' in self.data and 'coordinates' in self.data:
+            if 'mols' in self.data:
+                no_h_list = UniMolV2Feature(**params).transform_mols(self.data['mols'])
+                mols = None
+            elif 'atoms' in self.data and 'coordinates' in self.data:
                 no_h_list = UniMolV2Feature(**params).transform_raw(
                     self.data['atoms'], self.data['coordinates']
                 )
+                mols = None
             else:
                 smiles_list = self.data["smiles"]
-                no_h_list = UniMolV2Feature(**params).transform(smiles_list)
+                no_h_list, mols = UniMolV2Feature(**params).transform(smiles_list)
 
         self.data['unimol_input'] = no_h_list
 
+        if mols is not None:
+            self.save_mol2sdf(self.data['raw_data'], mols, params)
+
     def _init_split(self, **params):
 
         self.split_method = params.get('split_method', '5fold_random')
@@ -135,3 +149,50 @@ def _init_split(self, **params):
             nfolds[te_idx] = enu
         self.data['split_nfolds'] = split_nfolds
         return split_nfolds
+
+    def save_mol2sdf(self, data, mols, params):
+        """
+        Save the conformers to a SDF file.
+
+        :param data: DataFrame containing the raw data.
+        :param mols: List of RDKit molecule objects.
+        """
+        if isinstance(self.raw_data, str):
+            base_name = os.path.splitext(os.path.basename(self.raw_data))[0]
+        elif isinstance(self.raw_data, list) or isinstance(self.raw_data, np.ndarray):
+            # If the raw_data is a list of smiles, we can use a default name.
+            base_name = 'unimol_conformers'
+        else:
+            logger.warning('Warning: raw_data is not a path or list, cannot save sdf.')
+            return
+        if params.get('sdf_save_path') is None:
+            if self.save_path is not None:
+                params['sdf_save_path'] = self.save_path
+            else:
+                return
+        save_path = os.path.join(params.get('sdf_save_path'), f"{base_name}.sdf")
+        if self.conf_cache_level == 0:
+            logger.warning(f"conf_cache_level is 0, do not save conformers.")
+            return
+        elif self.conf_cache_level == 1 and os.path.exists(save_path):
+            logger.warning(f"conf_cache_level is 1, but {save_path} exists, so do not save conformers.")
+            return
+        elif self.conf_cache_level == 2 or not os.path.exists(save_path):
+            logger.info(f"conf_cache_level is {self.conf_cache_level}, saving conformers to {save_path}.")
+        else:
+            logger.warning(f"Unknown conf_cache_level: {self.conf_cache_level}, do not saving conformers.")
+            return
+        sdf_result = data.copy()
+        sdf_result['ROMol'] = mols
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        try:
+            PandasTools.WriteSDF(
+                sdf_result,
+                save_path,
+                properties=list(sdf_result.columns),
+                idName='RowID',
+            )
+            logger.info(f"Successfully saved sdf file to {save_path}")
+        except Exception as e:
+            logger.warning(f"Failed to write sdf file: {e}")
+        pass

unimol_tools/unimol_tools/data/datareader.py

@@ -10,6 +10,7 @@
 import numpy as np
 import pandas as pd
 from rdkit import Chem
+from rdkit.Chem import PandasTools
 from rdkit.Chem.Scaffolds import MurckoScaffold
 
 from ..utils import logger
@@ -49,7 +50,13 @@ def read_data(self, data=None, is_train=True, **params):
         if isinstance(data, str):
             # load from file
             self.data_path = data
-            data = pd.read_csv(self.data_path)
+            if data.endswith('.sdf'):
+                # load sdf file
+                data = PandasTools.LoadSDF(data)
+            elif data.endswith('.csv'):
+                data = pd.read_csv(self.data_path)
+            else:
+                raise ValueError('Unknown file type: {}'.format(data))
         elif isinstance(data, dict):
             # load from dict
             if 'target' in data:
@@ -137,6 +144,9 @@ def read_data(self, data=None, is_train=True, **params):
             dd['atoms'] = data['atoms'].tolist()
             dd['coordinates'] = data['coordinates'].tolist()
 
+        if 'ROMol' in data.columns:
+            dd['mols'] = data['ROMol'].tolist()
+
         return dd
 
     def check_smiles(self, smi, is_train, smi_strict):

unimol_tools/unimol_tools/data/datascaler.py

@@ -9,13 +9,20 @@
 import joblib
 import numpy as np
 from scipy.stats import kurtosis, skew
-from sklearn.preprocessing import (FunctionTransformer, MaxAbsScaler,
-                                   MinMaxScaler, Normalizer, PowerTransformer,
-                                   QuantileTransformer, RobustScaler,
-                                   StandardScaler)
+from sklearn.preprocessing import (
+    FunctionTransformer,
+    MaxAbsScaler,
+    MinMaxScaler,
+    Normalizer,
+    PowerTransformer,
+    QuantileTransformer,
+    RobustScaler,
+    StandardScaler,
+)
 
 from ..utils import logger
 
+
 class TargetScaler(object):
     '''
     A class to scale the target.
@@ -80,7 +87,9 @@ def fit(self, target, dump_dir):
         elif self.ss_method == 'auto':
             if self.task == 'regression':
                 if self.is_skewed(target):
-                    self.scaler = FunctionTransformer(func=np.log1p, inverse_func=np.expm1)
+                    self.scaler = FunctionTransformer(
+                        func=np.log1p, inverse_func=np.expm1
+                    )
                     logger.info('Auto select robust transformer.')
                 else:
                     self.scaler = StandardScaler()
@@ -90,7 +99,9 @@ def fit(self, target, dump_dir):
                 target = np.ma.masked_invalid(target)  # mask NaN value
                 for i in range(target.shape[1]):
                     if self.is_skewed(target[:, i]):
-                        self.scaler.append(FunctionTransformer(func=np.log1p, inverse_func=np.expm1))
+                        self.scaler.append(
+                            FunctionTransformer(func=np.log1p, inverse_func=np.expm1)
+                        )
                         logger.info('Auto select robust transformer.')
                     else:
                         self.scaler.append(StandardScaler())

unimol_tools/unimol_tools/models/nnmodel.py

@@ -14,8 +14,7 @@
 from torch.utils.data import Dataset
 
 from ..utils import logger
-from .loss import (FocalLossWithLogits, GHMC_Loss, MAEwithNan,
-                   myCrossEntropyLoss)
+from .loss import FocalLossWithLogits, GHMC_Loss, MAEwithNan, myCrossEntropyLoss
 from .unimol import UniMolModel
 from .unimolv2 import UniMolV2Model
 

unimol_tools/unimol_tools/models/unimol.py

@@ -5,6 +5,7 @@
 from __future__ import absolute_import, division, print_function
 
 import os
+
 # import argparse
 import pathlib