bayesianflow-for-chem 1.4.1__py3-none-any.whl → 1.4.3__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.1"
+__version__ = "1.4.3"
 __author__ = "Nianze A. Tao (Omozawa Sueno)"

bayesianflow_for_chem/model.py

@@ -55,8 +55,16 @@ class Linear(nn.Linear):
         :rtype: None
         """
         assert r > 0, "Rank should be larger than 0."
-        self.lora_A = nn.Parameter(self.weight.new_zeros((r, self.in_features)))
-        self.lora_B = nn.Parameter(self.weight.new_zeros((self.out_features, r)))
+        try:
+            self.lora_A = nn.Parameter(self.weight.new_zeros((r, self.in_features)))
+            self.lora_B = nn.Parameter(self.weight.new_zeros((self.out_features, r)))
+        except NotImplementedError:
+            self.lora_A = nn.Parameter(
+                torch.zeros((r, self.in_features), device=self.weight.device)
+            )
+            self.lora_B = nn.Parameter(
+                torch.zeros((self.out_features, r), device=self.weight.device)
+            )
         self.scaling = lora_alpha / r
         self.lora_dropout = lora_dropout
         self.lora_enabled = True
@@ -1207,23 +1215,23 @@ class EnsembleChemBFN(ChemBFN):
         )
 
     def quantise(
-        self, quantise_method: Optional[Callable[[ChemBFN], nn.Module]] = None
+        self, quantise_method: Optional[Callable[[ChemBFN], None]] = None
     ) -> None:
         """
         Quantise the submodels. \n
         This method should be called, if necessary, before `torch.compile()`.
 
-        :param quantise_method: quantisation method; default is `bayesianflow_for_chem.tool.quantise_model`
+        :param quantise_method: quantisation method; default is `bayesianflow_for_chem.tool.quantise_model_`
         :type quantise_method: callable | None
         :return:
         :rtype: None
         """
         if quantise_method is None:
-            from bayesianflow_for_chem.tool import quantise_model
+            from bayesianflow_for_chem.tool import quantise_model_
 
-            quantise_method = quantise_model
-        for k, v in self.models.items():
-            self.models[k] = quantise_method(v)
+            quantise_method = quantise_model_
+        for _, v in self.models.items():
+            quantise_method(v)
 
     def jit(self, freeze: bool = False) -> None:
         """

bayesianflow_for_chem/scorer.py

@@ -5,10 +5,12 @@ Define essential scorers.
 """
 from typing import List, Callable, Union, Optional
 import torch
+import numpy as np
 from torch import Tensor
 from rdkit import RDLogger
 from rdkit.Contrib.SA_Score import sascorer  # type: ignore
 from rdkit.Chem import MolFromSmiles, QED
+from scipy.stats import wasserstein_distance
 
 RDLogger.DisableLog("rdApp.*")  # type: ignore
 
@@ -49,6 +51,27 @@ def sa_score(smiles: str) -> float:
     return sascorer.calculateScore(MolFromSmiles(smiles))
 
 
+def spectra_wasserstein_score(
+    spectrum_u: np.ndarray, spectrum_v: np.ndarray, x_axis: np.ndarray
+) -> float:
+    """
+    Return the Wasserstein distance (earth mover's distance) between two
+    continuous spectra scaled by the area under the first spectrum curve `spectrum_u`.
+
+    :param spectrum_u: the reference spectrum
+    :param spectrum_v: the
+    :param x_axis: the shared x-axis of the spectra
+    :type spectrum_u: numpy.ndarray
+    :type spectrum_v: numpy.ndarray
+    :type x_axis: numpy.ndarray
+    :return: spectra Wasserstein score
+    :rtype: float
+    """
+    assert spectrum_u.size == spectrum_v.size, "Spectra sizes should be matched."
+    a = np.sqrt(np.trapezoid(spectrum_u, x_axis))
+    return (wasserstein_distance(spectrum_u, spectrum_v) / a).item()
+
+
 class Scorer:
     def __init__(
         self,
@@ -126,7 +149,7 @@ class Scorer:
             ]
             for i, scorer in enumerate(self.scorers)
         ]
-        loss = (e_k * p).sum(2).mean(1) * torch.tensor(scores, device=p.device).mean(0)
+        loss = (e_k * p).sum(2).mean(1) * p.new_tensor(scores).mean(0)
         return loss.mean()
 
 

bayesianflow_for_chem/tool.py

@@ -6,16 +6,12 @@ Essential tools.
 import csv
 import random
 import warnings
-from copy import deepcopy
 from pathlib import Path
 from typing import List, Dict, Tuple, Union, Optional
 import torch
 import numpy as np
-import torch.nn as nn
 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 import (
     rdDetermineBonds,
     GetFormalCharge,
@@ -37,7 +33,7 @@ from sklearn.metrics import (
     root_mean_squared_error,
 )
 from .data import VOCAB_KEYS
-from .model import ChemBFN, MLP, Linear, EnsembleChemBFN
+from .model import ChemBFN, MLP, EnsembleChemBFN
 
 
 def _find_device() -> torch.device:
@@ -386,145 +382,46 @@ def inpaint(
     ]
 
 
-def quantise_model(model: ChemBFN) -> nn.Module:
+def quantise_model_(model: ChemBFN) -> None:
     """
-    Dynamic quantisation of the trained model to `torch.qint8` data type.
+    In-place dynamic quantisation of the trained model to `int8` data type. \n
+    Due to some limitations of `torchao` module, not all layers will be quantised.
 
     :param model: trained ChemBFN model
     :type model: bayesianflow_for_chem.model.ChemBFN
-    :return: quantised model
-    :rtype: torch.nn.Module
+    :return:
+    :rtype: None
     """
-    from torch.ao.nn.quantized import dynamic
-    from torch.ao.nn.quantized.modules.utils import _quantize_weight
-    from torch.ao.quantization.qconfig import default_dynamic_qconfig
-
-    class QuantisedLinear(dynamic.Linear):
-        # Modified from https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/dynamic/modules/linear.py
-        # We made it compatible with our LoRA linear layer.
-        # LoRA parameters will not be quantised.
-        def __init__(
-            self,
-            in_features: int,
-            out_features: int,
-            bias_: bool = True,
-            dtype: torch.dtype = torch.qint8,
-        ) -> None:
-            super().__init__(in_features, out_features, bias_, dtype=dtype)
-            self.version = self._version
-            self.lora_enabled: bool = False
-            self.lora_A: Optional[nn.Parameter] = None
-            self.lora_B: Optional[nn.Parameter] = None
-            self.scaling: Optional[float] = None
-            self.lora_dropout: Optional[float] = None
-
-        def _get_name(self) -> str:
-            return "DynamicQuantizedLoRALinear"
-
-        def enable_lora(
-            self, r: int = 8, lora_alpha: int = 1, lora_dropout: float = 0.0
-        ) -> None:
-            assert r > 0, "Rank should be larger than 0."
-            device = self._weight_bias()[0].device
-            self.lora_A = nn.Parameter(
-                torch.zeros((r, self.in_features), device=device)
-            )
-            self.lora_B = nn.Parameter(
-                torch.zeros((self.out_features, r), device=device)
-            )
-            self.scaling = lora_alpha / r
-            self.lora_dropout = lora_dropout
-            self.lora_enabled = True
-            nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
-            nn.init.zeros_(self.lora_B)
-            self._packed_params.requires_grad_(False)
-
-        def forward(self, x: Tensor) -> Tensor:
-            if self._packed_params.dtype == torch.qint8:
-                if self.version is None or self.version < 4:
-                    Y = torch.ops.quantized.linear_dynamic(
-                        x, self._packed_params._packed_params
-                    )
-                else:
-                    Y = torch.ops.quantized.linear_dynamic(
-                        x, self._packed_params._packed_params, reduce_range=True
-                    )
-            elif self._packed_params.dtype == torch.float16:
-                Y = torch.ops.quantized.linear_dynamic_fp16(
-                    x, self._packed_params._packed_params
-                )
-            else:
-                raise RuntimeError("Unsupported dtype on dynamic quantized linear!")
-            result = Y.to(x.dtype)
-            if self.lora_enabled and isinstance(self.lora_dropout, float):
-                result += (
-                    nn.functional.dropout(x, self.lora_dropout, self.training)
-                    @ self.lora_A.transpose(0, 1)
-                    @ self.lora_B.transpose(0, 1)
-                ) * self.scaling
-            return result
+    from torchao.quantization.quant_api import (
+        quantize_,
+        Int8DynamicActivationInt8WeightConfig,
+    )
 
-        @classmethod
-        def from_float(
-            cls, mod: Linear, use_precomputed_fake_quant: bool = False
-        ) -> Self:
-            assert hasattr(
-                mod, "qconfig"
-            ), "Input float module must have qconfig defined"
-            if use_precomputed_fake_quant:
-                warnings.warn("Fake quantize operator is not implemented.")
-            if mod.qconfig is not None and mod.qconfig.weight is not None:
-                weight_observer = mod.qconfig.weight()
-            else:
-                weight_observer = default_dynamic_qconfig.weight()
-            dtype = weight_observer.dtype
-            assert dtype in [torch.qint8, torch.float16], (
-                "The only supported dtypes for "
-                f"dynamic quantized linear are qint8 and float16 got: {dtype}"
-            )
-            weight_observer(mod.weight)
-            if dtype == torch.qint8:
-                qweight = _quantize_weight(mod.weight.float(), weight_observer)
-            elif dtype == torch.float16:
-                qweight = mod.weight.float()
-            else:
-                raise RuntimeError(
-                    "Unsupported dtype specified for dynamic quantized Linear!"
-                )
-            qlinear = cls(mod.in_features, mod.out_features, dtype=dtype)
-            qlinear.set_weight_bias(qweight, mod.bias)
-            if mod.lora_enabled:
-                qlinear.lora_enabled = True
-                qlinear.lora_A = nn.Parameter(mod.lora_A.clone().detach_())
-                qlinear.lora_B = nn.Parameter(mod.lora_B.clone().detach_())
-                qlinear.scaling = deepcopy(mod.scaling)
-                qlinear.lora_dropout = deepcopy(mod.lora_dropout)
-            return qlinear
+    quantize_(model, Int8DynamicActivationInt8WeightConfig())
 
-        @classmethod
-        def from_reference(cls, ref_qlinear: Self) -> Self:
-            qlinear = cls(
-                ref_qlinear.in_features,
-                ref_qlinear.out_features,
-                dtype=ref_qlinear.weight_dtype,
-            )
-            qweight = ref_qlinear.get_quantized_weight()
-            bias = ref_qlinear.bias
-            qlinear.set_weight_bias(qweight, bias)
-            if ref_qlinear.lora_enabled:
-                qlinear.lora_enabled = True
-                qlinear.lora_A = nn.Parameter(ref_qlinear.lora_A.clone().detach_())
-                qlinear.lora_B = nn.Parameter(ref_qlinear.lora_B.clone().detach_())
-                qlinear.scaling = deepcopy(ref_qlinear.scaling)
-                qlinear.lora_dropout = deepcopy(ref_qlinear.lora_dropout)
-            return qlinear
 
-    mapping = deepcopy(quantization.DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS)
-    mapping[Linear] = QuantisedLinear
-    quantised_model = quantization.quantize_dynamic(
-        model, {nn.Linear, Linear}, torch.qint8, mapping
-    )
-    return quantised_model
+def build_uv_vis_sepctrum(
+    etoscs: np.ndarray, etenergies: np.ndarray, lambdas: np.ndarray
+) -> np.ndarray:
+    """
+    Build UV/Vis spectrum from calculated electron transtion energies and oscillator strengths. \n
+    This function follows the GaussView style: https://gaussian.com/uvvisplot/.
+
+    :param etoscs: oscillator strengths
+    :param etenergies: transtion energies
+    :param lambdas: wavelengths
+    :type etoscs: numpy.ndarray
+    :type etenergies: numpy.ndarray
+    :type lambdas: numpy.ndarray
+    :return: absorption coefficient corrospending to the wavelengths
+    :rtype: numpy.ndarray
+    """
+    return (
+        etoscs[:, None]
+        * np.exp(
+            -np.pow((1 / lambdas[None, :] - etenergies[:, None] / 45.5634) * 3099.6, 2)
+        )
+    ).sum(0) * 40489.99421
 
 
 class GeometryConverter:
@@ -543,7 +440,7 @@ class GeometryConverter:
     @staticmethod
     def smiles2cartesian(
         smiles: str,
-        num_conformers: int = 50,
+        num_conformers: int = 250,
         rdkit_ff_type: str = "MMFF",
         refine_with_crest: bool = False,
         spin: float = 0.0,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 1.4.1
+Version: 1.4.3
 Summary: Bayesian flow network framework for Chemistry
 Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
 Author: Nianze A. Tao
@@ -21,13 +21,15 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: rdkit>=2023.9.6
-Requires-Dist: torch>=2.3.1
-Requires-Dist: numpy>=1.26.4
+Requires-Dist: rdkit>=2025.3.5
+Requires-Dist: torch>=2.8.0
+Requires-Dist: torchao>=0.12
+Requires-Dist: numpy>=2.3.2
+Requires-Dist: scipy>=1.16.1
 Requires-Dist: loralib>=0.1.2
-Requires-Dist: lightning>=2.2.0
-Requires-Dist: scikit-learn>=1.5.0
-Requires-Dist: typing_extensions>=4.8.0
+Requires-Dist: lightning>=2.5.3
+Requires-Dist: scikit-learn>=1.7.1
+Requires-Dist: typing_extensions>=4.14.1
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

bayesianflow_for_chem-1.4.3.dist-info/RECORD

@@ -0,0 +1,12 @@
+bayesianflow_for_chem/__init__.py,sha256=lcVwcAWdsmdf8cFIromTsT2Hhn4hIBRlA47kCWWeHm0,329
+bayesianflow_for_chem/data.py,sha256=WoOCOVmJX4WeHa2WeO4i66J2FS8rvRaYRCdlBN7ZeOM,6576
+bayesianflow_for_chem/model.py,sha256=CIcFruMFg5IbaVRkc5geygcEV7LiBEN7X_2dqUqJi8w,50489
+bayesianflow_for_chem/scorer.py,sha256=i-1y6pd22n7ExzdXozCjOTXbEIJTkz_6Bb4g0RA5kQI,4928
+bayesianflow_for_chem/tool.py,sha256=AhW_dXuWEe7vl2epw_EZdEpmAYRd_3NMOQTuC25oHio,20558
+bayesianflow_for_chem/train.py,sha256=hGKyhGhLch-exSYPZdLXrLn3gf39Q1VLSJs2qtuikQE,9709
+bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
+bayesianflow_for_chem-1.4.3.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+bayesianflow_for_chem-1.4.3.dist-info/METADATA,sha256=Xvq_QOy7kk17IonxrE5onkJrgvWQ-Atcs4zCODpxTXI,5703
+bayesianflow_for_chem-1.4.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bayesianflow_for_chem-1.4.3.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
+bayesianflow_for_chem-1.4.3.dist-info/RECORD,,

bayesianflow_for_chem-1.4.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-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,,