test modification of intermediate structural representations

.github/workflows/codespell.yml

@@ -18,4 +18,4 @@ jobs:
       - name: Codespell
         uses: codespell-project/actions-codespell@v2
         with:
-          skip: ./.git,*.pdb,./tests/data
+          skip: ./.git,*.pdb,./tests/data,*.ipynb

.gitignore

@@ -4,3 +4,4 @@ __pycache__
 *.egg-info/
 *.swp
 
+data/*

requirements.txt

@@ -3,3 +3,6 @@ pandas
 numpy
 scipy
 periodictable
+biopandas
+seaborn
+alphafold-colabfold

scripts/preprocess_conformer_pairs.py

@@ -0,0 +1,63 @@
+import pandas as pd
+import pickle
+import glob
+import shutil
+
+from metfish.utils import get_alphafold_atom_positions, get_rmsd, convert_pdb_to_sequence
+from biopandas.pdb import PandasPdb
+
+# set up preprocessing parameters and directories
+calculate_less_similar = True  # calculates which conformer is less similar to the original AF prediction
+use_subset = True  # only selects subset of proteins to run preprocessing
+n_pairs = 6  # if using a subest, the # of pairs to look at, N pairs -> 2N structures predicted
+data_dir = "metfish/data"  # directory with pdb files / csv from Saldano et al., 2022
+output_dir = "metfish/data/output/no_manipulation/240213"  # directory with AF output files
+
+# load apo and holo id names
+pairs_df = pd.read_csv(f"{data_dir}/apo_holo_pairs.csv")
+if use_subset:
+    pairs_df = pairs_df.sample(n_pairs, random_state=1234)  # selects random rows as subsample
+
+# convert pdb files into sequences, AF structures to use as AF inputs
+holo_id = pairs_df["holo_id"].to_list()
+apo_id = pairs_df["apo_id"].to_list()
+for name in [*holo_id, *apo_id]:
+    # clean pdbs (extract only ATOM coordinates )
+    raw_pdb = f"{data_dir}/pdbs_raw/{name}.pdb"
+    clean_pdb = f"{data_dir}/pdbs/{name}_atom_only.pdb"
+    PandasPdb().read_pdb(raw_pdb).to_pdb(clean_pdb, records=["ATOM"])
+
+    # save pairs as fasta sequence files
+    seq = convert_pdb_to_sequence(f"{data_dir}/pdbs/{name}_atom_only.pdb")
+    seq = "\n".join([f">{name}", seq])
+    with open(f"{data_dir}/sequences/apo_and_holo/{name}.fasta", "w") as f:
+        f.write(seq)
+
+    # copy apo ids over to separate folder for AF input (AF only needs to run one of apo/holo pair bc same sequence)
+    if name in apo_id:
+        shutil.copyfile(
+            f"{data_dir}/sequences/apo_and_holo/{name}.fasta", f"{data_dir}/sequences/apo_only/{name}.fasta"
+        )
+
+    # save pairs as alphafold structure representations
+    struct = get_alphafold_atom_positions(f"{data_dir}/pdbs/{name}_atom_only.pdb")
+    with open(f"{data_dir}/af_structures/{name}.pickle", "wb") as f:
+        data = pickle.dump(struct, file=f)
+
+# calculate RMSD values between alphafold output and apo / holo conformers
+# NOTE - you need to run alphafold in advance to calculate these values
+if calculate_less_similar:
+    rmsd_h, rmsd_a = list(), list()
+    for h, a in zip(holo_id, apo_id):
+        af_output = glob.glob(f"{output_dir}/{a}_unrelaxed_rank_001_*_000.pdb")[0]
+        rmsd_h.append(get_rmsd(f"{data_dir}/pdbs/{h}_atom_only.pdb", af_output))  # same af output for same sequence
+        rmsd_a.append(get_rmsd(f"{data_dir}/pdbs/{a}_atom_only.pdb", af_output))  # same af output for same sequence
+
+    pairs_df["rmsd_apo_af"] = rmsd_a
+    pairs_df["rmsd_holo_af"] = rmsd_h
+    pairs_df["less_similar_conformer"] = pairs_df.apply(
+        lambda x: x["holo_id"] if x["rmsd_apo_af"] < x["rmsd_holo_af"] else x["apo_id"], axis=1
+    )
+
+    pairs_df.to_csv(f"{data_dir}/apo_holo_pairs_with_similarity.csv", index=False)
+

src/metfish/representation_manipulation.py

@@ -0,0 +1,115 @@
+import pickle
+import numpy as np
+import pandas as pd
+import warnings
+
+from pathlib import Path
+
+
+def modify_representations(prev: dict = None, method: str = "none", **kwargs):
+    """Modify the pair, position (structural), or MSA (first row only) representations
+    that will be used as inputs for the next recycling iteration of AlphaFold.
+
+    Args:
+        prev (dict): Dict of pair, position, and msa representations.
+        method (str): Method to use to modify representations.
+
+    Returns:
+        dict: modified pair, position and msa representations
+    """
+
+    # apply modification method (test examples, will eventually modify prev outputs based on SAXS data)
+    match method:
+        case "none":
+            repr_modified = prev
+        case "reinitialize":
+            repr_modified = reinitialize(prev)
+        case "add_noise":
+            repr_modified = add_noise(prev)
+        case "replace_structure":
+            repr_modified = replace_structure(prev, **kwargs)
+        case _:
+            warnings.warn("Representation modification method not supported - defaulting to no modification")
+            repr_modified = prev
+
+    return repr_modified
+
+
+def add_noise(prev):
+    """Adds gaussian noise to the pair, structure, and MSA representations"""
+    prev_with_noise = dict()
+    rng = np.random.default_rng()
+
+    for key, value in prev.items():
+        noise = rng.standard_normal(value.shape).astype("float16")  # gaussian noise with μ = 0, σ = 1
+        prev_with_noise[key] = value + noise
+
+    return prev_with_noise
+
+
+def reinitialize(prev):
+    """Reinitializes the pair, structure, and MSA, representations to zero arrays"""
+
+    L = np.shape(prev["prev_pair"])[0]
+
+    prev = {
+        "prev_msa_first_row": np.zeros([L, 256], dtype=np.float16),
+        "prev_pair": np.zeros([L, L, 128], dtype=np.float16),
+        "prev_pos": np.zeros([L, 37, 3], dtype=np.float16),
+    }
+
+    return prev
+
+
+def replace_structure(prev, job_name, input_dir=None, replacement_method="template"):
+    """Replace intermediate structure (atom position) representations from alphafold"""
+    # load in conformer pair information
+    input_dir = input_dir or Path(__file__).resolve().parents[2] / 'data'
+    conformer_pairs_fname = f"{input_dir}/apo_holo_pairs_with_similarity.csv"
+
+    pdb_name = job_name.split("_")[0]
+    conformer_df = pd.read_csv(conformer_pairs_fname)
+    pairs = list(zip(conformer_df["apo_id"], conformer_df["holo_id"]))
+
+    # get relevant pairs
+    index = [ind for ind, (a, h) in enumerate(pairs) if pdb_name in a or pdb_name in h]
+    pair_info = conformer_df.iloc[index, :]
+
+    # get structure name depending on replacement method
+    match replacement_method:
+        case "less_similar":
+            replacement_pdb_name = pair_info["less_similar_conformer"]
+        case "alternate":
+            replacement_pdb_name = (
+                pair_info["holo_id"] if pdb_name in pair_info["apo_id"].to_list()[0] else pair_info["apo_id"]
+            )  # get opposite conformer
+        case "template":
+            replacement_pdb_name = (
+                pair_info["apo_id"] if pdb_name in pair_info["apo_id"].to_list()[0] else pair_info["holo_id"]
+            )  # provide conformer experimental structure
+        case _:
+            replacement_pdb_name = []
+
+    if any(replacement_pdb_name):
+        # load replacement structure
+        print(f"Replacing {pdb_name} intermediate structure with {replacement_pdb_name.to_list()[0]}.")
+        with open(f"{input_dir}/af_structures/{replacement_pdb_name.to_list()[0]}.pickle", "rb") as f:
+            replacement_structure = pickle.load(f)
+
+        # NOTE - additional zero values seem to get added to the first dimension of the position array
+        # (n_res) when running multiple sequences. AF ignores anything longer than n_res when writing
+        # to a pdb file from the protein class, so replacing the first n_res values and adding a warning
+        if np.shape(prev["prev_pos"])[0] != np.shape(replacement_structure)[0]:
+            warnings.warn(
+                f"Alphafold intermediate {np.shape(prev['prev_pos'])} and modified conformer "
+                f"{np.shape(replacement_structure)} structures were not the same shape.",
+            )
+
+        # replace conformer with alternative option
+        n_res = np.shape(replacement_structure)[0]
+        prev["prev_pos"][:n_res, :, :] = replacement_structure.astype("float16")
+
+    else:
+        warnings.warn(f'No replacement option found for "{pdb_name}". Continuing without modification')
+
+    return prev

src/metfish/utils.py

@@ -1,14 +1,21 @@
 import os
 import warnings
+import numpy as np
 
 from Bio.PDB.MMCIFParser import FastMMCIFParser
 from Bio.PDB.PDBParser import PDBParser
-import numpy as np
+from Bio.PDB import PDBIO
+from Bio.SVDSuperimposer import SVDSuperimposer
+from Bio import SeqUtils, Align
+
+from biopandas.pdb import PandasPdb
 from periodictable import elements
 from scipy.spatial.distance import pdist, squareform
+from alphafold.common import protein
 
 
 n_elec_df = {el.symbol: el.number for el in elements}
+amino_acids = [a.upper() for a in SeqUtils.IUPACData.protein_letters_3to1.keys()]
 
 
 def get_Pr(structure, structure_id="", dmax=None, step=0.5):
@@ -69,3 +76,127 @@ def get_Pr(structure, structure_id="", dmax=None, step=0.5):
     p = np.concatenate(([0], hist / hist.sum()))
 
     return r, p
+
+
+def get_alphafold_atom_positions(fname: str):
+    """ Use alphafold protein module to convert pdb file to alphafold's atomic position representation
+
+    Args:
+        fname (str): path to the pdb file
+
+    Returns:
+        structure: An np.ndarray with cartesian coordinates of atoms in angstroms [num_res, num_atom_type, 3]. 
+        The atom types correspond to residue_constants.atom_types, i.e. the first three are N, CA, CB.
+    """
+    # read in file text and use af protein module to convert to position representation
+    with open(fname, 'r') as file:
+        pdb_str = file.read()
+    prot = protein.from_pdb_string(pdb_str)  # protein module uses pdb file contents as input
+
+    return prot.atom_positions
+
+
+def convert_pdb_to_sequence(fname: str):
+    """ Get single letter amino acid sequence from a pdb structure file """
+    structure = PDBParser(QUIET=True).get_structure('', fname)
+    residues = [res.resname for res in structure.get_residues() if res.resname in amino_acids]
+    sequence = get_single_letter_sequences(residues)
+
+    return sequence
+
+
+def get_single_letter_sequences(residues):
+    ret = list()
+    for res in residues:
+        res = res[0] + res[1:].lower()
+        ret.append(SeqUtils.IUPACData.protein_letters_3to1[res])
+    return "".join(ret)
+
+
+def align_sequences(ref_df, query_df):
+    """ Align protein sequences """
+    ref_seq = get_single_letter_sequences(ref_df['residue_name'])
+    query_seq = get_single_letter_sequences(query_df['residue_name'])
+
+    # if not the same sequence, align
+    if ref_seq != query_seq:
+        aligner = Align.PairwiseAligner()
+        alignments = aligner.align(ref_seq, query_seq)
+
+        ref_idx, query_idx = alignments[0].indices[:, ~(alignments[0].indices == -1).any(axis=0)]
+
+        ref_df = ref_df.iloc[ref_idx]
+        query_df = query_df.iloc[query_idx]
+
+    return ref_df, query_df
+
+
+def superimpose_structures(fname_fixed, fname_moving, atom_types=["CA", "N", "C", "O"]):
+    """ Superimpose two protein structures. 
+
+    Args:
+        fname_fixed (str): path to PDB file
+        fname_moving (str): path to PDB file
+        atom_types (list): atom types to align structures with, traditionally aligned with either
+        1) only alpha-carbon atoms (CA), or 2) the "protein backbone" atoms (CA, N, C, O), or all atoms
+
+    Returns:
+        superimposer: returns instance of BioPython SVDSuperImposer class
+    """
+
+    # read in structures
+    fixed_atom_df = PandasPdb().read_pdb(fname_fixed).df['ATOM']
+    moving_atom_df = PandasPdb().read_pdb(fname_moving).df['ATOM']
+
+    # filter for atom types and amino acide residues only
+    fixed_atom_df = fixed_atom_df.query(f"residue_name in {amino_acids} & atom_name in {atom_types}")
+    moving_atom_df = moving_atom_df.query(f"residue_name in {amino_acids} & atom_name in {atom_types}")
+
+    # align sequences (if already aligned, will return same df)
+    fixed_atom_df, moving_atom_df = align_sequences(fixed_atom_df, moving_atom_df)
+
+    # get coordinates of the atoms
+    fixed_coords = fixed_atom_df[['x_coord', 'y_coord', 'z_coord']].to_numpy()
+    moving_coords = moving_atom_df[['x_coord', 'y_coord', 'z_coord']].to_numpy()
+
+    # superimpose structures
+    si = SVDSuperimposer()
+    si.set(fixed_coords, moving_coords)
+    si.run() # Run the SVD alignment
+
+    return si
+
+def get_rmsd(fname_a, fname_b, atom_types=["CA", "N", "C", "O"]):
+    """ Calculate the RMSD between superimposed coordinates of two protein structures.
+    """
+
+    si = superimpose_structures(fname_a, fname_b, atom_types=atom_types)
+    return si.get_rms()
+
+def align_structures(fname_fixed, fname_moving):
+    """Align two protein structures from pdb files and save aligned structures
+
+    Args:
+        fname_fixed (str): path to PDB file
+        fname_moving (str): path to PDB file
+
+    Returns:
+        fname_aligned: path to aligned PDB file (rotate/translated version of fname_moving)
+    """
+
+    # load structure to transform
+    structure = PDBParser(QUIET=True).get_structure('', fname_moving)
+
+    # superimpose on experimental structure file 
+    si = superimpose_structures(fname_fixed, fname_moving)
+    rot, trans = si.get_rotran()
+    for atom in structure.get_atoms():
+        atom.transform(rot.astype("f"), trans.astype("f"))
+
+    # save modified outputs as PDB files
+    fname_aligned = f"{fname_moving.strip('.pdb')}_aligned.pdb"
+    io = PDBIO()
+    io.set_structure(structure) 
+    io.save(fname_aligned)
+
+    return fname_aligned