Merge pull request #29 from SWIFTSIM/bugfix_namedcolumn_mask

Bugfix: Copy named column data when deep copying

Author: kyleaoman

swiftgalaxy/iterator.py

@@ -239,7 +239,6 @@ def __init__(
             )
         else:
             num_targets = len(self.halo_catalogue._region_centre)
-            print(num_targets)
         # before evaluating optimized solutions:
         if num_targets == 0:
             # if we have 0 targets short-circuit

swiftgalaxy/reader.py

@@ -1910,7 +1910,9 @@ def _data_copy(
                         )
                         if data is not None:
                             setattr(
-                                new_named_columns_helper, f"_{named_column}", data[mask]
+                                new_named_columns_helper._named_column_dataset,
+                                f"_{named_column}",
+                                data[mask],
                             )
                 else:
                     data = getattr(

tests/test_copy.py

@@ -0,0 +1,32 @@
+from copy import deepcopy
+import unyt as u
+from unyt.testing import assert_allclose_units
+
+abstol_m = 1e2 * u.solMass
+reltol_m = 1.0e-4
+abstol_nd = 1.0e-4
+reltol_nd = 1.0e-4
+
+
+class TestCopySWIFTGalaxy:
+    def test_deepcopy_sg(self, sg):
+        """
+        Test that datasets get copied on deep copy.
+        """
+        # lazy load a dataset and a named column
+        sg.gas.masses
+        sg.gas.hydrogen_ionization_fractions.neutral
+        sg_copy = deepcopy(sg)
+        # check private attribute to not trigger lazy loading
+        assert_allclose_units(
+            sg.gas.masses,
+            sg_copy.gas._particle_dataset._masses,
+            rtol=reltol_m,
+            atol=abstol_m,
+        )
+        assert_allclose_units(
+            sg.gas.hydrogen_ionization_fractions.neutral,
+            sg_copy.gas.hydrogen_ionization_fractions._named_column_dataset._neutral,
+            rtol=reltol_nd,
+            atol=abstol_nd,
+        )

tests/test_masking.py

@@ -1,13 +1,9 @@
 import pytest
 import numpy as np
-import unyt as u
 from unyt.testing import assert_allclose_units
 from toysnap import present_particle_types
 from swiftgalaxy import MaskCollection
 
-abstol_c = 1 * u.pc  # less than this is ~0
-abstol_v = 10 * u.m / u.s  # less than this is ~0
-abstol_a = 1.0e-4 * u.rad
 abstol_nd = 1.0e-4
 reltol_nd = 1.0e-4
 
@@ -61,6 +57,21 @@ def test_bool_mask(self, sg, particle_name, before_load):
         ids = getattr(sg, particle_name).particle_ids
         assert_allclose_units(ids_before[mask], ids, rtol=0, atol=0)
 
+    @pytest.mark.parametrize("before_load", (True, False))
+    def test_namedcolumn_masked(self, sg, before_load):
+        """
+        Test that named columns get masked too.
+        """
+        neutral_before = sg.gas.hydrogen_ionization_fractions.neutral
+        mask = np.random.rand(neutral_before.size) > 0.5
+        if before_load:
+            sg.gas.hydrogen_ionization_fractions._named_column_dataset._neutral = None
+        sg.mask_particles(MaskCollection(**{"gas": mask}))
+        neutral = sg.gas.hydrogen_ionization_fractions.neutral
+        assert_allclose_units(
+            neutral_before[mask], neutral, rtol=reltol_nd, atol=abstol_nd
+        )
+
 
 class TestMaskingParticleDatasets:
     @pytest.mark.parametrize("particle_name", present_particle_types.values())