diff --git a/openfold/utils/feats.py b/openfold/utils/feats.py index 75845d0fe..2213ea0a4 100644 --- a/openfold/utils/feats.py +++ b/openfold/utils/feats.py @@ -258,8 +258,7 @@ def frames_and_literature_positions_to_atom14_pos( atom_mask, lit_positions, ): - # [*, N, 14, 4, 4] - default_4x4 = default_frames[aatype, ...] + num_rigid_groups = default_frames.shape[-3] # [*, N, 14] group_mask = group_idx[aatype, ...] @@ -267,7 +266,7 @@ def frames_and_literature_positions_to_atom14_pos( # [*, N, 14, 8] group_mask = nn.functional.one_hot( group_mask, - num_classes=default_frames.shape[-3], + num_classes=num_rigid_groups, ) # [*, N, 14, 8] diff --git a/tests/test_feats.py b/tests/test_feats.py index 7a1783b85..d28fa43e3 100644 --- a/tests/test_feats.py +++ b/tests/test_feats.py @@ -341,6 +341,39 @@ def test_frames_and_literature_positions_to_atom14_pos_shape(self): self.assertTrue(xyz.shape == (batch_size, n_res, 14, 3)) + def test_atom14_pos_does_not_gather_default_frames(self): + class DefaultFrameShapeOnly: + shape = torch.Size((21, 8, 4, 4)) + + def __getitem__(self, key): + raise AssertionError("default_frames should not be gathered") + + batch_size = 1 + n_res = 3 + + rots = torch.eye(3).expand(batch_size, n_res, 8, 3, 3) + trans = torch.zeros((batch_size, n_res, 8, 3)) + + if consts.is_multimer: + rotation = Rot3Array.from_array(rots) + translation = Vec3Array.from_array(trans) + ts = Rigid3Array(rotation, translation) + else: + ts = Rigid(Rotation(rot_mats=rots), trans) + + aatype = torch.randint(low=0, high=21, size=(batch_size, n_res)).long() + + xyz = feats.frames_and_literature_positions_to_atom14_pos( + ts, + aatype, + DefaultFrameShapeOnly(), + torch.tensor(restype_atom14_to_rigid_group), + torch.tensor(restype_atom14_mask), + torch.tensor(restype_atom14_rigid_group_positions), + ) + + self.assertTrue(xyz.shape == (batch_size, n_res, 14, 3)) + @compare_utils.skip_unless_alphafold_installed() def test_frames_and_literature_positions_to_atom14_pos_compare(self): def run_f(aatype, affines):