Upload 2 files

ProteinMPNN-main/protein_mpnn_run.py

@@ -0,0 +1,522 @@
+import argparse
+import os.path
+
+
+def p_m_r(ca_only, path_to_model_weights, model_name, seed, save_score, save_probs, score_only, conditional_probs_only,
+         conditional_probs_only_backbone
+         , unconditional_probs_only, backbone_noise, num_seq_per_target, batch_size, max_length, sampling_temp, out_folder,
+         pdb_path, pdb_path_chains, jsonl_path, chain_id_jsonl, fixed_positions_jsonl, omit_AAs, bias_AA_jsonl, bias_by_res_jsonl
+         , omit_AA_jsonl, pssm_jsonl, pssm_multi, pssm_threshold, pssm_log_odds_flag, pssm_bias_flag, tied_positions_jsonl):
+    seed = int(seed)
+    save_score = int(save_score)
+    save_probs = int(save_probs)
+    score_only = int(score_only)
+    conditional_probs_only = int(conditional_probs_only)
+    conditional_probs_only_backbone = int(conditional_probs_only_backbone)
+    unconditional_probs_only = int(unconditional_probs_only)
+    num_seq_per_target = int(num_seq_per_target)
+    batch_size = int(batch_size)
+    max_length = int(max_length)
+    pssm_log_odds_flag = int(pssm_log_odds_flag)
+    pssm_bias_flag = int(pssm_bias_flag)
+    import json, time, os, sys, glob
+    import shutil
+    import warnings
+    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 copy
+    import torch.nn as nn
+    import torch.nn.functional as F
+    import random
+    import os.path
+    import subprocess
+
+    from protein_mpnn_utils import loss_nll, loss_smoothed, gather_edges, gather_nodes, gather_nodes_t, \
+        cat_neighbors_nodes, _scores, _S_to_seq, tied_featurize, parse_PDB
+    from protein_mpnn_utils import StructureDataset, StructureDatasetPDB, ProteinMPNN
+
+    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)
+
+    hidden_dim = 128
+    num_layers = 3
+
+    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("\\")
+        if ca_only:
+            model_folder_path = file_path[:k] + '/ca_model_weights/'
+        else:
+            model_folder_path = file_path[:k] + '/vanilla_model_weights/'
+
+    checkpoint_path = model_folder_path + f'{model_name}.pt'
+    folder_for_outputs = out_folder
+
+    NUM_BATCHES = num_seq_per_target // batch_size
+    BATCH_COPIES = batch_size
+    temperatures = [float(item) for item in sampling_temp.split()]
+    omit_AAs_list = omit_AAs
+    alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
+
+    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")
+    # os.path.isfile()：判断某一对象(需提供绝对路径)是否为文件
+    # 改
+    if chain_id_jsonl:
+        if os.path.isfile(chain_id_jsonl.name):
+            with open(chain_id_jsonl.name, '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
+        print(40 * '-')
+        print('chain_id_jsonl is NOT loaded')
+    if fixed_positions_jsonl:
+        if os.path.isfile(fixed_positions_jsonl.name):
+            with open(fixed_positions_jsonl.name, 'r') as json_file:
+                json_list = list(json_file)
+            for json_str in json_list:
+                fixed_positions_dict = json.loads(json_str)
+    else:
+        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:
+        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:
+        print(40 * '-')
+        print('omit_AA_jsonl is NOT loaded')
+        omit_AA_dict = None
+    if bias_AA_jsonl:
+        if os.path.isfile(bias_AA_jsonl.name):
+            with open(bias_AA_jsonl.name, 'r') as json_file:
+                json_list = list(json_file)
+            for json_str in json_list:
+                bias_AA_dict = json.loads(json_str)
+    else:
+        print(40 * '-')
+        print('bias_AA_jsonl is NOT loaded')
+        bias_AA_dict = None
+    if tied_positions_jsonl:
+        if os.path.isfile(tied_positions_jsonl.name):
+            with open(tied_positions_jsonl.name, 'r') as json_file:
+                json_list = list(json_file)
+            for json_str in json_list:
+                tied_positions_dict = json.loads(json_str)
+    else:
+        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)
+        print('bias by residue dictionary is loaded')
+    else:
+        print(40 * '-')
+        print('bias by residue dictionary is not loaded, or not provided')
+        bias_by_res_dict = None
+
+    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.name, 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.name, truncate=None, max_length=max_length)
+
+    print(40 * '-')
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    print('Number of edges:', checkpoint['num_edges'])
+    noise_level_print = checkpoint['noise_level']
+    print(f'Training noise level: {noise_level_print}A')
+    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()
+
+    # Build paths for experiment
+    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 save_score:
+        if not os.path.exists(base_folder + 'scores'):
+            os.makedirs(base_folder + 'scores')
+
+    if score_only:
+        if not os.path.exists(base_folder + 'score_only'):
+            os.makedirs(base_folder + 'score_only')
+
+    if conditional_probs_only:
+        if not os.path.exists(base_folder + 'conditional_probs_only'):
+            os.makedirs(base_folder + 'conditional_probs_only')
+
+    if unconditional_probs_only:
+        if not os.path.exists(base_folder + 'unconditional_probs_only'):
+            os.makedirs(base_folder + 'unconditional_probs_only')
+
+    if 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
+    with torch.no_grad():
+        test_sum, test_weights = 0., 0.
+        # print('Generating sequences...')
+        # 改
+        results = []
+        # enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标，一般用在 for 循环当中。
+        for ix, protein in enumerate(dataset_valid):
+            score_list = []
+            global_score_list = []
+            all_probs_list = []
+            all_log_probs_list = []
+            S_sample_list = []
+            # deepcopy复制
+            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 score_only:
+                structure_sequence_score_file = base_folder + '/score_only/' + batch_clones[0]['name'] + '.npz'
+                native_score_list = []
+                global_native_score_list = []
+                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]
+                np.savez(structure_sequence_score_file, score=native_score, global_score=global_native_score)
+                print(
+                    f'Score for {name_}, mean: {ns_mean_print}, std: {ns_std_print}, sample size: {ns_sample_size},  Global Score for {name_}, mean: {global_ns_mean_print}, std: {global_ns_std_print}, sample size: {ns_sample_size}')
+                results.append(structure_sequence_score_file)
+            elif conditional_probs_only:
+                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,
+                                                                    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 unconditional_probs_only:
+                print(f'Calculating sequence unconditional probabilities for {name_}')
+                # 改
+                unconditional_probs_only_file = base_folder + '/unconditional_probs_only/' + batch_clones[0]['name'] + '.npz'
+                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)
+                results.append(unconditional_probs_only_file)
+            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
+                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'
+                print(f'Generating sequences for: {name_}')
+                t0 = time.time()
+                with open(ali_file, 'w') as f:
+                    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])
+                                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).decode().strip()
+                                    except subprocess.CalledProcessError:
+                                        commit_str = 'unknown'
+                                    if ca_only:
+                                        print_model_name = 'CA_model_name'
+                                    else:
+                                        print_model_name = 'model_name'
+                                    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
+                                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
+                results.append(ali_file)
+                if save_score:
+                    np.savez(score_file, score=np.array(score_list, np.float32),
+                             global_score=np.array(global_score_list, np.float32))
+                if 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]
+                print(f'{num_seqs} sequences of length {total_length} generated in {dt} seconds')
+
+        return results
+
+
+# if __name__ == "__main__":
+#     argparser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+#
+#     argparser.add_argument("--ca_only", action="store_true", default=False,
+#                            help="Parse CA-only structures and use CA-only models (default: false)")
+#     argparser.add_argument("--path_to_model_weights", type=str, default="", help="Path to model weights folder;")
+#     argparser.add_argument("--model_name", type=str, default="v_48_020",
+#                            help="ProteinMPNN model name: v_48_002, v_48_010, v_48_020, v_48_030; v_48_010=version with 48 edges 0.10A noise")
+#
+#     argparser.add_argument("--seed", type=int, default=0, help="If set to 0 then a random seed will be picked;")
+#
+#     argparser.add_argument("--save_score", type=int, default=0,
+#                            help="0 for False, 1 for True; save score=-log_prob to npy files")
+#     argparser.add_argument("--save_probs", type=int, default=0,
+#                            help="0 for False, 1 for True; save MPNN predicted probabilites per position")
+#
+#     argparser.add_argument("--score_only", type=int, default=0,
+#                            help="0 for False, 1 for True; score input backbone-sequence pairs")
+#
+#     argparser.add_argument("--conditional_probs_only", type=int, default=0,
+#                            help="0 for False, 1 for True; output conditional probabilities p(s_i given the rest of the sequence and backbone)")
+#     argparser.add_argument("--conditional_probs_only_backbone", type=int, default=0,
+#                            help="0 for False, 1 for True; if true output conditional probabilities p(s_i given backbone)")
+#     argparser.add_argument("--unconditional_probs_only", type=int, default=0,
+#                            help="0 for False, 1 for True; output unconditional probabilities p(s_i given backbone) in one forward pass")
+#
+#     argparser.add_argument("--backbone_noise", type=float, default=0.00,
+#                            help="Standard deviation of Gaussian noise to add to backbone atoms")
+#     argparser.add_argument("--num_seq_per_target", type=int, default=1,
+#                            help="Number of sequences to generate per target")
+#     argparser.add_argument("--batch_size", type=int, default=1,
+#                            help="Batch size; can set higher for titan, quadro GPUs, reduce this if running out of GPU memory")
+#     argparser.add_argument("--max_length", type=int, default=200000, help="Max sequence length")
+#     argparser.add_argument("--sampling_temp", type=str, default="0.1",
+#                            help="A string of temperatures, 0.2 0.25 0.5. Sampling temperature for amino acids. Suggested values 0.1, 0.15, 0.2, 0.25, 0.3. Higher values will lead to more diversity.")
+#
+#     argparser.add_argument("--out_folder", type=str, help="Path to a folder to output sequences, e.g. /home/out/")
+#     argparser.add_argument("--pdb_path", type=str, default='', help="Path to a single PDB to be designed")
+#     argparser.add_argument("--pdb_path_chains", type=str, default='',
+#                            help="Define which chains need to be designed for a single PDB ")
+#     argparser.add_argument("--jsonl_path", type=str, help="Path to a folder with parsed pdb into jsonl")
+#     argparser.add_argument("--chain_id_jsonl", type=str, default='',
+#                            help="Path to a dictionary specifying which chains need to be designed and which ones are fixed, if not specied all chains will be designed.")
+#     argparser.add_argument("--fixed_positions_jsonl", type=str, default='',
+#                            help="Path to a dictionary with fixed positions")
+#     argparser.add_argument("--omit_AAs", type=list, default='X',
+#                            help="Specify which amino acids should be omitted in the generated sequence, e.g. 'AC' would omit alanine and cystine.")
+#     argparser.add_argument("--bias_AA_jsonl", type=str, default='',
+#                            help="Path to a dictionary which specifies AA composion bias if neededi, e.g. {A: -1.1, F: 0.7} would make A less likely and F more likely.")
+#
+#     argparser.add_argument("--bias_by_res_jsonl", default='', help="Path to dictionary with per position bias.")
+#     argparser.add_argument("--omit_AA_jsonl", type=str, default='',
+#                            help="Path to a dictionary which specifies which amino acids need to be omited from design at specific chain indices")
+#     argparser.add_argument("--pssm_jsonl", type=str, default='', help="Path to a dictionary with pssm")
+#     argparser.add_argument("--pssm_multi", type=float, default=0.0,
+#                            help="A value between [0.0, 1.0], 0.0 means do not use pssm, 1.0 ignore MPNN predictions")
+#     argparser.add_argument("--pssm_threshold", type=float, default=0.0,
+#                            help="A value between -inf + inf to restric per position AAs")
+#     argparser.add_argument("--pssm_log_odds_flag", type=int, default=0, help="0 for False, 1 for True")
+#     argparser.add_argument("--pssm_bias_flag", type=int, default=0, help="0 for False, 1 for True")
+#
+#     argparser.add_argument("--tied_positions_jsonl", type=str, default='',
+#                            help="Path to a dictionary with tied positions")
+#
+#     args = argparser.parse_args()
+#     main(args)

ProteinMPNN-main/protein_mpnn_utils.py

@@ -0,0 +1,1363 @@
+from __future__ import print_function
+import json, time, os, sys, glob
+import shutil
+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 copy
+import torch.nn as nn
+import torch.nn.functional as F
+import random
+import itertools
+
+#A number of functions/classes are adopted from: https://github.com/jingraham/neurips19-graph-protein-design
+
+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
+        # g改
+        fi = biounit.rfind("\\")
+        my_dict['name']=biounit[(fi+1):(fi+5)]
+        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 = []
+    #shuffle all chains before the main loop
+    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 = []
+        num_chains = b['num_of_chains']
+        all_chains = masked_chains + visible_chains
+        #random.shuffle(all_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:
+                    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))
+
+            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):
+        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)
+
+        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):
+        """ 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.W3(self.act(self.W2(self.act(self.W1(h_EV)))))
+        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.W13(self.act(self.W12(self.act(self.W11(h_EV)))))
+        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):
+        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)
+
+        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):
+        """ 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.W3(self.act(self.W2(self.act(self.W1(h_EV)))))
+        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 
+
+
+
+
+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
+
+    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):
+        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.W_s = nn.Embedding(vocab, hidden_dim)
+
+        # Encoder layers
+        self.encoder_layers = nn.ModuleList([
+            EncLayer(hidden_dim, hidden_dim*2, dropout=dropout)
+            for _ in range(num_encoder_layers)
+        ])
+
+        # Decoder layers
+        self.decoder_layers = nn.ModuleList([
+            DecLayer(hidden_dim, hidden_dim*3, dropout=dropout)
+            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):
+        """ 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)
+
+        # 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, h_EX_encoder, E_idx)
+
+
+        chain_M = chain_M*mask #update chain_M to include missing regions
+        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]
+        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)
+        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
+