Update tool/PCE.py

tool/PCE.py

@@ -1,146 +1,206 @@
-# -*- coding: utf-8 -*-
-
-"""
-Created on Wed Sep 11 10:27:20 2024
-
-@author: BM109X32G-10GPU-02
-"""
-
-from langchain.tools import BaseTool
-from rdkit import Chem
-from rdkit.Chem import rdMolDescriptors
-from rdkit.Chem import Descriptors
-from .deepacceptor import RF
-from .deepdonor import sm, pm
-from .dap import run, screen
-import pandas as pd 
-
-class acceptor_predictor(BaseTool):
-    name:str  = "acceptor_predictor"
-    description:str  = (
-        "Input acceptor SMILES , returns the score of the acceptor."
-    )
-    
-    def __init__(self):
-        super().__init__()
-    def _run(self, smiles: str) -> str:
-        mol = Chem.MolFromSmiles(smiles)
-        if mol is None:
-            return "Invalid SMILES string"
-        smiles = Chem.MolToSmiles(mol) 
-        pce = RF.main( str(smiles) ) 
-        return f'The power conversion efficiency (PCE) is predicted to be {pce} (predicted by DeepAcceptor)  '  
- 
-    async def _arun(self, smiles: str) -> str:
-        """Use the tool asynchronously."""
-        raise NotImplementedError()
-        
-class donor_predictor(BaseTool):
-    name:str  = "donor_predictor"
-    description:str  = (
-        "Input donor SMILES , returns the score of the donor."
-    )
-    
-    def __init__(self):
-        super().__init__()
-    def _run(self, smiles: str) -> str:
-        mol = Chem.MolFromSmiles(smiles)
-        if mol is None:
-            return "Invalid SMILES string"
-     
-        mol = Chem.MolFromSmiles(smiles)
-        if mol is None:
-            return "Invalid SMILES string"
-        sdpce = sm.main( str(smiles) ) 
-        pdpce =  pm.main( str(smiles) ) 
-        return f'The power conversion efficiency (PCE) of the given molecule is predicted to be {sdpce} as a small molecule donor , and {pdpce} as a polymer donor(predicted by DeepDonor)  '  
-  
-    async def _arun(self, smiles: str) -> str:
-        """Use the tool asynchronously."""
-        raise NotImplementedError()
-        
- 
-
-class dap_predictor(BaseTool):
-    name:str  = "dap_predictor"
-    description :str = (
-        "Input SMILES of D/A pairs(separated by '.') , returns the performance of the D/A pairs ."
-    )
-
-
-    def __init__(self):
-        super().__init__()
-
-    def _run(self, smiles_pair: str) -> str:
-        smi_list = smiles_pair.split(".")
-        if len(smi_list) != 2:
- 
-            return "Input error, please input two smiles strings separated by '.'"  
-            
-        else:
-            smiles1, smiles2 = smi_list
-
-         
-        pce = run.smiles_aas_test( str(smiles1 ), str(smiles2) )
-         
-        return pce
-
-    async def _arun(self, smiles_pair: str) -> str:
-        """Use the tool asynchronously."""
-        raise NotImplementedError()
-        
-    
-
-class dap_screen(BaseTool):
-    name:str  = "dap_screen"
-    description :str = (
-        "Input dataset path containing D/A pairs, returns the files of prediction results."
-    )
-    return_direct: bool = True
-    def __init__(self):
-        super().__init__()
-
-    def _run(self, file_path: str) -> str:
-        smi_list = screen.smiles_aas_test(file_path)
-               
-        return smi_list
-
-    async def _arun(self, smiles_pair: str) -> str:
-        """Use the tool asynchronously."""
-        raise NotImplementedError()
-        
- 
-from .comget import generator  
- 
-class molgen(BaseTool):
-    name: str = "donorgen"
-    description: str = (
-
-        "Useful to generate polymer donor molecules with required PCE. "
-        "Input the values of PCE , return the SMILES"
-    )
- 
-    
-    def __init__(self
-                 ):
-        super().__init__(  )
- 
-
-    def _run(self, value ) -> str:
-        try:
-            results = generator.generation(value)
-            for i in results['smiles']:
-                pdpce =  pm.main( str(i) ) 
-                if abs(pdpce-float(value))<1.0:
-                    return f"The SMILES of generated donor is {i}, its predicted PCE is {pdpce}."
-                    break
-                
-            
-        
-        except Exception as e: 
-            return str(e)
-        
-
-    async def _arun(self, query) -> str:
-        """Use the tool asynchronously."""
-        raise NotImplementedError("this tool does not support async")
+# -*- coding: utf-8 -*-
+
+"""
+Created on Wed Sep 11 10:27:20 2024
+
+@author: BM109X32G-10GPU-02
+"""
+
+from langchain.tools import BaseTool
+from rdkit import Chem
+from rdkit.Chem import rdMolDescriptors
+from rdkit.Chem import Descriptors
+from .deepacceptor import RF
+from .deepdonor import sm, pm
+from .dap import run, screen
+import pandas as pd 
+
+class acceptor_predictor(BaseTool):
+    name:str  = "acceptor_predictor"
+    description:str  = (
+        "Input acceptor SMILES , returns the score of the acceptor."
+    )
+    
+    def __init__(self):
+        super().__init__()
+    def _run(self, smiles: str) -> str:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return "Invalid SMILES string"
+        smiles = Chem.MolToSmiles(mol) 
+        pce = RF.main( str(smiles) ) 
+        return f'The power conversion efficiency (PCE) is predicted to be {pce} (predicted by DeepAcceptor)  '  
+ 
+    async def _arun(self, smiles: str) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError()
+        
+class donor_predictor(BaseTool):
+    name:str  = "donor_predictor"
+    description:str  = (
+        "Input donor SMILES , returns the score of the donor."
+    )
+    
+    def __init__(self):
+        super().__init__()
+    def _run(self, smiles: str) -> str:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return "Invalid SMILES string"
+     
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return "Invalid SMILES string"
+        sdpce = sm.main( str(smiles) ) 
+        pdpce =  pm.main( str(smiles) ) 
+        return f'The power conversion efficiency (PCE) of the given molecule is predicted to be {sdpce} as a small molecule donor , and {pdpce} as a polymer donor(predicted by DeepDonor)  '  
+  
+    async def _arun(self, smiles: str) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError()
+        
+ 
+
+class dap_predictor(BaseTool):
+    name:str  = "dap_predictor"
+    description :str = (
+        "Input SMILES of D/A pairs(separated by '.') , returns the performance of the D/A pairs ."
+    )
+
+
+    def __init__(self):
+        super().__init__()
+
+    def _run(self, smiles_pair: str) -> str:
+        smi_list = smiles_pair.split(".")
+        if len(smi_list) != 2:
+ 
+            return "Input error, please input two smiles strings separated by '.'"  
+            
+        else:
+            smiles1, smiles2 = smi_list
+
+         
+        pce = run.smiles_aas_test( str(smiles1 ), str(smiles2) )
+         
+        return pce
+
+    async def _arun(self, smiles_pair: str) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError()
+        
+    
+
+class dap_screen(BaseTool):
+    name:str  = "dap_screen"
+    description :str = (
+        "Input dataset path containing D/A pairs, returns the files of prediction results."
+    )
+    return_direct: bool = True
+    def __init__(self):
+        super().__init__()
+
+    def _run(self, file_path: str) -> str:
+        smi_list = screen.smiles_aas_test(file_path)
+               
+        return smi_list
+
+    async def _arun(self, smiles_pair: str) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError()
+        
+ 
+from .comget import generator  
+ 
+class molgen(BaseTool):
+    name: str = "donorgen"
+    description: str = (
+
+        "Useful to generate polymer donor molecules with required PCE. "
+        "Input the values of PCE , return the SMILES"
+    )
+ 
+    
+    def __init__(self
+                 ):
+        super().__init__(  )
+ 
+
+    def _run(self, value ) -> str:
+        try:
+            results = generator.generation(value)
+            for i in results['smiles']:
+                pdpce =  pm.main( str(i) ) 
+                if abs(pdpce-float(value))<1.0:
+                    return f"The SMILES of generated donor is {i}, its predicted PCE is {pdpce}."
+                    break
+                
+            
+        
+        except Exception as e: 
+            return str(e)
+        
+
+    async def _arun(self, query) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError("this tool does not support async")
+from rdkit import Chem
+from rdkit.Chem import BRICS 
+from typing import List
+
+def connect_A_B_A(a_smiles: str, b_smiles: str) -> str:
+    mol_A = Chem.MolFromSmiles(a_smiles)
+    mol_B = Chem.MolFromSmiles(b_smiles)
+    final_product_smiles = set() 
+    ab_intermediates_gen = BRICS.BRICSBuild([mol_A, mol_B])
+    for intermediate_mol in ab_intermediates_gen:
+        Chem.SanitizeMol(intermediate_mol)
+        aba_products_gen = BRICS.BRICSBuild([intermediate_mol, mol_A])
+        for final_mol in aba_products_gen: 
+            try:
+                Chem.SanitizeMol(final_mol)
+                final_product_smiles.add(Chem.MolToSmiles(final_mol))
+            except Exception as e:
+                print(f" {e}")
+ 
+    final_mols = [Chem.MolFromSmiles(s) for s in sorted(list(final_product_smiles))]
+    smiles = Chem.MolToSmiles(final_mols[0])
+    return smiles
+
+ 
+class QMGen(BaseTool):
+    name: str = "QMGen"
+    description: str = (
+        "Useful to generateQuasi-macromolecule . "
+        "Input two molecule SMILES (separated by '.') , return the SMILES"
+    )  
+    def __init__(self):
+        super().__init__(  )
+    def _run(self, smiles_pair: str) -> str:
+     smi_list = smiles_pair.split(".")
+     if len(smi_list) != 2:
+         return "Input error, please input two smiles strings separated by '.'"
+     else:
+         smiles1, smiles2 = smi_list
+     
+     frag_1dummy_s = [smiles for smiles in smi_list if smiles.count('*')==1]
+     frag_2dummy_s = [smiles for smiles in smi_list if smiles.count('*')==2] 
+     if len(frag_1dummy_s) != len(frag_1dummy_s):
+         return "Input error, please input acceptor (A) unit smiles with one '[*]' and Π bridge unit with two '[*]'."
+     smiles1 =  frag_1dummy_s[0].replace('[*]','[16*]').replace('C[16*]','[16*]').replace('[16*]C','[16*]')
+     smiles2=  frag_2dummy_s[0].replace('[*]','[14*]')
+     
+     if len(frag_1dummy_s) != 1:
+        return "Input error, please input  smiles with '[*]' as the the key to be connected e.g.[14*]c1ccc([14*])s1."
+        try:
+            products = connect_A_B_A(smiles1, smiles2)
+            return products    
+            
+        
+        except Exception as e: 
+            return str(e)
+        
+
+    async def _arun(self, query) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError("this tool does not support async")