support reflection files in cif

SFC_Torch/Fmodel.py

@@ -57,6 +57,7 @@ def __init__(
 
         mtzdata: str | rs.Dataset 
             path to the mtz data or instance of rs.Dataset, will use the HKL list in the mtz instead, override dmin with an inference
+            also support reflection data in cif file
     
         n_bins: str, default 10
             Number of resolution bins used in the reciprocal space
@@ -320,11 +321,16 @@ def init_mtz(self, mtzdata, N_bins, expcolumns, set_experiment, freeflag, testse
         """
 
         if type(mtzdata) == str:
-            mtz_reference = rs.read_mtz(mtzdata)
+            if mtzdata.endswith(".mtz"):
+                mtz_reference = rs.read_mtz(mtzdata)
+            elif mtzdata.endswith(".cif"):
+                mtz_reference = rs.read_cif(mtzdata)
+            else:
+                raise ValueError("mtzdata should be rs.Dataset object or path str to a mtz/cif file!")
         elif type(mtzdata) == rs.DataSet:
             mtz_reference = mtzdata
         else:
-            raise TypeError("mtzdata should be rs.Dataset object or path str to a mtz file!")
+            raise TypeError("mtzdata should be rs.Dataset object or path str to a mtz/cif file!")
 
         try:
             mtz_reference.dropna(axis=0, subset=expcolumns, inplace=True)

SFC_Torch/io.py

@@ -409,8 +409,9 @@ def fetch_pdb(idlist, outpath):
         sequence_path = os.path.join(outpath, 'sequences/')
         model_path = os.path.join(outpath, 'model_pdbs/')
         mmcif_path = os.path.join(outpath, 'model_mmcifs/')
+        sfcif_path = os.path.join(outpath, 'model_sfcifs/')
         reflection_path = os.path.join(outpath, 'reflections/')
-        for folder in [sequence_path, model_path, mmcif_path, reflection_path]:
+        for folder in [sequence_path, model_path, mmcif_path, sfcif_path, reflection_path]:
             if os.path.exists(folder):
                 print(f"{folder:<80}" + f"{'already exists': >20}")
             else:
@@ -420,18 +421,21 @@ def fetch_pdb(idlist, outpath):
         sequence_path = outpath
         model_path = outpath
         mmcif_path = outpath
+        sfcif_path = outpath
         reflection_path = outpath
 
     codes = []
     with_sequence = []
     with_pdb = []
-    with_cif = []
+    with_mmcif = []
+    with_sfcif = []
     with_mtz = []
     for pdb_code in tqdm(idlist):
         valid_code = pdb_code.lower()
         seqlink = "https://www.rcsb.org/fasta/entry/" + valid_code.upper()
         pdblink = "https://files.rcsb.org/download/" + valid_code.upper() + ".pdb"
-        ciflink = "https://files.rcsb.org/download/" + valid_code.upper() + ".cif"
+        mmciflink = "https://files.rcsb.org/download/" + valid_code.upper() + ".cif"
+        sfciflink = "https://files.rcsb.org/download/" + valid_code.upper() + "-sf.cif"
         mtzlink = "https://edmaps.rcsb.org/coefficients/" + valid_code + ".mtz"
         codes.append(valid_code)
 
@@ -448,10 +452,16 @@ def fetch_pdb(idlist, outpath):
             with_pdb.append(0)
 
         try:
-            urllib.request.urlretrieve(ciflink, os.path.join(mmcif_path, valid_code+".cif"))
-            with_cif.append(1)
+            urllib.request.urlretrieve(mmciflink, os.path.join(mmcif_path, valid_code+".cif"))
+            with_mmcif.append(1)
         except:
-            with_cif.append(0)
+            with_mmcif.append(0)
+
+        try:
+            urllib.request.urlretrieve(sfciflink, os.path.join(sfcif_path, valid_code+"-sf.cif"))
+            with_sfcif.append(1)
+        except:
+            with_sfcif.append(0)
 
         try:
             urllib.request.urlretrieve(mtzlink, os.path.join(reflection_path, valid_code+".mtz"))
@@ -463,7 +473,8 @@ def fetch_pdb(idlist, outpath):
         "code" : codes,
         "with_sequence" : with_sequence,
         "with_pdb" : with_pdb,
-        "with_cif" : with_cif,
+        "with_mmcif" : with_mmcif,
+        "with_sfcif" : with_sfcif,
         "with_mtz" : with_mtz
     })
     stat_df.to_csv(os.path.join(outpath, "fetchpdb.csv"))

tests/conftest.py

@@ -13,6 +13,12 @@ def data_cif():
     filename = abspath(join(dirname(__file__), datapath))
     return filename
 
+@pytest.fixture
+def data_sfcif():
+    datapath = "data/1dur-sf.cif"
+    filename = abspath(join(dirname(__file__), datapath))
+    return filename
+
 @pytest.fixture
 def data_mtz_exp():
     datapath = "data/1dur.mtz"