tennetsac 0.1.1__tar.gz

tennetsac-0.1.1/PKG-INFO

@@ -0,0 +1,55 @@
+Metadata-Version: 2.4
+Name: tennetsac
+Version: 0.1.1
+Summary: Thermodynamics-embedded Neural Network for Segment Activity Coefficients
+Author: Yue Yang
+Description-Content-Type: text/markdown
+
+# TeNNet-SAC
+
+TeNNet-SAC (Thermodynamics-Embedded Neural Network for Segment Activity Coefficients) is a machine learning framework designed to predict molecular activity coefficients in multicomponent systems using only molecular SMILES strings, composition, and temperature as input. 
+
+This project provides:
+
+1. **σ-profile prediction model**, including surface area and molecular volume estimation  
+2. **Activity coefficient prediction** with two versions:  
+   - **Base model**: trained on synthetic data generated from the COSMO-SAC model  
+   - **Fine-tuned model**: further optimized using high-quality experimental data  
+
+## Features
+
+- Predicts activity coefficients beyond binary systems
+- Requires only SMILES strings, mole fractions, and temperature
+- Hard-constraint architecture ensures thermodynamic consistency  
+- Modular design: supports use of σ-profiles from QC calculations  
+- Robust and generalizable via two-stage training: synthetic COSMO-SAC pretraining followed by experimental fine-tuning, preserving physical consistency
+
+## Citation
+
+Yue Yang, Shiang-Tai Lin. *Physics-Embedded Machine Learning Model for Phase Equilibrium Prediction in Multicomponent Systems*. *Journal of Chemical Information and Modeling*, 2025. [DOI: 10.1021/acs.jcim.5c01804](https://doi.org/10.1021/acs.jcim.5c01804)
+
+## References
+
+This project builds upon the following foundational models. If you use this project in your research, we encourage you to cite them as well:
+
+- **ChemBERTa-2**  
+Ahmad, W.; Simon, E.; Chithrananda, S.; Grand, G.; Ramsundar, B. Chemberta-2: Towards chemical foundation models. arXiv preprint arXiv:2209.01712 2022.
+[https://arxiv.org/abs/2209.01712](https://arxiv.org/abs/2209.01712)
+
+- **SMI-TED**  
+Soares, E.; Shirasuna, V.; Brazil, E. V.; Cerqueira, R.; Zubarev, D.; Schmidt, K. A large encoder-decoder family of foundation models for chemical language. arXiv preprint arXiv:2407.20267 2024.
+[https://arxiv.org/abs/2407.20267](https://arxiv.org/abs/2407.20267)
+
+### External Code Acknowledgment
+
+The folder `smi_ted_light/` is adapted from the [SMI-TED](https://github.com/IBM/materials/tree/main/models/smi_ted) repository by Soares et al., with only minimal modifications. The core implementation remains unchanged. Full credit goes to the original authors.
+
+## License
+
+MIT License. See [LICENSE](LICENSE) for details.
+
+---
+
+Maintained by **Yue Yang** ([@yueyue2299](https://github.com/yueyue2299)).
+
+COMET, Department of Chemical Engineering, National Taiwan University  

tennetsac-0.1.1/README.md

@@ -0,0 +1,48 @@
+# TeNNet-SAC
+
+TeNNet-SAC (Thermodynamics-Embedded Neural Network for Segment Activity Coefficients) is a machine learning framework designed to predict molecular activity coefficients in multicomponent systems using only molecular SMILES strings, composition, and temperature as input. 
+
+This project provides:
+
+1. **σ-profile prediction model**, including surface area and molecular volume estimation  
+2. **Activity coefficient prediction** with two versions:  
+   - **Base model**: trained on synthetic data generated from the COSMO-SAC model  
+   - **Fine-tuned model**: further optimized using high-quality experimental data  
+
+## Features
+
+- Predicts activity coefficients beyond binary systems
+- Requires only SMILES strings, mole fractions, and temperature
+- Hard-constraint architecture ensures thermodynamic consistency  
+- Modular design: supports use of σ-profiles from QC calculations  
+- Robust and generalizable via two-stage training: synthetic COSMO-SAC pretraining followed by experimental fine-tuning, preserving physical consistency
+
+## Citation
+
+Yue Yang, Shiang-Tai Lin. *Physics-Embedded Machine Learning Model for Phase Equilibrium Prediction in Multicomponent Systems*. *Journal of Chemical Information and Modeling*, 2025. [DOI: 10.1021/acs.jcim.5c01804](https://doi.org/10.1021/acs.jcim.5c01804)
+
+## References
+
+This project builds upon the following foundational models. If you use this project in your research, we encourage you to cite them as well:
+
+- **ChemBERTa-2**  
+Ahmad, W.; Simon, E.; Chithrananda, S.; Grand, G.; Ramsundar, B. Chemberta-2: Towards chemical foundation models. arXiv preprint arXiv:2209.01712 2022.
+[https://arxiv.org/abs/2209.01712](https://arxiv.org/abs/2209.01712)
+
+- **SMI-TED**  
+Soares, E.; Shirasuna, V.; Brazil, E. V.; Cerqueira, R.; Zubarev, D.; Schmidt, K. A large encoder-decoder family of foundation models for chemical language. arXiv preprint arXiv:2407.20267 2024.
+[https://arxiv.org/abs/2407.20267](https://arxiv.org/abs/2407.20267)
+
+### External Code Acknowledgment
+
+The folder `smi_ted_light/` is adapted from the [SMI-TED](https://github.com/IBM/materials/tree/main/models/smi_ted) repository by Soares et al., with only minimal modifications. The core implementation remains unchanged. Full credit goes to the original authors.
+
+## License
+
+MIT License. See [LICENSE](LICENSE) for details.
+
+---
+
+Maintained by **Yue Yang** ([@yueyue2299](https://github.com/yueyue2299)).
+
+COMET, Department of Chemical Engineering, National Taiwan University  

tennetsac-0.1.1/pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+package-dir = {"" = "."}
+
+[tool.setuptools.packages.find]
+include = ["tennetsac*"]
+
+[tool.setuptools.package-data]
+tennetsac = [
+    "utils/smi_ted_light/*.txt",
+    "ckpt_files/*.ckpt",
+    "ckpt_files/fine-tuned/*.ckpt"
+]
+
+[project]
+name = "tennetsac"
+version = "0.1.1"
+description = "Thermodynamics-embedded Neural Network for Segment Activity Coefficients"
+authors = [{ name="Yue Yang" }]
+readme = "README.md"
+license = { file = "MIT" }
+

tennetsac-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tennetsac-0.1.1/tennetsac/__init__.py

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

tennetsac-0.1.1/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-0.1.1/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-0.1.1/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-0.1.1/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-0.1.1/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-0.1.1/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