Cannot unstack_anomalous with other columns in index

[kmdalton]

If a DataSet has additional columns besides ["H", "K", "L"] in the index, unstack anomalous will fail with a KeyError. Consider the following example:

import numpy as np
import reciprocalspaceship as rs

cell = [34., 45., 98., 90., 90., 90.]
spacegroup = 19
dmin = 4.
repeats = 1

h,k,l = rs.utils.generate_reciprocal_asu(cell, spacegroup, dmin, anomalous=True).T

ds = rs.DataSet({
    "H" : h,
    "K" : k,
    "L" : l,
    "I" : np.random.random(len(h)),
    "SIGI" : np.random.random(len(h)),
}, cell=cell, spacegroup=spacegroup, merged=True).infer_mtz_dtypes()
ds['M'] = 0

ds = ds.set_index(['M', 'H', 'K', 'L'])
print(ds)

ds.unstack_anomalous()

Output:

$ python test.py
                     I       SIGI
M H  K  L
0 -8 -3 -5  0.12483723 0.37942517
        -4    0.982229  0.4665011
        -3 0.055360574 0.43935636
        -2  0.24871129  0.0687793
        -1  0.22456451 0.72497946
...                ...        ...
  8  3  1  0.058091898  0.4236312
        2    0.7901098  0.5751216
        3   0.35781875 0.98029584
        4   0.30590278  0.5441381
        5    0.7120791 0.38891163

[2447 rows x 2 columns]
Traceback (most recent call last):
  File "/home/kmdalton/opt/anaconda/envs/careless/lib/python3.10/site-packages/pandas/core/indexes/base.py", line 3629, in get_loc
    return self._engine.get_loc(casted_key)
  File "pandas/_libs/index.pyx", line 136, in pandas._libs.index.IndexEngine.get_loc
  File "pandas/_libs/index.pyx", line 163, in pandas._libs.index.IndexEngine.get_loc
  File "pandas/_libs/hashtable_class_helper.pxi", line 5198, in pandas._libs.hashtable.PyObjectHashTable.get_item
  File "pandas/_libs/hashtable_class_helper.pxi", line 5206, in pandas._libs.hashtable.PyObjectHashTable.get_item
KeyError: 'M'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/kmdalton/opt/sandbox/stack/test.py", line 23, in <module>
    ds.unstack_anomalous()
  File "/home/kmdalton/opt/reciprocalspaceship/reciprocalspaceship/decorators.py", line 49, in wrapped
    result = result._index_from_names(names, inplace=True)
  File "/home/kmdalton/opt/reciprocalspaceship/reciprocalspaceship/decorators.py", line 23, in wrapped
    return f(ds, *args, **kwargs)
  File "/home/kmdalton/opt/reciprocalspaceship/reciprocalspaceship/dataset.py", line 193, in _index_from_names
    self.set_index(names, inplace=True)
  File "/home/kmdalton/opt/reciprocalspaceship/reciprocalspaceship/dataset.py", line 241, in set_index
    self._index_dtypes[key] = self[key].dtype.name
  File "/home/kmdalton/opt/anaconda/envs/careless/lib/python3.10/site-packages/pandas/core/frame.py", line 3505, in __getitem__
    indexer = self.columns.get_loc(key)
  File "/home/kmdalton/opt/anaconda/envs/careless/lib/python3.10/site-packages/pandas/core/indexes/base.py", line 3631, in get_loc
    raise KeyError(key) from err
KeyError: 'M'