Cleanup while checking unyt usage.

swiftsimio/_array_functions.py

@@ -111,7 +111,9 @@
 
 
 def _copy_cosmo_array_attributes(from_ca, to_ca):
-    if not isinstance(to_ca, objects.cosmo_array):
+    if not isinstance(to_ca, objects.cosmo_array) and isinstance(
+        from_ca, objects.cosmo_array
+    ):
         return to_ca
     if hasattr(from_ca, "cosmo_factor"):
         to_ca.cosmo_factor = from_ca.cosmo_factor

swiftsimio/snapshot_writer.py

@@ -14,6 +14,7 @@
 from functools import reduce
 
 from swiftsimio import metadata
+from swiftsimio.objects import cosmo_array
 from swiftsimio.metadata.cosmology.cosmology_fields import a_exponents
 
 
@@ -51,7 +52,7 @@ class __SWIFTWriterParticleDataset(object):
         checks if all required datasets are empty.
     check_consistent(self)
         performs consistency checks on dataset
-    generate_smoothing_lengths(self, boxsize: Union[List[unyt.unyt_quantity], unyt.unyt_quantity], dimension: int)
+    generate_smoothing_lengths(self, boxsize: cosmo_array, dimension: int)
         automatically generates the smoothing lengths
     write_particle_group(self, file_handle: h5py.File, compress: bool)
         writes the particle group's required properties to file.
@@ -164,7 +165,7 @@ def check_consistent(self) -> bool:
 
     def generate_smoothing_lengths(
         self,
-        boxsize: Union[List[unyt.unyt_quantity], unyt.unyt_quantity],
+        boxsize: cosmo_array,
         dimension: int,
     ):
         """
@@ -175,7 +176,7 @@ def generate_smoothing_lengths(
 
         Parameters
         ----------
-        boxsize : unyt.unyt_quantity or list of unyt.unyt_quantity
+        boxsize : cosmo_array or cosmo_quantity
             size of SWIFT computational box
         dimension : int
             number of box dimensions
@@ -271,7 +272,7 @@ def get_attributes(self, scale_factor: float) -> dict:
 
             # Find the scale factor associated quantities
             a_exp = a_exponents.get(name, 0)
-            a_factor = scale_factor ** a_exp
+            a_factor = scale_factor**a_exp
 
             attributes_dict[output_handle] = {
                 "Conversion factor to CGS (not including cosmological corrections)": [
@@ -506,7 +507,7 @@ class SWIFTSnapshotWriter(object):
     def __init__(
         self,
         unit_system: Union[unyt.UnitSystem, str],
-        box_size: Union[list, unyt.unyt_quantity],
+        box_size: cosmo_array,
         dimension=3,
         compress=True,
         extra_header: Union[None, dict] = None,
@@ -522,7 +523,7 @@ def __init__(
         ----------
         unit_system : unyt.UnitSystem or str
             unit system for dataset
-        boxsize : list or unyt.unyt_quantity
+        boxsize : cosmo_array
             size of simulation box and associated units
         dimension : int, optional
             dimensions of simulation

swiftsimio/visualisation/projection.py

@@ -3,53 +3,36 @@
 """
 
 from typing import Union
-from math import sqrt, ceil
 from numpy import (
-    float64,
     float32,
-    int32,
-    zeros,
     array,
-    arange,
-    ndarray,
     ones,
-    isclose,
     matmul,
     empty_like,
     logical_and,
     s_,
+    ceil,
 )
-from unyt import unyt_array, unyt_quantity, exceptions
+from unyt import unyt_quantity, exceptions
 from swiftsimio import SWIFTDataset, cosmo_array
 
 from swiftsimio.reader import __SWIFTGroupDataset
-from swiftsimio.accelerated import jit, NUM_THREADS, prange
 
 from swiftsimio.visualisation.projection_backends import backends, backends_parallel
 
-# Backwards compatability
-
-from swiftsimio.visualisation.projection_backends.kernels import (
-    kernel_gamma,
-    kernel_constant,
-)
-from swiftsimio.visualisation.projection_backends.kernels import (
-    kernel_single_precision as kernel,
-)
-
 scatter = backends["fast"]
 scatter_parallel = backends_parallel["fast"]
 
 
 def project_pixel_grid(
     data: __SWIFTGroupDataset,
-    boxsize: unyt_array,
+    boxsize: cosmo_array,
     resolution: int,
     project: Union[str, None] = "masses",
-    region: Union[None, unyt_array] = None,
+    region: Union[None, cosmo_array] = None,
     mask: Union[None, array] = None,
     rotation_matrix: Union[None, array] = None,
-    rotation_center: Union[None, unyt_array] = None,
+    rotation_center: Union[None, cosmo_array] = None,
     parallel: bool = False,
     backend: str = "fast",
     periodic: bool = True,
@@ -68,7 +51,7 @@ def project_pixel_grid(
     data: __SWIFTGroupDataset
         The SWIFT dataset that you wish to visualise (get this from ``load``)
 
-    boxsize: unyt_array
+    boxsize: cosmo_array
         The box-size of the simulation.
 
     resolution: int
@@ -81,7 +64,7 @@ def project_pixel_grid(
         always create it as ``data.gas.my_variable = data.gas.other_variable
         * data.gas.masses``.
 
-    region: unyt_array, optional
+    region: cosmo_array, optional
         Region, determines where the image will be created (this corresponds
         to the left and right-hand edges, and top and bottom edges) if it is
         not None. It should have a length of four or six, and take the form:
@@ -116,7 +99,7 @@ def project_pixel_grid(
     Returns
     -------
 
-    image: unyt_array
+    image: cosmo_array
         Projected image with units of project / length^2, of size ``res`` x ``res``.
 
 
@@ -140,7 +123,7 @@ def project_pixel_grid(
                 )
         except AttributeError:
             raise exceptions.InvalidUnitOperation(
-                "Ensure that rotation_center is a unyt array with the same units as coordinates"
+                "Ensure that rotation_center is a cosmo_array with the same units as coordinates"
             )
 
     number_of_particles = data.coordinates.shape[0]
@@ -202,12 +185,12 @@ def project_pixel_grid(
         if data.coordinates.comoving:
             if not hsml.compatible_with_comoving():
                 raise AttributeError(
-                    f"Physical smoothing length is not compatible with comoving coordinates!"
+                    "Physical smoothing length is not compatible with comoving coordinates!"
                 )
         else:
             if not hsml.compatible_with_physical():
                 raise AttributeError(
-                    f"Comoving smoothing length is not compatible with physical coordinates!"
+                    "Comoving smoothing length is not compatible with physical coordinates!"
                 )
     except AttributeError:
         # No hsml present. If they are using the 'histogram' backend, we
@@ -269,10 +252,10 @@ def project_gas_pixel_grid(
     data: SWIFTDataset,
     resolution: int,
     project: Union[str, None] = "masses",
-    region: Union[None, unyt_array] = None,
+    region: Union[None, cosmo_array] = None,
     mask: Union[None, array] = None,
     rotation_matrix: Union[None, array] = None,
-    rotation_center: Union[None, unyt_array] = None,
+    rotation_center: Union[None, cosmo_array] = None,
     parallel: bool = False,
     backend: str = "fast",
     periodic: bool = True,
@@ -303,7 +286,7 @@ def project_gas_pixel_grid(
         always create it as ``data.gas.my_variable = data.gas.other_variable
         * data.gas.masses``.
 
-    region: unyt_array, optional
+    region: cosmo_array, optional
         Region, determines where the image will be created (this corresponds
         to the left and right-hand edges, and top and bottom edges) if it is
         not None. It should have a length of four or six, and take the form:
@@ -374,9 +357,9 @@ def project_gas(
     data: SWIFTDataset,
     resolution: int,
     project: Union[str, None] = "masses",
-    region: Union[None, unyt_array] = None,
+    region: Union[None, cosmo_array] = None,
     mask: Union[None, array] = None,
-    rotation_center: Union[None, unyt_array] = None,
+    rotation_center: Union[None, cosmo_array] = None,
     rotation_matrix: Union[None, array] = None,
     parallel: bool = False,
     backend: str = "fast",
@@ -406,7 +389,7 @@ def project_gas(
         always create it as ``data.gas.my_variable = data.gas.other_variable
         * data.gas.masses``.
 
-    region: unyt_array, optional
+    region: cosmo_array, optional
         Region, determines where the image will be created (this corresponds
         to the left and right-hand edges, and top and bottom edges) if it is
         not None. It should have a length of four or six, and take the form:
@@ -442,7 +425,7 @@ def project_gas(
     Returns
     -------
 
-    image: unyt_array
+    image: cosmo_array
         Projected image with units of project / length^2, of size ``res`` x
         ``res``.
 
@@ -474,23 +457,23 @@ def project_gas(
         x_range = region[1] - region[0]
         y_range = region[3] - region[2]
         max_range = max(x_range, y_range)
-        units = 1.0 / (max_range ** 2)
+        units = 1.0 / (max_range**2)
         # Unfortunately this is required to prevent us from {over,under}flowing
         # the units...
         units.convert_to_units(1.0 / (x_range.units * y_range.units))
     else:
         max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1])
-        units = 1.0 / (max_range ** 2)
+        units = 1.0 / (max_range**2)
         # Unfortunately this is required to prevent us from {over,under}flowing
         # the units...
-        units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2)
+        units.convert_to_units(1.0 / data.metadata.boxsize.units**2)
 
     comoving = data.gas.coordinates.comoving
     coord_cosmo_factor = data.gas.coordinates.cosmo_factor
     if project is not None:
         units *= getattr(data.gas, project).units
         project_cosmo_factor = getattr(data.gas, project).cosmo_factor
-        new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2
+        new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2
     else:
         new_cosmo_factor = coord_cosmo_factor ** (-2)
 

swiftsimio/visualisation/ray_trace.py

@@ -217,15 +217,15 @@ def core_panels_parallel(
 
 def panel_pixel_grid(
     data: __SWIFTGroupDataset,
-    boxsize: unyt.unyt_array,
+    boxsize: cosmo_array,
     resolution: int,
     panels: int,
     project: Union[str, None] = "masses",
-    region: Union[None, unyt.unyt_array] = None,
+    region: Union[None, cosmo_array] = None,
     mask: Union[None, np.array] = None,
     rotation_matrix: Union[None, np.array] = None,
-    rotation_center: Union[None, unyt.unyt_array] = None,
-) -> unyt.unyt_array:
+    rotation_center: Union[None, cosmo_array] = None,
+) -> cosmo_array:
     if rotation_center is not None:
         try:
             if rotation_center.units == data.coordinates.units:
@@ -332,10 +332,10 @@ def panel_gas(
     resolution: int,
     panels: int,
     project: Union[str, None] = "masses",
-    region: Union[None, unyt.unyt_array] = None,
+    region: Union[None, cosmo_array] = None,
     mask: Union[None, np.array] = None,
     rotation_matrix: Union[None, np.array] = None,
-    rotation_center: Union[None, unyt.unyt_array] = None,
+    rotation_center: Union[None, cosmo_array] = None,
 ) -> cosmo_array:
     image = panel_pixel_grid(
         data=data.gas,
@@ -353,23 +353,23 @@ def panel_gas(
         x_range = region[1] - region[0]
         y_range = region[3] - region[2]
         max_range = max(x_range, y_range)
-        units = 1.0 / (max_range ** 2)
+        units = 1.0 / (max_range**2)
         # Unfortunately this is required to prevent us from {over,under}flowing
         # the units...
         units.convert_to_units(1.0 / (x_range.units * y_range.units))
     else:
         max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1])
-        units = 1.0 / (max_range ** 2)
+        units = 1.0 / (max_range**2)
         # Unfortunately this is required to prevent us from {over,under}flowing
         # the units...
-        units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2)
+        units.convert_to_units(1.0 / data.metadata.boxsize.units**2)
 
     comoving = data.gas.coordinates.comoving
     coord_cosmo_factor = data.gas.coordinates.cosmo_factor
     if project is not None:
         units *= getattr(data.gas, project).units
         project_cosmo_factor = getattr(data.gas, project).cosmo_factor
-        new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2
+        new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2
     else:
         new_cosmo_factor = coord_cosmo_factor ** (-2)