Merge pull request #136 from Hekstra-Lab/precog

Fixed `read_precognition()` as per #135 and updated tests

reciprocalspaceship/io/precognition.py

@@ -1,3 +1,5 @@
+import warnings
+
 import gemmi
 import pandas as pd
 
@@ -7,20 +9,21 @@
 def read_precognition(hklfile, spacegroup=None, cell=None, logfile=None):
     """
     Initialize attributes and populate the DataSet object with data from
-    a HKL file of reflections. This is the output format used by
+    a precognition hkl or ii file of reflections. This is the output format used by
     Precognition when processing Laue diffraction data.
 
     Parameters
     ----------
     hklfile : str or file
-        name of an hkl file or a file object
+        name of an hkl or ii file or a file object
     spacegroup : str or int
         If int, this should specify the space group number. If str,
         this should be a space group symbol
     cell : tuple or list of floats
         Unit cell parameters
     logfile : str or file
-        name of a log file to parse to get cell parameters and sg
+        name of a log file to parse to get cell parameters and spacegroup. Only
+        used when spacegroup and/or cell are not explicitly provided.
     """
     # Read data from HKL file
     if hklfile.endswith(".hkl"):
@@ -40,6 +43,7 @@ def read_precognition(hklfile, spacegroup=None, cell=None, logfile=None):
             F.rename(columns={"F(+)": "F", "SigF(+)": "SigF"}, inplace=True)
             mtztypes = ["H", "H", "H", "F", "Q"]
 
+    # Read data from II file
     elif hklfile.endswith(".ii"):
         usecols = range(10)
         F = pd.read_csv(
@@ -62,37 +66,48 @@ def read_precognition(hklfile, spacegroup=None, cell=None, logfile=None):
         )
         mtztypes = ["H", "H", "H", "I", "R", "R", "R", "R", "J", "Q"]
 
-        # If logfile is given, read cell parameters and spacegroup
-        if logfile:
-            from os.path import basename
+    # Limit use to supported file formats
+    else:
+        raise ValueError("rs.read_precognition() only supports .ii and .hkl files")
 
-            with open(logfile, "r") as log:
-                lines = log.readlines()
+    # If logfile is given, read cell parameters and spacegroup
+    # Assign these as temporary variables, and determine priority later.
 
-            # Read spacegroup
-            sgline = [l for l in lines if "space-group" in l][0]
-            sg = [s for s in sgline.split() if "#" in s][0].lstrip("#")
+    if logfile:
+        from os.path import basename
 
-            # Read cell parameters
-            block = [i for i, l in enumerate(lines) if basename(hklfile) in l][0]
-            lengths = lines[block - 19].split()[-3:]
-            a, b, c = map(float, lengths)
-            angles = lines[block - 18].split()[-3:]
-            alpha, beta, gamma = map(float, angles)
-            cell = (a, b, c, alpha, beta, gamma)
+        with open(logfile, "r") as log:
+            lines = log.readlines()
 
-    # GH#32: Limit use to supported file formats
-    else:
-        raise ValueError("rs.read_precognition() only supports .ii and .hkl files")
+        # Read spacegroup
+        sgline = [l for l in lines if "space-group" in l][0]
+        spacegroup_from_log = [s for s in sgline.split() if "#" in s][0].lstrip("#")
+
+        # Read cell parameters
+        block = [i for i, l in enumerate(lines) if basename(hklfile) in l][0]
+        lengths = lines[block - 19].split()[-3:]
+        a, b, c = map(float, lengths)
+        angles = lines[block - 18].split()[-3:]
+        alpha, beta, gamma = map(float, angles)
+        cell_from_log = (a, b, c, alpha, beta, gamma)
 
     dataset = DataSet(F)
     dataset = dataset.astype(dict(zip(dataset.columns, mtztypes)))
     dataset.set_index(["H", "K", "L"], inplace=True)
 
     # Set DataSet attributes
-    if cell and (len(cell) == 6):
+    # Prioritize explicitly supplied arguments
+    if cell:
         dataset.cell = cell
+    elif logfile:
+        dataset.cell = cell_from_log
+
     if spacegroup:
         dataset.spacegroup = spacegroup
+    elif logfile:
+        dataset.spacegroup = spacegroup_from_log
+
+    if cell and spacegroup and logfile:
+        warnings.warn("Ignoring logfile, as cell and spacegroup are both provided")
 
     return dataset

tests/io/test_precognition.py

@@ -59,9 +59,16 @@ def test_read_ii(IOtest_ii, IOtest_log, spacegroup, cell, log):
     if log == "log":
         log = IOtest_log
 
-    result = rs.read_precognition(
-        IOtest_ii, spacegroup=spacegroup, cell=cell, logfile=log
-    )
+    # Check warning for when log, cell, and spacegroup are all provided
+    if log and cell and spacegroup:
+        with pytest.warns(UserWarning):
+            result = rs.read_precognition(
+                IOtest_ii, spacegroup=spacegroup, cell=cell, logfile=log
+            )
+    else:
+        result = rs.read_precognition(
+            IOtest_ii, spacegroup=spacegroup, cell=cell, logfile=log
+        )
 
     # Check main DataSet features
     assert isinstance(result, rs.DataSet)
@@ -71,23 +78,27 @@ def test_read_ii(IOtest_ii, IOtest_log, spacegroup, cell, log):
 
     # Check _metadata
     assert result._index_dtypes == {"H": "HKL", "K": "HKL", "L": "HKL"}
+
     if spacegroup:
         assert result.spacegroup.xhm() == gemmi.SpaceGroup(spacegroup).xhm()
+    elif log:
+        assert result.spacegroup.xhm() == "P 21 21 21"
     else:
         assert result.spacegroup is None
-    if log:
-        assert result.cell.a == 34.4660
-        assert result.cell.b == 45.6000
-        assert result.cell.c == 99.5850
-        assert result.cell.alpha == 90.0
-        assert result.cell.beta == 90.0
-        assert result.cell.gamma == 90.0
-    elif cell:
+
+    if cell:
         assert result.cell.a == cell[0]
         assert result.cell.b == cell[1]
         assert result.cell.c == cell[2]
         assert result.cell.alpha == cell[3]
         assert result.cell.beta == cell[4]
         assert result.cell.gamma == cell[5]
+    elif log:
+        assert result.cell.a == 34.4660
+        assert result.cell.b == 45.6000
+        assert result.cell.c == 99.5850
+        assert result.cell.alpha == 90.0
+        assert result.cell.beta == 90.0
+        assert result.cell.gamma == 90.0
     else:
         assert result.cell is None