Add simultaneous physical and unit conversion

docs/source/cosmo_array/index.rst

@@ -45,6 +45,8 @@ is done with the :meth:`~swiftsimio.objects.cosmo_array.to_physical`, :meth:`~sw
    # Convert in-place
    rho_gas.convert_to_physical()
 
+If you instead want to get the raw array values in either physical or comoving units, the :meth:`~swiftsimio.objects.cosmo_array.to_physical_value` and :meth:`~swiftsimio.objects.cosmo_array.to_comoving_value` are provided, analogous to the :meth:`~unyt.unyt_array.unyt_array.to_value`.
+
 The ``valid_transform`` is a boolean flag that is set to ``False`` for some arrays that don't make sense to convert to comoving.
 
 :class:`~swiftsimio.objects.cosmo_array` supports array arithmetic and the entire :mod:`numpy` range of functions. Attempting to combine arrays (e.g. by addition) will validate the cosmology information first. The implementation is designed to be permissive: it will only raise exceptions when a genuinely invalid combination is encountered, but is tolerant of missing cosmology information. When one argument in a relevant operation (like addition, for example) is not a :class:`~swiftsimio.objects.cosmo_array` the attributes of the :class:`~swiftsimio.objects.cosmo_array` will be assumed for both arguments. In such cases a warning is produced stating that this assumption has been made.

docs/source/visualisation/projection.rst

@@ -38,15 +38,17 @@ Example
        periodic=True,
    )
 
-   # Let's say we wish to save it as msun / kpc^2,
+   # Let's say we wish to save it as msun / kpc^2 (physical, not comoving),
    from unyt import msun, kpc
-   mass_map.convert_to_units(msun / kpc**2)
-
    from matplotlib.pyplot import imsave
    from matplotlib.colors import LogNorm
 
    # Normalize and save
-   imsave("gas_surface_dens_map.png", LogNorm()(mass_map.value), cmap="viridis")
+   imsave(
+       "gas_surface_dens_map.png",
+       LogNorm()(mass_map.to_physical_value(msun / kpc**2)),
+       cmap="viridis",
+   )
 
 
 This basic demonstration creates a mass surface density map.
@@ -86,13 +88,15 @@ this:
    temp_map = mass_weighted_temp_map / mass_map
 
    from unyt import K
-   temp_map.convert_to_units(K)
-
    from matplotlib.pyplot import imsave
    from matplotlib.colors import LogNorm
 
    # Normalize and save
-   imsave("temp_map.png", LogNorm()(temp_map.value), cmap="twilight")
+   imsave(
+       "temp_map.png",
+       LogNorm()(temp_map.to_physical_value(K)),
+       cmap="twilight",
+   )
 
 
 The output from this example, when used with the example data provided in the
@@ -222,7 +226,7 @@ is shown in the ``velociraptor`` section.
 
    # The angular momentum vector will point perpendicular to the galaxy disk.
    # If your simulation contains stars, use lx_star
-   angular_momentum_vector = np.array([lx.value, ly.value, lz.value])
+   angular_momentum_vector = cosmo_array([lx, ly, lz])
    angular_momentum_vector /= np.linalg.norm(angular_momentum_vector)
 
    face_on_rotation_matrix = rotation_matrix_from_vector(angular_momentum_vector)

swiftsimio/conversions.py

@@ -103,11 +103,13 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology:
             # SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy.
             # This is an exact inversion of the procedure performed in astropy.
             critical_density_0 = astropy_units.Quantity(
-                critdens_const * H0.to("1/s").value ** 2,
+                critdens_const * H0.to_value("1/s") ** 2,
                 astropy_units.g / astropy_units.cm ** 3,
             )
 
-            Tcmb0 = (Omega_r * critical_density_0.value / a_B_c2) ** (1.0 / 4.0)
+            Tcmb0 = (Omega_r * critical_density_0.to_value("g/cm**3") / a_B_c2) ** (
+                1.0 / 4.0
+            )
 
         try:
             Neff = cosmo["N_eff"][0]

swiftsimio/objects.py

@@ -10,7 +10,7 @@
 """
 
 import unyt
-from unyt import unyt_array, unyt_quantity
+from unyt import unyt_array, unyt_quantity, Unit
 from unyt.array import multiple_output_operators, _iterable, POWER_MAPPING
 from numbers import Number as numeric_type
 from typing import Iterable, Union, Tuple, Callable, Optional
@@ -1454,6 +1454,36 @@ def to_comoving(self) -> "cosmo_array":
 
         return copied_data
 
+    def to_physical_value(self, units: Unit) -> np.ndarray:
+        """
+        Returns a copy of the array values in the specified physical units.
+
+        Parameters
+        ----------
+        units : unyt.unit_object.Unit
+
+        Returns
+        -------
+        out : np.ndarray
+            Copy of the array values in the specified physical units.
+        """
+        return self.to_physical().to_value(units)
+
+    def to_comoving_value(self, units: Unit) -> np.ndarray:
+        """
+        Returns a copy of the array values in the specified comoving units.
+
+        Parameters
+        ----------
+        units : unyt.unit_object.Unit
+
+        Returns
+        -------
+        out : np.ndarray
+            Copy of the array values in the specified comoving units.
+        """
+        return self.to_comoving().to_value(units)
+
     def compatible_with_comoving(self) -> bool:
         """
         Is this :class:`~swiftsimio.objects.cosmo_array` compatible with a comoving

swiftsimio/visualisation/rotation.py

@@ -2,12 +2,11 @@
 Rotation matrix calculation routines.
 """
 
-from numpy import float64, array, cross, identity, dot, matmul
-from numpy.linalg import norm, inv
-from math import sin, acos
+import numpy as np
+import unyt as u
 
 
-def rotation_matrix_from_vector(vector: float64, axis: str = "z") -> array:
+def rotation_matrix_from_vector(vector: np.float64, axis: str = "z") -> np.ndarray:
     """
     Calculate a rotation matrix from a vector. The comparison vector is
     assumed to be along an axis, x, y, or z (by default this is z). The
@@ -34,10 +33,12 @@ def rotation_matrix_from_vector(vector: float64, axis: str = "z") -> array:
         Rotation matrix (3x3).
     """
 
-    normed_vector = vector / norm(vector)
+    normed_vector = vector / np.linalg.norm(vector)
+    if isinstance(normed_vector, u.unyt_array):
+        normed_vector = normed_vector.to_value(u.dimensionless)
 
     # Directional vector describing the axis we wish to look 'down'
-    original_direction = array([0.0, 0.0, 0.0], dtype=float64)
+    original_direction = np.zeros(3, dtype=np.float64)
     switch = {"x": 0, "y": 1, "z": 2}
 
     try:
@@ -47,28 +48,30 @@ def rotation_matrix_from_vector(vector: float64, axis: str = "z") -> array:
             f"Parameter axis must be one of x, y, or z. You supplied {axis}."
         )
 
-    cross_product = cross(original_direction, normed_vector)
-    mod_cross_product = norm(cross_product)
-    cross_product /= mod_cross_product
+    cross_product = np.cross(original_direction, normed_vector)
+    mod_cross_product = np.linalg.norm(cross_product)
+    cross_product = cross_product / mod_cross_product
 
     if mod_cross_product <= 1e-6:
         # This case only covers when we point the vector
         # in the exact opposite direction (e.g. flip z).
-        output = identity(3)
+        output = np.identity(3)
         output[switch[axis], switch[axis]] = -1.0
 
         return output
     else:
-        cos_theta = dot(original_direction, normed_vector)
-        sin_theta = sin(acos(cos_theta))
+        cos_theta = np.dot(original_direction, normed_vector)
+        sin_theta = np.sin(np.arccos(cos_theta))
 
         # Skew symmetric matrix for cross product
-        A = array(
+        A = np.array(
             [
                 [0.0, -cross_product[2], cross_product[1]],
                 [cross_product[2], 0.0, -cross_product[0]],
                 [-cross_product[1], cross_product[0], 0.0],
             ]
         )
 
-        return inv(identity(3) + sin_theta * A + (1 - cos_theta) * matmul(A, A))
+        return np.linalg.inv(
+            np.identity(3) + sin_theta * A + (1 - cos_theta) * np.dot(A, A)
+        )

swiftsimio/visualisation/smoothing_length/generate.py

@@ -55,7 +55,10 @@ def generate_smoothing_lengths(
 
     number_of_parts = coordinates.shape[0]
 
-    tree = KDTree(coordinates.value, boxsize=boxsize.to(coordinates.units).value)
+    tree = KDTree(
+        coordinates.to_value(coordinates.units),
+        boxsize=boxsize.to_value(coordinates.units),
+    )
 
     smoothing_lengths = np.empty(number_of_parts, dtype=np.float32)
     smoothing_lengths[-1] = -0.1
@@ -86,15 +89,15 @@ def generate_smoothing_lengths(
 
         try:
             d, _ = tree.query(
-                coordinates[starting_index:ending_index].value,
+                coordinates[starting_index:ending_index].to_value(coordinates.units),
                 k=neighbours_search,
                 workers=-1,
             )
         except TypeError:
             # Backwards compatibility with older versions of
             # scipy.
             d, _ = tree.query(
-                coordinates[starting_index:ending_index].value,
+                coordinates[starting_index:ending_index].to_value(coordinates.units),
                 k=neighbours_search,
                 n_jobs=-1,
             )

tests/test_physical_conversion.py

@@ -1,21 +1,44 @@
 from tests.helper import requires
-from swiftsimio import load, cosmo_array
-from numpy import array_equal
+import unyt as u
+from swiftsimio import load
+from numpy import allclose
 
 
 @requires("cosmological_volume.hdf5")
 def test_convert(filename):
     """
-    Check that the conversion to physical units is done correctly
+    Check that the conversion to physical units is done correctly.
     """
     data = load(filename)
     coords = data.gas.coordinates
-    units = coords.units
+    units = u.kpc
+    assert units != coords.units  # ensure we make a non-trivial conversion
+    assert data.metadata.a != 1.0  # ensure we make a non-trivial conversion
     coords_physical = coords.to_physical()
 
     # array_equal applied to cosmo_array's is aware of physical & comoving
     # make sure to compare bare arrays:
-    assert array_equal(
-        coords.to_value(units) * data.metadata.a, coords_physical.to_value(units)
+    assert allclose(
+        coords.to_value(units) * data.metadata.a,
+        coords_physical.to_value(units),
+        rtol=1e-6,
     )
     return
+
+
+@requires("cosmological_volume.hdf5")
+def test_convert_to_value(filename):
+    """
+    Check that conversions to numerical values are correct.
+    """
+    data = load(filename)
+    coords = data.gas.coordinates
+    units = u.kpc
+    assert units != coords.units  # ensure we make a non-trivial conversion
+    assert data.metadata.a != 1.0  # ensure we make a non-trivial conversion
+    coords_physical_values = coords.to_physical_value(units)
+    coords_comoving_values = coords.to_comoving_value(units)
+    print(coords_physical_values / (coords_comoving_values * data.metadata.a))
+    assert allclose(
+        coords_physical_values, coords_comoving_values * data.metadata.a, rtol=1e-6
+    )

tests/test_smoothing_length_generation.py

@@ -2,7 +2,9 @@
 Tests the smoothing length generation code.
 """
 
-from swiftsimio import load
+import numpy as np
+import unyt as u
+from swiftsimio import load, cosmo_array
 from swiftsimio.visualisation.smoothing_length import generate_smoothing_lengths
 from tests.helper import requires
 
@@ -71,3 +73,22 @@ def test_generate_smoothing_length_faster(filename):
     ).all()
 
     return
+
+
+def test_generate_smoothing_length_return_type():
+    x = cosmo_array(
+        np.arange(20), u.Mpc, comoving=False, scale_factor=0.5, scale_exponent=1
+    )
+    xgrid, ygrid, zgrid = np.meshgrid(x, x, x)
+    coords = np.vstack((xgrid.flatten(), ygrid.flatten(), zgrid.flatten())).T
+    lbox = cosmo_array(
+        [20, 20, 20], u.Mpc, comoving=False, scale_factor=0.5, scale_exponent=1
+    )
+    from_ca_input = generate_smoothing_lengths(coords, lbox, 1)
+    assert from_ca_input.units == coords.units
+    assert from_ca_input.comoving == coords.comoving
+    assert from_ca_input.cosmo_factor == coords.cosmo_factor
+    from_ua_input = generate_smoothing_lengths(u.unyt_array(coords), lbox, 1)
+    assert isinstance(from_ua_input, u.unyt_array) and not isinstance(
+        from_ua_input, cosmo_array
+    )