bayesianflow-for-chem 1.4.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

bayesianflow_for_chem/__init__.py

@@ -7,5 +7,5 @@ from . import data, tool, train, scorer
 from .model import ChemBFN, MLP, EnsembleChemBFN
 
 __all__ = ["data", "tool", "train", "scorer", "ChemBFN", "MLP", "EnsembleChemBFN"]
-__version__ = "1.4.0"
+__version__ = "1.4.1"
 __author__ = "Nianze A. Tao (Omozawa Sueno)"

bayesianflow_for_chem/model.py

@@ -169,16 +169,6 @@ class Attention(nn.Module):
         k = k.view(split).permute(2, 0, 1, 3).contiguous()
         v = v.view(split).permute(2, 0, 1, 3).contiguous()
         q, k = self._rotate(q, k, pe)  # position embedding
-        """
-        # Original code. Maybe using `nn.functional.scaled_dot_product_attention(...)` is better.
-
-        k_t = k.transpose(-2, -1)
-        if mask is not None:
-            alpha = softmax((q @ k_t / self.tp).masked_fill_(mask, -torch.inf), -1)
-        else:
-            alpha = softmax(q @ k_t / self.tp, -1)
-        atten_out = (alpha @ v).permute(1, 2, 0, 3).contiguous().view(shape)
-        """
         atten_out = nn.functional.scaled_dot_product_attention(
             q, k, v, mask, 0.0, False, scale=1 / self.tp
         )
@@ -430,19 +420,14 @@ class ChemBFN(nn.Module):
             c += y
         pe = self.position(n_t)
         x = self.embedding(x)
-        attn_mask: Optional[Tensor] = None
         if self.semi_autoregressive:
             attn_mask = torch.tril(
                 torch.ones((1, n_b, n_t, n_t), device=x.device), diagonal=0
             )
+        elif mask is not None:
+            attn_mask = mask.transpose(-2, -1).repeat(1, n_t, 1)[None, ...] != 0
         else:
-            if mask is not None:
-                """
-                # Original Code.
-
-                attn_mask = mask.transpose(-2, -1).repeat(1, x.shape[1], 1)[None, ...] == 0
-                """
-                attn_mask = mask.transpose(-2, -1).repeat(1, n_t, 1)[None, ...] != 0
+            attn_mask = None
         for layer in self.encoder_layers:
             x = layer(x, pe, c, attn_mask)
         return self.final_layer(x, c, mask is None)

bayesianflow_for_chem/tool.py

@@ -16,15 +16,16 @@ from torch import cuda, Tensor, softmax
 from torch.ao import quantization
 from torch.utils.data import DataLoader
 from typing_extensions import Self
-from rdkit.Chem.rdchem import Mol, Bond
 from rdkit.Chem import (
     rdDetermineBonds,
+    GetFormalCharge,
     MolFromXYZBlock,
     MolFromSmiles,
     MolToSmiles,
     CanonSmiles,
     AllChem,
     AddHs,
+    Mol,
 )
 from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles  # type: ignore
 from sklearn.metrics import (
@@ -539,40 +540,88 @@ class GeometryConverter:
             xyz_block.append(f"{atom} {r[i][0]:.10f} {r[i][1]:.10f} {r[i][2]:.10f}")
         return MolFromXYZBlock("\n".join(xyz_block))
 
-    @staticmethod
-    def _bond_pair_idx(bonds: Bond) -> List[List[int]]:
-        return [[i.GetBeginAtomIdx(), i.GetEndAtomIdx()] for i in bonds]
-
     @staticmethod
     def smiles2cartesian(
-        smiles: str, num_conformers: int = 50, random_seed: int = 42
+        smiles: str,
+        num_conformers: int = 50,
+        rdkit_ff_type: str = "MMFF",
+        refine_with_crest: bool = False,
+        spin: float = 0.0,
     ) -> Tuple[List[str], np.ndarray]:
         """
         Guess the 3D geometry from SMILES string via MMFF conformer search.
 
         :param smiles: a valid SMILES string
         :param num_conformers: number of initial conformers
-        :param random_seed: random seed used to generate conformers
+        :param rdkit_ff_type: force field type chosen in `'MMFF'` and `'UFF'`
+        :param refine_with_crest: find the best conformer via CREST
+        :param spin: total spin; only required when `refine_with_crest=True`
         :type smiles: str
         :type num_conformers: int
-        :type random_seed: int
+        :type rdkit_ff_type: str
+        :type refine_with_crest: bool
+        :type spin: float
         :return: atomic symbols \n
                  cartesian coordinates;  shape: (n_a, 3)
         :rtype: tuple
         """
+        assert rdkit_ff_type.lower() in ("mmff", "uff")
+        if refine_with_crest:
+            from tempfile import TemporaryDirectory
+            from subprocess import run
+
+            # We need both CREST and xTB installed.
+            if run("crest --version", shell=True).returncode != 0:
+                raise RuntimeError(
+                    "`CREST` is not found! Make sure it is installed and added into the PATH."
+                )
+            if run("xtb --version", shell=True).returncode != 0:
+                raise RuntimeError(
+                    "`xTB` is not found! Make sure it is installed and added into the PATH."
+                )
         mol = MolFromSmiles(smiles)
         mol = AddHs(mol)
-        AllChem.EmbedMultipleConfs(mol, numConfs=num_conformers, randomSeed=random_seed)
+        AllChem.EmbedMultipleConfs(mol, numConfs=num_conformers, params=AllChem.ETKDG())
         symbols = [atom.GetSymbol() for atom in mol.GetAtoms()]
         energies = []
         for conf_id in range(num_conformers):
-            ff = AllChem.MMFFGetMoleculeForceField(
-                mol, AllChem.MMFFGetMoleculeProperties(mol), confId=conf_id
-            )
+            if rdkit_ff_type.lower() == "mmff":
+                ff = AllChem.MMFFGetMoleculeForceField(
+                    mol, AllChem.MMFFGetMoleculeProperties(mol), confId=conf_id
+                )
+            else:  # UFF
+                ff = AllChem.UFFGetMoleculeForceField(mol, confId=conf_id)
             energy = ff.CalcEnergy()
             energies.append((conf_id, energy))
         lowest_energy_conf = min(energies, key=lambda x: x[1])
         coordinates = mol.GetConformer(id=lowest_energy_conf[0]).GetPositions()
+        if refine_with_crest:
+            xyz = f"{len(symbols)}\n\n" + "\n".join(
+                f"{s} {coordinates[i][0]:.10f} {coordinates[i][1]:.10f} {coordinates[i][2]:.10f}"
+                for i, s in enumerate(symbols)
+            )
+            chrg = GetFormalCharge(mol)
+            uhf = int(spin * 2)
+            with TemporaryDirectory(dir=Path.cwd()) as temp_dir:
+                with open(Path(temp_dir) / "mol.xyz", "w", encoding="utf-8") as f:
+                    f.write(xyz)
+                s = run(
+                    f"crest mol.xyz -gfn2 -quick -prop ohess{f' --chrg {chrg}' if chrg != 0 else ''}{f' --uhf {uhf}' if uhf != 0 else ''}",
+                    shell=True,
+                    cwd=temp_dir,
+                )
+                if s.returncode == 0:
+                    with open(Path(temp_dir) / "crest_property.xyz", "r") as f:
+                        xyz = f.readlines()
+                    xyz_data = []
+                    for i in xyz[2:]:
+                        if i == xyz[0]:
+                            break
+                        xyz_data.append(i.strip().split())
+                    xyz_data = np.array(xyz_data)
+                    symbols, coordinates = np.split(xyz_data, [1], axis=-1)
+                    symbols = symbols.flatten().tolist()
+                    coordinates = coordinates.astype(np.float64)
         return symbols, coordinates
 
     def cartesian2smiles(

{bayesianflow_for_chem-1.4.0.dist-info → bayesianflow_for_chem-1.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 1.4.0
+Version: 1.4.1
 Summary: Bayesian flow network framework for Chemistry
 Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
 Author: Nianze A. Tao

bayesianflow_for_chem-1.4.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+bayesianflow_for_chem/__init__.py,sha256=N_7P9Ea0eUmdC0wQKXIHiuMzPK4p9_cBF_YOexjo5yo,329
+bayesianflow_for_chem/data.py,sha256=WoOCOVmJX4WeHa2WeO4i66J2FS8rvRaYRCdlBN7ZeOM,6576
+bayesianflow_for_chem/model.py,sha256=zJkcUnZcxFa4iTo9_-BHzAM1MkJm1pbEGiczVgyUcxo,50186
+bayesianflow_for_chem/scorer.py,sha256=7G1TVSwC0qONtNm6kiDZUWwvuFPzasNSjp4eJAk5TL0,4101
+bayesianflow_for_chem/tool.py,sha256=Ma4dEBWP5GFKk-Tj5vBJgxGs_WAp4F87-b1UcsqUAn4,25486
+bayesianflow_for_chem/train.py,sha256=hGKyhGhLch-exSYPZdLXrLn3gf39Q1VLSJs2qtuikQE,9709
+bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
+bayesianflow_for_chem-1.4.1.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+bayesianflow_for_chem-1.4.1.dist-info/METADATA,sha256=460yUOjHG9PTavIddJJ2Ufdq0bkLBZqbmMugyq6LVPQ,5643
+bayesianflow_for_chem-1.4.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bayesianflow_for_chem-1.4.1.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
+bayesianflow_for_chem-1.4.1.dist-info/RECORD,,

bayesianflow_for_chem-1.4.0.dist-info/RECORD

@@ -1,12 +0,0 @@
-bayesianflow_for_chem/__init__.py,sha256=3sP8nM4_idOX-ksvpBJEApxPAVAPijKvQHxidTO5790,329
-bayesianflow_for_chem/data.py,sha256=WoOCOVmJX4WeHa2WeO4i66J2FS8rvRaYRCdlBN7ZeOM,6576
-bayesianflow_for_chem/model.py,sha256=fUrXKhn2U9FrVPJyb4lqACqPTePkIgI0v6_1jPs5c0Q,50784
-bayesianflow_for_chem/scorer.py,sha256=7G1TVSwC0qONtNm6kiDZUWwvuFPzasNSjp4eJAk5TL0,4101
-bayesianflow_for_chem/tool.py,sha256=NMMRHk2FJY0fyA76zCrz6tkcylCuExMUMj5hohWTnkE,23155
-bayesianflow_for_chem/train.py,sha256=hGKyhGhLch-exSYPZdLXrLn3gf39Q1VLSJs2qtuikQE,9709
-bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
-bayesianflow_for_chem-1.4.0.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-bayesianflow_for_chem-1.4.0.dist-info/METADATA,sha256=1Y5mLIOaPsHcyCCm2SkWz7OCniQYVJ67-cVq3cUU0Mw,5643
-bayesianflow_for_chem-1.4.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bayesianflow_for_chem-1.4.0.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
-bayesianflow_for_chem-1.4.0.dist-info/RECORD,,