sequence_embedding.py

#Import all the packages.

import sys; print(sys.version)
import tensorflow as tf
from tensorflow import keras
print(tf.__version__)
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.io as pio
pio.renderers.default = 'plotly_mimetype'
from collections import Counter
import os
import xml.etree.ElementTree as ET
import lxml
import collections
from lxml import etree

import fasttext
from collections import defaultdict
from functools import partial
from itertools import repeat
def nested_defaultdict(default_factory, depth=1):
    result = partial(defaultdict, default_factory)
    for _ in repeat(None, depth - 1):
        result = partial(defaultdict, result)
    return result()
    
import os, psutil
np.set_printoptions(linewidth=15000)
def printmem():
    process = psutil.Process(os.getpid())
    print(round(process.memory_info().rss/(10**9),3),'Gbytes')  # in bytes 


from sklearn import metrics
from sklearn.metrics import classification_report
from sklearn.metrics import confusion_matrix,ConfusionMatrixDisplay
from sklearn.metrics import auc

from tabulate import tabulate
import matplotlib.pyplot as plt

import csv
import gzip
import re
import io
import json

import time
from tqdm import tqdm

import requests

import fastaparser
from Bio import SeqIO
from Bio.SeqUtils.ProtParam import ProteinAnalysis

import rdkit
from rdkit import Chem
from rdkit.Chem import Descriptors
from rdkit.Chem import Draw
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import IPythonConsole
from rdkit import DataStructs
from rdkit.Chem import rdFingerprintGenerator
print(rdkit.__version__)
#%pylab inline

#print current working directory.
print(os.getcwd())


from transformers import AutoTokenizer, AutoModel
import torch

##########################################################

#Import protein data for fasta file.

fasta_path = '/users/PAS2038/rosgaard1/osc_classes/PHYSICS_5680_OSU/final_project/protein.fasta'
sequences = []
invalid_seq = []
with open(fasta_path, 'r') as fasta:
    parser = fastaparser.Reader(fasta)
    for seq in parser:
        try:
            seq_id = seq.id[16:30]
            description = seq.description
            sequence = seq.sequence_as_string()

            # Skip sequences with invalid characters
            if any(aa not in "ACDEFGHIKLMNPQRSTVWY" for aa in sequence):
                invalid_seq.append(seq_id)
                continue

            seq_length = len(sequence)
            analyzed_seq = ProteinAnalysis(sequence)

            sequences.append({
                "uniprot_id": seq_id,
                "description": description,
                "length": seq_length,
                "sequence": sequence,
                "molecular_weight": analyzed_seq.molecular_weight(),
                "isoelectric_point": analyzed_seq.isoelectric_point(),
                "hydrophobicity": analyzed_seq.gravy(),
                "aromaticity": analyzed_seq.aromaticity(),
                "instability_index": analyzed_seq.instability_index(),
                "secondary_structure_fraction": analyzed_seq.secondary_structure_fraction(),
                "amino_acid_composition": analyzed_seq.get_amino_acids_percent()
            })
        except Exception as e:
            print(f"Error processing sequence {seq_id}: {e}")

print(f"Sequences {invalid_seq} were skipped due to invalid amino acids.")

#Create df for protein targets.

sequences_df = pd.DataFrame(sequences)
sequences_df.head()

##########################################################

#Use ProtBert to create embeddings for the amino acid sequences.

tokenizer = AutoTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
model = AutoModel.from_pretrained("Rostlab/prot_bert")
embedded_sequences = []
amino_acid_comp_vectors = []

max_length = 511

for idx, sequence in enumerate(sequences_df['sequence']):
    sequence = " ".join(sequence)
    sequence = f"[CLS] {sequence} [SEP]"

    chunks = [sequence[i:i+max_length] for i in range(0, len(sequence), max_length)]
    
    chunk_embeddings = []
    for chunk in chunks:
        tokens = tokens = tokenizer(chunk, return_tensors="pt", padding=True, truncation=True, max_length=max_length)
        with torch.no_grad():
            outputs = model(**tokens)
            chunk_embedding = outputs.last_hidden_state.mean(dim=1)
            chunk_embeddings.append(chunk_embedding)

    aggregated_embedding = torch.mean(torch.stack(chunk_embeddings), dim=0)
    
    embedded_sequences.append(aggregated_embedding.numpy())


    amino_acid_comp_df = pd.json_normalize(sequences_df['amino_acid_composition'].iloc[idx])
    comp_vector = amino_acid_comp_df.values.tolist()
    
    amino_acid_comp_vectors.append(comp_vector)

#Save it as a .h5 file.
import h5py

with h5py.File("processed_data.h5", "w") as f:
    f.create_dataset("embedded_sequences", data=np.array(embedded_sequences))
    f.create_dataset("amino_acid_comp_vectors", data=np.array(amino_acid_comp_vectors))