From 09fb49efcc47a7723481e2e26b910a2b7def528f Mon Sep 17 00:00:00 2001
From: Kohulan <kohulan.rajan@uni-jena.de>
Date: Tue, 27 Jun 2023 12:38:25 +0200
Subject: [PATCH] feat: add openbabel 3d depiction config

---
 app/modules/toolkits/openbabel_wrapper.py | 13 +++++++++----
 app/modules/toolkits/rdkit_wrapper.py     |  6 +++---
 app/routers/depict.py                     | 18 +++++++++++++++---
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/app/modules/toolkits/openbabel_wrapper.py b/app/modules/toolkits/openbabel_wrapper.py
index 2540dc4..d06aafd 100644
--- a/app/modules/toolkits/openbabel_wrapper.py
+++ b/app/modules/toolkits/openbabel_wrapper.py
@@ -49,10 +49,10 @@ def getOBInChI(smiles: str, InChIKey: bool = False):
     return inchi
 
 
-def getOBMol(smiles: str, threeD: bool = False):
+def getOBMol(smiles: str, threeD: bool = False, depict: bool = False):
     """This function takes an input as a SMILES string and
     returns a 2D/3D mol block.
-    Args (str): SMILES string.
+    Args (str,bool,bool): SMILES string, 3D mol block, 3D mol block for depiction
     Returns (str): Mol block (2D/3D).
     """
     if any(char.isspace() for char in smiles):
@@ -62,8 +62,13 @@ def getOBMol(smiles: str, threeD: bool = False):
         mol = pybel.readstring("smi", smiles)
         mol.addh()
         mol.make3D()
-        mol.removeh()
-        return mol.write("mol")
+        gen3d = ob.OBOp.FindType("gen3D")
+        gen3d.Do(mol.OBMol, "--best")
+        if depict:
+            return mol.write("mol")
+        else:
+            mol.removeh()
+            return mol.write("mol")
 
     # Create an Open Babel molecule object
     mol = ob.OBMol()
diff --git a/app/modules/toolkits/rdkit_wrapper.py b/app/modules/toolkits/rdkit_wrapper.py
index 7d13664..365e228 100644
--- a/app/modules/toolkits/rdkit_wrapper.py
+++ b/app/modules/toolkits/rdkit_wrapper.py
@@ -103,16 +103,16 @@ def get3Dconformers(smiles, depict=True):
         mol = Chem.MolFromSmiles(smiles)
     if mol:
         mol = Chem.AddHs(mol)
-        AllChem.EmbedMolecule(mol, randomSeed=0xF00D)
+        AllChem.EmbedMolecule(mol, maxAttempts=5000, useRandomCoords=True)
         try:
             AllChem.MMFFOptimizeMolecule(mol)
         except Exception as e:
             print(e)
-            AllChem.EmbedMolecule(mol, randomSeed=0xF00D)
+            AllChem.EmbedMolecule(mol, maxAttempts=5000, useRandomCoords=True)
         if depict:
             return Chem.MolToMolBlock(mol)
         else:
-            # mol = Chem.RemoveHs(mol)
+            mol = Chem.RemoveHs(mol)
             return Chem.MolToMolBlock(mol)
     else:
         return "Error reading SMILES string, check again."
diff --git a/app/routers/depict.py b/app/routers/depict.py
index 645a8f7..768773d 100644
--- a/app/routers/depict.py
+++ b/app/routers/depict.py
@@ -3,6 +3,7 @@
 from fastapi.responses import Response, HTMLResponse
 from app.modules.depiction import getRDKitDepiction, getCDKDepiction
 from app.modules.toolkits.rdkit_wrapper import get3Dconformers
+from app.modules.toolkits.openbabel_wrapper import getOBMol
 from fastapi.templating import Jinja2Templates
 
 templates = Jinja2Templates(directory="app/templates")
@@ -56,12 +57,23 @@ async def Depict2D_molecule(
 async def Depict3D_Molecule(
     request: Request,
     smiles: str,
+    generator: Optional[str] = "rdkit",
 ):
     """
-    Generate 3D Depictions using RDKit.
+    Generate 3D Depictions using OpenBabel/RDKit.
 
     - **SMILES**: required (query)
+    - **generator**: optional (query), default: rdkit, allowed: openbabel
     """
     if smiles:
-        content = {"request": request, "molecule": get3Dconformers(smiles)}
-        return templates.TemplateResponse("mol.html", content)
+        if generator == "openbabel":
+            content = {
+                "request": request,
+                "molecule": getOBMol(smiles, threeD=True, depict=True),
+            }
+            return templates.TemplateResponse("mol.html", content)
+        elif generator == "rdkit":
+            content = {"request": request, "molecule": get3Dconformers(smiles)}
+            return templates.TemplateResponse("mol.html", content)
+        else:
+            return "Check SMILES string or Toolkit configuration."