add hkls attribute rs.DataSet

reciprocalspaceship/dataset.py

@@ -43,6 +43,24 @@ class DataSet(pd.DataFrame):
     and attributes, please see the `Pandas.DataFrame documentation`_.
 
     .. _Pandas.DataFrame documentation: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html
+
+    Attributes
+    ----------
+    acentrics : rs.DataSet
+        Access only the acentric reflections in this dataset
+    cell : gemmi.UnitCell
+        The unit cell
+    centrics : rs.DataSet
+        Access only the centric reflections in this dataset
+    hkls : ndarray, shape=(n_reflections, 3)
+        Miller indices in DataSet. 
+    merged : bool
+        Whether this is a merged dataset or unmerged
+
+    spacegroup : gemmi.SpaceGroup
+        The space group
+    reindexing_ops : list
+        Possible reindexing ops consistent with the cell and spacegroup
     """
 
     _metadata = ["_spacegroup", "_cell", "_index_dtypes", "_merged"]
@@ -131,6 +149,36 @@ def merged(self):
     def merged(self, val):
         self._merged = val
 
+    @property
+    @range_indexed
+    def hkls(self):
+        """Miller indices"""
+        hkl = self[["H", "K", "L"]].to_numpy(dtype=np.int32)
+        return hkl
+
+    def get_hkls(self):
+        """ For backwards compatibility retain the get_hkls method in addition to the dataset.hkls attribute """
+        return self.hkls
+
+    @hkls.setter
+    @range_indexed
+    def hkls(self, hkls):
+        if isinstance(hkls, DataSet):
+            """ Convert to numpy if hkls is a dataset """
+            hkls = hkls.hkls
+        if isinstance(hkls, np.ndarray):
+            h,k,l = hkls[...,0], hkls[...,1], hkls[...,2]
+        else:
+            """ Try coercing to numpy """
+            try:
+                hkls = np.array(hkls)
+                h,k,l = hkls[...,0], hkls[...,1], hkls[...,2]
+            except:
+                raise ValueError("Unable to convert hkls to a suitable type. Please ensure hkls is a numpy array or rs.DataSet")
+        self['H'] = DataSeries(h, index = self.index, dtype='H')
+        self['K'] = DataSeries(k, index = self.index, dtype='H')
+        self['L'] = DataSeries(l, index = self.index, dtype='H')
+
     @property
     def centrics(self):
         """Access centric reflections in DataSet"""

tests/test_dataset.py

@@ -711,3 +711,42 @@ def test_select_mtzdtype_ValueError(data_merged, dtype):
     """
     with pytest.raises(ValueError):
         data_merged.select_mtzdtype(dtype)
+
+
+@pytest.mark.parametrize("merged", [True, False])
+@pytest.mark.parametrize("hkl_as_ds", [True, False])
+@pytest.mark.parametrize("range_index", [True, False])
+def test_hkls_property_setter(data_merged, data_unmerged, merged, hkl_as_ds, range_index):
+    """
+    Test the setter for the .hkls property of rs datasets
+    """
+    if merged:
+        input_ds = data_merged
+    else:
+        input_ds = data_unmerged
+
+    ds = input_ds.copy()
+
+    if range_index:
+        ds = ds.reset_index()
+
+    hmax = 20
+    n = len(ds)
+
+    hkls = np.random.randint(-hmax, hmax+1, size=(n, 3))
+    if hkl_as_ds:
+        hkls = rs.DataSet({
+            'H' : hkls[...,0],
+            'K' : hkls[...,1],
+            'L' : hkls[...,2],
+        })
+
+    ds.hkls = hkls
+    assert np.array_equal(hkls, ds.hkls)
+
+    # Test that all data remained the same
+    for k in input_ds:
+        if k not in ['H', 'K', 'L']:
+            assert np.array_equal(ds[k], input_ds[k])
+
+