Repository: rs-station/reciprocalspaceship
Issue: `DataSet.write_mtz()` with unmerged DataSet moves reflections to ASU

If a DataSet object contains unmerged reflections, a call to `DataSet.write_mtz()` currently has the side effect of calling `DataSet.hkl_to_asu(inplace=True)` on the object:

```
In [1]: ds = rs.read_mtz("data_unmerged.mtz")

In [2]: ds["BATCH"].head()
Out[2]: 
H    K    L 
19   -11  2    1424
3    18   -1    328
-17  21   17    576
21   26   -7     70
-9   17   11    798
Name: BATCH, dtype: Batch

In [3]: ds.write_mtz("/dev/null")

In [4]: ds["BATCH"].head()
Out[4]: 
H   K   L 
19  11  2    1424
18  3   1     328
21  17  17    576
26  21  7      70
17  9   11    798
Name: BATCH, dtype: Batch
```

I consider this to be a bug because writing data to a file should not change the underlying object. 

Changed files:
tests/io/test_mtz.py, reciprocalspaceship/io/mtz.py

diff --git a/reciprocalspaceship/io/mtz.py b/reciprocalspaceship/io/mtz.py
index 59a6167..b9722e4 100644
--- a/reciprocalspaceship/io/mtz.py
+++ b/reciprocalspaceship/io/mtz.py
@@ -4,6 +4,7 @@ import pandas as pd
 
 from reciprocalspaceship import DataSet
 from reciprocalspaceship.dtypes.base import MTZDtype
+from reciprocalspaceship.utils import in_asu
 
 
 def from_gemmi(gemmi_mtz):
@@ -97,7 +98,9 @@ def to_gemmi(dataset, skip_problem_mtztypes=False):
 
     # Handle Unmerged data
     if not dataset.merged:
-        dataset.hkl_to_asu(inplace=True)
+        all_in_asu = in_asu(dataset.get_hkls(), dataset.spacegroup).all()
+        if not all_in_asu:
+            dataset.hkl_to_asu(inplace=True)
 
     # Construct data for Mtz object.
     mtz.add_dataset("reciprocalspaceship")
@@ -122,6 +125,10 @@ def to_gemmi(dataset, skip_problem_mtztypes=False):
             )
     mtz.set_data(temp[columns].to_numpy(dtype="float32"))
 
+    # Handle Unmerged data
+    if not dataset.merged and not all_in_asu:
+        dataset.hkl_to_observed(m_isym="M/ISYM", inplace=True)
+
     return mtz
 
 
diff --git a/tests/io/test_mtz.py b/tests/io/test_mtz.py
index e3210d2..9c1ec94 100644
--- a/tests/io/test_mtz.py
+++ b/tests/io/test_mtz.py
@@ -1,7 +1,6 @@
 import filecmp
 import tempfile
-import unittest
-from os.path import abspath, dirname, exists, join
+from os.path import exists
 
 import gemmi
 import numpy as np
@@ -12,81 +11,81 @@ import reciprocalspaceship as rs
 from reciprocalspaceship.utils import in_asu
 
 
-class TestMTZ(unittest.TestCase):
-    def test_read(self):
+def test_read_merged(IOtest_mtz):
+    """Test rs.read_mtz() with merged MTZ file"""
+    dataset = rs.read_mtz(IOtest_mtz)
 
-        datadir = join(abspath(dirname(__file__)), "../data/fmodel")
-        data = rs.read_mtz(join(datadir, "9LYZ.mtz"))
+    assert dataset.spacegroup.number == 96
+    assert dataset.columns.to_list() == ["FMODEL", "PHIFMODEL"]
+    assert dataset.index.names == ["H", "K", "L"]
+    assert isinstance(dataset, rs.DataSet)
 
-        # Confirm columns, indices, and metadata
-        self.assertEqual(data.spacegroup.number, 96)
-        self.assertEqual(data.columns.to_list(), ["FMODEL", "PHIFMODEL"])
-        self.assertEqual(list(data.index.names), ["H", "K", "L"])
-        self.assertIsInstance(data.spacegroup, gemmi.SpaceGroup)
-        self.assertIsInstance(data.cell, gemmi.UnitCell)
-        self.assertIsInstance(data, rs.DataSet)
-        self.assertIsInstance(data["FMODEL"], rs.DataSeries)
 
-        return
+def test_write_merged(IOtest_mtz):
+    """Test DataSet.write_mtz() with merged MTZ file"""
+    dataset = rs.read_mtz(IOtest_mtz)
 
-    def test_write(self):
+    with tempfile.NamedTemporaryFile(suffix=".mtz") as temp:
+        dataset.write_mtz(temp.name)
+        assert exists(temp.name)
 
-        datadir = join(abspath(dirname(__file__)), "../data/fmodel")
-        data = rs.read_mtz(join(datadir, "9LYZ.mtz"))
 
-        temp = tempfile.NamedTemporaryFile(suffix=".mtz")
+def test_write_merged_nosg(IOtest_mtz):
+    """Test that DataSet.write_mtz() without spacegroup raises AttributeError"""
+    dataset = rs.read_mtz(IOtest_mtz)
+    dataset.spacegroup = None
 
-        # Missing cell should raise AttributeError
-        data_missingcell = data.copy()
-        data_missingcell.cell = None
-        with self.assertRaises(AttributeError):
-            data_missingcell.write_mtz(temp.name)
+    with tempfile.NamedTemporaryFile(suffix=".mtz") as temp:
+        with pytest.raises(AttributeError):
+            dataset.write_mtz(temp.name)
 
-        # Missing spacegroup should raise AttributeError
-        data_missingsg = data.copy()
-        data_missingsg.spacegroup = None
-        with self.assertRaises(AttributeError):
-            data_missingsg.write_mtz(temp.name)
 
-        # Writing MTZ should produce a file
-        data.write_mtz(temp.name)
-        self.assertTrue(exists(temp.name))
-        temp.close()
+def test_write_merged_nocell(IOtest_mtz):
+    """Test that DataSet.write_mtz() without cell raises AttributeError"""
+    dataset = rs.read_mtz(IOtest_mtz)
+    dataset.cell = None
 
-        # Having a non-MTZType should raise AttributeError, unless flag
-        # is set
-        temp = tempfile.NamedTemporaryFile(suffix=".mtz")
-        data["nonMTZ"] = 1
-        with self.assertRaises(ValueError):
-            data.write_mtz(temp.name)
-        data.write_mtz(temp.name, skip_problem_mtztypes=True)
-        self.assertTrue(exists(temp.name))
-        temp.close()
+    with tempfile.NamedTemporaryFile(suffix=".mtz") as temp:
+        with pytest.raises(AttributeError):
+            dataset.write_mtz(temp.name)
 
-        return
 
-    def test_roundtrip(self):
+@pytest.mark.parametrize("skip_problem_mtztypes", [True, False])
+def test_write_merged_nonMTZDtype(IOtest_mtz, skip_problem_mtztypes):
+    """
+    Test skip_problem_mtztypes flag of DataSet.write_mtz()
+    """
+    dataset = rs.read_mtz(IOtest_mtz)
+    dataset["nonMTZ"] = 1
+
+    with tempfile.NamedTemporaryFile(suffix=".mtz") as temp:
+        if not skip_problem_mtztypes:
+            with pytest.raises(ValueError):
+                dataset.write_mtz(
+                    temp.name, skip_problem_mtztypes=skip_problem_mtztypes
+                )
+        else:
+            dataset.write_mtz(temp.name, skip_problem_mtztypes=skip_problem_mtztypes)
+            assert exists(temp.name)
 
-        datadir = join(abspath(dirname(__file__)), "../data/fmodel")
-        data = rs.read_mtz(join(datadir, "9LYZ.mtz"))
 
-        # Write data, read data, write data again... shouldn't change
-        temp = tempfile.NamedTemporaryFile(suffix=".mtz")
-        data.write_mtz(temp.name)
-        data2 = rs.read_mtz(temp.name)
-        temp2 = tempfile.NamedTemporaryFile(suffix=".mtz")
-        data2.write_mtz(temp2.name)
+def test_roundtrip_merged(IOtest_mtz):
+    """Test roundtrip of rs.read_mtz() and DataSet.write_mtz() with merged MTZ file"""
+    expected = rs.read_mtz(IOtest_mtz)
 
-        self.assertTrue(data.equals(data2))
-        self.assertEqual(data.spacegroup.number, data2.spacegroup.number)
-        self.assertEqual(data.cell.a, data2.cell.a)
-        self.assertTrue(filecmp.cmp(temp.name, temp2.name))
+    temp1 = tempfile.NamedTemporaryFile(suffix=".mtz")
+    temp2 = tempfile.NamedTemporaryFile(suffix=".mtz")
+
+    expected.write_mtz(temp1.name)
+    result = rs.read_mtz(temp1.name)
+    result.write_mtz(temp2.name)
 
-        # Clean up
-        temp.close()
-        temp2.close()
+    assert_frame_equal(result, expected)
+    assert filecmp.cmp(temp1.name, temp2.name)
 
-        return
+    # Clean up
+    temp1.close()
+    temp2.close()
 
 
 def test_read_unmerged(data_unmerged):
@@ -113,7 +112,7 @@ def test_read_unmerged_2m_isym(data_unmerged):
 
 
 @pytest.mark.parametrize("label_centrics", [True, False])
-def test_roundtrip(data_unmerged, label_centrics):
+def test_roundtrip_unmerged(data_unmerged, label_centrics):
     """
     Test roundtrip of rs.read_mtz() and DataSet.write_mtz() with unmerged data
     """
@@ -124,8 +123,8 @@ def test_roundtrip(data_unmerged, label_centrics):
     temp2 = tempfile.NamedTemporaryFile(suffix=".mtz")
     data_unmerged.write_mtz(temp.name)
     data2 = rs.read_mtz(temp.name)
+    data2 = data2[data_unmerged.columns]  # Ensure consistent column ordering
     data2.write_mtz(temp2.name)
-
     assert filecmp.cmp(temp.name, temp2.name)
     assert_frame_equal(data_unmerged, data2)
     assert data_unmerged.merged == data2.merged
@@ -133,3 +132,16 @@ def test_roundtrip(data_unmerged, label_centrics):
     # Clean up
     temp.close()
     temp2.close()
+
+
+@pytest.mark.parametrize("in_asu", [True, False])
+def test_unmerged_after_write(data_unmerged, in_asu):
+    """
+    #110: Test that unmerged DataSet objects are unchanged following calls to
+    DataSet.write_mtz()
+    """
+    if in_asu:
+        data_unmerged.hkl_to_asu(inplace=True)
+    expected = data_unmerged.copy()
+    data_unmerged.write_mtz("/dev/null")
+    assert_frame_equal(data_unmerged, expected)