Merge pull request #23 from SWIFTSIM/cosmo_quantity

Changes for swiftsimio v10 support.

Author: kyleaoman

swiftgalaxy/halo_catalogues.py

@@ -215,10 +215,7 @@ def _get_user_spatial_mask(self, snapshot_filename: str) -> SWIFTMask:
         if region is not None:
             for ax in range(3):
                 region[ax] = (
-                    [
-                        self.centre[ax] + region[ax][0],
-                        self.centre[ax] + region[ax][1],
-                    ]
+                    [self.centre[ax] + region[ax][0], self.centre[ax] + region[ax][1]]
                     if region[ax] is not None
                     else None
                 )
@@ -659,10 +656,7 @@ def _load(self) -> None:
             sm.constrain_indices(self._soap_index)
         else:
             sm.constrain_index(self._soap_index)
-        self._catalogue = SWIFTDataset(
-            self.soap_file,
-            mask=sm,
-        )
+        self._catalogue = SWIFTDataset(self.soap_file, mask=sm)
 
     @property
     def soap_index(self) -> Union[int, List[int]]:
@@ -1303,7 +1297,14 @@ def centre(self) -> cosmo_array:
         if self.centre_type in ("_gas", "_stars"):
             # {XYZ}c_gas and {XYZ}c_stars are relative to {XYZ}c
             relative_to = np.hstack(
-                [getattr(self._catalogue.positions, "{:s}c".format(c)) for c in "xyz"]
+                [
+                    cosmo_array(
+                        getattr(self._catalogue.positions, "{:s}c".format(c)),
+                        comoving=self._catalogue.units.comoving,
+                        cosmo_factor=cosmo_factor(a**1, self._catalogue.units.a),
+                    )
+                    for c in "xyz"
+                ]
             ).T
         else:
             # {XYZ}cmbp, {XYZ}cminpot and {XYZ}c are absolute
@@ -1319,9 +1320,13 @@ def centre(self) -> cosmo_array:
                 relative_to
                 + np.hstack(
                     [
-                        getattr(
-                            self._catalogue.positions,
-                            "{:s}c{:s}".format(c, self.centre_type),
+                        cosmo_array(
+                            getattr(
+                                self._catalogue.positions,
+                                "{:s}c{:s}".format(c, self.centre_type),
+                            ),
+                            comoving=self._catalogue.units.comoving,
+                            cosmo_factor=cosmo_factor(a**1, self._catalogue.units.a),
                         )
                         for c in "xyz"
                     ]
@@ -1356,7 +1361,14 @@ def velocity_centre(self) -> cosmo_array:
         if self.centre_type in ("_gas", "_stars"):
             # V{XYZ}c_gas and V{XYZ}c_stars are relative to {XYZ}c
             relative_to = np.hstack(
-                [getattr(self._catalogue.velocities, "v{:s}c".format(c)) for c in "xyz"]
+                [
+                    cosmo_array(
+                        getattr(self._catalogue.velocities, "v{:s}c".format(c)),
+                        comoving=self._catalogue.units.comoving,
+                        cosmo_factor=cosmo_factor(a**0, self._catalogue.units.a),
+                    )
+                    for c in "xyz"
+                ]
             ).T
         else:
             # V{XYZ}cmbp, V{XYZ}cminpot and V{XYZ}c are absolute
@@ -1371,9 +1383,13 @@ def velocity_centre(self) -> cosmo_array:
                 relative_to
                 + np.hstack(
                     [
-                        getattr(
-                            self._catalogue.velocities,
-                            "v{:s}c{:s}".format(c, self.centre_type),
+                        cosmo_array(
+                            getattr(
+                                self._catalogue.velocities,
+                                "v{:s}c{:s}".format(c, self.centre_type),
+                            ),
+                            comoving=self._catalogue.units.comoving,
+                            cosmo_factor=cosmo_factor(a**0, self._catalogue.units.a),
                         )
                         for c in "xyz"
                     ]
@@ -1641,8 +1657,7 @@ def _generate_bound_only_mask(self, sg: "SWIFTGalaxy") -> MaskCollection:
         """
 
         def in_one_of_ranges(
-            ints: NDArray[np.int_],
-            int_ranges: NDArray[np.int_],
+            ints: NDArray[np.int_], int_ranges: NDArray[np.int_]
         ) -> NDArray[np.bool_]:
             """
             Produces a boolean mask corresponding to `ints`.

swiftgalaxy/iterator.py

@@ -373,12 +373,10 @@ def _eval_dense_optimized_solution(self):
         # grid should be at least 1 cell in size so that we efficiently group targets
         # in the same grid location
         grid_element_dim = (
-            np.max(
-                target_sizes // sm.cell_size[np.newaxis] + 1,
-                axis=0,
-            )
+            np.max(target_sizes // sm.cell_size[np.newaxis], axis=0)
             .to_value(u.dimensionless)
             .astype(int)
+            + 1
         )
         # Cells are not guaranteed to have a vertex at box coordinate (0, 0, 0)
         # but usually are. At least one of our grids is misaligned with the cells
@@ -393,16 +391,18 @@ def _eval_dense_optimized_solution(self):
             cells_dim % grid_element_dim > 0
         ).astype(int)
         target_grid_offsets_aligned, target_grid_indices_aligned = np.modf(
-            target_centres / (grid_element_dim * sm.cell_size)
+            (target_centres / (grid_element_dim * sm.cell_size)).to_value(
+                u.dimensionless
+            )
         )
         target_grid_offsets_offset, target_grid_indices_offset = np.modf(
-            (target_centres + 0.5 * grid_element_dim * sm.cell_size)
-            / (grid_element_dim * sm.cell_size)
+            (
+                (target_centres + 0.5 * grid_element_dim * sm.cell_size)
+                / (grid_element_dim * sm.cell_size)
+            ).to_value(u.dimensionless)
             - 1
         )
-        target_grid_indices_aligned = target_grid_indices_aligned.to_value(
-            u.dimensionless
-        ).astype(int)
+        target_grid_indices_aligned = target_grid_indices_aligned.astype(int)
         # need a "box wrap" of targets to the far side of the grid for the offset grid:
         target_grid_indices_offset = np.where(
             np.logical_and(
@@ -411,9 +411,7 @@ def _eval_dense_optimized_solution(self):
             grid_dim - 1,
             target_grid_indices_offset,
         )  # complains about missing cosmo_array info
-        target_grid_indices_offset = target_grid_indices_offset.to_value(
-            u.dimensionless
-        ).astype(int)
+        target_grid_indices_offset = target_grid_indices_offset.astype(int)
         target_grid_distances_aligned = np.sqrt(
             np.sum(
                 np.power(
@@ -603,8 +601,7 @@ def __iter__(self) -> Generator:
                         boost=False,
                     )
                     swift_galaxy._transform(
-                        self.coordinate_frame_from._velocity_like_transform,
-                        boost=True,
+                        self.coordinate_frame_from._velocity_like_transform, boost=True
                     )
                 elif self.auto_recentre:
                     swift_galaxy.recentre(self.halo_catalogue.centre)
@@ -757,8 +754,6 @@ def dm_median_position(
             kwargs = [dict()] * len(self.iteration_order)
         for sg, iteration_location in zip(self, self.iteration_order):
             result[sg.halo_catalogue._multi_galaxy_catalogue_mask] = func(
-                sg,
-                *args[iteration_location],
-                **kwargs[iteration_location],
+                sg, *args[iteration_location], **kwargs[iteration_location]
             )
         return result

swiftgalaxy/reader.py

@@ -55,10 +55,11 @@ def _apply_box_wrap(coords: cosmo_array, boxsize: Optional[cosmo_array]) -> cosm
     out : :class:`~swiftsimio.objects.cosmo_array`
         The coordinates wrapped to lie within the box dimensions.
     """
-    if boxsize is None:
-        return coords
-    retval = (coords + boxsize / 2.0) % boxsize - boxsize / 2.0
-    return retval
+    return (
+        (coords + boxsize / 2.0) % boxsize - boxsize / 2.0
+        if boxsize is not None
+        else coords
+    )
 
 
 def _apply_translation(coords: cosmo_array, offset: cosmo_array) -> cosmo_array:
@@ -1573,9 +1574,7 @@ def __init__(
                 getattr(self.metadata, f"{particle_name}_properties").group
             ]
             TypeDatasetHelper = type(
-                f"{nice_name}DatasetHelper",
-                (_SWIFTGroupDatasetHelper, object),
-                dict(),
+                f"{nice_name}DatasetHelper", (_SWIFTGroupDatasetHelper, object), dict()
             )
             self._particle_dataset_helpers[particle_name] = TypeDatasetHelper(
                 super().__getattribute__(particle_name), self

tests/test_coordinate_transformations.py

@@ -2,14 +2,18 @@
 import numpy as np
 import unyt as u
 from unyt.testing import assert_allclose_units
-from swiftsimio.objects import cosmo_array, cosmo_factor, a
+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 swiftgalaxy import SWIFTGalaxy
 
 reltol = 1.01  # allow some wiggle room for floating point roundoff
-abstol_c = 10 * u.pc  # less than this is ~0
-abstol_v = 10 * u.m / u.s  # less than this is ~0
+abstol_c = cosmo_quantity(
+    10, u.pc, comoving=True, cosmo_factor=cosmo_factor(a**1, 1.0)
+)  # less than this is ~0
+abstol_v = cosmo_quantity(
+    10, u.m / u.s, comoving=True, cosmo_factor=cosmo_factor(a**0, 1.0)
+)  # less than this is ~0
 
 expected_xy = {
     "gas": cosmo_array(
@@ -227,8 +231,24 @@ def test_auto_recentering_off(
         Positions should still be offcentre.
         """
         xyz = getattr(sg_autorecentre_off, particle_name).coordinates
-        assert (np.abs(xyz[:, :2] - 2 * u.Mpc) <= expected_xy).all()
-        assert (np.abs(xyz[:, 2] - 2 * u.Mpc) <= expected_z).all()
+        assert (
+            np.abs(
+                xyz[:, :2]
+                - cosmo_quantity(
+                    2, u.Mpc, comoving=True, cosmo_factor=cosmo_factor(a**1, 1.0)
+                )
+            )
+            <= expected_xy
+        ).all()
+        assert (
+            np.abs(
+                xyz[:, 2]
+                - cosmo_quantity(
+                    2, u.Mpc, comoving=True, cosmo_factor=cosmo_factor(a**1, 1.0)
+                )
+            )
+            <= expected_z
+        ).all()
 
     @pytest.mark.parametrize(
         "particle_name, expected_vxy, expected_vz",
@@ -241,8 +261,30 @@ def test_auto_recentering_off_velocity(
         Velocities should still be offcentre.
         """
         vxyz = getattr(sg_autorecentre_off, particle_name).velocities
-        assert (np.abs(vxyz[:, :2] - 200 * u.km / u.s) <= expected_vxy).all()
-        assert (np.abs(vxyz[:, 2] - 200 * u.km / u.s) <= expected_vz).all()
+        assert (
+            np.abs(
+                vxyz[:, :2]
+                - cosmo_quantity(
+                    200,
+                    u.km / u.s,
+                    comoving=True,
+                    cosmo_factor=cosmo_factor(a**0, 1.0),
+                )
+            )
+            <= expected_vxy
+        ).all()
+        assert (
+            np.abs(
+                vxyz[:, 2]
+                - cosmo_quantity(
+                    200,
+                    u.km / u.s,
+                    comoving=True,
+                    cosmo_factor=cosmo_factor(a**0, 1.0),
+                )
+            )
+            <= expected_vz
+        ).all()
 
 
 class TestManualCoordinateTransformations: