test_structure_module.py

# Copyright 2021 AlQuraishi Laboratory
#
# 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.

import torch
import numpy as np
import unittest

from openfold.np.residue_constants import (
    restype_rigid_group_default_frame,
    restype_atom14_to_rigid_group,
    restype_atom14_mask,
    restype_atom14_rigid_group_positions,
    restype_atom37_mask,
)    
from openfold.model.structure_module import (
    StructureModule,
    StructureModuleTransition,
    BackboneUpdate,
    AngleResnet,
    InvariantPointAttention,
)
from openfold.utils.affine_utils import T
import openfold.utils.feats as feats
import tests.compare_utils as compare_utils
from tests.config import consts
from tests.data_utils import (
    random_affines_4x4,
)

if(compare_utils.alphafold_is_installed()):
    alphafold = compare_utils.import_alphafold()
    import jax
    import haiku as hk


class TestStructureModule(unittest.TestCase):
    def test_structure_module_shape(self):
        batch_size = consts.batch_size
        n = consts.n_res
        c_s = consts.c_s
        c_z = consts.c_z
        c_ipa = 13
        c_resnet = 17
        no_heads_ipa = 6
        no_query_points = 4
        no_value_points = 4
        dropout_rate = 0.1
        no_layers = 3
        no_transition_layers = 3
        no_resnet_layers = 3
        ar_epsilon = 1e-6
        no_angles = 7
        trans_scale_factor = 10
        inf = 1e5

        sm = StructureModule(
            c_s,
            c_z,
            c_ipa,
            c_resnet,
            no_heads_ipa,
            no_query_points,
            no_value_points,
            dropout_rate,
            no_layers,
            no_transition_layers,
            no_resnet_layers,
            no_angles,
            trans_scale_factor,
            ar_epsilon,
            inf,
        )

        s = torch.rand((batch_size, n, c_s))
        z = torch.rand((batch_size, n, n, c_z))
        f = torch.randint(low=0, high=21, size=(batch_size, n)).long()

        out = sm(s, z, f)

        self.assertTrue(
            out["frames"].shape == (no_layers, batch_size, n, 4, 4)
        )
        self.assertTrue(
            out["angles"].shape == (no_layers, batch_size, n, no_angles, 2)
        )
        self.assertTrue(
            out["positions"].shape == (no_layers, batch_size, n, 14, 3)
        )

    def test_structure_module_transition_shape(self):
        batch_size = 2
        n = 5
        c = 7
        num_layers = 3
        dropout = 0.1

        smt = StructureModuleTransition(c, num_layers, dropout)

        s = torch.rand((batch_size, n, c))

        shape_before = s.shape
        s = smt(s)
        shape_after = s.shape

        self.assertTrue(shape_before == shape_after)

    @compare_utils.skip_unless_alphafold_installed()
    def test_structure_module_compare(self):
        config = compare_utils.get_alphafold_config()
        c_sm = config.model.heads.structure_module
        c_global = config.model.global_config
        def run_sm(representations, batch):
            sm = alphafold.model.folding.StructureModule(c_sm, c_global)
            representations = {
                k:jax.lax.stop_gradient(v) for k,v in representations.items()
            }
            batch = {
                k:jax.lax.stop_gradient(v) for k,v in batch.items()
            }
            return sm(representations, batch, is_training=False)
    
        f = hk.transform(run_sm)
    
        n_res = 200
    
        representations = {
            'single': np.random.rand(n_res, consts.c_s).astype(np.float32),
            'pair': 
                np.random.rand(n_res, n_res, consts.c_z).astype(np.float32),
        }
    
        batch = {
            'seq_mask': np.random.randint(0, 2, (n_res,)).astype(np.float32),
            'aatype': np.random.randint(0, 21, (n_res,)),
        }
    
        batch['atom14_atom_exists'] = np.take(
            restype_atom14_mask, 
            batch['aatype'],
            axis=0
        )
    
        batch['atom37_atom_exists'] = np.take(
            restype_atom37_mask, 
            batch['aatype'],
            axis=0
        )
    
        batch.update(feats.compute_residx_np(batch))
    
        params = compare_utils.fetch_alphafold_module_weights(
            "alphafold/alphafold_iteration/structure_module"
        )
        
        key = jax.random.PRNGKey(42)
        out_gt = f.apply(
            params, key, representations, batch
        )
        out_gt = torch.as_tensor(
            np.array(out_gt["final_atom14_positions"].block_until_ready())
        ) 
   
        model = compare_utils.get_global_pretrained_openfold()
        out_repro = model.structure_module(
            torch.as_tensor(representations["single"]).cuda(), 
            torch.as_tensor(representations["pair"]).cuda(), 
            torch.as_tensor(batch["aatype"]).cuda(), 
            mask=torch.as_tensor(batch["seq_mask"]).cuda(),
        )
        out_repro = out_repro["positions"][-1].cpu()
      
        # The structure module, thanks to angle normalization, is very volatile
        # We only assess the mean here. Heuristically speaking, it seems to 
        # have lower error in general on real rather than synthetic data.
        self.assertTrue(torch.mean(torch.abs(out_gt - out_repro)) < 0.01)


class TestBackboneUpdate(unittest.TestCase):
    def test_shape(self): 
        batch_size = 2
        n_res = 3
        c_in = 5
        
        bu = BackboneUpdate(c_in)

        s = torch.rand((batch_size, n_res, c_in))

        t = bu(s)
        rot, tra = t.rots, t.trans

        self.assertTrue(rot.shape == (batch_size, n_res, 3, 3))
        self.assertTrue(tra.shape == (batch_size, n_res, 3))


class TestInvariantPointAttention(unittest.TestCase):
    def test_shape(self):
        c_m = 13
        c_z = 17
        c_hidden = 19
        no_heads = 5
        no_qp = 7
        no_vp = 11

        batch_size = 2
        n_res = 23

        s = torch.rand((batch_size, n_res, c_m))
        z = torch.rand((batch_size, n_res, n_res, c_z))
        mask = torch.ones((batch_size, n_res))

        rots = torch.rand((batch_size, n_res, 3, 3))
        trans = torch.rand((batch_size, n_res, 3))

        t = T(rots, trans)

        ipa = InvariantPointAttention(
            c_m, c_z, c_hidden, no_heads, no_qp, no_vp
        )

        shape_before = s.shape
        s = ipa(s, z, t, mask)

        self.assertTrue(s.shape == shape_before)

    @compare_utils.skip_unless_alphafold_installed()
    def test_ipa_compare(self):
        def run_ipa(act, static_feat_2d, mask, affine):
            config = compare_utils.get_alphafold_config() 
            ipa = alphafold.model.folding.InvariantPointAttention(
                config.model.heads.structure_module, 
                config.model.global_config, 
            )
            attn = ipa(
                inputs_1d=act, 
                inputs_2d=static_feat_2d, 
                mask=mask, 
                affine=affine
            )
            return attn
        
        f = hk.transform(run_ipa)
        
        n_res = consts.n_res
        c_s = consts.c_s
        c_z = consts.c_z
        
        sample_act = np.random.rand(n_res, c_s)
        sample_2d = np.random.rand(n_res, n_res, c_z)
        sample_mask = np.ones((n_res, 1))

        affines = random_affines_4x4((n_res,))
        rigids = alphafold.model.r3.rigids_from_tensor4x4(affines)
        quats = alphafold.model.r3.rigids_to_quataffine(rigids)
        transformations = T.from_4x4(
            torch.as_tensor(affines).float().cuda()
        )
        
        sample_affine = quats
        
        ipa_params = compare_utils.fetch_alphafold_module_weights(
            "alphafold/alphafold_iteration/structure_module/" +
            "fold_iteration/invariant_point_attention"
        )

        out_gt = f.apply(
            ipa_params, None, sample_act, sample_2d, sample_mask, sample_affine
        ).block_until_ready()
        out_gt = torch.as_tensor(np.array(out_gt))

        with torch.no_grad():
            model = compare_utils.get_global_pretrained_openfold()
            out_repro = model.structure_module.ipa(
                torch.as_tensor(sample_act).float().cuda(), 
                torch.as_tensor(sample_2d).float().cuda(), 
                transformations, 
                torch.as_tensor(sample_mask.squeeze(-1)).float().cuda(),
            ).cpu()
   
        self.assertTrue(torch.max(torch.abs(out_gt - out_repro)) < consts.eps)


class TestAngleResnet(unittest.TestCase):
    def test_shape(self): 
        batch_size = 2
        n = 3
        c_s = 13
        c_hidden = 11
        no_layers = 5
        no_angles = 7
        epsilon = 1e-12
        
        ar = AngleResnet(c_s, c_hidden, no_layers, no_angles, epsilon)
        a = torch.rand((batch_size, n, c_s))
        a_initial = torch.rand((batch_size, n, c_s))

        _, a = ar(a, a_initial)

        self.assertTrue(a.shape == (batch_size, n, no_angles, 2))


if __name__ == "__main__":
    unittest.main()