tennetsac 0.1.1__py3-none-any.whl

tennetsac/__init__.py

@@ -0,0 +1 @@
+from .core import profile, binary_lng, multi_lng

tennetsac/core.py

@@ -0,0 +1,88 @@
+# === Standard Library ===
+import os
+import warnings
+
+# === Computing & Visualization ===
+import torch
+import pandas as pd
+import numpy as np
+
+# === Warning Suppression ===
+from tqdm import TqdmWarning
+warnings.simplefilter("ignore", category=TqdmWarning)
+warnings.filterwarnings("ignore", message="TypedStorage is deprecated")
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", message="Some weights of RobertaModel were not initialized")
+from transformers.utils import logging
+logging.set_verbosity_error()
+
+# === Model Architectures ===
+from .models.Emb2Profile import SigmaProfileGenerator
+from .models.Emb2Geometry import GeometryGenerator
+from .models.Prf2Gamma import Prf_to_Seg_Model
+
+# === Model Loading ===
+from .utils.model_io import load_model, load_all_Gamma_models
+
+# === Embedding Extraction ===
+
+from .utils.embedding import ChemBERTaEmbedder, SMITEDEmbedder
+
+# === Computation ===
+from .utils.property import get_sigma_profile, calc_ln_gamma, ensemble_segac, calc_ln_gamma_binary
+
+# === Embedding models ===
+cb_emb = ChemBERTaEmbedder()
+st_emb = SMITEDEmbedder()
+
+# === Load checkpoints ===
+here = os.path.dirname(__file__)  
+ckpt_path = os.path.join(here, "ckpt_files")
+
+prf_model = load_model(SigmaProfileGenerator(), os.path.join(ckpt_path, "prf.ckpt"))
+geometry_model = load_model(GeometryGenerator(), os.path.join(ckpt_path, "geo.ckpt"))
+Gamma_base_model = load_model(Prf_to_Seg_Model(), os.path.join(ckpt_path, "base.ckpt"))
+
+Gamma_finetuned_models = load_all_Gamma_models(Prf_to_Seg_Model, os.path.join(ckpt_path, "fine-tuned"))
+
+# === Define functions ===
+def sigma_profile_wrapper(smiles):
+    return get_sigma_profile(smiles, prf_model, geometry_model, cb_emb, st_emb)
+
+def single_model_predictor(sigma, temperature):
+    return Gamma_base_model(sigma, torch.tensor([temperature]))[1]
+
+def ensemble_predictor(sigma, temperature):
+    return ensemble_segac(Gamma_finetuned_models, sigma, temperature)
+
+def select_gamma_predictor(model_type: str):
+    if model_type == "base":
+        return single_model_predictor
+    elif model_type == "tuned":
+        return ensemble_predictor
+    else:
+        raise ValueError("Invalid model_type. Choose 'base' or 'tuned'.")
+
+def profile(smiles):
+    return sigma_profile_wrapper(smiles)
+
+def binary_lng(smiles:list, temperature:float, molefraction:list):
+    gamma_predictor = select_gamma_predictor("tuned")
+    
+    ln_gamma_1, ln_gamma_2 = calc_ln_gamma_binary(smiles[0], smiles[1], molefraction, temperature,
+                         gamma_predictor=gamma_predictor,
+                         get_sigma_profile_fn=sigma_profile_wrapper)
+    return ln_gamma_1.tolist(), ln_gamma_2.tolist()
+
+def multi_lng(smiles:list, temperature:float, composition:list):
+    gamma_predictor = select_gamma_predictor("tuned")
+    
+    lng_array = calc_ln_gamma(
+    smiles,
+    composition,
+    temperature,
+    gamma_predictor=gamma_predictor,
+    get_sigma_profile_fn=sigma_profile_wrapper
+    )
+    
+    return lng_array

tennetsac/models/Emb2Geometry.py

@@ -0,0 +1,105 @@
+import torch
+import torch.nn as nn
+
+class ResidualBlock(nn.Module): # ResNet v2
+    def __init__(self, in_features, hidden_features=None, activation=nn.GELU):
+        """
+        in_features: input feature
+        hidden_features: hidden feature, if None, use in_features
+        """
+        super(ResidualBlock, self).__init__()
+        hidden_features = hidden_features or in_features
+        self.act = activation()
+        self.bn1 = nn.BatchNorm1d(hidden_features)
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.bn2 = nn.BatchNorm1d(in_features)
+        self.fc2 = nn.Linear(hidden_features, in_features)
+
+    def forward(self, x):
+        identity = x
+
+        out = self.bn1(x)
+        out = self.act(out)
+        out = self.fc1(out)
+
+        out = self.bn2(out)
+        out = self.act(out)
+        out = self.fc2(out)
+
+        out += identity  # add residual
+        return out
+
+class GeometryGenerator(nn.Module):
+    def __init__(self, STemb_dim=768, CBemb_dim=384, CBout_dim=32, av_dim=2,
+                 hidden_dim=128, mw_dim=128, dropout_rate=0.1):
+        """
+        sig_dim: int, dimension of sigma profile input
+        """
+        super(GeometryGenerator, self).__init__()
+
+        self.isomerism_encoder = nn.Sequential(
+          nn.Linear(CBemb_dim, 256),
+          nn.GELU(),
+          nn.Linear(256, CBout_dim)
+        )
+
+        self.backbone = nn.Sequential(
+          nn.Linear(STemb_dim + CBout_dim, 512),
+          nn.GELU(),
+          nn.Linear(512, hidden_dim),
+          nn.GELU(),
+        )
+
+        self.mw_projector = nn.Sequential(
+          nn.Linear(1, 64),
+          nn.GELU(),
+          nn.Linear(64, mw_dim),
+          nn.GELU(),
+        )
+
+        self.comb_layers = nn.Sequential(
+          nn.Linear(hidden_dim + mw_dim, hidden_dim),
+          nn.GELU(),
+          nn.Linear(hidden_dim, hidden_dim),
+          nn.GELU(),
+        )
+
+        self.res_block = ResidualBlock(in_features=hidden_dim, hidden_features=hidden_dim, activation=nn.GELU)
+
+        self.area_head = nn.Sequential(
+            nn.Linear(hidden_dim, 64),
+            nn.GELU(),
+            # nn.Dropout(p=0.1),
+            nn.Linear(64, 16),
+            nn.GELU(),
+            nn.Linear(16, 1),
+        )
+
+        self.volume_head = nn.Sequential(
+            nn.Linear(hidden_dim, 64),
+            nn.GELU(),
+            # nn.Dropout(p=0.1),
+            nn.Linear(64, 16),
+            nn.GELU(),
+            nn.Linear(16, 1),
+        )
+
+    def forward(self, STemb, CBemb, mw):
+        mw = mw.unsqueeze(1)
+        conf_feat = self.isomerism_encoder(CBemb)
+
+        emb = torch.cat([STemb, conf_feat], dim=1)
+
+        hidden_emb = self.backbone(emb)
+        hidden_mw = self.mw_projector(mw)
+
+        hidden_comb = torch.cat([hidden_emb, hidden_mw], dim=1)
+        hidden = self.comb_layers(hidden_comb)
+        hidden = self.res_block(hidden)
+
+        area = self.area_head(hidden)
+        volume = self.volume_head(hidden)
+
+        output = torch.cat([area, volume], dim=1)
+
+        return output

tennetsac/models/Emb2Profile.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+
+class ResidualBlock(nn.Module): # ResNet v2
+    def __init__(self, in_features, hidden_features=None, activation=nn.GELU):
+        """
+        in_features: input feature
+        hidden_features: hidden feature, if None, use in_features
+        """
+        super(ResidualBlock, self).__init__()
+        hidden_features = hidden_features or in_features
+        self.act = activation()
+        self.bn1 = nn.BatchNorm1d(hidden_features)
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.bn2 = nn.BatchNorm1d(in_features)
+        self.fc2 = nn.Linear(hidden_features, in_features)
+
+    def forward(self, x):
+        identity = x
+
+        out = self.bn1(x)
+        out = self.act(out)
+        out = self.fc1(out)
+
+        out = self.bn2(out)
+        out = self.act(out)
+        out = self.fc2(out)
+
+        out += identity  # add residual
+        return out
+
+class SigmaProfileGenerator(nn.Module):
+    def __init__(self, STemb_dim=768, CBemb_dim=384, CBout_dim=32,
+                 prf_dim=51, hidden_dim=256, dropout_rate=0.3):
+        """
+        sig_dim: int, dimension of sigma profile input
+        """
+        super(SigmaProfileGenerator, self).__init__()
+
+        self.isomerism_encoder = nn.Sequential(
+          nn.Linear(CBemb_dim, 256),
+          nn.GELU(),
+          nn.Linear(256, CBout_dim)
+        )
+
+        self.backbone = nn.Sequential(
+          nn.Linear(STemb_dim + CBout_dim, 512),
+          nn.GELU(),
+          nn.Linear(512, 256),
+          nn.GELU(),
+          nn.Linear(256, hidden_dim),
+          nn.GELU(),
+          ResidualBlock(in_features=hidden_dim, hidden_features=hidden_dim, activation=nn.GELU),
+          nn.GELU(),
+          ResidualBlock(in_features=hidden_dim, hidden_features=hidden_dim, activation=nn.GELU),
+          nn.GELU(),
+          nn.Linear(hidden_dim, hidden_dim),
+          nn.GELU()
+        )
+
+        self.head = nn.Sequential(
+            nn.Linear(hidden_dim, 128),
+            nn.GELU(),
+            nn.Dropout(p=dropout_rate),
+            nn.Linear(128, 128),
+            nn.GELU(),
+            nn.Dropout(p=dropout_rate),
+            nn.Linear(128, 128),
+            nn.GELU(),
+            nn.Dropout(p=dropout_rate),
+            nn.Linear(128, 64),
+            nn.GELU(),
+            nn.Linear(64, prf_dim),
+            nn.ReLU()
+        )
+
+    def forward(self, STemb, CBemb):
+        conf_feat = self.isomerism_encoder(CBemb)
+        x = torch.cat([STemb, conf_feat], dim=1)
+        x = self.backbone(x)
+        output = self.head(x)
+        area = output.sum(dim=1, keepdim=True)
+        prf = output / area
+
+        return prf

tennetsac/models/Prf2Gamma.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+import torch.autograd as autograd
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features, hidden_features=None, activation=nn.GELU):
+        """
+        in_features: input feature
+        hidden_features: hidden feature, if None, use in_features
+        """
+        super(ResidualBlock, self).__init__()
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = activation()
+        self.fc2 = nn.Linear(hidden_features, in_features)
+
+    def forward(self, x):
+        identity = x
+        out = self.fc1(x)
+        out = self.act(out)
+        out = self.fc2(out)
+        out += identity  # add residual
+        out = self.act(out)
+        return out
+
+class Prf_to_Seg_Model(nn.Module):
+    def __init__(self, sig_dim=51, temp_hidden_dim=16, sig_hidden_dim=512, hidden_2_dim=256):
+        """
+        sig_dim: int, dimension of sigma profile input
+        """
+        super(Prf_to_Seg_Model, self).__init__()
+        # sigma model
+        self.model_sigma = nn.Sequential(
+            nn.Linear(sig_dim, 512),
+            nn.GELU(),
+            nn.Linear(512, 512),
+            nn.GELU(),
+            nn.Linear(512, sig_hidden_dim),
+            nn.GELU(),
+        )
+
+        self.bn_sig = nn.BatchNorm1d(num_features=sig_hidden_dim)
+
+        self.temp_embedding = nn.Sequential(
+            nn.Linear(1, 32),
+            nn.ReLU(),
+            nn.Linear(32, temp_hidden_dim),
+            nn.ReLU()
+        )
+
+        self.bn_t = nn.BatchNorm1d(num_features=temp_hidden_dim)
+
+        self.model_combined = nn.Sequential(
+            nn.Linear(sig_hidden_dim + temp_hidden_dim, 256),
+            nn.GELU(),
+            nn.Linear(256, hidden_2_dim),
+            nn.GELU(),
+        )
+
+        self.bn2 = nn.BatchNorm1d(num_features=hidden_2_dim)
+
+        # Resudual Block
+        self.res_block = ResidualBlock(in_features=hidden_2_dim, hidden_features=hidden_2_dim, activation=nn.GELU)
+
+        self.model_final = nn.Sequential(
+            nn.Linear(hidden_2_dim, 128),
+            nn.GELU(),
+            nn.Linear(128, 64),
+            nn.GELU(),
+            nn.Linear(64, 1),
+        )
+
+    def forward(self, sigs, t):
+        """
+        sigs: (batch_size, sig_dim)
+        t   : (batch_size,) or (batch_size, 1)
+        """
+        # Ensure sigs require gradients
+        sigs = sigs.requires_grad_(True)
+        # sigs.requires_grad = True
+        t = t.requires_grad_(False)
+
+        # sigma profile
+        sum_sigs = sigs.sum(dim=1, keepdim=True)  # (batch_size, 1)
+        # Send normalized profile into sigma model.
+        sigma_emb = self.model_sigma(sigs/sum_sigs)               # (batch_size, hidden_1_dim)
+        sigma_emb = self.bn_sig(sigma_emb)
+
+        # temperature
+        t_inv = 1.0 / t.view(-1, 1)   # (batch_size, 1)
+        t_emb = self.temp_embedding(t_inv) # (batch_size, temp_hidden_dim)
+        t_emb = self.bn_t(t_emb)
+
+        # Concatenate sigma_emb and t_emb
+        combined_in = torch.cat([sigma_emb, t_emb], dim=-1)
+        combined_out = self.model_combined(combined_in)
+
+        # Batch Normal 2
+        combined_out = self.bn2(combined_out)
+
+        # Residual connection
+        combined_out = self.res_block(combined_out)
+
+        # combined_out_t = torch.cat([combined_out, t_inv], dim=-1)
+
+        # Final FNN
+        gchg_part = self.model_final(combined_out)
+
+        # gchg = sum(sigs) * gchg_part 
+        gchg = sum_sigs * gchg_part                 # (batch_size, 1)
+
+        # Summing gchg and compute the gradient
+        grad = autograd.grad(
+            gchg.sum(),  # -> scalar
+            sigs,
+            create_graph=self.training
+          )[0]  # shape: (batch_size, sig_dim)
+        segs = grad
+        return gchg, segs
+

tennetsac/utils/embedding.py

@@ -0,0 +1,50 @@
+import os
+from transformers import RobertaTokenizer, RobertaModel
+from .smiles import canonicalize_smiles, compute_mol_weight
+import torch
+
+# === ChemBERTa2 Embedding  ===
+class ChemBERTaEmbedder:
+    def __init__(self, model_name="DeepChem/ChemBERTa-77M-MLM", max_length=128, device="cpu"):
+        self.tokenizer = RobertaTokenizer.from_pretrained(model_name)
+        self.model = RobertaModel.from_pretrained(model_name).to(device)
+        self.max_length = max_length
+        self.device = device
+
+    def __call__(self, smiles):
+        tokens = self.tokenizer(
+            smiles, return_tensors="pt", padding="max_length",
+            truncation=True, max_length=self.max_length
+        ).to(self.device)
+        with torch.no_grad():
+            output = self.model(tokens["input_ids"], tokens["attention_mask"])
+            hidden = output.last_hidden_state
+            mask = tokens["attention_mask"].unsqueeze(-1)
+            avg_emb = (hidden * mask).sum(dim=1) / mask.sum(dim=1)
+        return avg_emb
+
+# === SMI-TED Embedding  ===
+from .smi_ted_light.load import load_smi_ted
+
+class SMITEDEmbedder:
+    def __init__(self, ckpt_name="smi-ted-Light_40.pt", device="cpu"):
+        here = os.path.dirname(__file__)
+        model_dir = os.path.join(here, "smi_ted_light")
+        
+        self.model = load_smi_ted(folder=model_dir, ckpt_filename=ckpt_name).to(device)
+        self.device = device
+
+    def __call__(self, smiles):
+        with torch.no_grad():
+            emb = self.model.encode(smiles, return_torch=True).to("cpu")  # always return CPU tensor
+        return emb
+
+# === function  ===
+def get_input_embeddings(smiles, cb_embedder, st_embedder):
+    isomeric_cansmi = canonicalize_smiles(smiles, isomeric=True)
+    nonisomeric_cansmi = canonicalize_smiles(smiles, isomeric=False)
+
+    cb_emb = cb_embedder(isomeric_cansmi)
+    st_emb = st_embedder(nonisomeric_cansmi)
+    
+    return st_emb, cb_emb

tennetsac/utils/model_io.py

@@ -0,0 +1,19 @@
+import torch
+from pathlib import Path
+
+def load_model(model, ckpt_file, device='cpu'):
+    model.to(torch.device(device))
+    model.load_state_dict(torch.load(ckpt_file, map_location=torch.device(device)))
+    
+    return model.eval()
+
+def load_all_Gamma_models(model_class, ckpt_dir, num_models=10):
+    models = []
+    for i in range(1, num_models + 1):
+        ckpt_file = Path(ckpt_dir) / f"{i}.ckpt"
+        model = model_class()
+        model.to(torch.device('cpu'))
+        model.load_state_dict(torch.load(ckpt_file, map_location=torch.device('cpu')))
+        model.eval()
+        models.append(model)
+    return models

tennetsac/utils/plotting.py

@@ -0,0 +1,82 @@
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+import numpy as np
+import pandas as pd
+mpl.rcParams.update(mpl.rcParamsDefault)
+
+def plot_sigma_prf(smiles, prf_tensor, savepath=None):
+    prf = prf_tensor.squeeze().detach().cpu().numpy()
+    sigma = np.linspace(-0.025, 0.025, 51)
+    plt.figure(figsize=(6, 4))
+    plt.plot(sigma, prf, linestyle='-', linewidth=2)
+    
+    plt.xlim(-0.025, 0.025)
+    plt.xlabel("σ (e/Ų)", fontsize=12)
+    plt.ylabel("P(σ)", fontsize=12)
+    plt.title(f"σ-Profile: {smiles}", fontsize=13)
+    
+    plt.xticks(fontsize=10)
+    plt.yticks(fontsize=10)
+    
+    plt.tight_layout()
+
+    if savepath:
+        plt.savefig(savepath, dpi=300, bbox_inches="tight")
+    plt.show()
+    
+
+def plot_sigma_prf_pair(smiles_1, prf_tensor_1, smiles_2, prf_tensor_2, savepath=None):
+    prf_1 = prf_tensor_1.squeeze().detach().cpu().numpy()
+    prf_2 = prf_tensor_2.squeeze().detach().cpu().numpy()
+    sigma = np.linspace(-0.025, 0.025, 51)
+
+    fig, axs = plt.subplots(1, 2, figsize=(10, 4), sharey=False)
+
+    axs[0].plot(sigma, prf_1, linestyle='-', linewidth=2)
+    axs[0].set_title(f"σ-Profile: {smiles_1}", fontsize=13)
+    axs[0].set_xlabel("σ (e/Ų)", fontsize=12)
+    axs[0].set_ylabel("P(σ)", fontsize=12)
+    axs[0].tick_params(labelsize=10)
+
+    axs[1].plot(sigma, prf_2, linestyle='-', linewidth=2, color='tab:orange')
+    axs[1].set_title(f"σ-Profile: {smiles_2}", fontsize=13)
+    axs[1].set_xlabel("σ (e/Ų)", fontsize=12)
+    axs[1].tick_params(labelsize=10)
+
+    plt.tight_layout()
+    if savepath:
+        plt.savefig(savepath, dpi=300, bbox_inches="tight")
+    plt.show()
+
+def plot_binary_lng(x1_list, temperature, ln_gamma_1, ln_gamma_2, smiles_1, smiles_2, savepath=None):
+    plt.figure(figsize=(6, 4))
+    plt.plot(x1_list, ln_gamma_1, label='ln γ₁', linewidth=2)
+    plt.plot(x1_list, ln_gamma_2, label='ln γ₂', linewidth=2)
+    
+    plt.xlabel("x₁", fontsize=12)
+    plt.ylabel("ln γ", fontsize=12)
+    plt.title(f"Binary Mixture, T = {temperature} K:\n1: {smiles_1}, 2: {smiles_2}", fontsize=13)
+    plt.xticks(fontsize=10)
+    plt.yticks(fontsize=10)
+    plt.legend(fontsize=11)
+
+    plt.xlim(0, 1)
+    y_min = min(np.min(ln_gamma_1), np.min(ln_gamma_2)) - 0.01
+    y_max = max(np.max(ln_gamma_1), np.max(ln_gamma_2)) + 0.01
+    plt.ylim(y_min, y_max)
+
+    plt.tight_layout()
+    if savepath:
+        plt.savefig(savepath, dpi=300, bbox_inches="tight")
+    plt.show()
+
+def make_binary_df(x1_list, ln_gamma_1, ln_gamma_2, smiles_1, smiles_2, temperature):
+    df = pd.DataFrame({
+        "smiles_1": smiles_1,
+        "smiles_2": smiles_2,
+        "temperature (K)": temperature,
+        "x1": x1_list,
+        "ln_gamma_1": ln_gamma_1,
+        "ln_gamma_2": ln_gamma_2
+    })
+    return df