|
| 1 | +#!/usr/bin/env python |
| 2 | +import pickle |
| 3 | +import os |
| 4 | +import argparse |
| 5 | +import json |
| 6 | +#import torch moved to a conditional import since too bulky import if not used |
| 7 | +import numpy as np |
| 8 | +import csv |
| 9 | +from utils import plddt_from_struct_b_factor |
| 10 | + |
| 11 | +# TODO: add extraction of other values, iPTM, etc |
| 12 | +# TODO: look into have a --prog argument that could set filenames etc, logically seperate it? |
| 13 | +# {id}_{prog}_{metric}.tsv might be easier for MultiQC to parse a complex workdir, than without the .prog |
| 14 | + |
| 15 | +# Mapping of characters to integers for MSA parsing. |
| 16 | +# 20 is for unknown characters, and 21 is for gaps. |
| 17 | +AA_to_int = { |
| 18 | + "A": 0, "C": 1, "D": 2, "E": 3, "F": 4, "G": 5, "H": 6, "I": 7, "K": 8, "L": 9, |
| 19 | + "M": 10, "N": 11, "P": 12, "Q": 13, "R": 14, "S": 15, "T": 16, "V": 17, "W": 18, "Y": 19, |
| 20 | + ".": 20, "-": 21 |
| 21 | +} |
| 22 | + |
| 23 | +def a3m_to_int(a3m_file): |
| 24 | + """ |
| 25 | + Convert an A3M MSA representation into an integer representation (0-21). |
| 26 | + """ |
| 27 | + with open(a3m_file, "r") as f: |
| 28 | + msa = f.read() |
| 29 | + |
| 30 | + # Convert each sequence in the MSA |
| 31 | + int_sequences = [] |
| 32 | + for idx, line in enumerate(msa.splitlines()): |
| 33 | + if idx == 0 and not line.startswith(">"): # If there's an additional header (non-FASTA) skip it. E.g ColabFold |
| 34 | + continue |
| 35 | + |
| 36 | + if not line.startswith(">"): # Ignore header lines |
| 37 | + filtered_line = ''.join(char for char in line if not char.islower()) # Remove inserts (lower-case chars) in a3m |
| 38 | + int_sequence = [AA_to_int.get(char.upper(), 20) for char in filtered_line] |
| 39 | + int_sequences.append(int_sequence) |
| 40 | + |
| 41 | + int_sequences_array = np.array(int_sequences, dtype=object) |
| 42 | + return int_sequences_array |
| 43 | + |
| 44 | +def format_msa_rows(msa_data): |
| 45 | + return [[str(x) for x in val] for val in msa_data] |
| 46 | + |
| 47 | +def format_pae_rows(pae_data): |
| 48 | + return [[f"{num:.4f}" for num in row] for row in pae_data] |
| 49 | + |
| 50 | +def write_tsv(file_path, rows): |
| 51 | + with open(file_path, 'w') as out_f: |
| 52 | + writer = csv.writer(out_f, delimiter='\t') |
| 53 | + writer.writerows(rows) |
| 54 | + |
| 55 | +def extract_structs_plddt_to_tsv(id, structures): |
| 56 | + """ |
| 57 | + Write out a tsv file contain pLDDTs for reading by MultiQC in nf-core/proteinfold |
| 58 | + Uses utils function with BioPython PDB package to extract residue pLDDT values from the b-factor column. |
| 59 | + """ |
| 60 | + plddt_cols = [plddt_from_struct_b_factor(structure) for structure in structures] |
| 61 | + res_counts = [len(plddt_col) for plddt_col in plddt_cols] |
| 62 | + |
| 63 | + if len(set(res_counts)) != 1: |
| 64 | + raise ValueError("Not all structures have the same number of residues!") |
| 65 | + |
| 66 | + rank_names = [f"rank_{i}" for i in range(len(structures))] |
| 67 | + # Create header as the first row |
| 68 | + plddt_rows = [["Positions"] + rank_names] |
| 69 | + res_id_col = list(range(len(plddt_cols[0]))) |
| 70 | + plddt_rows.extend(zip(res_id_col, *plddt_cols)) # Combine lists column-wise to make rows |
| 71 | + write_tsv(f"{id}_plddt.tsv", plddt_rows) |
| 72 | + |
| 73 | +def read_pkl(id, pkl_files): |
| 74 | + """ |
| 75 | + Adapted from the Galaxy AlphaFold tool (https://github.com/usegalaxy-au/tools-au/blob/de94df520c8dc7b8652aedb92e90f6ebb312f95f/tools/alphafold/scripts/outputs.py), originally authored by @neoformit and @graceahall and funded by Australian Biocommons and QCIF Australia. |
| 76 | + """ |
| 77 | + for pkl_file in pkl_files: |
| 78 | + print(f"Processing {pkl_file}") |
| 79 | + data = pickle.load(open(pkl_file, "rb")) |
| 80 | + |
| 81 | + # Process MSA data |
| 82 | + if pkl_file.endswith("final_features.pkl"): # HelixFold3 - This one must be first |
| 83 | + write_tsv(f"{id}_msa.tsv", format_msa_rows(data["feat"]["msa"])) |
| 84 | + elif pkl_file.endswith("features.pkl"): # AlphaFold2.3 |
| 85 | + # TODO: AlphaFold2.3 fills end rows with 0s in AlpahFold muliter for an alanine nf-core/proteinfold Issue #300 |
| 86 | + write_tsv(f"{id}_msa.tsv", format_msa_rows(data["msa"])) |
| 87 | + # AlphaFold2.3 non-summary, for each pkl. TODO: Need to either read in ranking_debug.json to get the ranking order, or do it later in the workflow. |
| 88 | + else: |
| 89 | + model_id = os.path.basename(pkl_file).replace("result_model_", "").replace(".pkl", "") |
| 90 | + write_tsv(f"{id}_{model_id}_lddt.tsv", format_msa_rows(data["plddt"])) |
| 91 | + |
| 92 | + if 'predicted_aligned_error' not in data.keys(): |
| 93 | + print(f"No PAE output in {pkl_file}, it was likely a monomer calculation") |
| 94 | + write_tsv(f"{id}_{model_id}_pae.tsv", None) |
| 95 | + else: |
| 96 | + write_tsv(f"{id}_{model_id}_pae.tsv", format_pae_rows(data["predicted_aligned_error"])) |
| 97 | + |
| 98 | +def read_a3m(id, a3m_files): |
| 99 | + # ColabFold, RosettaFold-All-Atom, Boltz-1 |
| 100 | + for a3m_file in a3m_files: |
| 101 | + int_seqs = a3m_to_int(a3m_file) |
| 102 | + write_tsv(f"{id}_msa.tsv", format_msa_rows(int_seqs)) |
| 103 | + |
| 104 | +def read_npz(id, npz_files): |
| 105 | + for idx, npz_file in enumerate(npz_files): |
| 106 | + data = np.load(npz_file) |
| 107 | + #Boltz PAE files if --write_full_pae is used |
| 108 | + if npz_file.split('/')[-1].startswith('pae') and npz_file.endswith('.npz'): |
| 109 | + write_tsv(f"{id}_{idx}_pae.tsv", format_pae_rows(data["pae"])) |
| 110 | + |
| 111 | +def read_json(id, json_files): |
| 112 | + for idx, json_file in enumerate(json_files): |
| 113 | + with open(json_file, 'r') as f: |
| 114 | + data = json.load(f) |
| 115 | + if json_file.endswith("_data.json"): #AF3 output with MSA info |
| 116 | + # Can't just used format_msa_rows since there's FASTA headers in the json content |
| 117 | + unpaired_MSAs = data['sequences'][0]['protein']['unpairedMsa'] |
| 118 | + msa_lines = [''.join(c for c in line if not c.islower()) for line in unpaired_MSAs.split("\n") if line.strip() and not line.startswith(">")] |
| 119 | + msa_rows = [[str(AA_to_int.get(residue, 20)) for residue in line] for line in msa_lines] |
| 120 | + write_tsv(f"{id}_msa.tsv", msa_rows) |
| 121 | + #AF3 output with PAE info, or HF3 PAE data. TODO: Need to make sure the workflow points to [protein]/[protein]_rank1/all_results.json |
| 122 | + elif "pae" in data: |
| 123 | + write_tsv(f"{id}_{idx}_pae.tsv", format_pae_rows(data["pae"])) |
| 124 | + |
| 125 | +def read_pt(id, pt_files): |
| 126 | + import torch # moved to a conditional import since too bulky import if not used |
| 127 | + for pt_file in pt_files: |
| 128 | + with open(pt_file, 'rb') as f: # TODO: point to [protein]_aux.pt |
| 129 | + data = torch.load(f, map_location="cpu") |
| 130 | + if 'pae' in data: |
| 131 | + # The pt file contains a tensor that needs to be cast as an array |
| 132 | + # Squeeze leading dimension (batch?) |
| 133 | + write_tsv(f"{id}_pae.tsv", format_pae_rows(np.squeeze(data["pae"].numpy()))) |
| 134 | + |
| 135 | +def main(): |
| 136 | + parser = argparse.ArgumentParser() |
| 137 | + parser.add_argument("--pkls", dest="pkls", required=False, nargs="+") # For reading both HelixFold3 and AlphaFold2 MSA formats |
| 138 | + parser.add_argument("--npzs", dest="npzs", required=False, nargs="+") # For reading the Boltz-1 PAE formats. TODO: Boltz-1 MSA not implemented (go straight to .a3m file), implement |
| 139 | + parser.add_argument("--a3ms", dest="a3ms", required=False, nargs="+") # For reading the RosettaFold-All-Atom, ColabFold, and Boltz-1 MSA formats |
| 140 | + parser.add_argument("--jsons", dest="jsons", required=False, nargs="+") # For reading the AF3 MSA & PAE, HF3 PAE |
| 141 | + parser.add_argument("--pts", dest="pts", required=False, nargs="+") # For read RFAA pytorch model to get PAE data |
| 142 | + parser.add_argument("--structs", dest="structs", required=False, nargs="+") |
| 143 | + parser.add_argument("--name", default="untitled", dest="name") # might need a --name $meta.id |
| 144 | + parser.add_argument("--output_dir", default=".", dest="output_dir") |
| 145 | + args = parser.parse_args() |
| 146 | + |
| 147 | + if args.pkls: |
| 148 | + read_pkl(args.name, args.pkls) |
| 149 | + if args.a3ms: |
| 150 | + read_a3m(args.name, args.a3ms) |
| 151 | + if args.npzs: |
| 152 | + read_npz(args.name, args.npzs) |
| 153 | + if args.jsons: |
| 154 | + read_json(args.name, args.jsons) |
| 155 | + if args.pts: |
| 156 | + read_pt(args.name, args.pts) |
| 157 | + if args.structs: |
| 158 | + extract_structs_plddt_to_tsv(args.name, args.structs) |
| 159 | + |
| 160 | +if __name__ == "__main__": |
| 161 | + main() |
0 commit comments