Added tests for the data_transforms

tests/test_data_transforms.py

+import copy
+import gzip
+
+import os
+
+import pickle
+
+import numpy
+import torch
+import unittest
+
+from data.data_transforms import make_seq_mask
+from openfold.config import model_config
+
+
+class TestDataTransforms(unittest.TestCase):
+    def test_make_seq_mask(self):
+        seq = torch.tensor([range(20)], dtype=torch.int64).transpose(0,1)
+        seq_one_hot = torch.FloatTensor(seq.shape[0], 20).zero_()
+        seq_one_hot.scatter_(1, seq, 1)
+        protein_aatype = torch.tensor(seq_one_hot)
+        protein = {'aatype': protein_aatype}
+        protein = make_seq_mask(protein)
+        print(protein)
+        assert 'seq_mask' in protein
+        assert protein['seq_mask'].shape == torch.Size((seq.shape[0], 20))
+
+
+if __name__ == '__main__':
+    unittest.main()
+