Spaces:
Running
Running
File size: 6,284 Bytes
85ab89d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
# # Copyright (c) Meta Platforms, Inc. and affiliates.
# #
# # This source code is licensed under the MIT license found in the
# # LICENSE file in the root directory of this source tree.
#
# import biotite.structure
# import numpy as np
# import torch
# from typing import Sequence, Tuple, List
#
# from esm.inverse_folding.util import (
# load_structure,
# extract_coords_from_structure,
# load_coords,
# get_sequence_loss,
# get_encoder_output,
# )
#
#
# def extract_coords_from_complex(structure: biotite.structure.AtomArray):
# """
# Args:
# structure: biotite AtomArray
# Returns:
# Tuple (coords_list, seq_list)
# - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
# coordinates representing the backbone of each chain
# - seqs: Dictionary mapping chain ids to native sequences of each chain
# """
# coords = {}
# seqs = {}
# all_chains = biotite.structure.get_chains(structure)
# for chain_id in all_chains:
# chain = structure[structure.chain_id == chain_id]
# coords[chain_id], seqs[chain_id] = extract_coords_from_structure(chain)
# return coords, seqs
#
#
# def load_complex_coords(fpath, chains):
# """
# Args:
# fpath: filepath to either pdb or cif file
# chains: the chain ids (the order matters for autoregressive model)
# Returns:
# Tuple (coords_list, seq_list)
# - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
# coordinates representing the backbone of each chain
# - seqs: Dictionary mapping chain ids to native sequences of each chain
# """
# structure = load_structure(fpath, chains)
# return extract_coords_from_complex(structure)
#
#
# def _concatenate_coords(coords, target_chain_id, padding_length=10):
# """
# Args:
# coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
# coordinates representing the backbone of each chain
# target_chain_id: The chain id to sample sequences for
# padding_length: Length of padding between concatenated chains
# Returns:
# Tuple (coords, seq)
# - coords is an L x 3 x 3 array for N, CA, C coordinates, a
# concatenation of the chains with padding in between
# - seq is the extracted sequence, with padding tokens inserted
# between the concatenated chains
# """
# pad_coords = np.full((padding_length, 3, 3), np.nan, dtype=np.float32)
# # For best performance, put the target chain first in concatenation.
# coords_list = [coords[target_chain_id]]
# for chain_id in coords:
# if chain_id == target_chain_id:
# continue
# coords_list.append(pad_coords)
# coords_list.append(coords[chain_id])
# coords_concatenated = np.concatenate(coords_list, axis=0)
# return coords_concatenated
#
#
# def sample_sequence_in_complex(model, coords, target_chain_id, temperature=1.,
# padding_length=10):
# """
# Samples sequence for one chain in a complex.
# Args:
# model: An instance of the GVPTransformer model
# coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
# coordinates representing the backbone of each chain
# target_chain_id: The chain id to sample sequences for
# padding_length: padding length in between chains
# Returns:
# Sampled sequence for the target chain
# """
# target_chain_len = coords[target_chain_id].shape[0]
# all_coords = _concatenate_coords(coords, target_chain_id)
# device = next(model.parameters()).device
#
# # Supply padding tokens for other chains to avoid unused sampling for speed
# padding_pattern = ['<pad>'] * all_coords.shape[0]
# for i in range(target_chain_len):
# padding_pattern[i] = '<mask>'
# sampled = model.sample(all_coords, partial_seq=padding_pattern,
# temperature=temperature, device=device)
# sampled = sampled[:target_chain_len]
# return sampled
#
#
# def score_sequence_in_complex(model, alphabet, coords, target_chain_id,
# target_seq, padding_length=10):
# """
# Scores sequence for one chain in a complex.
# Args:
# model: An instance of the GVPTransformer model
# alphabet: Alphabet for the model
# coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
# coordinates representing the backbone of each chain
# target_chain_id: The chain id to sample sequences for
# target_seq: Target sequence for the target chain for scoring.
# padding_length: padding length in between chains
# Returns:
# Tuple (ll_fullseq, ll_withcoord)
# - ll_fullseq: Average log-likelihood over the full target chain
# - ll_withcoord: Average log-likelihood in target chain excluding those
# residues without coordinates
# """
# all_coords = _concatenate_coords(coords, target_chain_id)
#
# loss, target_padding_mask = get_sequence_loss(model, alphabet, all_coords,
# target_seq)
# ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(
# ~target_padding_mask)
#
# # Also calculate average when excluding masked portions
# coord_mask = np.all(np.isfinite(coords[target_chain_id]), axis=(-1, -2))
# ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
# return ll_fullseq, ll_withcoord
#
#
# def get_encoder_output_for_complex(model, alphabet, coords, target_chain_id):
# """
# Args:
# model: An instance of the GVPTransformer model
# alphabet: Alphabet for the model
# coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
# coordinates representing the backbone of each chain
# target_chain_id: The chain id to sample sequences for
# Returns:
# Dictionary mapping chain id to encoder output for each chain
# """
# all_coords = _concatenate_coords(coords, target_chain_id)
# all_rep = get_encoder_output(model, alphabet, all_coords)
# target_chain_len = coords[target_chain_id].shape[0]
# return all_rep[:target_chain_len]
|