heiwa 0.1.0__tar.gz

heiwa-0.1.0/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.4
+Name: heiwa
+Version: 0.1.0
+Requires-Dist: biotite
+Requires-Dist: biopython
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: torch>=2.5.1
+Requires-Dist: tqdm
+Requires-Dist: esm==3.2.1
+Requires-Dist: PyCifRW
+Requires-Dist: transformers
+Requires-Dist: tokenizers

heiwa-0.1.0/pyproject.toml

@@ -0,0 +1,40 @@
+[build-system]
+requires = ["setuptools>=75"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "heiwa"
+version = "0.1.0"
+
+dependencies = [
+    "biotite", 
+    "biopython", 
+    "numpy",
+    "pandas",
+    "torch>=2.5.1",
+    "tqdm", 
+    "esm==3.2.1",
+    "PyCifRW",
+    "transformers",
+    "tokenizers"
+]
+
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+heiwa = ["model_assets/*"]
+
+# requires-python = ">=3.11"
+
+
+# - python
+# - numpy
+# - pandas
+# - pytorch>=2.5.1
+# - pytorch-cuda>=11.8
+# - tqdm
+# - pycifrw
+# - transformers==4.46.3
+# - tokenizers==0.20

heiwa-0.1.0/setup.cfg

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

heiwa-0.1.0/src/heiwa/__init__.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2026 Ernest Ngai Hei Ho
+# Email: hoernesta@gmail.com
+# Licensed under the MIT License.
+__version__="0.1.0"
+from .loaders.model_loader import load_solv_encoder, load_thermo
+from .hardware import Hardware_Setting
+# from .esm_gen import gen_embed
+from .predict import predict
+from .generate.generate import de_novo, gen_embed
+
+
+## Replace esm3 script during import
+import shutil
+from importlib.resources import files
+
+_source_folder="replace_esm3"
+_package_path = files("heiwa")
+shutil.copy(
+    _package_path.joinpath(_source_folder, "esm3.py"), 
+    _package_path.parent.joinpath("esm/models/esm3.py")
+)
+shutil.copy(
+    _package_path.joinpath(_source_folder, "generation.py"), 
+    _package_path.parent.joinpath("esm/utils/generation.py")
+)
+# shutil.copy(
+#     _package_path.joinpath(_source_folder, "sequence_tokenizer.py"), 
+#     _package_path.parent.joinpath("esm/tokenization/sequence_tokenizer.py")
+# )
+# shutil.copy(
+#     _package_path.joinpath(_source_folder, "encoding.py"), 
+#     _package_path.parent.joinpath("esm/utils/encoding.py")
+# )
+shutil.copy(
+    _package_path.joinpath(_source_folder, "transformer_stack.py"), 
+    _package_path.parent.joinpath("esm/layers/transformer_stack.py")
+)
+shutil.copy(
+    _package_path.joinpath(_source_folder, "blocks.py"), 
+    _package_path.parent.joinpath("esm/layers/blocks.py")
+)
+
+del _source_folder, _package_path

heiwa-0.1.0/src/heiwa/batch_run.py

@@ -0,0 +1,67 @@
+import torch
+from heiwa.loaders.tensor_loader import concat_tensors_and_pad
+from heiwa.loaders.dataloader import compile_data_single
+from heiwa.hardware import Hardware_Setting
+from heiwa.model_run.epoch_run import EpochRunner
+from heiwa.utils import ensure_3d_tensor
+
+def batch_run(
+    x:torch.Tensor|list[torch.Tensor], 
+    hw_setting:Hardware_Setting|None=None, 
+    heiwa_model=None, 
+    heiwa_orthogonalizer=None, 
+    thermo_model=None, 
+    out_folder=".", 
+):
+    # assert target.lower() in ("rmsf", "thermo")
+    
+    # ========= Prepare dataset here =========
+    # m = None
+    # l = None
+    if isinstance(x, list):
+        print("hi")
+        dataset = concat_tensors_and_pad(
+            embeds=x, 
+            # create_l=target.lower()=="thermo", 
+            hw_setting=hw_setting, 
+        )
+        # print("x_unorthogonalized_list", len(x), x)
+        # print("x_unorthogonalized", dataset["X"].size(), dataset["X"])
+        # print("l", l.size(), l)
+    else:
+        print("hi single x")
+        x = ensure_3d_tensor(x=x)
+        # if target.lower()=="thermo":
+        #     print(f"Task is definitely thermo")
+        l = torch.tensor(
+            [x.size(1)], 
+            dtype=torch.float32, 
+            device='cpu'
+        ) 
+        dataset={"X":x, "L":l}
+        print("x_unorthogonalized", x.size(), x)
+    print("dataset", dataset.keys())
+
+    # ========= Prepare models, hardware, epoch_runner =========
+    assert heiwa_model is not None and thermo_model is not None
+
+    hw_setting.switch_on_gpu_rs()
+
+    epoch_runner = EpochRunner(out_folder=out_folder)
+    
+
+    data_loader = compile_data_single(
+        hw_setting=hw_setting, 
+        dataset=dataset, 
+    )
+        
+    # ========= Run ==========================
+    print(f"Start running batches")
+    
+    # print("dataset", data_loader.dataset)
+    epoch_runner.run(
+        data_source=data_loader,
+        heiwa_model=heiwa_model,
+        heiwa_orthogonalizer=heiwa_orthogonalizer, 
+        thermo_model=thermo_model, 
+    )

heiwa-0.1.0/src/heiwa/generate/esm3_gen.py

@@ -0,0 +1,159 @@
+import os
+import torch
+import torch.nn as  nn
+import pickle
+import pandas as pd
+# import traceback
+from heiwa.readers.reader import read_file
+from heiwa.generate.generate_utils import convert_idx_to_prompt
+from esm.sdk.api import GenerationConfig, ESMProtein
+# from heiwa.generate.modeller_script import fill_atom_to_pdb_in_single
+
+tracks = [
+    {'structure':16},
+    {'secondary_structure':8}, 
+    {'sasa':4}, 
+    {'structure':8}, 
+    {'sequence':8}
+]
+
+def esm3_generation(
+    real_tracks, protein, 
+    esm3_model, heiwa_model, heiwa_orthogonalizer, 
+    seq_prompt, crit_resi, 
+    out_filename
+):
+    if protein == "max_len_exceeded":
+        pass
+    fname = None
+    for i, track in enumerate(real_tracks):
+        # print(track)
+        if i == len(real_tracks)-1:
+            fname = out_filename
+            
+        track_k = list(track.keys())[0]
+        track_v = list(track.values())[0]
+        # print(track)
+        # print(fname)
+        condition_on_coordinates_only = True if track_k == 'structure' else False
+        
+        if track_k=='sequence':
+            protein.sequence = seq_prompt
+            # print(seq_prompt)
+        elif track_k=='structure' and track_v == 8:
+            protein = esm3_model.generate(protein, GenerationConfig(track=track_k, num_steps=track_v, \
+                                                               temperature=0.7, 
+                                                               temperature_annealing=True, \
+                                                               condition_on_coordinates_only = condition_on_coordinates_only), \
+                                     filename=fname, 
+                                     layer_i=-1, 
+                                     heiwa_model=heiwa_model, 
+                                     heiwa_normalizer=heiwa_orthogonalizer)
+            continue
+        # print(len(protein.sequence), len(protein.coordinates), protein.sequence)
+        protein = esm3_model.generate(protein, GenerationConfig(track=track_k, num_steps=track_v, \
+                                                           temperature=0.7, 
+                                                           temperature_annealing=True, \
+                                                           condition_on_coordinates_only = condition_on_coordinates_only), \
+                                 filename=fname, 
+                                 layer_i=-1)
+    return protein
+    
+def check_generation_success(outPath:str|list[str]):
+    """Check where pdb and pkl exists at the right path"""
+    if isinstance(outPath, str):
+        with open(f"./generation_report.txt", "a+") as file:
+            if os.path.exists(outPath):
+                file.write(f"Generation of {outPath} successful.\n")
+            else:
+                file.write(f"Generation of {outPath} failed.\n")
+    elif isinstance(outPath, list):
+        with open(f"./generation_report.txt", "a+") as file:
+            for path in outPath:
+                if os.path.exists(outPath):
+                    file.write(f"Generation of {path} successful.\n")
+                else:
+                    file.write(f"Generation of {path} failed.\n")
+
+def save_protein(protein, folder_out, out_file_name):
+    with open(f"{folder_out}/{out_file_name}.pkl", 'wb') as file:
+        pickle.dump(protein, file)
+    protein.to_pdb(f"{folder_out}/{out_file_name}.pdb")
+
+def generate_proteins(
+    task, 
+    hw_setting, 
+    esm3_model, 
+    heiwa_model=None, heiwa_orthogonalizer=None, seq_prompt=None, crit_resi=None, idx_style:int|None=None, 
+    inFile:str|None=None, 
+    protein:ESMProtein|None=None, 
+    cleave:int=0, 
+    out_file_name:str|None=None, 
+    folder_out=".", \
+    debug_path: str="", 
+):
+    assert task in ("gen_embed", "gen_de_novo")
+    assert esm3_model
+    assert out_file_name
+    stream = torch.cuda.Stream(device="cuda")
+    if heiwa_model:
+        assert heiwa_orthogonalizer
+        
+    if task == "gen_embed":
+        out_file_type="pt"
+        real_tracks = tracks[:-1]
+    else:
+        assert heiwa_model and heiwa_orthogonalizer
+        assert seq_prompt is not None or (crit_resi is not None and idx_style is not None)
+        out_file_type="pdb"
+        real_tracks = tracks
+        
+
+    if inFile:
+        protein = read_file(
+            filename=inFile, 
+            cleave=cleave, 
+            debug_path=debug_path, 
+        )
+        # print("b", type(protein))
+
+    if task=="gen_de_novo" and seq_prompt is None:
+        seq_prompt = convert_idx_to_prompt(
+            idx=crit_resi, 
+            cleave=cleave, 
+            seq_wt=protein.sequence, 
+            length=len(protein.sequence), 
+            idx_style=idx_style, 
+        )
+        # print(seq_prompt)
+
+    hw_setting.switch_on_gpu_rs()
+    
+    if out_file_name.count('.') > 0:
+        out_file_name = out_file_name[:out_file_name.rfind('.')]
+    if not os.path.exists(f"{folder_out}/{out_file_name}.{out_file_type}"):
+        protein = esm3_generation(
+            real_tracks=real_tracks, 
+            protein=protein, 
+            seq_prompt=seq_prompt, 
+            crit_resi=crit_resi, 
+            esm3_model=esm3_model, heiwa_model=heiwa_model, heiwa_orthogonalizer=heiwa_orthogonalizer, 
+            out_filename=f"{folder_out}/{out_file_name}"
+        )
+        
+        if task=="gen_de_novo":
+            save_protein(protein=protein, 
+                         folder_out=folder_out, 
+                         out_file_name=out_file_name)
+                
+            # fill_atom_to_pdb_in_single(
+            #     base_name=file_name, 
+            #     folder_out=folder_out, 
+            #     sequence=protein.sequence, 
+            #     length=len(protein.sequence)
+            # )
+        
+
+        
+
+    

heiwa-0.1.0/src/heiwa/generate/generate.py

@@ -0,0 +1,105 @@
+import torch.nn as nn
+from esm.sdk.api import ESMProtein
+from heiwa.generate.esm3_gen import generate_proteins
+from heiwa.utils import ensure_dirs
+
+def gen_embed(
+    esm3_model, 
+    hw_setting, 
+    inFile:str|None=None, 
+    protein:ESMProtein|None=None, 
+    cleave:int=0, 
+    outPath:str="./no_name.pt", 
+    debug_path:str="", 
+):
+    """1 protein at once"""
+    if outPath.count('/') > 0:
+        sep = outPath.rfind('/')
+        folder_out = outPath[:sep]
+        outPath  = outPath[sep+1:]
+    else:
+        folder_out = "."
+    ensure_dirs(folder_out)
+
+    if inFile:
+        generate_proteins(
+            task="gen_embed", 
+            hw_setting=hw_setting, 
+            inFile=inFile, 
+            out_file_name=outPath, 
+            cleave=cleave, 
+            esm3_model=esm3_model, 
+            folder_out=folder_out, 
+            debug_path=debug_path, 
+        )
+    else:
+        assert protein
+        generate_proteins(
+            task="gen_embed", 
+            hw_setting=hw_setting, 
+            protein=protein, 
+            cleave=cleave, 
+            out_file_name=outPath, 
+            esm3_model=esm3_model, 
+            folder_out=folder_out, 
+            debug_path=debug_path, 
+        )
+
+def de_novo(
+    esm3_model, 
+    hw_setting, 
+    heiwa_model:nn.Module|None=None, 
+    heiwa_orthogonalizer:nn.Module|None=None, 
+    inFile:str|None=None, 
+    seq_prompt:str|None=None, 
+    crit_resi:list[int]|None=None, 
+    idx_style:int|None=None, 
+    cleave:int=0, 
+    protein=None, 
+    outPath:str="./no_name.pdb", 
+    debug_path:str="", 
+):
+    """1 protein at once"""
+    assert not (seq_prompt is not None and (crit_resi is not None or idx_style))
+    
+    if outPath.count('/') > 0:
+        sep = outPath.rfind('/')
+        folder_out = outPath[:sep]
+        outPath  = outPath[sep+1:]
+    else:
+        folder_out = "."
+    ensure_dirs(folder_out)
+
+    if inFile:
+        generate_proteins(
+            task="gen_de_novo", 
+            hw_setting=hw_setting, 
+            inFile=inFile, 
+            cleave=cleave, 
+            seq_prompt=seq_prompt, 
+            crit_resi=crit_resi, 
+            idx_style=idx_style, 
+            out_file_name=outPath, 
+            esm3_model=esm3_model, 
+            heiwa_model=heiwa_model, 
+            heiwa_orthogonalizer=heiwa_orthogonalizer, 
+            folder_out=folder_out, 
+            debug_path=debug_path, 
+        )
+    else:
+        assert protein
+        generate_proteins(
+            task="gen_de_novo", 
+            hw_setting=hw_setting, 
+            protein=protein, 
+            cleave=cleave, 
+            seq_prompt=seq_prompt, 
+            crit_resi=crit_resi, 
+            idx_style=idx_style, 
+            out_file_name=outPath, 
+            esm3_model=esm3_model, 
+            heiwa_model=heiwa_model, 
+            heiwa_orthogonalizer=heiwa_orthogonalizer, 
+            folder_out=folder_out, 
+            debug_path=debug_path, 
+        )

heiwa-0.1.0/src/heiwa/generate/generate_utils.py

@@ -0,0 +1,18 @@
+
+def convert_idx_to_prompt(
+    idx:list[int]|tuple[int], # only this is affected by signal peptide
+    seq_wt:str, 
+    length:int, 
+    cleave:int=0, 
+    idx_style=0, 
+):
+    seq_prompt = ""
+    i_last=0
+    # print(seq_wt)
+    idx = sorted([i-cleave for i in idx])
+    for i in idx:
+        # print(i-i_last-idx_style, seq_wt[i-idx_style])
+        seq_prompt += "_"*(i-i_last-idx_style) + seq_wt[i-idx_style]
+        i_last = i
+    seq_prompt += "_"*(length - i_last + 1 - idx_style)
+    return seq_prompt

heiwa-0.1.0/src/heiwa/generate/modeller_script.py

@@ -0,0 +1,89 @@
+from modeller import *
+from modeller.automodel import *
+from modeller.parallel import *
+import os
+
+env = Environ()
+default_output = "target.B99990001.pdb"
+
+def alignment_writer(sequence, length, base_name, folder_out):
+    """Write alignment file for MODELLER
+    Args:
+        row: pandas DataFrame row containing 'Sequence_final', 'Filename', 'Y_thermo'
+    """
+    file_out_name = f"{folder_out}/{base_name}.ali"
+    
+    # Prepare the alignment content
+    alignment_content = f""">P1;backbone
+structureX:{base_name}.pdb:1:A:{length}:A::::
+{'-' * length}*
+
+>P1;target
+sequence:target::::::::
+{sequence}*
+"""
+    
+    # Write the alignment file
+    with open(file_out_name, 'w') as f:
+        f.write(alignment_content)
+    
+    print(f"Created alignment file: {file_out_name}")
+    return file_out_name
+
+
+def fill_atom_to_pdb(base_name, folder_out):
+    """Build full-atom model from backbone-only PDB
+    Args:
+        row: pandas DataFrame row containing 'Filename', 'Y_thermo'
+    """
+    pdb_path = f"{folder_out}/{base_name}.pdb"
+    ali_file = f"{folder_out}/{base_name}.ali"
+    ali_final = f"{folder_out}/{base_name}_final.ali"
+    
+    # Read the template structure
+    aln = Alignment(env)
+    mdl = Model(env, file=pdb_path.replace('.pdb', ''), model_segment=('FIRST:A', 'LAST:A'))
+    aln.append_model(mdl, align_codes='backbone', atom_files=pdb_path)
+    aln.append(file=ali_file, align_codes='target')
+    aln.align2d()
+    aln.write(file=ali_final, alignment_format='PIR')
+    
+    # Build the model with side chains
+    a = AutoModel(env, 
+                  alnfile=ali_final,
+                  knowns='backbone', 
+                  sequence='target')
+    a.starting_model = 1
+    a.ending_model = 1
+    
+    # Set output directory
+    a.outputs = ['MODELS']  # Only output final model
+    
+    a.make()
+    
+    # Rename output to something more meaningful
+    # MODELLER creates: target.B99990001.pdb
+    final_output = f"{folder_out}/{base_name}_full.pdb"
+    
+    if os.path.exists(default_output):
+        os.rename(default_output, final_output)
+        print(f"Created full-atom model: {final_output}")
+        for path in [pdb_path, ali_file, ali_final]:
+            os.remove(path)
+    else:
+        print(f"Warning: Expected output {default_output} not found!")
+    
+    return final_output
+
+def fill_atom_to_pdb_in_single(base_name, folder_out, sequence, length):
+    alignment_writer(
+        sequence=sequence, length=length, base_name=base_name, folder_out=folder_out
+    )
+    print("wrote alignment")
+    
+    fill_atom_to_pdb(
+        base_name, folder_out
+    )
+    print("pdb filled")
+
+

heiwa-0.1.0/src/heiwa/hardware.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+import torch
+# import numpy as np
+
+@dataclass
+class Hardware_Setting:
+    cpu: int = 1
+    batch_size : int = 2
+    gpu_rs: int|None = None
+    # cpu_rs: int
+
+    def __str__(self):
+        cout = "HEIWA hardware setting:\n"
+        cout += f"cpu = {self.cpu}\n"
+        cout += f"batch_size = {self.batch_size}\n"
+        cout += f"gpu_rs = {self.gpu_rs}\n"
+        # cout += f"cpu_rs = {self.cpu_rs}\n"
+        return cout
+
+    def switch_on_gpu_rs(self):
+        if self.gpu_rs is not None:
+            torch.manual_seed(self.gpu_rs)
+        
+    

heiwa-0.1.0/src/heiwa/loaders/dataloader.py

@@ -0,0 +1,71 @@
+import torch
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+from typing import Tuple, Dict
+from heiwa.hardware import Hardware_Setting
+
+class MyDataset(Dataset):
+    def __init__(self, prefix):
+        """ The real initlization is called by cls._from_dict """
+        self.DATAFILE_PREFIX = tuple(prefix)
+        print("Tensor prefix", self.DATAFILE_PREFIX)
+        
+    @classmethod
+    def _from_dict(cls, all_inputs:dict):
+        """ This triggers __init__ """
+        dataset = MyDataset(prefix=all_inputs.keys())
+        for key, value in all_inputs.items():
+            setattr(dataset, key, value)
+            if key == "X":
+                print(f"{key}: {len(value)}")
+            else:
+                print(f"{key}: {value}")
+        return dataset
+
+    def __str__(self):
+        cout = ""
+        for k in "XLM":
+            cout += f"{k}:{getattr(self, k, None)}\n"
+        return cout
+        
+    def __len__(self) -> int:
+        if hasattr(self, "X"):
+            return len(self.X)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        output = {}
+        output['model_input'] = {k: getattr(self, k)[idx] for k in self.DATAFILE_PREFIX}
+        return output
+
+
+
+def compile_data_single(
+    hw_setting:Hardware_Setting, 
+    dataset:dict
+) -> DataLoader:
+    """     Create DataLoader objects for train, validation, and test sets.
+    
+    Args:
+        target: 'test', 'inference', 'train' if fisher
+        dataset: dictionary of individual raw tensors
+        params: all Parameters
+    
+    Returns:
+        output_dataset      """
+    # Create datasets
+    # print(dataset)
+    output_dataset = MyDataset._from_dict(
+        dataset # {k:dataset[k] for k in dataset}
+    ) # 
+        
+    output_dataset = DataLoader(
+        output_dataset,
+        batch_size=hw_setting.batch_size,
+        shuffle=False,
+        num_workers=hw_setting.cpu,
+        pin_memory=True
+    )
+    print(f"Dataloader for compiled")
+    return output_dataset
+    

heiwa-0.1.0/src/heiwa/loaders/model_loader.py

@@ -0,0 +1,31 @@
+import torch
+from heiwa.models.model import HEIWA3H
+from heiwa.models.orthogonalizer import HEIWA_Orthogonalizer
+from heiwa.models.part_thermo import ThermoPred
+from heiwa.utils import asset_path
+
+def load_esm3():
+    with torch.no_grad():
+        return torch.load(
+            asset_path.joinpath("esm3.pth"), 
+            map_location="cuda", 
+            weights_only=False
+        )
+
+def load_solv_encoder(hw_setting):
+    with torch.no_grad():
+        heiwa = HEIWA3H(hw_setting=hw_setting)
+    
+        # heiwa = torch.load(
+        #     path, 
+        #     map_location="cuda", 
+        #     weights_only=False
+        # )
+    
+        heiwa.to("cuda")
+        orthogonalizer = HEIWA_Orthogonalizer()
+        return heiwa, orthogonalizer
+
+def load_thermo(hw_setting):
+    thermo = ThermoPred(hw_setting=hw_setting)
+    return thermo