# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""An implementation of the inference pipeline of AlphaMissense."""

alphamissense/__init__.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Common data types and constants used within AlphaMissense."""

alphamissense/common/__init__.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Constants used in AlphaFold."""

import collections
import functools
import os
from typing import Final, Mapping, Sequence

import numpy as np
import tree

# Internal import (35fd).


# Distance from one CA to next CA [trans configuration: omega = 180].
ca_ca = 3.80209737096

# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
# chi angles so their chi angle lists are empty.
chi_angles_atoms = {
    'ALA': [],
    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
    'ARG': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
            ['CB', 'CG', 'CD', 'NE'], ['CG', 'CD', 'NE', 'CZ']],
    'ASN': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'OD1']],
    'ASP': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'OD1']],
    'CYS': [['N', 'CA', 'CB', 'SG']],
    'GLN': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
            ['CB', 'CG', 'CD', 'OE1']],
    'GLU': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
            ['CB', 'CG', 'CD', 'OE1']],
    'GLY': [],
    'HIS': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'ND1']],
    'ILE': [['N', 'CA', 'CB', 'CG1'], ['CA', 'CB', 'CG1', 'CD1']],
    'LEU': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
    'LYS': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
            ['CB', 'CG', 'CD', 'CE'], ['CG', 'CD', 'CE', 'NZ']],
    'MET': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'SD'],
            ['CB', 'CG', 'SD', 'CE']],
    'PHE': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
    'PRO': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD']],
    'SER': [['N', 'CA', 'CB', 'OG']],
    'THR': [['N', 'CA', 'CB', 'OG1']],
    'TRP': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
    'TYR': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
    'VAL': [['N', 'CA', 'CB', 'CG1']],
}

# If chi angles given in fixed-length array, this matrix determines how to mask
# them for each AA type. The order is as per restype_order (see below).
chi_angles_mask = [
    [0.0, 0.0, 0.0, 0.0],  # ALA
    [1.0, 1.0, 1.0, 1.0],  # ARG
    [1.0, 1.0, 0.0, 0.0],  # ASN
    [1.0, 1.0, 0.0, 0.0],  # ASP
    [1.0, 0.0, 0.0, 0.0],  # CYS
    [1.0, 1.0, 1.0, 0.0],  # GLN
    [1.0, 1.0, 1.0, 0.0],  # GLU
    [0.0, 0.0, 0.0, 0.0],  # GLY
    [1.0, 1.0, 0.0, 0.0],  # HIS
    [1.0, 1.0, 0.0, 0.0],  # ILE
    [1.0, 1.0, 0.0, 0.0],  # LEU
    [1.0, 1.0, 1.0, 1.0],  # LYS
    [1.0, 1.0, 1.0, 0.0],  # MET
    [1.0, 1.0, 0.0, 0.0],  # PHE
    [1.0, 1.0, 0.0, 0.0],  # PRO
    [1.0, 0.0, 0.0, 0.0],  # SER
    [1.0, 0.0, 0.0, 0.0],  # THR
    [1.0, 1.0, 0.0, 0.0],  # TRP
    [1.0, 1.0, 0.0, 0.0],  # TYR
    [1.0, 0.0, 0.0, 0.0],  # VAL
]

# The following chi angles are pi periodic: they can be rotated by a multiple
# of pi without affecting the structure.
chi_pi_periodic = [
    [0.0, 0.0, 0.0, 0.0],  # ALA
    [0.0, 0.0, 0.0, 0.0],  # ARG
    [0.0, 0.0, 0.0, 0.0],  # ASN
    [0.0, 1.0, 0.0, 0.0],  # ASP
    [0.0, 0.0, 0.0, 0.0],  # CYS
    [0.0, 0.0, 0.0, 0.0],  # GLN
    [0.0, 0.0, 1.0, 0.0],  # GLU
    [0.0, 0.0, 0.0, 0.0],  # GLY
    [0.0, 0.0, 0.0, 0.0],  # HIS
    [0.0, 0.0, 0.0, 0.0],  # ILE
    [0.0, 0.0, 0.0, 0.0],  # LEU
    [0.0, 0.0, 0.0, 0.0],  # LYS
    [0.0, 0.0, 0.0, 0.0],  # MET
    [0.0, 1.0, 0.0, 0.0],  # PHE
    [0.0, 0.0, 0.0, 0.0],  # PRO
    [0.0, 0.0, 0.0, 0.0],  # SER
    [0.0, 0.0, 0.0, 0.0],  # THR
    [0.0, 0.0, 0.0, 0.0],  # TRP
    [0.0, 1.0, 0.0, 0.0],  # TYR
    [0.0, 0.0, 0.0, 0.0],  # VAL
    [0.0, 0.0, 0.0, 0.0],  # UNK
]

# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
# psi and chi angles:
# 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'
# The atom positions are relative to the axis-end-atom of the corresponding
# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
# is defined such that the dihedral-angle-defining atom (the last entry in
# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
# format: [atomname, group_idx, rel_position]
rigid_group_atom_positions = {
    'ALA': [
        ['N', 0, (-0.525, 1.363, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.526, -0.000, -0.000)],
        ['CB', 0, (-0.529, -0.774, -1.205)],
        ['O', 3, (0.627, 1.062, 0.000)],
    ],
    'ARG': [
        ['N', 0, (-0.524, 1.362, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.525, -0.000, -0.000)],
        ['CB', 0, (-0.524, -0.778, -1.209)],
        ['O', 3, (0.626, 1.062, 0.000)],
        ['CG', 4, (0.616, 1.390, -0.000)],
        ['CD', 5, (0.564, 1.414, 0.000)],
        ['NE', 6, (0.539, 1.357, -0.000)],
        ['NH1', 7, (0.206, 2.301, 0.000)],
        ['NH2', 7, (2.078, 0.978, -0.000)],
        ['CZ', 7, (0.758, 1.093, -0.000)],
    ],
    'ASN': [
        ['N', 0, (-0.536, 1.357, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.526, -0.000, -0.000)],
        ['CB', 0, (-0.531, -0.787, -1.200)],
        ['O', 3, (0.625, 1.062, 0.000)],
        ['CG', 4, (0.584, 1.399, 0.000)],
        ['ND2', 5, (0.593, -1.188, 0.001)],
        ['OD1', 5, (0.633, 1.059, 0.000)],
    ],
    'ASP': [
        ['N', 0, (-0.525, 1.362, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.527, 0.000, -0.000)],
        ['CB', 0, (-0.526, -0.778, -1.208)],
        ['O', 3, (0.626, 1.062, -0.000)],
        ['CG', 4, (0.593, 1.398, -0.000)],
        ['OD1', 5, (0.610, 1.091, 0.000)],
        ['OD2', 5, (0.592, -1.101, -0.003)],
    ],
    'CYS': [
        ['N', 0, (-0.522, 1.362, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.524, 0.000, 0.000)],
        ['CB', 0, (-0.519, -0.773, -1.212)],
        ['O', 3, (0.625, 1.062, -0.000)],
        ['SG', 4, (0.728, 1.653, 0.000)],
    ],
    'GLN': [
        ['N', 0, (-0.526, 1.361, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.526, 0.000, 0.000)],
        ['CB', 0, (-0.525, -0.779, -1.207)],
        ['O', 3, (0.626, 1.062, -0.000)],
        ['CG', 4, (0.615, 1.393, 0.000)],
        ['CD', 5, (0.587, 1.399, -0.000)],
        ['NE2', 6, (0.593, -1.189, -0.001)],
        ['OE1', 6, (0.634, 1.060, 0.000)],
    ],
    'GLU': [
        ['N', 0, (-0.528, 1.361, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.526, -0.000, -0.000)],
        ['CB', 0, (-0.526, -0.781, -1.207)],
        ['O', 3, (0.626, 1.062, 0.000)],
        ['CG', 4, (0.615, 1.392, 0.000)],
        ['CD', 5, (0.600, 1.397, 0.000)],
        ['OE1', 6, (0.607, 1.095, -0.000)],
        ['OE2', 6, (0.589, -1.104, -0.001)],
    ],
    'GLY': [
        ['N', 0, (-0.572, 1.337, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.517, -0.000, -0.000)],
        ['O', 3, (0.626, 1.062, -0.000)],
    ],
    'HIS': [
        ['N', 0, (-0.527, 1.360, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.525, 0.000, 0.000)],
        ['CB', 0, (-0.525, -0.778, -1.208)],
        ['O', 3, (0.625, 1.063, 0.000)],
        ['CG', 4, (0.600, 1.370, -0.000)],
        ['CD2', 5, (0.889, -1.021, 0.003)],
        ['ND1', 5, (0.744, 1.160, -0.000)],
        ['CE1', 5, (2.030, 0.851, 0.002)],
        ['NE2', 5, (2.145, -0.466, 0.004)],
    ],
    'ILE': [
        ['N', 0, (-0.493, 1.373, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.527, -0.000, -0.000)],
        ['CB', 0, (-0.536, -0.793, -1.213)],
        ['O', 3, (0.627, 1.062, -0.000)],
        ['CG1', 4, (0.534, 1.437, -0.000)],
        ['CG2', 4, (0.540, -0.785, -1.199)],
        ['CD1', 5, (0.619, 1.391, 0.000)],
    ],
    'LEU': [
        ['N', 0, (-0.520, 1.363, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.525, -0.000, -0.000)],
        ['CB', 0, (-0.522, -0.773, -1.214)],
        ['O', 3, (0.625, 1.063, -0.000)],
        ['CG', 4, (0.678, 1.371, 0.000)],
        ['CD1', 5, (0.530, 1.430, -0.000)],
        ['CD2', 5, (0.535, -0.774, 1.200)],
    ],
    'LYS': [
        ['N', 0, (-0.526, 1.362, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.526, 0.000, 0.000)],
        ['CB', 0, (-0.524, -0.778, -1.208)],
        ['O', 3, (0.626, 1.062, -0.000)],
        ['CG', 4, (0.619, 1.390, 0.000)],
        ['CD', 5, (0.559, 1.417, 0.000)],
        ['CE', 6, (0.560, 1.416, 0.000)],
        ['NZ', 7, (0.554, 1.387, 0.000)],
    ],
    'MET': [
        ['N', 0, (-0.521, 1.364, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.525, 0.000, 0.000)],
        ['CB', 0, (-0.523, -0.776, -1.210)],
        ['O', 3, (0.625, 1.062, -0.000)],
        ['CG', 4, (0.613, 1.391, -0.000)],
        ['SD', 5, (0.703, 1.695, 0.000)],
        ['CE', 6, (0.320, 1.786, -0.000)],
    ],
    'PHE': [
        ['N', 0, (-0.518, 1.363, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.524, 0.000, -0.000)],
        ['CB', 0, (-0.525, -0.776, -1.212)],
        ['O', 3, (0.626, 1.062, -0.000)],
        ['CG', 4, (0.607, 1.377, 0.000)],
        ['CD1', 5, (0.709, 1.195, -0.000)],
        ['CD2', 5, (0.706, -1.196, 0.000)],
        ['CE1', 5, (2.102, 1.198, -0.000)],
        ['CE2', 5, (2.098, -1.201, -0.000)],
        ['CZ', 5, (2.794, -0.003, -0.001)],
    ],
    'PRO': [
        ['N', 0, (-0.566, 1.351, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.527, -0.000, 0.000)],
        ['CB', 0, (-0.546, -0.611, -1.293)],
        ['O', 3, (0.621, 1.066, 0.000)],
        ['CG', 4, (0.382, 1.445, 0.0)],
        # ['CD', 5, (0.427, 1.440, 0.0)],
        ['CD', 5, (0.477, 1.424, 0.0)],  # manually made angle 2 degrees larger
    ],
    'SER': [
        ['N', 0, (-0.529, 1.360, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.525, -0.000, -0.000)],
        ['CB', 0, (-0.518, -0.777, -1.211)],
        ['O', 3, (0.626, 1.062, -0.000)],
        ['OG', 4, (0.503, 1.325, 0.000)],
    ],
    'THR': [
        ['N', 0, (-0.517, 1.364, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.526, 0.000, -0.000)],
        ['CB', 0, (-0.516, -0.793, -1.215)],
        ['O', 3, (0.626, 1.062, 0.000)],
        ['CG2', 4, (0.550, -0.718, -1.228)],
        ['OG1', 4, (0.472, 1.353, 0.000)],
    ],
    'TRP': [
        ['N', 0, (-0.521, 1.363, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.525, -0.000, 0.000)],
        ['CB', 0, (-0.523, -0.776, -1.212)],
        ['O', 3, (0.627, 1.062, 0.000)],
        ['CG', 4, (0.609, 1.370, -0.000)],
        ['CD1', 5, (0.824, 1.091, 0.000)],
        ['CD2', 5, (0.854, -1.148, -0.005)],
        ['CE2', 5, (2.186, -0.678, -0.007)],
        ['CE3', 5, (0.622, -2.530, -0.007)],
        ['NE1', 5, (2.140, 0.690, -0.004)],
        ['CH2', 5, (3.028, -2.890, -0.013)],
        ['CZ2', 5, (3.283, -1.543, -0.011)],
        ['CZ3', 5, (1.715, -3.389, -0.011)],
    ],
    'TYR': [
        ['N', 0, (-0.522, 1.362, 0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.524, -0.000, -0.000)],
        ['CB', 0, (-0.522, -0.776, -1.213)],
        ['O', 3, (0.627, 1.062, -0.000)],
        ['CG', 4, (0.607, 1.382, -0.000)],
        ['CD1', 5, (0.716, 1.195, -0.000)],
        ['CD2', 5, (0.713, -1.194, -0.001)],
        ['CE1', 5, (2.107, 1.200, -0.002)],
        ['CE2', 5, (2.104, -1.201, -0.003)],
        ['OH', 5, (4.168, -0.002, -0.005)],
        ['CZ', 5, (2.791, -0.001, -0.003)],
    ],
    'VAL': [
        ['N', 0, (-0.494, 1.373, -0.000)],
        ['CA', 0, (0.000, 0.000, 0.000)],
        ['C', 0, (1.527, -0.000, -0.000)],
        ['CB', 0, (-0.533, -0.795, -1.213)],
        ['O', 3, (0.627, 1.062, -0.000)],
        ['CG1', 4, (0.540, 1.429, -0.000)],
        ['CG2', 4, (0.533, -0.776, 1.203)],
    ],
}

# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
residue_atoms = {
    'ALA': ['C', 'CA', 'CB', 'N', 'O'],
    'ARG': ['C', 'CA', 'CB', 'CG', 'CD', 'CZ', 'N', 'NE', 'O', 'NH1', 'NH2'],
    'ASP': ['C', 'CA', 'CB', 'CG', 'N', 'O', 'OD1', 'OD2'],
    'ASN': ['C', 'CA', 'CB', 'CG', 'N', 'ND2', 'O', 'OD1'],
    'CYS': ['C', 'CA', 'CB', 'N', 'O', 'SG'],
    'GLU': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'O', 'OE1', 'OE2'],
    'GLN': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'NE2', 'O', 'OE1'],
    'GLY': ['C', 'CA', 'N', 'O'],
    'HIS': ['C', 'CA', 'CB', 'CG', 'CD2', 'CE1', 'N', 'ND1', 'NE2', 'O'],
    'ILE': ['C', 'CA', 'CB', 'CG1', 'CG2', 'CD1', 'N', 'O'],
    'LEU': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'N', 'O'],
    'LYS': ['C', 'CA', 'CB', 'CG', 'CD', 'CE', 'N', 'NZ', 'O'],
    'MET': ['C', 'CA', 'CB', 'CG', 'CE', 'N', 'O', 'SD'],
    'PHE': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'N', 'O'],
    'PRO': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'O'],
    'SER': ['C', 'CA', 'CB', 'N', 'O', 'OG'],
    'THR': ['C', 'CA', 'CB', 'CG2', 'N', 'O', 'OG1'],
    'TRP': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE2', 'CE3', 'CZ2', 'CZ3',
            'CH2', 'N', 'NE1', 'O'],
    'TYR': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'N', 'O',
            'OH'],
    'VAL': ['C', 'CA', 'CB', 'CG1', 'CG2', 'N', 'O']
}

# Naming swaps for ambiguous atom names.
# Due to symmetries in the amino acids the naming of atoms is ambiguous in
# 4 of the 20 amino acids.
# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
# in LEU, VAL and ARG can be resolved by using the 3d constellations of
# the 'ambiguous' atoms and their neighbours)
residue_atom_renaming_swaps = {
    'ASP': {'OD1': 'OD2'},
    'GLU': {'OE1': 'OE2'},
    'PHE': {'CD1': 'CD2', 'CE1': 'CE2'},
    'TYR': {'CD1': 'CD2', 'CE1': 'CE2'},
}

# Van der Waals radii [Angstrom] of the atoms (from Wikipedia)
van_der_waals_radius = {
    'C': 1.7,
    'N': 1.55,
    'O': 1.52,
    'S': 1.8,
}

Bond = collections.namedtuple(
    'Bond', ['atom1_name', 'atom2_name', 'length', 'stddev'])
BondAngle = collections.namedtuple(
    'BondAngle',
    ['atom1_name', 'atom2_name', 'atom3name', 'angle_rad', 'stddev'])


@functools.lru_cache(maxsize=None)
def load_stereo_chemical_props() -> tuple[Mapping[str, Sequence[Bond]],
                                          Mapping[str, Sequence[Bond]],
                                          Mapping[str, Sequence[BondAngle]]]:
  """Load stereo_chemical_props.txt into a nice structure.

  Load literature values for bond lengths and bond angles and translate
  bond angles into the length of the opposite edge of the triangle
  ("residue_virtual_bonds").

  Returns:
    residue_bonds: Dict that maps resname -> list of Bond tuples.
    residue_virtual_bonds: Dict that maps resname -> list of Bond tuples.
    residue_bond_angles: Dict that maps resname -> list of BondAngle tuples.
  """
  stereo_chemical_props_path = os.path.join(
      os.path.dirname(os.path.abspath(__file__)), 'stereo_chemical_props.txt'
  )
  with open(stereo_chemical_props_path, 'rt') as f:
    stereo_chemical_props = f.read()
  lines_iter = iter(stereo_chemical_props.splitlines())
  # Load bond lengths.
  residue_bonds = {}
  next(lines_iter)  # Skip header line.
  for line in lines_iter:
    if line.strip() == '-':
      break
    bond, resname, length, stddev = line.split()
    atom1, atom2 = bond.split('-')
    if resname not in residue_bonds:
      residue_bonds[resname] = []
    residue_bonds[resname].append(
        Bond(atom1, atom2, float(length), float(stddev)))
  residue_bonds['UNK'] = []

  # Load bond angles.
  residue_bond_angles = {}
  next(lines_iter)  # Skip empty line.
  next(lines_iter)  # Skip header line.
  for line in lines_iter:
    if line.strip() == '-':
      break
    bond, resname, angle_degree, stddev_degree = line.split()
    atom1, atom2, atom3 = bond.split('-')
    if resname not in residue_bond_angles:
      residue_bond_angles[resname] = []
    residue_bond_angles[resname].append(
        BondAngle(atom1, atom2, atom3,
                  float(angle_degree) / 180. * np.pi,
                  float(stddev_degree) / 180. * np.pi))
  residue_bond_angles['UNK'] = []

  def make_bond_key(atom1_name, atom2_name):
    """Unique key to lookup bonds."""
    return '-'.join(sorted([atom1_name, atom2_name]))

  # Translate bond angles into distances ("virtual bonds").
  residue_virtual_bonds = {}
  for resname, bond_angles in residue_bond_angles.items():
    # Create a fast lookup dict for bond lengths.
    bond_cache = {}
    for b in residue_bonds[resname]:
      bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
    residue_virtual_bonds[resname] = []
    for ba in bond_angles:
      bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
      bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]

      # Compute distance between atom1 and atom3 using the law of cosines
      # c^2 = a^2 + b^2 - 2ab*cos(gamma).
      gamma = ba.angle_rad
      length = np.sqrt(bond1.length**2 + bond2.length**2
                       - 2 * bond1.length * bond2.length * np.cos(gamma))

      # Propagation of uncertainty assuming uncorrelated errors.
      dl_outer = 0.5 / length
      dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
      dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
      dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
      stddev = np.sqrt((dl_dgamma * ba.stddev)**2 +
                       (dl_db1 * bond1.stddev)**2 +
                       (dl_db2 * bond2.stddev)**2)
      residue_virtual_bonds[resname].append(
          Bond(ba.atom1_name, ba.atom3name, length, stddev))

  return (residue_bonds,
          residue_virtual_bonds,
          residue_bond_angles)


# Between-residue bond lengths for general bonds (first element) and for Proline
# (second element).
between_res_bond_length_c_n = [1.329, 1.341]
between_res_bond_length_stddev_c_n = [0.014, 0.016]

# Between-residue cos_angles.
between_res_cos_angles_c_n_ca = [-0.5203, 0.0353]  # degrees: 121.352 +- 2.315
between_res_cos_angles_ca_c_n = [-0.4473, 0.0311]  # degrees: 116.568 +- 1.995

# This mapping is used when we need to store atom data in a format that requires
# fixed atom data size for every residue (e.g. a numpy array).
atom_types = [
    'N', 'CA', 'C', 'CB', 'O', 'CG', 'CG1', 'CG2', 'OG', 'OG1', 'SG', 'CD',
    'CD1', 'CD2', 'ND1', 'ND2', 'OD1', 'OD2', 'SD', 'CE', 'CE1', 'CE2', 'CE3',
    'NE', 'NE1', 'NE2', 'OE1', 'OE2', 'CH2', 'NH1', 'NH2', 'OH', 'CZ', 'CZ2',
    'CZ3', 'NZ', 'OXT'
]
atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
atom_type_num = len(atom_types)  # := 37.


# A compact atom encoding with 14 columns
# pylint: disable=line-too-long
# pylint: disable=bad-whitespace
restype_name_to_atom14_names = {
    'ALA': ['N', 'CA', 'C', 'O', 'CB', '',    '',    '',    '',    '',    '',    '',    '',    ''],
    'ARG': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'NE',  'CZ',  'NH1', 'NH2', '',    '',    ''],
    'ASN': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'ND2', '',    '',    '',    '',    '',    ''],
    'ASP': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'OD2', '',    '',    '',    '',    '',    ''],
    'CYS': ['N', 'CA', 'C', 'O', 'CB', 'SG',  '',    '',    '',    '',    '',    '',    '',    ''],
    'GLN': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'NE2', '',    '',    '',    '',    ''],
    'GLU': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'OE2', '',    '',    '',    '',    ''],
    'GLY': ['N', 'CA', 'C', 'O', '',   '',    '',    '',    '',    '',    '',    '',    '',    ''],
    'HIS': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'ND1', 'CD2', 'CE1', 'NE2', '',    '',    '',    ''],
    'ILE': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', 'CD1', '',    '',    '',    '',    '',    ''],
    'LEU': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', '',    '',    '',    '',    '',    ''],
    'LYS': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'CE',  'NZ',  '',    '',    '',    '',    ''],
    'MET': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'SD',  'CE',  '',    '',    '',    '',    '',    ''],
    'PHE': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  '',    '',    ''],
    'PRO': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  '',    '',    '',    '',    '',    '',    ''],
    'SER': ['N', 'CA', 'C', 'O', 'CB', 'OG',  '',    '',    '',    '',    '',    '',    '',    ''],
    'THR': ['N', 'CA', 'C', 'O', 'CB', 'OG1', 'CG2', '',    '',    '',    '',    '',    '',    ''],
    'TRP': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'NE1', 'CE2', 'CE3', 'CZ2', 'CZ3', 'CH2'],
    'TYR': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  'OH',  '',    ''],
    'VAL': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '',    '',    '',    '',    '',    '',    ''],
    'UNK': ['',  '',   '',  '',  '',   '',    '',    '',    '',    '',    '',    '',    '',    ''],

}
# pylint: enable=line-too-long
# pylint: enable=bad-whitespace


# This is the standard residue order when coding AA type as a number.
# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
restypes = [
    'A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P',
    'S', 'T', 'W', 'Y', 'V'
]
restype_order = {restype: i for i, restype in enumerate(restypes)}
restype_num = len(restypes)  # := 20.
unk_restype_index = restype_num  # Catch-all index for unknown restypes.

restypes_with_x = restypes + ['X']
restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}


def sequence_to_onehot(
    sequence: str,
    mapping: Mapping[str, int],
    map_unknown_to_x: bool = False) -> np.ndarray:
  """Maps the given sequence into a one-hot encoded matrix.

  Args:
    sequence: An amino acid sequence.
    mapping: A dictionary mapping amino acids to integers.
    map_unknown_to_x: If True, any amino acid that is not in the mapping will be
      mapped to the unknown amino acid 'X'. If the mapping doesn't contain
      amino acid 'X', an error will be thrown. If False, any amino acid not in
      the mapping will throw an error.

  Returns:
    A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
    the sequence.

  Raises:
    ValueError: If the mapping doesn't contain values from 0 to
      num_unique_aas - 1 without any gaps.
  """
  num_entries = max(mapping.values()) + 1

  if sorted(set(mapping.values())) != list(range(num_entries)):
    raise ValueError('The mapping must have values from 0 to num_unique_aas-1 '
                     'without any gaps. Got: %s' % sorted(mapping.values()))

  one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)

  for aa_index, aa_type in enumerate(sequence):
    if map_unknown_to_x:
      if aa_type.isalpha() and aa_type.isupper():
        aa_id = mapping.get(aa_type, mapping['X'])
      else:
        raise ValueError(f'Invalid character in the sequence: {aa_type}')
    else:
      aa_id = mapping[aa_type]
    one_hot_arr[aa_index, aa_id] = 1

  return one_hot_arr


restype_1to3 = {
    'A': 'ALA',
    'R': 'ARG',
    'N': 'ASN',
    'D': 'ASP',
    'C': 'CYS',
    'Q': 'GLN',
    'E': 'GLU',
    'G': 'GLY',
    'H': 'HIS',
    'I': 'ILE',
    'L': 'LEU',
    'K': 'LYS',
    'M': 'MET',
    'F': 'PHE',
    'P': 'PRO',
    'S': 'SER',
    'T': 'THR',
    'W': 'TRP',
    'Y': 'TYR',
    'V': 'VAL',
}

PROTEIN_CHAIN: Final[str] = 'polypeptide(L)'
POLYMER_CHAIN: Final[str] = 'polymer'


def atom_id_to_type(atom_id: str) -> str:
  """Convert atom ID to atom type, works only for standard protein residues.

  Args:
    atom_id: Atom ID to be converted.

  Returns:
    String corresponding to atom type.

  Raises:
    ValueError: If atom ID not recognized.
  """

  if atom_id.startswith('C'):
    return 'C'
  elif atom_id.startswith('N'):
    return 'N'
  elif atom_id.startswith('O'):
    return 'O'
  elif atom_id.startswith('H'):
    return 'H'
  elif atom_id.startswith('S'):
    return 'S'
  raise ValueError('Atom ID not recognized.')


# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
# many more, and less common, three letter names as keys and maps many of these
# to the same one letter name (including 'X' and 'U' which we don't use here).
restype_3to1 = {v: k for k, v in restype_1to3.items()}

# Define a restype name for all unknown residues.
unk_restype = 'UNK'

resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
resname_to_idx = {resname: i for i, resname in enumerate(resnames)}


# The mapping here uses hhblits convention, so that B is mapped to D, J and O
# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
# remaining 20 amino acids are kept in alphabetical order.
# There are 2 non-amino acid codes, X (representing any amino acid) and
# "-" representing a missing amino acid in an alignment.  The id for these
# codes is put at the end (20 and 21) so that they can easily be ignored if
# desired.
HHBLITS_AA_TO_ID = {
    'A': 0,
    'B': 2,
    'C': 1,
    'D': 2,
    'E': 3,
    'F': 4,
    'G': 5,
    'H': 6,
    'I': 7,
    'J': 20,
    'K': 8,
    'L': 9,
    'M': 10,
    'N': 11,
    'O': 20,
    'P': 12,
    'Q': 13,
    'R': 14,
    'S': 15,
    'T': 16,
    'U': 1,
    'V': 17,
    'W': 18,
    'X': 20,
    'Y': 19,
    'Z': 3,
    '-': 21,
}

# Partial inversion of HHBLITS_AA_TO_ID.
ID_TO_HHBLITS_AA = {
    0: 'A',
    1: 'C',  # Also U.
    2: 'D',  # Also B.
    3: 'E',  # Also Z.
    4: 'F',
    5: 'G',
    6: 'H',
    7: 'I',
    8: 'K',
    9: 'L',
    10: 'M',
    11: 'N',
    12: 'P',
    13: 'Q',
    14: 'R',
    15: 'S',
    16: 'T',
    17: 'V',
    18: 'W',
    19: 'Y',
    20: 'X',  # Includes J and O.
    21: '-',
}

restypes_with_x_and_gap = restypes + ['X', '-']
MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
    for i in range(len(restypes_with_x_and_gap)))


def _make_standard_atom_mask() -> np.ndarray:
  """Returns [num_res_types, num_atom_types] mask array."""
  # +1 to account for unknown (all 0s).
  mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
  for restype, restype_letter in enumerate(restypes):
    restype_name = restype_1to3[restype_letter]
    atom_names = residue_atoms[restype_name]
    for atom_name in atom_names:
      atom_type = atom_order[atom_name]
      mask[restype, atom_type] = 1
  return mask


STANDARD_ATOM_MASK = _make_standard_atom_mask()


# A one hot representation for the first and second atoms defining the axis
# of rotation for each chi-angle in each residue.
def chi_angle_atom(atom_index: int) -> np.ndarray:
  """Define chi-angle rigid groups via one-hot representations."""
  chi_angles_index = {}
  one_hots = []

  for k, v in chi_angles_atoms.items():
    indices = [atom_types.index(s[atom_index]) for s in v]
    indices.extend([-1]*(4-len(indices)))
    chi_angles_index[k] = indices

  for r in restypes:
    res3 = restype_1to3[r]
    one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
    one_hots.append(one_hot)

  one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
  one_hot = np.stack(one_hots, axis=0)
  one_hot = np.transpose(one_hot, [0, 2, 1])

  return one_hot

chi_atom_1_one_hot = chi_angle_atom(1)
chi_atom_2_one_hot = chi_angle_atom(2)

# An array like chi_angles_atoms but using indices rather than names.
chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
chi_angles_atom_indices = tree.map_structure(
    lambda atom_name: atom_order[atom_name], chi_angles_atom_indices)
chi_angles_atom_indices = np.array([
    chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms)))
    for chi_atoms in chi_angles_atom_indices])

# Mapping from (res_name, atom_name) pairs to the atom's chi group index
# and atom index within that group.
chi_groups_for_atom = collections.defaultdict(list)
for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
  for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
    for atom_i, atom in enumerate(chi_group):
      chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
chi_groups_for_atom = dict(chi_groups_for_atom)


def _make_rigid_transformation_4x4(ex, ey, translation):
  """Create a rigid 4x4 transformation matrix from two axes and transl."""
  # Normalize ex.
  ex_normalized = ex / np.linalg.norm(ex)

  # make ey perpendicular to ex
  ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
  ey_normalized /= np.linalg.norm(ey_normalized)

  # compute ez as cross product
  eznorm = np.cross(ex_normalized, ey_normalized)
  m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
  m = np.concatenate([m, [[0., 0., 0., 1.]]], axis=0)
  return m


# create an array with (restype, atomtype) --> rigid_group_idx
# and an array with (restype, atomtype, coord) for the atom positions
# and compute affine transformation matrices (4,4) from one rigid group to the
# previous group
restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int)
restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int)
restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)


def _make_rigid_group_constants():
  """Fill the arrays above."""
  for restype, restype_letter in enumerate(restypes):
    resname = restype_1to3[restype_letter]
    for atomname, group_idx, atom_position in rigid_group_atom_positions[
        resname]:
      atomtype = atom_order[atomname]
      restype_atom37_to_rigid_group[restype, atomtype] = group_idx
      restype_atom37_mask[restype, atomtype] = 1
      restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position

      atom14idx = restype_name_to_atom14_names[resname].index(atomname)
      restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
      restype_atom14_mask[restype, atom14idx] = 1
      restype_atom14_rigid_group_positions[restype,
                                           atom14idx, :] = atom_position

  for restype, restype_letter in enumerate(restypes):
    resname = restype_1to3[restype_letter]
    atom_positions = {name: np.array(pos) for name, _, pos
                      in rigid_group_atom_positions[resname]}

    # backbone to backbone is the identity transform
    restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)

    # pre-omega-frame to backbone (currently dummy identity matrix)
    restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)

    # phi-frame to backbone
    mat = _make_rigid_transformation_4x4(
        ex=atom_positions['N'] - atom_positions['CA'],
        ey=np.array([1., 0., 0.]),
        translation=atom_positions['N'])
    restype_rigid_group_default_frame[restype, 2, :, :] = mat

    # psi-frame to backbone
    mat = _make_rigid_transformation_4x4(
        ex=atom_positions['C'] - atom_positions['CA'],
        ey=atom_positions['CA'] - atom_positions['N'],
        translation=atom_positions['C'])
    restype_rigid_group_default_frame[restype, 3, :, :] = mat

    # chi1-frame to backbone
    if chi_angles_mask[restype][0]:
      base_atom_names = chi_angles_atoms[resname][0]
      base_atom_positions = [atom_positions[name] for name in base_atom_names]
      mat = _make_rigid_transformation_4x4(
          ex=base_atom_positions[2] - base_atom_positions[1],
          ey=base_atom_positions[0] - base_atom_positions[1],
          translation=base_atom_positions[2])
      restype_rigid_group_default_frame[restype, 4, :, :] = mat

    # chi2-frame to chi1-frame
    # chi3-frame to chi2-frame
    # chi4-frame to chi3-frame
    # luckily all rotation axes for the next frame start at (0,0,0) of the
    # previous frame
    for chi_idx in range(1, 4):
      if chi_angles_mask[restype][chi_idx]:
        axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
        axis_end_atom_position = atom_positions[axis_end_atom_name]
        mat = _make_rigid_transformation_4x4(
            ex=axis_end_atom_position,
            ey=np.array([-1., 0., 0.]),
            translation=axis_end_atom_position)
        restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat


_make_rigid_group_constants()


def make_atom14_dists_bounds(overlap_tolerance=1.5,
                             bond_length_tolerance_factor=15):
  """compute upper and lower bounds for bonds to assess violations."""
  restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
  restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
  restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
  residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
  for restype, restype_letter in enumerate(restypes):
    resname = restype_1to3[restype_letter]
    atom_list = restype_name_to_atom14_names[resname]

    # create lower and upper bounds for clashes
    for atom1_idx, atom1_name in enumerate(atom_list):
      if not atom1_name:
        continue
      atom1_radius = van_der_waals_radius[atom1_name[0]]
      for atom2_idx, atom2_name in enumerate(atom_list):
        if (not atom2_name) or atom1_idx == atom2_idx:
          continue
        atom2_radius = van_der_waals_radius[atom2_name[0]]
        lower = atom1_radius + atom2_radius - overlap_tolerance
        upper = 1e10
        restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
        restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
        restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
        restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper

    # overwrite lower and upper bounds for bonds and angles
    for b in (*residue_bonds[resname], *residue_virtual_bonds[resname]):
      atom1_idx = atom_list.index(b.atom1_name)
      atom2_idx = atom_list.index(b.atom2_name)
      lower = b.length - bond_length_tolerance_factor * b.stddev
      upper = b.length + bond_length_tolerance_factor * b.stddev
      restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
      restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
      restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
      restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
      restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
      restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
  return {'lower_bound': restype_atom14_bond_lower_bound,  # shape (21,14,14)
          'upper_bound': restype_atom14_bond_upper_bound,  # shape (21,14,14)
          'stddev': restype_atom14_bond_stddev,  # shape (21,14,14)
         }

alphamissense/common/residue_constants.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Parses the mmCIF file format."""
import collections
import dataclasses
import functools
import io
from typing import Any, Mapping, Optional, Sequence

from absl import logging
from Bio import PDB
from Bio.Data import SCOPData

# Type aliases:
ChainId = str
PdbHeader = Mapping[str, Any]
PdbStructure = PDB.Structure.Structure
SeqRes = str
MmCIFDict = Mapping[str, Sequence[str]]


@dataclasses.dataclass(frozen=True)
class Monomer:
  id: str
  num: int


# Note - mmCIF format provides no guarantees on the type of author-assigned
# sequence numbers. They need not be integers.
@dataclasses.dataclass(frozen=True)
class AtomSite:
  residue_name: str
  author_chain_id: str
  mmcif_chain_id: str
  author_seq_num: str
  mmcif_seq_num: int
  insertion_code: str
  hetatm_atom: str
  model_num: int


# Used to map SEQRES index to a residue in the structure.
@dataclasses.dataclass(frozen=True)
class ResiduePosition:
  chain_id: str
  residue_number: int
  insertion_code: str


@dataclasses.dataclass(frozen=True)
class ResidueAtPosition:
  position: Optional[ResiduePosition]
  name: str
  is_missing: bool
  hetflag: str


@dataclasses.dataclass(frozen=True)
class MmcifObject:
  """Representation of a parsed mmCIF file.

  Contains:
    file_id: A meaningful name, e.g. a pdb_id. Should be unique amongst all
      files being processed.
    header: Biopython header.
    structure: Biopython structure.
    chain_to_seqres: Dict mapping chain_id to 1 letter amino acid sequence. E.g.
      {'A': 'ABCDEFG'}
    seqres_to_structure: Dict; for each chain_id contains a mapping between
      SEQRES index and a ResidueAtPosition. e.g. {'A': {0: ResidueAtPosition,
                                                        1: ResidueAtPosition,
                                                        ...}}
    raw_string: The raw string used to construct the MmcifObject.
  """
  file_id: str
  header: PdbHeader
  structure: PdbStructure
  chain_to_seqres: Mapping[ChainId, SeqRes]
  seqres_to_structure: Mapping[ChainId, Mapping[int, ResidueAtPosition]]
  raw_string: Any


@dataclasses.dataclass(frozen=True)
class ParsingResult:
  """Returned by the parse function.

  Contains:
    mmcif_object: A MmcifObject, may be None if no chain could be successfully
      parsed.
    errors: A dict mapping (file_id, chain_id) to any exception generated.
  """
  mmcif_object: Optional[MmcifObject]
  errors: Mapping[tuple[str, str], Any]


class ParseError(Exception):
  """An error indicating that an mmCIF file could not be parsed."""


def mmcif_loop_to_list(prefix: str,
                       parsed_info: MmCIFDict) -> Sequence[Mapping[str, str]]:
  """Extracts loop associated with a prefix from mmCIF data as a list.

  Reference for loop_ in mmCIF:
    http://mmcif.wwpdb.org/docs/tutorials/mechanics/pdbx-mmcif-syntax.html

  Args:
    prefix: Prefix shared by each of the data items in the loop.
      e.g. '_entity_poly_seq.', where the data items are _entity_poly_seq.num,
      _entity_poly_seq.mon_id. Should include the trailing period.
    parsed_info: A dict of parsed mmCIF data, e.g. _mmcif_dict from a Biopython
      parser.

  Returns:
    Returns a list of dicts; each dict represents 1 entry from an mmCIF loop.
  """
  cols = []
  data = []
  for key, value in parsed_info.items():
    if key.startswith(prefix):
      cols.append(key)
      data.append(value)

  assert all([len(xs) == len(data[0]) for xs in data]), (
      'mmCIF error: Not all loops are the same length: %s' % cols)

  return [dict(zip(cols, xs)) for xs in zip(*data)]


def mmcif_loop_to_dict(prefix: str,
                       index: str,
                       parsed_info: MmCIFDict,
                       ) -> Mapping[str, Mapping[str, str]]:
  """Extracts loop associated with a prefix from mmCIF data as a dictionary.

  Args:
    prefix: Prefix shared by each of the data items in the loop.
      e.g. '_entity_poly_seq.', where the data items are _entity_poly_seq.num,
      _entity_poly_seq.mon_id. Should include the trailing period.
    index: Which item of loop data should serve as the key.
    parsed_info: A dict of parsed mmCIF data, e.g. _mmcif_dict from a Biopython
      parser.

  Returns:
    Returns a dict of dicts; each dict represents 1 entry from an mmCIF loop,
    indexed by the index column.
  """
  entries = mmcif_loop_to_list(prefix, parsed_info)
  return {entry[index]: entry for entry in entries}


@functools.lru_cache(16, typed=False)
def parse(*,
          file_id: str,
          mmcif_string: str,
          catch_all_errors: bool = True) -> ParsingResult:
  """Entry point, parses an mmcif_string.

  Args:
    file_id: A string identifier for this file. Should be unique within the
      collection of files being processed.
    mmcif_string: Contents of an mmCIF file.
    catch_all_errors: If True, all exceptions are caught and error messages are
      returned as part of the ParsingResult. If False exceptions will be allowed
      to propagate.

  Returns:
    A ParsingResult.
  """
  errors = {}
  try:
    parser = PDB.MMCIFParser(QUIET=True)
    handle = io.StringIO(mmcif_string)
    full_structure = parser.get_structure('', handle)
    first_model_structure = _get_first_model(full_structure)
    # Extract the _mmcif_dict from the parser, which contains useful fields not
    # reflected in the Biopython structure.
    parsed_info = parser._mmcif_dict  # pylint:disable=protected-access

    # Ensure all values are lists, even if singletons.
    for key, value in parsed_info.items():
      if not isinstance(value, list):
        parsed_info[key] = [value]

    header = _get_header(parsed_info)

    # Determine the protein chains, and their start numbers according to the
    # internal mmCIF numbering scheme (likely but not guaranteed to be 1).
    valid_chains = _get_protein_chains(parsed_info=parsed_info)
    if not valid_chains:
      return ParsingResult(
          None, {(file_id, ''): 'No protein chains found in this file.'})
    seq_start_num = {chain_id: min([monomer.num for monomer in seq])
                     for chain_id, seq in valid_chains.items()}

    # Loop over the atoms for which we have coordinates. Populate two mappings:
    # -mmcif_to_author_chain_id (maps internal mmCIF chain ids to chain ids used
    # the authors / Biopython).
    # -seq_to_structure_mappings (maps idx into sequence to ResidueAtPosition).
    mmcif_to_author_chain_id = {}
    seq_to_structure_mappings = {}
    for atom in _get_atom_site_list(parsed_info):
      if atom.model_num != '1':
        # We only process the first model at the moment.
        continue

      mmcif_to_author_chain_id[atom.mmcif_chain_id] = atom.author_chain_id

      if atom.mmcif_chain_id in valid_chains:
        hetflag = ' '
        if atom.hetatm_atom == 'HETATM':
          # Water atoms are assigned a special hetflag of W in Biopython. We
          # need to do the same, so that this hetflag can be used to fetch
          # a residue from the Biopython structure by id.
          if atom.residue_name in ('HOH', 'WAT'):
            hetflag = 'W'
          else:
            hetflag = 'H_' + atom.residue_name
        insertion_code = atom.insertion_code
        if not _is_set(atom.insertion_code):
          insertion_code = ' '
        position = ResiduePosition(chain_id=atom.author_chain_id,
                                   residue_number=int(atom.author_seq_num),
                                   insertion_code=insertion_code)
        seq_idx = int(atom.mmcif_seq_num) - seq_start_num[atom.mmcif_chain_id]
        current = seq_to_structure_mappings.get(atom.author_chain_id, {})
        current[seq_idx] = ResidueAtPosition(position=position,
                                             name=atom.residue_name,
                                             is_missing=False,
                                             hetflag=hetflag)
        seq_to_structure_mappings[atom.author_chain_id] = current

    # Add missing residue information to seq_to_structure_mappings.
    for chain_id, seq_info in valid_chains.items():
      author_chain = mmcif_to_author_chain_id[chain_id]
      current_mapping = seq_to_structure_mappings[author_chain]
      for idx, monomer in enumerate(seq_info):
        if idx not in current_mapping:
          current_mapping[idx] = ResidueAtPosition(position=None,
                                                   name=monomer.id,
                                                   is_missing=True,
                                                   hetflag=' ')

    author_chain_to_sequence = {}
    for chain_id, seq_info in valid_chains.items():
      author_chain = mmcif_to_author_chain_id[chain_id]
      seq = []
      for monomer in seq_info:
        code = SCOPData.protein_letters_3to1.get(monomer.id, 'X')
        seq.append(code if len(code) == 1 else 'X')
      seq = ''.join(seq)
      author_chain_to_sequence[author_chain] = seq

    mmcif_object = MmcifObject(
        file_id=file_id,
        header=header,
        structure=first_model_structure,
        chain_to_seqres=author_chain_to_sequence,
        seqres_to_structure=seq_to_structure_mappings,
        raw_string=parsed_info)

    return ParsingResult(mmcif_object=mmcif_object, errors=errors)
  except Exception as e:  # pylint:disable=broad-except
    errors[(file_id, '')] = e
    if not catch_all_errors:
      raise
    return ParsingResult(mmcif_object=None, errors=errors)


def _get_first_model(structure: PdbStructure) -> PdbStructure:
  """Returns the first model in a Biopython structure."""
  return next(structure.get_models())


def get_release_date(parsed_info: MmCIFDict) -> str:
  """Returns the oldest revision date."""
  revision_dates = parsed_info['_pdbx_audit_revision_history.revision_date']
  return min(revision_dates)


def _get_header(parsed_info: MmCIFDict) -> PdbHeader:
  """Returns a basic header containing method, release date and resolution."""
  header = {}

  experiments = mmcif_loop_to_list('_exptl.', parsed_info)
  header['structure_method'] = ','.join([
      experiment['_exptl.method'].lower() for experiment in experiments])

  # Note: The release_date here corresponds to the oldest revision. We prefer to
  # use this for dataset filtering over the deposition_date.
  if '_pdbx_audit_revision_history.revision_date' in parsed_info:
    header['release_date'] = get_release_date(parsed_info)
  else:
    logging.warning('Could not determine release_date: %s',
                    parsed_info['_entry.id'])

  header['resolution'] = 0.00
  for res_key in ('_refine.ls_d_res_high', '_em_3d_reconstruction.resolution',
                  '_reflns.d_resolution_high'):
    if res_key in parsed_info:
      try:
        raw_resolution = parsed_info[res_key][0]
        header['resolution'] = float(raw_resolution)
      except ValueError:
        logging.debug('Invalid resolution format: %s', parsed_info[res_key])

  return header


def _get_atom_site_list(parsed_info: MmCIFDict) -> Sequence[AtomSite]:
  """Returns list of atom sites; contains data not present in the structure."""
  return [AtomSite(*site) for site in zip(  # pylint:disable=g-complex-comprehension
      parsed_info['_atom_site.label_comp_id'],
      parsed_info['_atom_site.auth_asym_id'],
      parsed_info['_atom_site.label_asym_id'],
      parsed_info['_atom_site.auth_seq_id'],
      parsed_info['_atom_site.label_seq_id'],
      parsed_info['_atom_site.pdbx_PDB_ins_code'],
      parsed_info['_atom_site.group_PDB'],
      parsed_info['_atom_site.pdbx_PDB_model_num'],
      )]


def _get_protein_chains(
    *, parsed_info: Mapping[str, Any]) -> Mapping[ChainId, Sequence[Monomer]]:
  """Extracts polymer information for protein chains only.

  Args:
    parsed_info: _mmcif_dict produced by the Biopython parser.

  Returns:
    A dict mapping mmcif chain id to a list of Monomers.
  """
  # Get polymer information for each entity in the structure.
  entity_poly_seqs = mmcif_loop_to_list('_entity_poly_seq.', parsed_info)

  polymers = collections.defaultdict(list)
  for entity_poly_seq in entity_poly_seqs:
    polymers[entity_poly_seq['_entity_poly_seq.entity_id']].append(
        Monomer(id=entity_poly_seq['_entity_poly_seq.mon_id'],
                num=int(entity_poly_seq['_entity_poly_seq.num'])))

  # Get chemical compositions. Will allow us to identify which of these polymers
  # are proteins.
  chem_comps = mmcif_loop_to_dict('_chem_comp.', '_chem_comp.id', parsed_info)

  # Get chains information for each entity. Necessary so that we can return a
  # dict keyed on chain id rather than entity.
  struct_asyms = mmcif_loop_to_list('_struct_asym.', parsed_info)

  entity_to_mmcif_chains = collections.defaultdict(list)
  for struct_asym in struct_asyms:
    chain_id = struct_asym['_struct_asym.id']
    entity_id = struct_asym['_struct_asym.entity_id']
    entity_to_mmcif_chains[entity_id].append(chain_id)

  # Identify and return the valid protein chains.
  valid_chains = {}
  for entity_id, seq_info in polymers.items():
    chain_ids = entity_to_mmcif_chains[entity_id]

    # Reject polymers without any peptide-like components, such as DNA/RNA.
    if any(['peptide' in chem_comps[monomer.id]['_chem_comp.type'].lower()
            for monomer in seq_info]):
      for chain_id in chain_ids:
        valid_chains[chain_id] = seq_info
  return valid_chains


def _is_set(data: str) -> bool:
  """Returns False if data is a special mmCIF character indicating 'unset'."""
  return data not in ('.', '?')

alphamissense/data/mmcif_parsing.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Functions for parsing various file formats."""
import collections
import dataclasses
import itertools
import re
import string
from typing import Iterable, Optional, Sequence

# Internal import (7716).


DeletionMatrix = Sequence[Sequence[int]]


@dataclasses.dataclass(frozen=True)
class Msa:
  """Class representing a parsed MSA file."""
  sequences: Sequence[str]
  deletion_matrix: DeletionMatrix
  descriptions: Sequence[str]

  def __post_init__(self):
    if not (len(self.sequences) ==
            len(self.deletion_matrix) ==
            len(self.descriptions)):
      raise ValueError(
          'All fields for an MSA must have the same length. '
          f'Got {len(self.sequences)} sequences, '
          f'{len(self.deletion_matrix)} rows in the deletion matrix and '
          f'{len(self.descriptions)} descriptions.')

  def __len__(self):
    return len(self.sequences)

  def truncate(self, max_seqs: int):
    return Msa(sequences=self.sequences[:max_seqs],
               deletion_matrix=self.deletion_matrix[:max_seqs],
               descriptions=self.descriptions[:max_seqs])


@dataclasses.dataclass(frozen=True)
class TemplateHit:
  """Class representing a template hit."""
  index: int
  name: str
  aligned_cols: int
  sum_probs: Optional[float]
  query: str
  hit_sequence: str
  indices_query: list[int]
  indices_hit: list[int]


def parse_fasta(fasta_string: str) -> tuple[Sequence[str], Sequence[str]]:
  """Parses FASTA string and returns list of strings with amino-acid sequences.

  Arguments:
    fasta_string: The string contents of a FASTA file.

  Returns:
    A tuple of two lists:
    * A list of sequences.
    * A list of sequence descriptions taken from the comment lines. In the
      same order as the sequences.
  """
  sequences = []
  descriptions = []
  index = -1
  for line in fasta_string.splitlines():
    line = line.strip()
    if line.startswith('>'):
      index += 1
      descriptions.append(line[1:])  # Remove the '>' at the beginning.
      sequences.append('')
      continue
    elif not line:
      continue  # Skip blank lines.
    sequences[index] += line

  return sequences, descriptions


def parse_stockholm(stockholm_string: str) -> Msa:
  """Parses sequences and deletion matrix from stockholm format alignment.

  Args:
    stockholm_string: The string contents of a stockholm file. The first
      sequence in the file should be the query sequence.

  Returns:
    A tuple of:
      * A list of sequences that have been aligned to the query. These
        might contain duplicates.
      * The deletion matrix for the alignment as a list of lists. The element
        at `deletion_matrix[i][j]` is the number of residues deleted from
        the aligned sequence i at residue position j.
      * The names of the targets matched, including the jackhmmer subsequence
        suffix.
  """
  name_to_sequence = collections.OrderedDict()
  for line in stockholm_string.splitlines():
    line = line.strip()
    if not line or line.startswith(('#', '//')):
      continue
    name, sequence = line.split()
    if name not in name_to_sequence:
      name_to_sequence[name] = ''
    name_to_sequence[name] += sequence

  msa = []
  deletion_matrix = []

  query = ''
  keep_columns = []
  for seq_index, sequence in enumerate(name_to_sequence.values()):
    if seq_index == 0:
      # Gather the columns with gaps from the query
      query = sequence
      keep_columns = [i for i, res in enumerate(query) if res != '-']

    # Remove the columns with gaps in the query from all sequences.
    aligned_sequence = ''.join([sequence[c] for c in keep_columns])

    msa.append(aligned_sequence)

    # Count the number of deletions w.r.t. query.
    deletion_vec = []
    deletion_count = 0
    for seq_res, query_res in zip(sequence, query):
      if seq_res != '-' or query_res != '-':
        if query_res == '-':
          deletion_count += 1
        else:
          deletion_vec.append(deletion_count)
          deletion_count = 0
    deletion_matrix.append(deletion_vec)

  return Msa(sequences=msa,
             deletion_matrix=deletion_matrix,
             descriptions=list(name_to_sequence.keys()))


def parse_a3m(a3m_string: str) -> Msa:
  """Parses sequences and deletion matrix from a3m format alignment.

  Args:
    a3m_string: The string contents of a a3m file. The first sequence in the
      file should be the query sequence.

  Returns:
    A tuple of:
      * A list of sequences that have been aligned to the query. These
        might contain duplicates.
      * The deletion matrix for the alignment as a list of lists. The element
        at `deletion_matrix[i][j]` is the number of residues deleted from
        the aligned sequence i at residue position j.
      * A list of descriptions, one per sequence, from the a3m file.
  """
  sequences, descriptions = parse_fasta(a3m_string)
  deletion_matrix = []
  for msa_sequence in sequences:
    deletion_vec = []
    deletion_count = 0
    for j in msa_sequence:
      if j.islower():
        deletion_count += 1
      else:
        deletion_vec.append(deletion_count)
        deletion_count = 0
    deletion_matrix.append(deletion_vec)

  # Make the MSA matrix out of aligned (deletion-free) sequences.
  deletion_table = str.maketrans('', '', string.ascii_lowercase)
  aligned_sequences = [s.translate(deletion_table) for s in sequences]
  return Msa(sequences=aligned_sequences,
             deletion_matrix=deletion_matrix,
             descriptions=descriptions)


def _convert_sto_seq_to_a3m(
    query_non_gaps: Sequence[bool], sto_seq: str) -> Iterable[str]:
  for is_query_res_non_gap, sequence_res in zip(query_non_gaps, sto_seq):
    if is_query_res_non_gap:
      yield sequence_res
    elif sequence_res != '-':
      yield sequence_res.lower()


def convert_stockholm_to_a3m(stockholm_format: str,
                             max_sequences: Optional[int] = None,
                             remove_first_row_gaps: bool = True) -> str:
  """Converts MSA in Stockholm format to the A3M format."""
  descriptions = {}
  sequences = {}
  reached_max_sequences = False

  for line in stockholm_format.splitlines():
    reached_max_sequences = max_sequences and len(sequences) >= max_sequences
    if line.strip() and not line.startswith(('#', '//')):
      # Ignore blank lines, markup and end symbols - remainder are alignment
      # sequence parts.
      seqname, aligned_seq = line.split(maxsplit=1)
      if seqname not in sequences:
        if reached_max_sequences:
          continue
        sequences[seqname] = ''
      sequences[seqname] += aligned_seq

  for line in stockholm_format.splitlines():
    if line[:4] == '#=GS':
      # Description row - example format is:
      # #=GS UniRef90_Q9H5Z4/4-78            DE [subseq from] cDNA: FLJ22755 ...
      columns = line.split(maxsplit=3)
      seqname, feature = columns[1:3]
      value = columns[3] if len(columns) == 4 else ''
      if feature != 'DE':
        continue
      if reached_max_sequences and seqname not in sequences:
        continue
      descriptions[seqname] = value
      if len(descriptions) == len(sequences):
        break

  # Convert sto format to a3m line by line
  a3m_sequences = {}
  if remove_first_row_gaps:
    # query_sequence is assumed to be the first sequence
    query_sequence = next(iter(sequences.values()))
    query_non_gaps = [res != '-' for res in query_sequence]
  for seqname, sto_sequence in sequences.items():
    # Dots are optional in a3m format and are commonly removed.
    out_sequence = sto_sequence.replace('.', '')
    if remove_first_row_gaps:
      out_sequence = ''.join(
          _convert_sto_seq_to_a3m(query_non_gaps, out_sequence))
    a3m_sequences[seqname] = out_sequence

  fasta_chunks = (f">{k} {descriptions.get(k, '')}\n{a3m_sequences[k]}"
                  for k in a3m_sequences)
  return '\n'.join(fasta_chunks) + '\n'  # Include terminating newline.


def _keep_line(line: str, seqnames: set[str]) -> bool:
  """Function to decide which lines to keep."""
  if not line.strip():
    return True
  if line.strip() == '//':  # End tag
    return True
  if line.startswith('# STOCKHOLM'):  # Start tag
    return True
  if line.startswith('#=GC RF'):  # Reference Annotation Line
    return True
  if line[:4] == '#=GS':  # Description lines - keep if sequence in list.
    _, seqname, _ = line.split(maxsplit=2)
    return seqname in seqnames
  elif line.startswith('#'):  # Other markup - filter out
    return False
  else:  # Alignment data - keep if sequence in list.
    seqname = line.partition(' ')[0]
    return seqname in seqnames


def truncate_stockholm_msa(stockholm_msa_path: str, max_sequences: int) -> str:
  """Reads + truncates a Stockholm file while preventing excessive RAM usage."""
  seqnames = set()
  filtered_lines = []

  with open(stockholm_msa_path) as f:
    for line in f:
      if line.strip() and not line.startswith(('#', '//')):
        # Ignore blank lines, markup and end symbols - remainder are alignment
        # sequence parts.
        seqname = line.partition(' ')[0]
        seqnames.add(seqname)
        if len(seqnames) >= max_sequences:
          break

    f.seek(0)
    for line in f:
      if _keep_line(line, seqnames):
        filtered_lines.append(line)

  return ''.join(filtered_lines)


def remove_empty_columns_from_stockholm_msa(stockholm_msa: str) -> str:
  """Removes empty columns (dashes-only) from a Stockholm MSA."""
  processed_lines = {}
  unprocessed_lines = {}
  for i, line in enumerate(stockholm_msa.splitlines()):
    if line.startswith('#=GC RF'):
      reference_annotation_i = i
      reference_annotation_line = line
      # Reached the end of this chunk of the alignment. Process chunk.
      _, _, first_alignment = line.rpartition(' ')
      mask = []
      for j in range(len(first_alignment)):
        for _, unprocessed_line in unprocessed_lines.items():
          prefix, _, alignment = unprocessed_line.rpartition(' ')
          if alignment[j] != '-':
            mask.append(True)
            break
        else:  # Every row contained a hyphen - empty column.
          mask.append(False)
      # Add reference annotation for processing with mask.
      unprocessed_lines[reference_annotation_i] = reference_annotation_line

      if not any(mask):  # All columns were empty. Output empty lines for chunk.
        for line_index in unprocessed_lines:
          processed_lines[line_index] = ''
      else:
        for line_index, unprocessed_line in unprocessed_lines.items():
          prefix, _, alignment = unprocessed_line.rpartition(' ')
          masked_alignment = ''.join(itertools.compress(alignment, mask))
          processed_lines[line_index] = f'{prefix} {masked_alignment}'

      # Clear raw_alignments.
      unprocessed_lines = {}
    elif line.strip() and not line.startswith(('#', '//')):
      unprocessed_lines[i] = line
    else:
      processed_lines[i] = line
  return '\n'.join((processed_lines[i] for i in range(len(processed_lines))))


def deduplicate_stockholm_msa(stockholm_msa: str) -> str:
  """Remove duplicate sequences (ignoring insertions wrt query)."""
  sequence_dict = collections.defaultdict(str)

  # First we must extract all sequences from the MSA.
  for line in stockholm_msa.splitlines():
    # Only consider the alignments - ignore reference annotation, empty lines,
    # descriptions or markup.
    if line.strip() and not line.startswith(('#', '//')):
      line = line.strip()
      seqname, alignment = line.split()
      sequence_dict[seqname] += alignment

  seen_sequences = set()
  seqnames = set()
  # First alignment is the query.
  query_align = next(iter(sequence_dict.values()))
  mask = [c != '-' for c in query_align]  # Mask is False for insertions.
  for seqname, alignment in sequence_dict.items():
    # Apply mask to remove all insertions from the string.
    masked_alignment = ''.join(itertools.compress(alignment, mask))
    if masked_alignment in seen_sequences:
      continue
    else:
      seen_sequences.add(masked_alignment)
      seqnames.add(seqname)

  filtered_lines = []
  for line in stockholm_msa.splitlines():
    if _keep_line(line, seqnames):
      filtered_lines.append(line)

  return '\n'.join(filtered_lines) + '\n'


def _get_hhr_line_regex_groups(
    regex_pattern: str, line: str) -> Sequence[Optional[str]]:
  match = re.match(regex_pattern, line)
  if match is None:
    raise RuntimeError(f'Could not parse query line {line}')
  return match.groups()


def _update_hhr_residue_indices_list(
    sequence: str, start_index: int, indices_list: list[int]):
  """Computes the relative indices for each residue with respect to the original sequence."""
  counter = start_index
  for symbol in sequence:
    if symbol == '-':
      indices_list.append(-1)
    else:
      indices_list.append(counter)
      counter += 1


def _parse_hhr_hit(detailed_lines: Sequence[str]) -> TemplateHit:
  """Parses the detailed HMM HMM comparison section for a single Hit.

  This works on .hhr files generated from both HHBlits and HHSearch.

  Args:
    detailed_lines: A list of lines from a single comparison section between 2
      sequences (which each have their own HMM's)

  Returns:
    A dictionary with the information from that detailed comparison section

  Raises:
    RuntimeError: If a certain line cannot be processed
  """
  # Parse first 2 lines.
  number_of_hit = int(detailed_lines[0].split()[-1])
  name_hit = detailed_lines[1][1:]

  # Parse the summary line.
  pattern = (
      'Probab=(.*)[\t ]*E-value=(.*)[\t ]*Score=(.*)[\t ]*Aligned_cols=(.*)[\t'
      ' ]*Identities=(.*)%[\t ]*Similarity=(.*)[\t ]*Sum_probs=(.*)[\t '
      ']*Template_Neff=(.*)')
  match = re.match(pattern, detailed_lines[2])
  if match is None:
    raise RuntimeError(
        'Could not parse section: %s. Expected this: \n%s to contain summary.' %
        (detailed_lines, detailed_lines[2]))
  (_, _, _, aligned_cols, _, _, sum_probs, _) = [float(x)
                                                 for x in match.groups()]

  # The next section reads the detailed comparisons. These are in a 'human
  # readable' format which has a fixed length. The strategy employed is to
  # assume that each block starts with the query sequence line, and to parse
  # that with a regexp in order to deduce the fixed length used for that block.
  query = ''
  hit_sequence = ''
  indices_query = []
  indices_hit = []
  length_block = None

  for line in detailed_lines[3:]:
    # Parse the query sequence line
    if (line.startswith('Q ') and not line.startswith('Q ss_dssp') and
        not line.startswith('Q ss_pred') and
        not line.startswith('Q Consensus')):
      # Thus the first 17 characters must be 'Q <query_name> ', and we can parse
      # everything after that.
      #              start    sequence       end       total_sequence_length
      patt = r'[\t ]*([0-9]*) ([A-Z-]*)[\t ]*([0-9]*) \([0-9]*\)'
      groups = _get_hhr_line_regex_groups(patt, line[17:])

      # Get the length of the parsed block using the start and finish indices,
      # and ensure it is the same as the actual block length.
      start = int(groups[0]) - 1  # Make index zero based.
      delta_query = groups[1]
      end = int(groups[2])
      num_insertions = len([x for x in delta_query if x == '-'])
      length_block = end - start + num_insertions
      assert length_block == len(delta_query)

      # Update the query sequence and indices list.
      query += delta_query
      _update_hhr_residue_indices_list(delta_query, start, indices_query)

    elif line.startswith('T '):
      # Parse the hit sequence.
      if (not line.startswith('T ss_dssp') and
          not line.startswith('T ss_pred') and
          not line.startswith('T Consensus')):
        # Thus the first 17 characters must be 'T <hit_name> ', and we can
        # parse everything after that.
        #              start    sequence       end     total_sequence_length
        patt = r'[\t ]*([0-9]*) ([A-Z-]*)[\t ]*[0-9]* \([0-9]*\)'
        groups = _get_hhr_line_regex_groups(patt, line[17:])
        start = int(groups[0]) - 1  # Make index zero based.
        delta_hit_sequence = groups[1]
        assert length_block == len(delta_hit_sequence)

        # Update the hit sequence and indices list.
        hit_sequence += delta_hit_sequence
        _update_hhr_residue_indices_list(delta_hit_sequence, start, indices_hit)

  return TemplateHit(
      index=number_of_hit,
      name=name_hit,
      aligned_cols=int(aligned_cols),
      sum_probs=sum_probs,
      query=query,
      hit_sequence=hit_sequence,
      indices_query=indices_query,
      indices_hit=indices_hit,
  )


def parse_hhr(hhr_string: str) -> Sequence[TemplateHit]:
  """Parses the content of an entire HHR file."""
  lines = hhr_string.splitlines()

  # Each .hhr file starts with a results table, then has a sequence of hit
  # "paragraphs", each paragraph starting with a line 'No <hit number>'. We
  # iterate through each paragraph to parse each hit.

  block_starts = [i for i, line in enumerate(lines) if line.startswith('No ')]

  hits = []
  if block_starts:
    block_starts.append(len(lines))  # Add the end of the final block.
    for i in range(len(block_starts) - 1):
      hits.append(_parse_hhr_hit(lines[block_starts[i]:block_starts[i + 1]]))
  return hits


def parse_e_values_from_tblout(tblout: str) -> dict[str, float]:
  """Parse target to e-value mapping parsed from Jackhmmer tblout string."""
  e_values = {'query': 0}
  lines = [line for line in tblout.splitlines() if line[0] != '#']
  # As per http://eddylab.org/software/hmmer/Userguide.pdf fields are
  # space-delimited. Relevant fields are (1) target name:  and
  # (5) E-value (full sequence) (numbering from 1).
  for line in lines:
    fields = line.split()
    e_value = fields[4]
    target_name = fields[0]
    e_values[target_name] = float(e_value)
  return e_values


def _get_indices(sequence: str, start: int) -> list[int]:
  """Returns indices for non-gap/insert residues starting at the given index."""
  indices = []
  counter = start
  for symbol in sequence:
    # Skip gaps but add a placeholder so that the alignment is preserved.
    if symbol == '-':
      indices.append(-1)
    # Skip deleted residues, but increase the counter.
    elif symbol.islower():
      counter += 1
    # Normal aligned residue. Increase the counter and append to indices.
    else:
      indices.append(counter)
      counter += 1
  return indices


@dataclasses.dataclass(frozen=True)
class HitMetadata:
  pdb_id: str
  chain: str
  start: int
  end: int
  length: int
  text: str


def _parse_hmmsearch_description(description: str) -> HitMetadata:
  """Parses the hmmsearch A3M sequence description line."""
  # Example 1: >4pqx_A/2-217 [subseq from] mol:protein length:217  Free text
  # Example 2: >5g3r_A/1-55 [subseq from] mol:protein length:352
  match = re.match(
      r'^>?([a-z0-9]+)_(\w+)/([0-9]+)-([0-9]+).*protein length:([0-9]+) *(.*)$',
      description.strip())

  if not match:
    raise ValueError(f'Could not parse description: "{description}".')

  return HitMetadata(
      pdb_id=match[1],
      chain=match[2],
      start=int(match[3]),
      end=int(match[4]),
      length=int(match[5]),
      text=match[6])


def parse_hmmsearch_a3m(query_sequence: str,
                        a3m_string: str,
                        skip_first: bool = True) -> Sequence[TemplateHit]:
  """Parses an a3m string produced by hmmsearch.

  Args:
    query_sequence: The query sequence.
    a3m_string: The a3m string produced by hmmsearch.
    skip_first: Whether to skip the first sequence in the a3m string.

  Returns:
    A sequence of `TemplateHit` results.
  """
  # Zip the descriptions and MSAs together, skip the first query sequence.
  parsed_a3m = list(zip(*parse_fasta(a3m_string)))
  if skip_first:
    parsed_a3m = parsed_a3m[1:]

  indices_query = _get_indices(query_sequence, start=0)

  hits = []
  for i, (hit_sequence, hit_description) in enumerate(parsed_a3m, start=1):
    if 'mol:protein' not in hit_description:
      continue  # Skip non-protein chains.
    metadata = _parse_hmmsearch_description(hit_description)
    # Aligned columns are only the match states.
    aligned_cols = sum([r.isupper() and r != '-' for r in hit_sequence])
    indices_hit = _get_indices(hit_sequence, start=metadata.start - 1)

    hit = TemplateHit(
        index=i,
        name=f'{metadata.pdb_id}_{metadata.chain}',
        aligned_cols=aligned_cols,
        sum_probs=None,
        query=query_sequence,
        hit_sequence=hit_sequence.upper(),
        indices_query=indices_query,
        indices_hit=indices_hit,
    )
    hits.append(hit)

  return hits

alphamissense/data/parsers.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Functions for building the input features for the AlphaMissense model."""

import dataclasses
# Internal import (7717).
import os
import subprocess
import tempfile
from typing import Any, Mapping, MutableMapping, Optional, Sequence

from absl import logging
import numpy as np

# Internal import (7716).
from alphamissense.common import residue_constants
from alphamissense.data import mmcif_parsing
from alphamissense.data import msa_identifiers
from alphamissense.data import parsers
from alphamissense.data import templates
from alphamissense.data.tools import jackhmmer


FeatureDict = MutableMapping[str, np.ndarray]

# SEQ_FEATURES have shape [num_res, ...].
# NUM_SEQ_MSA_FEATURES have shape [num_msa, ...].
# NUM_SEQ_NUM_RES_MSA_FEATURES have shape [num_msa, num_res, ...].
NUM_SEQ_NUM_RES_MSA_FEATURES = ('msa', 'msa_mask', 'deletion_matrix',)
NUM_SEQ_MSA_FEATURES = ('cluster_bias_mask',)
SEQ_FEATURES = (
    'residue_index', 'all_atom_positions', 'msa_profile', 'deletion_mean',
    'aatype', 'all_atom_mask', 'seq_mask', 'variant_aatype', 'variant_mask')
TOKENIZER = {
    r: i for i, r in enumerate(residue_constants.restypes + ['X', '-'])}

AFDB_BUCKET = 'gs://public-datasets-deepmind-alphafold-v4'


@dataclasses.dataclass(frozen=True)
class ProteinVariant:
  """A class containing definition of protein variant."""
  sequence: str  # The original protein sequence as extracted from FASTA file.
  position: int  # 1-based aa-position within protein.
  reference_aa: str  # May be different from AA in `sequence`.
  alternate_aa: str
  protein_id: str
  pathogenicity: Optional[bool]

  def __post_init__(self):
    if len(self.reference_aa) != 1 or len(self.alternate_aa) != 1:
      raise ValueError(
          'Reference and alternate amino acids must be of length 1.')

    if self.position < 1 or self.position > len(self.sequence):
      raise ValueError(
          f'{self.protein_id} Variant position {self.position} is invalid. '
          f'Must be 1 to {len(self.sequence)} (both included).')

    if self.sequence[self.position - 1] != self.reference_aa:
      sequence_aa = self.sequence[self.position - 1]
      raise ValueError(f'Mismatch between {sequence_aa=} and input '
                       f'{self.reference_aa=} at position {self.position}. '
                       'Note that position must follow 1-based indexing.')

  @property
  def alternate_sequence(self) -> str:
    return self.sequence[:self.position - 1] + (
        self.alternate_aa + self.sequence[self.position:])

  @property
  def reference_sequence(self) -> str:
    return self.sequence


class Cropper:
  """Crops and pads features to enforce uniform input sizes."""

  def __init__(self, crop_size: int, msa_crop_size: int):
    self.crop_size = crop_size
    self.msa_crop_size = msa_crop_size

  def _apply_crop(self, features: FeatureDict) -> FeatureDict:
    residue_keep_mask = self._choose_crop_region(features)
    for k in features:
      if k in SEQ_FEATURES:
        features[k] = features[k][residue_keep_mask]
      elif k in NUM_SEQ_NUM_RES_MSA_FEATURES:
        features[k] = features[k][:self.msa_crop_size, residue_keep_mask]
      elif k in NUM_SEQ_MSA_FEATURES:
        features[k] = features[k][:self.msa_crop_size]
    return features

  def _apply_pad(self, features: FeatureDict) -> FeatureDict:
    """Pads num residues and num MSA sequences."""
    num_res_before_pad = features['aatype'].shape[0]
    num_msa_seq_before_pad = features['msa'].shape[0]
    paddable_features = frozenset(
        SEQ_FEATURES + NUM_SEQ_MSA_FEATURES + NUM_SEQ_NUM_RES_MSA_FEATURES)

    def _empty_padding_shape(feature: np.ndarray) -> list[list[int]]:
      padding = []
      for _ in range(feature.ndim):
        padding.append([0, 0])
      return padding

    for k, v in features.items():
      if k not in paddable_features:
        logging.info('Skipping padding for feature %s', k)
        continue

      pad_size = _empty_padding_shape(v)
      if k in frozenset(NUM_SEQ_MSA_FEATURES + NUM_SEQ_NUM_RES_MSA_FEATURES):
        assert v.shape[0] == num_msa_seq_before_pad, (
            f'Unexpected shape for MSA feature {k=}: {v.shape}.')
        pad_size[0][1] = self.msa_crop_size - v.shape[0]

      seq_axis = 1 if k in NUM_SEQ_NUM_RES_MSA_FEATURES else 0
      if k in frozenset(SEQ_FEATURES + NUM_SEQ_NUM_RES_MSA_FEATURES):
        assert v.shape[seq_axis] == num_res_before_pad, (
            f'Unexpected shape for MSA feature {k=}: {v.shape}.')
        pad_size[seq_axis][1] = self.crop_size - v.shape[seq_axis]

      features[k] = np.pad(v, pad_size, mode='constant')
    return features

  def crop(self, features: FeatureDict) -> FeatureDict:
    features = self._apply_crop(features)
    features = self._apply_pad(features)

    features['seq_length'] = np.asarray(self.crop_size, dtype=np.int32)
    features['num_alignments'] = np.asarray(self.msa_crop_size, dtype=np.int32)
    return features

  def _contiguous_crop_selection(self, chain: FeatureDict) -> np.ndarray:
    """Selects a contiguous crop around the variant position."""
    variant_pos = chain['variant_position'] + 1
    sequence_len = chain['seq_length']
    # Determine the crop by placing the variant in the center.
    if variant_pos > sequence_len // 2:
      crop_end = min(sequence_len, variant_pos + self.crop_size // 2)
      crop_start = max(0, crop_end - self.crop_size)
    else:
      crop_start = max(0, variant_pos - self.crop_size // 2)
      crop_end = min(sequence_len, crop_start + self.crop_size)
    keep_mask = np.zeros(chain['seq_length'], dtype=bool)
    keep_mask[crop_start:crop_end] = True
    return keep_mask

  def _spatial_crop_selection(self, chain: FeatureDict) -> np.ndarray:
    """Crops around the variant position."""
    all_positions = chain['all_atom_positions']
    all_masks = chain['all_atom_mask']
    variant_index = chain['variant_position']

    variant_positions = all_positions[variant_index]
    variant_mask = all_masks[variant_index]
    variant_positions = variant_positions[variant_mask.astype(bool)]

    distances = np.linalg.norm(
        all_positions[:, :, None] - variant_positions, axis=-1)
    # Do not include distance from masked atoms.
    distances = np.where(all_masks[:, :, None], distances, np.inf)
    # Sort by minimum over atoms in both residues.
    min_distances = np.min(distances, axis=(1, 2))

    # Take crop_size residues closest to the variant. The first argsort returns
    # a mapping from distance-rank (in ascending order) to residue-id. The
    # second returns the mapping from residue-id to distance-rank.
    include_mask = np.argsort(np.argsort(min_distances)) < self.crop_size
    included_res_ids = chain['residue_index'][include_mask]

    # Make the keep masks.
    chain_keep_mask = np.zeros(chain['seq_length'], dtype=bool)
    chain_keep_mask[included_res_ids] = True
    return chain_keep_mask

  def _choose_crop_region(self, chain: FeatureDict) -> np.ndarray:
    if (chain['all_atom_mask'] == 0).all():  # Structure not found.
      chain_keep_masks = self._contiguous_crop_selection(chain)
    else:
      chain_keep_masks = self._spatial_crop_selection(chain)

    total_residues = chain_keep_masks.sum()
    if total_residues == 0:
      raise ValueError('Selected keep masks are all false!')
    elif total_residues > self.crop_size:
      raise ValueError(f'Invalid crop: {total_residues=}, {self.crop_size=}.')
    return chain_keep_masks


def make_msa_features(msas: Sequence[parsers.Msa]) -> FeatureDict:
  """Constructs a feature dict of MSA features."""
  if not msas:
    raise ValueError('At least one MSA must be provided.')

  int_msa = []
  deletion_matrix = []
  species_ids = []
  seen_sequences = set()
  for msa_index, msa in enumerate(msas):
    if not msa:
      raise ValueError(f'MSA {msa_index} must contain at least one sequence.')
    for sequence_index, sequence in enumerate(msa.sequences):
      if sequence in seen_sequences:
        continue
      seen_sequences.add(sequence)
      int_msa.append([TOKENIZER[res] for res in sequence])
      deletion_matrix.append(msa.deletion_matrix[sequence_index])
      identifiers = msa_identifiers.get_identifiers(
          msa.descriptions[sequence_index])
      species_ids.append(identifiers.species_id.encode('utf-8'))

  num_res = len(msas[0].sequences[0])
  num_alignments = len(int_msa)
  features = {}
  features['deletion_matrix'] = np.array(deletion_matrix, dtype=np.float32)
  features['msa'] = np.array(int_msa, dtype=np.int32)
  features['num_alignments'] = np.array(
      [num_alignments] * num_res, dtype=np.int32)
  features['msa_mask'] = np.ones_like(features['msa'], dtype=bool)
  return features


def run_msa_tool(msa_runner, input_fasta_path: str, msa_out_path: str,
                 msa_format: str, use_precomputed_msas: bool,
                 max_sto_sequences: Optional[int] = None
                 ) -> Mapping[str, Any]:
  """Runs an MSA tool, checking if output already exists first."""
  if not use_precomputed_msas or not os.path.exists(msa_out_path):
    if msa_runner is None:
      raise ValueError(f'MSA results not found in {msa_out_path=} '
                       'and msa_runner is not set. Either pre-fill MSA cache '
                       'file or provide a valid jackhmmer_binary_path.')
    if msa_format == 'sto' and max_sto_sequences is not None:
      result = msa_runner.query(input_fasta_path, max_sto_sequences)[0]  # pytype: disable=wrong-arg-count
    else:
      result = msa_runner.query(input_fasta_path)[0]
    with open(msa_out_path, 'w') as f:
      f.write(result[msa_format])
  else:
    logging.warning('Reading MSA from file %s', msa_out_path)
    if msa_format == 'sto' and max_sto_sequences is not None:
      precomputed_msa = parsers.truncate_stockholm_msa(
          msa_out_path, max_sto_sequences)
      result = {'sto': precomputed_msa}
    else:
      with open(msa_out_path, 'r') as f:
        result = {msa_format: f.read()}
  return result


def variant_sequence_features(variant: ProteinVariant) -> FeatureDict:
  """Computes features for a single protein sequence."""

  def get_aatype(seq: str) -> np.ndarray:
    return np.array([
        residue_constants.restype_order_with_x[s] for s in seq], np.int32)

  variant_mask = np.zeros(len(variant.sequence), dtype=np.int32)
  variant_mask[variant.position - 1] = 1

  aatype = get_aatype(variant.reference_sequence)
  num_res = len(variant.sequence)
  features = {
      'seq_length': np.array(num_res, dtype=np.int32),
      'aatype': aatype,
      'seq_mask': np.ones_like(aatype, dtype=np.int32),
      'residue_index': np.arange(0, num_res, dtype=np.int32),
      'variant_aatype': get_aatype(variant.alternate_sequence),
      'variant_position': np.array(variant.position - 1, dtype=np.int32),
      'variant_mask': variant_mask,
  }
  if variant.pathogenicity is not None:
    features['pathogenicity'] = np.array(variant.pathogenicity,
                                         dtype=np.float32)
  return features


def make_variant_masked_msa(features: FeatureDict) -> FeatureDict:
  """Add masked mutant sequence in the second line of MSA."""
  # Remove the last line of the MSA to preserve cropped shape.
  features['msa'] = np.concatenate([
      features['msa'][:1],
      features['variant_aatype'][None],
      features['msa'][1:-1]])
  # For other features, duplicate the first row.
  for k in ['deletion_matrix', 'msa_mask']:
    features[k] = np.concatenate([
        features[k][:1], features[k][:1], features[k][1:-1]])

  features['cluster_bias_mask'] = np.zeros(features['msa'].shape[0])
  features['cluster_bias_mask'][:2] = 1.

  # Only mask target sequence at the variant position.
  features['msa'][1, features['variant_position']] = 22
  return features


def read_from_gcloud(source: str) -> str:
  """Download and read text file from gcloud."""
  with tempfile.NamedTemporaryFile(mode='w+t') as temp_file:
    cmd = ('gcloud', 'storage', 'cp', source, temp_file.name)

    logging.info('Launching subprocess "%s"', ' '.join(cmd))
    process = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    _, stderr = process.communicate()
    retcode = process.wait()

    if retcode:
      raise RuntimeError(
          'gcloud storage cp failed\nstderr:\n%s\n' % stderr.decode('utf-8'))

    with open(temp_file.name, 'r') as f:
      cif_string = f.read()
  return cif_string


def get_atom_positions(protein_id: str) -> FeatureDict:
  """Reads the positions predicted by AFDB. Used only to crop the sequence."""
  source_file = f'{AFDB_BUCKET}/AF-{protein_id}-F1-model_v4.cif'

  cif_string = read_from_gcloud(source_file)
  parsing_result = mmcif_parsing.parse(file_id=protein_id,
                                       mmcif_string=cif_string)
  if parsing_result.mmcif_object is None:
    raise ValueError('Unable to parse cif file.')

  chain_id = list(parsing_result.mmcif_object.chain_to_seqres.keys())[0]
  atom_positions, atom_mask = templates.get_atom_positions(
      parsing_result.mmcif_object, chain_id, max_ca_ca_distance=150.0)
  return {'all_atom_positions': atom_positions, 'all_atom_mask': atom_mask}


def get_empty_atom_positions(num_res: int) -> FeatureDict:
  return {'all_atom_positions': np.zeros([num_res, 37, 3], dtype=np.float32),
          'all_atom_mask': np.zeros([num_res, 37], dtype=np.int32)}


def make_msa_profile(batch: FeatureDict) -> FeatureDict:
  """Returns chains with the profile feature added."""
  def _one_hot(idx, n_classes, dtype=np.float32):
    """Numpy implementation of jax.nn.one_hot."""
    arr = np.concatenate([np.eye(n_classes, dtype=dtype),
                          np.zeros([1, n_classes], dtype=dtype)])
    idx = np.clip(idx, -1, n_classes)
    return arr[idx]

  batch['msa_profile'] = np.mean(_one_hot(batch['msa'], 30), axis=0)
  batch['deletion_mean'] = np.mean(batch['deletion_matrix'], axis=0)
  return batch


class DataPipeline:
  """Runs the alignment tools and assembles the input features."""

  def __init__(self,
               *,
               jackhmmer_binary_path: str,
               protein_sequence_file: str,
               uniref90_database_path: str,
               mgnify_database_path: str,
               small_bfd_database_path: str,
               mgnify_max_hits: int = 5_000,
               uniref_max_hits: int = 10_000,
               small_bfd_max_hits: int = 5_000,
               sequence_crop_size: int = 256,
               msa_crop_size: int = 2048,
               use_precomputed_msas: bool = True):
    """Initializes the data pipeline."""
    self.mgnify_max_hits = mgnify_max_hits
    self.uniref_max_hits = uniref_max_hits
    self.small_bfd_max_hits = small_bfd_max_hits
    self.use_precomputed_msas = use_precomputed_msas
    self.cropper = Cropper(crop_size=sequence_crop_size,
                           msa_crop_size=msa_crop_size)

    with open(protein_sequence_file) as f:
      input_fasta_str = f.read()
    input_seqs, input_descs = parsers.parse_fasta(input_fasta_str)
    self.protein_sequences = {k: v for k, v in zip(input_descs, input_seqs)}

    self.jackhmmer_uniref90_runner = None
    self.jackhmmer_small_bfd_runner = None
    self.jackhmmer_mgnify_runner = None
    if jackhmmer_binary_path:
      if uniref90_database_path:
        self.jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(
            binary_path=jackhmmer_binary_path,
            database_path=uniref90_database_path)
      if small_bfd_database_path:
        self.jackhmmer_small_bfd_runner = jackhmmer.Jackhmmer(
            binary_path=jackhmmer_binary_path,
            database_path=small_bfd_database_path)
      if mgnify_database_path:
        self.jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(
            binary_path=jackhmmer_binary_path,
            database_path=mgnify_database_path)

  def gather_msa_features(self, variant: ProteinVariant, msa_output_dir: str):
    """Runs all msa tools and returns the postprocessed features."""
    if not os.path.exists(msa_output_dir):
      os.makedirs(msa_output_dir)

    msa_results = []
    with tempfile.NamedTemporaryFile(mode='w+t', suffix='.fasta') as query_file:
      query_file.write(f'>{variant.protein_id}\n{variant.sequence}\n')
      query_file.seek(0)

      if self.jackhmmer_uniref90_runner is not None:
        jackhmmer_uniref90_result = run_msa_tool(
            msa_runner=self.jackhmmer_uniref90_runner,
            input_fasta_path=query_file.name,
            msa_out_path=os.path.join(msa_output_dir, 'uniref90_hits.sto'),
            msa_format='sto',
            use_precomputed_msas=self.use_precomputed_msas,
            max_sto_sequences=self.uniref_max_hits)
        uniref90_msa = parsers.parse_stockholm(jackhmmer_uniref90_result['sto'])
        logging.info('Uniref90 MSA size: %d sequences.', len(uniref90_msa))
        msa_results.append(uniref90_msa)

      if self.jackhmmer_mgnify_runner is not None:
        jackhmmer_mgnify_result = run_msa_tool(
            msa_runner=self.jackhmmer_mgnify_runner,
            input_fasta_path=query_file.name,
            msa_out_path=os.path.join(msa_output_dir, 'mgnify_hits.sto'),
            msa_format='sto',
            use_precomputed_msas=self.use_precomputed_msas,
            max_sto_sequences=self.mgnify_max_hits)
        mgnify_msa = parsers.parse_stockholm(jackhmmer_mgnify_result['sto'])
        logging.info('MGnify MSA size: %d sequences.', len(mgnify_msa))
        msa_results.append(mgnify_msa)

      if self.jackhmmer_small_bfd_runner is not None:
        jackhmmer_small_bfd_result = run_msa_tool(
            msa_runner=self.jackhmmer_small_bfd_runner,
            input_fasta_path=query_file.name,
            msa_out_path=os.path.join(msa_output_dir, 'small_bfd_hits.sto'),
            msa_format='sto',
            use_precomputed_msas=self.use_precomputed_msas,
            max_sto_sequences=self.small_bfd_max_hits)
        bfd_msa = parsers.parse_stockholm(jackhmmer_small_bfd_result['sto'])
        logging.info('BFD MSA size: %d sequences.', len(bfd_msa))
        msa_results.append(bfd_msa)

    msa_features = make_msa_features(msa_results)
    logging.info('Final (deduplicated) MSA size: %d sequences.',
                 msa_features['num_alignments'][0])
    return msa_features

  def process(self,
              *,
              protein_id: str,
              reference_aa: str,
              alternate_aa: str,
              position: int,
              msa_output_dir: str,
              pathogenicity: Optional[bool] = None,
              ) -> FeatureDict:
    """Runs alignment tools on the input sequence and creates features."""
    if protein_id not in self.protein_sequences:
      raise ValueError(f'{protein_id=} not in input protein_sequence_file')

    variant = ProteinVariant(sequence=self.protein_sequences[protein_id],
                             protein_id=protein_id,
                             position=position,
                             reference_aa=reference_aa,
                             alternate_aa=alternate_aa,
                             pathogenicity=pathogenicity)

    features = variant_sequence_features(variant=variant)
    features.update(self.gather_msa_features(variant, msa_output_dir))
    features = make_msa_profile(features)

    try:
      features.update(get_atom_positions(protein_id))
    except (ValueError, FileNotFoundError):
      logging.warning('Unable to find atom positions from AFDB.')
      features.update(get_empty_atom_positions(num_res=len(variant.sequence)))
    features = self.cropper.crop(features)

    # Correct variant_position after crop and size fix based on variant mask.
    features['variant_position'] = np.argmax(features['variant_mask'])
    # If variant is in second row of MSA, first row must be the target sequence.
    if not np.array_equal(features['msa'][0], features['aatype']):
      logging.warning('First row of MSA is not the target sequence diff=%s.',
                      np.sum(features['msa'][0] != features['aatype']))
      features['msa'][0] = features['aatype']
    # Mark that first and second row of MSA should not be affected by sampling.
    features['cluster_bias_mask'] = np.pad(
        np.zeros(features['msa'].shape[0] - 2), (2, 0), constant_values=1.)

    features = make_variant_masked_msa(features)

    return features

alphamissense/data/pipeline_missense.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Functions for getting templates and calculating template features."""

import numpy as np

from alphamissense.common import residue_constants
from alphamissense.data import mmcif_parsing


class Error(Exception):
  """Base class for exceptions."""


class NoChainsError(Error):
  """An error indicating that template mmCIF didn't have any chains."""


class SequenceNotInTemplateError(Error):
  """An error indicating that template mmCIF didn't contain the sequence."""


class NoAtomDataInTemplateError(Error):
  """An error indicating that template mmCIF didn't contain atom positions."""


class TemplateAtomMaskAllZerosError(Error):
  """An error indicating that template mmCIF had all atom positions masked."""


class QueryToTemplateAlignError(Error):
  """An error indicating that the query can't be aligned to the template."""


class CaDistanceError(Error):
  """An error indicating that a CA atom distance exceeds a threshold."""


class MultipleChainsError(Error):
  """An error indicating that multiple chains were found for a given ID."""


# Prefilter exceptions.
class PrefilterError(Exception):
  """A base class for template prefilter exceptions."""


class DateError(PrefilterError):
  """An error indicating that the hit date was after the max allowed date."""


class AlignRatioError(PrefilterError):
  """An error indicating that the hit align ratio to the query was too small."""


class DuplicateError(PrefilterError):
  """An error indicating that the hit was an exact subsequence of the query."""


class LengthError(PrefilterError):
  """An error indicating that the hit was too short."""


TEMPLATE_FEATURES = {
    'template_aatype': np.float32,
    'template_all_atom_masks': np.float32,
    'template_all_atom_positions': np.float32,
    'template_domain_names': object,
    'template_sequence': object,
    'template_sum_probs': np.float32,
}


def _check_residue_distances(all_positions: np.ndarray,
                             all_positions_mask: np.ndarray,
                             max_ca_ca_distance: float):
  """Checks if the distance between unmasked neighbor residues is ok."""
  ca_position = residue_constants.atom_order['CA']
  prev_is_unmasked = False
  prev_calpha = None
  for i, (coords, mask) in enumerate(zip(all_positions, all_positions_mask)):
    this_is_unmasked = bool(mask[ca_position])
    if this_is_unmasked:
      this_calpha = coords[ca_position]
      if prev_is_unmasked:
        distance = np.linalg.norm(this_calpha - prev_calpha)
        if distance > max_ca_ca_distance:
          raise CaDistanceError(
              f'The distance between residues {i} and {i + 1} is '
              f'{distance=} > limit {max_ca_ca_distance=}.')
      prev_calpha = this_calpha
    prev_is_unmasked = this_is_unmasked


def get_atom_positions(
    mmcif_object: mmcif_parsing.MmcifObject,
    auth_chain_id: str,
    max_ca_ca_distance: float) -> tuple[np.ndarray, np.ndarray]:
  """Gets atom positions and mask from a list of Biopython Residues."""
  num_res = len(mmcif_object.chain_to_seqres[auth_chain_id])

  relevant_chains = [c for c in mmcif_object.structure.get_chains()
                     if c.id == auth_chain_id]
  if len(relevant_chains) != 1:
    raise MultipleChainsError(
        f'Expected exactly one chain in structure with id {auth_chain_id}.')
  chain = relevant_chains[0]

  all_positions = np.zeros([num_res, residue_constants.atom_type_num, 3])
  all_positions_mask = np.zeros([num_res, residue_constants.atom_type_num],
                                dtype=np.int64)
  for res_index in range(num_res):
    pos = np.zeros([residue_constants.atom_type_num, 3], dtype=np.float32)
    mask = np.zeros([residue_constants.atom_type_num], dtype=np.float32)
    res_at_position = mmcif_object.seqres_to_structure[auth_chain_id][res_index]
    if not res_at_position.is_missing:
      assert res_at_position.position is not None
      res = chain[(res_at_position.hetflag,
                   res_at_position.position.residue_number,
                   res_at_position.position.insertion_code)]
      for atom in res.get_atoms():
        atom_name = atom.get_name()
        x, y, z = atom.get_coord()
        if atom_name in residue_constants.atom_order.keys():
          pos[residue_constants.atom_order[atom_name]] = [x, y, z]
          mask[residue_constants.atom_order[atom_name]] = 1.0
        elif atom_name.upper() == 'SE' and res.get_resname() == 'MSE':
          # Put the coordinates of the selenium atom in the sulphur column.
          pos[residue_constants.atom_order['SD']] = [x, y, z]
          mask[residue_constants.atom_order['SD']] = 1.0

      # Fix naming errors in arginine residues where NH2 is incorrectly
      # assigned to be closer to CD than NH1.
      cd = residue_constants.atom_order['CD']
      nh1 = residue_constants.atom_order['NH1']
      nh2 = residue_constants.atom_order['NH2']
      if (res.get_resname() == 'ARG' and
          all(mask[atom_index] for atom_index in (cd, nh1, nh2)) and
          (np.linalg.norm(pos[nh1] - pos[cd]) >
           np.linalg.norm(pos[nh2] - pos[cd]))):
        pos[nh1], pos[nh2] = pos[nh2].copy(), pos[nh1].copy()
        mask[nh1], mask[nh2] = mask[nh2].copy(), mask[nh1].copy()

    all_positions[res_index] = pos
    all_positions_mask[res_index] = mask
  _check_residue_distances(
      all_positions, all_positions_mask, max_ca_ca_distance)
  return all_positions, all_positions_mask

alphamissense/data/templates.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Library to run Jackhmmer from Python."""

from concurrent import futures
import glob
import os
import subprocess
from typing import Any, Callable, Mapping, Optional, Sequence
from urllib import request

from absl import logging

# Internal import (7716).
from alphamissense.data import parsers
from alphamissense.data.tools import utils


class Jackhmmer:
  """Python wrapper of the Jackhmmer binary."""

  def __init__(self,
               *,
               binary_path: str,
               database_path: str,
               n_cpu: int = 8,
               n_iter: int = 1,
               e_value: float = 0.0001,
               z_value: Optional[int] = None,
               get_tblout: bool = False,
               filter_f1: float = 0.0005,
               filter_f2: float = 0.00005,
               filter_f3: float = 0.0000005,
               incdom_e: Optional[float] = None,
               dom_e: Optional[float] = None,
               num_streamed_chunks: Optional[int] = None,
               streaming_callback: Optional[Callable[[int], None]] = None):
    """Initializes the Python Jackhmmer wrapper.

    Args:
      binary_path: The path to the jackhmmer executable.
      database_path: The path to the jackhmmer database (FASTA format).
      n_cpu: The number of CPUs to give Jackhmmer.
      n_iter: The number of Jackhmmer iterations.
      e_value: The E-value, see Jackhmmer docs for more details.
      z_value: The Z-value, see Jackhmmer docs for more details.
      get_tblout: Whether to save tblout string.
      filter_f1: MSV and biased composition pre-filter, set to >1.0 to turn off.
      filter_f2: Viterbi pre-filter, set to >1.0 to turn off.
      filter_f3: Forward pre-filter, set to >1.0 to turn off.
      incdom_e: Domain e-value criterion for inclusion of domains in MSA/next
        round.
      dom_e: Domain e-value criterion for inclusion in tblout.
      num_streamed_chunks: Number of database chunks to stream over.
      streaming_callback: Callback function run after each chunk iteration with
        the iteration number as argument.
    """
    self.binary_path = binary_path
    self.database_path = database_path
    self.num_streamed_chunks = num_streamed_chunks

    if not os.path.exists(self.database_path) and num_streamed_chunks is None:
      logging.error('Could not find Jackhmmer database %s', database_path)
      raise ValueError(f'Could not find Jackhmmer database {database_path}')

    self.n_cpu = n_cpu
    self.n_iter = n_iter
    self.e_value = e_value
    self.z_value = z_value
    self.filter_f1 = filter_f1
    self.filter_f2 = filter_f2
    self.filter_f3 = filter_f3
    self.incdom_e = incdom_e
    self.dom_e = dom_e
    self.get_tblout = get_tblout
    self.streaming_callback = streaming_callback

  def _query_chunk(self,
                   input_fasta_path: str,
                   database_path: str,
                   max_sequences: Optional[int] = None) -> Mapping[str, Any]:
    """Queries the database chunk using Jackhmmer."""
    with utils.tmpdir_manager() as query_tmp_dir:
      sto_path = os.path.join(query_tmp_dir, 'output.sto')

      # The F1/F2/F3 are the expected proportion to pass each of the filtering
      # stages (which get progressively more expensive), reducing these
      # speeds up the pipeline at the expensive of sensitivity.  They are
      # currently set very low to make querying Mgnify run in a reasonable
      # amount of time.
      cmd_flags = [
          # Don't pollute stdout with Jackhmmer output.
          '-o', '/dev/null',
          '-A', sto_path,
          '--noali',
          '--F1', str(self.filter_f1),
          '--F2', str(self.filter_f2),
          '--F3', str(self.filter_f3),
          '--incE', str(self.e_value),
          # Report only sequences with E-values <= x in per-sequence output.
          '-E', str(self.e_value),
          '--cpu', str(self.n_cpu),
          '-N', str(self.n_iter)
      ]
      if self.get_tblout:
        tblout_path = os.path.join(query_tmp_dir, 'tblout.txt')
        cmd_flags.extend(['--tblout', tblout_path])

      if self.z_value:
        cmd_flags.extend(['-Z', str(self.z_value)])

      if self.dom_e is not None:
        cmd_flags.extend(['--domE', str(self.dom_e)])

      if self.incdom_e is not None:
        cmd_flags.extend(['--incdomE', str(self.incdom_e)])

      cmd = [self.binary_path] + cmd_flags + [input_fasta_path,
                                              database_path]

      logging.info('Launching subprocess "%s"', ' '.join(cmd))
      process = subprocess.Popen(
          cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      with utils.timing(
          f'Jackhmmer ({os.path.basename(database_path)}) query'):
        _, stderr = process.communicate()
        retcode = process.wait()

      if retcode:
        raise RuntimeError(
            'Jackhmmer failed\nstderr:\n%s\n' % stderr.decode('utf-8'))

      # Get e-values for each target name
      tbl = ''
      if self.get_tblout:
        with open(tblout_path) as f:
          tbl = f.read()

      if max_sequences is None:
        with open(sto_path) as f:
          sto = f.read()
      else:
        sto = parsers.truncate_stockholm_msa(sto_path, max_sequences)

    raw_output = dict(
        sto=sto,
        tbl=tbl,
        stderr=stderr,
        n_iter=self.n_iter,
        e_value=self.e_value)

    return raw_output

  def query(self,
            input_fasta_path: str,
            max_sequences: Optional[int] = None) -> Sequence[Mapping[str, Any]]:
    """Queries the database using Jackhmmer."""
    return self.query_multiple([input_fasta_path], max_sequences)[0]

  def query_multiple(
      self,
      input_fasta_paths: Sequence[str],
      max_sequences: Optional[int] = None,
  ) -> Sequence[Sequence[Mapping[str, Any]]]:
    """Queries the database for multiple queries using Jackhmmer."""
    if self.num_streamed_chunks is None:
      single_chunk_results = []
      for input_fasta_path in input_fasta_paths:
        single_chunk_results.append([self._query_chunk(
            input_fasta_path, self.database_path, max_sequences)])
      return single_chunk_results

    db_basename = os.path.basename(self.database_path)
    db_remote_chunk = lambda db_idx: f'{self.database_path}.{db_idx}'
    db_local_chunk = lambda db_idx: f'/tmp/ramdisk/{db_basename}.{db_idx}'

    # Remove existing files to prevent OOM
    for f in glob.glob(db_local_chunk('[0-9]*')):
      try:
        os.remove(f)
      except OSError:
        print(f'OSError while deleting {f}')

    # Download the (i+1)-th chunk while Jackhmmer is running on the i-th chunk
    with futures.ThreadPoolExecutor(max_workers=2) as executor:
      chunked_outputs = [[] for _ in range(len(input_fasta_paths))]
      for i in range(1, self.num_streamed_chunks + 1):
        # Copy the chunk locally
        if i == 1:
          future = executor.submit(
              request.urlretrieve, db_remote_chunk(i), db_local_chunk(i))
        if i < self.num_streamed_chunks:
          next_future = executor.submit(
              request.urlretrieve, db_remote_chunk(i+1), db_local_chunk(i+1))

        # Run Jackhmmer with the chunk
        future.result()
        for fasta_index, input_fasta_path in enumerate(input_fasta_paths):
          chunked_outputs[fasta_index].append(self._query_chunk(
              input_fasta_path, db_local_chunk(i), max_sequences))
        # Remove the local copy of the chunk
        os.remove(db_local_chunk(i))
        # Do not set next_future for the last chunk so that this works even for
        # databases with only 1 chunk.
        if i < self.num_streamed_chunks:
          future = next_future
        if self.streaming_callback:
          self.streaming_callback(i)
    return chunked_outputs

alphamissense/data/tools/jackhmmer.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""AlphaMissense model."""

alphamissense/model/__init__.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Ops for all atom representations."""

from typing import Optional

import jax
import jax.numpy as jnp
import numpy as np
from alphamissense.common import residue_constants
from alphamissense.model import geometry
from alphamissense.model import utils


def squared_difference(x, y):
  return jnp.square(x - y)


def _make_chi_atom_indices():
  """Returns atom indices needed to compute chi angles for all residue types.

  Returns:
    A tensor of shape [residue_types=21, chis=4, atoms=4]. The residue types are
    in the order specified in residue_constants.restypes + unknown residue type
    at the end. For chi angles which are not defined on the residue, the
    positions indices are by default set to 0.
  """
  chi_atom_indices = []
  for residue_name in residue_constants.restypes:
    residue_name = residue_constants.restype_1to3[residue_name]
    residue_chi_angles = residue_constants.chi_angles_atoms[residue_name]
    atom_indices = []
    for chi_angle in residue_chi_angles:
      atom_indices.append(
          [residue_constants.atom_order[atom] for atom in chi_angle])
    for _ in range(4 - len(atom_indices)):
      atom_indices.append([0, 0, 0, 0])  # For chi angles not defined on the AA.
    chi_atom_indices.append(atom_indices)

  chi_atom_indices.append([[0, 0, 0, 0]] * 4)  # For UNKNOWN residue.

  return np.array(chi_atom_indices)


def _make_renaming_matrices():
  """Matrices to map atoms to symmetry partners in ambiguous case."""
  # As the atom naming is ambiguous for 7 of the 20 amino acids, provide
  # alternative groundtruth coordinates where the naming is swapped
  restype_3 = [
      residue_constants.restype_1to3[res] for res in residue_constants.restypes
  ]
  restype_3 += ['UNK']
  # Matrices for renaming ambiguous atoms.
  all_matrices = {res: np.eye(14, dtype=np.float32) for res in restype_3}
  for resname, swap in residue_constants.residue_atom_renaming_swaps.items():
    correspondences = np.arange(14)
    for source_atom_swap, target_atom_swap in swap.items():
      source_index = residue_constants.restype_name_to_atom14_names[
          resname].index(source_atom_swap)
      target_index = residue_constants.restype_name_to_atom14_names[
          resname].index(target_atom_swap)
      correspondences[source_index] = target_index
      correspondences[target_index] = source_index
      renaming_matrix = np.zeros((14, 14), dtype=np.float32)
      for index, correspondence in enumerate(correspondences):
        renaming_matrix[index, correspondence] = 1.
    all_matrices[resname] = renaming_matrix.astype(np.float32)
  renaming_matrices = np.stack([all_matrices[restype] for restype in restype_3])
  return renaming_matrices


def _make_restype_atom37_mask():
  """Mask of which atoms are present for which residue type in atom37."""
  # create the corresponding mask
  restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
  for restype, restype_letter in enumerate(residue_constants.restypes):
    restype_name = residue_constants.restype_1to3[restype_letter]
    atom_names = residue_constants.residue_atoms[restype_name]
    for atom_name in atom_names:
      atom_type = residue_constants.atom_order[atom_name]
      restype_atom37_mask[restype, atom_type] = 1
  return restype_atom37_mask


def _make_restype_atom14_mask():
  """Mask of which atoms are present for which residue type in atom14."""
  restype_atom14_mask = []

  for rt in residue_constants.restypes:
    atom_names = residue_constants.restype_name_to_atom14_names[
        residue_constants.restype_1to3[rt]]
    restype_atom14_mask.append([(1. if name else 0.) for name in atom_names])

  restype_atom14_mask.append([0.] * 14)
  restype_atom14_mask = np.array(restype_atom14_mask, dtype=np.float32)
  return restype_atom14_mask


def _make_restype_atom37_to_atom14():
  """Map from atom37 to atom14 per residue type."""
  restype_atom37_to_atom14 = []  # mapping (restype, atom37) --> atom14
  for rt in residue_constants.restypes:
    atom_names = residue_constants.restype_name_to_atom14_names[
        residue_constants.restype_1to3[rt]]
    atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
    restype_atom37_to_atom14.append([
        (atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0)
        for name in residue_constants.atom_types
    ])

  restype_atom37_to_atom14.append([0] * 37)
  restype_atom37_to_atom14 = np.array(restype_atom37_to_atom14, dtype=np.int32)
  return restype_atom37_to_atom14


def _make_restype_atom14_to_atom37():
  """Map from atom14 to atom37 per residue type."""
  restype_atom14_to_atom37 = []  # mapping (restype, atom14) --> atom37
  for rt in residue_constants.restypes:
    atom_names = residue_constants.restype_name_to_atom14_names[
        residue_constants.restype_1to3[rt]]
    restype_atom14_to_atom37.append([
        (residue_constants.atom_order[name] if name else 0)
        for name in atom_names
    ])
  # Add dummy mapping for restype 'UNK'
  restype_atom14_to_atom37.append([0] * 14)
  restype_atom14_to_atom37 = np.array(restype_atom14_to_atom37, dtype=np.int32)
  return restype_atom14_to_atom37


def _make_restype_atom14_is_ambiguous():
  """Mask which atoms are ambiguous in atom14."""
  # create an ambiguous atoms mask.  shape: (21, 14)
  restype_atom14_is_ambiguous = np.zeros((21, 14), dtype=np.float32)
  for resname, swap in residue_constants.residue_atom_renaming_swaps.items():
    for atom_name1, atom_name2 in swap.items():
      restype = residue_constants.restype_order[
          residue_constants.restype_3to1[resname]]
      atom_idx1 = residue_constants.restype_name_to_atom14_names[resname].index(
          atom_name1)
      atom_idx2 = residue_constants.restype_name_to_atom14_names[resname].index(
          atom_name2)
      restype_atom14_is_ambiguous[restype, atom_idx1] = 1
      restype_atom14_is_ambiguous[restype, atom_idx2] = 1

  return restype_atom14_is_ambiguous


def _make_restype_rigidgroup_base_atom37_idx():
  """Create Map from rigidgroups to atom37 indices."""
  # Create an array with the atom names.
  # shape (num_restypes, num_rigidgroups, 3_atoms): (21, 8, 3)
  base_atom_names = np.full([21, 8, 3], '', dtype=object)

  # 0: backbone frame
  base_atom_names[:, 0, :] = ['C', 'CA', 'N']

  # 3: 'psi-group'
  base_atom_names[:, 3, :] = ['CA', 'C', 'O']

  # 4,5,6,7: 'chi1,2,3,4-group'
  for restype, restype_letter in enumerate(residue_constants.restypes):
    resname = residue_constants.restype_1to3[restype_letter]
    for chi_idx in range(4):
      if residue_constants.chi_angles_mask[restype][chi_idx]:
        atom_names = residue_constants.chi_angles_atoms[resname][chi_idx]
        base_atom_names[restype, chi_idx + 4, :] = atom_names[1:]

  # Translate atom names into atom37 indices.
  lookuptable = residue_constants.atom_order.copy()
  lookuptable[''] = 0
  restype_rigidgroup_base_atom37_idx = np.vectorize(lambda x: lookuptable[x])(
      base_atom_names)
  return restype_rigidgroup_base_atom37_idx


CHI_ATOM_INDICES = _make_chi_atom_indices()
RENAMING_MATRICES = _make_renaming_matrices()
RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37()
RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14()
RESTYPE_ATOM37_MASK = _make_restype_atom37_mask()
RESTYPE_ATOM14_MASK = _make_restype_atom14_mask()
RESTYPE_ATOM14_IS_AMBIGUOUS = _make_restype_atom14_is_ambiguous()
RESTYPE_RIGIDGROUP_BASE_ATOM37_IDX = _make_restype_rigidgroup_base_atom37_idx()

# Create mask for existing rigid groups.
RESTYPE_RIGIDGROUP_MASK = np.zeros([21, 8], dtype=np.float32)
RESTYPE_RIGIDGROUP_MASK[:, 0] = 1
RESTYPE_RIGIDGROUP_MASK[:, 3] = 1
RESTYPE_RIGIDGROUP_MASK[:20, 4:] = residue_constants.chi_angles_mask


def get_atom37_mask(aatype):
  return utils.batched_gather(jnp.asarray(RESTYPE_ATOM37_MASK), aatype)


def get_atom14_mask(aatype):
  return utils.batched_gather(jnp.asarray(RESTYPE_ATOM14_MASK), aatype)


def get_atom14_is_ambiguous(aatype):
  return utils.batched_gather(jnp.asarray(RESTYPE_ATOM14_IS_AMBIGUOUS), aatype)


def get_atom14_to_atom37_map(aatype):
  return utils.batched_gather(jnp.asarray(RESTYPE_ATOM14_TO_ATOM37), aatype)


def get_atom37_to_atom14_map(aatype):
  return utils.batched_gather(jnp.asarray(RESTYPE_ATOM37_TO_ATOM14), aatype)


def atom14_to_atom37(atom14_data: jnp.ndarray,  # (N, 14, ...)
                     aatype: jnp.ndarray
                    ) -> jnp.ndarray:  # (N, 37, ...)
  """Convert atom14 to atom37 representation."""
  assert len(atom14_data.shape) in [2, 3]
  idx_atom37_to_atom14 = get_atom37_to_atom14_map(aatype)
  atom37_data = utils.batched_gather(
      atom14_data, idx_atom37_to_atom14, batch_dims=1)
  atom37_mask = get_atom37_mask(aatype)
  if len(atom14_data.shape) == 2:
    atom37_data *= atom37_mask
  elif len(atom14_data.shape) == 3:
    atom37_data *= atom37_mask[:, :, None].astype(atom37_data.dtype)
  return atom37_data


def atom37_to_atom14(aatype, all_atom_pos, all_atom_mask):
  """Convert Atom37 positions to Atom14 positions."""
  residx_atom14_to_atom37 = utils.batched_gather(
      jnp.asarray(RESTYPE_ATOM14_TO_ATOM37), aatype)
  atom14_mask = utils.batched_gather(
      all_atom_mask, residx_atom14_to_atom37, batch_dims=1).astype(jnp.float32)
  # create a mask for known groundtruth positions
  atom14_mask *= utils.batched_gather(jnp.asarray(RESTYPE_ATOM14_MASK), aatype)
  # gather the groundtruth positions
  atom14_positions = jax.tree_map(
      lambda x: utils.batched_gather(x, residx_atom14_to_atom37, batch_dims=1),
      all_atom_pos)
  atom14_positions = atom14_mask * atom14_positions
  return atom14_positions, atom14_mask


def get_alt_atom14(aatype, positions: geometry.Vec3Array, mask):
  """Get alternative atom14 positions."""
  # pick the transformation matrices for the given residue sequence
  # shape (num_res, 14, 14)
  renaming_transform = utils.batched_gather(
      jnp.asarray(RENAMING_MATRICES), aatype)

  alternative_positions = jax.tree_map(
      lambda x: jnp.sum(x, axis=1), positions[:, :, None] * renaming_transform)

  # Create the mask for the alternative ground truth (differs from the
  # ground truth mask, if only one of the atoms in an ambiguous pair has a
  # ground truth position)
  alternative_mask = jnp.sum(mask[..., None] * renaming_transform, axis=1)

  return alternative_positions, alternative_mask


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))

  # The frames for ambiguous rigid groups are just rotated by 180 degree around
  # the x-axis. The ambiguous group is always the last chi-group.
  restype_rigidgroup_is_ambiguous = np.zeros([21, 8], dtype=np.float32)
  restype_rigidgroup_rots = np.tile(np.eye(3, dtype=np.float32), [21, 8, 1, 1])

  for resname, _ in residue_constants.residue_atom_renaming_swaps.items():
    restype = residue_constants.restype_order[
        residue_constants.restype_3to1[resname]]
    chi_idx = int(sum(residue_constants.chi_angles_mask[restype]) - 1)
    restype_rigidgroup_is_ambiguous[restype, chi_idx + 4] = 1
    restype_rigidgroup_rots[restype, chi_idx + 4, 1, 1] = -1
    restype_rigidgroup_rots[restype, chi_idx + 4, 2, 2] = -1

  # Gather the ambiguity information for each residue.
  residx_rigidgroup_is_ambiguous = utils.batched_gather(
      restype_rigidgroup_is_ambiguous, aatype)
  ambiguity_rot = utils.batched_gather(restype_rigidgroup_rots, aatype)
  ambiguity_rot = geometry.Rot3Array.from_array(ambiguity_rot)

  # Create the alternative ground truth frames.
  alt_gt_frames = gt_frames.compose_rotation(ambiguity_rot)

  fix_shape = lambda x: jnp.reshape(x, aatype_in_shape + (8,))

  # reshape back to original residue layout
  gt_frames = jax.tree_map(fix_shape, gt_frames)
  gt_exists = fix_shape(gt_exists)
  group_exists = fix_shape(group_exists)
  residx_rigidgroup_is_ambiguous = fix_shape(residx_rigidgroup_is_ambiguous)
  alt_gt_frames = jax.tree_map(fix_shape, alt_gt_frames)

  return {
      'rigidgroups_gt_frames': gt_frames,  # Rigid (..., 8)
      'rigidgroups_gt_exists': gt_exists,  # (..., 8)
      'rigidgroups_group_exists': group_exists,  # (..., 8)
      'rigidgroups_group_is_ambiguous':
          residx_rigidgroup_is_ambiguous,  # (..., 8)
      'rigidgroups_alt_gt_frames': alt_gt_frames,  # Rigid (..., 8)
  }


def torsion_angles_to_frames(
    aatype: jnp.ndarray,  # (N)
    backb_to_global: geometry.Rigid3Array,  # (N)
    torsion_angles_sin_cos: jnp.ndarray  # (N, 7, 2)
) -> geometry.Rigid3Array:  # (N, 8)
  """Compute rigid group frames from torsion angles."""
  assert len(aatype.shape) == 1, (
      f'Expected array of rank 1, got array with shape: {aatype.shape}.')
  assert len(backb_to_global.rotation.shape) == 1, (
      f'Expected array of rank 1, got array with shape: '
      f'{backb_to_global.rotation.shape}')
  assert len(torsion_angles_sin_cos.shape) == 3, (
      f'Expected array of rank 3, got array with shape: '
      f'{torsion_angles_sin_cos.shape}')
  assert torsion_angles_sin_cos.shape[1] == 7, (
      f'wrong shape {torsion_angles_sin_cos.shape}')
  assert torsion_angles_sin_cos.shape[2] == 2, (
      f'wrong shape {torsion_angles_sin_cos.shape}')

  # Gather the default frames for all rigid groups.
  # geometry.Rigid3Array with shape (N, 8)
  m = utils.batched_gather(residue_constants.restype_rigid_group_default_frame,
                           aatype)
  default_frames = geometry.Rigid3Array.from_array4x4(m)

  # Create the rotation matrices according to the given angles (each frame is
  # defined such that its rotation is around the x-axis).
  sin_angles = torsion_angles_sin_cos[..., 0]
  cos_angles = torsion_angles_sin_cos[..., 1]

  # insert zero rotation for backbone group.
  num_residues, = aatype.shape
  sin_angles = jnp.concatenate([jnp.zeros([num_residues, 1]), sin_angles],
                               axis=-1)
  cos_angles = jnp.concatenate([jnp.ones([num_residues, 1]), cos_angles],
                               axis=-1)
  zeros = jnp.zeros_like(sin_angles)
  ones = jnp.ones_like(sin_angles)

  # all_rots are geometry.Rot3Array with shape (N, 8)
  all_rots = geometry.Rot3Array(ones, zeros, zeros,
                                zeros, cos_angles, -sin_angles,
                                zeros, sin_angles, cos_angles)

  # Apply rotations to the frames.
  all_frames = default_frames.compose_rotation(all_rots)

  # chi2, chi3, and chi4 frames do not transform to the backbone frame but to
  # the previous frame. So chain them up accordingly.

  chi1_frame_to_backb = all_frames[:, 4]
  chi2_frame_to_backb = chi1_frame_to_backb @ all_frames[:, 5]
  chi3_frame_to_backb = chi2_frame_to_backb @ all_frames[:, 6]
  chi4_frame_to_backb = chi3_frame_to_backb @ all_frames[:, 7]

  all_frames_to_backb = jax.tree_map(
      lambda *x: jnp.concatenate(x, axis=-1), all_frames[:, 0:5],
      chi2_frame_to_backb[:, None], chi3_frame_to_backb[:, None],
      chi4_frame_to_backb[:, None])

  # Create the global frames.
  # shape (N, 8)
  all_frames_to_global = backb_to_global[:, None] @ all_frames_to_backb

  return all_frames_to_global


def frames_and_literature_positions_to_atom14_pos(
    aatype: jnp.ndarray,  # (N)
    all_frames_to_global: geometry.Rigid3Array  # (N, 8)
) -> geometry.Vec3Array:  # (N, 14)
  """Put atom literature positions (atom14 encoding) in each rigid group."""

  # Pick the appropriate transform for every atom.
  residx_to_group_idx = utils.batched_gather(
      residue_constants.restype_atom14_to_rigid_group, aatype)
  group_mask = jax.nn.one_hot(
      residx_to_group_idx, num_classes=8)  # shape (N, 14, 8)

  # geometry.Rigid3Array with shape (N, 14)
  map_atoms_to_global = jax.tree_map(
      lambda x: jnp.sum(x[:, None, :] * group_mask, axis=-1),
      all_frames_to_global)

  # Gather the literature atom positions for each residue.
  # geometry.Vec3Array with shape (N, 14)
  lit_positions = geometry.Vec3Array.from_array(
      utils.batched_gather(
          residue_constants.restype_atom14_rigid_group_positions, aatype))

  # Transform each atom from its local frame to the global frame.
  # geometry.Vec3Array with shape (N, 14)
  pred_positions = map_atoms_to_global.apply_to_point(lit_positions)

  # Mask out non-existing atoms.
  mask = utils.batched_gather(residue_constants.restype_atom14_mask, aatype)
  pred_positions = pred_positions * mask

  return pred_positions


def extreme_ca_ca_distance_violations(
    positions: geometry.Vec3Array,  # (N, 37(14))
    mask: jnp.ndarray,  # (N, 37(14))
    residue_index: jnp.ndarray,  # (N)
    max_angstrom_tolerance=1.5
    ) -> jnp.ndarray:
  """Counts residues whose Ca is a large distance from its neighbor."""
  this_ca_pos = positions[:-1, 1]  # (N - 1,)
  this_ca_mask = mask[:-1, 1]         # (N - 1)
  next_ca_pos = positions[1:, 1]  # (N - 1,)
  next_ca_mask = mask[1:, 1]  # (N - 1)
  has_no_gap_mask = ((residue_index[1:] - residue_index[:-1]) == 1.0).astype(
      jnp.float32)
  ca_ca_distance = geometry.euclidean_distance(this_ca_pos, next_ca_pos, 1e-6)
  violations = (ca_ca_distance -
                residue_constants.ca_ca) > max_angstrom_tolerance
  mask = this_ca_mask * next_ca_mask * has_no_gap_mask
  return utils.mask_mean(mask=mask, value=violations)


def between_residue_bond_loss(
    pred_atom_positions: geometry.Vec3Array,  # (N, 37(14))
    pred_atom_mask: jnp.ndarray,  # (N, 37(14))
    residue_index: jnp.ndarray,  # (N)
    aatype: jnp.ndarray,  # (N)
    tolerance_factor_soft=12.0,
    tolerance_factor_hard=12.0) -> dict[str, jnp.ndarray]:
  """Flat-bottom loss to penalize structural violations between residues."""
  assert len(pred_atom_positions.shape) == 2
  assert len(pred_atom_mask.shape) == 2
  assert len(residue_index.shape) == 1
  assert len(aatype.shape) == 1

  # Get the positions of the relevant backbone atoms.
  this_ca_pos = pred_atom_positions[:-1, 1]  # (N - 1)
  this_ca_mask = pred_atom_mask[:-1, 1]         # (N - 1)
  this_c_pos = pred_atom_positions[:-1, 2]  # (N - 1)
  this_c_mask = pred_atom_mask[:-1, 2]          # (N - 1)
  next_n_pos = pred_atom_positions[1:, 0]  # (N - 1)
  next_n_mask = pred_atom_mask[1:, 0]           # (N - 1)
  next_ca_pos = pred_atom_positions[1:, 1]  # (N - 1)
  next_ca_mask = pred_atom_mask[1:, 1]          # (N - 1)
  has_no_gap_mask = ((residue_index[1:] - residue_index[:-1]) == 1.0).astype(
      jnp.float32)

  # Compute loss for the C--N bond.
  c_n_bond_length = geometry.euclidean_distance(this_c_pos, next_n_pos, 1e-6)

  # The C-N bond to proline has slightly different length because of the ring.
  next_is_proline = (
      aatype[1:] == residue_constants.restype_order['P']).astype(jnp.float32)
  gt_length = (
      (1. - next_is_proline) * residue_constants.between_res_bond_length_c_n[0]
      + next_is_proline * residue_constants.between_res_bond_length_c_n[1])
  gt_stddev = (
      (1. - next_is_proline) *
      residue_constants.between_res_bond_length_stddev_c_n[0] +
      next_is_proline * residue_constants.between_res_bond_length_stddev_c_n[1])
  c_n_bond_length_error = jnp.sqrt(1e-6 +
                                   jnp.square(c_n_bond_length - gt_length))
  c_n_loss_per_residue = jax.nn.relu(
      c_n_bond_length_error - tolerance_factor_soft * gt_stddev)
  mask = this_c_mask * next_n_mask * has_no_gap_mask
  c_n_loss = jnp.sum(mask * c_n_loss_per_residue) / (jnp.sum(mask) + 1e-6)
  c_n_violation_mask = mask * (
      c_n_bond_length_error > (tolerance_factor_hard * gt_stddev))

  # Compute loss for the angles.
  c_ca_unit_vec = (this_ca_pos - this_c_pos).normalized(1e-6)
  c_n_unit_vec = (next_n_pos - this_c_pos) / c_n_bond_length
  n_ca_unit_vec = (next_ca_pos - next_n_pos).normalized(1e-6)

  ca_c_n_cos_angle = c_ca_unit_vec.dot(c_n_unit_vec)
  gt_angle = residue_constants.between_res_cos_angles_ca_c_n[0]
  gt_stddev = residue_constants.between_res_bond_length_stddev_c_n[0]
  ca_c_n_cos_angle_error = jnp.sqrt(
      1e-6 + jnp.square(ca_c_n_cos_angle - gt_angle))
  ca_c_n_loss_per_residue = jax.nn.relu(
      ca_c_n_cos_angle_error - tolerance_factor_soft * gt_stddev)
  mask = this_ca_mask * this_c_mask * next_n_mask * has_no_gap_mask
  ca_c_n_loss = jnp.sum(mask * ca_c_n_loss_per_residue) / (jnp.sum(mask) + 1e-6)
  ca_c_n_violation_mask = mask * (ca_c_n_cos_angle_error >
                                  (tolerance_factor_hard * gt_stddev))

  c_n_ca_cos_angle = (-c_n_unit_vec).dot(n_ca_unit_vec)
  gt_angle, gt_stddev = residue_constants.between_res_cos_angles_c_n_ca
  c_n_ca_cos_angle_error = jnp.sqrt(
      1e-6 + jnp.square(c_n_ca_cos_angle - gt_angle))
  c_n_ca_loss_per_residue = jax.nn.relu(
      c_n_ca_cos_angle_error - tolerance_factor_soft * gt_stddev)
  mask = this_c_mask * next_n_mask * next_ca_mask * has_no_gap_mask
  c_n_ca_loss = jnp.sum(mask * c_n_ca_loss_per_residue) / (jnp.sum(mask) + 1e-6)
  c_n_ca_violation_mask = mask * (
      c_n_ca_cos_angle_error > (tolerance_factor_hard * gt_stddev))

  # Compute a per residue loss (equally distribute the loss to both
  # neighbouring residues).
  per_residue_loss_sum = (c_n_loss_per_residue +
                          ca_c_n_loss_per_residue +
                          c_n_ca_loss_per_residue)
  per_residue_loss_sum = 0.5 * (jnp.pad(per_residue_loss_sum, [[0, 1]]) +
                                jnp.pad(per_residue_loss_sum, [[1, 0]]))

  # Compute hard violations.
  violation_mask = jnp.max(
      jnp.stack([c_n_violation_mask,
                 ca_c_n_violation_mask,
                 c_n_ca_violation_mask]), axis=0)
  violation_mask = jnp.maximum(
      jnp.pad(violation_mask, [[0, 1]]),
      jnp.pad(violation_mask, [[1, 0]]))

  return {'c_n_loss_mean': c_n_loss,  # shape ()
          'ca_c_n_loss_mean': ca_c_n_loss,  # shape ()
          'c_n_ca_loss_mean': c_n_ca_loss,  # shape ()
          'per_residue_loss_sum': per_residue_loss_sum,  # shape (N)
          'per_residue_violation_mask': violation_mask  # shape (N)
         }


def between_residue_clash_loss(
    pred_positions: geometry.Vec3Array,  # (N, 14)
    atom_exists: jnp.ndarray,  # (N, 14)
    atom_radius: jnp.ndarray,  # (N, 14)
    residue_index: jnp.ndarray,  # (N)
    asym_id: jnp.ndarray,  # (N)
    overlap_tolerance_soft=1.5,
    overlap_tolerance_hard=1.5) -> dict[str, jnp.ndarray]:
  """Loss to penalize steric clashes between residues."""
  assert len(pred_positions.shape) == 2
  assert len(atom_exists.shape) == 2
  assert len(atom_radius.shape) == 2
  assert len(residue_index.shape) == 1

  # Create the distance matrix.
  # (N, N, 14, 14)
  dists = geometry.euclidean_distance(pred_positions[:, None, :, None],
                                      pred_positions[None, :, None, :], 1e-10)

  # Create the mask for valid distances.
  # shape (N, N, 14, 14)
  dists_mask = (atom_exists[:, None, :, None] * atom_exists[None, :, None, :])

  # Mask out all the duplicate entries in the lower triangular matrix.
  # Also mask out the diagonal (atom-pairs from the same residue) -- these atoms
  # are handled separately.
  dists_mask *= (
      residue_index[:, None, None, None] < residue_index[None, :, None, None])

  # Backbone C--N bond between subsequent residues is no clash.
  c_one_hot = jax.nn.one_hot(2, num_classes=14)
  n_one_hot = jax.nn.one_hot(0, num_classes=14)
  neighbour_mask = ((residue_index[:, None] + 1) == residue_index[None, :])
  neighbour_mask &= (asym_id[:, None] == asym_id[None, :])
  neighbour_mask = neighbour_mask[..., None, None]
  c_n_bonds = neighbour_mask * c_one_hot[None, None, :,
                                         None] * n_one_hot[None, None, None, :]
  dists_mask *= (1. - c_n_bonds)

  # Disulfide bridge between two cysteines is no clash.
  cys_sg_idx = residue_constants.restype_name_to_atom14_names['CYS'].index('SG')
  cys_sg_one_hot = jax.nn.one_hot(cys_sg_idx, num_classes=14)
  disulfide_bonds = (cys_sg_one_hot[None, None, :, None] *
                     cys_sg_one_hot[None, None, None, :])
  dists_mask *= (1. - disulfide_bonds)

  # Compute the lower bound for the allowed distances.
  # shape (N, N, 14, 14)
  dists_lower_bound = dists_mask * (
      atom_radius[:, None, :, None] + atom_radius[None, :, None, :])

  # Compute the error.
  # shape (N, N, 14, 14)
  dists_to_low_error = dists_mask * jax.nn.relu(
      dists_lower_bound - overlap_tolerance_soft - dists)

  # Compute the mean loss.
  # shape ()
  mean_loss = (jnp.sum(dists_to_low_error)
               / (1e-6 + jnp.sum(dists_mask)))

  # Compute the per atom loss sum.
  # shape (N, 14)
  per_atom_loss_sum = (jnp.sum(dists_to_low_error, axis=[0, 2]) +
                       jnp.sum(dists_to_low_error, axis=[1, 3]))

  # Compute the hard clash mask.
  # shape (N, N, 14, 14)
  clash_mask = dists_mask * (
      dists < (dists_lower_bound - overlap_tolerance_hard))

  # Compute the per atom clash.
  # shape (N, 14)
  per_atom_clash_mask = jnp.maximum(
      jnp.max(clash_mask, axis=[0, 2]),
      jnp.max(clash_mask, axis=[1, 3]))

  return {'mean_loss': mean_loss,  # shape ()
          'per_atom_loss_sum': per_atom_loss_sum,  # shape (N, 14)
          'per_atom_clash_mask': per_atom_clash_mask  # shape (N, 14)
         }


def within_residue_violations(
    pred_positions: geometry.Vec3Array,  # (N, 14)
    atom_exists: jnp.ndarray,  # (N, 14)
    dists_lower_bound: jnp.ndarray,  # (N, 14, 14)
    dists_upper_bound: jnp.ndarray,  # (N, 14, 14)
    tighten_bounds_for_loss=0.0,
) -> dict[str, jnp.ndarray]:
  """Find within-residue violations."""
  assert len(pred_positions.shape) == 2
  assert len(atom_exists.shape) == 2
  assert len(dists_lower_bound.shape) == 3
  assert len(dists_upper_bound.shape) == 3

  # Compute the mask for each residue.
  # shape (N, 14, 14)
  dists_masks = (1. - jnp.eye(14, 14)[None])
  dists_masks *= (atom_exists[:, :, None] * atom_exists[:, None, :])

  # Distance matrix
  # shape (N, 14, 14)
  dists = geometry.euclidean_distance(pred_positions[:, :, None],
                                      pred_positions[:, None, :], 1e-10)

  # Compute the loss.
  # shape (N, 14, 14)
  dists_to_low_error = jax.nn.relu(
      dists_lower_bound + tighten_bounds_for_loss - dists)
  dists_to_high_error = jax.nn.relu(
      dists + tighten_bounds_for_loss - dists_upper_bound)
  loss = dists_masks * (dists_to_low_error + dists_to_high_error)

  # Compute the per atom loss sum.
  # shape (N, 14)
  per_atom_loss_sum = (jnp.sum(loss, axis=1) +
                       jnp.sum(loss, axis=2))

  # Compute the violations mask.
  # shape (N, 14, 14)
  violations = dists_masks * ((dists < dists_lower_bound) |
                              (dists > dists_upper_bound))

  # Compute the per atom violations.
  # shape (N, 14)
  per_atom_violations = jnp.maximum(
      jnp.max(violations, axis=1), jnp.max(violations, axis=2))

  return {'per_atom_loss_sum': per_atom_loss_sum,  # shape (N, 14)
          'per_atom_violations': per_atom_violations  # shape (N, 14)
         }


def find_optimal_renaming(
    gt_positions: geometry.Vec3Array,  # (N, 14)
    alt_gt_positions: geometry.Vec3Array,  # (N, 14)
    atom_is_ambiguous: jnp.ndarray,  # (N, 14)
    gt_exists: jnp.ndarray,  # (N, 14)
    pred_positions: geometry.Vec3Array,  # (N, 14)
) -> jnp.ndarray:  # (N):
  """Find optimal renaming for ground truth that maximizes LDDT."""
  assert len(gt_positions.shape) == 2
  assert len(alt_gt_positions.shape) == 2
  assert len(atom_is_ambiguous.shape) == 2
  assert len(gt_exists.shape) == 2
  assert len(pred_positions.shape) == 2

  # Create the pred distance matrix.
  # shape (N, N, 14, 14)
  pred_dists = geometry.euclidean_distance(pred_positions[:, None, :, None],
                                           pred_positions[None, :, None, :],
                                           1e-10)

  # Compute distances for ground truth with original and alternative names.
  # shape (N, N, 14, 14)
  gt_dists = geometry.euclidean_distance(gt_positions[:, None, :, None],
                                         gt_positions[None, :, None, :], 1e-10)

  alt_gt_dists = geometry.euclidean_distance(alt_gt_positions[:, None, :, None],
                                             alt_gt_positions[None, :, None, :],
                                             1e-10)

  # Compute LDDT's.
  # shape (N, N, 14, 14)
  lddt = jnp.sqrt(1e-10 + squared_difference(pred_dists, gt_dists))
  alt_lddt = jnp.sqrt(1e-10 + squared_difference(pred_dists, alt_gt_dists))

  # Create a mask for ambiguous atoms in rows vs. non-ambiguous atoms
  # in cols.
  # shape (N ,N, 14, 14)
  mask = (
      gt_exists[:, None, :, None] *  # rows
      atom_is_ambiguous[:, None, :, None] *  # rows
      gt_exists[None, :, None, :] *  # cols
      (1. - atom_is_ambiguous[None, :, None, :]))  # cols

  # Aggregate distances for each residue to the non-amibuguous atoms.
  # shape (N)
  per_res_lddt = jnp.sum(mask * lddt, axis=[1, 2, 3])
  alt_per_res_lddt = jnp.sum(mask * alt_lddt, axis=[1, 2, 3])

  # Decide for each residue, whether alternative naming is better.
  # shape (N)
  alt_naming_is_better = (alt_per_res_lddt < per_res_lddt).astype(jnp.float32)

  return alt_naming_is_better  # shape (N)


def frame_aligned_point_error(
    pred_frames: geometry.Rigid3Array,  # shape (num_frames)
    target_frames: geometry.Rigid3Array,  # shape (num_frames)
    frames_mask: jnp.ndarray,  # shape (num_frames)
    pred_positions: geometry.Vec3Array,  # shape (num_positions)
    target_positions: geometry.Vec3Array,  # shape (num_positions)
    positions_mask: jnp.ndarray,  # shape (num_positions)
    pair_mask: Optional[jnp.ndarray],  # shape (num_frames, num_posiitons)
    l1_clamp_distance: float,
    length_scale=20.,
    epsilon=1e-4) -> jnp.ndarray:  # shape ()
  """Measure point error under different alignements.

  Computes error between two structures with B points
  under A alignments derived form the given pairs of frames.
  Args:
    pred_frames: num_frames reference frames for 'pred_positions'.
    target_frames: num_frames reference frames for 'target_positions'.
    frames_mask: Mask for frame pairs to use.
    pred_positions: num_positions predicted positions of the structure.
    target_positions: num_positions target positions of the structure.
    positions_mask: Mask on which positions to score.
    pair_mask: A (num_frames, num_positions) mask to use in the loss, useful
      for separating intra from inter chain losses.
    l1_clamp_distance: Distance cutoff on error beyond which gradients will
      be zero.
    length_scale: length scale to divide loss by.
    epsilon: small value used to regularize denominator for masked average.
  Returns:
    Masked Frame aligned point error.
  """
  # For now we do not allow any batch dimensions.
  assert len(pred_frames.rotation.shape) == 1
  assert len(target_frames.rotation.shape) == 1
  assert frames_mask.ndim == 1
  assert pred_positions.x.ndim == 1
  assert target_positions.x.ndim == 1
  assert positions_mask.ndim == 1

  # Compute array of predicted positions in the predicted frames.
  # geometry.Vec3Array (num_frames, num_positions)
  local_pred_pos = pred_frames[:, None].inverse().apply_to_point(
      pred_positions[None, :])

  # Compute array of target positions in the target frames.
  # geometry.Vec3Array (num_frames, num_positions)
  local_target_pos = target_frames[:, None].inverse().apply_to_point(
      target_positions[None, :])

  # Compute errors between the structures.
  # jnp.ndarray (num_frames, num_positions)
  error_dist = geometry.euclidean_distance(local_pred_pos, local_target_pos,
                                           epsilon)

  clipped_error_dist = jnp.clip(error_dist, 0, l1_clamp_distance)

  normed_error = clipped_error_dist / length_scale
  normed_error *= jnp.expand_dims(frames_mask, axis=-1)
  normed_error *= jnp.expand_dims(positions_mask, axis=-2)
  if pair_mask is not None:
    normed_error *= pair_mask

  mask = (jnp.expand_dims(frames_mask, axis=-1) *
          jnp.expand_dims(positions_mask, axis=-2))
  if pair_mask is not None:
    mask *= pair_mask
  normalization_factor = jnp.sum(mask, axis=(-1, -2))
  return (jnp.sum(normed_error, axis=(-2, -1)) /
          (epsilon + normalization_factor))


def get_chi_atom_indices():
  """Returns atom indices needed to compute chi angles for all residue types.

  Returns:
    A tensor of shape [residue_types=21, chis=4, atoms=4]. The residue types are
    in the order specified in residue_constants.restypes + unknown residue type
    at the end. For chi angles which are not defined on the residue, the
    positions indices are by default set to 0.
  """
  chi_atom_indices = []
  for residue_name in residue_constants.restypes:
    residue_name = residue_constants.restype_1to3[residue_name]
    residue_chi_angles = residue_constants.chi_angles_atoms[residue_name]
    atom_indices = []
    for chi_angle in residue_chi_angles:
      atom_indices.append(
          [residue_constants.atom_order[atom] for atom in chi_angle])
    for _ in range(4 - len(atom_indices)):
      atom_indices.append([0, 0, 0, 0])  # For chi angles not defined on the AA.
    chi_atom_indices.append(atom_indices)

  chi_atom_indices.append([[0, 0, 0, 0]] * 4)  # For UNKNOWN residue.

  return jnp.asarray(chi_atom_indices)


def compute_chi_angles(positions: geometry.Vec3Array,
                       mask: geometry.Vec3Array,
                       aatype: geometry.Vec3Array):
  """Computes the chi angles given all atom positions and the amino acid type.

  Args:
    positions: A Vec3Array of shape
      [num_res, residue_constants.atom_type_num], with positions of
      atoms needed to calculate chi angles. Supports up to 1 batch dimension.
    mask: An optional tensor of shape
      [num_res, residue_constants.atom_type_num] that masks which atom
      positions are set for each residue. If given, then the chi mask will be
      set to 1 for a chi angle only if the amino acid has that chi angle and all
      the chi atoms needed to calculate that chi angle are set. If not given
      (set to None), the chi mask will be set to 1 for a chi angle if the amino
      acid has that chi angle and whether the actual atoms needed to calculate
      it were set will be ignored.
    aatype: A tensor of shape [num_res] with amino acid type integer
      code (0 to 21). Supports up to 1 batch dimension.

  Returns:
    A tuple of tensors (chi_angles, mask), where both have shape
    [num_res, 4]. The mask masks out unused chi angles for amino acid
    types that have less than 4 chi angles. If atom_positions_mask is set, the
    chi mask will also mask out uncomputable chi angles.
  """

  # Don't assert on the num_res and batch dimensions as they might be unknown.
  assert positions.shape[-1] == residue_constants.atom_type_num
  assert mask.shape[-1] == residue_constants.atom_type_num

  # Compute the table of chi angle indices. Shape: [restypes, chis=4, atoms=4].
  chi_atom_indices = get_chi_atom_indices()
  # Select atoms to compute chis. Shape: [num_res, chis=4, atoms=4].
  atom_indices = utils.batched_gather(
      params=chi_atom_indices, indices=aatype, axis=0)
  # Gather atom positions. Shape: [num_res, chis=4, atoms=4, xyz=3].
  chi_angle_atoms = jax.tree_map(
      lambda x: utils.batched_gather(  # pylint: disable=g-long-lambda
          params=x, indices=atom_indices, axis=-1, batch_dims=1), positions)
  a, b, c, d = [chi_angle_atoms[..., i] for i in range(4)]

  chi_angles = geometry.dihedral_angle(a, b, c, d)

  # Copy the chi angle mask, add the UNKNOWN residue. Shape: [restypes, 4].
  chi_angles_mask = list(residue_constants.chi_angles_mask)
  chi_angles_mask.append([0.0, 0.0, 0.0, 0.0])
  chi_angles_mask = jnp.asarray(chi_angles_mask)
  # Compute the chi angle mask. Shape [num_res, chis=4].
  chi_mask = utils.batched_gather(params=chi_angles_mask, indices=aatype,
                                  axis=0)

  # The chi_mask is set to 1 only when all necessary chi angle atoms were set.
  # Gather the chi angle atoms mask. Shape: [num_res, chis=4, atoms=4].
  chi_angle_atoms_mask = utils.batched_gather(
      params=mask, indices=atom_indices, axis=-1, batch_dims=1)
  # Check if all 4 chi angle atoms were set. Shape: [num_res, chis=4].
  chi_angle_atoms_mask = jnp.prod(chi_angle_atoms_mask, axis=[-1])
  chi_mask = chi_mask * chi_angle_atoms_mask.astype(jnp.float32)

  return chi_angles, chi_mask


def make_transform_from_reference(
    a_xyz: geometry.Vec3Array,
    b_xyz: geometry.Vec3Array,
    c_xyz: geometry.Vec3Array) -> geometry.Rigid3Array:
  """Returns rotation and translation matrices to convert from reference.

  Note that this method does not take care of symmetries. If you provide the
  coordinates in the non-standard way, the A atom will end up in the negative
  y-axis rather than in the positive y-axis. You need to take care of such
  cases in your code.

  Args:
    a_xyz: A Vec3Array.
    b_xyz: A Vec3Array.
    c_xyz: A Vec3Array.

  Returns:
    A Rigid3Array which, when applied to coordinates in a canonicalized
    reference frame, will give coordinates approximately equal
    the original coordinates (in the global frame).
  """
  rotation = geometry.Rot3Array.from_two_vectors(c_xyz - b_xyz,
                                                 a_xyz - b_xyz)
  return geometry.Rigid3Array(rotation, b_xyz)

alphamissense/model/all_atom_multimer.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A collection of common Haiku modules for use in protein folding."""
import numbers
from typing import Optional, Sequence, Union

import haiku as hk
import jax.numpy as jnp
import numpy as np


# Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
TRUNCATED_NORMAL_STDDEV_FACTOR = np.asarray(.87962566103423978,
                                            dtype=np.float32)


def get_initializer_scale(initializer_name, input_shape):
  """Get Initializer for weights and scale to multiply activations by."""

  if initializer_name == 'zeros':
    w_init = hk.initializers.Constant(0.0)
  else:
    # fan-in scaling
    scale = 1.
    for channel_dim in input_shape:
      scale /= channel_dim
    if initializer_name == 'relu':
      scale *= 2

    noise_scale = scale

    stddev = np.sqrt(noise_scale)
    # Adjust stddev for truncation.
    stddev = stddev / TRUNCATED_NORMAL_STDDEV_FACTOR
    w_init = hk.initializers.TruncatedNormal(mean=0.0, stddev=stddev)

  return w_init


class Linear(hk.Module):
  """Protein folding specific Linear module.

  This differs from the standard Haiku Linear in a few ways:
    * It supports inputs and outputs of arbitrary rank
    * Initializers are specified by strings
  """

  def __init__(self,
               num_output: Union[int, Sequence[int]],
               initializer: str = 'linear',
               num_input_dims: int = 1,
               use_bias: bool = True,
               bias_init: float = 0.,
               precision: Optional[bool] = None,
               name: str = 'linear'):
    """Constructs Linear Module.

    Args:
      num_output: Number of output channels. Can be tuple when outputting
          multiple dimensions.
      initializer: What initializer to use, should be one of {'linear', 'relu',
        'zeros'}
      num_input_dims: Number of dimensions from the end to project.
      use_bias: Whether to include trainable bias
      bias_init: Value used to initialize bias.
      precision: What precision to use for matrix multiplication, defaults
        to None.
      name: Name of module, used for name scopes.
    """
    super().__init__(name=name)
    if isinstance(num_output, numbers.Integral):
      self.output_shape = (num_output,)
    else:
      self.output_shape = tuple(num_output)
    self.initializer = initializer
    self.use_bias = use_bias
    self.bias_init = bias_init
    self.num_input_dims = num_input_dims
    self.num_output_dims = len(self.output_shape)
    self.precision = precision

  def __call__(self, inputs):
    """Connects Module.

    Args:
      inputs: Tensor with at least num_input_dims dimensions.

    Returns:
      output of shape [...] + num_output.
    """

    if self.num_input_dims > 0:
      in_shape = inputs.shape[-self.num_input_dims:]
    else:
      in_shape = ()

    weight_init = get_initializer_scale(self.initializer, in_shape)

    in_letters = 'abcde'[:self.num_input_dims]
    out_letters = 'hijkl'[:self.num_output_dims]

    weight_shape = in_shape + self.output_shape
    weights = hk.get_parameter('weights', weight_shape, inputs.dtype,
                               weight_init)

    equation = f'...{in_letters}, {in_letters}{out_letters}->...{out_letters}'

    output = jnp.einsum(equation, inputs, weights, precision=self.precision)

    if self.use_bias:
      bias = hk.get_parameter('bias', self.output_shape, inputs.dtype,
                              hk.initializers.Constant(self.bias_init))
      output += bias

    return output


class LayerNorm(hk.LayerNorm):
  """LayerNorm module.

  Equivalent to hk.LayerNorm but with different parameter shapes: they are
  always vectors rather than possibly higher-rank tensors. This makes it easier
  to change the layout whilst keep the model weight-compatible.
  """

  def __init__(self,
               axis,
               create_scale: bool,
               create_offset: bool,
               eps: float = 1e-5,
               scale_init=None,
               offset_init=None,
               use_fast_variance: bool = False,
               name=None,
               param_axis=None):
    super().__init__(
        axis=axis,
        create_scale=False,
        create_offset=False,
        eps=eps,
        scale_init=None,
        offset_init=None,
        use_fast_variance=use_fast_variance,
        name=name,
        param_axis=param_axis)
    self._temp_create_scale = create_scale
    self._temp_create_offset = create_offset

  def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
    is_bf16 = (x.dtype == jnp.bfloat16)
    if is_bf16:
      x = x.astype(jnp.float32)

    param_axis = self.param_axis[0] if self.param_axis else -1
    param_shape = (x.shape[param_axis],)

    param_broadcast_shape = [1] * x.ndim
    param_broadcast_shape[param_axis] = x.shape[param_axis]
    scale = None
    offset = None
    if self._temp_create_scale:
      scale = hk.get_parameter(
          'scale', param_shape, x.dtype, init=self.scale_init)
      scale = scale.reshape(param_broadcast_shape)

    if self._temp_create_offset:
      offset = hk.get_parameter(
          'offset', param_shape, x.dtype, init=self.offset_init)
      offset = offset.reshape(param_broadcast_shape)

    out = super().__call__(x, scale=scale, offset=offset)

    if is_bf16:
      out = out.astype(jnp.bfloat16)

    return out

alphamissense/model/common_modules.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Model config."""

import copy
import ml_collections

NUM_RES = 'num residues placeholder'
NUM_MSA_SEQ = 'msa placeholder'
NUM_EXTRA_SEQ = 'extra msa placeholder'
NUM_TEMPLATES = 'num templates placeholder'


def model_config(name: str | None = None) -> ml_collections.ConfigDict:
  cfg = copy.deepcopy(CONFIG_MISSENSE)
  if name:
    if name not in CONFIG_DIFFS:
      raise ValueError(f'Invalid model name {name}.')
    cfg.update_from_flattened_dict(CONFIG_DIFFS[name])
  return cfg


CONFIG_DIFFS = {
    'model_3': {
        'model.embeddings_and_evoformer.enable_extra_msa_stack': True,
    },
}


CONFIG_MISSENSE = ml_collections.ConfigDict({
    'model': {
        'embeddings_and_evoformer': {
            'evoformer_num_block': 48,
            'evoformer': {
                'msa_column_attention': {
                    'dropout_rate': 0.0,
                    'gating': True,
                    'num_head': 8,
                    'orientation': 'per_column',
                    'shared_dropout': True,
                },
                'msa_row_attention_with_pair_bias': {
                    'dropout_rate': 0.15,
                    'gating': True,
                    'num_head': 8,
                    'orientation': 'per_row',
                    'shared_dropout': True,
                },
                'msa_transition': {
                    'dropout_rate': 0.0,
                    'num_intermediate_factor': 4,
                    'orientation': 'per_row',
                    'shared_dropout': True,
                },
                'outer_product_mean': {
                    'chunk_size': 128,
                    'dropout_rate': 0.0,
                    'first': True,
                    'num_outer_channel': 32,
                    'orientation': 'per_row',
                    'shared_dropout': True,
                },
                'pair_transition': {
                    'dropout_rate': 0.0,
                    'num_intermediate_factor': 4,
                    'orientation': 'per_row',
                    'shared_dropout': True,
                },
                'triangle_attention_ending_node': {
                    'dropout_rate': 0.25,
                    'gating': True,
                    'num_head': 4,
                    'orientation': 'per_column',
                    'shared_dropout': True,
                },
                'triangle_attention_starting_node': {
                    'dropout_rate': 0.25,
                    'gating': True,
                    'num_head': 4,
                    'orientation': 'per_row',
                    'shared_dropout': True,
                },
                'triangle_multiplication_incoming': {
                    'dropout_rate': 0.25,
                    'equation': 'kjc,kic->ijc',
                    'num_intermediate_channel': 128,
                    'orientation': 'per_row',
                    'shared_dropout': True,
                    'fuse_projection_weights': True,
                },
                'triangle_multiplication_outgoing': {
                    'dropout_rate': 0.25,
                    'equation': 'ikc,jkc->ijc',
                    'num_intermediate_channel': 128,
                    'orientation': 'per_row',
                    'shared_dropout': True,
                    'fuse_projection_weights': True,
                },
            },
            'extra_msa_channel': 64,
            'extra_msa_stack_num_block': 4,
            'num_msa': 256,
            'num_extra_msa': 1152,
            'enable_extra_msa_stack': False,
            'use_chain_relative': True,
            'max_relative_chain': 2,
            'max_relative_idx': 32,
            'seq_channel': 384,
            'msa_channel': 256,
            'pair_channel': 128,
            'prev_pos': {
                'max_bin': 20.75,
                'min_bin': 3.25,
                'num_bins': 15,
            },
            'recycle_features': True,
            'recycle_pos': True,
        },
        'global_config': {
            'bfloat16': True,
            'bfloat16_output': False,
            'deterministic': False,
            'multimer_mode': True,
            'subbatch_size': 4,
            'use_remat': False,
            'zero_init': True,
            'eval_dropout': False,
        },
        'heads': {
            'logit_diff': {
                'weight': 10.,
            },
            'distogram': {
                'first_break': 2.3125,
                'last_break': 21.6875,
                'num_bins': 64,
                'weight': 0.1,
            },
            'structure_module': {
                'angle_norm_weight': 0.01,
                'chi_weight': 0.5,
                'clash_overlap_tolerance': 1.5,
                'dropout': 0.1,
                'interface_fape': {
                    'atom_clamp_distance': 1000.0,
                    'loss_unit_distance': 20.0,
                },
                'intra_chain_fape': {
                    'atom_clamp_distance': 10.0,
                    'loss_unit_distance': 10.0,
                },
                'num_channel': 384,
                'num_head': 12,
                'num_layer': 8,
                'num_layer_in_transition': 3,
                'num_point_qk': 4,
                'num_point_v': 8,
                'num_scalar_qk': 16,
                'num_scalar_v': 16,
                'position_scale': 20.0,
                'sidechain': {
                    'atom_clamp_distance': 10.0,
                    'loss_unit_distance': 10.0,
                    'num_channel': 128,
                    'num_residual_block': 2,
                    'weight_frac': 0.5,
                },
                'structural_violation_loss_weight': 1.0,
                'violation_tolerance_factor': 12.0,
                'weight': 0.1,
            }
        },
        'num_recycle': 3,
    }
})

alphamissense/model/config.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Modules and utilities for the structure module in the multimer system."""

import numbers
from typing import Any, Iterable, Mapping, Optional, Union

import haiku as hk
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
from alphamissense.model import all_atom_multimer
from alphamissense.model import common_modules
from alphamissense.model import geometry
from alphamissense.model import modules
from alphamissense.model import prng


EPSILON = 1e-8
Float = Union[float, jnp.ndarray]


class QuatRigid(hk.Module):
  """Module for projecting Rigids via a quaternion."""

  def __init__(self,
               global_config: ml_collections.ConfigDict,
               rigid_shape: Union[int, Iterable[int]] = tuple(),
               full_quat: bool = False,
               init: str = 'zeros',
               name: str = 'quat_rigid'):
    """Module projecting a Rigid Object.

    For this Module the Rotation is parametrized as a quaternion,
    If 'full_quat' is True a 4 vector is produced for the rotation which is
    normalized and treated as a quaternion.
    When 'full_quat' is False a 3 vector is produced and the 1st component of
    the quaternion is set to 1.

    Args:
      global_config: Global Config, used to set certain properties of underlying
        Linear module, see common_modules.Linear for details.
      rigid_shape: Shape of Rigids relative to shape of activations, e.g. when
        activations have shape (n,) and this is (m,) output will be (n, m)
      full_quat: Whether to parametrize rotation using full quaternion.
      init: initializer to use, see common_modules.Linear for details
      name: Name to use for module.
    """
    self.init = init
    self.global_config = global_config
    if isinstance(rigid_shape, int):
      self.rigid_shape = (rigid_shape,)
    else:
      self.rigid_shape = tuple(rigid_shape)
    self.full_quat = full_quat
    super(QuatRigid, self).__init__(name=name)

  def __call__(self, activations: jnp.ndarray) -> geometry.Rigid3Array:
    """Executes Module.

    This returns a set of rigid with the same shape as activations, projecting
    the channel dimension, rigid_shape controls the trailing dimensions.
    For example when activations is shape (12, 5) and rigid_shape is (3, 2)
    then the shape of the output rigids will be (12, 3, 2).
    This also supports passing in an empty tuple for rigid shape, in that case
    the example would produce a rigid of shape (12,).

    Args:
      activations: Activations to use for projection, shape [..., num_channel]
    Returns:
      Rigid transformations with shape [...] + rigid_shape
    """
    if self.full_quat:
      rigid_dim = 7
    else:
      rigid_dim = 6
    linear_dims = self.rigid_shape + (rigid_dim,)
    rigid_flat = common_modules.Linear(
        linear_dims,
        initializer=self.init,
        precision=jax.lax.Precision.HIGHEST,
        name='rigid')(
            activations)
    rigid_flat = jnp.moveaxis(rigid_flat, source=-1, destination=0)
    if self.full_quat:
      qw, qx, qy, qz = rigid_flat[:4]
      translation = rigid_flat[4:]
    else:
      qx, qy, qz = rigid_flat[:3]
      qw = jnp.ones_like(qx)
      translation = rigid_flat[3:]
    rotation = geometry.Rot3Array.from_quaternion(
        qw, qx, qy, qz, normalize=True)
    translation = geometry.Vec3Array(*translation)
    return geometry.Rigid3Array(rotation, translation)


class PointProjection(hk.Module):
  """Given input reprensentation and frame produces points in global frame."""

  def __init__(self,
               num_points: Union[Iterable[int], int],
               global_config: ml_collections.ConfigDict,
               return_local_points: bool = False,
               name: str = 'point_projection'):
    """Constructs Linear Module.

    Args:
      num_points: number of points to project. Can be tuple when outputting
          multiple dimensions
      global_config: Global Config, passed through to underlying Linear
      return_local_points: Whether to return points in local frame as well.
      name: name of module, used for name scopes.
    """
    if isinstance(num_points, numbers.Integral):
      self.num_points = (num_points,)
    else:
      self.num_points = tuple(num_points)

    self.return_local_points = return_local_points

    self.global_config = global_config

    super().__init__(name=name)

  def __call__(
      self, activations: jnp.ndarray, rigids: geometry.Rigid3Array
  ) -> Union[geometry.Vec3Array, tuple[geometry.Vec3Array, geometry.Vec3Array]]:
    output_shape = self.num_points
    output_shape = output_shape[:-1] + (3 * output_shape[-1],)
    points_local = common_modules.Linear(
        output_shape,
        precision=jax.lax.Precision.HIGHEST,
        name='point_projection')(
            activations)
    points_local = jnp.split(points_local, 3, axis=-1)
    points_local = geometry.Vec3Array(*points_local)
    rigids = rigids[(...,) + (None,) * len(output_shape)]
    points_global = rigids.apply_to_point(points_local)
    if self.return_local_points:
      return points_global, points_local
    else:
      return points_global


class InvariantPointAttention(hk.Module):
  """Invariant point attention module.

  The high-level idea is that this attention module works over a set of points
  and associated orientations in 3D space (e.g. protein residues).

  Each residue outputs a set of queries and keys as points in their local
  reference frame.  The attention is then defined as the euclidean distance
  between the queries and keys in the global frame.
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               dist_epsilon: float = 1e-8,
               name: str = 'invariant_point_attention'):
    """Initialize.

    Args:
      config: iterative Fold Head Config
      global_config: Global Config of Model.
      dist_epsilon: Small value to avoid NaN in distance calculation.
      name: Sonnet name.
    """
    super().__init__(name=name)

    self._dist_epsilon = dist_epsilon
    self._zero_initialize_last = global_config.zero_init

    self.config = config

    self.global_config = global_config

  def __call__(
      self,
      inputs_1d: jnp.ndarray,
      inputs_2d: jnp.ndarray,
      mask: jnp.ndarray,
      rigid: geometry.Rigid3Array,
      ) -> jnp.ndarray:
    """Compute geometric aware attention.

    Given a set of query residues (defined by affines and associated scalar
    features), this function computes geometric aware attention between the
    query residues and target residues.

    The residues produce points in their local reference frame, which
    are converted into the global frame to get attention via euclidean distance.

    Equivalently the target residues produce points in their local frame to be
    used as attention values, which are converted into the query residues local
    frames.

    Args:
      inputs_1d: (N, C) 1D input embedding that is the basis for the
        scalar queries.
      inputs_2d: (N, M, C') 2D input embedding, used for biases values in the
        attention between query_inputs_1d and target_inputs_1d.
      mask: (N, 1) mask to indicate query_inputs_1d that participate in
        the attention.
      rigid: Rigid object describing the position and orientation of
        every element in query_inputs_1d.

    Returns:
      Transformation of the input embedding.
    """

    num_head = self.config.num_head

    attn_logits = 0.

    num_point_qk = self.config.num_point_qk
    # Each point pair (q, k) contributes Var [0.5 ||q||^2 - <q, k>] = 9 / 2
    point_variance = max(num_point_qk, 1) * 9. / 2
    point_weights = np.sqrt(1.0 / point_variance)

    # This is equivalent to jax.nn.softplus, but avoids a bug in the test...
    softplus = lambda x: jnp.logaddexp(x, jnp.zeros_like(x))
    raw_point_weights = hk.get_parameter(
        'trainable_point_weights',
        shape=[num_head],
        # softplus^{-1} (1)
        init=hk.initializers.Constant(np.log(np.exp(1.) - 1.)))

    # Trainable per-head weights for points.
    trainable_point_weights = softplus(raw_point_weights)
    point_weights *= trainable_point_weights
    q_point = PointProjection([num_head, num_point_qk],
                              self.global_config,
                              name='q_point_projection')(inputs_1d,
                                                         rigid)

    k_point = PointProjection([num_head, num_point_qk],
                              self.global_config,
                              name='k_point_projection')(inputs_1d,
                                                         rigid)

    dist2 = geometry.square_euclidean_distance(
        q_point[:, None, :, :], k_point[None, :, :, :], epsilon=0.)
    attn_qk_point = -0.5 * jnp.sum(point_weights[:, None] * dist2, axis=-1)
    attn_logits += attn_qk_point

    num_scalar_qk = self.config.num_scalar_qk
    # We assume that all queries and keys come iid from N(0, 1) distribution
    # and compute the variances of the attention logits.
    # Each scalar pair (q, k) contributes Var q*k = 1
    scalar_variance = max(num_scalar_qk, 1) * 1.
    scalar_weights = np.sqrt(1.0 / scalar_variance)
    q_scalar = common_modules.Linear([num_head, num_scalar_qk],
                                     use_bias=False,
                                     name='q_scalar_projection')(
                                         inputs_1d)

    k_scalar = common_modules.Linear([num_head, num_scalar_qk],
                                     use_bias=False,
                                     name='k_scalar_projection')(
                                         inputs_1d)
    q_scalar *= scalar_weights
    attn_logits += jnp.einsum('qhc,khc->qkh', q_scalar, k_scalar)

    attention_2d = common_modules.Linear(
        num_head, name='attention_2d')(inputs_2d)
    attn_logits += attention_2d

    mask_2d = mask * jnp.swapaxes(mask, -1, -2)
    attn_logits -= 1e5 * (1. - mask_2d[..., None])

    attn_logits *= np.sqrt(1. / 3)     # Normalize by number of logit terms (3)
    attn = jax.nn.softmax(attn_logits, axis=-2)

    num_scalar_v = self.config.num_scalar_v

    v_scalar = common_modules.Linear([num_head, num_scalar_v],
                                     use_bias=False,
                                     name='v_scalar_projection')(
                                         inputs_1d)

    # [num_query_residues, num_head, num_scalar_v]
    result_scalar = jnp.einsum('qkh, khc->qhc', attn, v_scalar)

    num_point_v = self.config.num_point_v
    v_point = PointProjection([num_head, num_point_v],
                              self.global_config,
                              name='v_point_projection')(inputs_1d,
                                                         rigid)

    result_point_global = jax.tree_map(
        lambda x: jnp.sum(attn[..., None] * x, axis=-3), v_point[None])

    # Features used in the linear output projection. Should have the size
    # [num_query_residues, ?]
    output_features = []
    num_query_residues, _ = inputs_1d.shape

    flat_shape = [num_query_residues, -1]

    result_scalar = jnp.reshape(result_scalar, flat_shape)
    output_features.append(result_scalar)

    result_point_global = jax.tree_map(lambda r: jnp.reshape(r, flat_shape),
                                       result_point_global)
    result_point_local = rigid[..., None].apply_inverse_to_point(
        result_point_global)
    output_features.extend(
        [result_point_local.x, result_point_local.y, result_point_local.z])

    point_norms = result_point_local.norm(self._dist_epsilon)
    output_features.append(point_norms)

    # Dimensions: h = heads, i and j = residues,
    # c = inputs_2d channels
    # Contraction happens over the second residue dimension, similarly to how
    # the usual attention is performed.
    result_attention_over_2d = jnp.einsum('ijh, ijc->ihc', attn, inputs_2d)
    output_features.append(jnp.reshape(result_attention_over_2d, flat_shape))

    final_init = 'zeros' if self._zero_initialize_last else 'linear'

    final_act = jnp.concatenate(output_features, axis=-1)

    return common_modules.Linear(
        self.config.num_channel,
        initializer=final_init,
        name='output_projection')(final_act)


class FoldIteration(hk.Module):
  """A single iteration of iterative folding.

  First, each residue attends to all residues using InvariantPointAttention.
  Then, we apply transition layers to update the hidden representations.
  Finally, we use the hidden representations to produce an update to the
  affine of each residue.
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'fold_iteration'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(
      self,
      activations: Mapping[str, Any],
      aatype: jnp.ndarray,
      sequence_mask: jnp.ndarray,
      update_rigid: bool,
      is_training: bool,
      initial_act: jnp.ndarray,
      safe_key: Optional[prng.SafeKey] = None,
      static_feat_2d: Optional[jnp.ndarray] = None,
  ) -> tuple[dict[str, Any], dict[str, Any]]:

    c = self.config

    if safe_key is None:
      safe_key = prng.SafeKey(hk.next_rng_key())

    def safe_dropout_fn(tensor, safe_key):
      return modules.apply_dropout(
          tensor=tensor,
          safe_key=safe_key,
          rate=0.0 if self.global_config.deterministic else c.dropout,
          is_training=is_training)

    rigid = activations['rigid']

    act = activations['act']
    attention_module = InvariantPointAttention(
        self.config, self.global_config)
    # Attention
    act += attention_module(
        inputs_1d=act,
        inputs_2d=static_feat_2d,
        mask=sequence_mask,
        rigid=rigid)

    unused_safe_key, *sub_keys = safe_key.split(3)
    sub_keys = iter(sub_keys)
    act = safe_dropout_fn(act, next(sub_keys))
    act = common_modules.LayerNorm(
        axis=-1,
        create_scale=True,
        create_offset=True,
        name='attention_layer_norm')(
            act)
    final_init = 'zeros' if self.global_config.zero_init else 'linear'

    # Transition
    input_act = act
    for i in range(c.num_layer_in_transition):
      init = 'relu' if i < c.num_layer_in_transition - 1 else final_init
      act = common_modules.Linear(
          c.num_channel,
          initializer=init,
          name='transition')(
              act)
      if i < c.num_layer_in_transition - 1:
        act = jax.nn.relu(act)
    act += input_act
    act = safe_dropout_fn(act, next(sub_keys))
    act = common_modules.LayerNorm(
        axis=-1,
        create_scale=True,
        create_offset=True,
        name='transition_layer_norm')(act)
    if update_rigid:
      # Rigid update
      rigid_update = QuatRigid(
          self.global_config, init=final_init)(
              act)
      rigid = rigid @ rigid_update

    sc = MultiRigidSidechain(c.sidechain, self.global_config)(
        rigid.scale_translation(c.position_scale), [act, initial_act], aatype)

    outputs = {'rigid': rigid, 'sc': sc}

    rotation = jax.tree_map(jax.lax.stop_gradient, rigid.rotation)
    rigid = geometry.Rigid3Array(rotation, rigid.translation)

    new_activations = {
        'act': act,
        'rigid': rigid
    }
    return new_activations, outputs


def generate_monomer_rigids(representations: Mapping[str, jnp.ndarray],
                            batch: Mapping[str, jnp.ndarray],
                            config: ml_collections.ConfigDict,
                            global_config: ml_collections.ConfigDict,
                            is_training: bool,
                            safe_key: prng.SafeKey
                            ) -> dict[str, Any]:
  """Generate predicted Rigid's for a single chain.

  This is the main part of the iterative fold head - it iteratively applies
  folding to produce a set of predicted residue positions.

  Args:
    representations: Embeddings dictionary.
    batch: Batch dictionary.
    config: config for the iterative fold head.
    global_config: global config.
    is_training: is training.
    safe_key: A prng.SafeKey object that wraps a PRNG key.

  Returns:
    A dictionary containing residue Rigid's and sidechain positions.
  """
  c = config
  sequence_mask = batch['seq_mask'][:, None]
  act = common_modules.LayerNorm(
      axis=-1, create_scale=True, create_offset=True, name='single_layer_norm')(
          representations['single'])

  initial_act = act
  act = common_modules.Linear(
      c.num_channel, name='initial_projection')(act)

  # Sequence Mask has extra 1 at the end.
  rigid = geometry.Rigid3Array.identity(sequence_mask.shape[:-1])

  fold_iteration = FoldIteration(
      c, global_config, name='fold_iteration')

  assert len(batch['seq_mask'].shape) == 1

  activations = {
      'act':
          act,
      'rigid':
          rigid
  }

  act_2d = common_modules.LayerNorm(
      axis=-1,
      create_scale=True,
      create_offset=True,
      name='pair_layer_norm')(
          representations['pair'])

  safe_keys = safe_key.split(c.num_layer)
  outputs = []
  for key in safe_keys:
    activations, output = fold_iteration(
        activations,
        initial_act=initial_act,
        static_feat_2d=act_2d,
        aatype=batch['aatype'],
        safe_key=key,
        sequence_mask=sequence_mask,
        update_rigid=True,
        is_training=is_training,
        )
    outputs.append(output)

  output = jax.tree_map(lambda *x: jnp.stack(x), *outputs)
  # Pass along for LDDT-Head.
  output['act'] = activations['act']

  return output


class StructureModule(hk.Module):
  """StructureModule as a network head.

  Jumper et al. (2021) Suppl. Alg. 20 "StructureModule"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'structure_module'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               representations: Mapping[str, jnp.ndarray],
               batch: Mapping[str, jnp.ndarray],
               is_training: bool,
               safe_key: Optional[prng.SafeKey] = None,
               compute_loss: bool = False
               ) -> dict[str, Any]:
    c = self.config
    ret = {}

    if safe_key is None:
      safe_key = prng.SafeKey(hk.next_rng_key())

    output = generate_monomer_rigids(
        representations=representations,
        batch=batch,
        config=self.config,
        global_config=self.global_config,
        is_training=is_training,
        safe_key=safe_key)

    ret['traj'] = output['rigid'].scale_translation(c.position_scale).to_array()
    ret['sidechains'] = output['sc']
    ret['sidechains']['atom_pos'] = ret['sidechains']['atom_pos'].to_array()
    ret['sidechains']['frames'] = ret['sidechains']['frames'].to_array()
    if 'local_atom_pos' in ret['sidechains']:
      ret['sidechains']['local_atom_pos'] = ret['sidechains'][
          'local_atom_pos'].to_array()
      ret['sidechains']['local_frames'] = ret['sidechains'][
          'local_frames'].to_array()

    aatype = batch['aatype']
    seq_mask = batch['seq_mask']

    atom14_pred_mask = all_atom_multimer.get_atom14_mask(
        aatype) * seq_mask[:, None]
    atom14_pred_positions = output['sc']['atom_pos'][-1]
    ret['final_atom14_positions'] = atom14_pred_positions  # (N, 14, 3)
    ret['final_atom14_mask'] = atom14_pred_mask  # (N, 14)

    atom37_mask = all_atom_multimer.get_atom37_mask(aatype) * seq_mask[:, None]
    atom37_pred_positions = all_atom_multimer.atom14_to_atom37(
        atom14_pred_positions, aatype)
    atom37_pred_positions *= atom37_mask[:, :, None]
    ret['final_atom_positions'] = atom37_pred_positions  # (N, 37, 3)
    ret['final_atom_mask'] = atom37_mask  # (N, 37)
    ret['final_rigids'] = ret['traj'][-1]

    ret['act'] = output['act']

    if compute_loss:
      return ret
    else:
      no_loss_features = ['final_atom_positions', 'final_atom_mask', 'act']
      no_loss_ret = {k: ret[k] for k in no_loss_features}
      return no_loss_ret

  def loss(self,
           value: Mapping[str, Any],
           batch: Mapping[str, Any]
           ) -> dict[str, Any]:

    raise NotImplementedError(
        'This function should be called on a batch with reordered chains (see '
        'Evans et al (2021) Section 7.3. Multi-Chain Permutation Alignment.')


def l2_normalize(x: jnp.ndarray,
                 axis: int = -1,
                 epsilon: float = 1e-12
                 ) -> jnp.ndarray:
  return x / jnp.sqrt(
      jnp.maximum(jnp.sum(x**2, axis=axis, keepdims=True), epsilon))


class MultiRigidSidechain(hk.Module):
  """Class to make side chain atoms."""

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'rigid_sidechain'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               rigid: geometry.Rigid3Array,
               representations_list: Iterable[jnp.ndarray],
               aatype: jnp.ndarray
               ) -> dict[str, Any]:
    """Predict sidechains using multi-rigid representations.

    Args:
      rigid: The Rigid's for each residue (translations in angstoms)
      representations_list: A list of activations to predict sidechains from.
      aatype: amino acid types.

    Returns:
      dict containing atom positions and frames (in angstrom)
    """
    act = [
        common_modules.Linear(  # pylint: disable=g-complex-comprehension
            self.config.num_channel,
            name='input_projection')(jax.nn.relu(x))
        for x in representations_list]
    # Sum the activation list (equivalent to concat then Conv1D)
    act = sum(act)

    final_init = 'zeros' if self.global_config.zero_init else 'linear'

    # Mapping with some residual blocks.
    for _ in range(self.config.num_residual_block):
      old_act = act
      act = common_modules.Linear(
          self.config.num_channel,
          initializer='relu',
          name='resblock1')(
              jax.nn.relu(act))
      act = common_modules.Linear(
          self.config.num_channel,
          initializer=final_init,
          name='resblock2')(
              jax.nn.relu(act))
      act += old_act

    # Map activations to torsion angles.
    # [batch_size, num_res, 14]
    num_res = act.shape[0]
    unnormalized_angles = common_modules.Linear(
        14, name='unnormalized_angles')(
            jax.nn.relu(act))
    unnormalized_angles = jnp.reshape(
        unnormalized_angles, [num_res, 7, 2])
    angles = l2_normalize(unnormalized_angles, axis=-1)

    outputs = {
        'angles_sin_cos': angles,  # jnp.ndarray (N, 7, 2)
        'unnormalized_angles_sin_cos':
            unnormalized_angles,  # jnp.ndarray (N, 7, 2)
    }

    # Map torsion angles to frames.
    # geometry.Rigid3Array with shape (N, 8)
    all_frames_to_global = all_atom_multimer.torsion_angles_to_frames(
        aatype,
        rigid,
        angles)

    # Use frames and literature positions to create the final atom coordinates.
    # geometry.Vec3Array with shape (N, 14)
    pred_positions = (
        all_atom_multimer.frames_and_literature_positions_to_atom14_pos(
            aatype, all_frames_to_global))

    outputs.update({
        'atom_pos': pred_positions,  # geometry.Vec3Array (N, 14)
        'frames': all_frames_to_global,  # geometry.Rigid3Array (N, 8)
    })
    return outputs

alphamissense/model/folding_multimer.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Geometry Module."""

from alphamissense.model.geometry import rigid_matrix_vector
from alphamissense.model.geometry import rotation_matrix
from alphamissense.model.geometry import struct_of_array
from alphamissense.model.geometry import vector

Rot3Array = rotation_matrix.Rot3Array
Rigid3Array = rigid_matrix_vector.Rigid3Array

StructOfArray = struct_of_array.StructOfArray

Vec3Array = vector.Vec3Array
square_euclidean_distance = vector.square_euclidean_distance
euclidean_distance = vector.euclidean_distance
dihedral_angle = vector.dihedral_angle
dot = vector.dot
cross = vector.cross

alphamissense/model/geometry/__init__.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Rigid3Array Transformations represented by a Matrix and a Vector."""

from __future__ import annotations
from typing import Union

import jax
import jax.numpy as jnp
from alphamissense.model.geometry import rotation_matrix
from alphamissense.model.geometry import struct_of_array
from alphamissense.model.geometry import vector

Float = Union[float, jnp.ndarray]

VERSION = '0.1'


@struct_of_array.StructOfArray(same_dtype=True)
class Rigid3Array:
  """Rigid Transformation, i.e. element of special euclidean group."""

  rotation: rotation_matrix.Rot3Array
  translation: vector.Vec3Array

  def __matmul__(self, other: Rigid3Array) -> Rigid3Array:
    new_rotation = self.rotation @ other.rotation
    new_translation = self.apply_to_point(other.translation)
    return Rigid3Array(new_rotation, new_translation)

  def inverse(self) -> Rigid3Array:
    """Return Rigid3Array corresponding to inverse transform."""
    inv_rotation = self.rotation.inverse()
    inv_translation = inv_rotation.apply_to_point(-self.translation)
    return Rigid3Array(inv_rotation, inv_translation)

  def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
    """Apply Rigid3Array transform to point."""
    return self.rotation.apply_to_point(point) + self.translation

  def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
    """Apply inverse Rigid3Array transform to point."""
    new_point = point - self.translation
    return self.rotation.apply_inverse_to_point(new_point)

  def compose_rotation(self, other_rotation):
    rot = self.rotation @ other_rotation
    trans = jax.tree_map(lambda x: jnp.broadcast_to(x, rot.shape),
                         self.translation)
    return Rigid3Array(rot, trans)

  @classmethod
  def identity(cls, shape, dtype=jnp.float32) -> Rigid3Array:
    """Return identity Rigid3Array of given shape."""
    return cls(
        rotation_matrix.Rot3Array.identity(shape, dtype=dtype),
        vector.Vec3Array.zeros(shape, dtype=dtype))  # pytype: disable=wrong-arg-count  # trace-all-classes

  def scale_translation(self, factor: Float) -> Rigid3Array:
    """Scale translation in Rigid3Array by 'factor'."""
    return Rigid3Array(self.rotation, self.translation * factor)

  def to_array(self):
    rot_array = self.rotation.to_array()
    vec_array = self.translation.to_array()
    return jnp.concatenate([rot_array, vec_array[..., None]], axis=-1)

  @classmethod
  def from_array(cls, array):
    rot = rotation_matrix.Rot3Array.from_array(array[..., :3])
    vec = vector.Vec3Array.from_array(array[..., -1])
    return cls(rot, vec)  # pytype: disable=wrong-arg-count  # trace-all-classes

  @classmethod
  def from_array4x4(cls, array: jnp.ndarray) -> Rigid3Array:
    """Construct Rigid3Array from homogeneous 4x4 array."""
    assert array.shape[-1] == 4
    assert array.shape[-2] == 4
    rotation = rotation_matrix.Rot3Array(
        array[..., 0, 0], array[..., 0, 1], array[..., 0, 2],
        array[..., 1, 0], array[..., 1, 1], array[..., 1, 2],
        array[..., 2, 0], array[..., 2, 1], array[..., 2, 2]
        )
    translation = vector.Vec3Array(
        array[..., 0, 3], array[..., 1, 3], array[..., 2, 3])
    return cls(rotation, translation)  # pytype: disable=wrong-arg-count  # trace-all-classes

  def __getstate__(self):
    return (VERSION, (self.rotation, self.translation))

  def __setstate__(self, state):
    version, (rot, trans) = state
    del version
    object.__setattr__(self, 'rotation', rot)
    object.__setattr__(self, 'translation', trans)

alphamissense/model/geometry/rigid_matrix_vector.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Rot3Array Matrix Class."""

from __future__ import annotations
import dataclasses
import itertools

import jax
import jax.numpy as jnp
import numpy as np
from alphamissense.model.geometry import struct_of_array
from alphamissense.model.geometry import vector

COMPONENTS = ['xx', 'xy', 'xz', 'yx', 'yy', 'yz', 'zx', 'zy', 'zz']

VERSION = '0.1'


@struct_of_array.StructOfArray(same_dtype=True)
class Rot3Array:
  """Rot3Array Matrix in 3 dimensional Space implemented as struct of arrays."""

  xx: jnp.ndarray = dataclasses.field(metadata={'dtype': jnp.float32})
  xy: jnp.ndarray
  xz: jnp.ndarray
  yx: jnp.ndarray
  yy: jnp.ndarray
  yz: jnp.ndarray
  zx: jnp.ndarray
  zy: jnp.ndarray
  zz: jnp.ndarray

  __array_ufunc__ = None

  def inverse(self) -> Rot3Array:
    """Returns inverse of Rot3Array."""
    return Rot3Array(self.xx, self.yx, self.zx,
                     self.xy, self.yy, self.zy,
                     self.xz, self.yz, self.zz)

  def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
    """Applies Rot3Array to point."""
    return vector.Vec3Array(
        self.xx * point.x + self.xy * point.y + self.xz * point.z,
        self.yx * point.x + self.yy * point.y + self.yz * point.z,
        self.zx * point.x + self.zy * point.y + self.zz * point.z)

  def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
    """Applies inverse Rot3Array to point."""
    return self.inverse().apply_to_point(point)

  def __matmul__(self, other: Rot3Array) -> Rot3Array:
    """Composes two Rot3Arrays."""
    c0 = self.apply_to_point(vector.Vec3Array(other.xx, other.yx, other.zx))
    c1 = self.apply_to_point(vector.Vec3Array(other.xy, other.yy, other.zy))
    c2 = self.apply_to_point(vector.Vec3Array(other.xz, other.yz, other.zz))
    return Rot3Array(c0.x, c1.x, c2.x, c0.y, c1.y, c2.y, c0.z, c1.z, c2.z)

  @classmethod
  def identity(cls, shape, dtype=jnp.float32) -> Rot3Array:
    """Returns identity of given shape."""
    ones = jnp.ones(shape, dtype=dtype)
    zeros = jnp.zeros(shape, dtype=dtype)
    return cls(ones, zeros, zeros, zeros, ones, zeros, zeros, zeros, ones)  # pytype: disable=wrong-arg-count  # trace-all-classes

  @classmethod
  def from_two_vectors(cls, e0: vector.Vec3Array,
                       e1: vector.Vec3Array) -> Rot3Array:
    """Construct Rot3Array from two Vectors.

    Rot3Array is constructed such that in the corresponding frame 'e0' lies on
    the positive x-Axis and 'e1' lies in the xy plane with positive sign of y.

    Args:
      e0: Vector
      e1: Vector
    Returns:
      Rot3Array
    """
    # Normalize the unit vector for the x-axis, e0.
    e0 = e0.normalized()
    # make e1 perpendicular to e0.
    c = e1.dot(e0)
    e1 = (e1 - c * e0).normalized()
    # Compute e2 as cross product of e0 and e1.
    e2 = e0.cross(e1)
    return cls(e0.x, e1.x, e2.x, e0.y, e1.y, e2.y, e0.z, e1.z, e2.z)  # pytype: disable=wrong-arg-count  # trace-all-classes

  @classmethod
  def from_array(cls, array: jnp.ndarray) -> Rot3Array:
    """Construct Rot3Array Matrix from array of shape. [..., 3, 3]."""
    return cls(*[
        array[..., i, j] for i, j in itertools.product(range(3), range(3))])

  def to_array(self) -> jnp.ndarray:
    """Convert Rot3Array to array of shape [..., 3, 3]."""
    return jnp.stack(
        [jnp.stack([self.xx, self.xy, self.xz], axis=-1),
         jnp.stack([self.yx, self.yy, self.yz], axis=-1),
         jnp.stack([self.zx, self.zy, self.zz], axis=-1)],
        axis=-2)

  @classmethod
  def from_quaternion(cls,
                      w: jnp.ndarray,
                      x: jnp.ndarray,
                      y: jnp.ndarray,
                      z: jnp.ndarray,
                      normalize: bool = True,
                      epsilon: float = 1e-6) -> Rot3Array:
    """Construct Rot3Array from components of quaternion."""
    if normalize:
      inv_norm = jax.lax.rsqrt(jnp.maximum(epsilon, w**2 + x**2 + y**2 + z**2))
      w *= inv_norm
      x *= inv_norm
      y *= inv_norm
      z *= inv_norm
    xx = 1 - 2 * (jnp.square(y) + jnp.square(z))
    xy = 2 * (x * y - w * z)
    xz = 2 * (x * z + w * y)
    yx = 2 * (x * y + w * z)
    yy = 1 - 2 * (jnp.square(x) + jnp.square(z))
    yz = 2 * (y * z - w * x)
    zx = 2 * (x * z - w * y)
    zy = 2 * (y * z + w * x)
    zz = 1 - 2 * (jnp.square(x) + jnp.square(y))
    return cls(xx, xy, xz, yx, yy, yz, zx, zy, zz)  # pytype: disable=wrong-arg-count  # trace-all-classes

  @classmethod
  def random_uniform(cls, key, shape, dtype=jnp.float32) -> Rot3Array:
    """Samples uniform random Rot3Array according to Haar Measure."""
    quat_array = jax.random.normal(key, tuple(shape) + (4,), dtype=dtype)
    return cls.from_quaternion(*[quat_array[..., i] for i in range(4)])

  def __getstate__(self):
    return (VERSION,
            [np.asarray(getattr(self, field)) for field in COMPONENTS])

  def __setstate__(self, state):
    version, state = state
    del version
    for i, field in enumerate(COMPONENTS):
      object.__setattr__(self, field, state[i])

alphamissense/model/geometry/rotation_matrix.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Shared utils for tests."""

import dataclasses

import jax.numpy as jnp
import numpy as np
from alphamissense.model.geometry import rigid_matrix_vector
from alphamissense.model.geometry import rotation_matrix
from alphamissense.model.geometry import vector


def assert_rotation_matrix_equal(matrix1: rotation_matrix.Rot3Array,
                                 matrix2: rotation_matrix.Rot3Array):
  for field in dataclasses.fields(rotation_matrix.Rot3Array):
    field = field.name
    np.testing.assert_array_equal(
        getattr(matrix1, field), getattr(matrix2, field))


def assert_rotation_matrix_close(mat1: rotation_matrix.Rot3Array,
                                 mat2: rotation_matrix.Rot3Array):
  np.testing.assert_array_almost_equal(mat1.to_array(), mat2.to_array(), 6)


def assert_array_equal_to_rotation_matrix(array: jnp.ndarray,
                                          matrix: rotation_matrix.Rot3Array):
  """Check that array and Matrix match."""
  np.testing.assert_array_equal(matrix.xx, array[..., 0, 0])
  np.testing.assert_array_equal(matrix.xy, array[..., 0, 1])
  np.testing.assert_array_equal(matrix.xz, array[..., 0, 2])
  np.testing.assert_array_equal(matrix.yx, array[..., 1, 0])
  np.testing.assert_array_equal(matrix.yy, array[..., 1, 1])
  np.testing.assert_array_equal(matrix.yz, array[..., 1, 2])
  np.testing.assert_array_equal(matrix.zx, array[..., 2, 0])
  np.testing.assert_array_equal(matrix.zy, array[..., 2, 1])
  np.testing.assert_array_equal(matrix.zz, array[..., 2, 2])


def assert_array_close_to_rotation_matrix(array: jnp.ndarray,
                                          matrix: rotation_matrix.Rot3Array):
  np.testing.assert_array_almost_equal(matrix.to_array(), array, 6)


def assert_vectors_equal(vec1: vector.Vec3Array, vec2: vector.Vec3Array):
  np.testing.assert_array_equal(vec1.x, vec2.x)
  np.testing.assert_array_equal(vec1.y, vec2.y)
  np.testing.assert_array_equal(vec1.z, vec2.z)


def assert_vectors_close(vec1: vector.Vec3Array, vec2: vector.Vec3Array):
  np.testing.assert_allclose(vec1.x, vec2.x, atol=1e-6, rtol=0.)
  np.testing.assert_allclose(vec1.y, vec2.y, atol=1e-6, rtol=0.)
  np.testing.assert_allclose(vec1.z, vec2.z, atol=1e-6, rtol=0.)


def assert_array_close_to_vector(array: jnp.ndarray, vec: vector.Vec3Array):
  np.testing.assert_allclose(vec.to_array(), array, atol=1e-6, rtol=0.)


def assert_array_equal_to_vector(array: jnp.ndarray, vec: vector.Vec3Array):
  np.testing.assert_array_equal(vec.to_array(), array)


def assert_rigid_equal_to_rigid(rigid1: rigid_matrix_vector.Rigid3Array,
                                rigid2: rigid_matrix_vector.Rigid3Array):
  assert_rot_trans_equal_to_rigid(rigid1.rotation, rigid1.translation, rigid2)


def assert_rigid_close_to_rigid(rigid1: rigid_matrix_vector.Rigid3Array,
                                rigid2: rigid_matrix_vector.Rigid3Array):
  assert_rot_trans_close_to_rigid(rigid1.rotation, rigid1.translation, rigid2)


def assert_rot_trans_equal_to_rigid(rot: rotation_matrix.Rot3Array,
                                    trans: vector.Vec3Array,
                                    rigid: rigid_matrix_vector.Rigid3Array):
  assert_rotation_matrix_equal(rot, rigid.rotation)
  assert_vectors_equal(trans, rigid.translation)


def assert_rot_trans_close_to_rigid(rot: rotation_matrix.Rot3Array,
                                    trans: vector.Vec3Array,
                                    rigid: rigid_matrix_vector.Rigid3Array):
  assert_rotation_matrix_close(rot, rigid.rotation)
  assert_vectors_close(trans, rigid.translation)

alphamissense/model/geometry/test_utils.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Vec3Array Class."""

from __future__ import annotations
import dataclasses
from typing import Union

import jax
import jax.numpy as jnp
import numpy as np
from alphamissense.model.geometry import struct_of_array

Float = Union[float, jnp.ndarray]

VERSION = '0.1'


@struct_of_array.StructOfArray(same_dtype=True)
class Vec3Array:
  """Vec3Array in 3 dimensional Space implemented as struct of arrays.

  This is done in order to improve performance and precision.
  On TPU small matrix multiplications are very suboptimal and will waste large
  compute ressources, furthermore any matrix multiplication on tpu happen in
  mixed bfloat16/float32 precision, which is often undesirable when handling
  physical coordinates.
  In most cases this will also be faster on cpu's/gpu's since it allows for
  easier use of vector instructions.
  """

  x: jnp.ndarray = dataclasses.field(metadata={'dtype': jnp.float32})
  y: jnp.ndarray
  z: jnp.ndarray

  def __post_init__(self):
    if hasattr(self.x, 'dtype'):
      assert self.x.dtype == self.y.dtype
      assert self.x.dtype == self.z.dtype
      assert all([x == y for x, y in zip(self.x.shape, self.y.shape)])
      assert all([x == z for x, z in zip(self.x.shape, self.z.shape)])

  def __add__(self, other: Vec3Array) -> Vec3Array:
    return jax.tree_map(lambda x, y: x + y, self, other)

  def __sub__(self, other: Vec3Array) -> Vec3Array:
    return jax.tree_map(lambda x, y: x - y, self, other)

  def __mul__(self, other: Float) -> Vec3Array:
    return jax.tree_map(lambda x: x * other, self)

  def __rmul__(self, other: Float) -> Vec3Array:
    return self * other

  def __truediv__(self, other: Float) -> Vec3Array:
    return jax.tree_map(lambda x: x / other, self)

  def __neg__(self) -> Vec3Array:
    return jax.tree_map(lambda x: -x, self)

  def __pos__(self) -> Vec3Array:
    return jax.tree_map(lambda x: x, self)

  def cross(self, other: Vec3Array) -> Vec3Array:
    """Compute cross product between 'self' and 'other'."""
    new_x = self.y * other.z - self.z * other.y
    new_y = self.z * other.x - self.x * other.z
    new_z = self.x * other.y - self.y * other.x
    return Vec3Array(new_x, new_y, new_z)

  def dot(self, other: Vec3Array) -> Float:
    """Compute dot product between 'self' and 'other'."""
    return self.x * other.x + self.y * other.y + self.z * other.z

  def norm(self, epsilon: float = 1e-6) -> Float:
    """Compute Norm of Vec3Array, clipped to epsilon."""
    # To avoid NaN on the backward pass, we must use maximum before the sqrt
    norm2 = self.dot(self)
    if epsilon:
      norm2 = jnp.maximum(norm2, epsilon**2)
    return jnp.sqrt(norm2)

  def norm2(self):
    return self.dot(self)

  def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
    """Return unit vector with optional clipping."""
    return self / self.norm(epsilon)

  @classmethod
  def zeros(cls, shape, dtype=jnp.float32):
    """Return Vec3Array corresponding to zeros of given shape."""
    return cls(
        jnp.zeros(shape, dtype), jnp.zeros(shape, dtype),
        jnp.zeros(shape, dtype))  # pytype: disable=wrong-arg-count  # trace-all-classes

  def to_array(self) -> jnp.ndarray:
    return jnp.stack([self.x, self.y, self.z], axis=-1)

  @classmethod
  def from_array(cls, array):
    return cls(*[row for row in jnp.moveaxis(array, source=-1, destination=0)])

  def __getstate__(self):
    return (VERSION,
            [np.asarray(self.x),
             np.asarray(self.y),
             np.asarray(self.z)])

  def __setstate__(self, state):
    version, state = state
    del version
    for i, letter in enumerate('xyz'):
      object.__setattr__(self, letter, state[i])


def square_euclidean_distance(vec1: Vec3Array,
                              vec2: Vec3Array,
                              epsilon: float = 1e-6) -> Float:
  """Computes square of euclidean distance between 'vec1' and 'vec2'.

  Args:
    vec1: Vec3Array to compute  distance to
    vec2: Vec3Array to compute  distance from, should be
          broadcast compatible with 'vec1'
    epsilon: distance is clipped from below to be at least epsilon

  Returns:
    Array of square euclidean distances;
    shape will be result of broadcasting 'vec1' and 'vec2'
  """
  difference = vec1 - vec2
  distance = difference.dot(difference)
  if epsilon:
    distance = jnp.maximum(distance, epsilon)
  return distance


def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float:
  return vector1.dot(vector2)


def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float:
  return vector1.cross(vector2)


def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float:
  return vector.norm(epsilon)


def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array:
  return vector.normalized(epsilon)


def euclidean_distance(vec1: Vec3Array,
                       vec2: Vec3Array,
                       epsilon: float = 1e-6) -> Float:
  """Computes euclidean distance between 'vec1' and 'vec2'.

  Args:
    vec1: Vec3Array to compute euclidean distance to
    vec2: Vec3Array to compute euclidean distance from, should be
          broadcast compatible with 'vec1'
    epsilon: distance is clipped from below to be at least epsilon

  Returns:
    Array of euclidean distances;
    shape will be result of broadcasting 'vec1' and 'vec2'
  """
  distance_sq = square_euclidean_distance(vec1, vec2, epsilon**2)
  distance = jnp.sqrt(distance_sq)
  return distance


def dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array,
                   d: Vec3Array) -> Float:
  """Computes torsion angle for a quadruple of points.

  For points (a, b, c, d), this is the angle between the planes defined by
  points (a, b, c) and (b, c, d). It is also known as the dihedral angle.

  Arguments:
    a: A Vec3Array of coordinates.
    b: A Vec3Array of coordinates.
    c: A Vec3Array of coordinates.
    d: A Vec3Array of coordinates.

  Returns:
    A tensor of angles in radians: [-pi, pi].
  """
  v1 = a - b
  v2 = b - c
  v3 = d - c

  c1 = v1.cross(v2)
  c2 = v3.cross(v2)
  c3 = c2.cross(c1)

  v2_mag = v2.norm()
  return jnp.arctan2(c3.dot(v2), v2_mag * c1.dot(c2))


def random_gaussian_vector(shape, key, dtype=jnp.float32):
  vec_array = jax.random.normal(key, shape + (3,), dtype)
  return Vec3Array.from_array(vec_array)

alphamissense/model/geometry/vector.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Function to stack repeats of a layer function without shared parameters."""

import collections
import contextlib
import functools
import inspect
from typing import Any, Callable, Optional, Union

import haiku as hk
import jax
import jax.numpy as jnp

LayerStackCarry = collections.namedtuple('LayerStackCarry', ['x', 'rng'])
LayerStackScanned = collections.namedtuple('LayerStackScanned',
                                           ['i', 'args_ys'])

# WrappedFn should take in arbitrarily nested `jnp.ndarray`, and return the
# exact same type. We cannot express this with `typing`. So we just use it
# to inform the user. In reality, the typing below will accept anything.
NestedArray = Any
WrappedFn = Callable[..., Union[NestedArray, tuple[NestedArray]]]


def _check_no_varargs(f):
  if list(inspect.signature(
      f).parameters.values())[0].kind == inspect.Parameter.VAR_POSITIONAL:
    raise ValueError(
        'The function `f` should not have any `varargs` (that is *args) '
        'argument. Instead, it should only use explicit positional'
        'arguments.')


@contextlib.contextmanager
def nullcontext():
  yield


def maybe_with_rng(key):
  if key is not None:
    return hk.with_rng(key)
  else:
    return nullcontext()


def maybe_fold_in(key, data):
  if key is not None:
    return jax.random.fold_in(key, data)
  else:
    return None


class _LayerStack(hk.Module):
  """Module to compose parameterized functions, implemented as a scan."""

  def __init__(self,
               count: int,
               unroll: int,
               name: Optional[str] = None):
    """Iterate a function `f` `count` times, with non-shared parameters."""
    super().__init__(name=name)
    self._count = count
    self._unroll = unroll

  def __call__(self, x, *args_ys):
    count = self._count
    if hk.running_init():
      # At initialization time, we run just one layer but add an extra first
      # dimension to every initialized tensor, making sure to use different
      # random keys for different slices.
      def creator(next_creator, shape, dtype, init, context):
        del context

        def multi_init(shape, dtype):
          assert shape[0] == count
          key = hk.maybe_next_rng_key()

          def rng_context_init(slice_idx):
            slice_key = maybe_fold_in(key, slice_idx)
            with maybe_with_rng(slice_key):
              return init(shape[1:], dtype)

          return jax.vmap(rng_context_init)(jnp.arange(count))

        return next_creator((count,) + tuple(shape), dtype, multi_init)

      def getter(next_getter, value, context):
        trailing_dims = len(context.original_shape) + 1
        sliced_value = jax.lax.index_in_dim(
            value, index=0, axis=value.ndim - trailing_dims, keepdims=False)
        return next_getter(sliced_value)

      with hk.experimental.custom_creator(
          creator), hk.experimental.custom_getter(getter):
        if len(args_ys) == 1 and args_ys[0] is None:
          args0 = (None,)
        else:
          args0 = [
              jax.lax.dynamic_index_in_dim(ys, 0, keepdims=False)
              for ys in args_ys
          ]
        x, z = self._call_wrapped(x, *args0)
        if z is None:
          return x, z

        # Broadcast state to hold each layer state.
        def broadcast_state(layer_state):
          return jnp.broadcast_to(
              layer_state, [count,] + list(layer_state.shape))
        zs = jax.tree_util.tree_map(broadcast_state, z)
        return x, zs
    else:
      # Use scan during apply, threading through random seed so that it's
      # unique for each layer.
      def layer(carry: LayerStackCarry, scanned: LayerStackScanned):
        rng = carry.rng

        def getter(next_getter, value, context):
          # Getter slices the full param at the current loop index.
          trailing_dims = len(context.original_shape) + 1
          assert value.shape[value.ndim - trailing_dims] == count, (
              f'Attempting to use a parameter stack of size '
              f'{value.shape[value.ndim - trailing_dims]} for a LayerStack of '
              f'size {count}.')

          sliced_value = jax.lax.dynamic_index_in_dim(
              value, scanned.i, axis=value.ndim - trailing_dims, keepdims=False)
          return next_getter(sliced_value)

        with hk.experimental.custom_getter(getter):
          if rng is None:
            out_x, z = self._call_wrapped(carry.x, *scanned.args_ys)
          else:
            rng, rng_ = jax.random.split(rng)
            with hk.with_rng(rng_):
              out_x, z = self._call_wrapped(carry.x, *scanned.args_ys)
        return LayerStackCarry(x=out_x, rng=rng), z

      carry = LayerStackCarry(x=x, rng=hk.maybe_next_rng_key())
      scanned = LayerStackScanned(i=jnp.arange(count, dtype=jnp.int32),
                                  args_ys=args_ys)

      carry, zs = hk.scan(
          layer, carry, scanned, length=count, unroll=self._unroll)
      return carry.x, zs

  def _call_wrapped(self,
                    x: jnp.ndarray,
                    *args,
                    ) -> tuple[jnp.ndarray, Optional[jnp.ndarray]]:
    raise NotImplementedError()


class _LayerStackNoState(_LayerStack):
  """_LayerStack impl with no per-layer state provided to the function."""

  def __init__(self,
               f: WrappedFn,
               count: int,
               unroll: int,
               name: Optional[str] = None):
    super().__init__(count=count, unroll=unroll, name=name)
    _check_no_varargs(f)
    self._f = f

  @hk.transparent
  def _call_wrapped(self, args, y):
    del y
    ret = self._f(*args)
    if len(args) == 1:
      # If the function takes a single argument, the wrapped function receives
      # a tuple of length 1, and therefore it must return a tuple of length 1.
      ret = (ret,)
    return ret, None


class _LayerStackWithState(_LayerStack):
  """_LayerStack impl with per-layer state provided to the function."""

  def __init__(self,
               f: WrappedFn,
               count: int,
               unroll: int,
               name: Optional[str] = None):
    super().__init__(count=count, unroll=unroll, name=name)
    self._f = f

  @hk.transparent
  def _call_wrapped(self, x, *args):
    return self._f(x, *args)


def layer_stack(num_layers: int,
                with_state=False,
                unroll: int = 1,
                name: Optional[str] = None):
  """Utility to wrap a Haiku function and recursively apply it to an input.

  A function is valid if it uses only explicit position parameters, and
  its return type matches its input type. The position parameters can be
  arbitrarily nested structures with `jnp.ndarray` at the leaf nodes. Note
  that kwargs are not supported, neither are functions with variable number
  of parameters (specified by `*args`).

  If `with_state=False` then the new, wrapped function can be understood as
  performing the following:
  ```
  for i in range(num_layers):
    x = f(x)
  return x
  ```

  And if `with_state=True`, assuming `f` takes two arguments on top of `x`:
  ```
  for i in range(num_layers):
    x, zs[i] = f(x, ys_0[i], ys_1[i])
  return x, zs
  ```
  The code using `layer_stack` for the above function would be:
  ```
  def f(x, y_0, y_1):
    ...
    return new_x, z
  x, zs = layer_stack.layer_stack(num_layers,
                                  with_state=True)(f)(x, ys_0, ys_1)
  ```

  Crucially, any parameters created inside `f` will not be shared across
  iterations.

  Args:
    num_layers: The number of times to iterate the wrapped function.
    with_state: Whether or not to pass per-layer state to the wrapped function.
    unroll: the unroll used by `scan`.
    name: Name of the Haiku context.

  Returns:
    Callable that will produce a layer stack when called with a valid function.
  """
  def iterate(f):
    if with_state:
      @functools.wraps(f)
      def wrapped(x, *args):
        for ys in args:
          assert ys.shape[0] == num_layers
        return _LayerStackWithState(
            f, num_layers, unroll=unroll, name=name)(x, *args)
    else:
      _check_no_varargs(f)
      @functools.wraps(f)
      def wrapped(*args):
        ret = _LayerStackNoState(
            f, num_layers, unroll=unroll, name=name)(args, None)[0]
        if len(args) == 1:
          # If the function takes a single argument, we must also return a
          # single value, and not a tuple of length 1.
          ret = ret[0]
        return ret

    return wrapped
  return iterate

alphamissense/model/layer_stack.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for layer_stack."""

import functools
from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
import scipy.stats
from alphamissense.model import layer_stack


# Suffixes applied by Haiku for repeated module names.
suffixes = [''] + [f'_{i}' for i in range(1, 100)]


def _slice_layers_params(layers_params):
  sliced_layers_params = {}
  for k, v in layers_params.items():
    for inner_k in v:
      for var_slice, suffix in zip(v[inner_k], suffixes):
        k_new = k.split('/')[-1] + suffix
        if k_new not in sliced_layers_params:
          sliced_layers_params[k_new] = {}
        sliced_layers_params[k_new][inner_k] = var_slice
  return sliced_layers_params


class LayerStackTest(parameterized.TestCase):

  @parameterized.parameters([1, 2, 4])
  def test_layer_stack(self, unroll):
    """Compare layer_stack to the equivalent unrolled stack.

    Tests that the layer_stack application of a Haiku layer function is
    equivalent to repeatedly applying the layer function in an unrolled loop.

    Args:
      unroll: Number of unrolled layers.
    """
    num_layers = 20

    def inner_fn(x):
      x += hk.Linear(100, name='linear1')(x)
      x += hk.Linear(100, name='linear2')(x)
      return x

    def outer_fn_unrolled(x):
      for _ in range(num_layers):
        x = inner_fn(x)
      return x

    def outer_fn_layer_stack(x):
      stack = layer_stack.layer_stack(num_layers, unroll=unroll)(inner_fn)
      return stack(x)

    unrolled_fn = hk.transform(outer_fn_unrolled)
    layer_stack_fn = hk.transform(outer_fn_layer_stack)

    x = jax.random.uniform(jax.random.PRNGKey(0), [10, 256, 100])

    rng_init = jax.random.PRNGKey(42)

    params = layer_stack_fn.init(rng_init, x)

    sliced_params = _slice_layers_params(params)

    unrolled_pred = unrolled_fn.apply(sliced_params, None, x)
    layer_stack_pred = layer_stack_fn.apply(params, None, x)

    np.testing.assert_allclose(unrolled_pred, layer_stack_pred)

  def test_layer_stack_multi_args(self):
    """Compare layer_stack to the equivalent unrolled stack.

    Similar to `test_layer_stack`, but use a function that takes more than one
    argument.
    """
    num_layers = 20

    def inner_fn(x, y):
      x_out = x + hk.Linear(100, name='linear1')(y)
      y_out = y + hk.Linear(100, name='linear2')(x)
      return x_out, y_out

    def outer_fn_unrolled(x, y):
      for _ in range(num_layers):
        x, y = inner_fn(x, y)
      return x, y

    def outer_fn_layer_stack(x, y):
      stack = layer_stack.layer_stack(num_layers)(inner_fn)
      return stack(x, y)

    unrolled_fn = hk.transform(outer_fn_unrolled)
    layer_stack_fn = hk.transform(outer_fn_layer_stack)

    x = jax.random.uniform(jax.random.PRNGKey(0), [10, 256, 100])
    y = jax.random.uniform(jax.random.PRNGKey(1), [10, 256, 100])

    rng_init = jax.random.PRNGKey(42)

    params = layer_stack_fn.init(rng_init, x, y)

    sliced_params = _slice_layers_params(params)

    unrolled_x, unrolled_y = unrolled_fn.apply(sliced_params, None, x, y)
    layer_stack_x, layer_stack_y = layer_stack_fn.apply(params, None, x, y)

    np.testing.assert_allclose(unrolled_x, layer_stack_x)
    np.testing.assert_allclose(unrolled_y, layer_stack_y)

  def test_layer_stack_no_varargs(self):
    """Test an error is raised when using a function with varargs."""

    class VarArgsModule(hk.Module):
      """When used, this module should cause layer_stack to raise an Error."""

      def __call__(self, *args):
        return args

    class NoVarArgsModule(hk.Module):
      """This module should be fine to use with layer_stack."""

      def __call__(self, x):
        return x

    def build_and_init_stack(module_class):
      def stack_fn(x):
        module = module_class()
        return layer_stack.layer_stack(1)(module)(x)

      stack = hk.without_apply_rng(hk.transform(stack_fn))
      stack.init(jax.random.PRNGKey(1729), jnp.ones([5]))

    build_and_init_stack(NoVarArgsModule)
    with self.assertRaisesRegex(
        ValueError, 'The function `f` should not have any `varargs`'):
      build_and_init_stack(VarArgsModule)

  @parameterized.parameters([1, 2, 4])
  def test_layer_stack_grads(self, unroll):
    """Compare layer_stack gradients to the equivalent unrolled stack.

    Tests that the layer_stack application of a Haiku layer function is
    equivalent to repeatedly applying the layer function in an unrolled loop.

    Args:
      unroll: Number of unrolled layers.
    """
    num_layers = 20

    def inner_fn(x):
      x += hk.Linear(100, name='linear1')(x)
      x += hk.Linear(100, name='linear2')(x)
      return x

    def outer_fn_unrolled(x):
      for _ in range(num_layers):
        x = inner_fn(x)
      return x

    def outer_fn_layer_stack(x):
      stack = layer_stack.layer_stack(num_layers, unroll=unroll)(inner_fn)
      return stack(x)

    unrolled_fn = hk.transform(outer_fn_unrolled)
    layer_stack_fn = hk.transform(outer_fn_layer_stack)

    x = jax.random.uniform(jax.random.PRNGKey(0), [10, 256, 100])

    rng_init = jax.random.PRNGKey(42)

    params = layer_stack_fn.init(rng_init, x)

    sliced_params = _slice_layers_params(params)

    unrolled_grad = jax.grad(
        lambda p, x: jnp.mean(unrolled_fn.apply(p, None, x)))(sliced_params, x)
    layer_stack_grad = jax.grad(
        lambda p, x: jnp.mean(layer_stack_fn.apply(p, None, x)))(params, x)

    assert_fn = functools.partial(
        np.testing.assert_allclose, atol=1e-4, rtol=1e-4)

    jax.tree_map(assert_fn, unrolled_grad,
                 _slice_layers_params(layer_stack_grad))

  def test_random(self):
    """Random numbers should be handled correctly."""
    n = 100

    @hk.transform
    @layer_stack.layer_stack(n)
    def add_random(x):
      x = x + jax.random.normal(hk.next_rng_key())
      return x

    # Evaluate a bunch of times
    key, *keys = jax.random.split(jax.random.PRNGKey(7), 1024 + 1)
    params = add_random.init(key, 0.)
    apply_fn = jax.jit(add_random.apply)
    values = [apply_fn(params, key, 0.) for key in keys]

    # Should be roughly N(0, sqrt(n))
    cdf = scipy.stats.norm(scale=np.sqrt(n)).cdf
    _, p = scipy.stats.kstest(values, cdf)
    self.assertLess(0.3, p)

  def test_threading(self):
    """Test @layer_stack when the function gets per-layer state."""
    n = 5

    @layer_stack.layer_stack(n, with_state=True)
    def f(x, y):
      x = x + y * jax.nn.one_hot(y, len(x)) / 10
      return x, 2 * y

    @hk.without_apply_rng
    @hk.transform
    def g(x, ys):
      x, zs = f(x, ys)
      # Check here to catch issues at init time
      self.assertEqual(zs.shape, (n,))
      return x, zs

    rng = jax.random.PRNGKey(7)
    x = np.zeros(n)
    ys = np.arange(n).astype(np.float32)
    params = g.init(rng, x, ys)
    x, zs = g.apply(params, x, ys)
    self.assertTrue(np.allclose(x, [0, .1, .2, .3, .4]))
    self.assertTrue(np.all(zs == 2 * ys))

  def test_nested_stacks(self):
    def stack_fn(x):
      def layer_fn(x):
        return hk.Linear(100)(x)

      outer_fn = layer_stack.layer_stack(10)(layer_fn)

      layer_outer = layer_stack.layer_stack(20)(outer_fn)
      return layer_outer(x)

    hk_mod = hk.transform(stack_fn)
    apply_rng, init_rng = jax.random.split(jax.random.PRNGKey(0))

    params = hk_mod.init(init_rng, jnp.zeros([10, 100]))

    hk_mod.apply(params, apply_rng, jnp.zeros([10, 100]))

    p, = params.values()

    assert p['w'].shape == (10, 20, 100, 100)
    assert p['b'].shape == (10, 20, 100)

  def test_with_state_multi_args(self):
    """Test layer_stack with state with multiple arguments."""
    width = 4
    batch_size = 5
    stack_height = 3

    def f_with_multi_args(x, a, b):
      return hk.Linear(
          width, w_init=hk.initializers.Constant(
              jnp.eye(width)))(x) * a + b, None

    @hk.without_apply_rng
    @hk.transform
    def hk_fn(x):
      return layer_stack.layer_stack(
          stack_height,
          with_state=True)(f_with_multi_args)(x, jnp.full([stack_height], 2.),
                                              jnp.ones([stack_height]))

    x = jnp.zeros([batch_size, width])
    key_seq = hk.PRNGSequence(19)
    params = hk_fn.init(next(key_seq), x)
    output, z = hk_fn.apply(params, x)
    self.assertIsNone(z)
    self.assertEqual(output.shape, (batch_size, width))
    np.testing.assert_equal(output, np.full([batch_size, width], 7.))

  def test_with_container_state(self):
    width = 2
    batch_size = 2
    stack_height = 3

    def f_with_container_state(x):
      hk_layer = hk.Linear(
          width, w_init=hk.initializers.Constant(jnp.eye(width)))
      layer_output = hk_layer(x)
      layer_state = {
          'raw_output': layer_output,
          'output_projection': jnp.sum(layer_output)
      }
      return layer_output + jnp.ones_like(layer_output), layer_state

    @hk.without_apply_rng
    @hk.transform
    def hk_fn(x):
      return layer_stack.layer_stack(
          stack_height,
          with_state=True)(f_with_container_state)(x)

    x = jnp.zeros([batch_size, width])
    key_seq = hk.PRNGSequence(19)
    params = hk_fn.init(next(key_seq), x)
    output, z = hk_fn.apply(params, x)
    self.assertEqual(z['raw_output'].shape, (stack_height, batch_size, width))
    self.assertEqual(output.shape, (batch_size, width))
    self.assertEqual(z['output_projection'].shape, (stack_height,))
    np.testing.assert_equal(np.sum(z['output_projection']), np.array(12.))
    np.testing.assert_equal(
        np.all(z['raw_output'] == np.array([0., 1., 2.])[..., None, None]),
        np.array(True))


if __name__ == '__main__':
  absltest.main()

alphamissense/model/layer_stack_test.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Specialized mapping functions."""

import functools
import inspect

from typing import Any, Callable, Optional, Sequence, Union

import haiku as hk
import jax
import jax.numpy as jnp


PYTREE = Any
PYTREE_JAX_ARRAY = Any

partial = functools.partial
PROXY = object()


def _maybe_slice(array, i, slice_size, axis):
  if axis is PROXY:
    return array
  else:
    return jax.lax.dynamic_slice_in_dim(
        array, i, slice_size=slice_size, axis=axis)


def _maybe_get_size(array, axis):
  if axis == PROXY:
    return -1
  else:
    return array.shape[axis]


def _expand_axes(axes, values, name='sharded_apply'):
  values_tree_def = jax.tree_util.tree_flatten(values)[1]
  flat_axes = jax.api_util.flatten_axes(name, values_tree_def, axes)
  # Replace None's with PROXY
  flat_axes = [PROXY if x is None else x for x in flat_axes]
  return jax.tree_util.tree_unflatten(values_tree_def, flat_axes)


def sharded_map(
    fun: Callable[..., PYTREE_JAX_ARRAY],
    shard_size: Union[int, None] = 1,
    in_axes: Union[int, PYTREE] = 0,
    out_axes: Union[int, PYTREE] = 0) -> Callable[..., PYTREE_JAX_ARRAY]:
  """Sharded vmap.

  Maps `fun` over axes, in a way similar to vmap, but does so in shards of
  `shard_size`. This allows a smooth trade-off between memory usage
  (as in a plain map) vs higher throughput (as in a vmap).

  Args:
    fun: Function to apply smap transform to.
    shard_size: Integer denoting shard size.
    in_axes: Either integer or pytree describing which axis to map over for each
      input to `fun`, None denotes broadcasting.
    out_axes: integer or pytree denoting to what axis in the output the mapped
      over axis maps.

  Returns:
    function with smap applied.
  """
  if 'split_rng' in inspect.signature(hk.vmap).parameters:
    vmapped_fun = hk.vmap(fun, in_axes, out_axes, split_rng=False)
  else:
    vmapped_fun = hk.vmap(fun, in_axes, out_axes)
  return sharded_apply(vmapped_fun, shard_size, in_axes, out_axes)


def sharded_apply(
    fun: Callable[..., PYTREE_JAX_ARRAY],
    shard_size: Union[int, None] = 1,
    in_axes: Union[int, PYTREE] = 0,
    out_axes: Union[int, PYTREE] = 0,
    new_out_axes: bool = False) -> Callable[..., PYTREE_JAX_ARRAY]:
  """Sharded apply.

  Applies `fun` over shards to axes, in a way similar to vmap,
  but does so in shards of `shard_size`. Shards are stacked after.
  This allows a smooth trade-off between
  memory usage (as in a plain map) vs higher throughput (as in a vmap).

  Args:
    fun: Function to apply smap transform to.
    shard_size: Integer denoting shard size.
    in_axes: Either integer or pytree describing which axis to map over for each
      input to `fun`, None denotes broadcasting.
    out_axes: integer or pytree denoting to what axis in the output the mapped
      over axis maps.
    new_out_axes: whether to stack outputs on new axes. This assumes that the
      output sizes for each shard (including the possible remainder shard) are
      the same.

  Returns:
    function with smap applied.
  """
  docstr = ('Mapped version of {fun}. Takes similar arguments to {fun} '
            'but with additional array axes over which {fun} is mapped.')
  if new_out_axes:
    raise NotImplementedError('New output axes not yet implemented.')

  # shard size None denotes no sharding
  if shard_size is None:
    return fun

  @jax.util.wraps(fun, docstr=docstr)
  def mapped_fn(*args):
    # Expand in axes and Determine Loop range
    in_axes_ = _expand_axes(in_axes, args)

    in_sizes = jax.tree_map(_maybe_get_size, args, in_axes_)
    flat_sizes = jax.tree_util.tree_flatten(in_sizes)[0]
    in_size = max(flat_sizes)
    assert all(i in {in_size, -1} for i in flat_sizes)

    num_extra_shards = (in_size - 1) // shard_size

    # Fix Up if necessary
    last_shard_size = in_size % shard_size
    last_shard_size = shard_size if last_shard_size == 0 else last_shard_size

    def apply_fun_to_slice(slice_start, slice_size):
      input_slice = jax.tree_map(
          lambda vec, axis: _maybe_slice(vec, slice_start, slice_size, axis),
          args, in_axes_)
      return fun(*input_slice)

    remainder_shape_dtype = hk.eval_shape(
        partial(apply_fun_to_slice, 0, last_shard_size))
    out_dtypes = jax.tree_map(lambda x: x.dtype, remainder_shape_dtype)
    out_shapes = jax.tree_map(lambda x: x.shape, remainder_shape_dtype)
    out_axes_ = _expand_axes(out_axes, remainder_shape_dtype)

    if num_extra_shards > 0:
      regular_shard_shape_dtype = hk.eval_shape(
          partial(apply_fun_to_slice, 0, shard_size))
      shard_shapes = jax.tree_map(lambda x: x.shape, regular_shard_shape_dtype)

      def make_output_shape(axis, shard_shape, remainder_shape):
        return shard_shape[:axis] + (
            shard_shape[axis] * num_extra_shards +
            remainder_shape[axis],) + shard_shape[axis + 1:]

      out_shapes = jax.tree_map(make_output_shape, out_axes_, shard_shapes,
                                out_shapes)

    # Calls dynamic Update slice with different argument order
    # This is here since tree_map only works with positional arguments
    def dynamic_update_slice_in_dim(full_array, update, axis, i):
      return jax.lax.dynamic_update_slice_in_dim(full_array, update, i, axis)

    def compute_shard(outputs, slice_start, slice_size):
      slice_out = apply_fun_to_slice(slice_start, slice_size)
      update_slice = partial(
          dynamic_update_slice_in_dim, i=slice_start)
      return jax.tree_map(update_slice, outputs, slice_out, out_axes_)

    def scan_iteration(outputs, i):
      new_outputs = compute_shard(outputs, i, shard_size)
      return new_outputs, ()

    slice_starts = jnp.arange(0, in_size - shard_size + 1, shard_size)

    def allocate_buffer(dtype, shape):
      return jnp.zeros(shape, dtype=dtype)

    outputs = jax.tree_map(allocate_buffer, out_dtypes, out_shapes)

    if slice_starts.shape[0] > 0:
      outputs, _ = hk.scan(scan_iteration, outputs, slice_starts)

    if last_shard_size != shard_size:
      remainder_start = in_size - last_shard_size
      outputs = compute_shard(outputs, remainder_start, last_shard_size)

    return outputs

  return mapped_fn


def inference_subbatch(
    module: Callable[..., PYTREE_JAX_ARRAY],
    subbatch_size: int,
    batched_args: Sequence[PYTREE_JAX_ARRAY],
    nonbatched_args: Sequence[PYTREE_JAX_ARRAY],
    low_memory: bool = True,
    input_subbatch_dim: int = 0,
    output_subbatch_dim: Optional[int] = None) -> PYTREE_JAX_ARRAY:
  """Run through subbatches (like batch apply but with split and concat)."""
  assert len(batched_args) > 0  # pylint: disable=g-explicit-length-test

  if not low_memory:
    args = [*batched_args, *nonbatched_args]
    return module(*args)

  if output_subbatch_dim is None:
    output_subbatch_dim = input_subbatch_dim

  def run_module(*batched_args):
    args = [*batched_args, *nonbatched_args]
    return module(*args)
  sharded_module = sharded_apply(run_module,
                                 shard_size=subbatch_size,
                                 in_axes=input_subbatch_dim,
                                 out_axes=output_subbatch_dim)
  return sharded_module(*batched_args)

alphamissense/model/mapping.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Modules and code used in the core part of AlphaFold.

The structure generation code is in 'folding.py'.
"""

import functools
from typing import Any, MutableMapping, Union

import haiku as hk
import jax
import jax.numpy as jnp
import ml_collections
from alphamissense.common import residue_constants
from alphamissense.model import common_modules
from alphamissense.model import mapping
from alphamissense.model import prng
from alphamissense.model import utils

_SOFTMAX_MASK = -1e9

FeatureDict = MutableMapping[str, jnp.ndarray]


def softmax_cross_entropy(logits: jnp.ndarray,
                          labels: jnp.ndarray
                          ) -> jnp.ndarray:
  """Computes softmax cross entropy given logits and one-hot class labels."""
  loss = -jnp.sum(labels * jax.nn.log_softmax(logits), axis=-1)
  return jnp.asarray(loss)


def apply_dropout(*,
                  tensor: jnp.ndarray,
                  safe_key: prng.SafeKey,
                  rate: float,
                  is_training: bool,
                  broadcast_dim: int | None = None,
                  ) -> jnp.ndarray:
  """Applies dropout to a tensor."""
  if is_training and rate != 0.0:
    shape = list(tensor.shape)
    if broadcast_dim is not None:
      shape[broadcast_dim] = 1
    keep_rate = 1.0 - rate
    keep = jax.random.bernoulli(safe_key.get(), keep_rate, shape=shape)
    return keep * tensor / keep_rate
  else:
    return tensor


def dropout_wrapper(module: Any,
                    input_act: jnp.ndarray,
                    mask: jnp.ndarray,
                    safe_key: prng.SafeKey,
                    global_config: ml_collections.ConfigDict,
                    output_act: jnp.ndarray | None = None,
                    is_training: bool = True,
                    **kwargs
                    ) -> jnp.ndarray:
  """Applies module + dropout + residual update."""
  if output_act is None:
    output_act = input_act

  gc = global_config
  residual = module(input_act, mask, is_training=is_training, **kwargs)
  dropout_rate = 0.0 if gc.deterministic else module.config.dropout_rate

  # Will override `is_training` to True if want to use dropout.
  should_apply_dropout = True if gc.eval_dropout else is_training

  if module.config.shared_dropout:
    if module.config.orientation == 'per_row':
      broadcast_dim = 0
    else:
      broadcast_dim = 1
  else:
    broadcast_dim = None

  residual = apply_dropout(tensor=residual,
                           safe_key=safe_key,
                           rate=dropout_rate,
                           is_training=should_apply_dropout,
                           broadcast_dim=broadcast_dim)

  new_act = output_act + residual

  return new_act


class Transition(hk.Module):
  """Transition layer.

  Jumper et al. (2021) Suppl. Alg. 9 "MSATransition"
  Jumper et al. (2021) Suppl. Alg. 15 "PairTransition"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'transition_block'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               act: jnp.ndarray,
               mask: jnp.ndarray,
               is_training: bool = True
               ) -> jnp.ndarray:
    """Builds Transition module.

    Arguments:
      act: A tensor of queries of size [batch_size, N_res, N_channel].
      mask: A tensor denoting the mask of size [batch_size, N_res].
      is_training: Whether the module is in training mode.

    Returns:
      A float32 tensor of size [batch_size, N_res, N_channel].
    """
    del mask  # Not applied since operation is element-wise.
    _, _, nc = act.shape

    num_intermediate = int(nc * self.config.num_intermediate_factor)

    act = common_modules.LayerNorm(
        axis=[-1],
        create_scale=True,
        create_offset=True,
        name='input_layer_norm')(
            act)

    transition_module = hk.Sequential([
        common_modules.Linear(
            num_intermediate,
            initializer='relu',
            name='transition1'), jax.nn.relu,
        common_modules.Linear(
            nc,
            initializer=utils.final_init(self.global_config),
            name='transition2')
    ])

    act = mapping.inference_subbatch(
        transition_module,
        self.global_config.subbatch_size,
        batched_args=[act],
        nonbatched_args=[],
        low_memory=not is_training)

    return act


def glorot_uniform():
  return hk.initializers.VarianceScaling(scale=1.0,
                                         mode='fan_avg',
                                         distribution='uniform')


class Attention(hk.Module):
  """Multihead attention."""

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               output_dim: int,
               name: str = 'attention'):
    super().__init__(name=name)

    self.config = config
    self.global_config = global_config
    self.output_dim = output_dim

  def __call__(self,
               q_data: jnp.ndarray,
               m_data: jnp.ndarray,
               mask: jnp.ndarray,
               nonbatched_bias: jnp.ndarray | None = None
               ) -> jnp.ndarray:
    """Builds Attention module.

    Arguments:
      q_data: A tensor of queries, shape [batch_size, N_queries, q_channels].
      m_data: A tensor of memories from which the keys and values are
        projected, shape [batch_size, N_keys, m_channels].
      mask: A mask for the attention, shape [batch_size, N_queries, N_keys].
      nonbatched_bias: Shared bias, shape [N_queries, N_keys].

    Returns:
      A float32 tensor of shape [batch_size, N_queries, output_dim].
    """
    # Sensible default for when the config keys are missing
    key_dim = self.config.get('key_dim', int(q_data.shape[-1]))
    value_dim = self.config.get('value_dim', int(m_data.shape[-1]))
    num_head = self.config.num_head
    assert key_dim % num_head == 0
    assert value_dim % num_head == 0
    key_dim = key_dim // num_head
    value_dim = value_dim // num_head

    q_weights = hk.get_parameter(
        'query_w', shape=(q_data.shape[-1], num_head, key_dim),
        dtype=q_data.dtype,
        init=glorot_uniform())
    k_weights = hk.get_parameter(
        'key_w', shape=(m_data.shape[-1], num_head, key_dim),
        dtype=q_data.dtype,
        init=glorot_uniform())
    v_weights = hk.get_parameter(
        'value_w', shape=(m_data.shape[-1], num_head, value_dim),
        dtype=q_data.dtype,
        init=glorot_uniform())

    q = jnp.einsum('bqa,ahc->bqhc', q_data, q_weights) * key_dim**(-0.5)
    k = jnp.einsum('bka,ahc->bkhc', m_data, k_weights)
    v = jnp.einsum('bka,ahc->bkhc', m_data, v_weights)
    logits = jnp.einsum('bqhc,bkhc->bhqk', q, k)
    if nonbatched_bias is not None:
      logits += jnp.expand_dims(nonbatched_bias, axis=0)
    logits = jnp.where(mask, logits, _SOFTMAX_MASK)
    weights = utils.stable_softmax(logits)
    weighted_avg = jnp.einsum('bhqk,bkhc->bqhc', weights, v)

    if self.global_config.zero_init:
      init = hk.initializers.Constant(0.0)
    else:
      init = glorot_uniform()

    if self.config.gating:
      gating_weights = hk.get_parameter(
          'gating_w',
          shape=(q_data.shape[-1], num_head, value_dim),
          dtype=q_data.dtype,
          init=hk.initializers.Constant(0.0))
      gating_bias = hk.get_parameter(
          'gating_b',
          shape=(num_head, value_dim),
          dtype=q_data.dtype,
          init=hk.initializers.Constant(1.0))

      gate_values = jnp.einsum('bqc, chv->bqhv', q_data,
                               gating_weights) + gating_bias

      gate_values = jax.nn.sigmoid(gate_values)

      weighted_avg *= gate_values

    o_weights = hk.get_parameter(
        'output_w', shape=(num_head, value_dim, self.output_dim),
        dtype=q_data.dtype,
        init=init)
    o_bias = hk.get_parameter(
        'output_b', shape=(self.output_dim,),
        dtype=q_data.dtype,
        init=hk.initializers.Constant(0.0))

    output = jnp.einsum('bqhc,hco->bqo', weighted_avg, o_weights) + o_bias

    return output


class GlobalAttention(hk.Module):
  """Global attention.

  Jumper et al. (2021) Suppl. Alg. 19 "MSAColumnGlobalAttention" lines 2-7
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               output_dim: int,
               name: str = 'attention'):
    super().__init__(name=name)

    self.config = config
    self.global_config = global_config
    self.output_dim = output_dim

  def __call__(self,
               q_data: jnp.ndarray,
               m_data: jnp.ndarray,
               q_mask: jnp.ndarray,
               ) -> jnp.ndarray:
    """Builds GlobalAttention module.

    Arguments:
      q_data: A tensor of queries with size [batch_size, N_queries,
        q_channels]
      m_data: A tensor of memories from which the keys and values
        projected. Size [batch_size, N_keys, m_channels]
      q_mask: A binary mask for q_data with zeros in the padded sequence
        elements and ones otherwise. Size [batch_size, N_queries, q_channels]
        (or broadcastable to this shape).

    Returns:
      A float32 tensor of size [batch_size, N_queries, output_dim].
    """
    # Sensible default for when the config keys are missing
    key_dim = self.config.get('key_dim', int(q_data.shape[-1]))
    value_dim = self.config.get('value_dim', int(m_data.shape[-1]))
    num_head = self.config.num_head
    assert key_dim % num_head == 0
    assert value_dim % num_head == 0
    key_dim = key_dim // num_head
    value_dim = value_dim // num_head

    q_weights = hk.get_parameter(
        'query_w', shape=(q_data.shape[-1], num_head, key_dim),
        dtype=q_data.dtype,
        init=glorot_uniform())
    k_weights = hk.get_parameter(
        'key_w', shape=(m_data.shape[-1], key_dim),
        dtype=q_data.dtype,
        init=glorot_uniform())
    v_weights = hk.get_parameter(
        'value_w', shape=(m_data.shape[-1], value_dim),
        dtype=q_data.dtype,
        init=glorot_uniform())

    v = jnp.einsum('bka,ac->bkc', m_data, v_weights)

    q_avg = utils.mask_mean(q_mask, q_data, axis=1)

    q = jnp.einsum('ba,ahc->bhc', q_avg, q_weights) * key_dim**(-0.5)
    k = jnp.einsum('bka,ac->bkc', m_data, k_weights)
    bias = q_mask[:, None, :, 0]
    logits = jnp.einsum('bhc,bkc->bhk', q, k)
    logits = jnp.where(bias, logits, _SOFTMAX_MASK)
    weights = utils.stable_softmax(logits)
    weighted_avg = jnp.einsum('bhk,bkc->bhc', weights, v)

    if self.global_config.zero_init:
      init = hk.initializers.Constant(0.0)
    else:
      init = glorot_uniform()

    o_weights = hk.get_parameter(
        'output_w', shape=(num_head, value_dim, self.output_dim),
        dtype=q_data.dtype,
        init=init)
    o_bias = hk.get_parameter(
        'output_b', shape=(self.output_dim,),
        dtype=q_data.dtype,
        init=hk.initializers.Constant(0.0))

    if self.config.gating:
      gating_weights = hk.get_parameter(
          'gating_w',
          shape=(q_data.shape[-1], num_head, value_dim),
          dtype=q_data.dtype,
          init=hk.initializers.Constant(0.0))
      gating_bias = hk.get_parameter(
          'gating_b',
          shape=(num_head, value_dim),
          dtype=q_data.dtype,
          init=hk.initializers.Constant(1.0))

      gate_values = jnp.einsum('bqc, chv->bqhv', q_data, gating_weights)
      gate_values = jax.nn.sigmoid(gate_values + gating_bias)
      weighted_avg = weighted_avg[:, None] * gate_values
      output = jnp.einsum('bqhc,hco->bqo', weighted_avg, o_weights) + o_bias
    else:
      output = jnp.einsum('bhc,hco->bo', weighted_avg, o_weights) + o_bias
      output = output[:, None]
    return output


class MSARowAttentionWithPairBias(hk.Module):
  """MSA per-row attention biased by the pair representation.

  Jumper et al. (2021) Suppl. Alg. 7 "MSARowAttentionWithPairBias"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'msa_row_attention_with_pair_bias'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               msa_act: jnp.ndarray,
               msa_mask: jnp.ndarray,
               pair_act: jnp.ndarray,
               is_training: bool = False
               ) -> jnp.ndarray:
    """Builds MSARowAttentionWithPairBias module.

    Arguments:
      msa_act: [N_seq, N_res, c_m] MSA representation.
      msa_mask: [N_seq, N_res] mask of non-padded regions.
      pair_act: [N_res, N_res, c_z] pair representation.
      is_training: Whether the module is in training mode.

    Returns:
      Update to msa_act, shape [N_seq, N_res, c_m].
    """
    c = self.config

    assert len(msa_act.shape) == 3
    assert len(msa_mask.shape) == 2
    assert c.orientation == 'per_row'

    mask = msa_mask[:, None, None, :]
    assert len(mask.shape) == 4

    msa_act = common_modules.LayerNorm(
        axis=[-1], create_scale=True, create_offset=True, name='query_norm')(
            msa_act)

    pair_act = common_modules.LayerNorm(
        axis=[-1],
        create_scale=True,
        create_offset=True,
        name='feat_2d_norm')(
            pair_act)

    init_factor = 1. / jnp.sqrt(int(pair_act.shape[-1]))
    weights = hk.get_parameter(
        'feat_2d_weights',
        shape=(pair_act.shape[-1], c.num_head),
        dtype=msa_act.dtype,
        init=hk.initializers.RandomNormal(stddev=init_factor))
    nonbatched_bias = jnp.einsum('qkc,ch->hqk', pair_act, weights)

    attn_mod = Attention(
        c, self.global_config, msa_act.shape[-1])
    msa_act = mapping.inference_subbatch(
        attn_mod,
        self.global_config.subbatch_size,
        batched_args=[msa_act, msa_act, mask],
        nonbatched_args=[nonbatched_bias],
        low_memory=not is_training)

    return msa_act


class MSAColumnAttention(hk.Module):
  """MSA per-column attention.

  Jumper et al. (2021) Suppl. Alg. 8 "MSAColumnAttention"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'msa_column_attention'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               msa_act: jnp.ndarray,
               msa_mask: jnp.ndarray,
               is_training: bool = False
               ) -> jnp.ndarray:
    """Builds MSAColumnAttention module.

    Arguments:
      msa_act: [N_seq, N_res, c_m] MSA representation.
      msa_mask: [N_seq, N_res] mask of non-padded regions.
      is_training: Whether the module is in training mode.

    Returns:
      Update to msa_act, shape [N_seq, N_res, c_m]
    """
    c = self.config

    assert len(msa_act.shape) == 3
    assert len(msa_mask.shape) == 2
    assert c.orientation == 'per_column'

    msa_act = jnp.swapaxes(msa_act, -2, -3)
    msa_mask = jnp.swapaxes(msa_mask, -1, -2)

    mask = msa_mask[:, None, None, :]
    assert len(mask.shape) == 4

    msa_act = common_modules.LayerNorm(
        axis=[-1], create_scale=True, create_offset=True, name='query_norm')(
            msa_act)

    attn_mod = Attention(
        c, self.global_config, msa_act.shape[-1])
    msa_act = mapping.inference_subbatch(
        attn_mod,
        self.global_config.subbatch_size,
        batched_args=[msa_act, msa_act, mask],
        nonbatched_args=[],
        low_memory=not is_training)

    msa_act = jnp.swapaxes(msa_act, -2, -3)

    return msa_act


class MSAColumnGlobalAttention(hk.Module):
  """MSA per-column global attention.

  Jumper et al. (2021) Suppl. Alg. 19 "MSAColumnGlobalAttention"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'msa_column_global_attention'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               msa_act: jnp.ndarray,
               msa_mask: jnp.ndarray,
               is_training: bool = False
               ) -> jnp.ndarray:
    """Builds MSAColumnGlobalAttention module.

    Arguments:
      msa_act: [N_seq, N_res, c_m] MSA representation.
      msa_mask: [N_seq, N_res] mask of non-padded regions.
      is_training: Whether the module is in training mode.

    Returns:
      Update to msa_act, shape [N_seq, N_res, c_m].
    """
    c = self.config

    assert len(msa_act.shape) == 3
    assert len(msa_mask.shape) == 2
    assert c.orientation == 'per_column'

    msa_act = jnp.swapaxes(msa_act, -2, -3)
    msa_mask = jnp.swapaxes(msa_mask, -1, -2)

    msa_act = common_modules.LayerNorm(
        axis=[-1], create_scale=True, create_offset=True, name='query_norm')(
            msa_act)

    attn_mod = GlobalAttention(
        c, self.global_config, msa_act.shape[-1],
        name='attention')
    # [N_seq, N_res, 1]
    msa_mask = jnp.expand_dims(msa_mask, axis=-1)
    msa_act = mapping.inference_subbatch(
        attn_mod,
        self.global_config.subbatch_size,
        batched_args=[msa_act, msa_act, msa_mask],
        nonbatched_args=[],
        low_memory=not is_training)

    msa_act = jnp.swapaxes(msa_act, -2, -3)

    return msa_act


class TriangleAttention(hk.Module):
  """Triangle Attention.

  Jumper et al. (2021) Suppl. Alg. 13 "TriangleAttentionStartingNode"
  Jumper et al. (2021) Suppl. Alg. 14 "TriangleAttentionEndingNode"
  """

  def __init__(self, config, global_config, name='triangle_attention'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self, pair_act, pair_mask, is_training=False):
    """Builds TriangleAttention module.

    Arguments:
      pair_act: [N_res, N_res, c_z] pair activations tensor
      pair_mask: [N_res, N_res] mask of non-padded regions in the tensor.
      is_training: Whether the module is in training mode.

    Returns:
      Update to pair_act, shape [N_res, N_res, c_z].
    """
    c = self.config

    assert len(pair_act.shape) == 3
    assert len(pair_mask.shape) == 2
    assert c.orientation in ['per_row', 'per_column']

    if c.orientation == 'per_column':
      pair_act = jnp.swapaxes(pair_act, -2, -3)
      pair_mask = jnp.swapaxes(pair_mask, -1, -2)

    mask = pair_mask[:, None, None, :]
    assert len(mask.shape) == 4

    pair_act = common_modules.LayerNorm(
        axis=[-1], create_scale=True, create_offset=True, name='query_norm')(
            pair_act)

    init_factor = 1. / jnp.sqrt(int(pair_act.shape[-1]))
    weights = hk.get_parameter(
        'feat_2d_weights',
        shape=(pair_act.shape[-1], c.num_head),
        dtype=pair_act.dtype,
        init=hk.initializers.RandomNormal(stddev=init_factor))
    nonbatched_bias = jnp.einsum('qkc,ch->hqk', pair_act, weights)

    attn_mod = Attention(
        c, self.global_config, pair_act.shape[-1])
    pair_act = mapping.inference_subbatch(
        attn_mod,
        self.global_config.subbatch_size,
        batched_args=[pair_act, pair_act, mask],
        nonbatched_args=[nonbatched_bias],
        low_memory=not is_training)

    if c.orientation == 'per_column':
      pair_act = jnp.swapaxes(pair_act, -2, -3)

    return pair_act


class MaskedMsaHead(hk.Module):
  """Head to predict MSA at the masked locations.

  The MaskedMsaHead employs a BERT-style objective to reconstruct a masked
  version of the full MSA, based on a linear projection of
  the MSA representation.
  Jumper et al. (2021) Suppl. Sec. 1.9.9 "Masked MSA prediction"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'masked_msa_head'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

    if global_config.multimer_mode:
      self.num_output = len(residue_constants.restypes_with_x_and_gap)
    else:
      self.num_output = config.num_output

  def __call__(self,
               representations: FeatureDict,
               batch: FeatureDict,
               is_training: bool,
               ) -> FeatureDict:
    """Builds MaskedMsaHead module.

    Arguments:
      representations: Dictionary of representations, must contain:
        * 'msa': MSA representation, shape [N_seq, N_res, c_m].
      batch: Batch, unused.
      is_training: Whether the module is in training mode.

    Returns:
      Dictionary containing:
        * 'logits': logits of shape [N_seq, N_res, N_aatype] with
            (unnormalized) log probabilies of predicted aatype at position.
    """
    del batch
    logits = common_modules.Linear(
        self.num_output,
        initializer=utils.final_init(self.global_config),
        name='logits')(
            representations['msa'])
    return dict(logits=logits)

  def loss(self, value: FeatureDict, batch: FeatureDict) -> jnp.ndarray:
    errors = softmax_cross_entropy(
        labels=jax.nn.one_hot(batch['true_msa'], num_classes=self.num_output),
        logits=value['logits'])
    loss = (jnp.sum(errors * batch['bert_mask'], axis=(-2, -1)) /
            (1e-8 + jnp.sum(batch['bert_mask'], axis=(-2, -1))))
    return loss


def _layer_norm(axis: int = -1, name: str = 'layer_norm'):
  return common_modules.LayerNorm(
      axis=axis,
      create_scale=True,
      create_offset=True,
      eps=1e-5,
      use_fast_variance=True,
      scale_init=hk.initializers.Constant(1.),
      offset_init=hk.initializers.Constant(0.),
      param_axis=axis,
      name=name)


class TriangleMultiplication(hk.Module):
  """Triangle multiplication layer ("outgoing" or "incoming").

  Jumper et al. (2021) Suppl. Alg. 11 "TriangleMultiplicationOutgoing"
  Jumper et al. (2021) Suppl. Alg. 12 "TriangleMultiplicationIncoming"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'triangle_multiplication'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               left_act: jnp.ndarray,
               left_mask: jnp.ndarray,
               is_training: bool = True
               ) -> jnp.ndarray:
    """Builds TriangleMultiplication module.

    Arguments:
      left_act: Pair activations, shape [N_res, N_res, c_z]
      left_mask: Pair mask, shape [N_res, N_res].
      is_training: Whether the module is in training mode.

    Returns:
      Outputs, same shape/type as left_act.
    """
    del is_training

    if self.config.fuse_projection_weights:
      return self._fused_triangle_multiplication(left_act, left_mask)
    else:
      return self._triangle_multiplication(left_act, left_mask)

  @hk.transparent
  def _triangle_multiplication(self,
                               left_act: jnp.ndarray,
                               left_mask: jnp.ndarray
                               ) -> jnp.ndarray:
    """Implementation of TriangleMultiplication used in AF2 and AF-M<2.3."""
    c = self.config
    gc = self.global_config

    mask = left_mask[..., None]

    act = common_modules.LayerNorm(
        axis=[-1], create_scale=True, create_offset=True,
        name='layer_norm_input')(left_act)
    input_act = act

    left_projection = common_modules.Linear(
        c.num_intermediate_channel,
        name='left_projection')
    left_proj_act = mask * left_projection(act)

    right_projection = common_modules.Linear(
        c.num_intermediate_channel,
        name='right_projection')
    right_proj_act = mask * right_projection(act)

    left_gate_values = jax.nn.sigmoid(common_modules.Linear(
        c.num_intermediate_channel,
        bias_init=1.,
        initializer=utils.final_init(gc),
        name='left_gate')(act))

    right_gate_values = jax.nn.sigmoid(common_modules.Linear(
        c.num_intermediate_channel,
        bias_init=1.,
        initializer=utils.final_init(gc),
        name='right_gate')(act))

    left_proj_act *= left_gate_values
    right_proj_act *= right_gate_values

    # "Outgoing" edges equation: 'ikc,jkc->ijc'
    # "Incoming" edges equation: 'kjc,kic->ijc'
    # Note on the Suppl. Alg. 11 & 12 notation:
    # For the "outgoing" edges, a = left_proj_act and b = right_proj_act
    # For the "incoming" edges, it's swapped:
    #   b = left_proj_act and a = right_proj_act
    act = jnp.einsum(c.equation, left_proj_act, right_proj_act)

    act = common_modules.LayerNorm(
        axis=[-1],
        create_scale=True,
        create_offset=True,
        name='center_layer_norm')(
            act)

    output_channel = int(input_act.shape[-1])

    act = common_modules.Linear(
        output_channel,
        initializer=utils.final_init(gc),
        name='output_projection')(act)

    gate_values = jax.nn.sigmoid(common_modules.Linear(
        output_channel,
        bias_init=1.,
        initializer=utils.final_init(gc),
        name='gating_linear')(input_act))
    act *= gate_values

    return act

  @hk.transparent
  def _fused_triangle_multiplication(self,
                                     left_act: jnp.ndarray,
                                     left_mask: jnp.ndarray,
                                     ) -> jnp.ndarray:
    """TriangleMultiplication with fused projection weights."""
    mask = left_mask[..., None]
    c = self.config
    gc = self.global_config

    left_act = _layer_norm(axis=-1, name='left_norm_input')(left_act)

    # Both left and right projections are fused into projection.
    projection = common_modules.Linear(
        2*c.num_intermediate_channel, name='projection')
    proj_act = mask * projection(left_act)

    # Both left + right gate are fused into gate_values.
    gate_values = common_modules.Linear(
        2 * c.num_intermediate_channel,
        name='gate',
        bias_init=1.,
        initializer=utils.final_init(gc))(left_act)
    proj_act *= jax.nn.sigmoid(gate_values)

    left_proj_act = proj_act[:, :, :c.num_intermediate_channel]
    right_proj_act = proj_act[:, :, c.num_intermediate_channel:]
    act = jnp.einsum(c.equation, left_proj_act, right_proj_act)

    act = _layer_norm(axis=-1, name='center_norm')(act)

    output_channel = int(left_act.shape[-1])

    act = common_modules.Linear(
        output_channel,
        initializer=utils.final_init(gc),
        name='output_projection')(act)

    gate_values = common_modules.Linear(
        output_channel,
        bias_init=1.,
        initializer=utils.final_init(gc),
        name='gating_linear')(left_act)
    act *= jax.nn.sigmoid(gate_values)

    return act


class DistogramHead(hk.Module):
  """Head to predict a distogram.

  Jumper et al. (2021) Suppl. Sec. 1.9.8 "Distogram prediction"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name='distogram_head'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               representations: FeatureDict,
               batch: FeatureDict,
               is_training: bool
               ) -> FeatureDict:
    """Builds DistogramHead module.

    Arguments:
      representations: Dictionary of representations, must contain:
        * 'pair': pair representation, shape [N_res, N_res, c_z].
      batch: Batch, unused.
      is_training: Whether the module is in training mode.

    Returns:
      Dictionary containing:
        * logits: logits for distogram, shape [N_res, N_res, N_bins].
        * bin_breaks: array containing bin breaks, shape [N_bins - 1,].
    """
    half_logits = common_modules.Linear(
        self.config.num_bins,
        initializer=utils.final_init(self.global_config),
        name='half_logits')(
            representations['pair'])

    logits = half_logits + jnp.swapaxes(half_logits, -2, -3)
    breaks = jnp.linspace(self.config.first_break, self.config.last_break,
                          self.config.num_bins - 1)

    return dict(logits=logits, bin_edges=breaks)

  def loss(self,
           value: FeatureDict,
           batch: FeatureDict
           ) -> jnp.ndarray:
    return _distogram_log_loss(value['logits'], value['bin_edges'],
                               batch, self.config.num_bins)


def _distogram_log_loss(logits: jnp.ndarray,
                        bin_edges: jnp.ndarray,
                        batch: FeatureDict,
                        num_bins: int
                        ) -> jnp.ndarray:
  """Log loss of a distogram."""

  assert len(logits.shape) == 3
  positions = batch['pseudo_beta']
  mask = batch['pseudo_beta_mask']

  assert positions.shape[-1] == 3

  sq_breaks = jnp.square(bin_edges)

  dist2 = jnp.sum(
      jnp.square(
          jnp.expand_dims(positions, axis=-2) -
          jnp.expand_dims(positions, axis=-3)),
      axis=-1,
      keepdims=True)

  true_bins = jnp.sum(dist2 > sq_breaks, axis=-1)

  errors = softmax_cross_entropy(
      labels=jax.nn.one_hot(true_bins, num_bins), logits=logits)

  square_mask = jnp.expand_dims(mask, axis=-2) * jnp.expand_dims(mask, axis=-1)

  avg_error = (
      jnp.sum(errors * square_mask, axis=(-2, -1)) /
      (1e-6 + jnp.sum(square_mask, axis=(-2, -1))))
  return avg_error


class OuterProductMean(hk.Module):
  """Computes mean outer product.

  Jumper et al. (2021) Suppl. Alg. 10 "OuterProductMean"
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               num_output_channel: int,
               name: str = 'outer_product_mean'):
    super().__init__(name=name)
    self.global_config = global_config
    self.config = config
    self.num_output_channel = num_output_channel

  def __call__(self,
               act: jnp.ndarray,
               mask: jnp.ndarray,
               is_training: bool = True
               ) -> jnp.ndarray:
    """Builds OuterProductMean module.

    Arguments:
      act: MSA representation, shape [N_seq, N_res, c_m].
      mask: MSA mask, shape [N_seq, N_res].
      is_training: Whether the module is in training mode.

    Returns:
      Update to pair representation, shape [N_res, N_res, c_z].
    """
    gc = self.global_config
    c = self.config

    mask = mask[..., None]
    act = common_modules.LayerNorm(
        [-1], True, True, name='layer_norm_input')(act)

    left_act = mask * common_modules.Linear(
        c.num_outer_channel,
        initializer='linear',
        name='left_projection')(
            act)

    right_act = mask * common_modules.Linear(
        c.num_outer_channel,
        initializer='linear',
        name='right_projection')(
            act)

    if gc.zero_init:
      init_w = hk.initializers.Constant(0.0)
    else:
      init_w = hk.initializers.VarianceScaling(scale=2., mode='fan_in')

    output_w = hk.get_parameter(
        'output_w',
        shape=(c.num_outer_channel, c.num_outer_channel,
               self.num_output_channel),
        dtype=act.dtype,
        init=init_w)
    output_b = hk.get_parameter(
        'output_b', shape=(self.num_output_channel,),
        dtype=act.dtype,
        init=hk.initializers.Constant(0.0))

    def compute_chunk(left_act):
      # This is equivalent to
      #
      # act = jnp.einsum('abc,ade->dceb', left_act, right_act)
      # act = jnp.einsum('dceb,cef->bdf', act, output_w) + output_b
      #
      # but faster.
      left_act = jnp.transpose(left_act, [0, 2, 1])
      act = jnp.einsum('acb,ade->dceb', left_act, right_act)
      act = jnp.einsum('dceb,cef->dbf', act, output_w) + output_b
      return jnp.transpose(act, [1, 0, 2])

    act = mapping.inference_subbatch(
        compute_chunk,
        c.chunk_size,
        batched_args=[left_act],
        nonbatched_args=[],
        low_memory=True,
        input_subbatch_dim=1,
        output_subbatch_dim=0)

    epsilon = 1e-3
    norm = jnp.einsum('abc,adc->bdc', mask, mask)
    act /= epsilon + norm

    return act


def dgram_from_positions(positions: jnp.ndarray,
                         num_bins: int, min_bin: float, max_bin: float
                         ) -> jnp.ndarray:
  """Compute distogram from amino acid positions.

  Arguments:
    positions: [N_res, 3] Position coordinates.
    num_bins: The number of bins in the distogram.
    min_bin: The left edge of the first bin.
    max_bin: The left edge of the final bin. The final bin catches
        everything larger than `max_bin`.

  Returns:
    Distogram with the specified number of bins.
  """

  def squared_difference(x, y):
    return jnp.square(x - y)

  lower_breaks = jnp.linspace(min_bin, max_bin, num_bins)
  lower_breaks = jnp.square(lower_breaks)
  upper_breaks = jnp.concatenate([lower_breaks[1:],
                                  jnp.array([1e8], dtype=jnp.float32)], axis=-1)
  dist2 = jnp.sum(
      squared_difference(
          jnp.expand_dims(positions, axis=-2),
          jnp.expand_dims(positions, axis=-3)),
      axis=-1, keepdims=True)

  dgram = ((dist2 > lower_breaks).astype(jnp.float32) *
           (dist2 < upper_breaks).astype(jnp.float32))
  return dgram


def pseudo_beta_fn(aatype: jnp.ndarray,
                   all_atom_positions: jnp.ndarray,
                   all_atom_masks: jnp.ndarray | None,
                   ) -> Union[jnp.ndarray, tuple[jnp.ndarray, jnp.ndarray]]:
  """Create pseudo beta features."""

  is_gly = jnp.equal(aatype, residue_constants.restype_order['G'])
  ca_idx = residue_constants.atom_order['CA']
  cb_idx = residue_constants.atom_order['CB']
  pseudo_beta = jnp.where(
      jnp.tile(is_gly[..., None], [1] * len(is_gly.shape) + [3]),
      all_atom_positions[..., ca_idx, :],
      all_atom_positions[..., cb_idx, :])

  if all_atom_masks is not None:
    pseudo_beta_mask = jnp.where(
        is_gly, all_atom_masks[..., ca_idx], all_atom_masks[..., cb_idx])
    pseudo_beta_mask = pseudo_beta_mask.astype(jnp.float32)
    return pseudo_beta, pseudo_beta_mask
  else:
    return pseudo_beta


class EvoformerIteration(hk.Module):
  """Single iteration (block) of Evoformer stack.

  Jumper et al. (2021) Suppl. Alg. 6 "EvoformerStack" lines 2-10
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               is_extra_msa: bool,
               name: str = 'evoformer_iteration'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config
    self.is_extra_msa = is_extra_msa

  def __call__(self,
               activations: FeatureDict,
               masks: FeatureDict,
               is_training: bool,
               safe_key: prng.SafeKey
               ) -> FeatureDict:
    """Builds EvoformerIteration module.

    Arguments:
      activations: Dictionary containing activations:
        * 'msa': MSA activations, shape [N_seq, N_res, c_m].
        * 'pair': pair activations, shape [N_res, N_res, c_z].
      masks: Dictionary of masks:
        * 'msa': MSA mask, shape [N_seq, N_res].
        * 'pair': pair mask, shape [N_res, N_res].
      is_training: Whether the module is in training mode.
      safe_key: prng.SafeKey encapsulating rng key.

    Returns:
      Outputs, same shape/type as act.
    """
    c = self.config
    gc = self.global_config

    msa_act, pair_act = activations['msa'], activations['pair']

    if safe_key is None:
      safe_key = prng.SafeKey(hk.next_rng_key())

    msa_mask, pair_mask = masks['msa'], masks['pair']

    dropout_wrapper_fn = functools.partial(
        dropout_wrapper,
        is_training=is_training,
        global_config=gc)

    unused_safe_key, *sub_keys = safe_key.split(10)
    sub_keys = iter(sub_keys)

    outer_module = OuterProductMean(
        config=c.outer_product_mean,
        global_config=self.global_config,
        num_output_channel=int(pair_act.shape[-1]),
        name='outer_product_mean')
    if c.outer_product_mean.first:
      pair_act = dropout_wrapper_fn(
          outer_module,
          msa_act,
          msa_mask,
          safe_key=next(sub_keys),
          output_act=pair_act)

    msa_act = dropout_wrapper_fn(
        MSARowAttentionWithPairBias(
            c.msa_row_attention_with_pair_bias, gc,
            name='msa_row_attention_with_pair_bias'),
        msa_act,
        msa_mask,
        safe_key=next(sub_keys),
        pair_act=pair_act)

    if not self.is_extra_msa:
      attn_mod = MSAColumnAttention(
          c.msa_column_attention, gc, name='msa_column_attention')
    else:
      attn_mod = MSAColumnGlobalAttention(
          c.msa_column_attention, gc, name='msa_column_global_attention')
    msa_act = dropout_wrapper_fn(
        attn_mod,
        msa_act,
        msa_mask,
        safe_key=next(sub_keys))

    msa_act = dropout_wrapper_fn(
        Transition(c.msa_transition, gc, name='msa_transition'),
        msa_act,
        msa_mask,
        safe_key=next(sub_keys))

    if not c.outer_product_mean.first:
      pair_act = dropout_wrapper_fn(
          outer_module,
          msa_act,
          msa_mask,
          safe_key=next(sub_keys),
          output_act=pair_act)

    pair_act = dropout_wrapper_fn(
        TriangleMultiplication(c.triangle_multiplication_outgoing, gc,
                               name='triangle_multiplication_outgoing'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))
    pair_act = dropout_wrapper_fn(
        TriangleMultiplication(c.triangle_multiplication_incoming, gc,
                               name='triangle_multiplication_incoming'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))

    pair_act = dropout_wrapper_fn(
        TriangleAttention(c.triangle_attention_starting_node, gc,
                          name='triangle_attention_starting_node'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))
    pair_act = dropout_wrapper_fn(
        TriangleAttention(c.triangle_attention_ending_node, gc,
                          name='triangle_attention_ending_node'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))

    pair_act = dropout_wrapper_fn(
        Transition(c.pair_transition, gc, name='pair_transition'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))

    return {'msa': msa_act, 'pair': pair_act}

alphamissense/model/modules.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Modules and code used in the core part of AlphaMissense."""

import math
from typing import Any, Mapping, MutableMapping, Optional, Sequence

from absl import logging
import haiku as hk
import jax
import jax.numpy as jnp
import ml_collections
import tree
from alphamissense.common import residue_constants
from alphamissense.model import common_modules
from alphamissense.model import folding_multimer
from alphamissense.model import layer_stack
from alphamissense.model import modules
from alphamissense.model import prng
from alphamissense.model import utils

FeatureDict = MutableMapping[str, jnp.ndarray]
OutputDict = MutableMapping[str, Any]


def gumbel_noise(key: jnp.ndarray, shape: Sequence[int]) -> jnp.ndarray:
  """Generate Gumbel Noise of given Shape.

  This generates samples from Gumbel(0, 1).

  Args:
    key: Jax random number key.
    shape: Shape of noise to return.

  Returns:
    Gumbel noise of given shape.
  """
  epsilon = 1e-6
  uniform = utils.padding_consistent_rng(jax.random.uniform)
  uniform_noise = uniform(
      key, shape=shape, dtype=jnp.float32, minval=0., maxval=1.)
  gumbel = -jnp.log(-jnp.log(uniform_noise + epsilon) + epsilon)
  return gumbel


def gumbel_argsort_sample_idx(key: jnp.ndarray,
                              logits: jnp.ndarray) -> jnp.ndarray:
  """Samples with replacement from a distribution given by 'logits'.

  This uses Gumbel trick to implement the sampling an efficient manner. For a
  distribution over k items this samples k times without replacement, so this
  is effectively sampling a random permutation with probabilities over the
  permutations derived from the logprobs.

  Args:
    key: prng key.
    logits: Logarithm of probabilities to sample from, probabilities can be
      unnormalized.

  Returns:
    Sample from logprobs in one-hot form.
  """
  z = gumbel_noise(key, logits.shape)
  # This construction is equivalent to jnp.argsort, but using a non stable sort,
  # since stable sort's aren't supported by jax2tf.
  axis = len(logits.shape) - 1
  iota = jax.lax.broadcasted_iota(jnp.int64, logits.shape, axis)
  _, perm = jax.lax.sort_key_val(
      logits + z, iota, dimension=-1, is_stable=False)
  return perm[::-1]


def sample_msa(key: prng.SafeKey, batch: FeatureDict, max_seq: int
               ) -> FeatureDict:
  """Sample MSA randomly, remaining sequences are stored as `extra_*`.

  Args:
    key: safe key for random number generation.
    batch: batch to sample msa from.
    max_seq: number of sequences to sample.
  Returns:
    Protein with sampled msa.
  """
  # Sample uniformly among sequences with at least one non-masked position.
  logits = (jnp.clip(jnp.sum(batch['msa_mask'], axis=-1), 0., 1.) - 1.) * 1e6
  # cluster_bias_mask can be used to ensure that the marked MSA rows are always
  # sampled and their positions are preserved. In AlphaMissense this ensures
  # that first row is always the target sequence and second row is the
  # masked variant row.
  cluster_bias_mask = batch['cluster_bias_mask']
  logits += cluster_bias_mask * 1e6
  index_order = gumbel_argsort_sample_idx(key.get(), logits)
  sel_idx = jnp.where(cluster_bias_mask[:max_seq],
                      jnp.arange(max_seq),
                      index_order[:max_seq])
  extra_idx = index_order[max_seq:]

  for k in ('msa', 'deletion_matrix', 'msa_mask', 'bert_mask', 'true_msa'):
    if k in batch:
      if k == 'true_msa':
        # extra_msa has no BERT mask
        batch['extra_msa'] = batch[k][extra_idx]
      if k not in ['msa']:
        batch['extra_' + k] = batch[k][extra_idx]
      batch[k] = batch[k][sel_idx]

  return batch


def create_extra_msa_feature(batch: FeatureDict, num_extra_msa: int
                             ) -> tuple[jnp.ndarray, jnp.ndarray]:
  """Expand extra_msa into 1hot and concat with other extra msa features.

  We do this as late as possible as the one_hot extra msa can be very large.

  Arguments:
    batch: a dictionary with the following keys:
     * 'extra_msa': [N_extra_seq, N_res] MSA that wasn't selected as a cluster
       centre. Note, that this is not one-hot encoded.
     * 'extra_has_deletion': [N_extra_seq, N_res] Whether there is a deletion to
       the left of each position in the extra MSA.
     * 'extra_deletion_value': [N_extra_seq, N_res] The number of deletions to
       the left of each position in the extra MSA.
    num_extra_msa: Number of sequences to include in the extra MSA features.

  Returns:
    Concatenated tensor of extra MSA features.
  """
  # 23 = 20 amino acids + 'X' for unknown + gap + bert mask
  deletion_matrix = batch['extra_deletion_matrix'][:num_extra_msa]
  extra_msa_1hot = jax.nn.one_hot(batch['extra_msa'][:num_extra_msa], 23)
  has_deletion = jnp.clip(deletion_matrix, 0., 1.)[..., None]
  deletion_value = (jnp.arctan(deletion_matrix / 3.) * (2. / jnp.pi))[..., None]
  extra_msa_mask = batch['extra_msa_mask'][:num_extra_msa]
  return jnp.concatenate([extra_msa_1hot, has_deletion, deletion_value],
                         axis=-1), extra_msa_mask


def clipped_sigmoid_cross_entropy(
    logits: jnp.ndarray,
    labels: jnp.ndarray,
    clip_negative_at_logit: float,
    clip_positive_at_logit: float,
    epsilon: float = 1e-07,
    ) -> jnp.ndarray:
  """Computes sigmoid xent loss with clipped input logits.

  Args:
    logits: The predicted values.
    labels: The ground truth values.
    clip_negative_at_logit: clip the loss to 0 if prediction smaller than this
      value for the negative class.
    clip_positive_at_logit: clip the loss to this value if prediction smaller
      than this value for the positive class.
    epsilon: A small increment to add to avoid taking a log of zero.

  Returns:
    Loss value.
  """
  prob = jax.nn.sigmoid(logits)
  prob = jnp.clip(prob, epsilon, 1. - epsilon)
  loss = -labels * jnp.log(
      prob) - (1. - labels) * jnp.log(1. - prob)
  loss_at_clip = math.log(math.exp(clip_negative_at_logit) + 1)
  loss = jnp.where(
      (1 - labels) * (logits < clip_negative_at_logit), loss_at_clip, loss)
  loss_at_clip = math.log(math.exp(-clip_positive_at_logit) + 1)
  loss = jnp.where(
      labels * (logits < clip_positive_at_logit), loss_at_clip, loss)
  return loss


class LogitDiffPathogenicityHead(hk.Module):
  """Variant pathogenicity classification head."""

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'logit_diff_head'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config
    self.num_output = len(residue_constants.restypes_with_x_and_gap)
    self.variant_row = 1

  def __call__(self,
               representations: FeatureDict,
               batch: FeatureDict,
               is_training: bool
               ) -> FeatureDict:
    logits = common_modules.Linear(
        self.num_output,
        initializer='linear',
        name='logits')(
            representations['msa'][self.variant_row])

    ref_score = jnp.einsum('ij, ij->i', logits, jax.nn.one_hot(
        batch['aatype'], num_classes=self.num_output))
    variant_score = jnp.einsum('ij, ij->i', logits, jax.nn.one_hot(
        batch['variant_aatype'], num_classes=self.num_output))
    logit_diff = ref_score - variant_score
    variant_pathogenicity = jnp.sum(logit_diff * batch['variant_mask'])
    return {'variant_row_logit_diff': logit_diff,
            'variant_pathogenicity': variant_pathogenicity}

  def loss(self, value: FeatureDict, batch: FeatureDict) -> jnp.ndarray:
    loss = clipped_sigmoid_cross_entropy(logits=value['variant_row_logit_diff'],
                                         labels=batch['pathogenicity'],
                                         clip_negative_at_logit=0.0,
                                         clip_positive_at_logit=-1.0)
    loss = (jnp.sum(loss * batch['variant_mask'], axis=(-2, -1)) /
            (1e-8 + jnp.sum(batch['variant_mask'], axis=(-2, -1))))
    return loss


class AlphaMissense(hk.Module):
  """AlphaMissense model with recycling.

  Changes relative to the original AlphaFold 2 implementation are described in
  the Methods section of Cheng et al. (2023). These changes include: (i) how
  the MSA and "pair" embeddings are initialized; (ii) how the "single"
  embedding is computed; (iii) the MSA embedding is not recycled. Therefore,
  checkpoints produced with AlphaFold 2 or AlphaFold Multimer are not
  compatible with AlphaMissense.
  """

  def __init__(self,
               config: ml_collections.ConfigDict,
               name: str = 'alphamissense'):
    super().__init__(name=name)
    self.config = config
    self.global_config = config.global_config

  def __call__(
      self,
      batch: FeatureDict,
      *,
      is_training: bool,
      return_representations: bool,
      ) -> OutputDict:

    impl = AlphaMissenseIteration(self.config, self.global_config)
    num_residues = batch['aatype'].shape[0]
    safe_key = prng.SafeKey(hk.next_rng_key())

    def get_prev(ret):
      new_prev = {
          'prev_pos':
              ret['structure_module']['final_atom_positions'],
          'prev_pair': ret['representations']['pair'],
      }
      return jax.tree_map(jax.lax.stop_gradient, new_prev)

    def apply_network(prev, safe_key):
      recycled_batch = {**batch, **prev}
      return impl(
          batch=recycled_batch,
          is_training=is_training,
          safe_key=safe_key)

    prev = {}
    emb_config = self.config.embeddings_and_evoformer
    if emb_config.recycle_pos:
      prev['prev_pos'] = jnp.zeros(
          [num_residues, residue_constants.atom_type_num, 3])
    if emb_config.recycle_features:
      prev['prev_pair'] = jnp.zeros(
          [num_residues, num_residues, emb_config.pair_channel])

    if self.config.num_recycle:
      if 'num_iter_recycling' in batch:
        # Training time: num_iter_recycling is in batch.
        # Value for each ensemble batch is the same, so arbitrarily taking 0-th.
        num_iter = batch['num_iter_recycling'][0]

        # Add insurance that even when ensembling, we will not run more
        # recyclings than the model is configured to run.
        num_iter = jnp.minimum(num_iter, self.config.num_recycle)
      else:
        # Eval mode or tests: use the maximum number of iterations.
        num_iter = self.config.num_recycle

      def recycle_body(x):
        i, prev, safe_key = x
        safe_key1, safe_key2 = safe_key.split()
        ret = apply_network(prev=prev, safe_key=safe_key2)
        return i + 1, get_prev(ret), safe_key1

      if hk.running_init():
        num_recycles, prev, safe_key = recycle_body((0, prev, safe_key))
      else:
        num_recycles, prev, safe_key = hk.while_loop(
            lambda x: x[0] < num_iter,
            recycle_body,
            (0, prev, safe_key))
    else:
      # No recycling.
      num_recycles = 0

    # Run extra iteration.
    ret = apply_network(prev=prev, safe_key=safe_key)

    if not return_representations:
      del ret['representations']
    ret['num_recycles'] = num_recycles

    return ret


class AlphaMissenseIteration(hk.Module):
  """A single recycling iteration of AlphaMissense architecture."""

  def __init__(self,
               config: ml_collections.ConfigDict,
               global_config: ml_collections.ConfigDict,
               name: str = 'alphafold_iteration'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def __call__(self,
               batch: FeatureDict,
               is_training: bool,
               safe_key: prng.SafeKey,
               ) -> OutputDict:

    # Compute representations for each batch element and average.
    evoformer_module = EmbeddingsAndEvoformer(
        self.config.embeddings_and_evoformer, self.global_config)
    representations = evoformer_module(batch, is_training, safe_key)

    self.representations = representations
    self.batch = batch
    self.heads = {}
    for head_name, head_config in sorted(self.config.heads.items()):
      if not head_config.weight:
        continue  # Do not instantiate zero-weight heads.

      head_factory = {
          'masked_msa': modules.MaskedMsaHead,
          'distogram': modules.DistogramHead,
          'structure_module': folding_multimer.StructureModule,
          'logit_diff': LogitDiffPathogenicityHead,
      }[head_name]
      self.heads[head_name] = (head_config,
                               head_factory(head_config, self.global_config))
    ret = {}
    ret['representations'] = representations
    for name, (_, module) in self.heads.items():
      ret[name] = module(representations, batch, is_training)
    return ret


class EmbeddingsAndEvoformer(hk.Module):
  """Embeds the input data and runs Evoformer."""

  def __init__(self, config, global_config, name='evoformer'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config

  def _relative_encoding(self, batch: FeatureDict) -> jnp.ndarray:
    """Add relative Position encodings.

    Feed the Pair stack a relative feature. For position (i, j), its value
    is (i-j) clipped to [-k, k] and one-hotted.

    When not using 'use_chain_relative' the residue indices are used as is, e.g.
    for heteromers relative positions will be computed using the positions in
    the corresponding chains.

    When using 'use_chain_relative' we add an extra bin that denotes
    'different chain'. Furthermore we also provide the relative chain index
    (i.e. sym_id) clipped and one-hotted to the network. And an extra feature
    which denotes whether they belong to the same chain type, i.e. it's 0 if
    they are in different heteromer chains and 1 otherwise.

    Args:
      batch: batch.
    Returns:
      Feature embedding using the features as described before.
    """
    c = self.config
    rel_feats = []
    pos = batch['residue_index']
    offset = pos[:, None] - pos[None, :]

    clipped_offset = jnp.clip(
        offset + c.max_relative_idx, a_min=0, a_max=2 * c.max_relative_idx)

    if c.use_chain_relative:
      asym_id = batch['seq_mask'] * batch.get(
          'asym_id', jnp.ones_like(batch['residue_index']))
      asym_id_same = jnp.equal(asym_id[:, None], asym_id[None, :])

      final_offset = jnp.where(asym_id_same, clipped_offset,
                               (2 * c.max_relative_idx + 1) *
                               jnp.ones_like(clipped_offset))

      rel_pos = jax.nn.one_hot(final_offset, 2 * c.max_relative_idx + 2)

      rel_feats.append(rel_pos)

      entity_id = batch['seq_mask'] * batch.get(
          'entity_id', jnp.ones_like(batch['residue_index']))
      entity_id_same = jnp.equal(entity_id[:, None], entity_id[None, :])
      rel_feats.append(entity_id_same.astype(rel_pos.dtype)[..., None])

      sym_id = batch['seq_mask'] * batch.get(
          'sym_id', jnp.ones_like(batch['residue_index']))
      rel_sym_id = sym_id[:, None] - sym_id[None, :]

      max_rel_chain = c.max_relative_chain

      clipped_rel_chain = jnp.clip(
          rel_sym_id + max_rel_chain, a_min=0, a_max=2 * max_rel_chain)

      final_rel_chain = jnp.where(entity_id_same, clipped_rel_chain,
                                  (2 * max_rel_chain + 1) *
                                  jnp.ones_like(clipped_rel_chain))
      rel_chain = jax.nn.one_hot(final_rel_chain, 2 * c.max_relative_chain + 2)

      rel_feats.append(rel_chain)

    else:
      rel_pos = jax.nn.one_hot(clipped_offset, 2 * c.max_relative_idx + 1)
      rel_feats.append(rel_pos)

    rel_feat = jnp.concatenate(rel_feats, axis=-1)
    return common_modules.Linear(
        c.pair_channel, name='position_activations')(
            rel_feat.astype(jnp.float32))

  def __call__(self,
               batch: FeatureDict,
               is_training: bool,
               safe_key: prng.SafeKey
               ) -> FeatureDict:

    c = self.config
    gc = self.global_config

    if safe_key is None:
      safe_key = prng.SafeKey(hk.next_rng_key())

    num_residues = batch['aatype'].shape[0]
    target_feat = jax.nn.one_hot(batch['aatype'], 21)
    preprocess_1d = jnp.zeros((num_residues, c.msa_channel))
    left_single = common_modules.Linear(
        c.pair_channel, name='left_single')(target_feat)
    right_single = common_modules.Linear(
        c.pair_channel, name='right_single')(target_feat)

    profile_feat = jnp.concatenate([
        batch['msa_profile'], batch['deletion_mean'][..., None]], axis=-1)
    preprocess_1d += common_modules.Linear(
        c.msa_channel, name='profile_preprocess_1d',
        use_bias=False, initializer='zeros')(profile_feat)
    left_single += common_modules.Linear(
        c.pair_channel, name='profile_left_single',
        use_bias=False, initializer='zeros')(profile_feat)
    right_single += common_modules.Linear(
        c.pair_channel, name='profile_right_single',
        use_bias=False, initializer='zeros')(profile_feat)

    safe_key, sample_key = safe_key.split()
    batch = sample_msa(sample_key, batch, c.num_msa)
    preprocess_msa = common_modules.Linear(
        c.msa_channel, name='preprocess_msa')(jax.nn.one_hot(batch['msa'], 23))
    msa_activations = jnp.expand_dims(preprocess_1d, axis=0) + preprocess_msa
    pair_activations = left_single[:, None] + right_single[None]
    mask_2d = batch['seq_mask'][:, None] * batch['seq_mask'][None, :]

    if 'prev_pos' in batch:
      prev_pseudo_beta = modules.pseudo_beta_fn(batch['aatype'],
                                                batch['prev_pos'],
                                                None)
      dgram = modules.dgram_from_positions(prev_pseudo_beta,
                                           **self.config.prev_pos)
      pair_activations += common_modules.Linear(c.pair_channel,
                                                name='prev_pos_linear')(
                                                    dgram)

    if 'prev_pair' in batch:
      pair_activations += common_modules.LayerNorm(axis=[-1],
                                                   create_scale=True,
                                                   create_offset=True,
                                                   name='prev_pair_norm')(
                                                       batch['prev_pair'])

    pair_activations += self._relative_encoding(batch)

    if c.enable_extra_msa_stack:
      extra_msa_feat, extra_msa_mask = create_extra_msa_feature(batch,
                                                                c.num_extra_msa)
      extra_msa_activations = common_modules.Linear(
          c.extra_msa_channel,
          name='extra_msa_activations')(
              extra_msa_feat)

      # Extra MSA Stack.
      # Jumper et al. (2021) Suppl. Alg. 18 "ExtraMsaStack"
      extra_msa_stack_input = dict(msa=extra_msa_activations,
                                   pair=pair_activations)
      extra_msa_stack_iteration = modules.EvoformerIteration(
          c.evoformer, gc, is_extra_msa=True, name='extra_msa_stack')

      def extra_msa_stack_fn(x):
        act, safe_key = x
        safe_key, safe_subkey = safe_key.split()
        extra_evoformer_output = extra_msa_stack_iteration(
            activations=act,
            masks=dict(msa=extra_msa_mask, pair=mask_2d),
            is_training=is_training,
            safe_key=safe_subkey)
        return (extra_evoformer_output, safe_key)

      if gc.use_remat:
        extra_msa_stack_fn = hk.remat(extra_msa_stack_fn)

      extra_msa_stack = layer_stack.layer_stack(
          c.extra_msa_stack_num_block)(
              extra_msa_stack_fn)
      extra_msa_output, safe_key = extra_msa_stack(
          (extra_msa_stack_input, safe_key))
      pair_activations = extra_msa_output['pair']

    evoformer_input = dict(msa=msa_activations, pair=pair_activations)
    evoformer_masks = dict(msa=batch['msa_mask'], pair=mask_2d)

    evoformer_iteration = modules.EvoformerIteration(
        c.evoformer, gc, is_extra_msa=False, name='evoformer_iteration')

    def evoformer_fn(x):
      act, safe_key = x
      safe_key, safe_subkey = safe_key.split()
      evoformer_output = evoformer_iteration(
          activations=act,
          masks=evoformer_masks,
          is_training=is_training,
          safe_key=safe_subkey)
      return (evoformer_output, safe_key)

    if gc.use_remat:
      evoformer_fn = hk.remat(evoformer_fn)

    evoformer_stack = layer_stack.layer_stack(c.evoformer_num_block)(
        evoformer_fn)
    evoformer_output, _ = evoformer_stack((evoformer_input, safe_key))

    msa_activations = evoformer_output['msa']
    pair_activations = evoformer_output['pair']
    diag_act = jnp.diagonal(pair_activations, axis1=0, axis2=1).T
    single_activations = common_modules.Linear(
        c.seq_channel, name='single_act_from_pair_diag')(diag_act)

    return dict(single=single_activations,
                pair=pair_activations,
                msa=msa_activations)


class RunModel:
  """Container for JAX model."""

  def __init__(self,
               config: ml_collections.ConfigDict,
               *,
               params: Optional[Mapping[str, Mapping[str, jax.Array]]]):
    self.config = config
    self.params = params
    def _forward_fn(batch):
      model = AlphaMissense(self.config.model)
      return model(batch, is_training=False, return_representations=False)

    self.apply = jax.jit(hk.transform(_forward_fn).apply)
    self.init = hk.transform(_forward_fn).init

  def init_params(self, feat: FeatureDict, random_seed: int = 0):
    """Initializes the model parameters.

    If none were provided when this class was instantiated then the parameters
    are randomly initialized.

    Args:
      feat: A dictionary of NumPy feature arrays as output by
        pipeline_missense.DataPipeline.process.
      random_seed: A random seed to use to initialize the parameters if none
        were set when this class was initialized.
    """
    if not self.params:
      # Init params randomly.
      rng = jax.random.PRNGKey(random_seed)
      self.params = hk.data_structures.to_mutable_dict(
          self.init(rng, feat))
      logging.warning('Initialized parameters randomly')

  def eval_shape(self, feat: FeatureDict) -> jax.ShapeDtypeStruct:
    self.init_params(feat)
    logging.info('Running eval_shape with shape(feat) = %s',
                 tree.map_structure(lambda x: x.shape, feat))
    shape = jax.eval_shape(self.apply, self.params, jax.random.PRNGKey(0), feat)
    logging.info('Output shape was %s', shape)
    return shape

  def predict(self,
              feat: FeatureDict,
              random_seed: int,
              ) -> OutputDict:
    """Makes a prediction by inferencing the model on the provided features.

    Args:
      feat: A dictionary of NumPy feature arrays as output by
        pipeline_missense.DataPipeline.process.
      random_seed: The random seed to use when running the model. In the
        multimer model this controls the MSA sampling.

    Returns:
      A dictionary of model outputs.
    """
    self.init_params(feat)
    logging.info('Running predict with shape(feat) = %s',
                 tree.map_structure(lambda x: x.shape, feat))
    result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)

    # This block is to ensure benchmark timings are accurate. Some blocking is
    # already happening when computing get_confidence_metrics, and this ensures
    # all outputs are blocked on.
    jax.tree_map(lambda x: x.block_until_ready(), result)
    return result

alphamissense/model/modules_missense.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for prng."""

from absl.testing import absltest
import jax
from alphamissense.model import prng


class PrngTest(absltest.TestCase):

  def test_key_reuse(self):

    init_key = jax.random.PRNGKey(42)
    safe_key = prng.SafeKey(init_key)
    _, safe_key = safe_key.split()

    raw_key = safe_key.get()

    self.assertFalse((raw_key == init_key).all())

    with self.assertRaises(RuntimeError):
      safe_key.get()

    with self.assertRaises(RuntimeError):
      safe_key.split()

    with self.assertRaises(RuntimeError):
      safe_key.duplicate()


if __name__ == '__main__':
  absltest.main()

alphamissense/model/prng_test.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A collection of JAX utility functions for use in protein folding."""

import collections
import contextlib
import functools
import numbers
from typing import Mapping

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np


def stable_softmax(logits: jax.Array) -> jax.Array:
  """Numerically stable softmax for (potential) bfloat 16."""
  if logits.dtype == jnp.float32:
    return jax.nn.softmax(logits)
  elif logits.dtype == jnp.bfloat16:
    # Need to explicitly do softmax in float32 to avoid numerical issues
    # with large negatives. Large negatives can occur if trying to mask
    # by adding on large negative logits so that things softmax to zero.
    return jax.nn.softmax(logits.astype(jnp.float32)).astype(jnp.bfloat16)
  else:
    raise ValueError(f'Unexpected input dtype {logits.dtype}')


def bfloat16_creator(next_creator, shape, dtype, init, context):
  """Creates float32 variables when bfloat16 is requested."""
  if context.original_dtype == jnp.bfloat16:
    dtype = jnp.float32
  return next_creator(shape, dtype, init)


def bfloat16_getter(next_getter, value, context):
  """Casts float32 to bfloat16 when bfloat16 was originally requested."""
  if context.original_dtype == jnp.bfloat16:
    assert value.dtype == jnp.float32
    value = value.astype(jnp.bfloat16)
  return next_getter(value)


@contextlib.contextmanager
def bfloat16_context():
  with hk.custom_creator(bfloat16_creator), hk.custom_getter(bfloat16_getter):
    yield


def final_init(config):
  """Converts the weight initialization config to the Haiku module argument."""
  if config.zero_init:
    return 'zeros'
  else:
    return 'linear'


def batched_gather(params, indices, axis=0, batch_dims=0):
  """Implements a JAX equivalent of `tf.gather` with `axis` and `batch_dims`."""
  take_fn = lambda p, i: jnp.take(p, i, axis=axis, mode='clip')
  for _ in range(batch_dims):
    take_fn = jax.vmap(take_fn)
  return take_fn(params, indices)


def mask_mean(mask, value, axis=None, drop_mask_channel=False, eps=1e-10):
  """Masked mean."""
  if drop_mask_channel:
    mask = mask[..., 0]

  mask_shape = mask.shape
  value_shape = value.shape

  assert len(mask_shape) == len(value_shape)

  if isinstance(axis, numbers.Integral):
    axis = [axis]
  elif axis is None:
    axis = list(range(len(mask_shape)))
  assert isinstance(axis, collections.abc.Iterable), (
      'axis needs to be either an iterable, integer or "None"')

  broadcast_factor = 1.
  for axis_ in axis:
    value_size = value_shape[axis_]
    mask_size = mask_shape[axis_]
    if mask_size == 1:
      broadcast_factor *= value_size
    else:
      assert mask_size == value_size

  return (jnp.sum(mask * value, axis=axis) /
          (jnp.sum(mask, axis=axis) * broadcast_factor + eps))


def flat_params_to_haiku(params: Mapping[str, np.ndarray]) -> hk.Params:
  """Convert a dictionary of NumPy arrays to Haiku parameters."""
  hk_params = {}
  for path, array in params.items():
    scope, name = path.split('//')
    if scope not in hk_params:
      hk_params[scope] = {}
    hk_params[scope][name] = jnp.array(array)

  return hk_params


def padding_consistent_rng(f):
  """Modify any element-wise random function to be consistent with padding.

  Normally if you take a function like jax.random.normal and generate an array,
  say of size (10,10), you will get a different set of random numbers to if you
  add padding and take the first (10,10) sub-array.

  This function makes a random function that is consistent regardless of the
  amount of padding added.

  Note: The padding-consistent function is likely to be slower to compile and
  run than the function it is wrapping, but these slowdowns are likely to be
  negligible in a large network.

  Args:
    f: Any element-wise function that takes (PRNG key, shape) as the first 2
      arguments.

  Returns:
    An equivalent function to f, that is now consistent for different amounts of
    padding.
  """
  def grid_keys(key, shape):
    """Generate a grid of rng keys that is consistent with different padding.

    Generate random keys such that the keys will be identical, regardless of
    how much padding is added to any dimension.

    Args:
      key: A PRNG key.
      shape: The shape of the output array of keys that will be generated.

    Returns:
      An array of shape `shape` consisting of random keys.
    """
    if not shape:
      return key
    new_keys = jax.vmap(functools.partial(jax.random.fold_in, key))(
        jnp.arange(shape[0]))
    return jax.vmap(functools.partial(grid_keys, shape=shape[1:]))(new_keys)

  def inner(key, shape, **kwargs):
    keys = grid_keys(key, shape)
    signature = (
        '()->()'
        if jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key)
        else '(2)->()'
    )
    return jnp.vectorize(
        functools.partial(f, shape=(), **kwargs), signature=signature
    )(keys)

  return inner

alphamissense/model/utils.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Single source of truth for the AlphaMissense version."""

__version__ = '0.0.1'

alphamissense/version.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Install script for setuptools."""

from setuptools import find_packages
from setuptools import setup
from alphamissense import version

setup(
    name='alphamissense',
    version=version.__version__,
    description=(
        'An implementation of the inference pipeline of AlphaMissense.'
    ),
    author='DeepMind',
    author_email='alphamissense@deepmind.com',
    license='Apache License, Version 2.0',
    packages=find_packages(),
    install_requires=[
        'absl-py',
        'biopython',
        'dm-haiku',
        'dm-tree',
        'immutabledict',
        'jax',
        'jaxlib',
        'ml-collections',
        'numpy',
        'scipy',
        'typing_extensions',
    ],
    tests_require=[
        'mock',
    ],
    classifiers=[
        'Development Status :: 5 - Production/Stable',
        'Intended Audience :: Science/Research',
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: POSIX :: Linux',
        'Programming Language :: Python :: 3.10',
        'Topic :: Scientific/Engineering :: Artificial Intelligence',
    ],
)

setup.py
google-deepmind/alphamissense
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Script to test installation of AlphaMissense."""

import pathlib

import jax
import jax.numpy as jnp
import tree

from alphamissense.data import pipeline_missense
from alphamissense.model import config
from alphamissense.model import modules_missense

CWD = str(pathlib.Path(__file__).parent)


def array_tree_summary(array_tree):
  """Creates flatten list of keys with shapes and dtype."""
  array_strs = []
  for path, array in tree.flatten_with_path(array_tree):
    path = '.'.join(map(str, path))
    if array is None:
      array_strs.append(f'{path}(None)')
    else:
      array_strs.append(f'{path}(shape={array.shape}, dtype={array.dtype})')
  return array_strs


def main():

  pipeline = pipeline_missense.DataPipeline(
      jackhmmer_binary_path='/usr/bin/jackhmmer',
      protein_sequence_file=CWD + '/test_sequence.fasta',
      uniref90_database_path=CWD + '/test_sequence.fasta',
      mgnify_database_path='',
      small_bfd_database_path='',
  )

  sample = pipeline.process(
      protein_id='Q08708',  # Sequence identifier in the FASTA file.
      reference_aa='A',
      alternate_aa='C',
      position=3,  # Note that the position is 1-based!
      msa_output_dir=CWD + '/msa_output_dir',
  )
  print('Pipeline produced sample with features:')
  for sample_str in sorted(array_tree_summary(sample)):
    print(sample_str)

  cfg = config.model_config()
  cfg.model.num_recycle = 1  # Speeds up testing.
  model_runner = modules_missense.RunModel(cfg, params=None)
  sample = jax.tree_map(jnp.asarray, sample)
  output = model_runner.predict(sample, random_seed=0)
  print('Model inizialized parameters:')
  for param_str in sorted(array_tree_summary(model_runner.params)):
    print(param_str)

  print('Model produced outputs:')
  for output_str in sorted(array_tree_summary(output)):
    print(output_str)


if __name__ == '__main__':
  main()

test/test_installation.py
google-deepmind/alphamissense