Merge branch 'main' of https://github.com/microsoft/protein-sequence-models

sequence_models/pdb_utils.py

@@ -1,3 +1,4 @@
+import gzip
 import numpy as np
 import scipy
 from scipy.spatial.distance import squareform, pdist
@@ -40,7 +41,8 @@ def parse_PDB(x, atoms=["N", "CA", "C"], chain=None):
     output: (length, atoms, coords=(x,y,z)), sequence
     """
     xyz, seq, min_resn, max_resn = {}, {}, np.inf, -np.inf
-    for line in open(x, "rb"):
+    open_func = gzip.open if x.endswith('.gz') else open
+    for line in open_func(x, "rb"):
         line = line.decode("utf-8", "ignore").rstrip()
 
         if line[:6] == "HETATM" and line[17 : 17 + 3] == "MSE":

tests/pdb_utils_test.py

@@ -0,0 +1,20 @@
+import os, sys
+import numpy as np
+
+from sequence_models import pdb_utils
+
+ex_dir = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "examples"
+)
+assert os.path.isdir(ex_dir)
+
+orig_coords, orig_atoms, orig_valid = pdb_utils.parse_PDB(
+    os.path.join(ex_dir, "gb1_a60fb_unrelaxed_rank_1_model_5.pdb")
+)
+gz_coords, gz_atoms, gz_valid = pdb_utils.parse_PDB(
+    os.path.join(ex_dir, "gb1_a60fb_unrelaxed_rank_1_model_5.pdb.gz")
+)
+
+assert np.all(np.isclose(orig_coords, gz_coords))
+assert orig_atoms == gz_atoms
+assert np.all(orig_valid == gz_valid)