Diffusion-based all-atom protein generative model.
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Dataloader from PDB files.
"""
import copy
import pickle
import json
import numpy as np
import torch
from torch.utils import data

from core import utils
from core import protein
from core import residue_constants


FEATURES_1D = (
    "coords_in",
    "torsions_in",
    "b_factors",
    "atom_positions",
    "aatype",
    "atom_mask",
    "residue_index",
    "chain_index",
)
FEATURES_FLOAT = (
    "coords_in",
    "torsions_in",
    "b_factors",
    "atom_positions",
    "atom_mask",
    "seq_mask",
)
FEATURES_LONG = ("aatype", "residue_index", "chain_index", "orig_size")


def make_fixed_size_1d(data, fixed_size=128):
    data_len = data.shape[0]
    if data_len >= fixed_size:
        extra_len = data_len - fixed_size
        start_idx = np.random.choice(np.arange(extra_len + 1))
        new_data = data[start_idx : (start_idx + fixed_size)]
        mask = torch.ones(fixed_size)
    if data_len < fixed_size:
        pad_size = fixed_size - data_len
        extra_shape = data.shape[1:]
        new_data = torch.cat([data, torch.zeros(pad_size, *extra_shape)], 0)
        mask = torch.cat([torch.ones(data_len), torch.zeros(pad_size)], 0)
    return new_data, mask


def apply_random_se3(coords_in, atom_mask=None, translation_scale=1.0):
    # unbatched. center on the mean of CA coords
    coords_mean = coords_in[:, 1:2].mean(-3, keepdim=True)
    coords_in -= coords_mean
    random_rot, _ = torch.linalg.qr(torch.randn(3, 3))
    coords_in = coords_in @ random_rot
    random_trans = torch.randn_like(coords_mean) * translation_scale
    coords_in += random_trans
    if atom_mask is not None:
        coords_in = coords_in * atom_mask[..., None]
    return coords_in


def get_masked_coords_array(coords, atom_mask):
    ma_mask = repeat(1 - atom_mask[..., None].cpu().numpy(), "... 1 -> ... 3")
    return np.ma.array(coords.cpu().numpy(), mask=ma_mask)


def make_crop_cond_mask_and_recenter_coords(
    atom_mask,
    atom_coords,
    contiguous_prob=0.05,
    discontiguous_prob=0.9,
    sidechain_only_prob=0.8,
    max_span_len=10,
    max_discontiguous_res=8,
    dist_threshold=8.0,
    recenter_coords=True,
):
    b, n, a = atom_mask.shape
    device = atom_mask.device
    seq_mask = atom_mask[..., 1]
    n_res = seq_mask.sum(-1)
    masks = []

    for i, nr in enumerate(n_res):
        nr = nr.int().item()
        mask = torch.zeros((n, a), device=device)
        conditioning_type = torch.distributions.Categorical(
            torch.tensor(
                [
                    contiguous_prob,
                    discontiguous_prob,
                    1.0 - contiguous_prob - discontiguous_prob,
                ]
            )
        ).sample()
        conditioning_type = ["contiguous", "discontiguous", "none"][conditioning_type]

        if conditioning_type == "contiguous":
            span_len = torch.randint(
                1, min(max_span_len, nr), (1,), device=device
            ).item()
            span_start = torch.randint(0, nr - span_len, (1,), device=device)
            mask[span_start : span_start + span_len, :] = 1
        elif conditioning_type == "discontiguous":
            # Extract CB atoms coordinates for the i-th example
            cb_atoms = atom_coords[i, :, 3]
            # Pairwise distances between CB atoms
            cb_distances = torch.cdist(cb_atoms, cb_atoms)
            close_mask = (
                cb_distances <= dist_threshold
            )  # Mask for selecting close CB atoms

            random_residue = torch.randint(0, nr, (1,), device=device).squeeze()
            cb_dist_i = cb_distances[random_residue] + 1e3 * (1 - seq_mask[i])
            close_mask = cb_dist_i <= dist_threshold
            n_neighbors = close_mask.sum().int()

            # pick how many neighbors (up to 10)
            n_sele = torch.randint(
                2,
                n_neighbors.clamp(min=3, max=max_discontiguous_res + 1),
                (1,),
                device=device,
            )

            # Select the indices of CB atoms that are close together
            idxs = torch.arange(n, device=device)[close_mask.bool()]
            idxs = idxs[torch.randperm(len(idxs))[:n_sele]]

            if len(idxs) > 0:
                mask[idxs] = 1

            if np.random.uniform() < sidechain_only_prob:
                mask[:, :5] = 0

        masks.append(mask)

    crop_cond_mask = torch.stack(masks)
    crop_cond_mask = crop_cond_mask * atom_mask
    if recenter_coords:
        motif_masked_array = get_masked_coords_array(atom_coords, crop_cond_mask)
        cond_coords_center = motif_masked_array.mean((1, 2))
        motif_mask = torch.Tensor(1 - cond_coords_center.mask).to(crop_cond_mask)
        means = torch.Tensor(cond_coords_center.data).to(atom_coords) * motif_mask
        coords_out = atom_coords - rearrange(means, "b c -> b 1 1 c")
    else:
        coords_out = atom_coords
    return coords_out, crop_cond_mask


class Dataset(data.Dataset):
    """Loads and processes PDBs into tensors."""

    def __init__(
        self,
        pdb_path,
        fixed_size,
        mode="train",
        overfit=-1,
        short_epoch=False,
        se3_data_augment=True,
    ):
        self.pdb_path = pdb_path
        self.fixed_size = fixed_size
        self.mode = mode
        self.overfit = overfit
        self.short_epoch = short_epoch
        self.se3_data_augment = se3_data_augment

        with open(f"{self.pdb_path}/{mode}_pdb_keys.list") as f:
            self.pdb_keys = np.array(f.read().split("\n")[:-1])

        if overfit > 0:
            n_data = len(self.pdb_keys)
            self.pdb_keys = np.random.choice(
                self.pdb_keys, min(n_data, overfit), replace=False
            ).repeat(n_data // overfit)

    def __len__(self):
        if self.short_epoch:
            return min(len(self.pdb_keys), 256)
        else:
            return len(self.pdb_keys)

    def __getitem__(self, idx):
        pdb_key = self.pdb_keys[idx]
        data = self.get_item(pdb_key)
        # For now, replace dataloading errors with a random pdb. 10 tries
        for _ in range(10):
            if data is not None:
                return data
            pdb_key = self.pdb_keys[np.random.randint(len(self.pdb_keys))]
            data = self.get_item(pdb_key)
        raise Exception("Failed to load data example after 10 tries.")

    def get_item(self, pdb_key):
        example = {}

        if self.pdb_path.endswith("cath_s40_dataset"):  # CATH pdbs
            data_file = f"{self.pdb_path}/dompdb/{pdb_key}"
        elif self.pdb_path.endswith("ingraham_cath_dataset"):  # ingraham splits
            data_file = f"{self.pdb_path}/pdb_store/{pdb_key}"
        else:
            raise Exception("Invalid pdb path.")

        try:
            example = utils.load_feats_from_pdb(data_file)
            coords_in = example["atom_positions"]
        except FileNotFoundError:
            raise Exception(f"File {pdb_key} not found. Check if dataset is corrupted?")
        except RuntimeError:
            return None

        # Apply data augmentation
        if self.se3_data_augment:
            coords_in = apply_random_se3(coords_in, atom_mask=example["atom_mask"])

        orig_size = coords_in.shape[0]
        example["coords_in"] = coords_in
        example["orig_size"] = torch.ones(1) * orig_size

        fixed_size_example = {}
        seq_mask = None
        for k, v in example.items():
            if k in FEATURES_1D:
                fixed_size_example[k], seq_mask = make_fixed_size_1d(
                    v, fixed_size=self.fixed_size
                )
            else:
                fixed_size_example[k] = v
        if seq_mask is not None:
            fixed_size_example["seq_mask"] = seq_mask

        example_out = {}
        for k, v in fixed_size_example.items():
            if k in FEATURES_FLOAT:
                example_out[k] = v.float()
            elif k in FEATURES_LONG:
                example_out[k] = v.long()

        return example_out

    def collate(self, example_list):
        out = {}
        for ex in example_list:
            for k, v in ex.items():
                out.setdefault(k, []).append(v)
        return {k: torch.stack(v) for k, v in out.items()}

    def sample(self, n=1, return_data=True, return_keys=False):
        keys = self.pdb_keys[torch.randperm(self.__len__())[:n].long()]

        if return_keys and not return_data:
            return keys

        if n == 1:
            data = self.collate([self.get_item(keys)])
        else:
            data = self.collate([self.get_item(key) for key in keys])

        if return_data and return_keys:
            return data, keys
        if return_data and not return_keys:
            return data

core/data.py
ProteinDesignLab/protpardelle
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Protein data type.
Adapted from original code by alexechu.
"""
import dataclasses
import io
from typing import Any, Mapping, Optional
from core import residue_constants
from Bio.PDB import PDBParser
import numpy as np

FeatureDict = Mapping[str, np.ndarray]
ModelOutput = Mapping[str, Any]  # Is a nested dict.

# Complete sequence of chain IDs supported by the PDB format.
PDB_CHAIN_IDS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
PDB_MAX_CHAINS = len(PDB_CHAIN_IDS)  # := 62.


@dataclasses.dataclass(frozen=True)
class Protein:
    """Protein structure representation."""

    # Cartesian coordinates of atoms in angstroms. The atom types correspond to
    # residue_constants.atom_types, i.e. the first three are N, CA, CB.
    atom_positions: np.ndarray  # [num_res, num_atom_type, 3]

    # Amino-acid type for each residue represented as an integer between 0 and
    # 20, where 20 is 'X'.
    aatype: np.ndarray  # [num_res]

    # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
    # is present and 0.0 if not. This should be used for loss masking.
    atom_mask: np.ndarray  # [num_res, num_atom_type]

    # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
    residue_index: np.ndarray  # [num_res]

    # 0-indexed number corresponding to the chain in the protein that this residue
    # belongs to.
    chain_index: np.ndarray  # [num_res]

    # B-factors, or temperature factors, of each residue (in sq. angstroms units),
    # representing the displacement of the residue from its ground truth mean
    # value.
    b_factors: np.ndarray  # [num_res, num_atom_type]

    def __post_init__(self):
        if len(np.unique(self.chain_index)) > PDB_MAX_CHAINS:
            raise ValueError(
                f"Cannot build an instance with more than {PDB_MAX_CHAINS} chains "
                "because these cannot be written to PDB format."
            )


def from_pdb_string(
    pdb_str: str, chain_id: Optional[str] = None, protein_only: bool = False
) -> Protein:
    """Takes a PDB string and constructs a Protein object.

    WARNING: All non-standard residue types will be converted into UNK. All
      non-standard atoms will be ignored.

    Args:
      pdb_str: The contents of the pdb file
      chain_id: If chain_id is specified (e.g. A), then only that chain
        is parsed. Otherwise all chains are parsed.

    Returns:
      A new `Protein` parsed from the pdb contents.
    """
    pdb_fh = io.StringIO(pdb_str)
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("none", pdb_fh)
    models = list(structure.get_models())
    if len(models) != 1:
        raise ValueError(
            f"Only single model PDBs are supported. Found {len(models)} models."
        )
    model = models[0]

    atom_positions = []
    aatype = []
    atom_mask = []
    residue_index = []
    chain_ids = []
    b_factors = []

    for chain in model:
        if chain_id is not None and chain.id != chain_id:
            continue
        for res in chain:
            if protein_only and res.id[0] != " ":
                continue
            if res.id[2] != " ":
                pass
                # raise ValueError(
                #     f"PDB contains an insertion code at chain {chain.id} and residue "
                #     f"index {res.id[1]}. These are not supported."
                # )
            res_shortname = residue_constants.restype_3to1.get(res.resname, "X")
            restype_idx = residue_constants.restype_order.get(
                res_shortname, residue_constants.restype_num
            )
            pos = np.zeros((residue_constants.atom_type_num, 3))
            mask = np.zeros((residue_constants.atom_type_num,))
            res_b_factors = np.zeros((residue_constants.atom_type_num,))
            for atom in res:
                if atom.name not in residue_constants.atom_types:
                    continue
                pos[residue_constants.atom_order[atom.name]] = atom.coord
                mask[residue_constants.atom_order[atom.name]] = 1.0
                res_b_factors[residue_constants.atom_order[atom.name]] = atom.bfactor
            if np.sum(mask) < 0.5:
                # If no known atom positions are reported for the residue then skip it.
                continue
            aatype.append(restype_idx)
            atom_positions.append(pos)
            atom_mask.append(mask)
            residue_index.append(res.id[1])
            chain_ids.append(chain.id)
            b_factors.append(res_b_factors)

    # Chain IDs are usually characters so map these to ints.
    unique_chain_ids = np.unique(chain_ids)
    chain_id_mapping = {cid: n for n, cid in enumerate(unique_chain_ids)}
    chain_index = np.array([chain_id_mapping[cid] for cid in chain_ids])

    return Protein(
        atom_positions=np.array(atom_positions),
        atom_mask=np.array(atom_mask),
        aatype=np.array(aatype),
        residue_index=np.array(residue_index),
        chain_index=chain_index,
        b_factors=np.array(b_factors),
    )


def _chain_end(atom_index, end_resname, chain_name, residue_index) -> str:
    chain_end = "TER"
    return (
        f"{chain_end:<6}{atom_index:>5}      {end_resname:>3} "
        f"{chain_name:>1}{residue_index:>4}"
    )


def are_atoms_bonded(res3name, atom1_name, atom2_name):
    lookup_table = residue_constants.standard_residue_bonds
    for bond in lookup_table[res3name]:
        if bond.atom1_name == atom1_name and bond.atom2_name == atom2_name:
            return True
        elif bond.atom1_name == atom2_name and bond.atom2_name == atom1_name:
            return True
    return False


def to_pdb(prot: Protein, conect=False) -> str:
    """Converts a `Protein` instance to a PDB string.

    Args:
      prot: The protein to convert to PDB.

    Returns:
      PDB string.
    """
    restypes = residue_constants.restypes + ["X"]
    res_1to3 = lambda r: residue_constants.restype_1to3.get(restypes[r], "UNK")
    atom_types = residue_constants.atom_types

    pdb_lines = []

    atom_mask = prot.atom_mask
    aatype = prot.aatype
    atom_positions = prot.atom_positions
    residue_index = prot.residue_index.astype(np.int32)
    chain_index = prot.chain_index.astype(np.int32)
    b_factors = prot.b_factors

    if np.any(aatype > residue_constants.restype_num):
        raise ValueError("Invalid aatypes.")

    # Construct a mapping from chain integer indices to chain ID strings.
    chain_ids = {}
    for i in np.unique(chain_index):  # np.unique gives sorted output.
        if i >= PDB_MAX_CHAINS:
            raise ValueError(
                f"The PDB format supports at most {PDB_MAX_CHAINS} chains."
            )
        chain_ids[i] = PDB_CHAIN_IDS[i]

    pdb_lines.append("MODEL     1")
    atom_index = 1
    last_chain_index = chain_index[0]
    conect_lines = []
    # Add all atom sites.
    for i in range(aatype.shape[0]):
        # Close the previous chain if in a multichain PDB.
        if last_chain_index != chain_index[i]:
            pdb_lines.append(
                _chain_end(
                    atom_index,
                    res_1to3(aatype[i - 1]),
                    chain_ids[chain_index[i - 1]],
                    residue_index[i - 1],
                )
            )
            last_chain_index = chain_index[i]
            atom_index += 1  # Atom index increases at the TER symbol.

        res_name_3 = res_1to3(aatype[i])
        atoms_appended_for_res = []
        for atom_name, pos, mask, b_factor in zip(
            atom_types, atom_positions[i], atom_mask[i], b_factors[i]
        ):
            if mask < 0.5:
                continue

            record_type = "ATOM"
            name = atom_name if len(atom_name) == 4 else f" {atom_name}"
            alt_loc = ""
            insertion_code = ""
            occupancy = 1.00
            element = atom_name[0]  # Protein supports only C, N, O, S, this works.
            charge = ""
            # PDB is a columnar format, every space matters here!
            atom_line = (
                f"{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"
                f"{res_name_3:>3} {chain_ids[chain_index[i]]:>1}"
                f"{residue_index[i]:>4}{insertion_code:>1}   "
                f"{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"
                f"{occupancy:>6.2f}{b_factor:>6.2f}          "
                f"{element:>2}{charge:>2}"
            )
            pdb_lines.append(atom_line)

            for prev_atom_idx, prev_atom in atoms_appended_for_res:
                if are_atoms_bonded(res_name_3, atom_name, prev_atom):
                    conect_line = f"CONECT{prev_atom_idx:5d}{atom_index:5d}\n"
                    conect_lines.append(conect_line)
            atoms_appended_for_res.append((atom_index, atom_name))
            if atom_name == "N":
                n_atom_idx = atom_index
            if atom_name == "C":
                c_atom_idx = atom_index

            atom_index += 1

        if i > 0:
            conect_line = f"CONECT{prev_c_atom_idx:5d}{n_atom_idx:5d}\n"
            conect_lines.append(conect_line)
        prev_c_atom_idx = c_atom_idx

    # Close the final chain.
    pdb_lines.append(
        _chain_end(
            atom_index,
            res_1to3(aatype[-1]),
            chain_ids[chain_index[-1]],
            residue_index[-1],
        )
    )
    pdb_lines.append("ENDMDL")
    pdb_lines.append("END")

    # Pad all lines to 80 characters.
    pdb_lines = [line.ljust(80) for line in pdb_lines]
    pdb_str = "\n".join(pdb_lines) + "\n"  # Add terminating newline.
    if conect:
        conect_str = "".join(conect_lines) + "\n"
        return pdb_str, conect_str
    return pdb_str


def ideal_atom_mask(prot: Protein) -> np.ndarray:
    """Computes an ideal atom mask.

    `Protein.atom_mask` typically is defined according to the atoms that are
    reported in the PDB. This function computes a mask according to heavy atoms
    that should be present in the given sequence of amino acids.

    Args:
      prot: `Protein` whose fields are `numpy.ndarray` objects.

    Returns:
      An ideal atom mask.
    """
    return residue_constants.STANDARD_ATOM_MASK[prot.aatype]


def from_prediction(
    features: FeatureDict,
    result: ModelOutput,
    b_factors: Optional[np.ndarray] = None,
    remove_leading_feature_dimension: bool = True,
) -> Protein:
    """Assembles a protein from a prediction.

    Args:
      features: Dictionary holding model inputs.
      result: Dictionary holding model outputs.
      b_factors: (Optional) B-factors to use for the protein.
      remove_leading_feature_dimension: Whether to remove the leading dimension
        of the `features` values.

    Returns:
      A protein instance.
    """
    fold_output = result["structure_module"]

    def _maybe_remove_leading_dim(arr: np.ndarray) -> np.ndarray:
        return arr[0] if remove_leading_feature_dimension else arr

    if "asym_id" in features:
        chain_index = _maybe_remove_leading_dim(features["asym_id"])
    else:
        chain_index = np.zeros_like(_maybe_remove_leading_dim(features["aatype"]))

    if b_factors is None:
        b_factors = np.zeros_like(fold_output["final_atom_mask"])

    return Protein(
        aatype=_maybe_remove_leading_dim(features["aatype"]),
        atom_positions=fold_output["final_atom_positions"],
        atom_mask=fold_output["final_atom_mask"],
        residue_index=_maybe_remove_leading_dim(features["residue_index"]) + 1,
        chain_index=chain_index,
        b_factors=b_factors,
    )

core/protein.py
ProteinDesignLab/protpardelle
# MIT License

# Copyright (c) 2022 Justas Dauparas

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

'''
Adapted from original code by alexechu.
'''
import json, time, os, sys, glob
import shutil
import warnings
import copy
import random
import os.path
import subprocess
import itertools

from einops.layers.torch import Rearrange
import numpy as np
import torch
from torch import optim
from torch.utils.data import DataLoader
from torch.utils.data.dataset import random_split, Subset
import torch.nn as nn
import torch.nn.functional as F


def get_mpnn_model(model_name='v_48_020', path_to_model_weights='', ca_only=False, backbone_noise=0.0, verbose=False, device=None):
    hidden_dim = 128
    num_layers = 3 
    if device is None:
        device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")

    if path_to_model_weights:
        model_folder_path = path_to_model_weights
        if model_folder_path[-1] != '/':
            model_folder_path = model_folder_path + '/'
    else: 
        file_path = os.path.realpath(__file__)
        k = file_path.rfind("/")
        k = file_path[:k].rfind("/")
        if ca_only:
            model_folder_path = file_path[:k] + '/ProteinMPNN/ca_model_weights/'
        else:
            model_folder_path = file_path[:k] + '/ProteinMPNN/vanilla_model_weights/'

    checkpoint_path = model_folder_path + f'{model_name}.pt'
    checkpoint = torch.load(checkpoint_path, map_location=device) 
    noise_level_print = checkpoint['noise_level']
    model = ProteinMPNN(ca_only=ca_only, num_letters=21, node_features=hidden_dim, edge_features=hidden_dim, hidden_dim=hidden_dim,
        num_encoder_layers=num_layers, num_decoder_layers=num_layers, augment_eps=backbone_noise, k_neighbors=checkpoint['num_edges'])
    model.to(device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()

    if verbose:
        print(40*'-')
        print('Model loaded...')
        print('Number of edges:', checkpoint['num_edges'])
        print(f'Training noise level: {noise_level_print}A')

    return model


def run_proteinmpnn(model=None, pdb_path='', pdb_path_chains='', path_to_model_weights='', model_name='v_48_020', seed=0, ca_only=False, out_folder='', num_seq_per_target=1, batch_size=1, sampling_temps=[0.1], backbone_noise=0.0, max_length=200000, omit_AAs=[], print_all=False,
    chain_id_jsonl='', fixed_positions_jsonl='', pssm_jsonl='', omit_AA_jsonl='', bias_AA_jsonl='', tied_positions_jsonl='', bias_by_res_jsonl='', jsonl_path='',
    pssm_threshold=0.0, pssm_multi=0.0, pssm_log_odds_flag=False, pssm_bias_flag=False, write_output_files=False):
    
    if model is None:
        model = get_mpnn_model(model_name=model_name, path_to_model_weights=path_to_model_weights, ca_only=ca_only, backbone_noise=backbone_noise, verbose=print_all)

    if seed:
        seed=seed
    else:
        seed=int(np.random.randint(0, high=999, size=1, dtype=int)[0])

    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)   

    

    NUM_BATCHES = num_seq_per_target//batch_size
    BATCH_COPIES = batch_size
    temperatures = sampling_temps
    omit_AAs_list = omit_AAs
    alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
    alphabet_dict = dict(zip(alphabet, range(21)))    
    omit_AAs_np = np.array([AA in omit_AAs_list for AA in alphabet]).astype(np.float32)
    device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
    if os.path.isfile(chain_id_jsonl):
        with open(chain_id_jsonl, 'r') as json_file:
            json_list = list(json_file)
        for json_str in json_list:
            chain_id_dict = json.loads(json_str)
    else:
        chain_id_dict = None
        if print_all:
            print(40*'-')
            print('chain_id_jsonl is NOT loaded')
        
    if os.path.isfile(fixed_positions_jsonl):
        with open(fixed_positions_jsonl, 'r') as json_file:
            json_list = list(json_file)
        for json_str in json_list:
            fixed_positions_dict = json.loads(json_str)
    else:
        if print_all:
            print(40*'-')
            print('fixed_positions_jsonl is NOT loaded')
        fixed_positions_dict = None


    if os.path.isfile(pssm_jsonl):
        with open(pssm_jsonl, 'r') as json_file:
            json_list = list(json_file)
        pssm_dict = {}
        for json_str in json_list:
            pssm_dict.update(json.loads(json_str))
    else:
        if print_all:
            print(40*'-')
            print('pssm_jsonl is NOT loaded')
        pssm_dict = None


    if os.path.isfile(omit_AA_jsonl):
        with open(omit_AA_jsonl, 'r') as json_file:
            json_list = list(json_file)
        for json_str in json_list:
            omit_AA_dict = json.loads(json_str)
    else:
        if print_all:
            print(40*'-')
            print('omit_AA_jsonl is NOT loaded')
        omit_AA_dict = None


    if os.path.isfile(bias_AA_jsonl):
        with open(bias_AA_jsonl, 'r') as json_file:
            json_list = list(json_file)
        for json_str in json_list:
            bias_AA_dict = json.loads(json_str)
    else:
        if print_all:
            print(40*'-')
            print('bias_AA_jsonl is NOT loaded')
        bias_AA_dict = None


    if os.path.isfile(tied_positions_jsonl):
        with open(tied_positions_jsonl, 'r') as json_file:
            json_list = list(json_file)
        for json_str in json_list:
            tied_positions_dict = json.loads(json_str)
    else:
        if print_all:
            print(40*'-')
            print('tied_positions_jsonl is NOT loaded')
        tied_positions_dict = None


    if os.path.isfile(bias_by_res_jsonl):
        with open(bias_by_res_jsonl, 'r') as json_file:
            json_list = list(json_file)

        for json_str in json_list:
            bias_by_res_dict = json.loads(json_str)
        if print_all:
            print('bias by residue dictionary is loaded')
    else:
        if print_all:
            print(40*'-')
            print('bias by residue dictionary is not loaded, or not provided')
        bias_by_res_dict = None


    if print_all: 
        print(40*'-')
    bias_AAs_np = np.zeros(len(alphabet))
    if bias_AA_dict:
            for n, AA in enumerate(alphabet):
                    if AA in list(bias_AA_dict.keys()):
                            bias_AAs_np[n] = bias_AA_dict[AA]

    if pdb_path:
        pdb_dict_list = parse_PDB(pdb_path, ca_only=ca_only)
        dataset_valid = StructureDatasetPDB(pdb_dict_list, truncate=None, max_length=max_length)
        all_chain_list = [item[-1:] for item in list(pdb_dict_list[0]) if item[:9]=='seq_chain'] #['A','B', 'C',...]
        if pdb_path_chains:
            designed_chain_list = [str(item) for item in pdb_path_chains.split()]
        else:
            designed_chain_list = all_chain_list
        fixed_chain_list = [letter for letter in all_chain_list if letter not in designed_chain_list]
        chain_id_dict = {}
        chain_id_dict[pdb_dict_list[0]['name']]= (designed_chain_list, fixed_chain_list)
    else:
        dataset_valid = StructureDataset(jsonl_path, truncate=None, max_length=max_length, verbose=print_all)

    # Build paths for experiment
    if write_output_files:
        folder_for_outputs = out_folder
        base_folder = folder_for_outputs
        if base_folder[-1] != '/':
            base_folder = base_folder + '/'
        if not os.path.exists(base_folder):
            os.makedirs(base_folder)
        if not os.path.exists(base_folder + 'seqs'):
            os.makedirs(base_folder + 'seqs')

    # if args.save_score:
    #     if not os.path.exists(base_folder + 'scores'):
    #         os.makedirs(base_folder + 'scores')

    # if args.score_only:
    #     if not os.path.exists(base_folder + 'score_only'):
    #         os.makedirs(base_folder + 'score_only')


    # if args.conditional_probs_only:
    #     if not os.path.exists(base_folder + 'conditional_probs_only'):
    #         os.makedirs(base_folder + 'conditional_probs_only')

    # if args.unconditional_probs_only:
    #     if not os.path.exists(base_folder + 'unconditional_probs_only'):
    #         os.makedirs(base_folder + 'unconditional_probs_only')

    # if args.save_probs:
    #     if not os.path.exists(base_folder + 'probs'):
    #         os.makedirs(base_folder + 'probs') 

    # Timing
    start_time = time.time()
    total_residues = 0
    protein_list = []
    total_step = 0
    # Validation epoch
    new_mpnn_seqs = []
    with torch.no_grad():
        test_sum, test_weights = 0., 0.
        for ix, protein in enumerate(dataset_valid):
            score_list = []
            global_score_list = []
            all_probs_list = []
            all_log_probs_list = []
            S_sample_list = []
            batch_clones = [copy.deepcopy(protein) for i in range(BATCH_COPIES)]
            X, S, mask, lengths, chain_M, chain_encoding_all, chain_list_list, visible_list_list, masked_list_list, masked_chain_length_list_list, chain_M_pos, omit_AA_mask, residue_idx, dihedral_mask, tied_pos_list_of_lists_list, pssm_coef, pssm_bias, pssm_log_odds_all, bias_by_res_all, tied_beta = tied_featurize(batch_clones, device, chain_id_dict, fixed_positions_dict, omit_AA_dict, tied_positions_dict, pssm_dict, bias_by_res_dict, ca_only=ca_only)
            pssm_log_odds_mask = (pssm_log_odds_all > pssm_threshold).float() #1.0 for true, 0.0 for false
            name_ = batch_clones[0]['name']
            if False:
                pass
            # if args.score_only:
            #     loop_c = 0 
            #     if args.path_to_fasta:
            #         fasta_names, fasta_seqs = parse_fasta(args.path_to_fasta, omit=["/"])
            #         loop_c = len(fasta_seqs)
            #     for fc in range(1+loop_c):
            #         if fc == 0:
            #             structure_sequence_score_file = base_folder + '/score_only/' + batch_clones[0]['name'] + f'_pdb'
            #         else:
            #             structure_sequence_score_file = base_folder + '/score_only/' + batch_clones[0]['name'] + f'_fasta_{fc}'
            #         native_score_list = []
            #         global_native_score_list = []
            #         if fc > 0:
            #             input_seq_length = len(fasta_seqs[fc-1])
            #             S_input = torch.tensor([alphabet_dict[AA] for AA in fasta_seqs[fc-1]], device=device)[None,:].repeat(X.shape[0], 1)
            #             S[:,:input_seq_length] = S_input #assumes that S and S_input are alphabetically sorted for masked_chains
            #         for j in range(NUM_BATCHES):
            #             randn_1 = torch.randn(chain_M.shape, device=X.device)
            #             log_probs = model(X, S, mask, chain_M*chain_M_pos, residue_idx, chain_encoding_all, randn_1)
            #             mask_for_loss = mask*chain_M*chain_M_pos
            #             scores = _scores(S, log_probs, mask_for_loss)
            #             native_score = scores.cpu().data.numpy()
            #             native_score_list.append(native_score)
            #             global_scores = _scores(S, log_probs, mask)
            #             global_native_score = global_scores.cpu().data.numpy()
            #             global_native_score_list.append(global_native_score)
            #         native_score = np.concatenate(native_score_list, 0)
            #         global_native_score = np.concatenate(global_native_score_list, 0)
            #         ns_mean = native_score.mean()
            #         ns_mean_print = np.format_float_positional(np.float32(ns_mean), unique=False, precision=4)
            #         ns_std = native_score.std()
            #         ns_std_print = np.format_float_positional(np.float32(ns_std), unique=False, precision=4)

            #         global_ns_mean = global_native_score.mean()
            #         global_ns_mean_print = np.format_float_positional(np.float32(global_ns_mean), unique=False, precision=4)
            #         global_ns_std = global_native_score.std()
            #         global_ns_std_print = np.format_float_positional(np.float32(global_ns_std), unique=False, precision=4)

            #         ns_sample_size = native_score.shape[0]
            #         seq_str = _S_to_seq(S[0,], chain_M[0,])
            #         np.savez(structure_sequence_score_file, score=native_score, global_score=global_native_score, S=S[0,].cpu().numpy(), seq_str=seq_str)
            #         if print_all:
            #             if fc == 0:
            #                 print(f'Score for {name_} from PDB, mean: {ns_mean_print}, std: {ns_std_print}, sample size: {ns_sample_size},  global score, mean: {global_ns_mean_print}, std: {global_ns_std_print}, sample size: {ns_sample_size}')
            #             else:
            #                 print(f'Score for {name_}_{fc} from FASTA, mean: {ns_mean_print}, std: {ns_std_print}, sample size: {ns_sample_size},  global score, mean: {global_ns_mean_print}, std: {global_ns_std_print}, sample size: {ns_sample_size}')
            # elif args.conditional_probs_only:
            #     if print_all:
            #         print(f'Calculating conditional probabilities for {name_}')
            #     conditional_probs_only_file = base_folder + '/conditional_probs_only/' + batch_clones[0]['name']
            #     log_conditional_probs_list = []
            #     for j in range(NUM_BATCHES):
            #         randn_1 = torch.randn(chain_M.shape, device=X.device)
            #         log_conditional_probs = model.conditional_probs(X, S, mask, chain_M*chain_M_pos, residue_idx, chain_encoding_all, randn_1, args.conditional_probs_only_backbone)
            #         log_conditional_probs_list.append(log_conditional_probs.cpu().numpy())
            #     concat_log_p = np.concatenate(log_conditional_probs_list, 0) #[B, L, 21]
            #     mask_out = (chain_M*chain_M_pos*mask)[0,].cpu().numpy()
            #     np.savez(conditional_probs_only_file, log_p=concat_log_p, S=S[0,].cpu().numpy(), mask=mask[0,].cpu().numpy(), design_mask=mask_out)
            # elif args.unconditional_probs_only:
            #     if print_all:
            #         print(f'Calculating sequence unconditional probabilities for {name_}')
            #     unconditional_probs_only_file = base_folder + '/unconditional_probs_only/' + batch_clones[0]['name']
            #     log_unconditional_probs_list = []
            #     for j in range(NUM_BATCHES):
            #         log_unconditional_probs = model.unconditional_probs(X, mask, residue_idx, chain_encoding_all)
            #         log_unconditional_probs_list.append(log_unconditional_probs.cpu().numpy())
            #     concat_log_p = np.concatenate(log_unconditional_probs_list, 0) #[B, L, 21]
            #     mask_out = (chain_M*chain_M_pos*mask)[0,].cpu().numpy()
            #     np.savez(unconditional_probs_only_file, log_p=concat_log_p, S=S[0,].cpu().numpy(), mask=mask[0,].cpu().numpy(), design_mask=mask_out)
            else:
                randn_1 = torch.randn(chain_M.shape, device=X.device)
                log_probs = model(X, S, mask, chain_M*chain_M_pos, residue_idx, chain_encoding_all, randn_1)
                mask_for_loss = mask*chain_M*chain_M_pos
                scores = _scores(S, log_probs, mask_for_loss) #score only the redesigned part
                native_score = scores.cpu().data.numpy()
                global_scores = _scores(S, log_probs, mask) #score the whole structure-sequence
                global_native_score = global_scores.cpu().data.numpy()
                # Generate some sequences
                if write_output_files:
                    ali_file = base_folder + '/seqs/' + batch_clones[0]['name'] + '.fa'
                    score_file = base_folder + '/scores/' + batch_clones[0]['name'] + '.npz'
                    probs_file = base_folder + '/probs/' + batch_clones[0]['name'] + '.npz'
                    f = open(ali_file, 'w')
                if print_all:
                    print(f'Generating sequences for: {name_}')
                t0 = time.time()
                for temp in temperatures:
                    for j in range(NUM_BATCHES):
                        randn_2 = torch.randn(chain_M.shape, device=X.device)
                        if tied_positions_dict == None:
                            sample_dict = model.sample(X, randn_2, S, chain_M, chain_encoding_all, residue_idx, mask=mask, temperature=temp, omit_AAs_np=omit_AAs_np, bias_AAs_np=bias_AAs_np, chain_M_pos=chain_M_pos, omit_AA_mask=omit_AA_mask, pssm_coef=pssm_coef, pssm_bias=pssm_bias, pssm_multi=pssm_multi, pssm_log_odds_flag=bool(pssm_log_odds_flag), pssm_log_odds_mask=pssm_log_odds_mask, pssm_bias_flag=bool(pssm_bias_flag), bias_by_res=bias_by_res_all)
                            S_sample = sample_dict["S"] 
                        else:
                            sample_dict = model.tied_sample(X, randn_2, S, chain_M, chain_encoding_all, residue_idx, mask=mask, temperature=temp, omit_AAs_np=omit_AAs_np, bias_AAs_np=bias_AAs_np, chain_M_pos=chain_M_pos, omit_AA_mask=omit_AA_mask, pssm_coef=pssm_coef, pssm_bias=pssm_bias, pssm_multi=pssm_multi, pssm_log_odds_flag=bool(pssm_log_odds_flag), pssm_log_odds_mask=pssm_log_odds_mask, pssm_bias_flag=bool(pssm_bias_flag), tied_pos=tied_pos_list_of_lists_list[0], tied_beta=tied_beta, bias_by_res=bias_by_res_all)
                        # Compute scores
                            S_sample = sample_dict["S"]
                        log_probs = model(X, S_sample, mask, chain_M*chain_M_pos, residue_idx, chain_encoding_all, randn_2, use_input_decoding_order=True, decoding_order=sample_dict["decoding_order"])
                        mask_for_loss = mask*chain_M*chain_M_pos
                        scores = _scores(S_sample, log_probs, mask_for_loss)
                        scores = scores.cpu().data.numpy()
                        
                        global_scores = _scores(S_sample, log_probs, mask) #score the whole structure-sequence
                        global_scores = global_scores.cpu().data.numpy()
                        
                        all_probs_list.append(sample_dict["probs"].cpu().data.numpy())
                        all_log_probs_list.append(log_probs.cpu().data.numpy())
                        S_sample_list.append(S_sample.cpu().data.numpy())
                        for b_ix in range(BATCH_COPIES):
                            masked_chain_length_list = masked_chain_length_list_list[b_ix]
                            masked_list = masked_list_list[b_ix]
                            seq_recovery_rate = torch.sum(torch.sum(torch.nn.functional.one_hot(S[b_ix], 21)*torch.nn.functional.one_hot(S_sample[b_ix], 21),axis=-1)*mask_for_loss[b_ix])/torch.sum(mask_for_loss[b_ix])
                            seq = _S_to_seq(S_sample[b_ix], chain_M[b_ix])
                            new_mpnn_seqs.append(seq)
                            score = scores[b_ix]
                            score_list.append(score)
                            global_score = global_scores[b_ix]
                            global_score_list.append(global_score)
                            native_seq = _S_to_seq(S[b_ix], chain_M[b_ix])
                            if b_ix == 0 and j==0 and temp==temperatures[0]:
                                start = 0
                                end = 0
                                list_of_AAs = []
                                for mask_l in masked_chain_length_list:
                                    end += mask_l
                                    list_of_AAs.append(native_seq[start:end])
                                    start = end
                                native_seq = "".join(list(np.array(list_of_AAs)[np.argsort(masked_list)]))
                                l0 = 0
                                for mc_length in list(np.array(masked_chain_length_list)[np.argsort(masked_list)])[:-1]:
                                    l0 += mc_length
                                    native_seq = native_seq[:l0] + '/' + native_seq[l0:]
                                    l0 += 1
                                sorted_masked_chain_letters = np.argsort(masked_list_list[0])
                                print_masked_chains = [masked_list_list[0][i] for i in sorted_masked_chain_letters]
                                sorted_visible_chain_letters = np.argsort(visible_list_list[0])
                                print_visible_chains = [visible_list_list[0][i] for i in sorted_visible_chain_letters]
                                native_score_print = np.format_float_positional(np.float32(native_score.mean()), unique=False, precision=4)
                                global_native_score_print = np.format_float_positional(np.float32(global_native_score.mean()), unique=False, precision=4)
                                script_dir = os.path.dirname(os.path.realpath(__file__))
                                try:
                                    commit_str = subprocess.check_output(f'git --git-dir {script_dir}/.git rev-parse HEAD', shell=True, stderr=subprocess.DEVNULL).decode().strip()
                                except subprocess.CalledProcessError:
                                    commit_str = 'unknown'
                                if ca_only:
                                    print_model_name = 'CA_model_name'
                                else:
                                    print_model_name = 'model_name'
                                if write_output_files:
                                    f.write('>{}, score={}, global_score={}, fixed_chains={}, designed_chains={}, {}={}, git_hash={}, seed={}\n{}\n'.format(name_, native_score_print, global_native_score_print, print_visible_chains, print_masked_chains, print_model_name, model_name, commit_str, seed, native_seq)) #write the native sequence
                            start = 0
                            end = 0
                            list_of_AAs = []
                            for mask_l in masked_chain_length_list:
                                end += mask_l
                                list_of_AAs.append(seq[start:end])
                                start = end

                            seq = "".join(list(np.array(list_of_AAs)[np.argsort(masked_list)]))
                            l0 = 0
                            for mc_length in list(np.array(masked_chain_length_list)[np.argsort(masked_list)])[:-1]:
                                l0 += mc_length
                                seq = seq[:l0] + '/' + seq[l0:]
                                l0 += 1
                            score_print = np.format_float_positional(np.float32(score), unique=False, precision=4)
                            global_score_print = np.format_float_positional(np.float32(global_score), unique=False, precision=4)
                            seq_rec_print = np.format_float_positional(np.float32(seq_recovery_rate.detach().cpu().numpy()), unique=False, precision=4)
                            sample_number = j*BATCH_COPIES+b_ix+1
                            if write_output_files:
                                f.write('>T={}, sample={}, score={}, global_score={}, seq_recovery={}\n{}\n'.format(temp,sample_number,score_print,global_score_print,seq_rec_print,seq)) #write generated sequence
                # if args.save_score:
                #     np.savez(score_file, score=np.array(score_list, np.float32), global_score=np.array(global_score_list, np.float32))
                # if args.save_probs:
                #     all_probs_concat = np.concatenate(all_probs_list)
                #     all_log_probs_concat = np.concatenate(all_log_probs_list)
                #     S_sample_concat = np.concatenate(S_sample_list)
                #     np.savez(probs_file, probs=np.array(all_probs_concat, np.float32), log_probs=np.array(all_log_probs_concat, np.float32), S=np.array(S_sample_concat, np.int32), mask=mask_for_loss.cpu().data.numpy(), chain_order=chain_list_list)
                t1 = time.time()
                dt = round(float(t1-t0), 4)
                num_seqs = len(temperatures)*NUM_BATCHES*BATCH_COPIES
                total_length = X.shape[1]
                if print_all:
                    print(f'{num_seqs} sequences of length {total_length} generated in {dt} seconds')
                if write_output_files:
                    f.close()

    return new_mpnn_seqs


def parse_fasta(filename,limit=-1, omit=[]):
    header = []
    sequence = []
    lines = open(filename, "r")
    for line in lines:
        line = line.rstrip()
        if line[0] == ">":
            if len(header) == limit:
                break
            header.append(line[1:])
            sequence.append([])
        else:
            if omit:
                line = [item for item in line if item not in omit]
                line = ''.join(line)
            line = ''.join(line)
            sequence[-1].append(line)
    lines.close()
    sequence = [''.join(seq) for seq in sequence]
    return np.array(header), np.array(sequence)

def _scores(S, log_probs, mask):
    """ Negative log probabilities """
    criterion = torch.nn.NLLLoss(reduction='none')
    loss = criterion(
        log_probs.contiguous().view(-1,log_probs.size(-1)),
        S.contiguous().view(-1)
    ).view(S.size())
    scores = torch.sum(loss * mask, dim=-1) / torch.sum(mask, dim=-1)
    return scores

def _S_to_seq(S, mask):
    alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
    seq = ''.join([alphabet[c] for c, m in zip(S.tolist(), mask.tolist()) if m > 0])
    return seq

def parse_PDB_biounits(x, atoms=['N','CA','C'], chain=None):
  '''
  input:  x = PDB filename
          atoms = atoms to extract (optional)
  output: (length, atoms, coords=(x,y,z)), sequence
  '''

  alpha_1 = list("ARNDCQEGHILKMFPSTWYV-")
  states = len(alpha_1)
  alpha_3 = ['ALA','ARG','ASN','ASP','CYS','GLN','GLU','GLY','HIS','ILE',
             'LEU','LYS','MET','PHE','PRO','SER','THR','TRP','TYR','VAL','GAP']
  
  aa_1_N = {a:n for n,a in enumerate(alpha_1)}
  aa_3_N = {a:n for n,a in enumerate(alpha_3)}
  aa_N_1 = {n:a for n,a in enumerate(alpha_1)}
  aa_1_3 = {a:b for a,b in zip(alpha_1,alpha_3)}
  aa_3_1 = {b:a for a,b in zip(alpha_1,alpha_3)}
  
  def AA_to_N(x):
    # ["ARND"] -> [[0,1,2,3]]
    x = np.array(x);
    if x.ndim == 0: x = x[None]
    return [[aa_1_N.get(a, states-1) for a in y] for y in x]
  
  def N_to_AA(x):
    # [[0,1,2,3]] -> ["ARND"]
    x = np.array(x);
    if x.ndim == 1: x = x[None]
    return ["".join([aa_N_1.get(a,"-") for a in y]) for y in x]

  xyz,seq,min_resn,max_resn = {},{},1e6,-1e6
  for line in open(x,"rb"):
    line = line.decode("utf-8","ignore").rstrip()

    if line[:6] == "HETATM" and line[17:17+3] == "MSE":
      line = line.replace("HETATM","ATOM  ")
      line = line.replace("MSE","MET")

    if line[:4] == "ATOM":
      ch = line[21:22]
      if ch == chain or chain is None:
        atom = line[12:12+4].strip()
        resi = line[17:17+3]
        resn = line[22:22+5].strip()
        x,y,z = [float(line[i:(i+8)]) for i in [30,38,46]]

        if resn[-1].isalpha(): 
            resa,resn = resn[-1],int(resn[:-1])-1
        else: 
            resa,resn = "",int(resn)-1
#         resn = int(resn)
        if resn < min_resn: 
            min_resn = resn
        if resn > max_resn: 
            max_resn = resn
        if resn not in xyz: 
            xyz[resn] = {}
        if resa not in xyz[resn]: 
            xyz[resn][resa] = {}
        if resn not in seq: 
            seq[resn] = {}
        if resa not in seq[resn]: 
            seq[resn][resa] = resi

        if atom not in xyz[resn][resa]:
          xyz[resn][resa][atom] = np.array([x,y,z])

  # convert to numpy arrays, fill in missing values
  seq_,xyz_ = [],[]
  try:
      for resn in range(min_resn,max_resn+1):
        if resn in seq:
          for k in sorted(seq[resn]): seq_.append(aa_3_N.get(seq[resn][k],20))
        else: seq_.append(20)
        if resn in xyz:
          for k in sorted(xyz[resn]):
            for atom in atoms:
              if atom in xyz[resn][k]: xyz_.append(xyz[resn][k][atom])
              else: xyz_.append(np.full(3,np.nan))
        else:
          for atom in atoms: xyz_.append(np.full(3,np.nan))
      return np.array(xyz_).reshape(-1,len(atoms),3), N_to_AA(np.array(seq_))
  except TypeError:
      return 'no_chain', 'no_chain'

def parse_PDB(path_to_pdb, input_chain_list=None, ca_only=False):
    c=0
    pdb_dict_list = []
    init_alphabet = ['A', 'B', 'C', 'D', 'E', 'F', 'G','H', 'I', 'J','K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T','U', 'V','W','X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g','h', 'i', 'j','k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't','u', 'v','w','x', 'y', 'z']
    extra_alphabet = [str(item) for item in list(np.arange(300))]
    chain_alphabet = init_alphabet + extra_alphabet
     
    if input_chain_list:
        chain_alphabet = input_chain_list  
 

    biounit_names = [path_to_pdb]
    for biounit in biounit_names:
        my_dict = {}
        s = 0
        concat_seq = ''
        concat_N = []
        concat_CA = []
        concat_C = []
        concat_O = []
        concat_mask = []
        coords_dict = {}
        for letter in chain_alphabet:
            if ca_only:
                sidechain_atoms = ['CA']
            else:
                sidechain_atoms = ['N', 'CA', 'C', 'O']
            xyz, seq = parse_PDB_biounits(biounit, atoms=sidechain_atoms, chain=letter)
            if type(xyz) != str:
                concat_seq += seq[0]
                my_dict['seq_chain_'+letter]=seq[0]
                coords_dict_chain = {}
                if ca_only:
                    coords_dict_chain['CA_chain_'+letter]=xyz.tolist()
                else:
                    coords_dict_chain['N_chain_' + letter] = xyz[:, 0, :].tolist()
                    coords_dict_chain['CA_chain_' + letter] = xyz[:, 1, :].tolist()
                    coords_dict_chain['C_chain_' + letter] = xyz[:, 2, :].tolist()
                    coords_dict_chain['O_chain_' + letter] = xyz[:, 3, :].tolist()
                my_dict['coords_chain_'+letter]=coords_dict_chain
                s += 1
        fi = biounit.rfind("/")
        my_dict['name']=biounit[(fi+1):-4]
        my_dict['num_of_chains'] = s
        my_dict['seq'] = concat_seq
        if s <= len(chain_alphabet):
            pdb_dict_list.append(my_dict)
            c+=1
    return pdb_dict_list



def tied_featurize(batch, device, chain_dict, fixed_position_dict=None, omit_AA_dict=None, tied_positions_dict=None, pssm_dict=None, bias_by_res_dict=None, ca_only=False):
    """ Pack and pad batch into torch tensors """
    alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
    B = len(batch)
    lengths = np.array([len(b['seq']) for b in batch], dtype=np.int32) #sum of chain seq lengths
    L_max = max([len(b['seq']) for b in batch])
    if ca_only:
        X = np.zeros([B, L_max, 1, 3])
    else:
        X = np.zeros([B, L_max, 4, 3])
    residue_idx = -100*np.ones([B, L_max], dtype=np.int32)
    chain_M = np.zeros([B, L_max], dtype=np.int32) #1.0 for the bits that need to be predicted
    pssm_coef_all = np.zeros([B, L_max], dtype=np.float32) #1.0 for the bits that need to be predicted
    pssm_bias_all = np.zeros([B, L_max, 21], dtype=np.float32) #1.0 for the bits that need to be predicted
    pssm_log_odds_all = 10000.0*np.ones([B, L_max, 21], dtype=np.float32) #1.0 for the bits that need to be predicted
    chain_M_pos = np.zeros([B, L_max], dtype=np.int32) #1.0 for the bits that need to be predicted
    bias_by_res_all = np.zeros([B, L_max, 21], dtype=np.float32)
    chain_encoding_all = np.zeros([B, L_max], dtype=np.int32) #1.0 for the bits that need to be predicted
    S = np.zeros([B, L_max], dtype=np.int32)
    omit_AA_mask = np.zeros([B, L_max, len(alphabet)], dtype=np.int32)
    # Build the batch
    letter_list_list = []
    visible_list_list = []
    masked_list_list = []
    masked_chain_length_list_list = []
    tied_pos_list_of_lists_list = []
    for i, b in enumerate(batch):
        if chain_dict != None:
            masked_chains, visible_chains = chain_dict[b['name']] #masked_chains a list of chain letters to predict [A, D, F]
        else:
            masked_chains = [item[-1:] for item in list(b) if item[:10]=='seq_chain_']
            visible_chains = []
        masked_chains.sort() #sort masked_chains 
        visible_chains.sort() #sort visible_chains 
        all_chains = masked_chains + visible_chains
    for i, b in enumerate(batch):
        mask_dict = {}
        a = 0
        x_chain_list = []
        chain_mask_list = []
        chain_seq_list = []
        chain_encoding_list = []
        c = 1
        letter_list = []
        global_idx_start_list = [0]
        visible_list = []
        masked_list = []
        masked_chain_length_list = []
        fixed_position_mask_list = []
        omit_AA_mask_list = []
        pssm_coef_list = []
        pssm_bias_list = []
        pssm_log_odds_list = []
        bias_by_res_list = []
        l0 = 0
        l1 = 0
        for step, letter in enumerate(all_chains):
            if letter in visible_chains:
                letter_list.append(letter)
                visible_list.append(letter)
                chain_seq = b[f'seq_chain_{letter}']
                chain_seq = ''.join([a if a!='-' else 'X' for a in chain_seq])
                chain_length = len(chain_seq)
                global_idx_start_list.append(global_idx_start_list[-1]+chain_length)
                chain_coords = b[f'coords_chain_{letter}'] #this is a dictionary
                chain_mask = np.zeros(chain_length) #0.0 for visible chains
                if ca_only:
                    x_chain = np.array(chain_coords[f'CA_chain_{letter}']) #[chain_lenght,1,3] #CA_diff
                    if len(x_chain.shape) == 2:
                        x_chain = x_chain[:,None,:]
                else:
                    x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_lenght,4,3]
                x_chain_list.append(x_chain)
                chain_mask_list.append(chain_mask)
                chain_seq_list.append(chain_seq)
                chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
                l1 += chain_length
                residue_idx[i, l0:l1] = 100*(c-1)+np.arange(l0, l1)
                l0 += chain_length
                c+=1
                fixed_position_mask = np.ones(chain_length)
                fixed_position_mask_list.append(fixed_position_mask)
                omit_AA_mask_temp = np.zeros([chain_length, len(alphabet)], np.int32)
                omit_AA_mask_list.append(omit_AA_mask_temp)
                pssm_coef = np.zeros(chain_length)
                pssm_bias = np.zeros([chain_length, 21])
                pssm_log_odds = 10000.0*np.ones([chain_length, 21])
                pssm_coef_list.append(pssm_coef)
                pssm_bias_list.append(pssm_bias)
                pssm_log_odds_list.append(pssm_log_odds)
                bias_by_res_list.append(np.zeros([chain_length, 21]))
            if letter in masked_chains:
                masked_list.append(letter)
                letter_list.append(letter)
                chain_seq = b[f'seq_chain_{letter}']
                chain_seq = ''.join([a if a!='-' else 'X' for a in chain_seq])
                chain_length = len(chain_seq)
                global_idx_start_list.append(global_idx_start_list[-1]+chain_length)
                masked_chain_length_list.append(chain_length)
                chain_coords = b[f'coords_chain_{letter}'] #this is a dictionary
                chain_mask = np.ones(chain_length) #1.0 for masked
                if ca_only:
                    x_chain = np.array(chain_coords[f'CA_chain_{letter}']) #[chain_lenght,1,3] #CA_diff
                    if len(x_chain.shape) == 2:
                        x_chain = x_chain[:,None,:]
                else:
                    x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_lenght,4,3]               
                x_chain_list.append(x_chain)
                chain_mask_list.append(chain_mask)
                chain_seq_list.append(chain_seq)
                chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
                l1 += chain_length
                residue_idx[i, l0:l1] = 100*(c-1)+np.arange(l0, l1)
                l0 += chain_length
                c+=1
                fixed_position_mask = np.ones(chain_length)
                if fixed_position_dict!=None:
                    fixed_pos_list = fixed_position_dict[b['name']][letter]
                    if fixed_pos_list:
                        fixed_position_mask[np.array(fixed_pos_list)-1] = 0.0
                fixed_position_mask_list.append(fixed_position_mask)
                omit_AA_mask_temp = np.zeros([chain_length, len(alphabet)], np.int32)
                if omit_AA_dict!=None:
                    for item in omit_AA_dict[b['name']][letter]:
                        idx_AA = np.array(item[0])-1
                        AA_idx = np.array([np.argwhere(np.array(list(alphabet))== AA)[0][0] for AA in item[1]]).repeat(idx_AA.shape[0])
                        idx_ = np.array([[a, b] for a in idx_AA for b in AA_idx])
                        omit_AA_mask_temp[idx_[:,0], idx_[:,1]] = 1
                omit_AA_mask_list.append(omit_AA_mask_temp)
                pssm_coef = np.zeros(chain_length)
                pssm_bias = np.zeros([chain_length, 21])
                pssm_log_odds = 10000.0*np.ones([chain_length, 21])
                if pssm_dict:
                    if pssm_dict[b['name']][letter]:
                        pssm_coef = pssm_dict[b['name']][letter]['pssm_coef']
                        pssm_bias = pssm_dict[b['name']][letter]['pssm_bias']
                        pssm_log_odds = pssm_dict[b['name']][letter]['pssm_log_odds']
                pssm_coef_list.append(pssm_coef)
                pssm_bias_list.append(pssm_bias)
                pssm_log_odds_list.append(pssm_log_odds)
                if bias_by_res_dict:
                    bias_by_res_list.append(bias_by_res_dict[b['name']][letter])
                else:
                    bias_by_res_list.append(np.zeros([chain_length, 21]))

       
        letter_list_np = np.array(letter_list)
        tied_pos_list_of_lists = []
        tied_beta = np.ones(L_max)
        if tied_positions_dict!=None:
            tied_pos_list = tied_positions_dict[b['name']]
            if tied_pos_list:
                set_chains_tied = set(list(itertools.chain(*[list(item) for item in tied_pos_list])))
                for tied_item in tied_pos_list:
                    one_list = []
                    for k, v in tied_item.items():
                        start_idx = global_idx_start_list[np.argwhere(letter_list_np == k)[0][0]]
                        if isinstance(v[0], list):
                            for v_count in range(len(v[0])):
                                one_list.append(start_idx+v[0][v_count]-1)#make 0 to be the first
                                tied_beta[start_idx+v[0][v_count]-1] = v[1][v_count]
                        else:
                            for v_ in v:
                                one_list.append(start_idx+v_-1)#make 0 to be the first
                    tied_pos_list_of_lists.append(one_list)
        tied_pos_list_of_lists_list.append(tied_pos_list_of_lists)


 
        x = np.concatenate(x_chain_list,0) #[L, 4, 3]
        all_sequence = "".join(chain_seq_list)
        m = np.concatenate(chain_mask_list,0) #[L,], 1.0 for places that need to be predicted
        chain_encoding = np.concatenate(chain_encoding_list,0)
        m_pos = np.concatenate(fixed_position_mask_list,0) #[L,], 1.0 for places that need to be predicted

        pssm_coef_ = np.concatenate(pssm_coef_list,0) #[L,], 1.0 for places that need to be predicted
        pssm_bias_ = np.concatenate(pssm_bias_list,0) #[L,], 1.0 for places that need to be predicted
        pssm_log_odds_ = np.concatenate(pssm_log_odds_list,0) #[L,], 1.0 for places that need to be predicted

        bias_by_res_ = np.concatenate(bias_by_res_list, 0)  #[L,21], 0.0 for places where AA frequencies don't need to be tweaked

        l = len(all_sequence)
        x_pad = np.pad(x, [[0,L_max-l], [0,0], [0,0]], 'constant', constant_values=(np.nan, ))
        X[i,:,:,:] = x_pad

        m_pad = np.pad(m, [[0,L_max-l]], 'constant', constant_values=(0.0, ))
        m_pos_pad = np.pad(m_pos, [[0,L_max-l]], 'constant', constant_values=(0.0, ))
        omit_AA_mask_pad = np.pad(np.concatenate(omit_AA_mask_list,0), [[0,L_max-l]], 'constant', constant_values=(0.0, ))
        chain_M[i,:] = m_pad
        chain_M_pos[i,:] = m_pos_pad
        omit_AA_mask[i,] = omit_AA_mask_pad

        chain_encoding_pad = np.pad(chain_encoding, [[0,L_max-l]], 'constant', constant_values=(0.0, ))
        chain_encoding_all[i,:] = chain_encoding_pad

        pssm_coef_pad = np.pad(pssm_coef_, [[0,L_max-l]], 'constant', constant_values=(0.0, ))
        pssm_bias_pad = np.pad(pssm_bias_, [[0,L_max-l], [0,0]], 'constant', constant_values=(0.0, ))
        pssm_log_odds_pad = np.pad(pssm_log_odds_, [[0,L_max-l], [0,0]], 'constant', constant_values=(0.0, ))

        pssm_coef_all[i,:] = pssm_coef_pad
        pssm_bias_all[i,:] = pssm_bias_pad
        pssm_log_odds_all[i,:] = pssm_log_odds_pad

        bias_by_res_pad = np.pad(bias_by_res_, [[0,L_max-l], [0,0]], 'constant', constant_values=(0.0, ))
        bias_by_res_all[i,:] = bias_by_res_pad

        # Convert to labels
        indices = np.asarray([alphabet.index(a) for a in all_sequence], dtype=np.int32)
        S[i, :l] = indices
        letter_list_list.append(letter_list)
        visible_list_list.append(visible_list)
        masked_list_list.append(masked_list)
        masked_chain_length_list_list.append(masked_chain_length_list)


    isnan = np.isnan(X)
    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32)
    X[isnan] = 0.

    # Conversion
    pssm_coef_all = torch.from_numpy(pssm_coef_all).to(dtype=torch.float32, device=device)
    pssm_bias_all = torch.from_numpy(pssm_bias_all).to(dtype=torch.float32, device=device)
    pssm_log_odds_all = torch.from_numpy(pssm_log_odds_all).to(dtype=torch.float32, device=device)

    tied_beta = torch.from_numpy(tied_beta).to(dtype=torch.float32, device=device)

    jumps = ((residue_idx[:,1:]-residue_idx[:,:-1])==1).astype(np.float32)
    bias_by_res_all = torch.from_numpy(bias_by_res_all).to(dtype=torch.float32, device=device)
    phi_mask = np.pad(jumps, [[0,0],[1,0]])
    psi_mask = np.pad(jumps, [[0,0],[0,1]])
    omega_mask = np.pad(jumps, [[0,0],[0,1]])
    dihedral_mask = np.concatenate([phi_mask[:,:,None], psi_mask[:,:,None], omega_mask[:,:,None]], -1) #[B,L,3]
    dihedral_mask = torch.from_numpy(dihedral_mask).to(dtype=torch.float32, device=device)
    residue_idx = torch.from_numpy(residue_idx).to(dtype=torch.long,device=device)
    S = torch.from_numpy(S).to(dtype=torch.long,device=device)
    X = torch.from_numpy(X).to(dtype=torch.float32, device=device)
    mask = torch.from_numpy(mask).to(dtype=torch.float32, device=device)
    chain_M = torch.from_numpy(chain_M).to(dtype=torch.float32, device=device)
    chain_M_pos = torch.from_numpy(chain_M_pos).to(dtype=torch.float32, device=device)
    omit_AA_mask = torch.from_numpy(omit_AA_mask).to(dtype=torch.float32, device=device)
    chain_encoding_all = torch.from_numpy(chain_encoding_all).to(dtype=torch.long, device=device)
    if ca_only:
        X_out = X[:,:,0]
    else:
        X_out = X
    return X_out, S, mask, lengths, chain_M, chain_encoding_all, letter_list_list, visible_list_list, masked_list_list, masked_chain_length_list_list, chain_M_pos, omit_AA_mask, residue_idx, dihedral_mask, tied_pos_list_of_lists_list, pssm_coef_all, pssm_bias_all, pssm_log_odds_all, bias_by_res_all, tied_beta



def loss_nll(S, log_probs, mask):
    """ Negative log probabilities """
    criterion = torch.nn.NLLLoss(reduction='none')
    loss = criterion(
        log_probs.contiguous().view(-1, log_probs.size(-1)), S.contiguous().view(-1)
    ).view(S.size())
    loss_av = torch.sum(loss * mask) / torch.sum(mask)
    return loss, loss_av


def loss_smoothed(S, log_probs, mask, weight=0.1):
    """ Negative log probabilities """
    S_onehot = torch.nn.functional.one_hot(S, 21).float()

    # Label smoothing
    S_onehot = S_onehot + weight / float(S_onehot.size(-1))
    S_onehot = S_onehot / S_onehot.sum(-1, keepdim=True)

    loss = -(S_onehot * log_probs).sum(-1)
    loss_av = torch.sum(loss * mask) / torch.sum(mask)
    return loss, loss_av

class StructureDataset():
    def __init__(self, jsonl_file, verbose=True, truncate=None, max_length=100,
        alphabet='ACDEFGHIKLMNPQRSTVWYX-'):
        alphabet_set = set([a for a in alphabet])
        discard_count = {
            'bad_chars': 0,
            'too_long': 0,
            'bad_seq_length': 0
        }

        with open(jsonl_file) as f:
            self.data = []

            lines = f.readlines()
            start = time.time()
            for i, line in enumerate(lines):
                entry = json.loads(line)
                seq = entry['seq'] 
                name = entry['name']

                # Convert raw coords to np arrays
                #for key, val in entry['coords'].items():
                #    entry['coords'][key] = np.asarray(val)

                # Check if in alphabet
                bad_chars = set([s for s in seq]).difference(alphabet_set)
                if len(bad_chars) == 0:
                    if len(entry['seq']) <= max_length:
                        if True:
                            self.data.append(entry)
                        else:
                            discard_count['bad_seq_length'] += 1
                    else:
                        discard_count['too_long'] += 1
                else:
                    if verbose:
                        print(name, bad_chars, entry['seq'])
                    discard_count['bad_chars'] += 1

                # Truncate early
                if truncate is not None and len(self.data) == truncate:
                    return

                if verbose and (i + 1) % 1000 == 0:
                    elapsed = time.time() - start
                    print('{} entries ({} loaded) in {:.1f} s'.format(len(self.data), i+1, elapsed))
            if verbose:
                print('discarded', discard_count)
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]
    

class StructureDatasetPDB():
    def __init__(self, pdb_dict_list, verbose=True, truncate=None, max_length=100,
        alphabet='ACDEFGHIKLMNPQRSTVWYX-'):
        alphabet_set = set([a for a in alphabet])
        discard_count = {
            'bad_chars': 0,
            'too_long': 0,
            'bad_seq_length': 0
        }

        self.data = []

        start = time.time()
        for i, entry in enumerate(pdb_dict_list):
            seq = entry['seq']
            name = entry['name']

            bad_chars = set([s for s in seq]).difference(alphabet_set)
            if len(bad_chars) == 0:
                if len(entry['seq']) <= max_length:
                    self.data.append(entry)
                else:
                    discard_count['too_long'] += 1
            else:
                discard_count['bad_chars'] += 1

            # Truncate early
            if truncate is not None and len(self.data) == truncate:
                return

            if verbose and (i + 1) % 1000 == 0:
                elapsed = time.time() - start

            #print('Discarded', discard_count)
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


    
class StructureLoader():
    def __init__(self, dataset, batch_size=100, shuffle=True,
        collate_fn=lambda x:x, drop_last=False):
        self.dataset = dataset
        self.size = len(dataset)
        self.lengths = [len(dataset[i]['seq']) for i in range(self.size)]
        self.batch_size = batch_size
        sorted_ix = np.argsort(self.lengths)

        # Cluster into batches of similar sizes
        clusters, batch = [], []
        batch_max = 0
        for ix in sorted_ix:
            size = self.lengths[ix]
            if size * (len(batch) + 1) <= self.batch_size:
                batch.append(ix)
                batch_max = size
            else:
                clusters.append(batch)
                batch, batch_max = [], 0
        if len(batch) > 0:
            clusters.append(batch)
        self.clusters = clusters

    def __len__(self):
        return len(self.clusters)

    def __iter__(self):
        np.random.shuffle(self.clusters)
        for b_idx in self.clusters:
            batch = [self.dataset[i] for i in b_idx]
            yield batch
            
            
            
# The following gather functions
def gather_edges(edges, neighbor_idx):
    # Features [B,N,N,C] at Neighbor indices [B,N,K] => Neighbor features [B,N,K,C]
    neighbors = neighbor_idx.unsqueeze(-1).expand(-1, -1, -1, edges.size(-1))
    edge_features = torch.gather(edges, 2, neighbors)
    return edge_features

def gather_nodes(nodes, neighbor_idx):
    # Features [B,N,C] at Neighbor indices [B,N,K] => [B,N,K,C]
    # Flatten and expand indices per batch [B,N,K] => [B,NK] => [B,NK,C]
    neighbors_flat = neighbor_idx.view((neighbor_idx.shape[0], -1))
    neighbors_flat = neighbors_flat.unsqueeze(-1).expand(-1, -1, nodes.size(2))
    # Gather and re-pack
    neighbor_features = torch.gather(nodes, 1, neighbors_flat)
    neighbor_features = neighbor_features.view(list(neighbor_idx.shape)[:3] + [-1])
    return neighbor_features

def gather_nodes_t(nodes, neighbor_idx):
    # Features [B,N,C] at Neighbor index [B,K] => Neighbor features[B,K,C]
    idx_flat = neighbor_idx.unsqueeze(-1).expand(-1, -1, nodes.size(2))
    neighbor_features = torch.gather(nodes, 1, idx_flat)
    return neighbor_features

def cat_neighbors_nodes(h_nodes, h_neighbors, E_idx):
    h_nodes = gather_nodes(h_nodes, E_idx)
    h_nn = torch.cat([h_neighbors, h_nodes], -1)
    return h_nn


class EncLayer(nn.Module):
    def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30, time_cond_dim=None):
        super(EncLayer, self).__init__()
        self.num_hidden = num_hidden
        self.num_in = num_in
        self.scale = scale
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(num_hidden)
        self.norm2 = nn.LayerNorm(num_hidden)
        self.norm3 = nn.LayerNorm(num_hidden)

        if time_cond_dim is not None:
            self.time_block1 = nn.Sequential(
                Rearrange('b 1 d -> b 1 1 d'),
                nn.SiLU(),
                nn.Linear(time_cond_dim, num_hidden * 2))
            self.time_block2 = nn.Sequential(
                Rearrange('b 1 d -> b 1 1 d'),
                nn.SiLU(),
                nn.Linear(time_cond_dim, num_hidden * 2))

        self.W1 = nn.Linear(num_hidden + num_in, num_hidden, bias=True)
        self.W2 = nn.Linear(num_hidden, num_hidden, bias=True)
        self.W3 = nn.Linear(num_hidden, num_hidden, bias=True)
        self.W11 = nn.Linear(num_hidden + num_in, num_hidden, bias=True)
        self.W12 = nn.Linear(num_hidden, num_hidden, bias=True)
        self.W13 = nn.Linear(num_hidden, num_hidden, bias=True)
        self.act = torch.nn.GELU()
        self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4)

    def forward(self, h_V, h_E, E_idx, mask_V=None, mask_attend=None, time_cond=None):
        """ Parallel computation of full transformer layer """

        h_EV = cat_neighbors_nodes(h_V, h_E, E_idx)
        h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_EV.size(-2),-1)
        h_EV = torch.cat([h_V_expand, h_EV], -1)
        
        h_message = self.act(self.W2(self.act(self.W1(h_EV))))
        if time_cond is not None:
            scale, shift = self.time_block1(time_cond).chunk(2, dim=-1)
            h_message = h_message * (scale + 1) + shift
        h_message = self.W3(h_message)
        
        if mask_attend is not None:
            h_message = mask_attend.unsqueeze(-1) * h_message
        dh = torch.sum(h_message, -2) / self.scale
        h_V = self.norm1(h_V + self.dropout1(dh))

        dh = self.dense(h_V)
        h_V = self.norm2(h_V + self.dropout2(dh))
        if mask_V is not None:
            mask_V = mask_V.unsqueeze(-1)
            h_V = mask_V * h_V

        h_EV = cat_neighbors_nodes(h_V, h_E, E_idx)
        h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_EV.size(-2),-1)
        h_EV = torch.cat([h_V_expand, h_EV], -1)
        
        h_message = self.act(self.W12(self.act(self.W11(h_EV))))
        if time_cond is not None:
            scale, shift = self.time_block2(time_cond).chunk(2, dim=-1)
            h_message = h_message * (scale + 1) + shift
        h_message = self.W13(h_message)

        h_E = self.norm3(h_E + self.dropout3(h_message))
        return h_V, h_E


class DecLayer(nn.Module):
    def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30, time_cond_dim=None):
        super(DecLayer, self).__init__()
        self.num_hidden = num_hidden
        self.num_in = num_in
        self.scale = scale
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(num_hidden)
        self.norm2 = nn.LayerNorm(num_hidden)

        if time_cond_dim is not None:
            self.time_block = nn.Sequential(
                Rearrange('b 1 d -> b 1 1 d'),
                nn.SiLU(),
                nn.Linear(time_cond_dim, num_hidden * 2))

        self.W1 = nn.Linear(num_hidden + num_in, num_hidden, bias=True)
        self.W2 = nn.Linear(num_hidden, num_hidden, bias=True)
        self.W3 = nn.Linear(num_hidden, num_hidden, bias=True)
        self.act = torch.nn.GELU()
        self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4)

    def forward(self, h_V, h_E, mask_V=None, mask_attend=None, time_cond=None):
        """ Parallel computation of full transformer layer """

        # Concatenate h_V_i to h_E_ij
        h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_E.size(-2),-1)
        h_EV = torch.cat([h_V_expand, h_E], -1)

        h_message = self.act(self.W2(self.act(self.W1(h_EV))))
        if time_cond is not None:
            scale, shift = self.time_block(time_cond).chunk(2, dim=-1)
            h_message = h_message * (scale + 1) + shift
        h_message = self.W3(h_message)
 
        if mask_attend is not None:
            h_message = mask_attend.unsqueeze(-1) * h_message
        dh = torch.sum(h_message, -2) / self.scale

        h_V = self.norm1(h_V + self.dropout1(dh))

        # Position-wise feedforward
        dh = self.dense(h_V)
        h_V = self.norm2(h_V + self.dropout2(dh))

        if mask_V is not None:
            mask_V = mask_V.unsqueeze(-1)
            h_V = mask_V * h_V
        return h_V 



class PositionWiseFeedForward(nn.Module):
    def __init__(self, num_hidden, num_ff):
        super(PositionWiseFeedForward, self).__init__()
        self.W_in = nn.Linear(num_hidden, num_ff, bias=True)
        self.W_out = nn.Linear(num_ff, num_hidden, bias=True)
        self.act = torch.nn.GELU()
    def forward(self, h_V):
        h = self.act(self.W_in(h_V))
        h = self.W_out(h)
        return h

class PositionalEncodings(nn.Module):
    def __init__(self, num_embeddings, max_relative_feature=32):
        super(PositionalEncodings, self).__init__()
        self.num_embeddings = num_embeddings
        self.max_relative_feature = max_relative_feature
        self.linear = nn.Linear(2*max_relative_feature+1+1, num_embeddings)

    def forward(self, offset, mask):
        d = torch.clip(offset + self.max_relative_feature, 0, 2*self.max_relative_feature)*mask + (1-mask)*(2*self.max_relative_feature+1)
        d_onehot = torch.nn.functional.one_hot(d, 2*self.max_relative_feature+1+1)
        E = self.linear(d_onehot.float())
        return E



class CA_ProteinFeatures(nn.Module):
    def __init__(self, edge_features, node_features, num_positional_embeddings=16,
        num_rbf=16, top_k=30, augment_eps=0., num_chain_embeddings=16):
        """ Extract protein features """
        super(CA_ProteinFeatures, self).__init__()
        self.edge_features = edge_features
        self.node_features = node_features
        self.top_k = top_k
        self.augment_eps = augment_eps 
        self.num_rbf = num_rbf
        self.num_positional_embeddings = num_positional_embeddings

        # Positional encoding
        self.embeddings = PositionalEncodings(num_positional_embeddings)
        # Normalization and embedding
        node_in, edge_in = 3, num_positional_embeddings + num_rbf*9 + 7
        self.node_embedding = nn.Linear(node_in,  node_features, bias=False) #NOT USED
        self.edge_embedding = nn.Linear(edge_in, edge_features, bias=False)
        self.norm_nodes = nn.LayerNorm(node_features)
        self.norm_edges = nn.LayerNorm(edge_features)


    def _quaternions(self, R):
        """ Convert a batch of 3D rotations [R] to quaternions [Q]
            R [...,3,3]
            Q [...,4]
        """
        # Simple Wikipedia version
        # en.wikipedia.org/wiki/Rotation_matrix#Quaternion
        # For other options see math.stackexchange.com/questions/2074316/calculating-rotation-axis-from-rotation-matrix
        diag = torch.diagonal(R, dim1=-2, dim2=-1)
        Rxx, Ryy, Rzz = diag.unbind(-1)
        magnitudes = 0.5 * torch.sqrt(torch.abs(1 + torch.stack([
              Rxx - Ryy - Rzz, 
            - Rxx + Ryy - Rzz, 
            - Rxx - Ryy + Rzz
        ], -1)))
        _R = lambda i,j: R[:,:,:,i,j]
        signs = torch.sign(torch.stack([
            _R(2,1) - _R(1,2),
            _R(0,2) - _R(2,0),
            _R(1,0) - _R(0,1)
        ], -1))
        xyz = signs * magnitudes
        # The relu enforces a non-negative trace
        w = torch.sqrt(F.relu(1 + diag.sum(-1, keepdim=True))) / 2.
        Q = torch.cat((xyz, w), -1)
        Q = F.normalize(Q, dim=-1)
        return Q

    def _orientations_coarse(self, X, E_idx, eps=1e-6):
        dX = X[:,1:,:] - X[:,:-1,:]
        dX_norm = torch.norm(dX,dim=-1)
        dX_mask = (3.6<dX_norm) & (dX_norm<4.0) #exclude CA-CA jumps
        dX = dX*dX_mask[:,:,None]
        U = F.normalize(dX, dim=-1)
        u_2 = U[:,:-2,:]
        u_1 = U[:,1:-1,:]
        u_0 = U[:,2:,:]
        # Backbone normals
        n_2 = F.normalize(torch.cross(u_2, u_1), dim=-1)
        n_1 = F.normalize(torch.cross(u_1, u_0), dim=-1)

        # Bond angle calculation
        cosA = -(u_1 * u_0).sum(-1)
        cosA = torch.clamp(cosA, -1+eps, 1-eps)
        A = torch.acos(cosA)
        # Angle between normals
        cosD = (n_2 * n_1).sum(-1)
        cosD = torch.clamp(cosD, -1+eps, 1-eps)
        D = torch.sign((u_2 * n_1).sum(-1)) * torch.acos(cosD)
        # Backbone features
        AD_features = torch.stack((torch.cos(A), torch.sin(A) * torch.cos(D), torch.sin(A) * torch.sin(D)), 2)
        AD_features = F.pad(AD_features, (0,0,1,2), 'constant', 0)

        # Build relative orientations
        o_1 = F.normalize(u_2 - u_1, dim=-1)
        O = torch.stack((o_1, n_2, torch.cross(o_1, n_2)), 2)
        O = O.view(list(O.shape[:2]) + [9])
        O = F.pad(O, (0,0,1,2), 'constant', 0)
        O_neighbors = gather_nodes(O, E_idx)
        X_neighbors = gather_nodes(X, E_idx)
        
        # Re-view as rotation matrices
        O = O.view(list(O.shape[:2]) + [3,3])
        O_neighbors = O_neighbors.view(list(O_neighbors.shape[:3]) + [3,3])

        # Rotate into local reference frames
        dX = X_neighbors - X.unsqueeze(-2)
        dU = torch.matmul(O.unsqueeze(2), dX.unsqueeze(-1)).squeeze(-1)
        dU = F.normalize(dU, dim=-1)
        R = torch.matmul(O.unsqueeze(2).transpose(-1,-2), O_neighbors)
        Q = self._quaternions(R)

        # Orientation features
        O_features = torch.cat((dU,Q), dim=-1)
        return AD_features, O_features



    def _dist(self, X, mask, eps=1E-6):
        """ Pairwise euclidean distances """
        # Convolutional network on NCHW
        mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2)
        dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
        D = mask_2D * torch.sqrt(torch.sum(dX**2, 3) + eps)

        # Identify k nearest neighbors (including self)
        D_max, _ = torch.max(D, -1, keepdim=True)
        D_adjust = D + (1. - mask_2D) * D_max
        D_neighbors, E_idx = torch.topk(D_adjust, np.minimum(self.top_k, X.shape[1]), dim=-1, largest=False)
        mask_neighbors = gather_edges(mask_2D.unsqueeze(-1), E_idx)
        return D_neighbors, E_idx, mask_neighbors

    def _rbf(self, D):
        # Distance radial basis function
        device = D.device
        D_min, D_max, D_count = 2., 22., self.num_rbf
        D_mu = torch.linspace(D_min, D_max, D_count).to(device)
        D_mu = D_mu.view([1,1,1,-1])
        D_sigma = (D_max - D_min) / D_count
        D_expand = torch.unsqueeze(D, -1)
        RBF = torch.exp(-((D_expand - D_mu) / D_sigma)**2)
        return RBF

    def _get_rbf(self, A, B, E_idx):
        D_A_B = torch.sqrt(torch.sum((A[:,:,None,:] - B[:,None,:,:])**2,-1) + 1e-6) #[B, L, L]
        D_A_B_neighbors = gather_edges(D_A_B[:,:,:,None], E_idx)[:,:,:,0] #[B,L,K]
        RBF_A_B = self._rbf(D_A_B_neighbors)
        return RBF_A_B

    def forward(self, Ca, mask, residue_idx, chain_labels):
        """ Featurize coordinates as an attributed graph """
        if self.augment_eps > 0:
            Ca = Ca + self.augment_eps * torch.randn_like(Ca)

        D_neighbors, E_idx, mask_neighbors = self._dist(Ca, mask)

        Ca_0 = torch.zeros(Ca.shape, device=Ca.device)
        Ca_2 = torch.zeros(Ca.shape, device=Ca.device)
        Ca_0[:,1:,:] = Ca[:,:-1,:]
        Ca_1 = Ca
        Ca_2[:,:-1,:] = Ca[:,1:,:]

        V, O_features = self._orientations_coarse(Ca, E_idx)
        
        RBF_all = []
        RBF_all.append(self._rbf(D_neighbors)) #Ca_1-Ca_1
        RBF_all.append(self._get_rbf(Ca_0, Ca_0, E_idx)) 
        RBF_all.append(self._get_rbf(Ca_2, Ca_2, E_idx))

        RBF_all.append(self._get_rbf(Ca_0, Ca_1, E_idx))
        RBF_all.append(self._get_rbf(Ca_0, Ca_2, E_idx))

        RBF_all.append(self._get_rbf(Ca_1, Ca_0, E_idx))
        RBF_all.append(self._get_rbf(Ca_1, Ca_2, E_idx))

        RBF_all.append(self._get_rbf(Ca_2, Ca_0, E_idx))
        RBF_all.append(self._get_rbf(Ca_2, Ca_1, E_idx))


        RBF_all = torch.cat(tuple(RBF_all), dim=-1)


        offset = residue_idx[:,:,None]-residue_idx[:,None,:]
        offset = gather_edges(offset[:,:,:,None], E_idx)[:,:,:,0] #[B, L, K]

        d_chains = ((chain_labels[:, :, None] - chain_labels[:,None,:])==0).long()
        E_chains = gather_edges(d_chains[:,:,:,None], E_idx)[:,:,:,0]
        E_positional = self.embeddings(offset.long(), E_chains)
        E = torch.cat((E_positional, RBF_all, O_features), -1)
        

        E = self.edge_embedding(E)
        E = self.norm_edges(E)
        
        return E, E_idx 


def get_closest_neighbors(X, mask, top_k, eps=1e-6):
    # X is ca coords (b, n, 3), mask is seq mask
    mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2)
    dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
    D = mask_2D * torch.sqrt(torch.sum(dX**2, 3) + eps)
    D_max, _ = torch.max(D, -1, keepdim=True)
    D_adjust = D + (1. - mask_2D) * D_max
    sampled_top_k = top_k
    D_neighbors, E_idx = torch.topk(D_adjust, np.minimum(top_k, X.shape[1]), dim=-1, largest=False)
    return D_neighbors, E_idx


class ProteinFeatures(nn.Module):
    def __init__(self, edge_features, node_features, num_positional_embeddings=16,
        num_rbf=16, top_k=30, augment_eps=0., num_chain_embeddings=16):
        """ Extract protein features """
        super(ProteinFeatures, self).__init__()
        self.edge_features = edge_features
        self.node_features = node_features
        self.top_k = top_k
        self.augment_eps = augment_eps 
        self.num_rbf = num_rbf
        self.num_positional_embeddings = num_positional_embeddings

        self.embeddings = PositionalEncodings(num_positional_embeddings)
        node_in, edge_in = 6, num_positional_embeddings + num_rbf*25
        self.edge_embedding = nn.Linear(edge_in, edge_features, bias=False)
        self.norm_edges = nn.LayerNorm(edge_features)

    def _dist(self, X, mask, eps=1E-6):
        # mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2)
        # dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
        # D = mask_2D * torch.sqrt(torch.sum(dX**2, 3) + eps)
        # D_max, _ = torch.max(D, -1, keepdim=True)
        # D_adjust = D + (1. - mask_2D) * D_max
        # sampled_top_k = self.top_k
        # D_neighbors, E_idx = torch.topk(D_adjust, np.minimum(self.top_k, X.shape[1]), dim=-1, largest=False)
        # return D_neighbors, E_idx
        return get_closest_neighbors(X, mask, self.top_k, eps=eps)

    def _rbf(self, D):
        device = D.device
        D_min, D_max, D_count = 2., 22., self.num_rbf
        D_mu = torch.linspace(D_min, D_max, D_count, device=device)
        D_mu = D_mu.view([1,1,1,-1])
        D_sigma = (D_max - D_min) / D_count
        D_expand = torch.unsqueeze(D, -1)
        RBF = torch.exp(-((D_expand - D_mu) / D_sigma)**2)
        return RBF

    def _get_rbf(self, A, B, E_idx):
        D_A_B = torch.sqrt(torch.sum((A[:,:,None,:] - B[:,None,:,:])**2,-1) + 1e-6) #[B, L, L]
        D_A_B_neighbors = gather_edges(D_A_B[:,:,:,None], E_idx)[:,:,:,0] #[B,L,K]
        RBF_A_B = self._rbf(D_A_B_neighbors)
        return RBF_A_B

    def forward(self, X, mask, residue_idx, chain_labels):
        if self.augment_eps > 0:
            X = X + self.augment_eps * torch.randn_like(X)
        
        b = X[:,:,1,:] - X[:,:,0,:]
        c = X[:,:,2,:] - X[:,:,1,:]
        a = torch.cross(b, c, dim=-1)
        Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + X[:,:,1,:]
        Ca = X[:,:,1,:]
        N = X[:,:,0,:]
        C = X[:,:,2,:]
        O = X[:,:,3,:]
 
        D_neighbors, E_idx = self._dist(Ca, mask)

        RBF_all = []
        RBF_all.append(self._rbf(D_neighbors)) #Ca-Ca
        RBF_all.append(self._get_rbf(N, N, E_idx)) #N-N
        RBF_all.append(self._get_rbf(C, C, E_idx)) #C-C
        RBF_all.append(self._get_rbf(O, O, E_idx)) #O-O
        RBF_all.append(self._get_rbf(Cb, Cb, E_idx)) #Cb-Cb
        RBF_all.append(self._get_rbf(Ca, N, E_idx)) #Ca-N
        RBF_all.append(self._get_rbf(Ca, C, E_idx)) #Ca-C
        RBF_all.append(self._get_rbf(Ca, O, E_idx)) #Ca-O
        RBF_all.append(self._get_rbf(Ca, Cb, E_idx)) #Ca-Cb
        RBF_all.append(self._get_rbf(N, C, E_idx)) #N-C
        RBF_all.append(self._get_rbf(N, O, E_idx)) #N-O
        RBF_all.append(self._get_rbf(N, Cb, E_idx)) #N-Cb
        RBF_all.append(self._get_rbf(Cb, C, E_idx)) #Cb-C
        RBF_all.append(self._get_rbf(Cb, O, E_idx)) #Cb-O
        RBF_all.append(self._get_rbf(O, C, E_idx)) #O-C
        RBF_all.append(self._get_rbf(N, Ca, E_idx)) #N-Ca
        RBF_all.append(self._get_rbf(C, Ca, E_idx)) #C-Ca
        RBF_all.append(self._get_rbf(O, Ca, E_idx)) #O-Ca
        RBF_all.append(self._get_rbf(Cb, Ca, E_idx)) #Cb-Ca
        RBF_all.append(self._get_rbf(C, N, E_idx)) #C-N
        RBF_all.append(self._get_rbf(O, N, E_idx)) #O-N
        RBF_all.append(self._get_rbf(Cb, N, E_idx)) #Cb-N
        RBF_all.append(self._get_rbf(C, Cb, E_idx)) #C-Cb
        RBF_all.append(self._get_rbf(O, Cb, E_idx)) #O-Cb
        RBF_all.append(self._get_rbf(C, O, E_idx)) #C-O
        RBF_all = torch.cat(tuple(RBF_all), dim=-1)

        offset = residue_idx[:,:,None]-residue_idx[:,None,:]
        offset = gather_edges(offset[:,:,:,None], E_idx)[:,:,:,0] #[B, L, K]

        d_chains = ((chain_labels[:, :, None] - chain_labels[:,None,:])==0).long() #find self vs non-self interaction
        E_chains = gather_edges(d_chains[:,:,:,None], E_idx)[:,:,:,0]
        E_positional = self.embeddings(offset.long(), E_chains)
        E = torch.cat((E_positional, RBF_all), -1)
        E = self.edge_embedding(E)
        E = self.norm_edges(E)
        return E, E_idx 



class ProteinMPNN(nn.Module):
    def __init__(self, num_letters, node_features, edge_features,
        hidden_dim, num_encoder_layers=3, num_decoder_layers=3,
        vocab=21, k_neighbors=64, augment_eps=0.05, dropout=0.1, ca_only=False, time_cond_dim=None, input_S_is_embeddings=False):
        super(ProteinMPNN, self).__init__()

        # Hyperparameters
        self.node_features = node_features
        self.edge_features = edge_features
        self.hidden_dim = hidden_dim

        # Featurization layers
        if ca_only:
            self.features = CA_ProteinFeatures(node_features, edge_features, top_k=k_neighbors, augment_eps=augment_eps)
            self.W_v = nn.Linear(node_features, hidden_dim, bias=True)
        else:
            self.features = ProteinFeatures(node_features, edge_features, top_k=k_neighbors, augment_eps=augment_eps)

        self.W_e = nn.Linear(edge_features, hidden_dim, bias=True)
        self.input_S_is_embeddings = input_S_is_embeddings
        if not self.input_S_is_embeddings:
            self.W_s = nn.Embedding(vocab, hidden_dim)

        if time_cond_dim is not None:
            self.time_block = nn.Sequential(
                nn.SiLU(),
                nn.Linear(time_cond_dim, hidden_dim)
            )

        # Encoder layers
        self.encoder_layers = nn.ModuleList([
            EncLayer(hidden_dim, hidden_dim*2, dropout=dropout, time_cond_dim=time_cond_dim)
            for _ in range(num_encoder_layers)
        ])

        # Decoder layers
        self.decoder_layers = nn.ModuleList([
            DecLayer(hidden_dim, hidden_dim*3, dropout=dropout, time_cond_dim=time_cond_dim)
            for _ in range(num_decoder_layers)
        ])
        self.W_out = nn.Linear(hidden_dim, num_letters, bias=True)

        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, X, S, mask, chain_M, residue_idx, chain_encoding_all, randn, use_input_decoding_order=False, decoding_order=None, causal_mask=True, time_cond=None, return_node_embs=False):
        """ Graph-conditioned sequence model """
        device=X.device
        # Prepare node and edge embeddings
        E, E_idx = self.features(X, mask, residue_idx, chain_encoding_all)
        h_V = torch.zeros((E.shape[0], E.shape[1], E.shape[-1]), device=E.device)
        if time_cond is not None:
            time_cond_nodes = self.time_block(time_cond)
            h_V += time_cond_nodes  # time_cond is b, 1, c
        h_E = self.W_e(E)

        # Encoder is unmasked self-attention
        mask_attend = gather_nodes(mask.unsqueeze(-1),  E_idx).squeeze(-1)
        mask_attend = mask.unsqueeze(-1) * mask_attend
        for layer in self.encoder_layers:
            h_V, h_E = layer(h_V, h_E, E_idx, mask, mask_attend, time_cond=time_cond)

        encoder_embs = h_V

        # Concatenate sequence embeddings for autoregressive decoder
        if self.input_S_is_embeddings:
            h_S = S
        else:
            h_S = self.W_s(S)
        h_ES = cat_neighbors_nodes(h_S, h_E, E_idx)

        # Build encoder embeddings
        h_EX_encoder = cat_neighbors_nodes(torch.zeros_like(h_S), h_E, E_idx)
        h_EXV_encoder = cat_neighbors_nodes(h_V, h_EX_encoder, E_idx)


        chain_M = chain_M*mask #update chain_M to include missing regions
        mask_size = E_idx.shape[1]
        if causal_mask:
            if not use_input_decoding_order:
                decoding_order = torch.argsort((chain_M+0.0001)*(torch.abs(randn))) #[numbers will be smaller for places where chain_M = 0.0 and higher for places where chain_M = 1.0]
            permutation_matrix_reverse = torch.nn.functional.one_hot(decoding_order, num_classes=mask_size).float()
            order_mask_backward = torch.einsum('ij, biq, bjp->bqp',(1-torch.triu(torch.ones(mask_size,mask_size, device=device))), permutation_matrix_reverse, permutation_matrix_reverse)
        else:
            order_mask_backward = torch.ones(X.shape[0], mask_size, mask_size, device=device)
        mask_attend = torch.gather(order_mask_backward, 2, E_idx).unsqueeze(-1)
        mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
        mask_bw = mask_1D * mask_attend
        mask_fw = mask_1D * (1. - mask_attend)

        h_EXV_encoder_fw = mask_fw * h_EXV_encoder
        for layer in self.decoder_layers:
            # Masked positions attend to encoder information, unmasked see. 
            h_ESV = cat_neighbors_nodes(h_V, h_ES, E_idx)
            h_ESV = mask_bw * h_ESV + h_EXV_encoder_fw
            h_V = layer(h_V, h_ESV, mask, time_cond=time_cond)

        if return_node_embs:
            return h_V, encoder_embs
        else:
            logits = self.W_out(h_V)
            log_probs = F.log_softmax(logits, dim=-1)
            return log_probs


    def sample(self, X, randn, S_true, chain_mask, chain_encoding_all, residue_idx, mask=None, temperature=1.0, omit_AAs_np=None, bias_AAs_np=None, chain_M_pos=None, omit_AA_mask=None, pssm_coef=None, pssm_bias=None, pssm_multi=None, pssm_log_odds_flag=None, pssm_log_odds_mask=None, pssm_bias_flag=None, bias_by_res=None):
        device = X.device
        # Prepare node and edge embeddings
        E, E_idx = self.features(X, mask, residue_idx, chain_encoding_all)
        h_V = torch.zeros((E.shape[0], E.shape[1], E.shape[-1]), device=device)
        h_E = self.W_e(E)

        # Encoder is unmasked self-attention
        mask_attend = gather_nodes(mask.unsqueeze(-1),  E_idx).squeeze(-1)
        mask_attend = mask.unsqueeze(-1) * mask_attend
        for layer in self.encoder_layers:
            h_V, h_E = layer(h_V, h_E, E_idx, mask, mask_attend)

        # Decoder uses masked self-attention
        chain_mask = chain_mask*chain_M_pos*mask #update chain_M to include missing regions
        decoding_order = torch.argsort((chain_mask+0.0001)*(torch.abs(randn))) #[numbers will be smaller for places where chain_M = 0.0 and higher for places where chain_M = 1.0]
        mask_size = E_idx.shape[1]
        permutation_matrix_reverse = torch.nn.functional.one_hot(decoding_order, num_classes=mask_size).float()
        order_mask_backward = torch.einsum('ij, biq, bjp->bqp',(1-torch.triu(torch.ones(mask_size,mask_size, device=device))), permutation_matrix_reverse, permutation_matrix_reverse)
        mask_attend = torch.gather(order_mask_backward, 2, E_idx).unsqueeze(-1)
        mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
        mask_bw = mask_1D * mask_attend
        mask_fw = mask_1D * (1. - mask_attend)

        N_batch, N_nodes = X.size(0), X.size(1)
        log_probs = torch.zeros((N_batch, N_nodes, 21), device=device)
        all_probs = torch.zeros((N_batch, N_nodes, 21), device=device, dtype=torch.float32)
        h_S = torch.zeros_like(h_V, device=device)
        S = torch.zeros((N_batch, N_nodes), dtype=torch.int64, device=device)
        h_V_stack = [h_V] + [torch.zeros_like(h_V, device=device) for _ in range(len(self.decoder_layers))]
        constant = torch.tensor(omit_AAs_np, device=device)
        constant_bias = torch.tensor(bias_AAs_np, device=device)
        #chain_mask_combined = chain_mask*chain_M_pos 
        omit_AA_mask_flag = omit_AA_mask != None


        h_EX_encoder = cat_neighbors_nodes(torch.zeros_like(h_S), h_E, E_idx)
        h_EXV_encoder = cat_neighbors_nodes(h_V, h_EX_encoder, E_idx)
        h_EXV_encoder_fw = mask_fw * h_EXV_encoder
        for t_ in range(N_nodes):
            t = decoding_order[:,t_] #[B]
            chain_mask_gathered = torch.gather(chain_mask, 1, t[:,None]) #[B]
            mask_gathered = torch.gather(mask, 1, t[:,None]) #[B]
            bias_by_res_gathered = torch.gather(bias_by_res, 1, t[:,None,None].repeat(1,1,21))[:,0,:] #[B, 21]
            if (mask_gathered==0).all(): #for padded or missing regions only
                S_t = torch.gather(S_true, 1, t[:,None])
            else:
                # Hidden layers
                E_idx_t = torch.gather(E_idx, 1, t[:,None,None].repeat(1,1,E_idx.shape[-1]))
                h_E_t = torch.gather(h_E, 1, t[:,None,None,None].repeat(1,1,h_E.shape[-2], h_E.shape[-1]))
                h_ES_t = cat_neighbors_nodes(h_S, h_E_t, E_idx_t)
                h_EXV_encoder_t = torch.gather(h_EXV_encoder_fw, 1, t[:,None,None,None].repeat(1,1,h_EXV_encoder_fw.shape[-2], h_EXV_encoder_fw.shape[-1]))
                mask_t = torch.gather(mask, 1, t[:,None])
                for l, layer in enumerate(self.decoder_layers):
                    # Updated relational features for future states
                    h_ESV_decoder_t = cat_neighbors_nodes(h_V_stack[l], h_ES_t, E_idx_t)
                    h_V_t = torch.gather(h_V_stack[l], 1, t[:,None,None].repeat(1,1,h_V_stack[l].shape[-1]))
                    h_ESV_t = torch.gather(mask_bw, 1, t[:,None,None,None].repeat(1,1,mask_bw.shape[-2], mask_bw.shape[-1])) * h_ESV_decoder_t + h_EXV_encoder_t
                    h_V_stack[l+1].scatter_(1, t[:,None,None].repeat(1,1,h_V.shape[-1]), layer(h_V_t, h_ESV_t, mask_V=mask_t))
                # Sampling step
                h_V_t = torch.gather(h_V_stack[-1], 1, t[:,None,None].repeat(1,1,h_V_stack[-1].shape[-1]))[:,0]
                logits = self.W_out(h_V_t) / temperature
                probs = F.softmax(logits-constant[None,:]*1e8+constant_bias[None,:]/temperature+bias_by_res_gathered/temperature, dim=-1)
                if pssm_bias_flag:
                    pssm_coef_gathered = torch.gather(pssm_coef, 1, t[:,None])[:,0]
                    pssm_bias_gathered = torch.gather(pssm_bias, 1, t[:,None,None].repeat(1,1,pssm_bias.shape[-1]))[:,0]
                    probs = (1-pssm_multi*pssm_coef_gathered[:,None])*probs + pssm_multi*pssm_coef_gathered[:,None]*pssm_bias_gathered
                if pssm_log_odds_flag:
                    pssm_log_odds_mask_gathered = torch.gather(pssm_log_odds_mask, 1, t[:,None, None].repeat(1,1,pssm_log_odds_mask.shape[-1]))[:,0] #[B, 21]
                    probs_masked = probs*pssm_log_odds_mask_gathered
                    probs_masked += probs * 0.001
                    probs = probs_masked/torch.sum(probs_masked, dim=-1, keepdim=True) #[B, 21]
                if omit_AA_mask_flag:
                    omit_AA_mask_gathered = torch.gather(omit_AA_mask, 1, t[:,None, None].repeat(1,1,omit_AA_mask.shape[-1]))[:,0] #[B, 21]
                    probs_masked = probs*(1.0-omit_AA_mask_gathered)
                    probs = probs_masked/torch.sum(probs_masked, dim=-1, keepdim=True) #[B, 21]
                S_t = torch.multinomial(probs, 1)
                all_probs.scatter_(1, t[:,None,None].repeat(1,1,21), (chain_mask_gathered[:,:,None,]*probs[:,None,:]).float())
            S_true_gathered = torch.gather(S_true, 1, t[:,None])
            S_t = (S_t*chain_mask_gathered+S_true_gathered*(1.0-chain_mask_gathered)).long()
            temp1 = self.W_s(S_t)
            h_S.scatter_(1, t[:,None,None].repeat(1,1,temp1.shape[-1]), temp1)
            S.scatter_(1, t[:,None], S_t)
        output_dict = {"S": S, "probs": all_probs, "decoding_order": decoding_order}
        return output_dict


    def tied_sample(self, X, randn, S_true, chain_mask, chain_encoding_all, residue_idx, mask=None, temperature=1.0, omit_AAs_np=None, bias_AAs_np=None, chain_M_pos=None, omit_AA_mask=None, pssm_coef=None, pssm_bias=None, pssm_multi=None, pssm_log_odds_flag=None, pssm_log_odds_mask=None, pssm_bias_flag=None, tied_pos=None, tied_beta=None, bias_by_res=None):
        device = X.device
        # Prepare node and edge embeddings
        E, E_idx = self.features(X, mask, residue_idx, chain_encoding_all)
        h_V = torch.zeros((E.shape[0], E.shape[1], E.shape[-1]), device=device)
        h_E = self.W_e(E)
        # Encoder is unmasked self-attention
        mask_attend = gather_nodes(mask.unsqueeze(-1),  E_idx).squeeze(-1)
        mask_attend = mask.unsqueeze(-1) * mask_attend
        for layer in self.encoder_layers:
            h_V, h_E = layer(h_V, h_E, E_idx, mask, mask_attend)

        # Decoder uses masked self-attention
        chain_mask = chain_mask*chain_M_pos*mask #update chain_M to include missing regions
        decoding_order = torch.argsort((chain_mask+0.0001)*(torch.abs(randn))) #[numbers will be smaller for places where chain_M = 0.0 and higher for places where chain_M = 1.0]

        new_decoding_order = []
        for t_dec in list(decoding_order[0,].cpu().data.numpy()):
            if t_dec not in list(itertools.chain(*new_decoding_order)):
                list_a = [item for item in tied_pos if t_dec in item]
                if list_a:
                    new_decoding_order.append(list_a[0])
                else:
                    new_decoding_order.append([t_dec])
        decoding_order = torch.tensor(list(itertools.chain(*new_decoding_order)), device=device)[None,].repeat(X.shape[0],1)

        mask_size = E_idx.shape[1]
        permutation_matrix_reverse = torch.nn.functional.one_hot(decoding_order, num_classes=mask_size).float()
        order_mask_backward = torch.einsum('ij, biq, bjp->bqp',(1-torch.triu(torch.ones(mask_size,mask_size, device=device))), permutation_matrix_reverse, permutation_matrix_reverse)
        mask_attend = torch.gather(order_mask_backward, 2, E_idx).unsqueeze(-1)
        mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
        mask_bw = mask_1D * mask_attend
        mask_fw = mask_1D * (1. - mask_attend)

        N_batch, N_nodes = X.size(0), X.size(1)
        log_probs = torch.zeros((N_batch, N_nodes, 21), device=device)
        all_probs = torch.zeros((N_batch, N_nodes, 21), device=device, dtype=torch.float32)
        h_S = torch.zeros_like(h_V, device=device)
        S = torch.zeros((N_batch, N_nodes), dtype=torch.int64, device=device)
        h_V_stack = [h_V] + [torch.zeros_like(h_V, device=device) for _ in range(len(self.decoder_layers))]
        constant = torch.tensor(omit_AAs_np, device=device)
        constant_bias = torch.tensor(bias_AAs_np, device=device)
        omit_AA_mask_flag = omit_AA_mask != None

        h_EX_encoder = cat_neighbors_nodes(torch.zeros_like(h_S), h_E, E_idx)
        h_EXV_encoder = cat_neighbors_nodes(h_V, h_EX_encoder, E_idx)
        h_EXV_encoder_fw = mask_fw * h_EXV_encoder
        for t_list in new_decoding_order:
            logits = 0.0
            logit_list = []
            done_flag = False
            for t in t_list:
                if (mask[:,t]==0).all():
                    S_t = S_true[:,t]
                    for t in t_list:
                        h_S[:,t,:] = self.W_s(S_t)
                        S[:,t] = S_t
                    done_flag = True
                    break
                else:
                    E_idx_t = E_idx[:,t:t+1,:]
                    h_E_t = h_E[:,t:t+1,:,:]
                    h_ES_t = cat_neighbors_nodes(h_S, h_E_t, E_idx_t)
                    h_EXV_encoder_t = h_EXV_encoder_fw[:,t:t+1,:,:]
                    mask_t = mask[:,t:t+1]
                    for l, layer in enumerate(self.decoder_layers):
                        h_ESV_decoder_t = cat_neighbors_nodes(h_V_stack[l], h_ES_t, E_idx_t)
                        h_V_t = h_V_stack[l][:,t:t+1,:]
                        h_ESV_t = mask_bw[:,t:t+1,:,:] * h_ESV_decoder_t + h_EXV_encoder_t
                        h_V_stack[l+1][:,t,:] = layer(h_V_t, h_ESV_t, mask_V=mask_t).squeeze(1)
                    h_V_t = h_V_stack[-1][:,t,:]
                    logit_list.append((self.W_out(h_V_t) / temperature)/len(t_list))
                    logits += tied_beta[t]*(self.W_out(h_V_t) / temperature)/len(t_list)
            if done_flag:
                pass
            else:
                bias_by_res_gathered = bias_by_res[:,t,:] #[B, 21]
                probs = F.softmax(logits-constant[None,:]*1e8+constant_bias[None,:]/temperature+bias_by_res_gathered/temperature, dim=-1)
                if pssm_bias_flag:
                    pssm_coef_gathered = pssm_coef[:,t]
                    pssm_bias_gathered = pssm_bias[:,t]
                    probs = (1-pssm_multi*pssm_coef_gathered[:,None])*probs + pssm_multi*pssm_coef_gathered[:,None]*pssm_bias_gathered
                if pssm_log_odds_flag:
                    pssm_log_odds_mask_gathered = pssm_log_odds_mask[:,t]
                    probs_masked = probs*pssm_log_odds_mask_gathered
                    probs_masked += probs * 0.001
                    probs = probs_masked/torch.sum(probs_masked, dim=-1, keepdim=True) #[B, 21]
                if omit_AA_mask_flag:
                    omit_AA_mask_gathered = omit_AA_mask[:,t]
                    probs_masked = probs*(1.0-omit_AA_mask_gathered)
                    probs = probs_masked/torch.sum(probs_masked, dim=-1, keepdim=True) #[B, 21]
                S_t_repeat = torch.multinomial(probs, 1).squeeze(-1)
                S_t_repeat = (chain_mask[:,t]*S_t_repeat + (1-chain_mask[:,t])*S_true[:,t]).long() #hard pick fixed positions
                for t in t_list:
                    h_S[:,t,:] = self.W_s(S_t_repeat)
                    S[:,t] = S_t_repeat
                    all_probs[:,t,:] = probs.float()
        output_dict = {"S": S, "probs": all_probs, "decoding_order": decoding_order}
        return output_dict


    def conditional_probs(self, X, S, mask, chain_M, residue_idx, chain_encoding_all, randn, backbone_only=False):
        """ Graph-conditioned sequence model """
        device=X.device
        # Prepare node and edge embeddings
        E, E_idx = self.features(X, mask, residue_idx, chain_encoding_all)
        h_V_enc = torch.zeros((E.shape[0], E.shape[1], E.shape[-1]), device=E.device)
        h_E = self.W_e(E)

        # Encoder is unmasked self-attention
        mask_attend = gather_nodes(mask.unsqueeze(-1),  E_idx).squeeze(-1)
        mask_attend = mask.unsqueeze(-1) * mask_attend
        for layer in self.encoder_layers:
            h_V_enc, h_E = layer(h_V_enc, h_E, E_idx, mask, mask_attend)

        # Concatenate sequence embeddings for autoregressive decoder
        h_S = self.W_s(S)
        h_ES = cat_neighbors_nodes(h_S, h_E, E_idx)

        # Build encoder embeddings
        h_EX_encoder = cat_neighbors_nodes(torch.zeros_like(h_S), h_E, E_idx)
        h_EXV_encoder = cat_neighbors_nodes(h_V_enc, h_EX_encoder, E_idx)


        chain_M = chain_M*mask #update chain_M to include missing regions
  
        chain_M_np = chain_M.cpu().numpy()
        idx_to_loop = np.argwhere(chain_M_np[0,:]==1)[:,0]
        log_conditional_probs = torch.zeros([X.shape[0], chain_M.shape[1], 21], device=device).float()

        for idx in idx_to_loop:
            h_V = torch.clone(h_V_enc)
            order_mask = torch.zeros(chain_M.shape[1], device=device).float()
            if backbone_only:
                order_mask = torch.ones(chain_M.shape[1], device=device).float()
                order_mask[idx] = 0.
            else:
                order_mask = torch.zeros(chain_M.shape[1], device=device).float()
                order_mask[idx] = 1.
            decoding_order = torch.argsort((order_mask[None,]+0.0001)*(torch.abs(randn))) #[numbers will be smaller for places where chain_M = 0.0 and higher for places where chain_M = 1.0]
            mask_size = E_idx.shape[1]
            permutation_matrix_reverse = torch.nn.functional.one_hot(decoding_order, num_classes=mask_size).float()
            order_mask_backward = torch.einsum('ij, biq, bjp->bqp',(1-torch.triu(torch.ones(mask_size,mask_size, device=device))), permutation_matrix_reverse, permutation_matrix_reverse)
            mask_attend = torch.gather(order_mask_backward, 2, E_idx).unsqueeze(-1)
            mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
            mask_bw = mask_1D * mask_attend
            mask_fw = mask_1D * (1. - mask_attend)

            h_EXV_encoder_fw = mask_fw * h_EXV_encoder
            for layer in self.decoder_layers:
                # Masked positions attend to encoder information, unmasked see. 
                h_ESV = cat_neighbors_nodes(h_V, h_ES, E_idx)
                h_ESV = mask_bw * h_ESV + h_EXV_encoder_fw
                h_V = layer(h_V, h_ESV, mask)

            logits = self.W_out(h_V)
            log_probs = F.log_softmax(logits, dim=-1)
            log_conditional_probs[:,idx,:] = log_probs[:,idx,:]
        return log_conditional_probs


    def unconditional_probs(self, X, mask, residue_idx, chain_encoding_all):
        """ Graph-conditioned sequence model """
        device=X.device
        # Prepare node and edge embeddings
        E, E_idx = self.features(X, mask, residue_idx, chain_encoding_all)
        h_V = torch.zeros((E.shape[0], E.shape[1], E.shape[-1]), device=E.device)
        h_E = self.W_e(E)

        # Encoder is unmasked self-attention
        mask_attend = gather_nodes(mask.unsqueeze(-1),  E_idx).squeeze(-1)
        mask_attend = mask.unsqueeze(-1) * mask_attend
        for layer in self.encoder_layers:
            h_V, h_E = layer(h_V, h_E, E_idx, mask, mask_attend)

        # Build encoder embeddings
        h_EX_encoder = cat_neighbors_nodes(torch.zeros_like(h_V), h_E, E_idx)
        h_EXV_encoder = cat_neighbors_nodes(h_V, h_EX_encoder, E_idx)

        order_mask_backward = torch.zeros([X.shape[0], X.shape[1], X.shape[1]], device=device)
        mask_attend = torch.gather(order_mask_backward, 2, E_idx).unsqueeze(-1)
        mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
        mask_bw = mask_1D * mask_attend
        mask_fw = mask_1D * (1. - mask_attend)

        h_EXV_encoder_fw = mask_fw * h_EXV_encoder
        for layer in self.decoder_layers:
            h_V = layer(h_V, h_EXV_encoder_fw, mask)

        logits = self.W_out(h_V)
        log_probs = F.log_softmax(logits, dim=-1)
        return log_probs

core/protein_mpnn.py
ProteinDesignLab/protpardelle
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Constants used in AlphaFold.
Adapted from original code by alexechu.
"""

import collections
import functools
import os
from typing import List, Mapping, Tuple

import numpy as np
import tree

# Internal import (35fd).


# Distance from one CA to next CA [trans configuration: omega = 180].
ca_ca = 3.80209737096

# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
# chi angles so their chi angle lists are empty.
chi_angles_atoms = {
    "ALA": [],
    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
    "ARG": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "NE"],
        ["CG", "CD", "NE", "CZ"],
    ],
    "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
    "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
    "CYS": [["N", "CA", "CB", "SG"]],
    "GLN": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "OE1"],
    ],
    "GLU": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "OE1"],
    ],
    "GLY": [],
    "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]],
    "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]],
    "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "LYS": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "CE"],
        ["CG", "CD", "CE", "NZ"],
    ],
    "MET": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "SD"],
        ["CB", "CG", "SD", "CE"],
    ],
    "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]],
    "SER": [["N", "CA", "CB", "OG"]],
    "THR": [["N", "CA", "CB", "OG1"]],
    "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "VAL": [["N", "CA", "CB", "CG1"]],
}

# If chi angles given in fixed-length array, this matrix determines how to mask
# them for each AA type. The order is as per restype_order (see below).
chi_angles_mask = [
    [0.0, 0.0, 0.0, 0.0],  # ALA
    [1.0, 1.0, 1.0, 1.0],  # ARG
    [1.0, 1.0, 0.0, 0.0],  # ASN
    [1.0, 1.0, 0.0, 0.0],  # ASP
    [1.0, 0.0, 0.0, 0.0],  # CYS
    [1.0, 1.0, 1.0, 0.0],  # GLN
    [1.0, 1.0, 1.0, 0.0],  # GLU
    [0.0, 0.0, 0.0, 0.0],  # GLY
    [1.0, 1.0, 0.0, 0.0],  # HIS
    [1.0, 1.0, 0.0, 0.0],  # ILE
    [1.0, 1.0, 0.0, 0.0],  # LEU
    [1.0, 1.0, 1.0, 1.0],  # LYS
    [1.0, 1.0, 1.0, 0.0],  # MET
    [1.0, 1.0, 0.0, 0.0],  # PHE
    [1.0, 1.0, 0.0, 0.0],  # PRO
    [1.0, 0.0, 0.0, 0.0],  # SER
    [1.0, 0.0, 0.0, 0.0],  # THR
    [1.0, 1.0, 0.0, 0.0],  # TRP
    [1.0, 1.0, 0.0, 0.0],  # TYR
    [1.0, 0.0, 0.0, 0.0],  # VAL
]

# The following chi angles are pi periodic: they can be rotated by a multiple
# of pi without affecting the structure.
chi_pi_periodic = [
    [0.0, 0.0, 0.0, 0.0],  # ALA
    [0.0, 0.0, 0.0, 0.0],  # ARG
    [0.0, 0.0, 0.0, 0.0],  # ASN
    [0.0, 1.0, 0.0, 0.0],  # ASP
    [0.0, 0.0, 0.0, 0.0],  # CYS
    [0.0, 0.0, 0.0, 0.0],  # GLN
    [0.0, 0.0, 1.0, 0.0],  # GLU
    [0.0, 0.0, 0.0, 0.0],  # GLY
    [0.0, 0.0, 0.0, 0.0],  # HIS
    [0.0, 0.0, 0.0, 0.0],  # ILE
    [0.0, 0.0, 0.0, 0.0],  # LEU
    [0.0, 0.0, 0.0, 0.0],  # LYS
    [0.0, 0.0, 0.0, 0.0],  # MET
    [0.0, 1.0, 0.0, 0.0],  # PHE
    [0.0, 0.0, 0.0, 0.0],  # PRO
    [0.0, 0.0, 0.0, 0.0],  # SER
    [0.0, 0.0, 0.0, 0.0],  # THR
    [0.0, 0.0, 0.0, 0.0],  # TRP
    [0.0, 1.0, 0.0, 0.0],  # TYR
    [0.0, 0.0, 0.0, 0.0],  # VAL
    [0.0, 0.0, 0.0, 0.0],  # UNK
]

# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
# psi and chi angles:
# 0: 'backbone group',
# 1: 'pre-omega-group', (empty)
# 2: 'phi-group', (currently empty, because it defines only hydrogens)
# 3: 'psi-group',
# 4,5,6,7: 'chi1,2,3,4-group'
# The atom positions are relative to the axis-end-atom of the corresponding
# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
# is defined such that the dihedral-angle-defining atom (the last entry in
# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
# format: [atomname, group_idx, rel_position]
rigid_group_atom_positions = {
    "ALA": [
        ["N", 0, (-0.525, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, -0.000, -0.000)],
        ["CB", 0, (-0.529, -0.774, -1.205)],
        ["O", 3, (0.627, 1.062, 0.000)],
    ],
    "ARG": [
        ["N", 0, (-0.524, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, -0.000)],
        ["CB", 0, (-0.524, -0.778, -1.209)],
        ["O", 3, (0.626, 1.062, 0.000)],
        ["CG", 4, (0.616, 1.390, -0.000)],
        ["CD", 5, (0.564, 1.414, 0.000)],
        ["NE", 6, (0.539, 1.357, -0.000)],
        ["NH1", 7, (0.206, 2.301, 0.000)],
        ["NH2", 7, (2.078, 0.978, -0.000)],
        ["CZ", 7, (0.758, 1.093, -0.000)],
    ],
    "ASN": [
        ["N", 0, (-0.536, 1.357, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, -0.000, -0.000)],
        ["CB", 0, (-0.531, -0.787, -1.200)],
        ["O", 3, (0.625, 1.062, 0.000)],
        ["CG", 4, (0.584, 1.399, 0.000)],
        ["ND2", 5, (0.593, -1.188, 0.001)],
        ["OD1", 5, (0.633, 1.059, 0.000)],
    ],
    "ASP": [
        ["N", 0, (-0.525, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, 0.000, -0.000)],
        ["CB", 0, (-0.526, -0.778, -1.208)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.593, 1.398, -0.000)],
        ["OD1", 5, (0.610, 1.091, 0.000)],
        ["OD2", 5, (0.592, -1.101, -0.003)],
    ],
    "CYS": [
        ["N", 0, (-0.522, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.524, 0.000, 0.000)],
        ["CB", 0, (-0.519, -0.773, -1.212)],
        ["O", 3, (0.625, 1.062, -0.000)],
        ["SG", 4, (0.728, 1.653, 0.000)],
    ],
    "GLN": [
        ["N", 0, (-0.526, 1.361, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, 0.000, 0.000)],
        ["CB", 0, (-0.525, -0.779, -1.207)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.615, 1.393, 0.000)],
        ["CD", 5, (0.587, 1.399, -0.000)],
        ["NE2", 6, (0.593, -1.189, -0.001)],
        ["OE1", 6, (0.634, 1.060, 0.000)],
    ],
    "GLU": [
        ["N", 0, (-0.528, 1.361, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, -0.000, -0.000)],
        ["CB", 0, (-0.526, -0.781, -1.207)],
        ["O", 3, (0.626, 1.062, 0.000)],
        ["CG", 4, (0.615, 1.392, 0.000)],
        ["CD", 5, (0.600, 1.397, 0.000)],
        ["OE1", 6, (0.607, 1.095, -0.000)],
        ["OE2", 6, (0.589, -1.104, -0.001)],
    ],
    "GLY": [
        ["N", 0, (-0.572, 1.337, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.517, -0.000, -0.000)],
        ["O", 3, (0.626, 1.062, -0.000)],
    ],
    "HIS": [
        ["N", 0, (-0.527, 1.360, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, 0.000, 0.000)],
        ["CB", 0, (-0.525, -0.778, -1.208)],
        ["O", 3, (0.625, 1.063, 0.000)],
        ["CG", 4, (0.600, 1.370, -0.000)],
        ["CD2", 5, (0.889, -1.021, 0.003)],
        ["ND1", 5, (0.744, 1.160, -0.000)],
        ["CE1", 5, (2.030, 0.851, 0.002)],
        ["NE2", 5, (2.145, -0.466, 0.004)],
    ],
    "ILE": [
        ["N", 0, (-0.493, 1.373, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, -0.000, -0.000)],
        ["CB", 0, (-0.536, -0.793, -1.213)],
        ["O", 3, (0.627, 1.062, -0.000)],
        ["CG1", 4, (0.534, 1.437, -0.000)],
        ["CG2", 4, (0.540, -0.785, -1.199)],
        ["CD1", 5, (0.619, 1.391, 0.000)],
    ],
    "LEU": [
        ["N", 0, (-0.520, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, -0.000)],
        ["CB", 0, (-0.522, -0.773, -1.214)],
        ["O", 3, (0.625, 1.063, -0.000)],
        ["CG", 4, (0.678, 1.371, 0.000)],
        ["CD1", 5, (0.530, 1.430, -0.000)],
        ["CD2", 5, (0.535, -0.774, 1.200)],
    ],
    "LYS": [
        ["N", 0, (-0.526, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, 0.000, 0.000)],
        ["CB", 0, (-0.524, -0.778, -1.208)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.619, 1.390, 0.000)],
        ["CD", 5, (0.559, 1.417, 0.000)],
        ["CE", 6, (0.560, 1.416, 0.000)],
        ["NZ", 7, (0.554, 1.387, 0.000)],
    ],
    "MET": [
        ["N", 0, (-0.521, 1.364, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, 0.000, 0.000)],
        ["CB", 0, (-0.523, -0.776, -1.210)],
        ["O", 3, (0.625, 1.062, -0.000)],
        ["CG", 4, (0.613, 1.391, -0.000)],
        ["SD", 5, (0.703, 1.695, 0.000)],
        ["CE", 6, (0.320, 1.786, -0.000)],
    ],
    "PHE": [
        ["N", 0, (-0.518, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.524, 0.000, -0.000)],
        ["CB", 0, (-0.525, -0.776, -1.212)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.607, 1.377, 0.000)],
        ["CD1", 5, (0.709, 1.195, -0.000)],
        ["CD2", 5, (0.706, -1.196, 0.000)],
        ["CE1", 5, (2.102, 1.198, -0.000)],
        ["CE2", 5, (2.098, -1.201, -0.000)],
        ["CZ", 5, (2.794, -0.003, -0.001)],
    ],
    "PRO": [
        ["N", 0, (-0.566, 1.351, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, -0.000, 0.000)],
        ["CB", 0, (-0.546, -0.611, -1.293)],
        ["O", 3, (0.621, 1.066, 0.000)],
        ["CG", 4, (0.382, 1.445, 0.0)],
        # ['CD', 5, (0.427, 1.440, 0.0)],
        ["CD", 5, (0.477, 1.424, 0.0)],  # manually made angle 2 degrees larger
    ],
    "SER": [
        ["N", 0, (-0.529, 1.360, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, -0.000)],
        ["CB", 0, (-0.518, -0.777, -1.211)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["OG", 4, (0.503, 1.325, 0.000)],
    ],
    "THR": [
        ["N", 0, (-0.517, 1.364, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, 0.000, -0.000)],
        ["CB", 0, (-0.516, -0.793, -1.215)],
        ["O", 3, (0.626, 1.062, 0.000)],
        ["CG2", 4, (0.550, -0.718, -1.228)],
        ["OG1", 4, (0.472, 1.353, 0.000)],
    ],
    "TRP": [
        ["N", 0, (-0.521, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, 0.000)],
        ["CB", 0, (-0.523, -0.776, -1.212)],
        ["O", 3, (0.627, 1.062, 0.000)],
        ["CG", 4, (0.609, 1.370, -0.000)],
        ["CD1", 5, (0.824, 1.091, 0.000)],
        ["CD2", 5, (0.854, -1.148, -0.005)],
        ["CE2", 5, (2.186, -0.678, -0.007)],
        ["CE3", 5, (0.622, -2.530, -0.007)],
        ["NE1", 5, (2.140, 0.690, -0.004)],
        ["CH2", 5, (3.028, -2.890, -0.013)],
        ["CZ2", 5, (3.283, -1.543, -0.011)],
        ["CZ3", 5, (1.715, -3.389, -0.011)],
    ],
    "TYR": [
        ["N", 0, (-0.522, 1.362, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.524, -0.000, -0.000)],
        ["CB", 0, (-0.522, -0.776, -1.213)],
        ["O", 3, (0.627, 1.062, -0.000)],
        ["CG", 4, (0.607, 1.382, -0.000)],
        ["CD1", 5, (0.716, 1.195, -0.000)],
        ["CD2", 5, (0.713, -1.194, -0.001)],
        ["CE1", 5, (2.107, 1.200, -0.002)],
        ["CE2", 5, (2.104, -1.201, -0.003)],
        ["OH", 5, (4.168, -0.002, -0.005)],
        ["CZ", 5, (2.791, -0.001, -0.003)],
    ],
    "VAL": [
        ["N", 0, (-0.494, 1.373, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, -0.000, -0.000)],
        ["CB", 0, (-0.533, -0.795, -1.213)],
        ["O", 3, (0.627, 1.062, -0.000)],
        ["CG1", 4, (0.540, 1.429, -0.000)],
        ["CG2", 4, (0.533, -0.776, 1.203)],
    ],
}

# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
residue_atoms = {
    "ALA": ["C", "CA", "CB", "N", "O"],
    "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
    "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
    "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
    "CYS": ["C", "CA", "CB", "N", "O", "SG"],
    "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
    "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
    "GLY": ["C", "CA", "N", "O"],
    "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
    "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
    "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
    "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
    "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
    "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
    "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
    "SER": ["C", "CA", "CB", "N", "O", "OG"],
    "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
    "TRP": [
        "C",
        "CA",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "CE2",
        "CE3",
        "CZ2",
        "CZ3",
        "CH2",
        "N",
        "NE1",
        "O",
    ],
    "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"],
    "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
}

# Naming swaps for ambiguous atom names.
# Due to symmetries in the amino acids the naming of atoms is ambiguous in
# 4 of the 20 amino acids.
# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
# in LEU, VAL and ARG can be resolved by using the 3d constellations of
# the 'ambiguous' atoms and their neighbours)
residue_atom_renaming_swaps = {
    "ASP": {"OD1": "OD2"},
    "GLU": {"OE1": "OE2"},
    "PHE": {"CD1": "CD2", "CE1": "CE2"},
    "TYR": {"CD1": "CD2", "CE1": "CE2"},
}

# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
van_der_waals_radius = {
    "C": 1.7,
    "N": 1.55,
    "O": 1.52,
    "S": 1.8,
}

Bond = collections.namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"])
BondAngle = collections.namedtuple(
    "BondAngle", ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"]
)


@functools.lru_cache(maxsize=None)
def load_stereo_chemical_props() -> (
    Tuple[
        Mapping[str, List[Bond]],
        Mapping[str, List[Bond]],
        Mapping[str, List[BondAngle]],
    ]
):
    """Load stereo_chemical_props.txt into a nice structure.

    Load literature values for bond lengths and bond angles and translate
    bond angles into the length of the opposite edge of the triangle
    ("residue_virtual_bonds").

    Returns:
      residue_bonds: Dict that maps resname -> list of Bond tuples.
      residue_virtual_bonds: Dict that maps resname -> list of Bond tuples.
      residue_bond_angles: Dict that maps resname -> list of BondAngle tuples.
    """
    stereo_chemical_props_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "stereo_chemical_props.txt"
    )
    with open(stereo_chemical_props_path, "rt") as f:
        stereo_chemical_props = f.read()
    lines_iter = iter(stereo_chemical_props.splitlines())
    # Load bond lengths.
    residue_bonds = {}
    next(lines_iter)  # Skip header line.
    for line in lines_iter:
        if line.strip() == "-":
            break
        bond, resname, length, stddev = line.split()
        atom1, atom2 = bond.split("-")
        if resname not in residue_bonds:
            residue_bonds[resname] = []
        residue_bonds[resname].append(Bond(atom1, atom2, float(length), float(stddev)))
    residue_bonds["UNK"] = []

    # Load bond angles.
    residue_bond_angles = {}
    next(lines_iter)  # Skip empty line.
    next(lines_iter)  # Skip header line.
    for line in lines_iter:
        if line.strip() == "-":
            break
        bond, resname, angle_degree, stddev_degree = line.split()
        atom1, atom2, atom3 = bond.split("-")
        if resname not in residue_bond_angles:
            residue_bond_angles[resname] = []
        residue_bond_angles[resname].append(
            BondAngle(
                atom1,
                atom2,
                atom3,
                float(angle_degree) / 180.0 * np.pi,
                float(stddev_degree) / 180.0 * np.pi,
            )
        )
    residue_bond_angles["UNK"] = []

    def make_bond_key(atom1_name, atom2_name):
        """Unique key to lookup bonds."""
        return "-".join(sorted([atom1_name, atom2_name]))

    # Translate bond angles into distances ("virtual bonds").
    residue_virtual_bonds = {}
    for resname, bond_angles in residue_bond_angles.items():
        # Create a fast lookup dict for bond lengths.
        bond_cache = {}
        for b in residue_bonds[resname]:
            bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
        residue_virtual_bonds[resname] = []
        for ba in bond_angles:
            bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
            bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]

            # Compute distance between atom1 and atom3 using the law of cosines
            # c^2 = a^2 + b^2 - 2ab*cos(gamma).
            gamma = ba.angle_rad
            length = np.sqrt(
                bond1.length**2
                + bond2.length**2
                - 2 * bond1.length * bond2.length * np.cos(gamma)
            )

            # Propagation of uncertainty assuming uncorrelated errors.
            dl_outer = 0.5 / length
            dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
            dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
            dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
            stddev = np.sqrt(
                (dl_dgamma * ba.stddev) ** 2
                + (dl_db1 * bond1.stddev) ** 2
                + (dl_db2 * bond2.stddev) ** 2
            )
            residue_virtual_bonds[resname].append(
                Bond(ba.atom1_name, ba.atom3name, length, stddev)
            )

    return (residue_bonds, residue_virtual_bonds, residue_bond_angles)


# Between-residue bond lengths for general bonds (first element) and for Proline
# (second element).
between_res_bond_length_c_n = [1.329, 1.341]
between_res_bond_length_stddev_c_n = [0.014, 0.016]

# Between-residue cos_angles.
between_res_cos_angles_c_n_ca = [-0.5203, 0.0353]  # degrees: 121.352 +- 2.315
between_res_cos_angles_ca_c_n = [-0.4473, 0.0311]  # degrees: 116.568 +- 1.995

# This mapping is used when we need to store atom data in a format that requires
# fixed atom data size for every residue (e.g. a numpy array).
atom_types = [
    "N",
    "CA",
    "C",
    "CB",
    "O",
    "CG",
    "CG1",
    "CG2",
    "OG",
    "OG1",
    "SG",
    "CD",
    "CD1",
    "CD2",
    "ND1",
    "ND2",
    "OD1",
    "OD2",
    "SD",
    "CE",
    "CE1",
    "CE2",
    "CE3",
    "NE",
    "NE1",
    "NE2",
    "OE1",
    "OE2",
    "CH2",
    "NH1",
    "NH2",
    "OH",
    "CZ",
    "CZ2",
    "CZ3",
    "NZ",
    "OXT",
]
atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
atom_type_num = len(atom_types)  # := 37.

# A compact atom encoding with 14 columns
# pylint: disable=line-too-long
# pylint: disable=bad-whitespace
restype_name_to_atom14_names = {
    "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
    "ARG": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD",
        "NE",
        "CZ",
        "NH1",
        "NH2",
        "",
        "",
        "",
    ],
    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""],
    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""],
    "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""],
    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""],
    "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
    "HIS": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "ND1",
        "CD2",
        "CE1",
        "NE2",
        "",
        "",
        "",
        "",
    ],
    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""],
    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""],
    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""],
    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""],
    "PHE": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "CE1",
        "CE2",
        "CZ",
        "",
        "",
        "",
    ],
    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
    "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""],
    "TRP": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "NE1",
        "CE2",
        "CE3",
        "CZ2",
        "CZ3",
        "CH2",
    ],
    "TYR": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "CE1",
        "CE2",
        "CZ",
        "OH",
        "",
        "",
    ],
    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""],
    "UNK": ["", "", "", "", "", "", "", "", "", "", "", "", "", ""],
}
# pylint: enable=line-too-long
# pylint: enable=bad-whitespace


# This is the standard residue order when coding AA type as a number.
# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
restypes = [
    "A",
    "R",
    "N",
    "D",
    "C",
    "Q",
    "E",
    "G",
    "H",
    "I",
    "L",
    "K",
    "M",
    "F",
    "P",
    "S",
    "T",
    "W",
    "Y",
    "V",
]
restype_order = {restype: i for i, restype in enumerate(restypes)}
restype_num = len(restypes)  # := 20.
unk_restype_index = restype_num  # Catch-all index for unknown restypes.

restypes_with_x = restypes + ["X"]
restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}


def sequence_to_onehot(
    sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False
) -> np.ndarray:
    """Maps the given sequence into a one-hot encoded matrix.

    Args:
      sequence: An amino acid sequence.
      mapping: A dictionary mapping amino acids to integers.
      map_unknown_to_x: If True, any amino acid that is not in the mapping will be
        mapped to the unknown amino acid 'X'. If the mapping doesn't contain
        amino acid 'X', an error will be thrown. If False, any amino acid not in
        the mapping will throw an error.

    Returns:
      A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
      the sequence.

    Raises:
      ValueError: If the mapping doesn't contain values from 0 to
        num_unique_aas - 1 without any gaps.
    """
    num_entries = max(mapping.values()) + 1

    if sorted(set(mapping.values())) != list(range(num_entries)):
        raise ValueError(
            "The mapping must have values from 0 to num_unique_aas-1 "
            "without any gaps. Got: %s" % sorted(mapping.values())
        )

    one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)

    for aa_index, aa_type in enumerate(sequence):
        if map_unknown_to_x:
            if aa_type.isalpha() and aa_type.isupper():
                aa_id = mapping.get(aa_type, mapping["X"])
            else:
                raise ValueError(f"Invalid character in the sequence: {aa_type}")
        else:
            aa_id = mapping[aa_type]
        one_hot_arr[aa_index, aa_id] = 1

    return one_hot_arr


restype_1to3 = {
    "A": "ALA",
    "R": "ARG",
    "N": "ASN",
    "D": "ASP",
    "C": "CYS",
    "Q": "GLN",
    "E": "GLU",
    "G": "GLY",
    "H": "HIS",
    "I": "ILE",
    "L": "LEU",
    "K": "LYS",
    "M": "MET",
    "F": "PHE",
    "P": "PRO",
    "S": "SER",
    "T": "THR",
    "W": "TRP",
    "Y": "TYR",
    "V": "VAL",
}


# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
# many more, and less common, three letter names as keys and maps many of these
# to the same one letter name (including 'X' and 'U' which we don't use here).
restype_3to1 = {v: k for k, v in restype_1to3.items()}

# Define a restype name for all unknown residues.
unk_restype = "UNK"

resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
resname_to_idx = {resname: i for i, resname in enumerate(resnames)}


# Define exploded all-atom representation (atom73)
atom73_names = ['N', 'CA', 'C', 'CB', 'O']
for aa1 in restypes:
    aa3 = restype_1to3[aa1]
    atom_list = residue_atoms[aa3]
    for atom in atom_types:
        if atom in atom_list and atom not in atom73_names:
            atom73_names.append(f'{aa1}{atom}')

atom73_names_to_idx = {a: i for i, a in enumerate(atom73_names)}

restype_atom73_mask = np.zeros((22, 73))
for i, restype in enumerate(restypes):
    for atom_name in atom_types:
        atom73_name = atom_name
        if atom_name not in ['N', 'CA', 'C', 'CB', 'O']:
            atom73_name = restype + atom_name
        if atom73_name in atom73_names_to_idx:
            atom73_idx = atom73_names_to_idx[atom73_name]
            restype_atom73_mask[i, atom73_idx] = 1
# Remove CB for glycine
restype_atom73_mask[restype_order["G"], 3] = 0
# Backbone atoms for unk and mask
restype_atom73_mask[-2:, [0, 1, 2, 4]] = 1


# The mapping here uses hhblits convention, so that B is mapped to D, J and O
# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
# remaining 20 amino acids are kept in alphabetical order.
# There are 2 non-amino acid codes, X (representing any amino acid) and
# "-" representing a missing amino acid in an alignment.  The id for these
# codes is put at the end (20 and 21) so that they can easily be ignored if
# desired.
HHBLITS_AA_TO_ID = {
    "A": 0,
    "B": 2,
    "C": 1,
    "D": 2,
    "E": 3,
    "F": 4,
    "G": 5,
    "H": 6,
    "I": 7,
    "J": 20,
    "K": 8,
    "L": 9,
    "M": 10,
    "N": 11,
    "O": 20,
    "P": 12,
    "Q": 13,
    "R": 14,
    "S": 15,
    "T": 16,
    "U": 1,
    "V": 17,
    "W": 18,
    "X": 20,
    "Y": 19,
    "Z": 3,
    "-": 21,
}

# Partial inversion of HHBLITS_AA_TO_ID.
ID_TO_HHBLITS_AA = {
    0: "A",
    1: "C",  # Also U.
    2: "D",  # Also B.
    3: "E",  # Also Z.
    4: "F",
    5: "G",
    6: "H",
    7: "I",
    8: "K",
    9: "L",
    10: "M",
    11: "N",
    12: "P",
    13: "Q",
    14: "R",
    15: "S",
    16: "T",
    17: "V",
    18: "W",
    19: "Y",
    20: "X",  # Includes J and O.
    21: "-",
}

restypes_with_x_and_gap = restypes + ["X", "-"]
MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
    for i in range(len(restypes_with_x_and_gap))
)


def _make_standard_atom_mask() -> np.ndarray:
    """Returns [num_res_types, num_atom_types] mask array."""
    # +1 to account for unknown (all 0s).
    mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
    for restype, restype_letter in enumerate(restypes):
        restype_name = restype_1to3[restype_letter]
        atom_names = residue_atoms[restype_name]
        for atom_name in atom_names:
            atom_type = atom_order[atom_name]
            mask[restype, atom_type] = 1
    return mask


STANDARD_ATOM_MASK = _make_standard_atom_mask()


# A one hot representation for the first and second atoms defining the axis
# of rotation for each chi-angle in each residue.
def chi_angle_atom(atom_index: int) -> np.ndarray:
    """Define chi-angle rigid groups via one-hot representations."""
    chi_angles_index = {}
    one_hots = []

    for k, v in chi_angles_atoms.items():
        indices = [atom_types.index(s[atom_index]) for s in v]
        indices.extend([-1] * (4 - len(indices)))
        chi_angles_index[k] = indices

    for r in restypes:
        res3 = restype_1to3[r]
        one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
        one_hots.append(one_hot)

    one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
    one_hot = np.stack(one_hots, axis=0)
    one_hot = np.transpose(one_hot, [0, 2, 1])

    return one_hot


chi_atom_1_one_hot = chi_angle_atom(1)
chi_atom_2_one_hot = chi_angle_atom(2)

# An array like chi_angles_atoms but using indices rather than names.
chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
chi_angles_atom_indices = tree.map_structure(
    lambda atom_name: atom_order[atom_name], chi_angles_atom_indices
)
chi_angles_atom_indices = np.array(
    [
        chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms)))
        for chi_atoms in chi_angles_atom_indices
    ]
)

# Mapping from (res_name, atom_name) pairs to the atom's chi group index
# and atom index within that group.
chi_groups_for_atom = collections.defaultdict(list)
for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
    for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
        for atom_i, atom in enumerate(chi_group):
            chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
chi_groups_for_atom = dict(chi_groups_for_atom)


def _make_rigid_transformation_4x4(ex, ey, translation):
    """Create a rigid 4x4 transformation matrix from two axes and transl."""
    # Normalize ex.
    ex_normalized = ex / np.linalg.norm(ex)

    # make ey perpendicular to ex
    ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
    ey_normalized /= np.linalg.norm(ey_normalized)

    # compute ez as cross product
    eznorm = np.cross(ex_normalized, ey_normalized)
    m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
    m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0)
    return m


# create an array with (restype, atomtype) --> rigid_group_idx
# and an array with (restype, atomtype, coord) for the atom positions
# and compute affine transformation matrices (4,4) from one rigid group to the
# previous group
restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int)
restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int)
restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)


def _make_rigid_group_constants():
    """Fill the arrays above."""
    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
            atomtype = atom_order[atomname]
            restype_atom37_to_rigid_group[restype, atomtype] = group_idx
            restype_atom37_mask[restype, atomtype] = 1
            restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position

            atom14idx = restype_name_to_atom14_names[resname].index(atomname)
            restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
            restype_atom14_mask[restype, atom14idx] = 1
            restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position

    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        atom_positions = {
            name: np.array(pos) for name, _, pos in rigid_group_atom_positions[resname]
        }

        # backbone to backbone is the identity transform
        restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)

        # pre-omega-frame to backbone (currently dummy identity matrix)
        restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)

        # phi-frame to backbone
        mat = _make_rigid_transformation_4x4(
            ex=atom_positions["N"] - atom_positions["CA"],
            ey=np.array([1.0, 0.0, 0.0]),
            translation=atom_positions["N"],
        )
        restype_rigid_group_default_frame[restype, 2, :, :] = mat

        # psi-frame to backbone
        mat = _make_rigid_transformation_4x4(
            ex=atom_positions["C"] - atom_positions["CA"],
            ey=atom_positions["CA"] - atom_positions["N"],
            translation=atom_positions["C"],
        )
        restype_rigid_group_default_frame[restype, 3, :, :] = mat

        # chi1-frame to backbone
        if chi_angles_mask[restype][0]:
            base_atom_names = chi_angles_atoms[resname][0]
            base_atom_positions = [atom_positions[name] for name in base_atom_names]
            mat = _make_rigid_transformation_4x4(
                ex=base_atom_positions[2] - base_atom_positions[1],
                ey=base_atom_positions[0] - base_atom_positions[1],
                translation=base_atom_positions[2],
            )
            restype_rigid_group_default_frame[restype, 4, :, :] = mat

        # chi2-frame to chi1-frame
        # chi3-frame to chi2-frame
        # chi4-frame to chi3-frame
        # luckily all rotation axes for the next frame start at (0,0,0) of the
        # previous frame
        for chi_idx in range(1, 4):
            if chi_angles_mask[restype][chi_idx]:
                axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
                axis_end_atom_position = atom_positions[axis_end_atom_name]
                mat = _make_rigid_transformation_4x4(
                    ex=axis_end_atom_position,
                    ey=np.array([-1.0, 0.0, 0.0]),
                    translation=axis_end_atom_position,
                )
                restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat


_make_rigid_group_constants()


def make_atom14_dists_bounds(overlap_tolerance=1.5, bond_length_tolerance_factor=15):
    """compute upper and lower bounds for bonds to assess violations."""
    restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
    restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
    restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
    residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        atom_list = restype_name_to_atom14_names[resname]

        # create lower and upper bounds for clashes
        for atom1_idx, atom1_name in enumerate(atom_list):
            if not atom1_name:
                continue
            atom1_radius = van_der_waals_radius[atom1_name[0]]
            for atom2_idx, atom2_name in enumerate(atom_list):
                if (not atom2_name) or atom1_idx == atom2_idx:
                    continue
                atom2_radius = van_der_waals_radius[atom2_name[0]]
                lower = atom1_radius + atom2_radius - overlap_tolerance
                upper = 1e10
                restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
                restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
                restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
                restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper

        # overwrite lower and upper bounds for bonds and angles
        for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
            atom1_idx = atom_list.index(b.atom1_name)
            atom2_idx = atom_list.index(b.atom2_name)
            lower = b.length - bond_length_tolerance_factor * b.stddev
            upper = b.length + bond_length_tolerance_factor * b.stddev
            restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
            restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
            restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
            restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
            restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
            restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
    return {
        "lower_bound": restype_atom14_bond_lower_bound,  # shape (21,14,14)
        "upper_bound": restype_atom14_bond_upper_bound,  # shape (21,14,14)
        "stddev": restype_atom14_bond_stddev,  # shape (21,14,14)
    }


standard_residue_bonds, _, standard_residue_bond_angles = load_stereo_chemical_props()

core/residue_constants.py
ProteinDesignLab/protpardelle
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Various utils for handling protein data.
"""

import os
import shlex
import subprocess
import sys
import torch
import yaml
import argparse

from einops import rearrange, repeat
import numpy as np
import torch
import torch.nn.functional as F
import Bio
from Bio.PDB.DSSP import DSSP

from core import protein
from core import protein_mpnn
from core import residue_constants


PATH_TO_TMALIGN = "/home/alexechu/essentials_kit/ml_utils/align/TMalign/TMalign"


################ STRUCTURE/FORMAT UTILS #############################


def aatype_to_seq(aatype, seq_mask=None):
    if seq_mask is None:
        seq_mask = torch.ones_like(aatype)

    mapping = residue_constants.restypes_with_x
    mapping = mapping + ["<mask>"]

    unbatched = False
    if len(aatype.shape) == 1:
        unbatched = True
        aatype = [aatype]
        seq_mask = [seq_mask]

    seqs = []
    for i, ai in enumerate(aatype):
        seq = []
        for j, aa in enumerate(ai):
            if seq_mask[i][j] == 1:
                try:
                    seq.append(mapping[aa])
                except IndexError:
                    print(aatype[i])
                    raise Exception(f"Error in mapping {aa} at {i},{j}")
        seqs.append("".join(seq))

    if unbatched:
        seqs = seqs[0]
    return seqs


def seq_to_aatype(seq, num_tokens=21):
    if num_tokens == 20:
        mapping = residue_constants.restype_order
    if num_tokens == 21:
        mapping = residue_constants.restype_order_with_x
    if num_tokens == 22:
        mapping = residue_constants.restype_order_with_x
        mapping["<mask>"] = 21
    return torch.Tensor([mapping[aa] for aa in seq]).long()


def batched_seq_to_aatype_and_mask(seqs, max_len=None):
    if max_len is None:
        max_len = max([len(s) for s in seqs])
    aatypes = []
    seq_mask = []
    for s in seqs:
        pad_size = max_len - len(s)
        aatype = seq_to_aatype(s)
        aatypes.append(F.pad(aatype, (0, pad_size)))
        mask = torch.ones_like(aatype).float()
        seq_mask.append(F.pad(mask, (0, pad_size)))
    return torch.stack(aatypes), torch.stack(seq_mask)


def atom37_mask_from_aatype(aatype, seq_mask=None):
    # source_mask is (21,37) originally
    source_mask = torch.Tensor(residue_constants.restype_atom37_mask).to(aatype.device)
    bb_atoms = source_mask[residue_constants.restype_order["G"]][None]
    # Use only the first 20 plus bb atoms for X, mask
    source_mask = torch.cat([source_mask[:-1], bb_atoms, bb_atoms], 0)
    atom_mask = source_mask[aatype]
    if seq_mask is not None:
        atom_mask *= seq_mask[..., None]
    return atom_mask


def atom37_coords_from_atom14(atom14_coords, aatype, return_mask=False):
    # Unbatched
    device = atom14_coords.device
    atom37_coords = torch.zeros((atom14_coords.shape[0], 37, 3)).to(device)
    for i in range(atom14_coords.shape[0]):  # per residue
        aa = aatype[i]
        aa_3name = residue_constants.restype_1to3[residue_constants.restypes[aa]]
        atom14_atoms = residue_constants.restype_name_to_atom14_names[aa_3name]
        for j in range(14):
            atom_name = atom14_atoms[j]
            if atom_name != "":
                atom37_idx = residue_constants.atom_order[atom_name]
                atom37_coords[i, atom37_idx, :] = atom14_coords[i, j, :]

    if return_mask:
        atom37_mask = atom37_mask_from_aatype(aatype)
        return atom37_coords, atom37_mask
    return atom37_coords


def atom73_mask_from_aatype(aatype, seq_mask=None):
    source_mask = torch.Tensor(residue_constants.restype_atom73_mask).to(aatype.device)
    atom_mask = source_mask[aatype]
    if seq_mask is not None:
        atom_mask *= seq_mask[..., None]
    return atom_mask


def atom37_to_atom73(atom37, aatype, return_mask=False):
    # Unbatched
    atom73 = torch.zeros((atom37.shape[0], 73, 3)).to(atom37)
    for i in range(atom37.shape[0]):
        aa = aatype[i]
        aa1 = residue_constants.restypes[aa]
        for j, atom37_name in enumerate(residue_constants.atom_types):
            atom73_name = atom37_name
            if atom37_name not in ["N", "CA", "C", "O", "CB"]:
                atom73_name = aa1 + atom73_name
            if atom73_name in residue_constants.atom73_names_to_idx:
                atom73_idx = residue_constants.atom73_names_to_idx[atom73_name]
                atom73[i, atom73_idx, :] = atom37[i, j, :]

    if return_mask:
        atom73_mask = atom73_mask_from_aatype(aatype)
        return atom73, atom73_mask
    return atom73


def atom73_to_atom37(atom73, aatype, return_mask=False):
    # Unbatched
    atom37_coords = torch.zeros((atom73.shape[0], 37, 3)).to(atom73)
    for i in range(atom73.shape[0]):  # per residue
        aa = aatype[i]
        aa1 = residue_constants.restypes[aa]
        for j, atom_type in enumerate(residue_constants.atom_types):
            atom73_name = atom_type
            if atom73_name not in ["N", "CA", "C", "O", "CB"]:
                atom73_name = aa1 + atom73_name
            if atom73_name in residue_constants.atom73_names_to_idx:
                atom73_idx = residue_constants.atom73_names_to_idx[atom73_name]
                atom37_coords[i, j, :] = atom73[i, atom73_idx, :]

    if return_mask:
        atom37_mask = atom37_mask_from_aatype(aatype)
        return atom37_coords, atom37_mask
    return atom37_coords


def get_dmap(pdb, atoms=["N", "CA", "C", "O"], batched=True, out="torch", device=None):
    def _dmap_from_coords(coords):
        coords = coords.contiguous()
        dmaps = torch.cdist(coords, coords).unsqueeze(1)
        if out == "numpy":
            return dmaps.detach().cpu().numpy()
        elif out == "torch":
            if device is not None:
                return dmaps.to(device)
            else:
                return dmaps

    if isinstance(pdb, str):  # input is pdb file
        coords = load_coords_from_pdb(pdb, atoms=atoms).view(1, -1, 3)
        return _dmap_from_coords(coords)
    elif len(pdb.shape) == 2:  # single set of coords
        if isinstance(pdb, np.ndarray):
            pdb = torch.Tensor(pdb)
        return _dmap_from_coords(pdb.unsqueeze(0))
    elif len(pdb.shape) == 3 and batched:
        return _dmap_from_coords(pdb)
    elif len(pdb.shape) == 3 and not batched:
        return _dmap_from_coords(pdb.view(1, -1, 3))
    elif len(pdb.shape) == 4:
        return _dmap_from_coords(pdb.view(pdb.size(0), -1, 3))


def get_channeled_dmap(coords):
    # coords is b, nres, natom, 3
    coords = coords.permute(0, 2, 1, 3)
    dvecs = coords[..., None, :] - coords[..., None, :, :]  # b, natom, nres, nres, 3
    dists = torch.sqrt(dvecs.pow(2).sum(-1) + 1e-8)
    return dists


def fill_in_cbeta_for_atom37(coords):
    b = coords[..., 1, :] - coords[..., 0, :]
    c = coords[..., 2, :] - coords[..., 1, :]
    a = torch.cross(b, c, dim=-1)
    cbeta = -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + coords[..., 1, :]
    new_coords = torch.clone(coords)
    new_coords[..., 3, :] = cbeta
    return new_coords


def get_distogram(coords, n_bins=20, start=2, return_onehot=True, seq_mask=None):
    # coords is b, nres, natom, 3
    # distogram for cb atom (assume 3rd atom)
    coords_with_cb = fill_in_cbeta_for_atom37(coords)
    dists = get_channeled_dmap(coords_with_cb[:, :, 3:4]).squeeze(1)
    bins = torch.arange(start, start + n_bins - 1).to(dists.device)
    dgram = torch.bucketize(dists, bins)
    dgram_oh = F.one_hot(dgram, n_bins)
    if seq_mask is not None:
        mask_2d = seq_mask[:, :, None] * seq_mask[:, None, :]
        dgram = dgram * mask_2d
        dgram_oh = dgram_oh * mask_2d[..., None]

    if return_onehot:
        return dgram_oh
    return dgram


def get_contacts(coords=None, distogram=None, seq_mask=None):
    if distogram is None:
        distogram = get_distogram(coords)
    contacts = (distogram.argmax(-1) < 6).float()
    if seq_mask is not None:
        contacts *= seq_mask[..., None] * seq_mask[..., None, :]
    return contacts


def dihedral(a, b, c, d):
    # inputs can be (1,3), (n,3), or (bs,n,3)
    b1 = a - b
    b2 = b - c
    b3 = c - d
    n1 = F.normalize(torch.cross(b1, b2), dim=-1)
    n2 = F.normalize(torch.cross(b2, b3), dim=-1)
    m1 = torch.cross(n1, b2 / b2.norm(dim=-1).unsqueeze(-1))
    y = (m1 * n2).sum(dim=-1)
    x = (n1 * n2).sum(dim=-1)
    return torch.atan2(y, x)


def get_torsions_from_coords(
    coords, atoms=["N", "CA", "C", "O"], batched=True, out="torch", device=None
):
    """
    Returns a n-dim array of shape (bs, nres, ntors), where ntors is the
    number of torsion angles (e.g. 2 if using phi and psi), with units of radians.
    """
    if isinstance(coords, np.ndarray):
        coords = torch.Tensor(coords)
    if len(coords.shape) == 2:
        coords = coords.unsqueeze(0)
    if len(coords.shape) == 4:
        coords = coords.view(coords.size(0), -1, 3)
    if len(coords.shape) == 3 and not batched:
        coords = coords.view(1, -1, 3)
    if len(coords.shape) == 3:
        bs = coords.size(0)
        if "O" in atoms:
            idxs = [
                i for i in range(coords.size(1)) if i % 4 != 3
            ]  # deselect O atoms for N-Ca-C-O coords
            coords = coords[:, idxs, :]
        a, b, c, d = (
            coords[:, :-3, :],
            coords[:, 1:-2, :],
            coords[:, 2:-1, :],
            coords[:, 3:, :],
        )
        torsions = dihedral(
            a, b, c, d
        )  # output order is psi-omega-phi, reorganize to (bs, nres, 3)
        torsions = torsions.view(bs, torsions.size(1) // 3, 3)
        omegaphi = torch.cat(
            (torch.zeros(bs, 1, 2).to(coords.device), torsions[:, :, 1:]), 1
        )
        psi = torch.cat((torsions[:, :, 0], torch.zeros(bs, 1).to(coords.device)), 1)
        torsions = torch.cat(
            (
                omegaphi[:, :, 1].unsqueeze(-1),
                psi.unsqueeze(-1),
                omegaphi[:, :, 0].unsqueeze(-1),
            ),
            -1,
        )
    else:
        raise Exception("input coords not of correct dims")

    if out == "numpy":
        return torsions.detach().cpu().numpy()
    elif out == "torch":
        if device is not None:
            return torsions.to(device)
        else:
            return torsions


def get_trig_from_torsions(torsions, out="torch", device=None):
    """
    Calculate unit circle projections from coords input.

    Returns a n-dim array of shape (bs, nres, ntors, 2), where ntors is the
    number of torsion angles (e.g. 2 if using phi and psi), and the last
    dimension is the xy unit-circle coordinates of the corresponding angle.
    """
    if isinstance(torsions, np.ndarray):
        torsions = torch.Tensor(torsions)
    x = torsions.cos()
    y = torsions.sin()
    trig = torch.cat((x.unsqueeze(-1), y.unsqueeze(-1)), -1)
    if out == "numpy":
        return trig.detach().cpu().numpy()
    elif out == "torch":
        if device is not None:
            return trig.to(device)
        else:
            return trig


def get_abego_string_from_torsions(torsions):
    A_bin = (-75, 50)
    G_bin = (-100, 100)
    torsions = torsions * 180.0 / np.pi
    phi, psi = torsions[:, :, 0], torsions[:, :, 1]
    abego_vec = np.zeros((torsions.size(0), torsions.size(1))).astype(str)
    A = (phi <= 0) & (psi <= A_bin[1]) & (psi > A_bin[0])
    B = (phi <= 0) & ((psi > A_bin[1]) | (psi <= A_bin[0]))
    G = (phi > 0) & (psi <= G_bin[1]) & (psi > G_bin[0])
    E = (phi > 0) & ((psi > G_bin[1]) | (psi <= G_bin[0]))
    abego_vec[A] = "A"
    abego_vec[B] = "B"
    abego_vec[G] = "G"
    abego_vec[E] = "E"
    abego_strs = ["".join(v) for v in abego_vec]
    return abego_strs


def get_bond_lengths_from_coords(coords, batched=True, out="torch", device=None):
    """
    Returns array of shape (bs, n_res, 4), where final dim is bond lengths
    in order of N-Ca, Ca-C, C-O, C-N (none for last residue)
    """
    if isinstance(coords, np.ndarray):
        coords = torch.Tensor(coords)
    if len(coords.shape) == 2:
        coords = coords.unsqueeze(0)
    if len(coords.shape) == 3 and not batched:
        coords = coords.view(1, -1, 3)
    if len(coords.shape) == 4:
        coords = coords.view(coords.size(0), -1, 3)
    N = coords[:, ::4, :]
    Ca = coords[:, 1::4, :]
    C = coords[:, 2::4, :]
    O = coords[:, 3::4, :]
    NCa = (Ca - N).norm(dim=-1).unsqueeze(-1)
    CaC = (C - Ca).norm(dim=-1).unsqueeze(-1)
    CO = (O - C).norm(dim=-1).unsqueeze(-1)
    CN = (N[:, 1:] - C[:, :-1]).norm(dim=-1)
    CN = torch.cat([CN, torch.zeros(CN.size(0), 1).to(CN.device)], 1).unsqueeze(-1)
    blengths = torch.cat((NCa, CaC, CO, CN), -1)
    if out == "numpy":
        return blengths.detach().cpu().numpy()
    elif out == "torch":
        if device is not None:
            return blengths.to(device)
        else:
            return blengths


def get_bond_angles_from_coords(coords, batched=True, out="torch", device=None):
    """
    Returns array of shape (bs, n_res, 5), where final dim is bond angles
    in order of N-Ca-C, Ca-C-O, Ca-C-N, O-C-N, C-N-Ca (none for last residue)
    """

    def _angle(v1, v2):
        cos = (v1 * v2).sum(-1) / (v1.norm(dim=-1) * v2.norm(dim=-1))
        return cos.acos()

    if isinstance(coords, np.ndarray):
        coords = torch.Tensor(coords)
    if len(coords.shape) == 2:
        coords = coords.unsqueeze(0)
    if len(coords.shape) == 3 and not batched:
        coords = coords.view(1, -1, 3)
    if len(coords.shape) == 4:
        coords = coords.view(coords.size(0), -1, 3)
    N = coords[:, ::4, :]
    Nnext = coords[:, 4::4, :]
    Ca = coords[:, 1::4, :]
    Canext = coords[:, 5::4, :]
    C = coords[:, 2::4, :]
    O = coords[:, 3::4, :]
    CaN = N - Ca
    CaC = C - Ca
    CCa = Ca - C
    CO = O - C
    CNnext = Nnext - C[:, :-1, :]
    NnextC = -1 * CNnext
    NnextCanext = Canext - Nnext
    NCaC = _angle(CaN, CaC).unsqueeze(-1)
    CaCO = _angle(CCa, CO).unsqueeze(-1)
    CaCN = _angle(CCa[:, :-1], CNnext).unsqueeze(-1)
    CaCN = _extend(CaCN)
    OCN = _angle(CO[:, :-1], CNnext).unsqueeze(-1)
    OCN = _extend(OCN)
    CNCa = _angle(NnextC, NnextCanext).unsqueeze(-1)
    # CNCa = torch.cat([CNCa, torch.zeros(CNCa.size(0), 1).to(CNCa.device)], 1).unsqueeze(-1)
    CNCa = _extend(CNCa)
    bangles = torch.cat((NCaC, CaCO, CaCN, OCN, CNCa), -1)
    if out == "numpy":
        return bangles.detach().cpu().numpy()
    elif out == "torch":
        if device is not None:
            return bangles.to(device)
        else:
            return bangles


def get_buried_positions_mask(coords, seq_mask=None, threshold=6.0):
    ca_idx = residue_constants.atom_order["CA"]  # typically 1
    cb_idx = residue_constants.atom_order["CB"]  # typically 3
    if seq_mask is None:
        seq_mask = torch.ones_like(coords)[..., 0, 0]
    coords = fill_in_cbeta_for_atom37(coords)

    # get 8 closest neighbors by CB
    neighbor_coords = coords[:, :, cb_idx]

    ca_neighbor_dists, edge_index = protein_mpnn.get_closest_neighbors(
        neighbor_coords, seq_mask, 9
    )
    edge_index = edge_index[..., 1:].contiguous()

    # compute avg CB distance
    cb_coords = coords[:, :, cb_idx]
    neighbor_cb = protein_mpnn.gather_nodes(cb_coords, edge_index)
    avg_cb_dist = (neighbor_cb - cb_coords[..., None, :]).pow(2).sum(-1).sqrt().mean(-1)

    buried_positions_mask = (avg_cb_dist < threshold).float() * seq_mask
    return buried_positions_mask


def get_fullatom_bond_lengths_from_coords(
    coords, aatype, atom_mask=None, return_format="per_aa"
):
    # Also return sidechain bond angles. All unbatched. return list of dicts
    def dist(xyz1, xyz2):
        return (xyz1 - xyz2).pow(2).sum().sqrt().detach().cpu().item()

    assert aatype.max() <= 19
    seq = aatype_to_seq(aatype)
    # residue-wise list of dicts [{'N-CA': a, 'CA-C': b}, {'N-CA': a, 'CA-C': b}]
    all_bond_lens_by_pos = []
    # aa-wise dict of dicts of lists {'A': {'N-CA': [a, b, c], 'CA-C': [a, b, c]}}
    all_bond_lens_by_aa = {aa: {} for aa in residue_constants.restypes}
    for i, res in enumerate(coords):
        aa3 = residue_constants.restype_1to3[seq[i]]
        res_bond_lens = {}
        for j, atom1 in enumerate(residue_constants.atom_types):
            for k, atom2 in enumerate(residue_constants.atom_types):
                if j < k and protein.are_atoms_bonded(aa3, atom1, atom2):
                    if atom_mask is None or (
                        atom_mask[i, j] > 0.5 and atom_mask[i, k] > 0.5
                    ):
                        bond_name = f"{atom1}-{atom2}"
                        bond_len = dist(res[j], res[k])
                        res_bond_lens[bond_name] = bond_len
        all_bond_lens_by_pos.append(res_bond_lens)
        for key, val in res_bond_lens.items():
            all_bond_lens_by_aa[seq[i]].setdefault(key, []).append(val)

    if return_format == "per_aa":
        return all_bond_lens_by_aa
    elif return_format == "per_position":
        return all_bond_lens_by_pos


def batched_fullatom_bond_lengths_from_coords(
    coords, aatype, atom_mask=None, return_format="per_aa"
):
    # Expects trimmed coords (no mask)
    if return_format == "per_position":
        batched_bond_lens = []
    elif return_format == "per_aa":
        batched_bond_lens = {aa: {} for aa in residue_constants.restypes}
    for i, c in enumerate(coords):
        atom_mask_i = None if atom_mask is None else atom_mask[i]
        bond_lens = get_fullatom_bond_lengths_from_coords(
            c, aatype[i], atom_mask=atom_mask_i, return_format=return_format
        )
        if return_format == "per_position":
            batched_bond_lens.extend(bond_lens)
        elif return_format == "per_aa":
            for aa, d in bond_lens.items():
                for bond, lengths in d.items():
                    batched_bond_lens[aa].setdefault(bond, []).extend(lengths)
    return batched_bond_lens


def batched_fullatom_bond_angles_from_coords(coords, aatype, return_format="per_aa"):
    # Expects trimmed coords (no mask)
    if return_format == "per_position":
        batched_bond_angles = []
    elif return_format == "per_aa":
        batched_bond_angles = {aa: {} for aa in residue_constants.restypes}
    for i, c in enumerate(coords):
        bond_angles = get_fullatom_bond_angles_from_coords(
            c, aatype[i], return_format=return_format
        )
        if return_format == "per_position":
            batched_bond_angles.extend(bond_angles)
        elif return_format == "per_aa":
            for aa, d in bond_angles.items():
                for bond, lengths in d.items():
                    batched_bond_angles[aa].setdefault(bond, []).extend(lengths)
    return batched_bond_angles


def get_chi_angles(coords, aatype, atom_mask=None, seq_mask=None):
    # unbatched
    # return (n, 4) chis in degrees and mask
    chis = []
    chi_mask = []
    atom_order = residue_constants.atom_order

    seq = aatype_to_seq(aatype, seq_mask=seq_mask)

    for i, aa1 in enumerate(seq):  # per residue
        if seq_mask is not None and seq_mask[i] == 0:
            chis.append([0, 0, 0, 0])
            chi_mask.append([0, 0, 0, 0])
        else:
            chi = []
            mask = []
            chi_atoms = residue_constants.chi_angles_atoms[
                residue_constants.restype_1to3[aa1]
            ]
            for j in range(4):  # per chi angle
                if j > len(chi_atoms) - 1:
                    chi.append(0)
                    mask.append(0)
                elif atom_mask is not None and any(
                    [atom_mask[i, atom_order[a]] < 0.5 for a in chi_atoms[j]]
                ):
                    chi.append(0)
                    mask.append(0)
                else:
                    # Four atoms per dihedral
                    xyz4 = [coords[i, atom_order[a]] for a in chi_atoms[j]]
                    angle = dihedral(*xyz4) * 180 / np.pi
                    chi.append(angle)
                    mask.append(1)
            chis.append(chi)
            chi_mask.append(mask)

    chis = torch.Tensor(chis)
    chi_mask = torch.Tensor(chi_mask)

    return chis, chi_mask


def fill_Os_from_NCaC_coords(
    coords: torch.Tensor, out: str = "torch", device: str = None
):
    """Given NCaC coords, add O atom coordinates in.
    (bs, 3n, 3) -> (bs, 4n, 3)
    """
    CO_LEN = 1.231
    if len(coords.shape) == 2:
        coords = coords.unsqueeze(0)
    Cs = coords[:, 2:-1:3, :]  # all but last C
    CCa_norm = F.normalize(coords[:, 1:-2:3, :] - Cs, dim=-1)  # all but last Ca
    CN_norm = F.normalize(coords[:, 3::3, :] - Cs, dim=-1)  # all but first N
    Os = F.normalize(CCa_norm + CN_norm, dim=-1) * -CO_LEN
    Os += Cs
    # TODO place C-term O atom properly
    Os = torch.cat([Os, coords[:, -1, :].view(-1, 1, 3) + 1], 1)
    coords_out = []
    for i in range(Os.size(1)):
        coords_out.append(coords[:, i * 3 : (i + 1) * 3, :])
        coords_out.append(Os[:, i, :].view(-1, 1, 3))
    coords_out = torch.cat(coords_out, 1)
    if out == "numpy":
        return coords_out.detach().cpu().numpy()
    elif out == "torch":
        if device is not None:
            return coords_out.to(device)
        else:
            return coords_out


def _extend(x, axis=1, n=1, prepend=False):
    # Add an extra zeros 'residue' to the end (or beginning, prepend=True) of a Tensor
    # Used to extend torsions when there is no 'psi' for last residue
    shape = list(x.shape)
    shape[axis] = n
    if prepend:
        return torch.cat([torch.zeros(shape).to(x.device), x], axis)
    else:
        return torch.cat([x, torch.zeros(shape).to(x.device)], axis)


def trim_coords(coords, n_res, batched=True):
    if batched:  # Return list of tensors
        front = (coords.shape[1] - n_res) // 2
        return [
            coords[i, front[i] : front[i] + n_res[i]] for i in range(coords.shape[0])
        ]
    else:
        if isinstance(n_res, torch.Tensor):
            n_res = n_res.int()
        front_pad = (coords.shape[0] - n_res) // 2
        return coords[front_pad : front_pad + n_res]


def batch_align_on_calpha(x, y):
    aligned_x = []
    for i, xi in enumerate(x):
        xi_calpha = xi[:, 1, :]
        _, (R, t) = kabsch_align(xi_calpha, y[i, :, 1, :])
        xi_ctr = xi - xi_calpha.mean(0, keepdim=True)
        xi_aligned = xi_ctr @ R.t() + t
        aligned_x.append(xi_aligned)
    return torch.stack(aligned_x)


def kabsch_align(p, q):
    if len(p.shape) > 2:
        p = p.reshape(-1, 3)
    if len(q.shape) > 2:
        q = q.reshape(-1, 3)
    p_ctr = p - p.mean(0, keepdim=True)
    t = q.mean(0, keepdim=True)
    q_ctr = q - t
    H = p_ctr.t() @ q_ctr
    U, S, V = torch.svd(H)
    R = V @ U.t()
    I_ = torch.eye(3).to(p)
    I_[-1, -1] = R.det().sign()
    R = V @ I_ @ U.t()
    p_aligned = p_ctr @ R.t() + t
    return p_aligned, (R, t)


def get_dssp_string(pdb):
    try:
        structure = Bio.PDB.PDBParser(QUIET=True).get_structure(pdb[:-3], pdb)
        dssp = DSSP(structure[0], pdb, dssp="mkdssp")
        dssp_string = "".join([dssp[k][2] for k in dssp.keys()])
        return dssp_string
    except Exception as e:
        print(e)
        return None


def pool_dssp_symbols(dssp_string, newchar=None, chars=["-", "T", "S", "C", " "]):
    """Replaces all instances of chars with newchar. DSSP chars are helix=GHI, strand=EB, loop=- TSC"""
    if newchar is None:
        newchar = chars[0]
    string_out = dssp_string
    for c in chars:
        string_out = string_out.replace(c, newchar)
    return string_out


def get_3state_dssp(pdb=None, coords=None):
    if coords is not None:
        pdb = "temp_dssp.pdb"
        write_coords_to_pdb(coords, pdb, batched=False)
    dssp_string = get_dssp_string(pdb)
    if dssp_string is not None:
        dssp_string = pool_dssp_symbols(dssp_string, newchar="L")
        dssp_string = pool_dssp_symbols(dssp_string, chars=["H", "G", "I"])
        dssp_string = pool_dssp_symbols(dssp_string, chars=["E", "B"])
    if coords is not None:
        subprocess.run(shlex.split(f"rm {pdb}"))
    return dssp_string


############## SAVE/LOAD UTILS #################################


def load_feats_from_pdb(
    pdb, bb_atoms=["N", "CA", "C", "O"], load_atom73=False, **kwargs
):
    feats = {}
    with open(pdb, "r") as f:
        pdb_str = f.read()
    protein_obj = protein.from_pdb_string(pdb_str, **kwargs)
    bb_idxs = [residue_constants.atom_order[a] for a in bb_atoms]
    bb_coords = torch.from_numpy(protein_obj.atom_positions[:, bb_idxs])
    feats["bb_coords"] = bb_coords.float()
    for k, v in vars(protein_obj).items():
        feats[k] = torch.Tensor(v)
    feats["aatype"] = feats["aatype"].long()
    if load_atom73:
        feats["atom73_coords"], feats["atom73_mask"] = atom37_to_atom73(
            feats["atom_positions"], feats["aatype"], return_mask=True
        )
    return feats


def load_coords_from_pdb(
    pdb,
    atoms=["N", "CA", "C", "O"],
    method="raw",
    also_bfactors=False,
    normalize_bfactors=True,
):
    """Returns array of shape (1, n_res, len(atoms), 3)"""
    coords = []
    bfactors = []
    if method == "raw":  # Raw numpy implementation, faster than biopdb
        # Indexing into PDB format, allowing XXXX.XXX
        coords_in_pdb = [slice(30, 38), slice(38, 46), slice(46, 54)]
        # Indexing into PDB format, allowing XXX.XX
        bfactor_in_pdb = slice(60, 66)

        with open(pdb, "r") as f:
            resi_prev = 1
            counter = 0
            for l in f:
                l_split = l.rstrip("\n").split()
                if len(l_split) > 0 and l_split[0] == "ATOM" and l_split[2] in atoms:
                    resi = l_split[5]
                    if resi == resi_prev:
                        counter += 1
                    else:
                        counter = 0
                    if counter < len(atoms):
                        xyz = [
                            np.array(l[s].strip()).astype(float) for s in coords_in_pdb
                        ]
                        coords.append(xyz)
                        if also_bfactors:
                            bfactor = np.array(l[bfactor_in_pdb].strip()).astype(float)
                            bfactors.append(bfactor)
                    resi_prev = resi
            coords = torch.Tensor(np.array(coords)).view(1, -1, len(atoms), 3)
            if also_bfactors:
                bfactors = torch.Tensor(np.array(bfactors)).view(1, -1, len(atoms))
    elif method == "biopdb":
        structure = Bio.PDB.PDBParser(QUIET=True).get_structure(pdb[:-3], pdb)
        for model in structure:
            for chain in model:
                for res in chain:
                    for atom in atoms:
                        try:
                            coords.append(np.asarray(res[atom].get_coord()))
                            if also_bfactors:
                                bfactors.append(np.asarray(res[atom].get_bfactor()))
                        except:
                            continue
    else:
        raise NotImplementedError(f"Invalid method for reading coords: {method}")
    if also_bfactors:
        if normalize_bfactors:  # Normalize over Calphas
            mean_b = bfactors[..., 1].mean()
            std_b = bfactors[..., 1].var().sqrt()
            bfactors = (bfactors - mean_b) / (std_b + 1e-6)
        return coords, bfactors
    return coords


def feats_to_pdb_str(
    atom_positions,
    aatype=None,
    atom_mask=None,
    residue_index=None,
    chain_index=None,
    b_factors=None,
    atom_lines_only=True,
    conect=False,
    **kwargs,
):
    # Expects unbatched, cropped inputs. needs at least one of atom_mask, aatype
    # Uses all-GLY aatype if aatype not given: does not infer from atom_mask
    assert aatype is not None or atom_mask is not None
    if atom_mask is None:
        aatype = aatype.cpu()
        atom_mask = atom37_mask_from_aatype(aatype, torch.ones_like(aatype))
    if aatype is None:
        seq_mask = atom_mask[:, residue_constants.atom_order["CA"]].cpu()
        aatype = seq_mask * residue_constants.restype_order["G"]
    if residue_index is None:
        residue_index = torch.arange(aatype.shape[-1])
    if chain_index is None:
        chain_index = torch.ones_like(aatype)
    if b_factors is None:
        b_factors = torch.ones_like(atom_mask)

    cast = lambda x: np.array(x.detach().cpu()) if isinstance(x, torch.Tensor) else x
    prot = protein.Protein(
        atom_positions=cast(atom_positions),
        atom_mask=cast(atom_mask),
        aatype=cast(aatype),
        residue_index=cast(residue_index),
        chain_index=cast(chain_index),
        b_factors=cast(b_factors),
    )
    pdb_str = protein.to_pdb(prot, conect=conect)
    if conect:
        pdb_str, conect_str = pdb_str
    if atom_lines_only:
        pdb_lines = pdb_str.split("\n")
        atom_lines = [
            l for l in pdb_lines if len(l.split()) > 1 and l.split()[0] == "ATOM"
        ]
        pdb_str = "\n".join(atom_lines) + "\n"
    if conect:
        pdb_str = pdb_str + conect_str
    return pdb_str


def bb_coords_to_pdb_str(coords, atoms=["N", "CA", "C", "O"]):
    def _bb_pdb_line(atom, atomnum, resnum, coords, elem, res="GLY"):
        atm = "ATOM".ljust(6)
        atomnum = str(atomnum).rjust(5)
        atomname = atom.center(4)
        resname = res.ljust(3)
        chain = "A".rjust(1)
        resnum = str(resnum).rjust(4)
        x = str("%8.3f" % (float(coords[0]))).rjust(8)
        y = str("%8.3f" % (float(coords[1]))).rjust(8)
        z = str("%8.3f" % (float(coords[2]))).rjust(8)
        occ = str("%6.2f" % (float(1))).rjust(6)
        temp = str("%6.2f" % (float(20))).ljust(6)
        elname = elem.rjust(12)
        return "%s%s %s %s %s%s    %s%s%s%s%s%s\n" % (
            atm,
            atomnum,
            atomname,
            resname,
            chain,
            resnum,
            x,
            y,
            z,
            occ,
            temp,
            elname,
        )

    n = coords.shape[0]
    na = len(atoms)
    pdb_str = ""
    for j in range(0, n, na):
        for idx, atom in enumerate(atoms):
            pdb_str += _bb_pdb_line(
                atom,
                j + idx + 1,
                (j + na) // na,
                coords[j + idx],
                atom[0],
            )
    return pdb_str


def write_coords_to_pdb(
    coords_in,
    filename,
    batched=True,
    write_to_frames=False,
    conect=False,
    **all_atom_feats,
):
    def _write_pdb_string(pdb_str, filename, append=False):
        write_mode = "a" if append else "w"
        with open(filename, write_mode) as f:
            if write_to_frames:
                f.write("MODEL\n")
            f.write(pdb_str)
            if write_to_frames:
                f.write("ENDMDL\n")

    if not (batched or write_to_frames):
        coords_in = [coords_in]
        filename = [filename]
        all_atom_feats = {k: [v] for k, v in all_atom_feats.items()}

    n_atoms_in = coords_in[0].shape[-2]
    is_bb_or_ca_pdb = n_atoms_in <= 4
    for i, c in enumerate(coords_in):
        n_res = c.shape[0]
        if isinstance(filename, list):
            fname = filename[i]
        elif write_to_frames or len(coords_in) == 1:
            fname = filename
        else:
            fname = f"{filename[:-4]}_{i}.pdb"

        if is_bb_or_ca_pdb:
            c_flat = rearrange(c, "n a c -> (n a) c")
            if n_atoms_in == 1:
                atoms = ["CA"]
            if n_atoms_in == 3:
                atoms = ["N", "CA", "C"]
            if n_atoms_in == 4:
                atoms = ["N", "CA", "C", "O"]
            pdb_str = bb_coords_to_pdb_str(c_flat, atoms)
        else:
            feats_i = {k: v[i][:n_res] for k, v in all_atom_feats.items()}
            pdb_str = feats_to_pdb_str(c, conect=conect, **feats_i)
        _write_pdb_string(pdb_str, fname, append=write_to_frames and i > 0)


###################### LOSSES ###################################


def masked_cross_entropy(logprobs, target, loss_mask):
    # target is onehot
    cel = -(target * logprobs)
    cel = cel * loss_mask[..., None]
    cel = cel.sum((-1, -2)) / loss_mask.sum(-1).clamp(min=1e-6)
    return cel


def masked_mse(x, y, mask, weight=None):
    data_dims = tuple(range(1, len(x.shape)))
    mse = (x - y).pow(2) * mask
    if weight is not None:
        mse = mse * expand(weight, mse)
    mse = mse.sum(data_dims) / mask.sum(data_dims).clamp(min=1e-6)
    return mse


###################### ALIGN ###################################


def quick_tmalign(
    p, p_sele, q_sele, tmscore_type="avg", differentiable_rmsd=False, rmsd_type="ca"
):
    # sota 210712
    write_coords_to_pdb(p_sele[:, 1:2], "temp_p.pdb", atoms=["CA"], batched=False)
    write_coords_to_pdb(q_sele[:, 1:2], "temp_q.pdb", atoms=["CA"], batched=False)
    cmd = f"{PATH_TO_TMALIGN} temp_p.pdb temp_q.pdb -m temp_matrix.txt"
    outputs = subprocess.run(shlex.split(cmd), capture_output=True, text=True)

    # Get RMSD and TM scores
    tmout = outputs.stdout.split("\n")
    rmsd = float(tmout[16].split()[4][:-1])
    tmscore1 = float(tmout[17].split()[1])
    tmscore2 = float(tmout[18].split()[1])
    if tmscore_type == "avg":
        tmscore = (tmscore1 + tmscore2) / 2
    elif tmscore_type == "1" or tmscore_type == "query":
        tmscore = tmscore1
    elif tmscore_type == "2":
        tmscore = tmscore2
    elif tmscore_type == "both":
        tmscore = (tmscore1, tmscore2)

    # Get R, t and transform p coords
    m = open("temp_matrix.txt", "r").readlines()[2:5]
    m = [l.strip()[1:].strip() for l in m]
    m = torch.Tensor([[float(i) for i in l.split()] for l in m]).to(p_sele.device)
    R = m[:, 1:].t()
    t = m[:, 0]
    aligned_psele = p_sele @ R + t
    aligned = p @ R + t

    # Option 2 for rms - MSE of aligned against target coords using TMalign seq alignment. Differentiable
    if differentiable_rmsd:
        pi, qi = 0, 0
        p_idxs, q_idxs = [], []
        for i, c in enumerate(tmout[23]):
            if c in [":", "."]:
                p_idxs.append(pi)
                q_idxs.append(qi)
            if tmout[22][i] != "-":
                pi += 1
            if tmout[24][i] != "-":
                qi += 1
        tmalign_seq_p = p_sele[p_idxs]
        tmalign_seq_q = q_sele[q_idxs]
        if rmsd_type == "ca":
            tmalign_seq_p = tmalign_seq_p[:, 1]
            tmalign_seq_q = tmalign_seq_q[:, 1]
        elif rmsd_type == "bb":
            pass
        rmsd = (tmalign_seq_p - tmalign_seq_q).pow(2).sum(-1).sqrt().mean()

    # Delete temp files: p.pdb, q.pdb, matrix.txt, tmalign.out
    subprocess.run(shlex.split("rm temp_p.pdb"))
    subprocess.run(shlex.split("rm temp_q.pdb"))
    subprocess.run(shlex.split("rm temp_matrix.txt"))

    return {"aligned": aligned, "rmsd": rmsd, "tm_score": tmscore, "R": R, "t": t}


###################### OTHER ###################################


def expand(x, tgt=None, dim=1):
    if tgt is None:
        for _ in range(dim):
            x = x[..., None]
    else:
        while len(x.shape) < len(tgt.shape):
            x = x[..., None]
    return x


def hookfn(name, verbose=False):
    def f(grad):
        if check_nan_inf(grad) > 0:
            print(name, "grad nan/infs", grad.shape, check_nan_inf(grad), grad)
        if verbose:
            print(name, "grad shape", grad.shape, "norm", grad.norm())

    return f


def trigger_nan_check(name, x):
    if check_nan_inf(x) > 0:
        print(name, check_nan_inf(x))
        raise Exception


def check_nan_inf(x):
    return torch.isinf(x).sum() + torch.isnan(x).sum()


def directory_find(atom, root="."):
    for path, dirs, files in os.walk(root):
        if atom in dirs:
            return os.path.join(path, atom)


def dict2namespace(config):
    namespace = argparse.Namespace()
    for key, value in config.items():
        if isinstance(value, dict):
            new_value = dict2namespace(value)
        else:
            new_value = value
        setattr(namespace, key, new_value)
    return namespace


def load_config(path, return_dict=False):
    with open(path, "r") as f:
        config_dict = yaml.safe_load(f)
    config = dict2namespace(config_dict)
    if return_dict:
        return config, config_dict
    else:
        return config

core/utils.py
ProteinDesignLab/protpardelle
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Noise and diffusion utils.
"""
from scipy.stats import norm
import torch
from torchtyping import TensorType

from core import utils


def noise_schedule(
    time: TensorType[float],
    function: str = "uniform",
    sigma_data: float = 10.0,
    psigma_mean: float = -1.2,
    psigma_std: float = 1.2,
    s_min: float = 0.001,
    s_max: float = 60,
    rho: float = 7.0,
    time_power: float = 4.0,
    constant_val: float = 0.0,
):
    def sampling_noise(time):
        # high noise = 1; low noise = 0. opposite of Karras et al. schedule
        term1 = s_max ** (1 / rho)
        term2 = (1 - time) * (s_min ** (1 / rho) - s_max ** (1 / rho))
        noise_level = sigma_data * ((term1 + term2) ** rho)
        return noise_level

    if function == "lognormal":
        normal_sample = torch.Tensor(norm.ppf(time.cpu())).to(time)
        noise_level = sigma_data * torch.exp(psigma_mean + psigma_std * normal_sample)
    elif function == "uniform":
        noise_level = sampling_noise(time)
    elif function == "mpnn":
        time = time**time_power
        noise_level = sampling_noise(time)
    elif function == "constant":
        noise_level = torch.ones_like(time) * constant_val
    return noise_level


def noise_coords(
    coords: TensorType["b n a x", float],
    noise_level: TensorType["b", float],
    dummy_fill_masked_atoms: bool = False,
    atom_mask: TensorType["b n a"] = None,
):
    # Does not apply atom mask after adding noise
    if dummy_fill_masked_atoms:
        assert atom_mask is not None
        dummy_fill_mask = 1 - atom_mask
        dummy_fill_value = coords[..., 1:2, :]  # CA
        # dummy_fill_value = utils.fill_in_cbeta_for_atom37(coords)[..., 3:4, :]  # CB
        coords = (
            coords * atom_mask[..., None]
            + dummy_fill_value * dummy_fill_mask[..., None]
        )

    noise = torch.randn_like(coords) * utils.expand(noise_level, coords)
    noisy_coords = coords + noise
    return noisy_coords

diffusion.py
ProteinDesignLab/protpardelle
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Entry point for unconditional or simple conditional sampling.
"""
import argparse
from datetime import datetime
import json
import os
import shlex
import subprocess
import sys
import time

from einops import repeat
import torch

from core import data
from core import residue_constants
from core import utils
import diffusion
import models
import sampling


def draw_and_save_samples(
    model,
    samples_per_len=8,
    lengths=range(50, 512),
    save_dir="./",
    mode="backbone",
    **sampling_kwargs,
):
    device = model.device
    if mode == "backbone":
        total_sampling_time = 0
        for l in lengths:
            prot_lens = torch.ones(samples_per_len).long() * l
            seq_mask = model.make_seq_mask_for_sampling(prot_lens=prot_lens)
            aux = sampling.draw_backbone_samples(
                model,
                seq_mask=seq_mask,
                pdb_save_path=f"{save_dir}/len{format(l, '03d')}_samp",
                return_aux=True,
                return_sampling_runtime=True,
                **sampling_kwargs,
            )
            total_sampling_time += aux["runtime"]
            print("Samples drawn for length", l)
        return total_sampling_time
    elif mode == "allatom":
        total_sampling_time = 0
        for l in lengths:
            prot_lens = torch.ones(samples_per_len).long() * l
            seq_mask = model.make_seq_mask_for_sampling(prot_lens=prot_lens)
            aux = sampling.draw_allatom_samples(
                model,
                seq_mask=seq_mask,
                pdb_save_path=f"{save_dir}/len{format(l, '03d')}",
                return_aux=True,
                **sampling_kwargs,
            )
            total_sampling_time += aux["runtime"]
            print("Samples drawn for length", l)
        return total_sampling_time


def parse_idx_string(idx_str):
    spans = idx_str.split(",")
    idxs = []
    for s in spans:
        if "-" in s:
            start, stop = s.split("-")
            idxs.extend(list(range(int(start), int(stop))))
        else:
            idxs.append(int(s))
    return idxs


class Manager(object):
    def __init__(self):
        self.parser = argparse.ArgumentParser(
            formatter_class=argparse.RawTextHelpFormatter
        )

        self.parser.add_argument(
            "--model_checkpoint",
            type=str,
            default="checkpoints",
            help="Path to denoiser model weights and config",
        )
        self.parser.add_argument(
            "--mpnnpath",
            type=str,
            default="checkpoints/minimpnn_state_dict.pth",
            help="Path to minimpnn model weights",
        )
        self.parser.add_argument(
            "--modeldir",
            type=str,
            help="Model base directory, ex 'training_logs/other/lemon-shape-51'",
        )
        self.parser.add_argument("--modelepoch", type=int, help="Model epoch, ex 1000")
        self.parser.add_argument(
            "--type", type=str, default="allatom", help="Type of model"
        )
        self.parser.add_argument(
            "--param", type=str, default=None, help="Which sampling param to vary"
        )
        self.parser.add_argument(
            "--paramval", type=str, default=None, help="Which param val to use"
        )
        self.parser.add_argument(
            "--parampath",
            type=str,
            default=None,
            help="Path to json file with params, either use param/paramval or parampath, not both",
        )
        self.parser.add_argument(
            "--perlen", type=int, default=2, help="How many samples per sequence length"
        )
        self.parser.add_argument(
            "--minlen", type=int, default=50, help="Minimum sequence length"
        )
        self.parser.add_argument(
            "--maxlen",
            type=int,
            default=60,
            help="Maximum sequence length, not inclusive",
        )
        self.parser.add_argument(
            "--steplen",
            type=int,
            default=5,
            help="How frequently to select sequence length, for steplen 2, would be 50, 52, 54, etc",
        )
        self.parser.add_argument(
            "--num_lens",
            type=int,
            required=False,
            help="If steplen not provided, how many random lengths to sample at",
        )
        self.parser.add_argument(
            "--targetdir", type=str, default=".", help="Directory to save results"
        )
        self.parser.add_argument(
            "--input_pdb", type=str, required=False, help="PDB file to condition on"
        )
        self.parser.add_argument(
            "--resample_idxs",
            type=str,
            required=False,
            help="Indices from PDB file to resample. Zero-indexed, comma-delimited, can use dashes, eg 0,2-5,7",
        )

    def add_argument(self, *args, **kwargs):
        self.parser.add_argument(*args, **kwargs)

    def parse_args(self):
        self.args = self.parser.parse_args()

        return self.args


def main():
    # Set up params, arguments, sampling config
    ####################
    manager = Manager()
    manager.parse_args()
    args = manager.args
    print(args)
    is_test_run = False
    seed = 0
    samples_per_len = args.perlen
    min_len = args.minlen
    max_len = args.maxlen
    len_step_size = args.steplen
    device = "cuda:0"

    # setting default sampling config
    if args.type == "backbone":
        sampling_config = sampling.default_backbone_sampling_config()
    elif args.type == "allatom":
        sampling_config = sampling.default_allatom_sampling_config()

    sampling_kwargs = vars(sampling_config)

    # Parse conditioning inputs
    input_pdb_len = None
    if args.input_pdb:
        input_feats = utils.load_feats_from_pdb(args.input_pdb, protein_only=True)
        input_pdb_len = input_feats["aatype"].shape[0]
        if args.resample_idxs:
            print(
                f"Warning: when sampling conditionally, the input pdb length ({input_pdb_len} residues) is used automatically for the sampling lengths."
            )
            resample_idxs = parse_idx_string(args.resample_idxs)
        else:
            resample_idxs = list(range(input_pdb_len))
        cond_idxs = [i for i in range(input_pdb_len) if i not in resample_idxs]
        to_batch_size = lambda x: repeat(x, "... -> b ...", b=samples_per_len).to(
            device
        )

        # For unconditional model, center coords on whole structure
        centered_coords = data.apply_random_se3(
            input_feats["atom_positions"],
            atom_mask=input_feats["atom_mask"],
            translation_scale=0.0,
        )
        cond_kwargs = {}
        cond_kwargs["gt_coords"] = to_batch_size(centered_coords)
        cond_kwargs["gt_cond_atom_mask"] = to_batch_size(input_feats["atom_mask"])
        cond_kwargs["gt_cond_atom_mask"][:, resample_idxs] = 0
        cond_kwargs["gt_aatype"] = to_batch_size(input_feats["aatype"])
        cond_kwargs["gt_cond_seq_mask"] = torch.zeros_like(cond_kwargs["gt_aatype"])
        cond_kwargs["gt_cond_seq_mask"][:, cond_idxs] = 1
        sampling_kwargs.update(cond_kwargs)

    # Determine lengths to sample at
    if min_len is not None and max_len is not None:
        if len_step_size is not None:
            sampling_lengths = range(min_len, max_len, len_step_size)
        else:
            sampling_lengths = list(
                torch.randint(min_len, max_len, size=(args.num_lens,))
            )
    elif input_pdb_len is not None:
        sampling_lengths = [input_pdb_len]
    else:
        raise Exception("Need to provide a set of protein lengths or an input pdb.")

    total_num_samples = len(list(sampling_lengths)) * samples_per_len

    model_directory = args.modeldir
    epoch = args.modelepoch
    base_dir = args.targetdir

    date_string = datetime.now().strftime("%y-%m-%d-%H-%M-%S")
    if is_test_run:
        date_string = f"test-{date_string}"

    # Update sampling config with arguments
    if args.param:
        var_param = args.param
        var_value = args.paramval
        sampling_kwargs[var_param] = (
            None
            if var_value == "None"
            else int(var_value)
            if var_param == "n_steps"
            else float(var_value)
        )
    elif args.parampath:
        with open(args.parampath) as f:
            var_params = json.loads(f.read())
            sampling_kwargs.update(var_params)

    # this is only used for the readme, keep s_min and s_max as params instead of struct_noise_schedule
    sampling_kwargs_readme = list(sampling_kwargs.items())

    print("Base directory:", base_dir)
    save_dir = f"{base_dir}/samples"
    save_init_dir = f"{base_dir}/samples_inits"

    print("Samples saved to:", save_dir)
    ####################

    torch.manual_seed(seed)
    if not os.path.exists(save_dir):
        subprocess.run(shlex.split(f"mkdir -p {save_dir}"))

    if not os.path.exists(save_init_dir):
        subprocess.run(shlex.split(f"mkdir -p {save_init_dir}"))

    # Load model
    if args.type == "backbone":
        if args.model_checkpoint:
            checkpoint = f"{args.model_checkpoint}/backbone_state_dict.pth"
            cfg_path = f"{args.model_checkpoint}/backbone_pretrained.yml"
        else:
            checkpoint = (
                f"{model_directory}/checkpoints/epoch{epoch}_training_state.pth"
            )
            cfg_path = f"{model_directory}/configs/backbone.yml"
        config = utils.load_config(cfg_path)
        weights = torch.load(checkpoint, map_location=device)["model_state_dict"]
        model = models.Protpardelle(config, device=device)
        model.load_state_dict(weights)
        model.to(device)
        model.eval()
        model.device = device
    elif args.type == "allatom":
        if args.model_checkpoint:
            checkpoint = f"{args.model_checkpoint}/allatom_state_dict.pth"
            cfg_path = f"{args.model_checkpoint}/allatom_pretrained.yml"
        else:
            checkpoint = (
                f"{model_directory}/checkpoints/epoch{epoch}_training_state.pth"
            )
            cfg_path = f"{model_directory}/configs/allatom.yml"
        config = utils.load_config(cfg_path)
        weights = torch.load(checkpoint, map_location=device)["model_state_dict"]
        model = models.Protpardelle(config, device=device)
        model.load_state_dict(weights)
        model.load_minimpnn(args.mpnnpath)
        model.to(device)
        model.eval()
        model.device = device

    if config.train.home_dir == '':
        config.train.home_dir = os.path.dirname(os.getcwd())

    # Sampling
    with open(save_dir + "/readme.txt", "w") as f:
        f.write(f"Sampling run for {date_string}\n")
        f.write(f"Random seed {seed}\n")
        f.write(f"Model checkpoint: {checkpoint}\n")
        f.write(
            f"{samples_per_len} samples per length from {min_len}:{max_len}:{len_step_size}\n"
        )
        f.write("Sampling params:\n")
        for k, v in sampling_kwargs_readme:
            f.write(f"{k}\t{v}\n")

    print(f"Model loaded from {checkpoint}")
    print(f"Beginning sampling for {date_string}...")

    # Draw samples
    start_time = time.time()
    sampling_time = draw_and_save_samples(
        model,
        samples_per_len=samples_per_len,
        lengths=sampling_lengths,
        save_dir=save_dir,
        mode=args.type,
        **sampling_kwargs,
    )
    time_elapsed = time.time() - start_time

    print(f"Sampling concluded after {time_elapsed} seconds.")
    print(f"Of this, {sampling_time} seconds were for actual sampling.")
    print(f"{total_num_samples} total samples were drawn.")

    with open(save_dir + "/readme.txt", "a") as f:
        f.write(f"Total job time: {time_elapsed} seconds\n")
        f.write(f"Model run time: {sampling_time} seconds\n")
        f.write(f"Total samples drawn: {total_num_samples}\n")

    return


if __name__ == "__main__":
    main()

draw_samples.py
ProteinDesignLab/protpardelle
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Utils for computing evaluation metrics.
"""
import argparse
import os
import warnings
from typing import Tuple

from Bio.Align import substitution_matrices
import numpy as np
import torch
from transformers import AutoTokenizer, EsmForProteinFolding
from torchtyping import TensorType

from core import residue_constants
from core import utils
from core import protein_mpnn as mpnn
import modules
import sampling


def mean(x):
    if len(x) == 0:
        return 0
    return sum(x) / len(x)


def calculate_seq_identity(seq1, seq2, seq_mask=None):
    identity = (seq1 == seq2.to(seq1)).float()
    if seq_mask is not None:
        identity *= seq_mask.to(seq1)
        return identity.sum(-1) / seq_mask.to(seq1).sum(-1).clamp(min=1)
    else:
        return identity.mean(-1)


def design_sequence(coords, model=None, num_seqs=1, disallow_aas=["C"]):
    # Returns list of strs; seqs like 'MKRLLDS', not aatypes
    if model is None:
        model = mpnn.get_mpnn_model()
    if isinstance(coords, str):
        temp_pdb = False
        pdb_fn = coords
    else:
        temp_pdb = True
        pdb_fn = f"tmp{np.random.randint(0, 1e8)}.pdb"
        gly_idx = residue_constants.restype_order["G"]
        gly_aatype = (torch.ones(coords.shape[0]) * gly_idx).long()
        utils.write_coords_to_pdb(coords, pdb_fn, batched=False, aatype=gly_aatype)

    with torch.no_grad():
        designed_seqs = mpnn.run_proteinmpnn(
            model=model,
            pdb_path=pdb_fn,
            num_seq_per_target=num_seqs,
            omit_AAs=disallow_aas,
        )

    if temp_pdb:
        os.system("rm " + pdb_fn)
    return designed_seqs


def get_esmfold_model(device=None):
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1").to(device)
    model.esm = model.esm.half()
    return model


def inference_esmfold(sequence_list, model, tokenizer):
    inputs = tokenizer(
        sequence_list,
        return_tensors="pt",
        padding=True,
        add_special_tokens=False,
    ).to(model.device)
    outputs = model(**inputs)
    # positions is shape (l, b, n, a, c)
    pred_coords = outputs.positions[-1].contiguous()
    plddts = (outputs.plddt[:, :, 1] * inputs.attention_mask).sum(
        -1
    ) / inputs.attention_mask.sum(-1).clamp(min=1e-3)
    return pred_coords, plddts


def predict_structures(sequences, model="esmfold", tokenizer=None, force_unk_to_X=True):
    # Expects seqs like 'MKRLLDS', not aatypes
    # model can be a model, or a string describing which pred model to load
    if isinstance(sequences, str):
        sequences = [sequences]
    if model == "esmfold":
        model = get_esmfold_model()
    device = model.device
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")

    aatype = [utils.seq_to_aatype(seq).to(device) for seq in sequences]

    with torch.no_grad():
        if isinstance(model, EsmForProteinFolding):
            pred_coords, plddts = inference_esmfold(sequences, model, tokenizer)

    seq_lens = [len(s) for s in sequences]
    trimmed_coords = [c[: seq_lens[i]] for i, c in enumerate(pred_coords)]
    trimmed_coords_atom37 = [
        utils.atom37_coords_from_atom14(c, aatype[i])
        for i, c in enumerate(trimmed_coords)
    ]
    return trimmed_coords_atom37, plddts


def compute_structure_metric(coords1, coords2, metric="ca_rmsd", atom_mask=None):
    # coords1 tensor[l][a][3]
    def _tmscore(a, b, mask=None):
        length = len(b)
        dists = (a - b).pow(2).sum(-1)
        d0 = 1.24 * ((length - 15) ** (1 / 3)) - 1.8
        term = 1 / (1 + ((dists) / (d0**2)))
        if mask is None:
            return term.mean()
        else:
            term = term * mask
            return term.sum() / mask.sum().clamp(min=1)

    aligned_coords1_ca, (R, t) = utils.kabsch_align(coords1[:, 1], coords2[:, 1])
    aligned_coords1 = coords1 - coords1[:, 1:2].mean(0, keepdim=True)
    aligned_coords1 = aligned_coords1 @ R.t() + t
    if metric == "ca_rmsd":
        return (aligned_coords1_ca - coords2[:, 1]).pow(2).sum(-1).sqrt().mean()
    elif metric == "tm_score":
        tm = _tmscore(aligned_coords1_ca, coords2[:, 1])
        # TODO: return 1 - tm score for now so sorts work properly
        return 1 - tm
    elif metric == "allatom_tm":
        # Align on Ca, compute allatom TM
        assert atom_mask is not None
        return _tmscore(aligned_coords1, coords2, mask=atom_mask)
    elif metric == "allatom_lddt":
        assert atom_mask is not None
        lddt = modules.lddt(
            coords1.reshape(-1, 3),
            coords2.reshape(-1, 3),
            atom_mask.reshape(-1, 1),
            per_residue=False,
        )
        return lddt
    else:
        raise NotImplementedError


def compute_self_consistency(
    comparison_structures,  # can be sampled or ground truth
    sampled_sequences=None,
    mpnn_model=None,
    struct_pred_model=None,
    tokenizer=None,
    num_seqs=1,
    return_aux=False,
    metric="ca_rmsd",
    output_file=None,
):
    # Typically used for eval of backbone sampling or sequence design or joint sampling
    # (Maybe MPNN) + Fold + TM/RMSD
    # Expects seqs like 'MKRLLDS', not aatypes
    per_sample_primary_metrics = []
    per_sample_secondary_metrics = []
    per_sample_plddts = []
    per_sample_coords = []
    per_sample_seqs = []
    aux = {}
    for i, coords in enumerate(comparison_structures):
        if sampled_sequences is None:
            seqs_to_predict = design_sequence(
                coords, model=mpnn_model, num_seqs=num_seqs
            )
        else:
            seqs_to_predict = sampled_sequences[i]
        pred_coords, plddts = predict_structures(
            seqs_to_predict, model=struct_pred_model, tokenizer=tokenizer
        )
        primary_metric_name = "tm_score" if metric == "tm_score" else "ca_rmsd"
        secondary_metric_name = "tm_score" if metric == "both" else None
        primary_metrics = [
            compute_structure_metric(coords.to(pred), pred, metric=primary_metric_name)
            for pred in pred_coords
        ]
        if secondary_metric_name:
            secondary_metrics = [
                compute_structure_metric(
                    coords.to(pred), pred, metric=secondary_metric_name
                )
                for pred in pred_coords
            ]
            aux.setdefault(secondary_metric_name, []).extend(secondary_metrics)
        else:
            secondary_metrics = primary_metrics

        aux.setdefault("pred", []).extend(pred_coords)
        seqs_to_predict_arr = seqs_to_predict
        if isinstance(seqs_to_predict_arr, str):
            seqs_to_predict_arr = [seqs_to_predict_arr]

        aux.setdefault("seqs", []).extend(seqs_to_predict_arr)
        aux.setdefault("plddt", []).extend(plddts)
        aux.setdefault("rmsd", []).extend(primary_metrics)

        # Report best rmsd design only among MPNN reps
        all_designs = [
            (m, p, t, c, s)
            for m, p, t, c, s in zip(
                primary_metrics,
                plddts,
                secondary_metrics,
                pred_coords,
                seqs_to_predict_arr,
            )
        ]
        best_rmsd_design = min(all_designs, key=lambda x: x[0])
        per_sample_primary_metrics.append(best_rmsd_design[0].detach().cpu())
        per_sample_plddts.append(best_rmsd_design[1].detach().cpu())
        per_sample_secondary_metrics.append(best_rmsd_design[2].detach().cpu())
        per_sample_coords.append(best_rmsd_design[3])
        per_sample_seqs.append(best_rmsd_design[4])
    best_idx = np.argmin(per_sample_primary_metrics)
    metrics = {
        "sc_rmsd_best": per_sample_primary_metrics[best_idx],
        "sc_plddt_best": per_sample_plddts[best_idx],
        "sc_rmsd_mean": mean(per_sample_primary_metrics),
        "sc_plddt_mean": mean(per_sample_plddts),
    }
    if metric == "both":
        metrics["sc_tmscore_best"] = per_sample_secondary_metrics[best_idx]
        metrics["sc_tmscore_mean"] = mean(per_sample_secondary_metrics)

    if output_file:
        pred_coords = per_sample_coords
        designed_seqs = per_sample_seqs

        if torch.isnan(pred_coords[best_idx]).sum() == 0:
            designed_seq = utils.seq_to_aatype(designed_seqs[best_idx])
            utils.write_coords_to_pdb(
                pred_coords[best_idx],
                output_file,
                batched=False,
                aatype=designed_seq,
            )

    if return_aux:
        return metrics, best_idx, aux
    else:
        return metrics, best_idx


def compute_secondary_structure_content(coords_batch):
    dssp_sample = []
    for i, c in enumerate(coords_batch):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            dssp_str = utils.get_3state_dssp(coords=c)
        if dssp_str is None or len(dssp_str) == 0:
            pass
        else:
            dssp_sample.append(dssp_str)
    dssp_sample = "".join(dssp_sample)
    metrics = {}
    metrics["sample_pct_beta"] = mean([c == "E" for c in dssp_sample])
    metrics["sample_pct_alpha"] = mean([c == "H" for c in dssp_sample])
    return metrics


def compute_bond_length_metric(
    cropped_coords_list, cropped_aatypes_list, atom_mask=None
):
    bond_length_dict = utils.batched_fullatom_bond_lengths_from_coords(
        cropped_coords_list, cropped_aatypes_list, atom_mask=atom_mask
    )
    all_errors = {}
    for aa1, d in bond_length_dict.items():
        aa3 = residue_constants.restype_1to3[aa1]
        per_bond_errors = []
        for bond, lengths in d.items():
            a1, a2 = bond.split("-")
            ideal_val = None
            for bond in residue_constants.standard_residue_bonds[aa3]:
                if (
                    bond.atom1_name == a1
                    and bond.atom2_name == a2
                    or bond.atom1_name == a2
                    and bond.atom2_name == a1
                ):
                    ideal_val = bond.length
                    break
            error = (np.array(lengths) - ideal_val) ** 2
            per_bond_errors.append(error.mean() ** 0.5)
        if len(per_bond_errors) > 0:  # often no Cys
            per_res_errors = np.mean(per_bond_errors)
            all_errors[aa1] = per_res_errors
    return np.mean(list(all_errors.values()))


def evaluate_backbone_generation(
    model,
    n_samples=1,
    mpnn_model=None,
    struct_pred_model=None,
    tokenizer=None,
    sample_length_range=(50, 512),
):
    sampling_config = sampling.default_backbone_sampling_config()
    trimmed_coords, seq_mask = sampling.draw_backbone_samples(
        model,
        n_samples=n_samples,
        sample_length_range=sample_length_range,
        **vars(sampling_config),
    )
    sc_metrics, best_idx, aux = compute_self_consistency(
        trimmed_coords,
        mpnn_model=mpnn_model,
        struct_pred_model=struct_pred_model,
        tokenizer=tokenizer,
        return_aux=True,
    )
    dssp_metrics = compute_secondary_structure_content(trimmed_coords)
    all_metrics = {**sc_metrics, **dssp_metrics}
    all_metrics = {f"bb_{k}": v for k, v in all_metrics.items()}
    return all_metrics, (trimmed_coords, seq_mask, best_idx, aux["pred"], aux["seqs"])


def evaluate_allatom_generation(
    model,
    n_samples,
    two_stage_sampling=True,
    struct_pred_model=None,
    tokenizer=None,
    sample_length_range=(50, 512),
):
    # Convert allatom model to codesign model by loading miniMPNN
    model.task = "codesign"
    model.load_minimpnn()
    model.eval()

    sampling_config = sampling.default_allatom_sampling_config()
    ret = sampling.draw_allatom_samples(
        model,
        n_samples=n_samples,
        two_stage_sampling=two_stage_sampling,
        **vars(sampling_config),
    )
    (
        cropped_samp_coords,
        cropped_samp_aatypes,
        samp_atom_mask,
        stage1_coords,
        seq_mask,
    ) = ret

    # Compute self consistency
    if struct_pred_model is None:
        struct_pred_model = EsmForProteinFolding.from_pretrained(
            "facebook/esmfold_v1"
        ).to(device)
        struct_pred_model.esm = struct_pred_model.esm.half()
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
    designed_seqs = [utils.aatype_to_seq(a) for a in cropped_samp_aatypes]
    sc_metrics, best_idx, sc_aux = compute_self_consistency(
        comparison_structures=cropped_samp_coords,
        sampled_sequences=designed_seqs,
        struct_pred_model=struct_pred_model,
        tokenizer=tokenizer,
        return_aux=True,
    )
    aa_metrics_out = {f"aa_{k}": v for k, v in sc_metrics.items()}

    # Compute secondary structure content
    cropped_bb_coords = [c[..., [0, 1, 2, 4], :] for c in cropped_samp_coords]
    dssp_metrics = compute_secondary_structure_content(cropped_bb_coords)
    aa_metrics_out = {**aa_metrics_out, **dssp_metrics}

    # Compute bond length RMSE
    if two_stage_sampling:  # compute on original sample
        bond_rmse_coords = stage1_coords
    else:
        bond_rmse_coords = cropped_samp_coords
    bond_rmse = compute_bond_length_metric(
        bond_rmse_coords, cropped_samp_aatypes, samp_atom_mask
    )
    aa_metrics_out["aa_bond_rmse"] = bond_rmse

    # Convert codesign model back to allatom model and return metrics
    model.task = "allatom"
    model.remove_minimpnn()
    aa_aux_out = (
        cropped_samp_coords,
        cropped_samp_aatypes,
        samp_atom_mask,
        sc_aux["pred"],
        best_idx,
    )
    return aa_metrics_out, aa_aux_out

evaluation.py
ProteinDesignLab/protpardelle
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Top-level model definitions.
Typically these are initialized with config rather than arguments.
"""
import argparse
from functools import partial
import os
from typing import Callable, List, Optional

from einops import rearrange, repeat
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchtyping import TensorType

from core import protein_mpnn
from core import residue_constants
from core import utils
import diffusion
import evaluation
import modules


class MiniMPNN(nn.Module):
    """Wrapper for ProteinMPNN network to predict sequence from structure."""

    def __init__(self, config: argparse.Namespace):
        super().__init__()
        self.config = config
        self.model_config = cfg = config.model.mpnn_model
        self.n_tokens = config.data.n_aatype_tokens
        self.seq_emb_dim = cfg.n_channel
        time_cond_dim = cfg.n_channel * cfg.noise_cond_mult

        self.noise_block = modules.NoiseConditioningBlock(cfg.n_channel, time_cond_dim)
        self.token_embedding = nn.Linear(self.n_tokens, self.seq_emb_dim)
        self.mpnn_net = modules.NoiseConditionalProteinMPNN(
            n_channel=cfg.n_channel,
            n_layers=cfg.n_layers,
            n_neighbors=cfg.n_neighbors,
            time_cond_dim=time_cond_dim,
            vocab_size=config.data.n_aatype_tokens,
            input_S_is_embeddings=True,
        )
        self.proj_out = nn.Linear(cfg.n_channel, self.n_tokens)

    def forward(
        self,
        denoised_coords: TensorType["b n a x", float],
        coords_noise_level: TensorType["b", float],
        seq_mask: TensorType["b n", float],
        residue_index: TensorType["b n", int],
        seq_self_cond: Optional[TensorType["b n t", float]] = None,  # logprobs
        return_embeddings: bool = False,
    ):
        coords_noise_level_scaled = 0.25 * torch.log(coords_noise_level)
        noise_cond = self.noise_block(coords_noise_level_scaled)

        b, n, _, _ = denoised_coords.shape
        if seq_self_cond is None or not self.model_config.use_self_conditioning:
            seq_emb_in = torch.zeros(b, n, self.seq_emb_dim).to(denoised_coords)
        else:
            seq_emb_in = self.token_embedding(seq_self_cond.exp())

        node_embs, encoder_embs = self.mpnn_net(
            denoised_coords, seq_emb_in, seq_mask, residue_index, noise_cond
        )

        logits = self.proj_out(node_embs)
        pred_logprobs = F.log_softmax(logits, -1)

        if return_embeddings:
            return pred_logprobs, node_embs, encoder_embs
        return pred_logprobs


class CoordinateDenoiser(nn.Module):
    """Wrapper for U-ViT module to denoise structure coordinates."""

    def __init__(self, config: argparse.Namespace):
        super().__init__()
        self.config = config

        # Configuration
        self.sigma_data = config.data.sigma_data
        m_cfg = config.model.struct_model
        nc = m_cfg.n_channel
        bb_atoms = ["N", "CA", "C", "O"]
        n_atoms = config.model.struct_model.n_atoms
        self.use_conv = len(m_cfg.uvit.n_filt_per_layer) > 0
        if self.use_conv and n_atoms == 37:
            n_atoms += 1  # make it an even number
        self.n_atoms = n_atoms
        self.bb_idxs = [residue_constants.atom_order[a] for a in bb_atoms]
        n_xyz = 9 if config.model.crop_conditional else 6
        nc_in = n_xyz * n_atoms  # xyz + selfcond xyz + maybe cropcond xyz

        # Neural networks
        n_noise_channel = nc * m_cfg.noise_cond_mult
        self.net = modules.TimeCondUViT(
            seq_len=config.data.fixed_size,
            patch_size=m_cfg.uvit.patch_size,
            dim=nc,
            depth=m_cfg.uvit.n_layers,
            n_filt_per_layer=m_cfg.uvit.n_filt_per_layer,
            heads=m_cfg.uvit.n_heads,
            dim_head=m_cfg.uvit.dim_head,
            conv_skip_connection=m_cfg.uvit.conv_skip_connection,
            n_atoms=n_atoms,
            channels_per_atom=n_xyz,
            time_cond_dim=n_noise_channel,
            position_embedding_type=m_cfg.uvit.position_embedding_type,
        )
        self.noise_block = modules.NoiseConditioningBlock(nc, n_noise_channel)

    def forward(
        self,
        noisy_coords: TensorType["b n a x", float],
        noise_level: TensorType["b", float],
        seq_mask: TensorType["b n", float],
        residue_index: Optional[TensorType["b n", int]] = None,
        struct_self_cond: Optional[TensorType["b n a x", float]] = None,
        struct_crop_cond: Optional[TensorType["b n a x", float]] = None,
    ):
        # Prep inputs and time conditioning
        actual_var_data = self.sigma_data**2
        var_noisy_coords = noise_level**2 + actual_var_data
        emb = noisy_coords / utils.expand(var_noisy_coords.sqrt(), noisy_coords)
        struct_noise_scaled = 0.25 * torch.log(noise_level)
        noise_cond = self.noise_block(struct_noise_scaled)

        # Prepare self- and crop-conditioning and concatenate along channels
        if struct_self_cond is None:
            struct_self_cond = torch.zeros_like(noisy_coords)
        if self.config.model.crop_conditional:
            if struct_crop_cond is None:
                struct_crop_cond = torch.zeros_like(noisy_coords)
            else:
                struct_crop_cond = struct_crop_cond / self.sigma_data
            emb = torch.cat([emb, struct_self_cond, struct_crop_cond], -1)
        else:
            emb = torch.cat([emb, struct_self_cond], -1)

        # Run neural network
        emb = self.net(emb, noise_cond, seq_mask=seq_mask, residue_index=residue_index)

        # Preconditioning from Karras et al.
        out_scale = noise_level * actual_var_data**0.5 / torch.sqrt(var_noisy_coords)
        skip_scale = actual_var_data / var_noisy_coords
        emb = emb * utils.expand(out_scale, emb)
        skip_info = noisy_coords * utils.expand(skip_scale, noisy_coords)
        denoised_coords_x0 = emb + skip_info

        # Don't use atom mask; denoise all atoms
        denoised_coords_x0 *= utils.expand(seq_mask, denoised_coords_x0)
        return denoised_coords_x0


class Protpardelle(nn.Module):
    """All-atom protein diffusion-based generative model.

    This class wraps a structure denoising network and a sequence prediction network
    to do structure/sequence co-design (for all-atom generation), or backbone generation.

    It can be trained for one of four main tasks. To produce the all-atom (co-design)
    Protpardelle model, we will typically pretrain an 'allatom' model, then use this
    to train a 'seqdes' model. A 'seqdes' model can be trained with either a backbone
    or allatom denoiser. The two can be combined to yield all-atom (co-design) Protpardelle
    without further training.
        'backbone': train only a backbone coords denoiser.
        'seqdes': train only a mini-MPNN, using a pretrained coords denoiser.
        'allatom': train only an allatom coords denoiser (cannot do all-atom generation
            by itself).
        'codesign': train both an allatom denoiser and mini-MPNN at once.

    """

    def __init__(self, config: argparse.Namespace, device: str = "cpu"):
        super().__init__()
        self.config = config
        self.device = device
        self.task = config.model.task
        self.n_tokens = config.data.n_aatype_tokens

        self.use_mpnn_model = self.task in ["seqdes", "codesign"]

        # Modules
        self.all_modules = {}
        self.bb_idxs = [0, 1, 2, 4]
        self.n_atoms = 37
        self.struct_model = CoordinateDenoiser(config)
        self.all_modules["struct_model"] = self.struct_model
        self.bb_idxs = self.struct_model.bb_idxs
        self.n_atoms = self.struct_model.n_atoms

        if self.use_mpnn_model:
            self.mpnn_model = MiniMPNN(config)
            self.all_modules["mpnn_model"] = self.mpnn_model

        # Load any pretrained modules
        for module_name in self.config.model.pretrained_modules:
            self.load_pretrained_module(module_name)

        # Diffusion-related
        self.sigma_data = self.struct_model.sigma_data
        self.training_noise_schedule = partial(
            diffusion.noise_schedule,
            sigma_data=self.sigma_data,
            **vars(config.diffusion.training),
        )
        self.sampling_noise_schedule_default = self.make_sampling_noise_schedule()

    def load_pretrained_module(self, module_name: str, ckpt_path: Optional[str] = None):
        """Load pretrained weights for a given module name."""
        assert module_name in ["struct_model", "mpnn_model"], module_name

        # Load pretrained checkpoint
        if ckpt_path is None:
            ckpt_path = getattr(self.config.model, f"{module_name}_checkpoint")
            ckpt_path = os.path.join(self.config.train.home_dir, ckpt_path)
        ckpt_dict = torch.load(ckpt_path, map_location=self.device)
        model_state_dict = ckpt_dict["model_state_dict"]

        # Get only submodule state_dict
        submodule_state_dict = {
            sk[len(module_name) + 1 :]: sv
            for sk, sv in model_state_dict.items()
            if sk.startswith(module_name)
        }

        # Load into module
        module = dict(self.named_modules())[module_name]
        module.load_state_dict(submodule_state_dict)

        # Freeze unneeded modules
        if module_name == "struct_model":
            self.struct_model = module
            if self.task == "seqdes":
                for p in module.parameters():
                    p.requires_grad = False
        if module_name == "mpnn_model":
            self.mpnn_model = module
            if self.task not in ["codesign", "seqdes"]:
                for p in module.parameters():
                    p.requires_grad = False

        return module

    def load_minimpnn(self, mpnn_ckpt_path: Optional[str] = None):
        """Convert an allatom model to a codesign model."""
        if mpnn_ckpt_path is None:
            mpnn_ckpt_path = "checkpoints/minimpnn_state_dict.pth"
        self.mpnn_model = MiniMPNN(self.config).to(self.device)
        self.load_pretrained_module("mpnn_model", ckpt_path=mpnn_ckpt_path)
        self.use_mpnn_model = True
        return

    def remove_minimpnn(self):
        """Revert a codesign model to an allatom model to a codesign model."""
        self.use_mpnn_model = False
        self.mpnn_model = None
        self.all_modules["mpnn_model"] = None

    def make_sampling_noise_schedule(self, **noise_kwargs):
        """Make the default sampling noise schedule function."""
        noise_schedule_kwargs = vars(self.config.diffusion.sampling)
        if len(noise_kwargs) > 0:
            noise_schedule_kwargs.update(noise_kwargs)
        return partial(diffusion.noise_schedule, **noise_schedule_kwargs)

    def forward(
        self,
        *,
        noisy_coords: TensorType["b n a x", float],
        noise_level: TensorType["b", float],
        seq_mask: TensorType["b n", float],
        residue_index: TensorType["b n", int],
        struct_self_cond: Optional[TensorType["b n a x", float]] = None,
        struct_crop_cond: Optional[TensorType["b n a x", float]] = None,
        seq_self_cond: Optional[TensorType["b n t", float]] = None,  # logprobs
        run_struct_model: bool = True,
        run_mpnn_model: bool = True,
    ):
        """Main forward function for denoising/co-design.

        Arguments:
            noisy_coords: noisy array of xyz coordinates.
            noise_level: std of noise for each example in the batch.
            seq_mask: mask indicating which indexes contain data.
            residue_index: residue ordering. This is used by proteinMPNN, but currently
                only used by the diffusion model when the 'absolute_residx' or
                'relative' position_embedding_type is specified.
            struct_self_cond: denoised coordinates from the previous step, scaled
                down by sigma data.
            struct_crop_cond: unnoised coordinates. unscaled (scaled down by sigma
                data inside the denoiser)
            seq_self_cond: mpnn-predicted sequence logprobs from the previous step.
            run_struct_model: flag to optionally not run structure denoiser.
            run_mpnn_model: flag to optionally not run mini-mpnn.
        """

        # Coordinate denoiser
        denoised_x0 = noisy_coords
        if run_struct_model:
            denoised_x0 = self.struct_model(
                noisy_coords,
                noise_level,
                seq_mask,
                residue_index=residue_index,
                struct_self_cond=struct_self_cond,
                struct_crop_cond=struct_crop_cond,
            )

        # Mini-MPNN
        aatype_logprobs = None
        if self.use_mpnn_model and run_mpnn_model:
            aatype_logprobs = self.mpnn_model(
                denoised_x0.detach(),
                noise_level,
                seq_mask,
                residue_index,
                seq_self_cond=seq_self_cond,
                return_embeddings=False,
            )
            aatype_logprobs = aatype_logprobs * seq_mask[..., None]

        # Process outputs
        if aatype_logprobs is None:
            aatype_logprobs = repeat(seq_mask, "b n -> b n t", t=self.n_tokens)
            aatype_logprobs = torch.ones_like(aatype_logprobs)
            aatype_logprobs = F.log_softmax(aatype_logprobs, -1)
        struct_self_cond_out = denoised_x0.detach() / self.sigma_data
        seq_self_cond_out = aatype_logprobs.detach()

        return denoised_x0, aatype_logprobs, struct_self_cond_out, seq_self_cond_out

    def make_seq_mask_for_sampling(
        self,
        prot_lens: Optional[TensorType["b", int]] = None,
        n_samples: int = 1,
        min_len: int = 50,
        max_len: Optional[int] = None,
    ):
        """Makes a sequence mask of varying protein lengths (only input required
        to begin sampling).
        """
        if max_len is None:
            max_len = self.config.data.fixed_size
        if prot_lens is None:
            possible_lens = np.arange(min_len, max_len)
            prot_lens = torch.Tensor(np.random.choice(possible_lens, n_samples))
        else:
            n_samples = len(prot_lens)
            max_len = max(prot_lens)
        mask = repeat(torch.arange(max_len), "n -> b n", b=n_samples)
        mask = (mask < prot_lens[:, None]).float().to(self.device)
        return mask

    def sample(
        self,
        *,
        seq_mask: TensorType["b n", float] = None,
        n_samples: int = 1,
        min_len: int = 50,
        max_len: int = 512,
        residue_index: TensorType["b n", int] = None,
        gt_coords: TensorType["b n a x", float] = None,
        gt_coords_traj: List[TensorType["b n a x", float]] = None,
        gt_cond_atom_mask: TensorType["b n a", float] = None,
        gt_aatype: TensorType["b n", int] = None,
        gt_cond_seq_mask: TensorType["b n", float] = None,
        apply_cond_proportion: float = 1.0,
        n_steps: int = 200,
        step_scale: float = 1.2,
        s_churn: float = 50.0,
        noise_scale: float = 1.0,
        s_t_min: float = 0.01,
        s_t_max: float = 50.0,
        temperature: float = 1.0,
        top_p: float = 1.0,
        disallow_aas: List[int] = [4, 20],  # cys, unk
        sidechain_mode: bool = False,
        skip_mpnn_proportion: float = 0.7,
        anneal_seq_resampling_rate: Optional[str] = None,  # linear, cosine
        use_fullmpnn: bool = False,
        use_fullmpnn_for_final: bool = True,
        use_reconstruction_guidance: bool = False,
        use_classifier_free_guidance: bool = False,  # defaults to replacement guidance if these are all false
        guidance_scale: float = 1.0,
        noise_schedule: Optional[Callable] = None,
        tqdm_pbar: Optional[Callable] = None,
        return_last: bool = True,
        return_aux: bool = False,
    ):
        """Sampling function for backbone or all-atom diffusion. All arguments are optional.

        Arguments:
            seq_mask: mask defining the number and lengths of proteins to be sampled.
            n_samples: number of samples to draw (if seq_mask not provided).
            min_len: minimum length of proteins to be sampled (if seq_mask not provided).
            max_len: maximum length of proteins to be sampled (if seq_mask not provided).
            residue_index: residue index of proteins to be sampled.
            gt_coords: conditioning information for coords.
            gt_coords_traj: conditioning information for coords specified for each timestep
                (if gt_coords is not provided).
            gt_cond_atom_mask: mask identifying atoms to apply gt_coords.
            gt_aatype: conditioning information for sequence.
            gt_cond_seq_mask: sequence positions to apply gt_aatype.
            apply_cond_proportion: the proportion of timesteps to apply the conditioning.
                e.g. if 0.5, then the first 50% of steps use conditioning, and the last 50%
                are unconditional.
            n_steps: number of denoising steps (ODE discretizations).
            step_scale: scale to apply to the score.
            s_churn: gamma = s_churn / n_steps describes the additional noise to add
                relatively at each denoising step. Use 0.0 for deterministic sampling or
                0.2 * n_steps as a rough default for stochastic sampling.
            noise_scale: scale to apply to gamma.
            s_t_min: don't apply s_churn below this noise level.
            s_t_max: don't apply s_churn above this noise level.
            temperature: scale to apply to aatype logits.
            top_p: don't tokens which fall outside this proportion of the total probability.
            disallow_aas: don't sample these token indices.
            sidechain_mode: whether to do all-atom sampling (False for backbone-only).
            skip_mpnn_proportion: proportion of timesteps from the start to skip running
                mini-MPNN.
            anneal_seq_resampling_rate: whether and how to decay the probability of
                running mini-MPNN. None, 'linear', or 'cosine'
            use_fullmpnn: use "full" ProteinMPNN at each step.
            use_fullmpnn_for_final: use "full" ProteinMPNN at the final step.
            use_reconstruction_guidance: use reconstruction guidance on the conditioning.
            use_classifier_free_guidance: use classifier-free guidance on the conditioning.
            guidance_scale: weight for reconstruction/classifier-free guidance.
            noise_schedule: specify the noise level timesteps for sampling.
            tqdm_pbar: progress bar in interactive contexts.
            return_last: return only the sampled structure and sequence.
            return_aux: return a dict of everything associated with the sampling run.
        """

        def ode_step(sigma_in, sigma_next, xt_in, x0_pred, gamma, guidance_in=None):
            if gamma > 0:
                t_hat = sigma_in + gamma * sigma_in
                sigma_delta = torch.sqrt(t_hat**2 - sigma_in**2)
                noisier_x = xt_in + utils.expand(
                    sigma_delta, xt_in
                ) * noise_scale * torch.randn_like(xt_in).to(xt_in)
                xt_in = noisier_x * utils.expand(seq_mask, noisier_x)
                sigma_in = t_hat

            mask = (sigma_in > 0).float()
            score = (xt_in - x0_pred) / utils.expand(sigma_in.clamp(min=1e-6), xt_in)
            score = score * utils.expand(mask, score)
            if use_reconstruction_guidance:
                guidance, guidance_mask = guidance_in
                guidance = guidance * guidance_mask[..., None]
                guidance_std = guidance[guidance_mask.bool()].var().sqrt()
                score_std = score[guidance_mask.bool()].var().sqrt()
                score = score + guidance * guidance_scale
            if use_classifier_free_guidance:
                # guidance_in is the unconditional x0 (x0_pred is the conditional x0)
                # guidance_scale = 1 + w from Ho paper
                # ==0: use only unconditional score; <1: interpolate the scores;
                # ==1: use only conditional score; >1: skew towards conditional score
                uncond_x0 = guidance_in
                uncond_score = (xt_in - uncond_x0) / utils.expand(
                    sigma_in.clamp(min=1e-6), xt_in
                )
                uncond_score = uncond_score * utils.expand(mask, uncond_score)
                score = guidance_scale * score + (1 - guidance_scale) * uncond_score
            step = score * step_scale * utils.expand(sigma_next - sigma_in, score)
            new_xt = xt_in + step
            return new_xt

        def sample_aatype(logprobs):
            # Top-p truncation
            probs = F.softmax(logprobs.clone(), dim=-1)
            sorted_prob, sorted_idxs = torch.sort(probs, descending=True)
            cumsum_prob = torch.cumsum(sorted_prob, dim=-1)
            sorted_indices_to_remove = cumsum_prob > top_p
            sorted_indices_to_remove[..., 0] = 0
            sorted_prob[sorted_indices_to_remove] = 0
            orig_probs = torch.scatter(
                torch.zeros_like(sorted_prob),
                dim=-1,
                index=sorted_idxs,
                src=sorted_prob,
            )

            # Apply temperature and disallowed AAs and sample
            assert temperature >= 0.0
            scaled_logits = orig_probs.clamp(min=1e-9).log() / (temperature + 1e-4)
            if disallow_aas:
                unwanted_mask = torch.zeros(scaled_logits.shape[-1]).to(scaled_logits)
                unwanted_mask[disallow_aas] = 1
                scaled_logits -= unwanted_mask * 1e10
            orig_probs = F.softmax(scaled_logits, dim=-1)
            categorical = torch.distributions.Categorical(probs=orig_probs)
            samp_aatype = categorical.sample()
            return samp_aatype

        def design_with_fullmpnn(batched_coords, seq_mask):
            seq_lens = seq_mask.sum(-1).long()
            designed_seqs = [
                evaluation.design_sequence(c[: seq_lens[i]], model=fullmpnn_model)[0]
                for i, c in enumerate(batched_coords)
            ]
            designed_aatypes, _ = utils.batched_seq_to_aatype_and_mask(
                designed_seqs, max_len=seq_mask.shape[-1]
            )
            return designed_aatypes

        # Initialize masks/features
        if seq_mask is None:  # Sample random lengths
            assert gt_aatype is None  # Don't condition on aatype without seq_mask
            seq_mask = self.make_seq_mask_for_sampling(
                n_samples=n_samples,
                min_len=min_len,
                max_len=max_len,
            )
        if residue_index is None:
            residue_index = torch.arange(seq_mask.shape[-1])
            residue_index = repeat(residue_index, "n -> b n", b=seq_mask.shape[0])
            residue_index = residue_index.to(seq_mask) * seq_mask
        if use_fullmpnn or use_fullmpnn_for_final:
            fullmpnn_model = protein_mpnn.get_mpnn_model(
                path_to_model_weights=self.config.train.home_dir
                + "/ProteinMPNN/vanilla_model_weights",
                device=self.device,
            )

        # Initialize noise schedule/parameters
        to_batch_size = lambda x: x * torch.ones(seq_mask.shape[0]).to(self.device)
        s_t_min = s_t_min * self.sigma_data
        s_t_max = s_t_max * self.sigma_data
        if noise_schedule is None:
            noise_schedule = self.sampling_noise_schedule_default
        sigma = noise_schedule(1)
        timesteps = torch.linspace(1, 0, n_steps + 1)

        # Set up conditioning/guidance information
        crop_cond_coords = None
        if gt_coords is None:
            coords_shape = seq_mask.shape + (self.n_atoms, 3)
            xt = torch.randn(*coords_shape).to(self.device) * sigma
            xt *= utils.expand(seq_mask, xt)
        else:
            assert gt_coords_traj is None
            noise_levels = [to_batch_size(noise_schedule(t)) for t in timesteps]
            gt_coords_traj = [
                diffusion.noise_coords(gt_coords, nl) for nl in noise_levels
            ]
            xt = gt_coords_traj[0]
            if gt_cond_atom_mask is not None:
                crop_cond_coords = gt_coords * gt_cond_atom_mask[..., None]
        gt_atom_mask = None
        if gt_aatype is not None:
            gt_atom_mask = utils.atom37_mask_from_aatype(gt_aatype, seq_mask)
        fake_logits = repeat(seq_mask, "b n -> b n t", t=self.n_tokens)
        s_hat = (sample_aatype(fake_logits) * seq_mask).long()

        # Initialize superposition for all-atom sampling
        if sidechain_mode:
            b, n = seq_mask.shape[:2]

            # Latest predicted x0 for sidechain superpositions
            atom73_state_0 = torch.zeros(b, n, 73, 3).to(xt)

            # Current state xt for sidechain superpositions (denoised to different levels)
            atom73_state_t = torch.randn(b, n, 73, 3).to(xt) * sigma

            # Noise level of xt
            sigma73_last = torch.ones(b, n, 73).to(xt) * sigma

            # Seqhat and mask used to choose sidechains for euler step (b, n)
            s_hat = (seq_mask * 7).long()
            mask37 = utils.atom37_mask_from_aatype(s_hat, seq_mask).bool()
            mask73 = utils.atom73_mask_from_aatype(s_hat, seq_mask).bool()
            begin_mpnn_step = int(n_steps * skip_mpnn_proportion)

        # Prepare to run sampling trajectory
        sigma = to_batch_size(sigma)
        x0 = None
        x0_prev = None
        x_self_cond = None
        s_logprobs = None
        s_self_cond = None
        if tqdm_pbar is None:
            tqdm_pbar = lambda x: x
        torch.set_grad_enabled(False)

        # *t_traj is the denoising trajectory; *0_traj is the evolution of predicted clean data
        # s0 are aatype probs of shape (b n t); s_hat are discrete aatype of shape (b n)
        xt_traj, x0_traj, st_traj, s0_traj = [], [], [], []

        # Sampling trajectory
        for i, t in tqdm_pbar(enumerate(iter(timesteps[1:]))):
            # Set up noise levels
            sigma_next = noise_schedule(t)
            if i == n_steps - 1:
                sigma_next *= 0
            gamma = (
                s_churn / n_steps
                if (sigma_next >= s_t_min and sigma_next <= s_t_max)
                else 0.0
            )
            sigma_next = to_batch_size(sigma_next)

            if sidechain_mode:
                # Fill in noise for masked positions since xt is initialized to zeros at each step
                dummy_fill_noise = torch.randn_like(xt) * utils.expand(sigma, xt)
                zero_atom_mask = utils.atom37_mask_from_aatype(s_hat, seq_mask)
                dummy_fill_mask = 1 - zero_atom_mask[..., None]
                xt = xt * zero_atom_mask[..., None] + dummy_fill_noise * dummy_fill_mask
            else:  # backbone only
                bb_seq = (seq_mask * residue_constants.restype_order["G"]).long()
                bb_atom_mask = utils.atom37_mask_from_aatype(bb_seq, seq_mask)
                xt *= bb_atom_mask[..., None]

            # Enable grad for reconstruction guidance
            if use_reconstruction_guidance:
                torch.set_grad_enabled(True)
                xt.requires_grad = True

            # Run denoising network
            run_mpnn = not sidechain_mode or i > begin_mpnn_step
            x0, s_logprobs, x_self_cond, s_self_cond = self.forward(
                noisy_coords=xt,
                noise_level=sigma,
                seq_mask=seq_mask,
                residue_index=residue_index,
                struct_self_cond=x_self_cond,
                struct_crop_cond=crop_cond_coords,
                seq_self_cond=s_self_cond,
                run_mpnn_model=run_mpnn,
            )

            # Compute additional stuff for guidance
            if use_reconstruction_guidance:
                loss = (x0 - gt_coords).pow(2).sum(-1)
                loss = loss * gt_cond_atom_mask
                loss = loss.sum() / gt_cond_atom_mask.sum().clamp(min=1)
                xt.retain_grad()
                loss.backward()
                guidance = xt.grad.clone()
                xt.grad *= 0
                torch.set_grad_enabled(False)
            if use_classifier_free_guidance:
                assert not use_reconstruction_guidance
                uncond_x0, _, _, _ = self.forward(
                    noisy_coords=xt,
                    noise_level=sigma,
                    seq_mask=seq_mask,
                    residue_index=residue_index,
                    struct_self_cond=x_self_cond,
                    seq_self_cond=s_self_cond,
                    run_mpnn_model=run_mpnn,
                )

            # Structure denoising step
            if not sidechain_mode:  # backbone
                if sigma[0] > 0:
                    xt = ode_step(sigma, sigma_next, xt, x0, gamma)
                else:
                    xt = x0
            else:  # allatom
                # Write x0 into atom73_state_0 for atoms corresponding to old seqhat
                atom73_state_0[mask73] = x0[mask37]

                # Determine sequence resampling probability
                if anneal_seq_resampling_rate is not None:
                    step_time = 1 - (i - begin_mpnn_step) / max(
                        1, n_steps - begin_mpnn_step
                    )
                    if anneal_seq_resampling_rate == "linear":
                        resampling_rate = step_time
                    elif anneal_seq_resampling_rate == "cosine":
                        k = 2
                        resampling_rate = (
                            1 + np.cos(2 * np.pi * (step_time - 0.5))
                        ) / k
                    resample_this_step = np.random.uniform() < resampling_rate

                # Resample sequence or design with full ProteinMPNN
                if i == n_steps - 1 and use_fullmpnn_for_final:
                    s_hat = design_with_fullmpnn(x0, seq_mask).to(x0.device)
                elif anneal_seq_resampling_rate is None or resample_this_step:
                    if run_mpnn and use_fullmpnn:
                        s_hat = design_with_fullmpnn(x0, seq_mask).to(x0.device)
                    else:
                        s_hat = sample_aatype(s_logprobs)

                # Overwrite s_hat with any conditioning information
                if (i + 1) / n_steps <= apply_cond_proportion:
                    if gt_cond_seq_mask is not None and gt_aatype is not None:
                        s_hat = (
                            1 - gt_cond_seq_mask
                        ) * s_hat + gt_cond_seq_mask * gt_aatype
                        s_hat = s_hat.long()

                # Set masks for collapsing superposition using new sequence
                mask37 = utils.atom37_mask_from_aatype(s_hat, seq_mask).bool()
                mask73 = utils.atom73_mask_from_aatype(s_hat, seq_mask).bool()

                # Determine prev noise levels for atoms corresponding to new sequence
                step_sigma_prev = (
                    torch.ones(*xt.shape[:-1]).to(xt) * sigma[..., None, None]
                )
                step_sigma_prev[mask37] = sigma73_last[mask73]  # b, n, 37
                step_sigma_next = sigma_next[..., None, None]  # b, 1, 1

                # Denoising step on atoms corresponding to new sequence
                b, n = mask37.shape[:2]
                step_xt = torch.zeros(b, n, 37, 3).to(xt)
                step_x0 = torch.zeros(b, n, 37, 3).to(xt)
                step_xt[mask37] = atom73_state_t[mask73]
                step_x0[mask37] = atom73_state_0[mask73]

                guidance_in = None
                if (i + 1) / n_steps <= apply_cond_proportion:
                    if use_reconstruction_guidance:
                        guidance_in = (guidance, mask37.float())
                    elif use_classifier_free_guidance:
                        guidance_in = uncond_x0

                step_xt = ode_step(
                    step_sigma_prev,
                    step_sigma_next,
                    step_xt,
                    step_x0,
                    gamma,
                    guidance_in=guidance_in,
                )
                xt = step_xt

                # Write new xt into atom73_state_t for atoms corresponding to new seqhat and update sigma_last
                atom73_state_t[mask73] = step_xt[mask37]
                sigma73_last[mask73] = step_sigma_next[0].item()

            # Replacement guidance if conditioning information provided
            if (i + 1) / n_steps <= apply_cond_proportion:
                if gt_coords_traj is not None:
                    if gt_cond_atom_mask is None:
                        xt = gt_coords_traj[i + 1]
                    else:
                        xt = (1 - gt_cond_atom_mask)[
                            ..., None
                        ] * xt + gt_cond_atom_mask[..., None] * gt_coords_traj[i + 1]

            sigma = sigma_next

            # Logging
            xt_scale = self.sigma_data / utils.expand(
                torch.sqrt(sigma_next**2 + self.sigma_data**2), xt
            )
            scaled_xt = xt * xt_scale
            xt_traj.append(scaled_xt.cpu())
            x0_traj.append(x0.cpu())
            st_traj.append(s_hat.cpu())
            s0_traj.append(s_logprobs.cpu())

        if return_last:
            return xt, s_hat, seq_mask
        elif return_aux:
            return {
                "x": xt,
                "s": s_hat,
                "seq_mask": seq_mask,
                "xt_traj": xt_traj,
                "x0_traj": x0_traj,
                "st_traj": st_traj,
                "s0_traj": s0_traj,
            }
        else:
            return xt_traj, x0_traj, st_traj, s0_traj, seq_mask

models.py
ProteinDesignLab/protpardelle
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Neural network modules. Many of these are adapted from open source modules.
"""
from typing import List, Sequence, Optional

from einops import rearrange, reduce, repeat
from einops.layers.torch import Rearrange
import numpy as np
from rotary_embedding_torch import RotaryEmbedding
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, EsmModel

from core import protein_mpnn
from core import residue_constants
from core import utils


########################################
# Adapted from https://github.com/ermongroup/ddim


def downsample(x):
    return nn.functional.avg_pool2d(x, 2, 2, ceil_mode=True)


def upsample_coords(x, shape):
    new_l, new_w = shape
    return nn.functional.interpolate(x, size=(new_l, new_w), mode="nearest")


########################################
# Adapted from https://github.com/aqlaboratory/openfold


def permute_final_dims(tensor: torch.Tensor, inds: List[int]):
    zero_index = -1 * len(inds)
    first_inds = list(range(len(tensor.shape[:zero_index])))
    return tensor.contiguous().permute(first_inds + [zero_index + i for i in inds])


def lddt(
    all_atom_pred_pos: torch.Tensor,
    all_atom_positions: torch.Tensor,
    all_atom_mask: torch.Tensor,
    cutoff: float = 15.0,
    eps: float = 1e-10,
    per_residue: bool = True,
) -> torch.Tensor:
    n = all_atom_mask.shape[-2]
    dmat_true = torch.sqrt(
        eps
        + torch.sum(
            (all_atom_positions[..., None, :] - all_atom_positions[..., None, :, :])
            ** 2,
            dim=-1,
        )
    )

    dmat_pred = torch.sqrt(
        eps
        + torch.sum(
            (all_atom_pred_pos[..., None, :] - all_atom_pred_pos[..., None, :, :]) ** 2,
            dim=-1,
        )
    )
    dists_to_score = (
        (dmat_true < cutoff)
        * all_atom_mask
        * permute_final_dims(all_atom_mask, (1, 0))
        * (1.0 - torch.eye(n, device=all_atom_mask.device))
    )

    dist_l1 = torch.abs(dmat_true - dmat_pred)

    score = (
        (dist_l1 < 0.5).type(dist_l1.dtype)
        + (dist_l1 < 1.0).type(dist_l1.dtype)
        + (dist_l1 < 2.0).type(dist_l1.dtype)
        + (dist_l1 < 4.0).type(dist_l1.dtype)
    )
    score = score * 0.25

    dims = (-1,) if per_residue else (-2, -1)
    norm = 1.0 / (eps + torch.sum(dists_to_score, dim=dims))
    score = norm * (eps + torch.sum(dists_to_score * score, dim=dims))

    return score


class RelativePositionalEncoding(nn.Module):
    def __init__(self, attn_dim=8, max_rel_idx=32):
        super().__init__()
        self.max_rel_idx = max_rel_idx
        self.n_rel_pos = 2 * self.max_rel_idx + 1
        self.linear = nn.Linear(self.n_rel_pos, attn_dim)

    def forward(self, residue_index):
        d_ij = residue_index[..., None] - residue_index[..., None, :]
        v_bins = torch.arange(self.n_rel_pos).to(d_ij.device) - self.max_rel_idx
        idxs = (d_ij[..., None] - v_bins[None, None]).abs().argmin(-1)
        p_ij = nn.functional.one_hot(idxs, num_classes=self.n_rel_pos)
        embeddings = self.linear(p_ij.float())
        return embeddings


########################################
# Adapted from https://github.com/NVlabs/edm


class Noise_Embedding(nn.Module):
    def __init__(self, num_channels, max_positions=10000, endpoint=False):
        super().__init__()
        self.num_channels = num_channels
        self.max_positions = max_positions
        self.endpoint = endpoint

    def forward(self, x):
        freqs = torch.arange(
            start=0, end=self.num_channels // 2, dtype=torch.float32, device=x.device
        )
        freqs = freqs / (self.num_channels // 2 - (1 if self.endpoint else 0))
        freqs = (1 / self.max_positions) ** freqs
        x = x.outer(freqs.to(x.dtype))
        x = torch.cat([x.cos(), x.sin()], dim=1)
        return x


########################################
# Adapted from github.com/lucidrains
# https://github.com/lucidrains/denoising-diffusion-pytorch
# https://github.com/lucidrains/recurrent-interface-network-pytorch


def exists(x):
    return x is not None


def default(val, d):
    if exists(val):
        return val
    return d() if callable(d) else d


def posemb_sincos_1d(patches, temperature=10000, residue_index=None):
    _, n, dim, device, dtype = *patches.shape, patches.device, patches.dtype

    n = torch.arange(n, device=device) if residue_index is None else residue_index
    assert (dim % 2) == 0, "feature dimension must be multiple of 2 for sincos emb"
    omega = torch.arange(dim // 2, device=device) / (dim // 2 - 1)
    omega = 1.0 / (temperature**omega)

    n = n[..., None] * omega
    pe = torch.cat((n.sin(), n.cos()), dim=-1)
    return pe.type(dtype)


class LayerNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(dim))
        self.register_buffer("beta", torch.zeros(dim))

    def forward(self, x):
        return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)


class NoiseConditioningBlock(nn.Module):
    def __init__(self, n_in_channel, n_out_channel):
        super().__init__()
        self.block = nn.Sequential(
            Noise_Embedding(n_in_channel),
            nn.Linear(n_in_channel, n_out_channel),
            nn.SiLU(),
            nn.Linear(n_out_channel, n_out_channel),
            Rearrange("b d -> b 1 d"),
        )

    def forward(self, noise_level):
        return self.block(noise_level)


class TimeCondResnetBlock(nn.Module):
    def __init__(
        self, nic, noc, cond_nc, conv_layer=nn.Conv2d, dropout=0.1, n_norm_in_groups=4
    ):
        super().__init__()
        self.block1 = nn.Sequential(
            nn.GroupNorm(num_groups=nic // n_norm_in_groups, num_channels=nic),
            nn.SiLU(),
            conv_layer(nic, noc, 3, 1, 1),
        )
        self.cond_proj = nn.Linear(cond_nc, noc * 2)
        self.mid_norm = nn.GroupNorm(num_groups=noc // 4, num_channels=noc)
        self.dropout = dropout if dropout is None else nn.Dropout(dropout)
        self.block2 = nn.Sequential(
            nn.GroupNorm(num_groups=noc // 4, num_channels=noc),
            nn.SiLU(),
            conv_layer(noc, noc, 3, 1, 1),
        )
        self.mismatch = False
        if nic != noc:
            self.mismatch = True
            self.conv_match = conv_layer(nic, noc, 1, 1, 0)

    def forward(self, x, time=None):
        h = self.block1(x)

        if time is not None:
            h = self.mid_norm(h)
            scale, shift = self.cond_proj(time).chunk(2, dim=-1)
            h = (h * (utils.expand(scale, h) + 1)) + utils.expand(shift, h)

        if self.dropout is not None:
            h = self.dropout(h)

        h = self.block2(h)

        if self.mismatch:
            x = self.conv_match(x)

        return x + h


class TimeCondAttention(nn.Module):
    def __init__(
        self,
        dim,
        dim_context=None,
        heads=4,
        dim_head=32,
        norm=False,
        norm_context=False,
        time_cond_dim=None,
        attn_bias_dim=None,
        rotary_embedding_module=None,
    ):
        super().__init__()
        hidden_dim = dim_head * heads
        dim_context = default(dim_context, dim)

        self.time_cond = None

        if exists(time_cond_dim):
            self.time_cond = nn.Sequential(nn.SiLU(), nn.Linear(time_cond_dim, dim * 2))

            nn.init.zeros_(self.time_cond[-1].weight)
            nn.init.zeros_(self.time_cond[-1].bias)

        self.scale = dim_head**-0.5
        self.heads = heads

        self.norm = LayerNorm(dim) if norm else nn.Identity()
        self.norm_context = LayerNorm(dim_context) if norm_context else nn.Identity()

        self.attn_bias_proj = None
        if attn_bias_dim is not None:
            self.attn_bias_proj = nn.Sequential(
                Rearrange("b a i j -> b i j a"),
                nn.Linear(attn_bias_dim, heads),
                Rearrange("b i j a -> b a i j"),
            )

        self.to_q = nn.Linear(dim, hidden_dim, bias=False)
        self.to_kv = nn.Linear(dim_context, hidden_dim * 2, bias=False)
        self.to_out = nn.Linear(hidden_dim, dim, bias=False)
        nn.init.zeros_(self.to_out.weight)

        self.use_rope = False
        if rotary_embedding_module is not None:
            self.use_rope = True
            self.rope = rotary_embedding_module

    def forward(self, x, context=None, time=None, attn_bias=None, seq_mask=None):
        # attn_bias is b, c, i, j
        h = self.heads
        has_context = exists(context)

        context = default(context, x)

        if x.shape[-1] != self.norm.gamma.shape[-1]:
            print(context.shape, x.shape, self.norm.gamma.shape)

        x = self.norm(x)

        if exists(time):
            scale, shift = self.time_cond(time).chunk(2, dim=-1)
            x = (x * (scale + 1)) + shift

        if has_context:
            context = self.norm_context(context)

        if seq_mask is not None:
            x = x * seq_mask[..., None]

        qkv = (self.to_q(x), *self.to_kv(context).chunk(2, dim=-1))
        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv)

        q = q * self.scale

        if self.use_rope:
            q = self.rope.rotate_queries_or_keys(q)
            k = self.rope.rotate_queries_or_keys(k)

        sim = torch.einsum("b h i d, b h j d -> b h i j", q, k)
        if attn_bias is not None:
            if self.attn_bias_proj is not None:
                attn_bias = self.attn_bias_proj(attn_bias)
            sim += attn_bias
        if seq_mask is not None:
            attn_mask = torch.einsum("b i, b j -> b i j", seq_mask, seq_mask)[:, None]
            sim -= (1 - attn_mask) * 1e6
        attn = sim.softmax(dim=-1)

        out = torch.einsum("b h i j, b h j d -> b h i d", attn, v)
        out = rearrange(out, "b h n d -> b n (h d)")
        out = self.to_out(out)
        if seq_mask is not None:
            out = out * seq_mask[..., None]
        return out


class TimeCondFeedForward(nn.Module):
    def __init__(self, dim, mult=4, dim_out=None, time_cond_dim=None, dropout=0.1):
        super().__init__()
        if dim_out is None:
            dim_out = dim
        self.norm = LayerNorm(dim)

        self.time_cond = None
        self.dropout = None
        inner_dim = int(dim * mult)

        if exists(time_cond_dim):
            self.time_cond = nn.Sequential(
                nn.SiLU(),
                nn.Linear(time_cond_dim, inner_dim * 2),
            )

            nn.init.zeros_(self.time_cond[-1].weight)
            nn.init.zeros_(self.time_cond[-1].bias)

        self.linear_in = nn.Linear(dim, inner_dim)
        self.nonlinearity = nn.SiLU()
        if dropout is not None:
            self.dropout = nn.Dropout(dropout)
        self.linear_out = nn.Linear(inner_dim, dim_out)
        nn.init.zeros_(self.linear_out.weight)
        nn.init.zeros_(self.linear_out.bias)

    def forward(self, x, time=None):
        x = self.norm(x)
        x = self.linear_in(x)
        x = self.nonlinearity(x)

        if exists(time):
            scale, shift = self.time_cond(time).chunk(2, dim=-1)
            x = (x * (scale + 1)) + shift

        if exists(self.dropout):
            x = self.dropout(x)

        return self.linear_out(x)


class TimeCondTransformer(nn.Module):
    def __init__(
        self,
        dim,
        depth,
        heads,
        dim_head,
        time_cond_dim,
        attn_bias_dim=None,
        mlp_inner_dim_mult=4,
        position_embedding_type: str = "rotary",
    ):
        super().__init__()

        self.rope = None
        self.pos_emb_type = position_embedding_type
        if position_embedding_type == "rotary":
            self.rope = RotaryEmbedding(dim=32)
        elif position_embedding_type == "relative":
            self.relpos = nn.Sequential(
                RelativePositionalEncoding(attn_dim=heads),
                Rearrange("b i j d -> b d i j"),
            )

        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        TimeCondAttention(
                            dim,
                            heads=heads,
                            dim_head=dim_head,
                            norm=True,
                            time_cond_dim=time_cond_dim,
                            attn_bias_dim=attn_bias_dim,
                            rotary_embedding_module=self.rope,
                        ),
                        TimeCondFeedForward(
                            dim, mlp_inner_dim_mult, time_cond_dim=time_cond_dim
                        ),
                    ]
                )
            )

    def forward(
        self,
        x,
        time=None,
        attn_bias=None,
        context=None,
        seq_mask=None,
        residue_index=None,
    ):
        if self.pos_emb_type == "absolute":
            pos_emb = posemb_sincos_1d(x)
            x = x + pos_emb
        elif self.pos_emb_type == "absolute_residx":
            assert residue_index is not None
            pos_emb = posemb_sincos_1d(x, residue_index=residue_index)
            x = x + pos_emb
        elif self.pos_emb_type == "relative":
            assert residue_index is not None
            pos_emb = self.relpos(residue_index)
            attn_bias = pos_emb if attn_bias is None else attn_bias + pos_emb
        if seq_mask is not None:
            x = x * seq_mask[..., None]

        for i, (attn, ff) in enumerate(self.layers):
            x = x + attn(
                x, context=context, time=time, attn_bias=attn_bias, seq_mask=seq_mask
            )
            x = x + ff(x, time=time)
            if seq_mask is not None:
                x = x * seq_mask[..., None]

        return x


class TimeCondUViT(nn.Module):
    def __init__(
        self,
        *,
        seq_len: int,
        dim: int,
        patch_size: int = 1,
        depth: int = 6,
        heads: int = 8,
        dim_head: int = 32,
        n_filt_per_layer: List[int] = [],
        n_blocks_per_layer: int = 2,
        n_atoms: int = 37,
        channels_per_atom: int = 6,
        attn_bias_dim: int = None,
        time_cond_dim: int = None,
        conv_skip_connection: bool = False,
        position_embedding_type: str = "rotary",
    ):
        super().__init__()

        # Initialize configuration params
        if time_cond_dim is None:
            time_cond_dim = dim * 4
        self.position_embedding_type = position_embedding_type
        channels = channels_per_atom
        self.n_conv_layers = n_conv_layers = len(n_filt_per_layer)
        if n_conv_layers > 0:
            post_conv_filt = n_filt_per_layer[-1]
        self.conv_skip_connection = conv_skip_connection and n_conv_layers == 1
        transformer_seq_len = seq_len // (2**n_conv_layers)
        assert transformer_seq_len % patch_size == 0
        num_patches = transformer_seq_len // patch_size
        dim_a = post_conv_atom_dim = max(1, n_atoms // (2 ** (n_conv_layers - 1)))
        if n_conv_layers == 0:
            patch_dim = patch_size * n_atoms * channels_per_atom
            patch_dim_out = patch_size * n_atoms * 3
            dim_a = n_atoms
        elif conv_skip_connection and n_conv_layers == 1:
            patch_dim = patch_size * (channels + post_conv_filt) * post_conv_atom_dim
            patch_dim_out = patch_size * post_conv_filt * post_conv_atom_dim
        elif n_conv_layers > 0:
            patch_dim = patch_dim_out = patch_size * post_conv_filt * post_conv_atom_dim

        # Make downsampling conv
        # Downsamples n-1 times where n is n_conv_layers
        down_conv = []
        block_in = channels
        for i, nf in enumerate(n_filt_per_layer):
            block_out = nf
            layer = []
            for j in range(n_blocks_per_layer):
                n_groups = 2 if i == 0 and j == 0 else 4
                layer.append(
                    TimeCondResnetBlock(
                        block_in, block_out, time_cond_dim, n_norm_in_groups=n_groups
                    )
                )
                block_in = block_out
            down_conv.append(nn.ModuleList(layer))
        self.down_conv = nn.ModuleList(down_conv)

        # Make transformer
        self.to_patch_embedding = nn.Sequential(
            Rearrange("b c (n p) a -> b n (p c a)", p=patch_size),
            nn.Linear(patch_dim, dim),
            LayerNorm(dim),
        )
        self.transformer = TimeCondTransformer(
            dim,
            depth,
            heads,
            dim_head,
            time_cond_dim,
            attn_bias_dim=attn_bias_dim,
            position_embedding_type=position_embedding_type,
        )
        self.from_patch = nn.Sequential(
            LayerNorm(dim),
            nn.Linear(dim, patch_dim_out),
            Rearrange("b n (p c a) -> b c (n p) a", p=patch_size, a=dim_a),
        )
        nn.init.zeros_(self.from_patch[-2].weight)
        nn.init.zeros_(self.from_patch[-2].bias)

        # Make upsampling conv
        up_conv = []
        for i, nf in enumerate(reversed(n_filt_per_layer)):
            skip_in = nf
            block_out = nf
            layer = []
            for j in range(n_blocks_per_layer):
                layer.append(
                    TimeCondResnetBlock(block_in + skip_in, block_out, time_cond_dim)
                )
                block_in = block_out
            up_conv.append(nn.ModuleList(layer))
        self.up_conv = nn.ModuleList(up_conv)

        # Conv out
        if n_conv_layers > 0:
            self.conv_out = nn.Sequential(
                nn.GroupNorm(num_groups=block_out // 4, num_channels=block_out),
                nn.SiLU(),
                nn.Conv2d(block_out, channels // 2, 3, 1, 1),
            )

    def forward(
        self, coords, time_cond, pair_bias=None, seq_mask=None, residue_index=None
    ):
        if self.n_conv_layers > 0:  # pad up to even dims
            coords = F.pad(coords, (0, 0, 0, 0, 0, 1, 0, 0))

        x = rearr_coords = rearrange(coords, "b n a c -> b c n a")
        hiddens = []
        for i, layer in enumerate(self.down_conv):
            for block in layer:
                x = block(x, time=time_cond)
                hiddens.append(x)
            if i != self.n_conv_layers - 1:
                x = downsample(x)

        if self.conv_skip_connection:
            x = torch.cat([x, rearr_coords], 1)

        x = self.to_patch_embedding(x)
        # if self.position_embedding_type == 'absolute':
        #     pos_emb = posemb_sincos_1d(x)
        #     x = x + pos_emb
        if seq_mask is not None and x.shape[1] == seq_mask.shape[1]:
            x *= seq_mask[..., None]
        x = self.transformer(
            x,
            time=time_cond,
            attn_bias=pair_bias,
            seq_mask=seq_mask,
            residue_index=residue_index,
        )
        x = self.from_patch(x)

        for i, layer in enumerate(self.up_conv):
            for block in layer:
                x = torch.cat([x, hiddens.pop()], 1)
                x = block(x, time=time_cond)
            if i != self.n_conv_layers - 1:
                x = upsample_coords(x, hiddens[-1].shape[2:])

        if self.n_conv_layers > 0:
            x = self.conv_out(x)
            x = x[..., :-1, :]  # drop even-dims padding

        x = rearrange(x, "b c n a -> b n a c")
        return x


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


class LinearWarmupCosineDecay(torch.optim.lr_scheduler._LRScheduler):
    def __init__(
        self,
        optimizer,
        max_lr,
        warmup_steps=1000,
        decay_steps=int(1e6),
        min_lr=1e-6,
        **kwargs,
    ):
        self.max_lr = max_lr
        self.min_lr = min_lr
        self.warmup_steps = warmup_steps
        self.decay_steps = decay_steps
        self.total_steps = warmup_steps + decay_steps
        super(LinearWarmupCosineDecay, self).__init__(optimizer, **kwargs)

    def get_lr(self):
        # TODO double check for off-by-one errors
        if self.last_epoch < self.warmup_steps:
            curr_lr = self.last_epoch / self.warmup_steps * self.max_lr
            return [curr_lr for group in self.optimizer.param_groups]
        elif self.last_epoch < self.total_steps:
            time = (self.last_epoch - self.warmup_steps) / self.decay_steps * np.pi
            curr_lr = self.min_lr + (self.max_lr - self.min_lr) * 0.5 * (
                1 + np.cos(time)
            )
            return [curr_lr for group in self.optimizer.param_groups]
        else:
            return [self.min_lr for group in self.optimizer.param_groups]


class NoiseConditionalProteinMPNN(nn.Module):
    def __init__(
        self,
        n_channel=128,
        n_layers=3,
        n_neighbors=32,
        time_cond_dim=None,
        vocab_size=21,
        input_S_is_embeddings=False,
    ):
        super().__init__()
        self.n_channel = n_channel
        self.n_layers = n_layers
        self.n_neighbors = n_neighbors
        self.time_cond_dim = time_cond_dim
        self.vocab_size = vocab_size
        self.bb_idxs_if_atom37 = [
            residue_constants.atom_order[a] for a in ["N", "CA", "C", "O"]
        ]

        self.mpnn = protein_mpnn.ProteinMPNN(
            num_letters=vocab_size,
            node_features=n_channel,
            edge_features=n_channel,
            hidden_dim=n_channel,
            num_encoder_layers=n_layers,
            num_decoder_layers=n_layers,
            vocab=vocab_size,
            k_neighbors=n_neighbors,
            augment_eps=0.0,
            dropout=0.1,
            ca_only=False,
            time_cond_dim=time_cond_dim,
            input_S_is_embeddings=input_S_is_embeddings,
        )

    def forward(
        self, denoised_coords, noisy_aatype, seq_mask, residue_index, time_cond
    ):
        if denoised_coords.shape[-2] == 37:
            denoised_coords = denoised_coords[:, :, self.bb_idxs_if_atom37]

        node_embs, encoder_embs = self.mpnn(
            X=denoised_coords,
            S=noisy_aatype,
            mask=seq_mask,
            chain_M=seq_mask,
            residue_idx=residue_index,
            chain_encoding_all=seq_mask,
            randn=None,
            use_input_decoding_order=False,
            decoding_order=None,
            causal_mask=False,
            time_cond=time_cond,
            return_node_embs=True,
        )
        return node_embs, encoder_embs

modules.py
ProteinDesignLab/protpardelle
"""
Protpardelle Pymol Extension 

load protpardelle_pymol
(protpardelle_setlink link ) (if you want to run a local server)

Conditional:
fetch 1pga
select some residues
generate 5 samples resampling the residues in sele
protpardelle 1pga, sele, 5  

unconditional: 
Sample sequences between length 50 and 60 with step size 5 and generate 1 sample per sampled length. Use the allatom model
protpardelle_uncond 50, 60, 5, 1, allatom
"""

from pymol import cmd

import os
import json
import time
import threading

try:
    from gradio_client import Client
except ImportError:
    print("gradio_client not installed, trying install:")
    import pip 
    pip.main(['install', 'gradio_client'])
    from gradio_client import Client


if os.environ.get("GRADIO_LOCAL") != None:
    public_link = "http://127.0.0.1:7860"
else:
    public_link = "ProteinDesignLab/protpardelle"

    


def thread_protpardelle(input_pdb,
    resample_idxs,
    modeltype,
    mode,
    minlen=50,
    maxlen= 60,
    steplen = 2,
    per_len = 2):
    client = Client(public_link)

    job = client.submit(
				input_pdb,	# str in 'PDB Content' Textbox component
				modeltype,	# str in 'Choose a Mode' Radio component
				f'"{resample_idxs}"',	# str in 'Resampled Idxs' Textbox component
				mode,	# str (Option from: ['backbone', 'allatom'])
				minlen,	# int | float (numeric value between 2 and 200) minlen
				maxlen,	# int | float (numeric value between 3 and 200) in 'maxlen' Slider component
				steplen,	# int | float (numeric value between 1 and 50) in 'steplen' Slider component
				per_len,	# int | float (numeric value between 1 and 200) in 'perlen' Slider component
				api_name="/protpardelle"
)
    #start time
    start = time.time()

    while (job.done() == False):
        status =  job.status()
        elapsed = time.time()-start 
        # format as hh:mm:ss
        elapsed = time.strftime("%H:%M:%S", time.gmtime(elapsed))

        print(f"\r protpardelle running since {elapsed}", end="")
        time.sleep(1)
    results = job.result()

    # load each result into pymol 
    results = json.loads(results)

    for (name,pdb_content) in results:
        print(name)
        cmd.read_pdbstr(pdb_content, os.path.basename(name))


def query_protpardelle(
    name_of_input: str,
    selection_resample_idxs: str="",
    per_len: int = 2,
    mode: str="allatom",
):
    """
    AUTHOR
    Simon Duerr
    https://twitter.com/simonduerr
    DESCRIPTION
    Run Protpardelle
    USAGE
    protpardelle name_of_input, selection_resampled_idx, modeltype, mode, per_len
    PARAMETERS
    name_of_input = string: name of input object
    selection_resampled_idx = string: selection of resampled protein residues
    per_len = int: per_len (default: 2)
    mode = string: mode (default: 'allatom')
    """
    if name_of_input != "":
        input_pdb = cmd.get_pdbstr(name_of_input)

        all_aa = cmd.index(name_of_input+" and name CA")
        idx = cmd.index(selection_resample_idxs+" and name CA")

        #map to zero indexed values
        aa_mapping = {aa[1]:i for i,aa in enumerate(all_aa)}

        idx = ",".join([str(aa_mapping[aa[1]]) for aa in idx])

        print("resampling", idx , "(zero indexed) from", name_of_input)

    t = threading.Thread(target=thread_protpardelle,
                         args=(input_pdb, idx, "conditional",mode ),
                         kwargs={'per_len':per_len},
                         daemon=True)
    t.start()

def query_protpardelle_uncond(
    
    minlen: int = 50,
    maxlen: int = 60,
    steplen: int = 2,
    per_len: int = 2,
    mode: str="allatom",
):
    """
    AUTHOR
    Simon Duerr
    https://twitter.com/simonduerr
    DESCRIPTION
    Run Protpardelle
    USAGE
    protpardelle_uncond minlen, maxlen, steplen, per_len,mode
    PARAMETERS
    minlen = int: minlen
    maxlen = int: maxlen
    steplen = int: steplen
    per_len = int: per_len
    mode = string: mode (default: 'allatom')
    """

    modeltype = "unconditional"
    idx = None
    input_pdb = None

    t = threading.Thread(target=thread_protpardelle,
                         args=(input_pdb, idx, modeltype, mode),
                         kwargs={'minlen':minlen, 'maxlen':maxlen, 'steplen':steplen,'per_len':per_len},
                         daemon=True)
    t.start()



def setprotpardellelink(link:str):
    """
    AUTHOR
    Simon Duerr
    https://twitter.com/simonduerr
    DESCRIPTION
    Set a public link to use a locally hosted version of this space
    USAGE
    protpardelle_setlink link_or_username/spacename
    """

    global public_link
    try:
        client = Client(link)
    except:
        print("could not connect to:", public_link)
    
    public_link = link




cmd.extend("protpardelle_setlink", setprotpardellelink)

cmd.extend("protpardelle", query_protpardelle)

cmd.extend("protpardelle_uncond", query_protpardelle_uncond)


protpardelle_pymol.py
ProteinDesignLab/protpardelle
"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Configs and convenience functions for wrapping the model sample() function.
"""
import argparse
import time
from typing import Optional, Tuple

import torch
from torchtyping import TensorType

from core import residue_constants
from core import utils
import diffusion


def default_backbone_sampling_config():
    config = argparse.Namespace(
        n_steps=500,
        s_churn=200,
        step_scale=1.2,
        sidechain_mode=False,
        noise_schedule=lambda t: diffusion.noise_schedule(t, s_max=80, s_min=0.001),
    )
    return config


def default_allatom_sampling_config():
    noise_schedule = lambda t: diffusion.noise_schedule(t, s_max=80, s_min=0.001)
    stage2 = argparse.Namespace(
        apply_cond_proportion=1.0,
        n_steps=200,
        s_churn=100,
        step_scale=1.2,
        sidechain_mode=True,
        skip_mpnn_proportion=1.0,
        noise_schedule=noise_schedule,
    )
    config = argparse.Namespace(
        n_steps=500,
        s_churn=200,
        step_scale=1.2,
        sidechain_mode=True,
        skip_mpnn_proportion=0.6,
        use_fullmpnn=False,
        use_fullmpnn_for_final=True,
        anneal_seq_resampling_rate="linear",
        noise_schedule=noise_schedule,
        stage_2=stage2,
    )
    return config


def draw_backbone_samples(
    model: torch.nn.Module,
    seq_mask: TensorType["b n", float] = None,
    n_samples: int = None,
    sample_length_range: Tuple[int] = (50, 512),
    pdb_save_path: Optional[str] = None,
    return_aux: bool = False,
    return_sampling_runtime: bool = False,
    **sampling_kwargs,
):
    device = model.device
    if seq_mask is None:
        assert n_samples is not None
        seq_mask = model.make_seq_mask_for_sampling(
            n_samples=n_samples,
            min_len=sample_length_range[0],
            max_len=sample_length_range[1],
        )

    start = time.time()
    aux = model.sample(
        seq_mask=seq_mask, return_last=False, return_aux=True, **sampling_kwargs
    )
    aux["runtime"] = time.time() - start
    seq_lens = seq_mask.sum(-1).long()
    cropped_samp_coords = [
        s[: seq_lens[i], model.bb_idxs] for i, s in enumerate(aux["xt_traj"][-1])
    ]

    if pdb_save_path is not None:
        gly_aatype = (seq_mask * residue_constants.restype_order["G"]).long()
        trimmed_aatype = [a[: seq_lens[i]] for i, a in enumerate(gly_aatype)]
        atom_mask = utils.atom37_mask_from_aatype(gly_aatype, seq_mask).cpu()
        for i in range(len(cropped_samp_coords)):
            utils.write_coords_to_pdb(
                cropped_samp_coords[i],
                f"{pdb_save_path}{i}.pdb",
                batched=False,
                aatype=trimmed_aatype[i],
                atom_mask=atom_mask[i],
            )

    if return_aux:
        return aux
    else:
        if return_sampling_runtime:
            return cropped_samp_coords, seq_mask, aux["runtime"]
        else:
            return cropped_samp_coords, seq_mask


def draw_allatom_samples(
    model: torch.nn.Module,
    seq_mask: TensorType["b n", float] = None,
    n_samples: int = None,
    sample_length_range: Tuple[int] = (50, 512),
    two_stage_sampling: bool = True,
    pdb_save_path: Optional[str] = None,
    return_aux: bool = False,
    return_sampling_runtime: bool = False,
    **sampling_kwargs,
):
    """Implement the default 2-stage all-atom sampling routine."""

    def save_allatom_samples(aux, path):
        seq_lens = aux["seq_mask"].sum(-1).long()
        cropped_samp_coords = [
            c[: seq_lens[i]] for i, c in enumerate(aux["xt_traj"][-1])
        ]
        cropped_samp_aatypes = [
            s[: seq_lens[i]] for i, s in enumerate(aux["st_traj"][-1])
        ]
        samp_atom_mask = utils.atom37_mask_from_aatype(
            aux["st_traj"][-1].to(device), seq_mask
        )
        samp_atom_mask = [m[: seq_lens[i]] for i, m in enumerate(samp_atom_mask)]
        for i, c in enumerate(cropped_samp_coords):
            utils.write_coords_to_pdb(
                c,
                f"{path}{i}.pdb",
                batched=False,
                aatype=cropped_samp_aatypes[i],
                atom_mask=samp_atom_mask[i],
                conect=True,
            )

    device = model.device
    if seq_mask is None:
        assert n_samples is not None
        seq_mask = model.make_seq_mask_for_sampling(
            n_samples=n_samples,
            min_len=sample_length_range[0],
            max_len=sample_length_range[1],
        )
    sampling_runtime = 0.0

    # Stage 1 sampling
    start = time.time()
    if "stage_2" in sampling_kwargs:
        stage_2_kwargs = vars(sampling_kwargs.pop("stage_2"))
    aux = model.sample(
        seq_mask=seq_mask,
        return_last=False,
        return_aux=True,
        **sampling_kwargs,
    )
    sampling_runtime = time.time() - start
    if pdb_save_path is not None and two_stage_sampling:
        save_allatom_samples(aux, pdb_save_path + "_init")

    # Stage 2 sampling (sidechain refinement only)
    if two_stage_sampling:
        samp_seq = aux["st_traj"][-1]
        samp_coords = aux["xt_traj"][-1]
        cond_atom_mask = utils.atom37_mask_from_aatype((seq_mask * 7).long(), seq_mask)
        aux = {f"stage1_{k}": v for k, v in aux.items()}
        start = time.time()
        stage2_aux = model.sample(
            gt_cond_atom_mask=cond_atom_mask.to(device),  # condition on backbone
            gt_cond_seq_mask=seq_mask.to(device),
            gt_coords=samp_coords.to(device),
            gt_aatype=samp_seq.to(device),
            seq_mask=seq_mask,
            return_last=False,
            return_aux=True,
            **stage_2_kwargs,
        )
        sampling_runtime += time.time() - start
        aux = {**aux, **stage2_aux}
    if pdb_save_path is not None:
        save_allatom_samples(aux, pdb_save_path + "_samp")
    aux["runtime"] = sampling_runtime

    # Process outputs, crop to correct length
    if return_aux:
        return aux
    else:
        xt_traj = aux["xt_traj"]
        st_traj = aux["st_traj"]
        seq_mask = aux["seq_mask"]
        seq_lens = seq_mask.sum(-1).long()
        cropped_samp_coords = [c[: seq_lens[i]] for i, c in enumerate(xt_traj[-1])]
        cropped_samp_aatypes = [s[: seq_lens[i]] for i, s in enumerate(st_traj[-1])]
        samp_atom_mask = utils.atom37_mask_from_aatype(st_traj[-1].to(device), seq_mask)
        samp_atom_mask = [m[: seq_lens[i]] for i, m in enumerate(samp_atom_mask)]
        orig_xt_traj = aux["stage1_xt_traj"]
        stage1_coords = [c[: seq_lens[i]] for i, c in enumerate(orig_xt_traj[-1])]
        ret = (
            cropped_samp_coords,
            cropped_samp_aatypes,
            samp_atom_mask,
            stage1_coords,
            seq_mask,
        )
        if return_sampling_runtime:
            ret = ret + (sampling_runtime,)
        return ret

sampling.py
ProteinDesignLab/protpardelle