Function atom37_to_frames

alphafold/model/all_atom_multimer.py:275–371

Computes the frames for the up to 8 rigid groups for each residue.

(
    aatype: jnp.ndarray,  # (...)
    all_atom_positions: geometry.Vec3Array,  # (..., 37)
    all_atom_mask: jnp.ndarray,  # (..., 37)
)

def atom37_to_frames(
    aatype: jnp.ndarray,  # (...)
    all_atom_positions: geometry.Vec3Array,  # (..., 37)
    all_atom_mask: jnp.ndarray,  # (..., 37)
) -> Dict[str, jnp.ndarray]:
  """Computes the frames for the up to 8 rigid groups for each residue."""
  # 0: 'backbone group',
  # 1: 'pre-omega-group', (empty)
  # 2: 'phi-group', (currently empty, because it defines only hydrogens)
  # 3: 'psi-group',
  # 4,5,6,7: 'chi1,2,3,4-group'
  aatype_in_shape = aatype.shape

  # If there is a batch axis, just flatten it away, and reshape everything
  # back at the end of the function.
  aatype = jnp.reshape(aatype, [-1])
  all_atom_positions = jax.tree_map(lambda x: jnp.reshape(x, [-1, 37]),
                                    all_atom_positions)
  all_atom_mask = jnp.reshape(all_atom_mask, [-1, 37])

  # Compute the gather indices for all residues in the chain.
  # shape (N, 8, 3)
  residx_rigidgroup_base_atom37_idx = utils.batched_gather(
      RESTYPE_RIGIDGROUP_BASE_ATOM37_IDX, aatype)

  # Gather the base atom positions for each rigid group.
  base_atom_pos = jax.tree_map(
      lambda x: utils.batched_gather(  # pylint: disable=g-long-lambda
          x, residx_rigidgroup_base_atom37_idx, batch_dims=1),
      all_atom_positions)

  # Compute the Rigids.
  point_on_neg_x_axis = base_atom_pos[:, :, 0]
  origin = base_atom_pos[:, :, 1]
  point_on_xy_plane = base_atom_pos[:, :, 2]
  gt_rotation = geometry.Rot3Array.from_two_vectors(
      origin - point_on_neg_x_axis, point_on_xy_plane - origin)

  gt_frames = geometry.Rigid3Array(gt_rotation, origin)

  # Compute a mask whether the group exists.
  # (N, 8)
  group_exists = utils.batched_gather(RESTYPE_RIGIDGROUP_MASK, aatype)

  # Compute a mask whether ground truth exists for the group
  gt_atoms_exist = utils.batched_gather(  # shape (N, 8, 3)
      all_atom_mask.astype(jnp.float32),
      residx_rigidgroup_base_atom37_idx,
      batch_dims=1)
  gt_exists = jnp.min(gt_atoms_exist, axis=-1) * group_exists  # (N, 8)

  # Adapt backbone frame to old convention (mirror x-axis and z-axis).
  rots = np.tile(np.eye(3, dtype=np.float32), [8, 1, 1])
  rots[0, 0, 0] = -1
  rots[0, 2, 2] = -1
  gt_frames = gt_frames.compose_rotation(
      geometry.Rot3Array.from_array(rots))