Continue increasing test coverage.

Author: kyleaoman

swiftgalaxy/reader.py

@@ -103,7 +103,7 @@ def _apply_translation(coords: cosmo_array, offset: cosmo_array) -> cosmo_array:
         offset = offset.to_comoving()
     elif hasattr(offset, "comoving") and not coords.comoving:
         offset = offset.to_physical()
-    elif not hasattr(offset, "comoving"):
+    else:  # not hasattr(offset, "comoving")
         msg = (
             "Translation assumed to be in comoving (not physical) coordinates."
             if coords.comoving
@@ -1125,10 +1125,13 @@ def spherical_velocities(self) -> _CoordinateHelper:
                 + _sin_t * _sin_p * self.cartesian_velocities.y
                 - _cos_t * self.cartesian_velocities.z
             )
-            v_p = (
-                -_sin_p * self.cartesian_velocities.x
-                + _cos_p * self.cartesian_velocities.y
-            )
+            if self._cylindrical_velocities is not None:
+                v_p = self.cylindrical_velocities.phi
+            else:
+                v_p = (
+                    -_sin_p * self.cartesian_velocities.x
+                    + _cos_p * self.cartesian_velocities.y
+                )
             self._spherical_velocities = dict(_v_r=v_r, _v_t=v_t, _v_p=v_p)
         return _CoordinateHelper(
             self._spherical_velocities,
@@ -1531,7 +1534,7 @@ def __init__(
             self._velocity_like_transform = np.eye(4)
         if self.halo_catalogue is None:
             # in server mode we don't have a halo_catalogue yet
-            pass
+            self._spatial_mask = getattr(self, "_spatial_mask", None)
         elif self.halo_catalogue._user_spatial_offsets is not None:
             self._spatial_mask = self.halo_catalogue._get_user_spatial_mask(
                 self.snapshot_filename
@@ -2022,13 +2025,18 @@ def _transform(self, transform4: cosmo_array, boost: bool = False) -> None:
             if boost
             else self.transforms_like_coordinates
         )
+        transform_units = (
+            self.metadata.units.length / self.metadata.units.time
+            if boost
+            else self.metadata.units.length
+        )
         for particle_name in self.metadata.present_group_names:
             dataset = getattr(self, particle_name)._particle_dataset
             for field_name in transformable:
                 field_data = getattr(dataset, f"_{field_name}")
                 if field_data is not None:
                     field_data = _apply_4transform(
-                        field_data, transform4.to_value(), transform4.units
+                        field_data, transform4, transform_units
                     )
                     setattr(dataset, f"_{field_name}", field_data)
         if boost:

tests/test_coordinate_transformations.py

@@ -363,6 +363,18 @@ def test_translate(self, sg, particle_name, coordinate_name, before_load):
         xyz = getattr(getattr(sg, particle_name), f"{coordinate_name}")
         assert_allclose_units(xyz_before + translation, xyz, rtol=1.0e-4, atol=abstol_c)
 
+    def test_translate_warn_comoving_missing(self, sg):
+        """
+        If the translation does not have comoving information issue a warning.
+        """
+        translation = u.unyt_array(
+            [1, 1, 1],
+            units=u.Mpc,
+        )
+        msg = "Translation assumed to be in comoving"
+        with pytest.warns(RuntimeWarning, match=msg):
+            sg.translate(translation)
+
     @pytest.mark.parametrize("velocity_name", ("velocities", "extra_velocities"))
     @pytest.mark.parametrize("particle_name", present_particle_types.values())
     @pytest.mark.parametrize("before_load", (True, False))
@@ -425,6 +437,64 @@ def test_box_wrap(self, sg, particle_name, coordinate_name):
         xyz = getattr(getattr(sg, particle_name), f"{coordinate_name}")
         assert_allclose_units(xyz_before, xyz, rtol=1.0e-4, atol=abstol_c)
 
+    @pytest.mark.parametrize("coordinate_name", ("coordinates", "extra_coordinates"))
+    @pytest.mark.parametrize("particle_name", present_particle_types.values())
+    @pytest.mark.parametrize("before_load", (True, False))
+    def test_transform(self, sg, particle_name, coordinate_name, before_load):
+        """
+        Check that affine transformation works.
+        """
+        xyz_before = getattr(getattr(sg, particle_name), f"{coordinate_name}")
+        if before_load:
+            setattr(
+                getattr(sg, particle_name)._particle_dataset,
+                f"_{coordinate_name}",
+                None,
+            )
+        translation = cosmo_array(
+            [1, 1, 1],
+            units=u.Mpc,
+            comoving=True,
+            cosmo_factor=cosmo_factor(a**1, scale_factor=1.0),
+        )
+        transform = np.eye(4)
+        transform[:3, :3] = rot
+        transform[3, :3] = translation.to_comoving_value(u.Mpc)
+        sg._transform(transform)
+        xyz = getattr(getattr(sg, particle_name), f"{coordinate_name}")
+        assert_allclose_units(
+            xyz_before.dot(rot) + translation, xyz, rtol=1.0e-4, atol=abstol_c
+        )
+
+    @pytest.mark.parametrize("coordinate_name", ("velocities", "extra_velocities"))
+    @pytest.mark.parametrize("particle_name", present_particle_types.values())
+    @pytest.mark.parametrize("before_load", (True, False))
+    def test_transform_velocity(self, sg, particle_name, coordinate_name, before_load):
+        """
+        Check that affine transformation works.
+        """
+        xyz_before = getattr(getattr(sg, particle_name), f"{coordinate_name}")
+        if before_load:
+            setattr(
+                getattr(sg, particle_name)._particle_dataset,
+                f"_{coordinate_name}",
+                None,
+            )
+        translation = cosmo_array(
+            [100, 100, 100],
+            units=u.km / u.s,
+            comoving=True,
+            cosmo_factor=cosmo_factor(a**0, scale_factor=1.0),
+        )
+        transform = np.eye(4)
+        transform[:3, :3] = rot
+        transform[3, :3] = translation.to_comoving_value(u.km / u.s)
+        sg._transform(transform, boost=True)
+        xyz = getattr(getattr(sg, particle_name), f"{coordinate_name}")
+        assert_allclose_units(
+            xyz_before.dot(rot) + translation, xyz, rtol=1.0e-4, atol=abstol_v
+        )
+
 
 class TestSequentialTransformations:
     @pytest.mark.parametrize("before_load", (True, False))
@@ -640,7 +710,48 @@ def test_comoving_physical_conversion(self, comoving):
             comoving=comoving,
             cosmo_factor=cosmo_factor(a**1, scale_factor=1.0),
         )
-        # identity 4transform:
-        transform = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]])
+        transform = np.eye(4)  # identity 4transform
         result = _apply_4transform(coords, transform, transform_units=u.Mpc)
         assert result.comoving == comoving
+
+
+class TestCoordinateProperties:
+
+    def test_centre(self, sg):
+        """
+        Check the centre attribute.
+        """
+        new_centre = cosmo_array(
+            [1, 2, 3],
+            units=u.Mpc,
+            comoving=True,
+            cosmo_factor=cosmo_factor(a**1, scale_factor=1.0),
+        )
+        sg.recentre(new_centre)
+        assert_allclose_units(
+            sg.halo_catalogue.centre + new_centre,
+            sg.centre,
+        )
+
+    def test_velocity_centre(self, sg):
+        """
+        Check the velocity_centre attribute.
+        """
+        new_centre = cosmo_array(
+            [100, 200, 300],
+            units=u.km / u.s,
+            comoving=True,
+            cosmo_factor=cosmo_factor(a**0, scale_factor=1.0),
+        )
+        sg.recentre_velocity(new_centre)
+        assert_allclose_units(
+            sg.halo_catalogue.velocity_centre + new_centre,
+            sg.velocity_centre,
+        )
+
+    def test_rotation(self, sg):
+        """
+        Check the rotation attribute.
+        """
+        sg.rotate(Rotation.from_matrix(rot))
+        assert np.allclose(sg.rotation.as_matrix(), rot)

tests/test_copy.py

@@ -1,4 +1,5 @@
 from copy import copy, deepcopy
+import pytest
 import unyt as u
 from unyt.testing import assert_allclose_units
 
@@ -24,13 +25,19 @@ def test_copy_sg(self, sg):
             is None
         )
 
-    def test_deepcopy_sg(self, sg):
+    @pytest.mark.parametrize("derived_coords_initialized", [True, False])
+    def test_deepcopy_sg(self, sg, derived_coords_initialized):
         """
         Test that dataset arrays get copied on deep copy.
         """
         # lazy load a dataset and a named column
         sg.gas.masses
         sg.gas.hydrogen_ionization_fractions.neutral
+        if derived_coords_initialized:
+            sg.gas.spherical_coordinates
+            sg.gas.spherical_velocities
+            sg.gas.cylindrical_coordinates
+            sg.gas.cylindrical_velocities
         sg_copy = deepcopy(sg)
         # check private attribute to not trigger lazy loading
         assert_allclose_units(
@@ -45,6 +52,28 @@ def test_deepcopy_sg(self, sg):
             rtol=reltol_nd,
             atol=abstol_nd,
         )
+        if derived_coords_initialized:
+            assert_allclose_units(
+                sg.gas.spherical_coordinates.r,
+                sg_copy.gas._spherical_coordinates["_r"],
+            )
+            assert_allclose_units(
+                sg.gas.spherical_velocities.r,
+                sg_copy.gas._spherical_velocities["_v_r"],
+            )
+            assert_allclose_units(
+                sg.gas.cylindrical_coordinates.rho,
+                sg_copy.gas._cylindrical_coordinates["_rho"],
+            )
+            assert_allclose_units(
+                sg.gas.cylindrical_velocities.rho,
+                sg_copy.gas._cylindrical_velocities["_v_rho"],
+            )
+        else:
+            assert sg_copy.gas._spherical_coordinates is None
+            assert sg_copy.gas._spherical_velocities is None
+            assert sg_copy.gas._cylindrical_coordinates is None
+            assert sg_copy.gas._cylindrical_velocities is None
 
 
 class TestCopyDataset:

tests/test_creation.py

@@ -1,5 +1,4 @@
-import pytest
-from swiftgalaxy import SWIFTGalaxy, Standalone
+from swiftgalaxy import SWIFTGalaxy
 from toysnap import (
     create_toysnap,
     remove_toysnap,
@@ -9,8 +8,6 @@
     n_s_all,
     n_bh_all,
 )
-from swiftsimio.objects import cosmo_array, cosmo_factor, a
-import unyt as u
 
 
 class TestSWIFTGalaxyCreation:
@@ -90,26 +87,9 @@ def test_no_masks(self):
         """
         try:
             create_toysnap()
-            with pytest.warns(UserWarning, match="No spatial_offsets provided."):
-                sa = Standalone(
-                    extra_mask=None,
-                    centre=cosmo_array(
-                        [0, 0, 0],
-                        u.Mpc,
-                        comoving=True,
-                        cosmo_factor=cosmo_factor(a**1, 1.0),
-                    ),
-                    velocity_centre=cosmo_array(
-                        [0, 0, 0],
-                        u.km / u.s,
-                        comoving=True,
-                        cosmo_factor=cosmo_factor(a**0, 1.0),
-                    ),
-                    spatial_offsets=None,
-                )
             sg = SWIFTGalaxy(
                 toysnap_filename,
-                sa,
+                None,  # no halo_catalogue is easiest way to get no mask
                 transforms_like_coordinates={"coordinates", "extra_coordinates"},
                 transforms_like_velocities={"velocities", "extra_velocities"},
             )
@@ -119,8 +99,7 @@ def test_no_masks(self):
             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:
-            for cell_mask in sg._spatial_mask.cell_mask.values():
-                assert cell_mask.all()
+            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

tests/test_derived_coordinates.py

@@ -206,6 +206,17 @@ def test_spherical_velocity_phi(self, sg, particle_name, alias):
             spherical_v_phi, v_phi_from_cartesian, rtol=1.0e-4, atol=abstol_v
         )
 
+    @pytest.mark.parametrize("ctype", ["coordinates", "velocities"])
+    def test_copy_from_cylindrical(self, sg, ctype):
+        """
+        Check that copying the azimuth from cylindrical if already evaluated works.
+        """
+        getattr(sg.gas, f"cylindrical_{ctype}").phi  # trigger evaluation
+        assert (
+            getattr(sg.gas, f"spherical_{ctype}").phi
+            is getattr(sg.gas, f"cylindrical_{ctype}").phi
+        )
+
 
 class TestCylindricalCoordinates:
     @pytest.mark.parametrize("particle_name", present_particle_types.values())
@@ -329,6 +340,17 @@ def test_cylindrical_velocity_z(self, sg, particle_name, alias):
             cylindrical_v_z, v_z_from_cartesian, rtol=1.0e-4, atol=abstol_v
         )
 
+    @pytest.mark.parametrize("ctype", ["coordinates", "velocities"])
+    def test_copy_from_spherical(self, sg, ctype):
+        """
+        Check that copying the azimuth from spherical if already evaluated works.
+        """
+        getattr(sg.gas, f"spherical_{ctype}").phi  # trigger evaluation
+        assert (
+            getattr(sg.gas, f"cylindrical_{ctype}").phi
+            is getattr(sg.gas, f"spherical_{ctype}").phi
+        )
+
 
 class TestInteractionWithCoordinateTransformations:
     @pytest.mark.parametrize("coordinate_type", ("coordinates", "velocities"))

tests/test_str.py

@@ -24,3 +24,8 @@ def test_namedcolumn_fullname(self, sg):
             sg.gas.hydrogen_ionization_fractions._fullname
             == "gas.hydrogen_ionization_fractions"
         )
+
+    def test_sg_string(self, sg):
+        string = str(sg)
+        assert "SWIFTGalaxy at" in string
+        assert repr(sg) == string