File size: 2,790 Bytes
35cdf53 | 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 |
"""Methods for merging existing features to create a new example.
Covers:
- Merging features across chains.
- Merging the paired and unpaired parts of the MSA.
"""
from typing import TypeAlias
from flax_model.alphafold3.model import data_constants
import jax.numpy as jnp
import numpy as np
NUM_SEQ_NUM_RES_MSA_FEATURES = data_constants.NUM_SEQ_NUM_RES_MSA_FEATURES
NUM_SEQ_MSA_FEATURES = data_constants.NUM_SEQ_MSA_FEATURES
MSA_PAD_VALUES = data_constants.MSA_PAD_VALUES
xnp_ndarray: TypeAlias = np.ndarray | jnp.ndarray # pylint: disable=invalid-name
BatchDict: TypeAlias = dict[str, xnp_ndarray]
def _pad_features_to_max(feat_name: str, chains: list[BatchDict], axis: int):
"""Pad a set of features to the maximum size amongst all chains.
Args:
feat_name: The feature name to pad.
chains: A list of chains with associated features.
axis: Which axis to pad to the max.
Returns:
A list of features, all with the same size on the given axis.
"""
max_num_seq = np.max([chain[feat_name].shape[axis] for chain in chains])
padded_feats = []
for chain in chains:
feat = chain[feat_name]
padding = np.zeros_like(feat.shape) # pytype: disable=attribute-error
padding[axis] = max_num_seq - feat.shape[axis] # pytype: disable=attribute-error
padding = [(0, p) for p in padding]
padded_feats.append(
np.pad(
feat,
padding,
mode='constant',
constant_values=MSA_PAD_VALUES[feat_name],
)
)
return padded_feats
def merge_msa_features(feat_name: str, chains: list[BatchDict]) -> np.ndarray:
"""Merges MSA features with shape (NUM_SEQ, NUM_RES) across chains."""
expected_dtype = chains[0][feat_name].dtype
if '_all_seq' in feat_name:
return np.concatenate(
[c.get(feat_name, np.array([], expected_dtype)) for c in chains], axis=1
)
else:
# Since each MSA can be of different lengths, we first need to pad them
# all to the size of the largest MSA before concatenating.
padded_feats = _pad_features_to_max(feat_name, chains, axis=0)
return np.concatenate(padded_feats, axis=1)
def merge_paired_and_unpaired_msa(example: BatchDict) -> BatchDict:
"""Concatenates the paired (all_seq) MSA features with the unpaired ones."""
new_example = dict(example)
for feature_name in NUM_SEQ_NUM_RES_MSA_FEATURES + NUM_SEQ_MSA_FEATURES:
if feature_name in example and feature_name + '_all_seq' in example:
feat = example[feature_name]
feat_all_seq = example[feature_name + '_all_seq']
merged_feat = np.concatenate([feat_all_seq, feat], axis=0)
new_example[feature_name] = merged_feat
new_example['num_alignments'] = np.array(
new_example['msa'].shape[0], dtype=np.int32
)
return new_example
|