diff --git a/docs/source/cosmo_array/index.rst b/docs/source/cosmo_array/index.rst index 44ef2fd3..e283cdc8 100644 --- a/docs/source/cosmo_array/index.rst +++ b/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. diff --git a/docs/source/visualisation/projection.rst b/docs/source/visualisation/projection.rst index e33ef60f..d99a3ec2 100644 --- a/docs/source/visualisation/projection.rst +++ b/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) diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index 0ff3878f..4c54d2f6 100644 --- a/swiftsimio/conversions.py +++ b/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] diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 3538012e..b69fcc04 100644 --- a/swiftsimio/objects.py +++ b/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 diff --git a/swiftsimio/visualisation/rotation.py b/swiftsimio/visualisation/rotation.py index 5fab663d..28c73d1c 100644 --- a/swiftsimio/visualisation/rotation.py +++ b/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,23 +48,23 @@ 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]], @@ -71,4 +72,6 @@ def rotation_matrix_from_vector(vector: float64, axis: str = "z") -> array: ] ) - 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) + ) diff --git a/swiftsimio/visualisation/smoothing_length/generate.py b/swiftsimio/visualisation/smoothing_length/generate.py index e151fc0b..2a63fc44 100644 --- a/swiftsimio/visualisation/smoothing_length/generate.py +++ b/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,7 +89,7 @@ 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, ) @@ -94,7 +97,7 @@ def generate_smoothing_lengths( # 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, ) diff --git a/tests/test_physical_conversion.py b/tests/test_physical_conversion.py index 2ade790f..03e09944 100644 --- a/tests/test_physical_conversion.py +++ b/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 + ) diff --git a/tests/test_smoothing_length_generation.py b/tests/test_smoothing_length_generation.py index 00e6a05f..0bba5b3d 100644 --- a/tests/test_smoothing_length_generation.py +++ b/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 + )