generate_stru_seq.py

import os
import sys

import pandas as pd

sys.path.append("code")
from foldseek_util import get_struc_seq


def process_cif_files(directory, batch_size=100):
    data = []
    failed_ids = []
    count = 0

    for filename in os.listdir(directory):
        if filename.endswith(".cif"):
            path = os.path.join(directory, filename)
            uniprot_id = filename.split(".")[0]
            chain_found = False

            for chain in [chr(i) for i in range(ord("A"), ord("Z") + 1)]:
                try:
                    parsed_seqs = get_struc_seq(path, chains=[chain])

                    if chain in parsed_seqs:
                        seq, _, struc_seq = parsed_seqs[chain]
                        data.append(
                            {
                                "Uniprot_ID": uniprot_id,
                                "Alpha_seq": seq,
                                "Alpha_stru_seq": struc_seq,
                            }
                        )

                        count += 1
                        if count % batch_size == 0:
                            save_to_csv(data, append=True)
                            data = []

                        chain_found = True
                        break
                except Exception as e:
                    print(f"Error processing {filename}: {e}")

            if not chain_found:
                failed_ids.append(uniprot_id)

    if data:
        save_to_csv(data, append=True)

    if failed_ids:
        with open("failed_pdb_ids_lit.txt", "w") as f:
            for id in failed_ids:
                f.write(f"{id}\n")


def save_to_csv(data, append=False):
    df = pd.DataFrame(data)
    mode = "a" if append else "w"
    header = not append
    # df.to_csv('Alphafold_pdb_seq.csv', mode=mode, index=False, header=header)
    # df.to_csv('PDB_seq_ext.csv', mode=mode, index=False, header=header)


directory = "alphafold_approved"
process_cif_files(directory, batch_size=100)