Continue increasing test coverage.

Author: kyleaoman

swiftgalaxy/reader.py

@@ -848,9 +848,7 @@ def _mask_dataset(self, mask: slice) -> None:
         else:
             if self._swiftgalaxy._spatial_mask is None:
                 # get a count of particles in the box
-                num_part = self._particle_dataset.metadata.num_part[
-                    particle_metadata.particle_type
-                ]
+                num_part = getattr(self.metadata, f"n_{particle_metadata.group_name}")
             else:
                 # get a count of particles in the spatial mask region
                 num_part = np.sum(
@@ -1029,20 +1027,16 @@ def spherical_coordinates(self) -> _CoordinateHelper:
                     a**0, scale_factor=r.cosmo_factor.scale_factor
                 ),
             )
-            if self.cylindrical_coordinates is not None:
+            if self._cylindrical_coordinates is not None:
                 phi = self.cylindrical_coordinates.phi
             else:
-                phi = cosmo_array(
+                phi = (
                     np.arctan2(
                         self.cartesian_coordinates.y, self.cartesian_coordinates.x
-                    ),
-                    units=unyt.rad,
-                    comoving=r.comoving,
-                    cosmo_factor=cosmo_factor(
-                        a**0, scale_factor=r.cosmo_factor.scale_factor
-                    ),
-                )
-                phi[phi < 0] = phi[phi < 0] + 2 * np.pi * unyt.rad
+                    )
+                    * unyt.rad
+                )  # arctan2 returns dimensionless
+                phi[phi < 0] += 2 * np.pi * np.ones_like(phi)[phi < 0]
             self._spherical_coordinates = dict(_r=r, _theta=theta, _phi=phi)
         return _CoordinateHelper(
             self._spherical_coordinates,
@@ -1198,19 +1192,13 @@ def cylindrical_coordinates(self) -> _CoordinateHelper:
             if self._spherical_coordinates is not None:
                 phi = self.spherical_coordinates.phi
             else:
-                # np.where returns ndarray
-                phi = np.arctan2(
-                    self.cartesian_coordinates.y, self.cartesian_coordinates.x
-                ).view(np.ndarray)
-                phi = np.where(phi < 0, phi + 2 * np.pi, phi)
-                phi = cosmo_array(
-                    phi,
-                    units=unyt.rad,
-                    comoving=rho.comoving,
-                    cosmo_factor=cosmo_factor(
-                        a**0, scale_factor=rho.cosmo_factor.scale_factor
-                    ),
-                )
+                phi = (
+                    np.arctan2(
+                        self.cartesian_coordinates.y, self.cartesian_coordinates.x
+                    )
+                    * unyt.rad
+                )  # arctan2 returns dimensionless
+                phi[phi < 0] += 2 * np.pi * np.ones_like(phi)[phi < 0]
             z = self.cartesian_coordinates.z
             self._cylindrical_coordinates = dict(_rho=rho, _phi=phi, _z=z)
         return _CoordinateHelper(

tests/test_coordinate_transformations.py

@@ -4,7 +4,13 @@
 from unyt.testing import assert_allclose_units
 from swiftsimio.objects import cosmo_array, cosmo_factor, a, cosmo_quantity
 from scipy.spatial.transform import Rotation
-from toysnap import present_particle_types, toysnap_filename, ToyHF
+from toysnap import (
+    present_particle_types,
+    toysnap_filename,
+    ToyHF,
+    create_toysnap,
+    remove_toysnap,
+)
 from swiftgalaxy import SWIFTGalaxy
 from swiftgalaxy.reader import _apply_translation, _apply_4transform
 
@@ -604,6 +610,27 @@ def test_auto_recentering_with_copied_coordinate_frame(self, sg):
                 toysnap_filename, ToyHF(), auto_recentre=True, coordinate_frame_from=sg
             )
 
+    def test_invalid_coordinate_frame_from(self, sg):
+        """
+        Check that we get an error if coordinate_frame_from has mismatched internal units.
+        """
+        new_time_unit = u.s
+        assert sg.metadata.units.time != new_time_unit
+        sg.metadata.units.time = new_time_unit
+        try:
+            create_toysnap()
+            with pytest.raises(
+                ValueError, match="Internal units \\(length and time\\) of"
+            ):
+                SWIFTGalaxy(
+                    toysnap_filename,
+                    ToyHF(),
+                    coordinate_frame_from=sg,
+                    auto_recentre=False,
+                )
+        finally:
+            remove_toysnap()
+
     def test_copied_coordinate_transform(self, sg):
         """
         Check that a SWIFTGalaxy initialised to copy the coordinate frame

tests/test_copy.py

@@ -1,7 +1,9 @@
+import numpy as np
 from copy import copy, deepcopy
 import pytest
 import unyt as u
 from unyt.testing import assert_allclose_units
+from swiftgalaxy.masks import MaskCollection
 
 abstol_m = 1e2 * u.solMass
 reltol_m = 1.0e-4
@@ -139,3 +141,44 @@ def test_deepcopy_namedcolumn(self, sg):
             rtol=reltol_nd,
             atol=abstol_nd,
         )
+
+
+class TestCopyMaskCollection:
+    def test_copy_mask_collection(self):
+        """
+        Test that masks get copied.
+        """
+        mc = MaskCollection(
+            gas=np.ones(100, dtype=bool),
+            dark_matter=np.s_[:20],
+            stars=None,
+            black_holes=np.arange(3),
+        )
+        mc_copy = copy(mc)
+        assert set(mc_copy.__dict__.keys()) == set(mc.__dict__.keys())
+        for k in ("gas", "dark_matter", "stars", "black_holes"):
+            comparison = getattr(mc, k) == getattr(mc_copy, k)
+            if type(comparison) is bool:
+                assert comparison
+            else:
+                assert all(comparison)
+
+    def test_deepcopy_mask_collection(self):
+        """
+        Test that masks get copied along with values. Since the object isn't
+        really "deep", shallow copy and deepcopy have the same expectation.
+        """
+        mc = MaskCollection(
+            gas=np.ones(100, dtype=bool),
+            dark_matter=np.s_[:20],
+            stars=None,
+            black_holes=np.arange(3),
+        )
+        mc_copy = deepcopy(mc)
+        assert set(mc_copy.__dict__.keys()) == set(mc.__dict__.keys())
+        for k in ("gas", "dark_matter", "stars", "black_holes"):
+            comparison = getattr(mc, k) == getattr(mc_copy, k)
+            if type(comparison) is bool:
+                assert comparison
+            else:
+                assert all(comparison)

tests/test_creation.py

@@ -1,13 +1,10 @@
+"""
+Tests checking that we can create objects, if these fail something
+fundamental has gone wrong.
+"""
+
 from swiftgalaxy import SWIFTGalaxy
-from toysnap import (
-    create_toysnap,
-    remove_toysnap,
-    toysnap_filename,
-    n_g_all,
-    n_dm_all,
-    n_s_all,
-    n_bh_all,
-)
+from toysnap import create_toysnap, remove_toysnap, toysnap_filename, n_g_1
 
 
 class TestSWIFTGalaxyCreation:
@@ -81,30 +78,36 @@ def test_tab_completion(self, sg):
         # finally, check that we didn't lazy-load everything!
         assert sg.gas._particle_dataset._coordinates is None
 
-    def test_no_masks(self):
+    def test_mask_preloaded_namedcolumn(self):
         """
-        Check that if we have no masks we read everything in the box (and warn about it).
+        If namedcolumn data was loaded during evaluation of a mask, it needs to be masked
+        during initialization.
         """
+        from toysnap import ToyHF
+
+        def load_namedcolumn(method):
+
+            def wrapper(self, sg):
+                sg.gas.hydrogen_ionization_fractions.neutral
+                return method(self, sg)
+
+            return wrapper
+
+        # decorate the mask evaluation to load an (unused) namedcolumn
+        ToyHF._generate_bound_only_mask = load_namedcolumn(
+            ToyHF._generate_bound_only_mask
+        )
+
         try:
             create_toysnap()
-            sg = SWIFTGalaxy(
-                toysnap_filename,
-                None,  # no halo_catalogue is easiest way to get no mask
-                transforms_like_coordinates={"coordinates", "extra_coordinates"},
-                transforms_like_velocities={"velocities", "extra_velocities"},
+            sg = SWIFTGalaxy(toysnap_filename, ToyHF())
+            # confirm that we loaded a namedcolumn during initialization:
+            assert (
+                sg.gas.hydrogen_ionization_fractions._internal_dataset._neutral
+                is not None
             )
-            # check that extra mask is blank for all particle types:
-            assert sg._extra_mask.gas is None
-            assert sg._extra_mask.dark_matter is None
-            assert sg._extra_mask.stars is None
-            assert sg._extra_mask.black_holes is None
-            # check that cell mask is blank for all particle types:
-            assert sg._spatial_mask is None
-            # check that we read all the particles:
-            assert sg.gas.masses.size == n_g_all
-            assert sg.dark_matter.masses.size == n_dm_all
-            assert sg.stars.masses.size == n_s_all
-            assert sg.black_holes.masses.size == n_bh_all
+            # confirm that it got masked:
+            assert sg.gas.hydrogen_ionization_fractions.neutral.size == n_g_1
         finally:
             remove_toysnap()
 

tests/test_masking.py

@@ -1,8 +1,21 @@
+"""
+Tests for applying masks to swiftgalaxy, datasets and named columns.
+"""
+
 import pytest
 import numpy as np
 from unyt.testing import assert_allclose_units
 from toysnap import present_particle_types
-from swiftgalaxy import MaskCollection
+from swiftgalaxy import SWIFTGalaxy, MaskCollection
+from toysnap import (
+    create_toysnap,
+    remove_toysnap,
+    toysnap_filename,
+    n_g_all,
+    n_dm_all,
+    n_s_all,
+    n_bh_all,
+)
 
 abstol_nd = 1.0e-4
 reltol_nd = 1.0e-4
@@ -72,6 +85,41 @@ def test_namedcolumn_masked(self, sg, before_load):
             neutral_before[mask], neutral, rtol=reltol_nd, atol=abstol_nd
         )
 
+    def test_mask_without_spatial_mask(self):
+        """
+        Check that if we have no masks we read everything in the box (and warn about it).
+        Then that we can still apply an extra mask, and a second one (there's specific
+        logic for applying two consecutively).
+        """
+        try:
+            create_toysnap()
+            sg = SWIFTGalaxy(
+                toysnap_filename,
+                None,  # no halo_catalogue is easiest way to get no mask
+                transforms_like_coordinates={"coordinates", "extra_coordinates"},
+                transforms_like_velocities={"velocities", "extra_velocities"},
+            )
+            # check that extra mask is blank for all particle types:
+            assert sg._extra_mask.gas is None
+            assert sg._extra_mask.dark_matter is None
+            assert sg._extra_mask.stars is None
+            assert sg._extra_mask.black_holes is None
+            # check that cell mask is blank for all particle types:
+            assert sg._spatial_mask is None
+            # check that we read all the particles:
+            assert sg.gas.masses.size == n_g_all
+            assert sg.dark_matter.masses.size == n_dm_all
+            assert sg.stars.masses.size == n_s_all
+            assert sg.black_holes.masses.size == n_bh_all
+            # now apply an extra mask
+            sg.mask_particles(MaskCollection(gas=np.s_[:1000]))
+            assert sg.gas.masses.size == 1000
+            # and the second consecutive one
+            sg.mask_particles(MaskCollection(gas=np.s_[:100]))
+            assert sg.gas.masses.size == 100
+        finally:
+            remove_toysnap()
+
 
 class TestMaskingParticleDatasets:
     @pytest.mark.parametrize("particle_name", present_particle_types.values())