File size: 2,096 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 |
"""Distogram head."""
from typing import Final
from flax_model.alphafold3.common import base_config
from flax_model.alphafold3.model import feat_batch
from flax_model.alphafold3.model import model_config
from flax_model.alphafold3.model.components import haiku_modules as hm
import haiku as hk
import jax
import jax.numpy as jnp
_CONTACT_THRESHOLD: Final[float] = 8.0
_CONTACT_EPSILON: Final[float] = 1e-3
class DistogramHead(hk.Module):
"""Distogram head."""
class Config(base_config.BaseConfig):
first_break: float = 2.3125
last_break: float = 21.6875
num_bins: int = 64
def __init__(
self,
config: Config,
global_config: model_config.GlobalConfig,
name='distogram_head',
):
super().__init__(name=name)
self.config = config
self.global_config = global_config
def __call__(
self,
batch: feat_batch.Batch,
embeddings: dict[str, jnp.ndarray],
return_distogram: bool = False,
) -> dict[str, jnp.ndarray]:
pair_act = embeddings['pair']
seq_mask = batch.token_features.mask.astype(bool)
pair_mask = seq_mask[:, None] * seq_mask[None, :]
left_half_logits = hm.Linear(
self.config.num_bins,
initializer=self.global_config.final_init,
name='half_logits',
)(pair_act)
right_half_logits = left_half_logits
logits = left_half_logits + jnp.swapaxes(right_half_logits, -2, -3)
probs = jax.nn.softmax(logits, axis=-1)
breaks = jnp.linspace(
self.config.first_break,
self.config.last_break,
self.config.num_bins - 1,
)
bin_tops = jnp.append(breaks, breaks[-1] + (breaks[-1] - breaks[-2]))
threshold = _CONTACT_THRESHOLD + _CONTACT_EPSILON
is_contact_bin = 1.0 * (bin_tops <= threshold)
contact_probs = jnp.einsum(
'ijk,k->ij', probs, is_contact_bin, precision=jax.lax.Precision.HIGHEST
)
contact_probs = pair_mask * contact_probs
return_dict = {'bin_edges': breaks, 'contact_probs': contact_probs}
if return_distogram:
return_dict['distogram'] = logits
return return_dict
|