diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index aab5bc05..309cfa2a 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -16,7 +16,7 @@ jobs: with: python-version: ${{ matrix.python-version }} - name: Cache pip - uses: actions/cache@v2 + uses: actions/cache@v4 with: # This path is specific to Ubuntu path: ~/.cache/pip diff --git a/docs/source/conf.py b/docs/source/conf.py index 3088ca6d..d0f61711 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -40,6 +40,7 @@ "sphinx.ext.napoleon", "sphinx.ext.mathjax", "sphinx.ext.autosummary", + "sphinx.ext.intersphinx", ] # Add any paths that contain templates here, relative to this directory. @@ -98,3 +99,13 @@ def run_apidoc(_): def setup(app): app.connect("builder-inited", run_apidoc) + + +intersphinx_mapping = dict( + numpy=("https://numpy.org/doc/stable/", None), + numba=("https://numba.readthedocs.io/en/stable/", None), + unyt=("https://unyt.readthedocs.io/en/stable/", None), + scipy=("https://docs.scipy.org/doc/scipy/", None), + swiftgalaxy=("https://swiftsimio.readthedocs.io/en/latest/", None), + velociraptor=("https://velociraptor-python.readthedocs.io/en/latest/", None), +) diff --git a/docs/source/cosmo_array/index.rst b/docs/source/cosmo_array/index.rst new file mode 100644 index 00000000..44ef2fd3 --- /dev/null +++ b/docs/source/cosmo_array/index.rst @@ -0,0 +1,91 @@ +The ``cosmo_array`` +=================== + +:mod:`swiftsimio` uses a customized class based on the :class:`~unyt.array.unyt_array` +to store data arrays. The :class:`~swiftsimio.objects.cosmo_array` has all of the same +functionality as the :class:`~unyt.array.unyt_array`, but also adds information about +data transformation between physical and comoving coordinates, and descriptive metadata. + +For instance, should you ever need to know what a dataset represents, you can +ask for a description by accessing the ``name`` attribute: + +.. code-block:: python + + print(rho_gas.name) + +which will output ``Co-moving mass densities of the particles``. + +The cosmology information is stored in three attributes: + + + ``comoving`` + + ``cosmo_factor`` + + ``valid_transform`` + +The ``comoving`` attribute specifies whether the array is a physical (``True``) or +comoving (``False``) quantity, while the ``cosmo_factor`` stores the expression needed +to convert back and forth between comoving and physical quantities and the value of +the scale factor. The conversion factors can be accessed like this: + +.. code-block:: python + + # Conversion factor to make the densities a physical quantity + print(rho_gas.cosmo_factor.a_factor) + physical_rho_gas = rho_gas.cosmo_factor.a_factor * rho_gas + + # Symbolic scale-factor expression + print(rho_gas.cosmo_factor.expr) + +which will output ``132651.002785671`` and ``a**(-3.0)``. Converting an array to/from physical/comoving +is done with the :meth:`~swiftsimio.objects.cosmo_array.to_physical`, :meth:`~swiftsimio.objects.cosmo_array.to_comoving`, :meth:`~swiftsimio.objects.cosmo_array.convert_to_physical` and :meth:`~swiftsimio.objects.cosmo_array.to_comoving` methods, for instance: + +.. code-block:: python + + physical_rho_gas = rho_gas.to_physical() + + # Convert in-place + rho_gas.convert_to_physical() + +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. + +.. note:: + + :class:`~swiftsimio.objects.cosmo_array` and the related :class:`~swiftsimio.objects.cosmo_quantity` are now intended to support all :mod:`numpy` functions, propagating units and cosmology information correctly through mathematical operations. Try making a histogram with weights and ``density=True`` with :func:`numpy.histogram`! There are a large number of functions and a very large number of possible parameter combinations, so some corner cases may have been missed in development. Please report any errors or unexpected results using github issues or other channels so that they can be fixed. Currently :mod:`scipy` functions are not supported (although some might "just work"). Requests to support specific functions can be accommodated. + +To make the most of the utility offered by the :class:`~swiftsimio.objects.cosmo_array` class, it is helpful to know how to create your own. A good template for this looks like: + +.. code-block:: python + + import unyt as u + from swiftsimio.objects import cosmo_array, cosmo_factor + + # suppose the scale factor is 0.5 and it scales as a**1, then: + my_cosmo_array = cosmo_array( + [1, 2, 3], + u.Mpc, + comoving=True, + scale_factor=0.5, # a=0.5, i.e. z=1 + scale_exponent=1, # distances scale as a**1, so the scale exponent is 1 + ) + # consider getting the scale factor from metadata when applicable, i.e. replace: + # scale_factor=0.5 + # with: + # scale_factor=data.metadata.a + +There is also a very similar :class:`~swiftsimio.objects.cosmo_quantity` class designed for scalar values, +analogous to the :class:`~unyt.array.unyt_quantity`. You may encounter this being returned by :mod:`numpy` functions. Cosmology-aware scalar values can be initialized similarly: + +.. code-block:: python + + import unyt as u + from swiftsimio.objects import cosmo_quantity, cosmo_factor + + my_cosmo_quantity = cosmo_quantity( + 2, + u.Mpc, + comoving=False, + scale_factor=0.5, + cosmo_factor=1, + ) + diff --git a/docs/source/getting_started/index.rst b/docs/source/getting_started/index.rst index 9715fdaa..a4366936 100644 --- a/docs/source/getting_started/index.rst +++ b/docs/source/getting_started/index.rst @@ -109,8 +109,8 @@ In the above it's important to note the following: + Only the density and temperatures (corresponding to the ``PartType0/Densities`` and ``PartType0/Temperatures``) datasets are read in. + That data is only read in once the - :meth:`swiftsimio.objects.cosmo_array.convert_to_cgs` method is called. -+ :meth:`swiftsimio.objects.cosmo_array.convert_to_cgs` converts data in-place; + :meth:`~swiftsimio.objects.cosmo_array.convert_to_cgs` method is called. ++ :meth:`~swiftsimio.objects.cosmo_array.convert_to_cgs` converts data in-place; i.e. it returns `None`. + The data is cached and not re-read in when ``plt.scatter`` is called. diff --git a/docs/source/index.rst b/docs/source/index.rst index f7ecc438..9470b8ea 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -17,8 +17,10 @@ snapshots to enable partial reading. getting_started/index loading_data/index + cosmo_array/index masking/index visualisation/index + soap/index velociraptor/index creating_initial_conditions/index statistics/index diff --git a/docs/source/loading_data/index.rst b/docs/source/loading_data/index.rst index e62f7b30..2a901069 100644 --- a/docs/source/loading_data/index.rst +++ b/docs/source/loading_data/index.rst @@ -247,47 +247,6 @@ Then, to access individual columns (in this case element abundances): data.gas.element_mass_fractions.silicon -Non-unyt properties -------------------- - -Each data array has some custom properties that are not present within the base -:obj:`unyt.unyt_array` class. We create our own version of this in -:obj:`swiftsimio.objects.cosmo_array`, which allows each dataset to contain -its own cosmology and name properties. - -For instance, should you ever need to know what a dataset represents, you can -ask for a description: - -.. code-block:: python - - print(rho_gas.name) - -which will output ``Co-moving mass densities of the particles``. They include -scale-factor information, too, through the ``cosmo_factor`` object, - -.. code-block:: python - - # Conversion factor to make the densities a physical quantity - print(rho_gas.cosmo_factor.a_factor) - physical_rho_gas = rho_gas.cosmo_factor.a_factor * rho_gas - - # Symbolic scale-factor expression - print(rho_gas.cosmo_factor.expr) - -which will output ``132651.002785671`` and ``a**(-3.0)``. This is an easy way -to convert your co-moving values to physical ones. - -An even easier way to convert your properties to physical is to use the -built-in ``to_physical`` and ``convert_to_physical`` methods, as follows: - -.. code-block:: python - - physical_rho_gas = rho_gas.to_physical() - - # Convert in-place - rho_gas.convert_to_physical() - - User-defined particle types --------------------------- @@ -310,38 +269,3 @@ in SWIFT will be automatically read. "extra_test.hdf5", ) - -Halo Catalogues ---------------- - -SWIFT-compatible halo catalogues, such as those written with SOAP, can be -loaded entirely transparently with ``swiftsimio``. It is generally possible -to use all of the functionality (masking, visualisation, etc.) that is used -with snapshots with these files, assuming the files conform to the -correct metadata standard. - -An example SOAP file is available at -``http://virgodb.cosma.dur.ac.uk/swift-webstorage/IOExamples/soap_example.hdf5`` - -You can load SOAP files as follows: - -.. code-block:: python - - from swiftsimio import load - - catalogue = load("soap_example.hdf5") - - print(catalogue.spherical_overdensity_200_mean.total_mass) - - # >>> [ 591. 328.5 361. 553. 530. 507. 795. - # 574. 489.5 233.75 0. 1406. 367.5 2308. - # ... - # 0. 534. 0. 191.75 1450. 600. 290. ] 10000000000.0*Msun (Physical) - -What's going on here? Under the hood, ``swiftsimio`` has a discrimination function -between different metadata types, based upon a property stored in the HDF5 file, -``Header/OutputType``. If this is set to ``FullVolume``, we have a snapshot, -and use the :obj:`swiftsimio.metadata.objects.SWIFTSnapshotMetadata` -class. If it is ``SOAP``, we use -:obj:`swiftsimio.metadata.objects.SWIFTSOAPMetadata`, which instructs -``swiftsimio`` to read slightly different properties from the HDF5 file. diff --git a/docs/source/soap/index.rst b/docs/source/soap/index.rst new file mode 100644 index 00000000..c143c9ec --- /dev/null +++ b/docs/source/soap/index.rst @@ -0,0 +1,41 @@ +Halo Catalogues & SOAP integration +================================== + +SWIFT-compatible halo catalogues, such as those written with SOAP, can be +loaded entirely transparently with ``swiftsimio``. It is generally possible +to use all of the functionality (masking, visualisation, etc.) that is used +with snapshots with these files, assuming the files conform to the +correct metadata standard. + +An example SOAP file is available at +``http://virgodb.cosma.dur.ac.uk/swift-webstorage/IOExamples/soap_example.hdf5`` + +You can load SOAP files as follows: + +.. code-block:: python + + from swiftsimio import load + + catalogue = load("soap_example.hdf5") + + print(catalogue.spherical_overdensity_200_mean.total_mass) + + # >>> [ 591. 328.5 361. 553. 530. 507. 795. + # 574. 489.5 233.75 0. 1406. 367.5 2308. + # ... + # 0. 534. 0. 191.75 1450. 600. 290. ] 10000000000.0*Msun (Physical) + +What's going on here? Under the hood, ``swiftsimio`` has a discrimination function +between different metadata types, based upon a property stored in the HDF5 file, +``Header/OutputType``. If this is set to ``FullVolume``, we have a snapshot, +and use the :obj:`swiftsimio.metadata.objects.SWIFTSnapshotMetadata` +class. If it is ``SOAP``, we use +:obj:`swiftsimio.metadata.objects.SWIFTSOAPMetadata`, which instructs +``swiftsimio`` to read slightly different properties from the HDF5 file. + +swiftgalaxy +----------- + +The :mod:`swiftgalaxy` companion package to :mod:`swiftsimio` offers further integration with halo catalogues in SOAP, Caesar and Velociraptor formats (so far). It greatly simplifies efficient loading of particles belonging to an object from a catalogue, and additional tools that are useful when working with a galaxy or other localized collection of particles. Refer to the `swiftgalaxy documentation`_ for details. + +.. _swiftgalaxy documentation: https://swiftgalaxy.readthedocs.io/en/latest/ diff --git a/docs/source/visualisation/index.rst b/docs/source/visualisation/index.rst index 2b6ad60c..036e7be2 100644 --- a/docs/source/visualisation/index.rst +++ b/docs/source/visualisation/index.rst @@ -1,12 +1,12 @@ Visualisation ============= -:mod:`swiftsimio` provides visualisation routines accelerated with the -:mod:`numba` module. They work without this module, but we strongly recommend -installing it for the best performance (1000x+ speedups). These are provided -in the :mod:`swiftismio.visualisation` sub-modules. +:mod:`swiftsimio` provides visualisation routines in the +:mod:`swiftsimio.visualisation` sub-module. They are accelerated with the +:mod:`numba` module. They can work without :mod:`numba`, but we strongly recommend +installing it for the best performance (1000x+ speedups). -The three built-in rendering types (described below) have the following +The four built-in rendering types (described below) have the following common interface: .. code-block:: python @@ -30,6 +30,7 @@ additional functionality. projection slice volume_render + ray_trace power_spectra tools diff --git a/docs/source/visualisation/power_spectra.rst b/docs/source/visualisation/power_spectra.rst index 7a892356..6fee7d92 100644 --- a/docs/source/visualisation/power_spectra.rst +++ b/docs/source/visualisation/power_spectra.rst @@ -6,15 +6,15 @@ runs on-the-fly, and as such after a run has completed you may wish to create a number of more non-standard power spectra. These tools are available as part of the :mod:`swiftsimio.visualisation.power_spectrum` -package. Making a power spectrum consists of two major steps: depositing the particles +module. Making a power spectrum consists of two major steps: depositing the particles on grid(s), and then binning their fourier transform to get the one-dimensional power. -Depositing on a Grid +Depositing on a grid -------------------- Depositing your particles on a grid is performed using -:meth:`swiftsimio.visualisation.power_spectrum.render_to_deposit`. This function +:func:`swiftsimio.visualisation.power_spectrum.render_to_deposit`. This function performs a nearest-grid-point (NGP) of all particles in the provided particle dataset. For example: @@ -35,17 +35,17 @@ dataset. For example: The specific field being depositied can be controlled with the ``project`` keyword argument. The ``resolution``` argument gives the one-dimensional resolution of the 3D grid, so in this case you would recieve a ``512x512x512`` -grid. Note that the ``gas_mass_deposit`` is a :obj:`swiftsimio.cosmo_array`, +grid. Note that the ``gas_mass_deposit`` is a :obj:`~swiftsimio.objects.cosmo_array`, and as such includes cosmological and unit information that is used later in the process. -Generating a Power Spectrum +Generating a power spectrum --------------------------- Once you have your grid deposited, you can easily generate a power spectrum using the -:meth:`swiftsimio.visualisation.power_spectrum.deposition_to_power_spectrum` +:func:`~swiftsimio.visualisation.power_spectrum.deposition_to_power_spectrum` function. For example, using the above deposit: .. code-block:: python @@ -54,7 +54,7 @@ function. For example, using the above deposit: wavenumbers, power_spectrum, _ = deposition_to_power_spectrum( deposition=gas_mass_deposit, - box_size=data.metadata.box_size, + boxsize=data.metadata.boxsize, ) This power spectrum can then be plotted. Units are included on both the wavenumbers @@ -64,7 +64,7 @@ Wavenumbers are calculated to be at the weighted mean of the k-values in each bin, rather than representing the center of the bin. -More Complex Scenarios +More complex scenarios ---------------------- In a realistic simualted power spectrum, you will need to perform 'folding' @@ -89,38 +89,44 @@ The ``folding`` parameter is available for both ``render_to_deposit`` and ``deposition_to_power_spectrum``, but it may be easier to use the utility functions provided for automatically stitching together the folded spectra. The function -:meth:`swiftsimio.visualsation.power_spectrum.folded_depositions_to_power_spectrum` +:func:`~swiftsimio.visualsation.power_spectrum.folded_depositions_to_power_spectrum` allows you to do this easily: .. code-block:: python - from swiftsimio.visualisation.power_spectrum import folded_depositions_to_power_spectrum - import unyt - - folded_depositions = {} - - for folding in [x * 2 for x in range(5)]: - folded_depositions[folding] = render_to_deposit( - data.gas, - resolution=512, - project="masses", - parallel=True, - folding=folding, - ) - - bins, centers, power_spectrum, foldings = folded_depositions_to_power_spectrum( - depositions=folded_depositions, - box_size=data.metadata.box_size, - number_of_wavenumber_bins=128, - wavenumber_range=[1e-2 / unyt.Mpc, 1e2 / unyt.Mpc], - log_wavenumber_bins=True, - workers=4, - minimal_sample_modes=8192, - cutoff_above_wavenumber_fraction=0.75, - shot_noise_norm=len(gas_mass_deposit), - - ) - + import unyt as u + from swiftsimio.visualisation.power_spectrum import folded_depositions_to_power_spectrum + from swiftsimio.objects import cosmo_array + + folded_depositions = {} + + for folding in [x * 2 for x in range(5)]: + folded_depositions[folding] = render_to_deposit( + data.gas, + resolution=512, + project="masses", + parallel=True, + folding=folding, + ) + + bins, centers, power_spectrum, foldings = folded_depositions_to_power_spectrum( + depositions=folded_depositions, + boxsize=data.metadata.boxsize, + number_of_wavenumber_bins=128, + wavenumber_range=cosmo_array( + [1e-2, 1e2], + u.Mpc**-1, + comoving=True, + scale_factor=data.metadata.a, + scale_exponent=-1, + ), + log_wavenumber_bins=True, + workers=4, + minimal_sample_modes=8192, + cutoff_above_wavenumber_fraction=0.75, + shot_noise_norm=len(gas_mass_deposit), + ) + The 'used' foldings of the power spectrum are shown in the ``foldings`` return vaule, which is an array containing the folding that was used for each given bin. This is useful for debugging and @@ -129,19 +135,19 @@ visualisation. There are a few crucial parameters to this function: 1. ``workers`` is the number of threads to use for the calculation of - the fourier transforms. + the Fourier transforms. 2. ``minimal_sample_modes`` is the minimum number of modes that must be - present in a bin for it to be included in the final power spectrum. - Generally for a big simulation you want to set this to around 10'000, - and this number is ignored for the lowest wavenumber bin. + present in a bin for it to be included in the final power spectrum. + Generally for a big simulation you want to set this to around 10,000, + and this number is ignored for the lowest wavenumber bin. 3. ``cutoff_above_wavenumber_fraction`` is the fraction of the individual fold's (as represented by the FFT itself) maximally sampled wavenumber. Ignored for the last fold, and we always cap the maximal - wavenumber to the nyquist frequency. + wavenumber to the Nyquist frequency. 4. ``shot_noise_norm`` is the number of particles in the simulation - that contribute to the power spectrum. This is used to normalise - the power spectrum to the shot noise level. This is very - important in this case because of the use of NGP deposition. + that contribute to the power spectrum. This is used to normalise + the power spectrum to the shot noise level. This is very + important in this case because of the use of NGP deposition. Foldings are stitched using a simple method where the 'better sampled' -foldings are used preferentially, up to the cutoff value. \ No newline at end of file +foldings are used preferentially, up to the cutoff value. diff --git a/docs/source/visualisation/projection.rst b/docs/source/visualisation/projection.rst index e92428e1..e33ef60f 100644 --- a/docs/source/visualisation/projection.rst +++ b/docs/source/visualisation/projection.rst @@ -15,7 +15,7 @@ with :math:`\tilde{A}_i` the smoothed quantity in pixel :math:`i`, and Here we use the Wendland-C2 kernel. The primary function here is -:meth:`swiftsimio.visualisation.projection.project_gas`, which allows you to +:func:`swiftsimio.visualisation.projection.project_gas`, which allows you to create a gas projection of any field. See the example below. Example @@ -29,7 +29,7 @@ Example data = load("cosmo_volume_example.hdf5") # This creates a grid that has units msun / Mpc^2, and can be transformed like - # any other unyt quantity + # any other cosmo_array mass_map = project_gas( data, resolution=1024, @@ -52,7 +52,7 @@ Example This basic demonstration creates a mass surface density map. To create, for example, a projected temperature map, we need to remove the -surface density dependence (i.e. :meth:`project_gas` returns a surface +surface density dependence (i.e. :func:`~swiftsimio.visualisation.projection.project_gas` returns a surface temperature in units of K / kpc^2 and we just want K) by dividing out by this: @@ -114,11 +114,11 @@ backends are as follows: + ``fast``: The default backend - this is extremely fast, and provides very basic smoothing, with a return type of single precision floating point numbers. + ``histogram``: This backend provides zero smoothing, and acts in a similar way - to the ``np.hist2d`` function but with the same arguments as ``scatter``. -+ ``reference``: The same backend as ``fast`` but with two distinguishing features; + to the :func:`~numpy.histogram2d` function but with the same arguments as ``scatter``. ++ ``reference``: The same backend as ``fast`` but with two distinguishing features: all calculations are performed in double precision, and it will return early with a warning message if there are not enough pixels to fully resolve each kernel. - Regular users should not use this mode. + Intended for developer usage, regular users should not use this mode. + ``renormalised``: The same as ``fast``, but each kernel is evaluated twice and renormalised to ensure mass conservation within floating point precision. Returns single precision arrays. @@ -166,7 +166,7 @@ All visualisation functions by default assume a periodic box. Rather than simply projecting each individual particle once, four additional periodic copies of each particle are also projected. Most copies will project outside the valid pixel range, but the copies that do not ensure that pixels close to the edge -receive all necessary contributions. Thanks to Numba optimisations, the overhead +receive all necessary contributions. Thanks to :mod:`numba` optimisations, the overhead of these additional copies is relatively small. There are some caveats with this approach. If you try to visualise a subset of @@ -183,8 +183,8 @@ Rotations Sometimes you will need to visualise a galaxy from a different perspective. The :mod:`swiftsimio.visualisation.rotation` sub-module provides routines to generate rotation matrices corresponding to vectors, which can then be -provided to the ``rotation_matrix`` argument of :meth:`project_gas` (and -:meth:`project_gas_pixel_grid`). You will also need to supply the +provided to the ``rotation_matrix`` argument of :func:`~swiftsimio.visualisation.projection.project_gas` (and +:func:`~swiftsimio.visualisation.projection.project_gas_pixel_grid`). You will also need to supply the ``rotation_center`` argument, as the rotation takes place around this given point. The example code below loads a snapshot, and a halo catalogue, and creates an edge-on and face-on projection using the integration in @@ -193,87 +193,96 @@ is shown in the ``velociraptor`` section. .. code-block:: python - from swiftsimio import load, mask + from swiftsimio import load, mask, cosmo_array from velociraptor import load as load_catalogue from swiftsimio.visualisation.rotation import rotation_matrix_from_vector - from swiftsimio.visualisation.projection import project_gas_pixel_grid - + from swiftsimio.visualisation.projection import project_gas + import unyt import numpy as np - import matplotlib.pyplot as plt - from matplotlib.colors import LogNorm - + # Radius around which to load data, we will visualise half of this size = 1000 * unyt.kpc - + snapshot_filename = "cosmo_volume_example.hdf5" catalogue_filename = "cosmo_volume_example.properties" - + catalogue = load_catalogue(catalogue_filename) - + # Which halo should we visualise? halo = 0 - + x = catalogue.positions.xcmbp[halo] y = catalogue.positions.ycmbp[halo] z = catalogue.positions.zcmbp[halo] - + lx = catalogue.angular_momentum.lx[halo] ly = catalogue.angular_momentum.ly[halo] lz = catalogue.angular_momentum.lz[halo] - + # 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 /= np.linalg.norm(angular_momentum_vector) - - face_on_rotation_matrix = rotation_matrix_from_vector( - angular_momentum_vector + + face_on_rotation_matrix = rotation_matrix_from_vector(angular_momentum_vector) + edge_on_rotation_matrix = rotation_matrix_from_vector(angular_momentum_vector, axis="y") + + data_mask = mask(snapshot_filename) + region = cosmo_array( + [ + [x - size, x + size], + [y - size, y + size], + [z - size, z + size], + ], + x.units, + comoving=True, + scale_factor=data_mask.metadata.a, + scale_exponent=1, ) - edge_on_rotation_matrix = rotation_matrix_from_vector( - angular_momentum_vector, - axis="y" + + visualise_region = cosmo_array( + [ + x - 0.5 * size, + x + 0.5 * size, + y - 0.5 * size, + y + 0.5 * size, + ], + comoving=True, + scale_factor=data_mask.metadata.a, + scale_exponent=1, ) - - region = [ - [x - size, x + size], - [y - size, y + size], - [z - size, z + size], - ] - - visualise_region = [ - x - 0.5 * size, x + 0.5 * size, - y - 0.5 * size, y + 0.5 * size, - ] - - data_mask = mask(snapshot_filename) + data_mask.constrain_spatial(region) data = load(snapshot_filename, mask=data_mask) - + # Use project_gas_pixel_grid to generate projected images - + common_arguments = dict( data=data, resolution=512, parallel=True, region=visualise_region, - periodic=False, # disable periodic boundaries when using rotations + periodic=False, # disable periodic boundaries when using rotations ) - - un_rotated = project_gas_pixel_grid(**common_arguments) - - face_on = project_gas_pixel_grid( - **common_arguments, - rotation_center=unyt.unyt_array([x, y, z]), - rotation_matrix=face_on_rotation_matrix, + + un_rotated = project_gas(**common_arguments) + + rotation_center = cosmo_array( + [x, y, z], comoving=True, scale_factor=data_mask.metadata.a, scale_exponent=1 ) - - edge_on = project_gas_pixel_grid( - **common_arguments, - rotation_center=unyt.unyt_array([x, y, z]), - rotation_matrix=edge_on_rotation_matrix, + face_on = project_gas( + **common_arguments, + rotation_center=rotation_center, + rotation_matrix=face_on_rotation_matrix, ) - + + edge_on = project_gas( + **common_arguments, + rotation_center=rotation_center, + rotation_matrix=edge_on_rotation_matrix, + ) + Using this with the provided example data will just show blobs due to its low resolution nature. Using one of the EAGLE volumes (``examples/EAGLE_ICs``) will produce much nicer galaxies, but that data is too large to provide as an example in this tutorial. @@ -287,15 +296,12 @@ Other particle types -------------------- Other particle types are able to be visualised through the use of the -:meth:`swiftsimio.visualisation.projection.project_pixel_grid` function. This -does not attach correct symbolic units, so you will have to work those out -yourself, but it does perform the smoothing. We aim to introduce the feature -of correctly applied units to these projections soon. +:func:`swiftsimio.visualisation.projection.project_pixel_grid` function. To use this feature for particle types that do not have smoothing lengths, you will need to generate them, as in the example below where we create a mass density map for dark matter. We provide a utility to do this through -:meth:`swiftsimio.visualisation.smoothing_length.generate_smoothing_lengths`. +:func:`~swiftsimio.visualisation.smoothing_length.generate.generate_smoothing_lengths`. .. code-block:: python @@ -320,7 +326,6 @@ mass density map for dark matter. We provide a utility to do this through # Note here that we pass in the dark matter dataset not the whole # data object, to specify what particle type we wish to visualise data=data.dark_matter, - boxsize=data.metadata.boxsize, resolution=1024, project="masses", parallel=True, @@ -348,15 +353,15 @@ smoothing lengths, and smoothed quantities, to generate a pixel grid that represents the smoothed version of the data. This API is available through -:meth:`swiftsimio.visualisation.projection.scatter` and -:meth:`swiftsimio.visualisation.projection.scatter_parallel` for the parallel +:obj:`swiftsimio.visualisation.projection_backends.backends["scatter"]` and +:obj:`swiftsimio.visualisation.projection_backends.backends_parallel["scatter"]` for the parallel version. The parallel version uses significantly more memory as it allocates a thread-local image array for each thread, summing them in the end. Here we will only describe the ``scatter`` variant, but they behave in the exact same way. By default this uses the "fast" backend. To use the others, you can select them manually from the module, or by using the ``backends`` and ``backends_parallel`` -dictionaries in :mod:`swiftsimio.visualisation.projection`. +dictionaries in :mod:`swiftsimio.visualisation.projection_backends`. To use this function, you will need: @@ -374,7 +379,9 @@ The key here is that only particles in the domain [0, 1] in x, and [0, 1] in y will be visible in the image. You may have particles outside of this range; they will not crash the code, and may even contribute to the image if their smoothing lengths overlap with [0, 1]. You will need to re-scale your data -such that it lives within this range. Then you may use the function as follows: +such that it lives within this range. You should also pass raw numpy arrays (not +:class:`~swiftsimio.objects.cosmo_array` or :class:`~unyt.array.unyt_array`, the +inputs are dimensionless). Then you may use the function as follows: .. code-block:: python @@ -383,12 +390,12 @@ such that it lives within this range. Then you may use the function as follows: # Using the variable names from above out = scatter(x=x, y=y, h=h, m=m, res=res) -``out`` will be a 2D :mod:`numpy` grid of shape ``[res, res]``. You will need +``out`` will be a 2D :class:`~numpy.ndarray` grid of shape ``[res, res]``. You will need to re-scale this back to your original dimensions to get it in the correct units, and do not forget that it now represents the smoothed quantity per surface area. If the optional arguments ``box_x`` and ``box_y`` are provided, they should contain the simulation box size in the same re-scaled coordinates as ``x`` and ``y``. The projection backend will then correctly apply periodic boundary -wrapping. If ``box_x`` and ``box_y`` are not provided or set to 0, no +wrapping. If ``box_x`` and ``box_y`` are not provided or set to ``0``, no periodic boundaries are applied. diff --git a/docs/source/visualisation/ray_trace.rst b/docs/source/visualisation/ray_trace.rst new file mode 100644 index 00000000..9830b113 --- /dev/null +++ b/docs/source/visualisation/ray_trace.rst @@ -0,0 +1,4 @@ +Ray tracing +=========== + +Documentation to be completed... diff --git a/docs/source/visualisation/slice.rst b/docs/source/visualisation/slice.rst index 7fb7ca8c..4345a396 100644 --- a/docs/source/visualisation/slice.rst +++ b/docs/source/visualisation/slice.rst @@ -5,7 +5,7 @@ The :mod:`swiftsimio.visualisation.slice` sub-module provides an interface to render SWIFT data onto a slice. This takes your 3D data and finds the 3D density at fixed z-position, slicing through the box. -The default :code:`"sph"` backend effectively solves the equation: +The default ``"sph"`` backend effectively solves the equation: :math:`\tilde{A}_i = \sum_j A_j W_{ij, 3D}` @@ -19,7 +19,7 @@ nearest-neighbour interpolation to compute the densities at each pixel. This backend is more suited for use with moving-mesh hydrodynamics schemes. The primary function here is -:meth:`swiftsimio.visualisation.slice.slice_gas`, which allows you to +:func:`swiftsimio.visualisation.slice.slice_gas`, which allows you to create a gas slice of any field. See the example below. Example @@ -58,8 +58,9 @@ Example This basic demonstration creates a mass density map. To create, for example, a projected temperature map, we need to remove the -density dependence (i.e. :meth:`slice_gas` returns a volumetric temperature -in units of K / kpc^3 and we just want K) by dividing out by this: +density dependence (i.e. :func:`~swiftsimio.visualisation.slice.slice_gas` +returns a volumetric temperature in units of K / kpc^3 and we just want K) +by dividing out by this: .. code-block:: python @@ -121,7 +122,7 @@ All visualisation functions by default assume a periodic box. Rather than simply summing each individual particle once, eight additional periodic copies of each particle are also accounted for. Most copies will contribute outside the valid pixel range, but the copies that do not ensure that pixels close to the -edge receive all necessary contributions. Thanks to Numba optimisations, the +edge receive all necessary contributions. Thanks to :mod:`numba` optimisations, the overhead of these additional copies is relatively small. There are some caveats with this approach. If you try to visualise a subset of @@ -139,11 +140,12 @@ Rotations Rotations of the box prior to slicing are provided in a similar fashion to the :mod:`swiftsimio.visualisation.projection` sub-module, by using the :mod:`swiftsimio.visualisation.rotation` sub-module. To rotate the perspective -prior to slicing a ``rotation_center`` argument in :meth:`slice_gas` needs +prior to slicing a ``rotation_center`` argument in +:func:`~swiftsimio.visualisation.slice.slice_gas` needs to be provided, specifying the point around which the rotation takes place. The angle of rotation is specified with a matrix, supplied by ``rotation_matrix`` -in :meth:`slice_gas`. The rotation matrix may be computed with -:meth:`rotation_matrix_from_vector`. This will result in the perspective being +in :func:`~swiftsimio.visualisation.slice.slice_gas`. The rotation matrix may be computed with +:func:`~swiftsimio.visualisation.rotation.rotation_matrix_from_vector`. This will result in the perspective being rotated to be along the provided vector. This approach to rotations applied to the above example is shown below. @@ -209,8 +211,8 @@ smoothing lengths, and smoothed quantities, to generate a pixel grid that represents the smoothed, sliced, version of the data. This API is available through -:meth:`swiftsimio.visualisation.slice.slice_scatter` and -:meth:`swiftsimio.visualisation.slice.slice_scatter_parallel` for the parallel +:func:`swiftsimio.visualisation.slice_backends.backends["sph"]` and +:func:`swiftsimio.visualisation.slice_backends.backends_parallel["sph"]` for the parallel version. The parallel version uses significantly more memory as it allocates a thread-local image array for each thread, summing them in the end. Here we will only describe the ``scatter`` variant, but they behave in the exact same way. @@ -235,7 +237,9 @@ not crash the code, and may even contribute to the image if their smoothing lengths overlap with [0, 1]. You will need to re-scale your data such that it lives within this range. Smoothing lengths and z coordinates need to be re-scaled in the same way (using the same scaling factor), but z coordinates do -not need to lie in the domain [0, 1]. Then you may use the function as follows: +not need to lie in the domain [0, 1]. You should provide inputs as raw numpy arrays +(not :class:`~swiftsimio.objects.cosmo_array` or :class:`~unyt.array.unyt_array`). +Then you may use the function as follows: .. code-block:: python diff --git a/docs/source/visualisation/tools.rst b/docs/source/visualisation/tools.rst index d4244927..6f92a8d4 100644 --- a/docs/source/visualisation/tools.rst +++ b/docs/source/visualisation/tools.rst @@ -9,8 +9,8 @@ Tools The :mod:`swiftsimio.visualisation.tools.cmaps` module includes three objects that can be used to deploy two dimensional colour maps. The first, -:class:`swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2D`, and second -:class:`swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2DHSV`, allow +:class:`~swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2D`, and second +:class:`~swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2DHSV`, allow you to generate new color maps from sets of colors and coordinates. .. code-block:: python @@ -22,7 +22,7 @@ you to generate new color maps from sets of colors and coordinates. ) This generates a color map that is a quasi-linear interpolation between all -of the points. The map can be displayed using the ``plot`` method, +of the points. The map can be displayed using the :func:`~matplotlib.pyplot.plot` function, .. code-block:: python @@ -30,7 +30,7 @@ of the points. The map can be displayed using the ``plot`` method, bower.plot(ax) -Which generates: +which generates: .. image:: bower_cmap.png diff --git a/docs/source/visualisation/volume_render.rst b/docs/source/visualisation/volume_render.rst index 2c3d8ee5..1a223581 100644 --- a/docs/source/visualisation/volume_render.rst +++ b/docs/source/visualisation/volume_render.rst @@ -15,7 +15,7 @@ with :math:`\tilde{A}_i` the smoothed quantity in pixel :math:`i`, and Here we use the Wendland-C2 kernel. The primary function here is -:meth:`swiftsimio.visualisation.volume_render.render_gas`, which allows you +:func:`swiftsimio.visualisation.volume_render.render_gas`, which allows you to create a gas density grid of any field, see the example below. Example @@ -41,7 +41,7 @@ Example This basic demonstration creates a mass density cube. To create, for example, a projected temperature cube, we need to remove the -density dependence (i.e. :meth:`render_gas` returns a volumetric +density dependence (i.e. :func:`~swiftsimio.visualisation.volume_render.render_gas` returns a volumetric temperature in units of K / kpc^3 and we just want K) by dividing out by this: @@ -89,7 +89,7 @@ All visualisation functions by default assume a periodic box. Rather than simply summing each individual particle once, eight additional periodic copies of each particle are also taken into account. Most copies will contribute outside the valid voxel range, but the copies that do not ensure that voxels -close to the edge receive all necessary contributions. Thanks to Numba +close to the edge receive all necessary contributions. Thanks to :mod:`numba` optimisations, the overhead of these additional copies is relatively small. There are some caveats with this approach. If you try to visualise a subset of @@ -106,11 +106,11 @@ Rotations Rotations of the box prior to volume rendering are provided in a similar fashion to the :mod:`swiftsimio.visualisation.projection` sub-module, by using the :mod:`swiftsimio.visualisation.rotation` sub-module. To rotate the perspective -prior to slicing a ``rotation_center`` argument in :meth:`render_gas` needs +prior to slicing a ``rotation_center`` argument in :func:`~swiftsimio.visualisation.volume_render.render_gas` needs to be provided, specifying the point around which the rotation takes place. The angle of rotation is specified with a matrix, supplied by ``rotation_matrix`` -in :meth:`render_gas`. The rotation matrix may be computed with -:meth:`rotation_matrix_from_vector`. This will result in the perspective being +in :func:`~swiftsimio.visualisation.volume_render.render_gas`. The rotation matrix may be computed with +:func:`~swiftsimio.visualisation.rotation.rotation_matrix_from_vector`. This will result in the perspective being rotated to be along the provided vector. This approach to rotations applied to the above example is shown below. @@ -160,8 +160,8 @@ Rendering --------- We provide a volume rendering function that can be used to make images highlighting -specific density contours. The notable function here is -:meth:``swiftsimio.visualisation.volume_render.visualise_render``. This takes +specific density contours. The key function here is +:func:`swiftsimio.visualisation.volume_render.visualise_render`. This takes in your volume rendering, along with a colour map and centers, to create these highlights. The example below shows how to use this. @@ -175,7 +175,7 @@ these highlights. The example below shows how to use this. from swiftsimio.visualisation import volume_render # Load the data - data = load("test_data/eagle_6.hdf5") + data = load("eagle_6.hdf5") # Rough location of an interesting galaxy in the volume. region = [ @@ -238,7 +238,7 @@ Here we can see the quick view of this image. It's just a regular density projec plt.savefig("volume_render_options.png") -This function :meth:`swiftsimio.visualisation.volume_render.visualise_render_options` allows +This function :func:`swiftsimio.visualisation.volume_render.visualise_render_options` allows you to see what densities your rendering is picking out: .. image:: volume_render_options.png @@ -280,8 +280,8 @@ smoothing lengths, and smoothed quantities, to generate a pixel grid that represents the smoothed, volume rendered, version of the data. This API is available through -:meth:`swiftsimio.visualisation.volume_render.scatter` and -:meth:`swiftsimio.visualisation.volume_render.scatter_parallel` for the parallel +:func:`swiftsimio.visualisation.volume_render_backends.backends["scatter"]` and +:func:`swiftsimio.visualisation.volume_render_backends.backends_parallel["scatter"]` for the parallel version. The parallel version uses significantly more memory as it allocates a thread-local image array for each thread, summing them in the end. Here we will only describe the ``scatter`` variant, but they behave in the exact same way. @@ -303,8 +303,9 @@ The key here is that only particles in the domain [0, 1] in x, [0, 1] in y, and [0, 1] in z. will be visible in the cube. You may have particles outside of this range; they will not crash the code, and may even contribute to the image if their smoothing lengths overlap with [0, 1]. You will need to -re-scale your data such that it lives within this range. Then you may use the -function as follows: +re-scale your data such that it lives within this range. You should pass in +raw numpy array (not :class:`~swiftsimio.objects.cosmo_array` or :class:`~unyt.array.unyt_array`). +Then you may use the function as follows: .. code-block:: python @@ -313,7 +314,7 @@ function as follows: # Using the variable names from above out = scatter(x=x, y=y, z=z, h=h, m=m, res=res) -``out`` will be a 3D :mod:`numpy` grid of shape ``[res, res, res]``. You will +``out`` will be a 3D :class:`~numpy.ndarray` grid of shape ``[res, res, res]``. You will need to re-scale this back to your original dimensions to get it in the correct units, and do not forget that it now represents the smoothed quantity per volume. @@ -322,4 +323,4 @@ If the optional arguments ``box_x``, ``box_y`` and ``box_z`` are provided, they should contain the simulation box size in the same re-scaled coordinates as ``x``, ``y`` and ``z``. The rendering function will then correctly apply periodic boundary wrapping. If ``box_x``, ``box_y`` and ``box_z`` are not -provided or set to 0, no periodic boundaries are applied +provided or set to 0, no periodic boundaries are applied. diff --git a/pyproject.toml b/pyproject.toml index 83482749..d357269c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,8 @@ packages = [ "swiftsimio.visualisation", "swiftsimio.visualisation.projection_backends", "swiftsimio.visualisation.slice_backends", + "swiftsimio.visualisation.volume_render_backends", + "swiftsimio.visualisation.ray_trace_backends", "swiftsimio.visualisation.tools", "swiftsimio.visualisation.smoothing_length", ] @@ -37,9 +39,9 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - "numpy", + "numpy>=2.1.0", "h5py", - "unyt>=3.0.2", + "unyt>=3.0.4", "numba>=0.50.0", ] diff --git a/swiftsimio/__init__.py b/swiftsimio/__init__.py index f9fa573e..90efe849 100644 --- a/swiftsimio/__init__.py +++ b/swiftsimio/__init__.py @@ -8,6 +8,7 @@ import swiftsimio.metadata as metadata import swiftsimio.accelerated as accelerated import swiftsimio.objects as objects +from swiftsimio.objects import cosmo_array, cosmo_quantity import swiftsimio.visualisation as visualisation import swiftsimio.units as units import swiftsimio.subset_writer as subset_writer diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py new file mode 100644 index 00000000..c8510576 --- /dev/null +++ b/swiftsimio/_array_functions.py @@ -0,0 +1,2172 @@ +""" +Overloaded implementations of unyt and numpy functions to correctly handle +:class:`~swiftsimio.objects.cosmo_array` input. + +This module also defines wrappers and helper functions to facilitate overloading +functions and handling the processing of our custom array attributes. + +Nothing in this module is intended to be user-facing, but the helpers and wrappers +are documented to assist in maintenance and development of swiftsimio. +""" + +import warnings +from functools import reduce +import numpy as np +from typing import Callable, Tuple, Optional +import unyt +from unyt import unyt_quantity, unyt_array +from swiftsimio import objects +from unyt._array_functions import ( + dot as unyt_dot, + vdot as unyt_vdot, + inner as unyt_inner, + outer as unyt_outer, + kron as unyt_kron, + histogram_bin_edges as unyt_histogram_bin_edges, + linalg_inv as unyt_linalg_inv, + linalg_tensorinv as unyt_linalg_tensorinv, + linalg_pinv as unyt_linalg_pinv, + linalg_svd as unyt_linalg_svd, + histogram as unyt_histogram, + histogram2d as unyt_histogram2d, + histogramdd as unyt_histogramdd, + concatenate as unyt_concatenate, + intersect1d as unyt_intersect1d, + union1d as unyt_union1d, + norm as unyt_linalg_norm, # not linalg_norm, doesn't follow usual pattern + vstack as unyt_vstack, + hstack as unyt_hstack, + dstack as unyt_dstack, + column_stack as unyt_column_stack, + stack as unyt_stack, + around as unyt_around, + block as unyt_block, + fft_fft as unyt_fft_fft, + fft_fft2 as unyt_fft_fft2, + fft_fftn as unyt_fft_fftn, + fft_hfft as unyt_fft_hfft, + fft_rfft as unyt_fft_rfft, + fft_rfft2 as unyt_fft_rfft2, + fft_rfftn as unyt_fft_rfftn, + fft_ifft as unyt_fft_ifft, + fft_ifft2 as unyt_fft_ifft2, + fft_ifftn as unyt_fft_ifftn, + fft_ihfft as unyt_fft_ihfft, + fft_irfft as unyt_fft_irfft, + fft_irfft2 as unyt_fft_irfft2, + fft_irfftn as unyt_fft_irfftn, + fft_fftshift as unyt_fft_fftshift, + fft_ifftshift as unyt_fft_ifftshift, + sort_complex as unyt_sort_complex, + isclose as unyt_isclose, + allclose as unyt_allclose, + array2string as unyt_array2string, + cross as unyt_cross, + array_equal as unyt_array_equal, + array_equiv as unyt_array_equiv, + linspace as unyt_linspace, + logspace as unyt_logspace, + geomspace as unyt_geomspace, + copyto as unyt_copyto, + prod as unyt_prod, + var as unyt_var, + trace as unyt_trace, + percentile as unyt_percentile, + quantile as unyt_quantile, + nanpercentile as unyt_nanpercentile, + nanquantile as unyt_nanquantile, + linalg_det as unyt_linalg_det, + diff as unyt_diff, + ediff1d as unyt_ediff1d, + ptp as unyt_ptp, + pad as unyt_pad, + choose as unyt_choose, + insert as unyt_insert, + linalg_lstsq as unyt_linalg_lstsq, + linalg_solve as unyt_linalg_solve, + linalg_tensorsolve as unyt_linalg_tensorsolve, + linalg_eig as unyt_linalg_eig, + linalg_eigh as unyt_linalg_eigh, + linalg_eigvals as unyt_linalg_eigvals, + linalg_eigvalsh as unyt_linalg_eigvalsh, + savetxt as unyt_savetxt, + fill_diagonal as unyt_fill_diagonal, + isin as unyt_isin, + place as unyt_place, + put as unyt_put, + put_along_axis as unyt_put_along_axis, + putmask as unyt_putmask, + searchsorted as unyt_searchsorted, + select as unyt_select, + setdiff1d as unyt_setdiff1d, + sinc as unyt_sinc, + clip as unyt_clip, + where as unyt_where, + triu as unyt_triu, + tril as unyt_tril, + einsum as unyt_einsum, + convolve as unyt_convolve, + correlate as unyt_correlate, + tensordot as unyt_tensordot, + unwrap as unyt_unwrap, + interp as unyt_interp, + array_repr as unyt_array_repr, + linalg_outer as unyt_linalg_outer, + trapezoid as unyt_trapezoid, + isin as unyt_in1d, + take as unyt_take, +) + +_HANDLED_FUNCTIONS = {} + +# first we define helper functions to handle repetitive operations in wrapping unyt & +# numpy functions (we will actually wrap the functions below): + + +def _copy_cosmo_array_attributes_if_present( + from_ca: object, to_ca: object, copy_units=False +) -> object: + """ + Copy :class:`~swiftsimio.objects.cosmo_array` attributes across two objects. + + Copies the ``cosmo_factor``, ``comoving``, ``valid_transform`` and ``compression`` + attributes across if both the source and destination objects are + :class:`~swiftsimio.objects.cosmo_array` instances (else returns input). + + Parameters + ---------- + from_ca : :obj:`object` + The source object. + + to_ca : :obj:`object` + The destination object. + + copy_units : bool + If ``True`` also copy ``units`` attribute (usually let :mod:`unyt` handle this). + + Returns + ------- + out : :obj:`object` + The destination object (with attributes copied if copy occurred). + """ + if not ( + isinstance(to_ca, objects.cosmo_array) + and isinstance(from_ca, objects.cosmo_array) + ): + return to_ca + if copy_units: + to_ca.units = from_ca.units + to_ca.cosmo_factor = from_ca.cosmo_factor + to_ca.comoving = from_ca.comoving + to_ca.valid_transform = from_ca.valid_transform + to_ca.compression = from_ca.compression + return to_ca + + +def _propagate_cosmo_array_attributes_to_result(func: Callable) -> Callable: + """ + Wrapper that copies :class:`~swiftsimio.objects.cosmo_array` attributes from first + input argument to first output. + + Many functions take one input (or have a first input that has a close correspondance + to the output) and one output. This helper copies the ``cosmo_factor``, ``comoving``, + ``valid_transform`` and ``compression`` attributes from the first input argument to + the output. Can be used as a decorator on functions (the first argument is then the + first argument of the function) or methods (the first argument is then ``self``). + If the output is not a :class:`~swiftsimio.objects.cosmo_array` it is not promoted + (and then no attributes are copied). + + Parameters + ---------- + func : callable + The function whose argument attributes will be copied to its result. + + Returns + ------- + out : callable + The wrapped function. + """ + + def wrapped(obj, *args, **kwargs): + # omit docstring so that sphinx picks up docstring of wrapped function + return _copy_cosmo_array_attributes_if_present(obj, func(obj, *args, **kwargs)) + + return wrapped + + +def _promote_unyt_to_cosmo(input_object: object) -> object: + """ + Upgrades the input unyt instance to its cosmo equivalent. + + In many cases we can obtain a unyt class instance and want to promote it to its cosmo + equivalent to attach our cosmo attributes. This helper promotes an input + :class:`~unyt.array.unyt_array` to a :class:`~swiftsimio.objects.cosmo_array` or an + input :class:`~unyt.array.unyt_quantity` to a + :class:`~swiftsimio.objects.cosmo_quantity`. If the input is neither type, it is just + returned. + + Parameters + ---------- + input_object : :obj:`object` + Object to consider for promotion from unyt instance to cosmo instance. + """ + if isinstance(input_object, unyt_quantity) and not isinstance( + input_object, objects.cosmo_quantity + ): + return input_object.view(objects.cosmo_quantity) + elif isinstance(input_object, unyt_array) and not isinstance( + input_object, objects.cosmo_array + ): + return input_object.view(objects.cosmo_array) + else: + return input_object + + +def _ensure_array_or_quantity_matches_shape(input_object: object) -> object: + """ + Convert scalars to :class:`~swiftsimio.objects.cosmo_quantity` and arrays to + :class:`~swiftsimio.objects.cosmo_array`. + + Scalar quantities are meant to be contained in + :class:`~swiftsimio.objects.cosmo_quantity` and arrays in + :class:`~swiftsimio.objects.cosmo_array`. Many functions (e.g. from numpy) can change + the data contents without changing the containing class. This helper checks the input + to make sure the data match the container type and converts if not. + + Parameters + ---------- + input_object : :obj:`object` + The object whose data is to be checked against its type. + + Returns + ------- + out : :obj:`object` + A version of the input with container type matching data contents. + """ + if ( + isinstance(input_object, objects.cosmo_array) + and not isinstance(input_object, objects.cosmo_quantity) + and input_object.shape == () + ): + return input_object.view(objects.cosmo_quantity) + elif isinstance(input_object, objects.cosmo_quantity) and input_object.shape != (): + return input_object.view(objects.cosmo_array) + else: + return input_object + + +def _ensure_result_is_cosmo_array_or_quantity(func: Callable) -> Callable: + """ + Wrapper that converts any :class:`~unyt.array.unyt_array` or + :class:`~unyt.array.unyt_quantity` instances in function output to cosmo equivalents. + + If the wrapped function returns a :obj:`tuple` (as many numpy functions do) it is + iterated over (but not recursively) and each element with a unyt class type is + upgraded to its cosmo equivalent. If anything but a :obj:`tuple` is returned, that + object is promoted to the cosmo equivalent if it is of a unyt class type. + + Parameters + ---------- + func : Callable + The function whose result(s) will be upgraded to + :class:`~swiftsimio.objects.cosmo_array` or + :class:`~swifsimio.objects.cosmo_quantity`. + Returns + ------- + out : Callable + The wrapped function. + """ + + def wrapped(*args, **kwargs) -> object: + # omit docstring so that sphinx picks up docstring of wrapped function + result = func(*args, **kwargs) + if isinstance(result, tuple): + return tuple( + _ensure_array_or_quantity_matches_shape(_promote_unyt_to_cosmo(item)) + for item in result + ) + else: + return _ensure_array_or_quantity_matches_shape( + _promote_unyt_to_cosmo(result) + ) + + return wrapped + + +def _sqrt_cosmo_factor(cf: "objects.cosmo_factor", **kwargs) -> "objects.cosmo_factor": + """ + Take the square root of a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` whose square root should be taken. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The square root of the input :class:`~swiftsimio.objects.cosmo_factor`. + """ + return _power_cosmo_factor(cf, None, power=0.5) + + +def _multiply_cosmo_factor( + *cfs: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Recursively multiply :class:`~swiftsimio.objects.cosmo_factor`s. + + All arguments are expected to be of type :class:`~swiftsimio.objects.cosmo_factor`. + They are cumumatively multipled together and the result returned. + + Parameters + ---------- + cfs : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor`s to be multiplied. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the input :class:`~swiftsimio.objects.cosmo_factor`s. + """ + return reduce(__binary_multiply_cosmo_factor, cfs) + + +def __binary_multiply_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Multiply two :class:`~swiftsimio.objects.cosmo_factor`s. + + Not intended for direct use but only as a helper for + :func:`~swiftsimio._array_functions._multiply_cosmo_factor`. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + The first :class:`~swiftsimio.objects.cosmo_factor`. + + cf2 : swiftsimio.objects.cosmo_factor + The second :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the :class:`~swiftsimio.objects.cosmo_factor`s. + """ + if (cf1 is None) and (cf2 is None): + # neither has cosmo_factor information: + return None + elif (cf1 is None) and (cf2 is not None): + # first has no cosmo information, allow e.g. multiplication by constants: + return cf2 + elif (cf1 is not None) and (cf2 is None): + # second has no cosmo information, allow e.g. multiplication by constants: + return cf1 + elif (cf1 is not None) and (cf2 is not None): + # both cosmo_array and both with cosmo_factor: + return cf1 * cf2 # cosmo_factor.__mul__ raises if scale factors differ + + +def _preserve_cosmo_factor( + *cfs: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Helper to preserve the :class:`~swiftsimio.objects.cosmo_factor` of input. + + If there is a single argument, return its ``cosmo_factor``. If there are multiple + arguments, check that they all have matching ``cosmo_factor``. Any arguments that + are not :class:`~swiftsimio.objects.cosmo_array`s are ignored for this purpose. + + Parameters + ---------- + cfs : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor`s to be preserved. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The preserved :class:`~swiftsimio.objects.cosmo_factor`. + """ + return reduce(__binary_preserve_cosmo_factor, cfs) + + +def __binary_preserve_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Given two :class:`~swiftsimio.objects.cosmo_factor`s, get it if they match. + + Not intended for direct use but only as a helper for + :func:`~swiftsimio._array_functions._preserve_cosmo_factor`. If the two inputs + are compatible, return the compatible :class:`~swiftsimio.objects.cosmo_factor`. + If one of them is ``None``, produce a warning. If they are incompatible, raise. + + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + The first :class:`~swiftsimio.objects.cosmo_factor`. + + cf2 : swiftsimio.objects.cosmo_factor + The second :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The preserved :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + Raised if the two arguments are :class:`~swiftsimio.objects.cosmo_factor`s + that do not have matching attributes. + """ + if (cf1 is None) and (cf2 is None): + # neither has cosmo_factor information: + return None + elif (cf1 is not None) and (cf2 is None): + # only one is cosmo_array + warnings.warn( + f"Mixing arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf1}) for all arguments.", + RuntimeWarning, + ) + return cf1 + elif (cf1 is None) and (cf2 is not None): + # only one is cosmo_array + warnings.warn( + f"Mixing arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf2}) for all arguments.", + RuntimeWarning, + ) + return cf2 + elif cf1 != cf2: + raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") + else: # cf1 == cf2 + return cf1 # or cf2, they're equal + + +def _power_cosmo_factor( + cf1: "objects.cosmo_factor", + cf2: "objects.cosmo_factor", + inputs: "Optional[Tuple[objects.cosmo_array]]" = None, + power: Optional[float] = None, +) -> "objects.cosmo_factor": + """ + Raise a :class:`~swiftsimio.objects.cosmo_factor` to a power of another + :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` attached to the base. + + cf2 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` attached to the exponent. + + inputs : :obj:`tuple` + The objects that ``cf1`` and ``cf2`` are attached to. Give ``inputs`` or + ``power``, not both. + + power : float + The power to raise ``cf1`` to. Give ``inputs`` or ``power``, not both. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The exponentiated :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the exponent is not a dimensionless quantity, or the exponent has a + ``cosmo_factor`` whose scaling with scale factor is not ``1.0``. + """ + if (inputs is not None and power is not None) or (inputs is None and power is None): + raise ValueError + power = inputs[1] if inputs else power + if hasattr(power, "units"): + if not power.units.is_dimensionless: + raise ValueError("Exponent must be dimensionless.") + elif power.units is not unyt.dimensionless: + power = power.to_value(unyt.dimensionless) + # else power.units is unyt.dimensionless, do nothing + exp_afactor = getattr(cf2, "a_factor", None) + if (exp_afactor is not None) and (exp_afactor != 1.0): + raise ValueError("Exponent has scaling with scale factor != 1.") + if cf1 is None: + return None + return np.power(cf1, power) + + +def _square_cosmo_factor( + cf: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Square a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` to square. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The squared :class:`~swiftsimio.objects.cosmo_factor`. + """ + return _power_cosmo_factor(cf, None, power=2) + + +def _cbrt_cosmo_factor(cf: "objects.cosmo_factor", **kwargs) -> "objects.cosmo_factor": + """ + Take the cube root of a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` whose cube root should be taken. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The cube root of the input :class:`~swiftsimio.objects.cosmo_factor`. + """ + return _power_cosmo_factor(cf, None, power=1.0 / 3.0) + + +def _divide_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Divide two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + Numerator :class:`~swiftsimio.objects.cosmo_factor`. + + cf1 : swiftsimio.objects.cosmo_factor + Denominator :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The ratio of the input :class:`~swiftsimio.objects.cosmo_factor`s. + """ + return _multiply_cosmo_factor(cf1, _reciprocal_cosmo_factor(cf2)) + + +def _reciprocal_cosmo_factor( + cf: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Take the inverse of a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` to be inverted. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The inverted :class:`~swiftsimio.objects.cosmo_factor`. + """ + return _power_cosmo_factor(cf, None, power=-1) + + +def _passthrough_cosmo_factor( + cf: "objects.cosmo_factor", cf2: "Optional[objects.cosmo_factor]" = None, **kwargs +) -> "objects.cosmo_factor": + """ + Preserve a :class:`~swiftsimio.objects.cosmo_factor`, optionally checking that it + matches a second :class:`~swiftsimio.objects.cosmo_factor`. + + This helper is intended for e.g. numpy ufuncs with a second dimensionless argument + so it's ok if ``cf2`` is ``None`` and ``cf1`` is not. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` to pass through. + + cf2 : swiftsimio.objects.cosmo_factor + Optional second :class:`~swiftsimio.objects.cosmo_factor` to check matches. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The input :class:`~swiftsimio.objects.cosmo_factor`. + + Raises + ------ + ValueError + If ``cf2`` is provided, is not ``None`` and does not match ``cf1``. + """ + if (cf2 is not None) and cf != cf2: + # if both have cosmo_factor information and it differs this is an error + raise ValueError(f"Arguments have cosmo_factors that differ: {cf} and {cf2}.") + else: + return cf + + +def _return_without_cosmo_factor( + cf: "objects.cosmo_factor", + cf2: "objects.cosmo_factor" = np._NoValue, + zero_comparison: Optional[bool] = None, + **kwargs, +) -> None: + """ + Return ``None``, but first check that argument + :class:`~swiftsimio.objects.cosmo_factor`s match, raising or warning if not. + + Comparisons are a special case that wraps around this wrapper, see + :func:`~swiftsimio._array_functions._comparison_cosmo_factor`. + + We borrow ``np._NoValue`` as a default for ``cf2`` because ``None`` here + represents the absence of a :class:`~swiftsimio.objects.cosmo_factor`, and + we need to handle that case. + + Warnings are produced when one argument has no cosmo factor, relevant e.g. + when comparing to constants. We handle comparison with zero, that is unambiguous, + as a special case. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to discard. + + cf2 : swiftsimio.objects.cosmo_factor + Optional second :class:`~swiftsimio.objects.cosmo_factor` to check for a + match with ``cf1``. + + zero_comparison : bool + If ``True``, silences warnings when exactly one of ``cf1`` and ``cf2`` is + ``None``. Enables comparing with zero without warning. + + Returns + ------- + out : None + The :class:`~swiftsimio.objects.cosmo_factor` is discarded. + + Raises + ------ + ValueError + If ``cf2`` is provided, is not ``None`` and does not match ``cf1``. + """ + if cf2 is np._NoValue: + # there was no second argument, return promptly + return None + if (cf is not None) and (cf2 is None): + # one is not a cosmo_array, warn on e.g. comparison to constants: + if not zero_comparison: + warnings.warn( + f"Mixing arguments with and without cosmo_factors, continuing" + f" assuming provided cosmo_factor ({cf}) for all arguments.", + RuntimeWarning, + ) + elif (cf is None) and (cf2 is not None): + # two is not a cosmo_array, warn on e.g. comparison to constants: + if not zero_comparison: + warnings.warn( + f"Mixing arguments with and without cosmo_factors, continuing" + f" assuming provided cosmo_factor ({cf2}) for all arguments.", + RuntimeWarning, + ) + elif (cf is not None) and (cf2 is not None) and (cf != cf2): + # both have cosmo_factor, don't match: + raise ValueError(f"Arguments have cosmo_factors that differ: {cf} and {cf2}.") + elif (cf is not None) and (cf2 is not None) and (cf == cf2): + # both have cosmo_factor, and they match: + pass + # return without cosmo_factor + return None + + +def _arctan2_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Helper specifically to handle the :class:`~swiftsimio.objects.cosmo_factor`s for the + ``arctan2`` ufunc from numpy. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` for the first ``arctan2`` argument. + + cf2 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` for the second ``arctan2`` argument. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` for the ``arctan2`` result. + + Raises + ------ + ValueError + If the input :class:`~swiftsimio.objects.cosmo_factor`s differ, they will + not cancel out and this is an error. + + """ + if (cf1 is None) and (cf2 is None): + return None + elif (cf1 is None) and (cf2 is not None): + warnings.warn( + f"Mixing arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf2}) for all arguments.", + RuntimeWarning, + ) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) + elif (cf1 is not None) and (cf2 is None): + warnings.warn( + f"Mixing arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf1}) for all arguments.", + RuntimeWarning, + ) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): + raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") + elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + + +def _comparison_cosmo_factor( + cf1: "objects.cosmo_factor", + cf2: "objects.cosmo_factor", + inputs: "Optional[Tuple[objects.cosmo_array]]" = None, +) -> None: + """ + Helper to enable comparisons involving :class:`~swiftsimio.objects.cosmo_factor`s. + + Warnings are emitted when the comparison is ambiguous, for instance if comparing to a + bare :obj:`float` or similar. Comparison to zero is a special case where we suppress + warnings. + + See also :func:`~swiftsimio._array_functions._return_without_cosmo_factor` that is + used in implementing this function. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + First :class:`~swiftsimio.objects.cosmo_factor` to compare. + + cf2 : swiftsimio.objects.cosmo_factor + Second :class:`~swiftsimio.objects.cosmo_factor` to compare. + + inputs : :obj:`tuple` + The objects that ``cf1`` and ``cf2`` are attached to. + + Returns + ------- + out : None + The :class:`~swiftsimio.objects.cosmo_factor` is discarded. + """ + try: + iter(inputs[0]) + except TypeError: + input1_iszero = ( + not getattr(inputs[0], "value", inputs[0]) and inputs[0] is not False + ) + else: + input1_iszero = not getattr(inputs[0], "value", inputs[0]).any() + try: + iter(inputs[1]) + except IndexError: + input2_iszero = None + except TypeError: + input2_iszero = ( + not getattr(inputs[1], "value", inputs[1]) and inputs[1] is not False + ) + else: + input2_iszero = not getattr(inputs[1], "value", inputs[1]).any() + zero_comparison = input1_iszero or input2_iszero + return _return_without_cosmo_factor(cf1, cf2=cf2, zero_comparison=zero_comparison) + + +def _prepare_array_func_args(*args, _default_cm: bool = True, **kwargs) -> dict: + """ + Coerce args and kwargs to a common ``comoving`` and collect ``cosmo_factor``s. + + This helper function is mostly intended for writing wrappers for unyt and numpy + functions. It checks for consistency for all args and kwargs, but is not recursive + so mixed cosmo attributes could be passed in the first argument to + :func:`numpy.concatenate`, for instance. This function can be used "recursively" in a + limited way manually: in functions like :func:`numpy.concatenate` where a list of + arrays is expected, it makes sense to pass the first argument to this function + to check consistency and attempt to coerce to comoving if needed. + + Note that unyt allows creating a :class:`~unyt.array.unyt_array` from e.g. arrays with + heterogenous units (it probably shouldn't...). Example: + + :: + + >>> u.unyt_array([np.arange(3), np.arange(3) * u.m]) + unyt_array([[0, 1, 2], + [0, 1, 2]], '(dimensionless)') + + It's impractical for us to try to cover all possible invalid user input without + unyt being stricter. + + The best way to understand the usage of this helper function is to look at the many + wrapped unyt and numpy functions in :mod:`~swiftsimio._array_functions`. + + Parameters + ---------- + _default_cm: bool + If mixed ``comoving`` attributes are found, their data are converted such that + their ``comoving`` has the value of this argument. (Default: ``True``) + + Returns + ------- + out : dict + A dictionary containing the input `args`` and ``kwargs`` coerced to a common + state, and lists of their ``cosmo_factor`` attributes, and ``comoving`` and + ``compression`` values that can be used in return values for wrapped functions, + when relevant. + + Raises + ------ + ValueError + If the input arrays cannot be coerced to a consistent state of ``comoving``. + """ + cms = [(hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args] + cfs = [getattr(arg, "cosmo_factor", None) for arg in args] + comps = [ + (hasattr(arg, "compression"), getattr(arg, "compression", None)) for arg in args + ] + kw_cms = { + k: (hasattr(kwarg, "comoving"), getattr(kwarg, "comoving", None)) + for k, kwarg in kwargs.items() + } + kw_cfs = {k: getattr(kwarg, "cosmo_factor", None) for k, kwarg in kwargs.items()} + kw_comps = { + k: (hasattr(kwarg, "compression"), getattr(kwarg, "compression", None)) + for k, kwarg in kwargs.items() + } + if len([cm[1] for cm in cms + list(kw_cms.values()) if cm[0]]) == 0: + # no cosmo inputs + ret_cm = None + elif all([cm[1] for cm in cms + list(kw_cms.values()) if cm[0]]): + # all cosmo inputs are comoving + ret_cm = True + elif all([cm[1] is None for cm in cms + list(kw_cms.values()) if cm[0]]): + # all cosmo inputs have comoving=None + ret_cm = None + elif any([cm[1] is None for cm in cms + list(kw_cms.values()) if cm[0]]): + # only some cosmo inputs have comoving=None + raise ValueError( + "Some arguments have comoving=None and others have comoving=True|False. " + "Result is undefined!" + ) + elif all([cm[1] is False for cm in cms + list(kw_cms.values()) if cm[0]]): + # all cosmo_array inputs are physical + ret_cm = False + else: + # mix of comoving and physical inputs + # better to modify inplace (convert_to_comoving)? + if _default_cm: + args = [ + arg.to_comoving() if cm[0] and not cm[1] else arg + for arg, cm in zip(args, cms) + ] + kwargs = { + k: kwarg.to_comoving() if kw_cms[k][0] and not kw_cms[k][1] else kwarg + for k, kwarg in kwargs.items() + } + ret_cm = True + else: + args = [ + arg.to_physical() if cm[0] and not cm[1] else arg + for arg, cm in zip(args, cms) + ] + kwargs = { + k: kwarg.to_physical() if kw_cms[k][0] and not kw_cms[k][1] else kwarg + for k, kwarg in kwargs.items() + } + ret_cm = False + if len(set(comps + list(kw_comps.values()))) == 1: + # all compressions identical, preserve it + ret_comp = (comps + list(kw_comps.values()))[0] + else: + # mixed compressions, strip it off + ret_comp = None + return dict( + args=args, + kwargs=kwargs, + cfs=cfs, + kw_cfs=kw_cfs, + comoving=ret_cm, + compression=ret_comp, + ) + + +def implements(numpy_function: Callable) -> Callable: + """ + Register an __array_function__ implementation for cosmo_array objects. + + Intended for use as a decorator. + + Parameters + ---------- + numpy_function : Callable + A function handle from numpy (not ufuncs). + + Returns + ------- + out : Callable + The wrapped function. + """ + + # See NEP 18 https://numpy.org/neps/nep-0018-array-function-protocol.html + def decorator(func: Callable) -> Callable: + """ + Actually register the specified function. + + Parameters + ---------- + func: Callable + The function wrapping the numpy equivalent. + + Returns + ------- + out : Callable + The input ``func``. + """ + _HANDLED_FUNCTIONS[numpy_function] = func + return func + + return decorator + + +def _return_helper( + res: np.ndarray, + helper_result: dict, + ret_cf: "objects.cosmo_factor", + out: Optional[np.ndarray] = None, +) -> "objects.cosmo_array": + """ + Helper function to attach our cosmo attributes to return values of wrapped functions. + + The return value is first promoted to be a :class:`~swiftsimio.objects.cosmo_array` + (or quantity) if necessary. If the return value is still not one of our cosmo + types, we don't attach attributes. + + Parameters + ---------- + res : numpy.ndarray + The output array of a function to attach our attributes to. + helper_result : dict + A helper :obj:`dict` returned by + :func:`~swiftsimio._array_functions._prepare_array_func_args`. + ret_cf : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to attach to the result. + out : numpy.ndarray + For functions that can place output in an ``out`` argument, a reference to + the output array (optional). + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The input return value of a wrapped function with our cosmo attributes applied. + """ + res = _promote_unyt_to_cosmo(res) + if isinstance(res, objects.cosmo_array): # also recognizes cosmo_quantity + res.comoving = helper_result["comoving"] + res.cosmo_factor = ret_cf + res.compression = helper_result["compression"] + if isinstance(out, objects.cosmo_array): # also recognizes cosmo_quantity + out.comoving = helper_result["comoving"] + out.cosmo_factor = ret_cf + out.compression = helper_result["compression"] + return res + + +def _default_unary_wrapper( + unyt_func: Callable, cosmo_factor_handler: Callable +) -> Callable: + """ + Wrapper helper for unary functions with typical behaviour. + + For many numpy and unyt functions with one (main) input argument, the wrapping + code that we need to apply is repetitive. Just prepare the arguments, apply + the chosen processing to the ``cosmo_factor``, and attach our cosmo attributes + to the return value. This function facilitates writing these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) function to be wrapped. + cosmo_factor_handler : Callable + The function that handles the ``cosmo_factor``s, chosen from those defined + in :mod:`~swiftsimio._array_functions`. + + Returns + ------- + out : Callable + The wrapped function. + """ + + def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor`` attriubtes, and attach attributes to + output. + + Returns + ------- + out : Callable + The wrapped function. + """ + helper_result = _prepare_array_func_args(*args, **kwargs) + ret_cf = cosmo_factor_handler(helper_result["cfs"][0]) + res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) + if "out" in kwargs: + return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) + else: + return _return_helper(res, helper_result, ret_cf) + + return wrapper + + +def _default_binary_wrapper( + unyt_func: Callable, cosmo_factor_handler: Callable +) -> Callable: + """ + Wrapper helper for binary functions with typical behaviour. + + For many numpy and unyt functions with two (main) input arguments, the wrapping + code that we need to apply is repetitive. Just prepare the arguments, apply + the chosen processing to the cosmo_factors, and attach our cosmo attributes + to the return value. This function facilitates writing these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) function to be wrapped. + cosmo_factor_handler : Callable + The function that handles the ``cosmo_factor``s, chosen from those defined + in :mod:`~swiftsimio._array_functions`. + + Returns + ------- + out : Callable + The wrapped function. + """ + + def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor`` attributes, and attach attributes to + output. + + Returns + ------- + out : Callable + The wrapped function. + """ + helper_result = _prepare_array_func_args(*args, **kwargs) + ret_cf = cosmo_factor_handler(helper_result["cfs"][0], helper_result["cfs"][1]) + res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) + if "out" in kwargs: + return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) + else: + return _return_helper(res, helper_result, ret_cf) + + return wrapper + + +def _default_comparison_wrapper(unyt_func: Callable) -> Callable: + """ + Wrapper helper for binary comparison functions with typical behaviour. + + For many numpy and unyt comparison functions with two (main) input arguments, the + wrapping code that we need to apply is repetitive. Just prepare the arguments, + process the ``cosmo_factors``, and attach our cosmo attributes + to the return value. This function facilitates writing these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) comparison function to be wrapped. + + Returns + ------- + out : Callable + The wrapped function. + """ + + # assumes we have two primary arguments that will be handled with + # _comparison_cosmo_factor with them as the inputs + def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor`` attributes, and attach attributes to + output. + + Returns + ------- + out : Callable + The wrapped function. + """ + helper_result = _prepare_array_func_args(*args, **kwargs) + ret_cf = _comparison_cosmo_factor( + helper_result["cfs"][0], helper_result["cfs"][1], inputs=args[:2] + ) + res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + return wrapper + + +def _default_oplist_wrapper(unyt_func: Callable) -> Callable: + """ + Wrapper helper for functions accepting a list of operands with typical behaviour. + + For many numpy and unyt functions taking a list of operands as an argument, the + wrapping code that we need to apply is repetitive. Just prepare the arguments, + preserve the ``cosmo_factor`` (after checking that it's common across the list), and + attach our cosmo attributes to the return value. This function facilitates writing + these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) function to be wrapped. + + Returns + ------- + out : Callable + The wrapped function. + """ + + def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor`` attributes, and attach attributes to + output. + + Returns + ------- + out : Callable + The wrapped function. + """ + helper_result = _prepare_array_func_args(*args, **kwargs) + helper_result_oplist = _prepare_array_func_args(*args[0]) + ret_cf = _preserve_cosmo_factor(helper_result_oplist["cfs"][0]) + res = unyt_func( + helper_result_oplist["args"], + *helper_result["args"][1:], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result_oplist, ret_cf) + + return wrapper + + +# Next we wrap functions from unyt and numpy. There's not much point in writing docstrings +# or type hints for all of these. + +# Now we wrap functions that unyt handles explicitly (below that will be those not handled +# explicitly): + + +@implements(np.array2string) +def array2string( + a, + max_line_width=None, + precision=None, + suppress_small=None, + separator=" ", + prefix="", + style=np._NoValue, + formatter=None, + threshold=None, + edgeitems=None, + sign=None, + floatmode=None, + suffix="", + *, + legacy=None, +): + + res = unyt_array2string( + a, + max_line_width=max_line_width, + precision=precision, + suppress_small=suppress_small, + separator=separator, + prefix=prefix, + style=style, + formatter=formatter, + threshold=threshold, + edgeitems=edgeitems, + sign=sign, + floatmode=floatmode, + suffix=suffix, + legacy=legacy, + ) + if a.comoving: + append = " (comoving)" + elif a.comoving is False: + append = " (physical)" + elif a.comoving is None: + append = "" + return res + append + + +implements(np.dot)(_default_binary_wrapper(unyt_dot, _multiply_cosmo_factor)) +implements(np.vdot)(_default_binary_wrapper(unyt_vdot, _multiply_cosmo_factor)) +implements(np.inner)(_default_binary_wrapper(unyt_inner, _multiply_cosmo_factor)) +implements(np.outer)(_default_binary_wrapper(unyt_outer, _multiply_cosmo_factor)) +implements(np.kron)(_default_binary_wrapper(unyt_kron, _multiply_cosmo_factor)) + + +@implements(np.histogram_bin_edges) +def histogram_bin_edges(a, bins=10, range=None, weights=None): + + helper_result = _prepare_array_func_args(a, bins=bins, range=range, weights=weights) + if not isinstance(bins, str) and np.ndim(bins) == 1: + # we got bin edges as input + ret_cf = _preserve_cosmo_factor(helper_result["kw_cfs"]["bins"]) + else: + # bins based on values in a + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) + res = unyt_histogram_bin_edges( + *helper_result["args"], **helper_result["kwargs"] + ) + return _return_helper(res, helper_result, ret_cf) + + +implements(np.linalg.inv)( + _default_unary_wrapper(unyt_linalg_inv, _reciprocal_cosmo_factor) +) +implements(np.linalg.tensorinv)( + _default_unary_wrapper(unyt_linalg_tensorinv, _reciprocal_cosmo_factor) +) +implements(np.linalg.pinv)( + _default_unary_wrapper(unyt_linalg_pinv, _reciprocal_cosmo_factor) +) +implements(np.linalg.svd)( + _default_unary_wrapper(unyt_linalg_svd, _preserve_cosmo_factor) +) + + +@implements(np.histogram) +def histogram(a, bins=10, range=None, density=None, weights=None): + + helper_result = _prepare_array_func_args( + a, bins=bins, range=range, density=density, weights=weights + ) + ret_cf_bins = _preserve_cosmo_factor(helper_result["cfs"][0]) + ret_cf_dens = _reciprocal_cosmo_factor(helper_result["cfs"][0]) + counts, bins = unyt_histogram(*helper_result["args"], **helper_result["kwargs"]) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"]) + ret_cf_counts = ( + _multiply_cosmo_factor(ret_cf_w, ret_cf_dens) if density else ret_cf_w + ) + else: + ret_cf_counts = ret_cf_dens if density else None + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity + counts.comoving = helper_result["comoving"] + counts.cosmo_factor = ret_cf_counts + counts.compression = helper_result["compression"] + return counts, _return_helper(bins, helper_result, ret_cf_bins) + + +@implements(np.histogram2d) +def histogram2d(x, y, bins=10, range=None, density=None, weights=None): + + if range is not None: + xrange, yrange = range + else: + xrange, yrange = None, None + + try: + N = len(bins) + except TypeError: + N = 1 + if N != 2: + xbins = ybins = bins + elif N == 2: + xbins, ybins = bins + helper_result_x = _prepare_array_func_args(x, bins=xbins, range=xrange) + helper_result_y = _prepare_array_func_args(y, bins=ybins, range=yrange) + if not density: + helper_result_w = _prepare_array_func_args(weights=weights) + ret_cf_x = _preserve_cosmo_factor(helper_result_x["cfs"][0]) + ret_cf_y = _preserve_cosmo_factor(helper_result_y["cfs"][0]) + if (helper_result_x["kwargs"]["range"] is None) and ( + helper_result_y["kwargs"]["range"] is None + ): + safe_range = None + else: + safe_range = ( + helper_result_x["kwargs"]["range"], + helper_result_y["kwargs"]["range"], + ) + counts, xbins, ybins = unyt_histogram2d( + helper_result_x["args"][0], + helper_result_y["args"][0], + bins=(helper_result_x["kwargs"]["bins"], helper_result_y["kwargs"]["bins"]), + range=safe_range, + density=density, + weights=helper_result_w["kwargs"]["weights"], + ) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_cfs"]["weights"]) + counts = _promote_unyt_to_cosmo(counts) + if isinstance( + counts, objects.cosmo_array + ): # also recognizes cosmo_quantity + counts.comoving = helper_result_w["comoving"] + counts.cosmo_factor = ret_cf_w + counts.compression = helper_result_w["compression"] + else: # density=True + # now x, y and weights must be compatible because they will combine + # we unpack input to the helper to get everything checked for compatibility + helper_result = _prepare_array_func_args( + x, + y, + xbins=xbins, + ybins=ybins, + xrange=xrange, + yrange=yrange, + weights=weights, + ) + ret_cf_x = _preserve_cosmo_factor(helper_result_x["cfs"][0]) + ret_cf_y = _preserve_cosmo_factor(helper_result_y["cfs"][0]) + if (helper_result["kwargs"]["xrange"] is None) and ( + helper_result["kwargs"]["yrange"] is None + ): + safe_range = None + else: + safe_range = ( + helper_result["kwargs"]["xrange"], + helper_result["kwargs"]["yrange"], + ) + counts, xbins, ybins = unyt_histogram2d( + helper_result["args"][0], + helper_result["args"][1], + bins=(helper_result["kwargs"]["xbins"], helper_result["kwargs"]["ybins"]), + range=safe_range, + density=density, + weights=helper_result["kwargs"]["weights"], + ) + ret_cf_xy = _multiply_cosmo_factor( + helper_result["cfs"][0], helper_result["cfs"][1] + ) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"]) + inv_ret_cf_xy = _reciprocal_cosmo_factor(ret_cf_xy) + ret_cf_counts = _multiply_cosmo_factor(ret_cf_w, inv_ret_cf_xy) + else: + ret_cf_counts = _reciprocal_cosmo_factor(ret_cf_xy) + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity + counts.comoving = helper_result["comoving"] + counts.cosmo_factor = ret_cf_counts + counts.compression = helper_result["compression"] + return ( + counts, + _return_helper(xbins, helper_result_x, ret_cf_x), + _return_helper(ybins, helper_result_y, ret_cf_y), + ) + + +@implements(np.histogramdd) +def histogramdd(sample, bins=10, range=None, density=None, weights=None): + + D = len(sample) + if range is not None: + ranges = range + else: + ranges = D * [None] + + try: + len(bins) + except TypeError: + # bins is an integer + bins = D * [bins] + helper_results = [ + _prepare_array_func_args(s, bins=b, range=r) + for s, b, r in zip(sample, bins, ranges) + ] + if not density: + helper_result_w = _prepare_array_func_args(weights=weights) + ret_cfs = [ + _preserve_cosmo_factor(helper_result["cfs"][0]) + for helper_result in helper_results + ] + if all( + [ + helper_result["kwargs"]["range"] is None + for helper_result in helper_results + ] + ): + safe_range = None + else: + safe_range = [ + helper_result["kwargs"]["range"] for helper_result in helper_results + ] + counts, bins = unyt_histogramdd( + [helper_result["args"][0] for helper_result in helper_results], + bins=[helper_result["kwargs"]["bins"] for helper_result in helper_results], + range=safe_range, + density=density, + weights=helper_result_w["kwargs"]["weights"], + ) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_cfs"]["weights"]) + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): + counts.comoving = helper_result_w["comoving"] + counts.cosmo_factor = ret_cf_w + counts.compression = helper_result_w["compression"] + else: # density=True + # now sample and weights must be compatible because they will combine + # we unpack input to the helper to get everything checked for compatibility + helper_result = _prepare_array_func_args( + *sample, bins=bins, range=range, weights=weights + ) + ret_cfs = D * [_preserve_cosmo_factor(helper_result["cfs"][0])] + counts, bins = unyt_histogramdd( + helper_result["args"], + bins=helper_result["kwargs"]["bins"], + range=helper_result["kwargs"]["range"], + density=density, + weights=helper_result["kwargs"]["weights"], + ) + if len(helper_result["cfs"]) == 1: + ret_cf_sample = _preserve_cosmo_factor(helper_result["cfs"][0]) + else: + ret_cf_sample = _multiply_cosmo_factor(*helper_result["cfs"]) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"]) + inv_ret_cf_sample = _reciprocal_cosmo_factor(ret_cf_sample) + ret_cf_counts = _multiply_cosmo_factor(ret_cf_w, inv_ret_cf_sample) + else: + ret_cf_counts = _reciprocal_cosmo_factor(ret_cf_sample) + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity + counts.comoving = helper_result["comoving"] + counts.cosmo_factor = ret_cf_counts + counts.compression = helper_result["compression"] + return ( + counts, + tuple( + _return_helper(b, helper_result, ret_cf) + for b, helper_result, ret_cf in zip(bins, helper_results, ret_cfs) + ), + ) + + +implements(np.concatenate)(_default_oplist_wrapper(unyt_concatenate)) +implements(np.cross)(_default_binary_wrapper(unyt_cross, _multiply_cosmo_factor)) +implements(np.intersect1d)( + _default_binary_wrapper(unyt_intersect1d, _preserve_cosmo_factor) +) +implements(np.union1d)(_default_binary_wrapper(unyt_union1d, _preserve_cosmo_factor)) +implements(np.linalg.norm)( + _default_unary_wrapper(unyt_linalg_norm, _preserve_cosmo_factor) +) +implements(np.vstack)(_default_oplist_wrapper(unyt_vstack)) +implements(np.hstack)(_default_oplist_wrapper(unyt_hstack)) +implements(np.dstack)(_default_oplist_wrapper(unyt_dstack)) +implements(np.column_stack)(_default_oplist_wrapper(unyt_column_stack)) +implements(np.stack)(_default_oplist_wrapper(unyt_stack)) +implements(np.around)(_default_unary_wrapper(unyt_around, _preserve_cosmo_factor)) + + +def _recursive_to_comoving(lst): + ret_lst = list() + for item in lst: + if isinstance(item, list): + ret_lst.append(_recursive_to_comoving(item)) + else: + ret_lst.append(item.to_comoving()) + return ret_lst + + +def _prepare_array_block_args(lst, recursing=False): + """ + Block accepts only a nested list of array "blocks". We need to recurse on this. + """ + helper_results = list() + if isinstance(lst, list): + for item in lst: + if isinstance(item, list): + helper_results += _prepare_array_block_args(item, recursing=True) + else: + helper_results.append(_prepare_array_func_args(item)) + if recursing: + return helper_results + cms = [hr["comoving"] for hr in helper_results] + comps = [hr["compression"] for hr in helper_results] + cfs = [hr["cfs"] for hr in helper_results] + convert_to_cm = False + if all(cms): + ret_cm = True + elif all([cm is None for cm in cms]): + ret_cm = None + elif any([cm is None for cm in cms]) and not all([cm is None for cm in cms]): + raise ValueError( + "Some input has comoving=None and others have " + "comoving=True|False. Result is undefined!" + ) + elif all([cm is False for cm in cms]): + ret_cm = False + else: + # mix of True and False only + ret_cm = True + convert_to_cm = True + if len(set(comps)) == 1: + ret_comp = comps[0] + else: + ret_comp = None + ret_cf = cfs[0] + for cf in cfs[1:]: + if cf != ret_cf: + raise ValueError("Mixed cosmo_factor values in input.") + if convert_to_cm: + ret_lst = _recursive_to_comoving(lst) + else: + ret_lst = lst + return dict( + args=ret_lst, + kwargs=dict(), + comoving=ret_cm, + cosmo_factor=ret_cf, + compression=ret_comp, + ) + + +@implements(np.block) +def block(arrays): + # block is a special case since we need to recurse more than one level + # down the list of arrays. + helper_result_block = _prepare_array_block_args(arrays) + ret_cf = helper_result_block["cosmo_factor"] + res = unyt_block(helper_result_block["args"]) + return _return_helper(res, helper_result_block, ret_cf) + + +implements(np.fft.fft)(_default_unary_wrapper(unyt_fft_fft, _reciprocal_cosmo_factor)) +implements(np.fft.fft2)(_default_unary_wrapper(unyt_fft_fft2, _reciprocal_cosmo_factor)) +implements(np.fft.fftn)(_default_unary_wrapper(unyt_fft_fftn, _reciprocal_cosmo_factor)) +implements(np.fft.hfft)(_default_unary_wrapper(unyt_fft_hfft, _reciprocal_cosmo_factor)) +implements(np.fft.rfft)(_default_unary_wrapper(unyt_fft_rfft, _reciprocal_cosmo_factor)) +implements(np.fft.rfft2)( + _default_unary_wrapper(unyt_fft_rfft2, _reciprocal_cosmo_factor) +) +implements(np.fft.rfftn)( + _default_unary_wrapper(unyt_fft_rfftn, _reciprocal_cosmo_factor) +) +implements(np.fft.ifft)(_default_unary_wrapper(unyt_fft_ifft, _reciprocal_cosmo_factor)) +implements(np.fft.ifft2)( + _default_unary_wrapper(unyt_fft_ifft2, _reciprocal_cosmo_factor) +) +implements(np.fft.ifftn)( + _default_unary_wrapper(unyt_fft_ifftn, _reciprocal_cosmo_factor) +) +implements(np.fft.ihfft)( + _default_unary_wrapper(unyt_fft_ihfft, _reciprocal_cosmo_factor) +) +implements(np.fft.irfft)( + _default_unary_wrapper(unyt_fft_irfft, _reciprocal_cosmo_factor) +) +implements(np.fft.irfft2)( + _default_unary_wrapper(unyt_fft_irfft2, _reciprocal_cosmo_factor) +) +implements(np.fft.irfftn)( + _default_unary_wrapper(unyt_fft_irfftn, _reciprocal_cosmo_factor) +) +implements(np.fft.fftshift)( + _default_unary_wrapper(unyt_fft_fftshift, _preserve_cosmo_factor) +) +implements(np.fft.ifftshift)( + _default_unary_wrapper(unyt_fft_ifftshift, _preserve_cosmo_factor) +) + +implements(np.sort_complex)( + _default_unary_wrapper(unyt_sort_complex, _preserve_cosmo_factor) +) +implements(np.isclose)(_default_comparison_wrapper(unyt_isclose)) +implements(np.allclose)(_default_comparison_wrapper(unyt_allclose)) +implements(np.array_equal)(_default_comparison_wrapper(unyt_array_equal)) +implements(np.array_equiv)(_default_comparison_wrapper(unyt_array_equiv)) + + +@implements(np.linspace) +def linspace( + start, + stop, + num=50, + endpoint=True, + retstep=False, + dtype=None, + axis=0, + *, + device=None, +): + + helper_result = _prepare_array_func_args( + start, + stop, + num=num, + endpoint=endpoint, + retstep=retstep, + dtype=dtype, + axis=axis, + device=device, + ) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][1]) + ress = unyt_linspace(*helper_result["args"], **helper_result["kwargs"]) + if retstep: + return tuple(_return_helper(res, helper_result, ret_cf) for res in ress) + else: + return _return_helper(ress, helper_result, ret_cf) + + +@implements(np.logspace) +def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0): + + helper_result = _prepare_array_func_args( + start, stop, num=num, endpoint=endpoint, base=base, dtype=dtype, axis=axis + ) + ret_cf = _preserve_cosmo_factor(helper_result["kw_cfs"]["base"]) + res = unyt_logspace(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +implements(np.geomspace)( + _default_binary_wrapper(unyt_geomspace, _preserve_cosmo_factor) +) + + +@implements(np.copyto) +def copyto(dst, src, casting="same_kind", where=True): + + helper_result = _prepare_array_func_args(dst, src, casting=casting, where=where) + if isinstance(src, objects.cosmo_array) and isinstance(dst, objects.cosmo_array): + # if we're copyting across two + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][1]) + # must pass dst directly here because it's modified in-place + if isinstance(src, objects.cosmo_array): + comoving = getattr(dst, "comoving", None) + if comoving: + src.convert_to_comoving() + elif comoving is False: + src.convert_to_physical() + unyt_copyto(dst, src, **helper_result["kwargs"]) + + +@implements(np.prod) +def prod( + a, + axis=None, + dtype=None, + out=None, + keepdims=np._NoValue, + initial=np._NoValue, + where=np._NoValue, +): + + helper_result = _prepare_array_func_args( + a, + axis=axis, + dtype=dtype, + out=out, + keepdims=keepdims, + initial=initial, + where=where, + ) + res = unyt_prod(*helper_result["args"], **helper_result["kwargs"]) + ret_cf = _power_cosmo_factor( + helper_result["cfs"][0], None, power=a.size // res.size + ) + return _return_helper(res, helper_result, ret_cf, out=out) + + +implements(np.var)(_default_unary_wrapper(unyt_var, _preserve_cosmo_factor)) +implements(np.trace)(_default_unary_wrapper(unyt_trace, _preserve_cosmo_factor)) +implements(np.percentile)( + _default_unary_wrapper(unyt_percentile, _preserve_cosmo_factor) +) +implements(np.quantile)(_default_unary_wrapper(unyt_quantile, _preserve_cosmo_factor)) +implements(np.nanpercentile)( + _default_unary_wrapper(unyt_nanpercentile, _preserve_cosmo_factor) +) +implements(np.nanquantile)( + _default_unary_wrapper(unyt_nanquantile, _preserve_cosmo_factor) +) + + +@implements(np.linalg.det) +def linalg_det(a): + + helper_result = _prepare_array_func_args(a) + ret_cf = _power_cosmo_factor(helper_result["cfs"][0], None, power=a.shape[0]) + res = unyt_linalg_det(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +implements(np.diff)(_default_unary_wrapper(unyt_diff, _preserve_cosmo_factor)) +implements(np.ediff1d)(_default_unary_wrapper(unyt_ediff1d, _preserve_cosmo_factor)) +implements(np.ptp)(_default_unary_wrapper(unyt_ptp, _preserve_cosmo_factor)) +# implements(np.cumprod)(...) Omitted because unyt just raises if called. + + +@implements(np.pad) +def pad(array, pad_width, mode="constant", **kwargs): + + helper_result = _prepare_array_func_args(array, pad_width, mode=mode, **kwargs) + # the number of options is huge, including user defined functions to handle data + # let's just preserve the cosmo_factor of the input `array` and trust the user... + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) + res = unyt_pad(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.choose) +def choose(a, choices, out=None, mode="raise"): + + helper_result = _prepare_array_func_args(a, choices, out=out, mode=mode) + helper_result_choices = _prepare_array_func_args(*choices) + ret_cf = _preserve_cosmo_factor(*helper_result_choices["cfs"]) + res = unyt_choose(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) + + +@implements(np.insert) +def insert(arr, obj, values, axis=None): + + helper_result = _prepare_array_func_args(arr, obj, values, axis=axis) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) + res = unyt_insert(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.linalg.lstsq) +def linalg_lstsq(a, b, rcond=None): + + helper_result = _prepare_array_func_args(a, b, rcond=rcond) + ret_cf = _divide_cosmo_factor(helper_result["cfs"][1], helper_result["cfs"][0]) + resid_cf = _power_cosmo_factor(helper_result["cfs"][1], None, power=2) + sing_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) + ress = unyt_linalg_lstsq(*helper_result["args"], **helper_result["kwargs"]) + return ( + _return_helper(ress[0], helper_result, ret_cf), + _return_helper(ress[1], helper_result, resid_cf), + ress[2], + _return_helper(ress[3], helper_result, sing_cf), + ) + + +@implements(np.linalg.solve) +def linalg_solve(a, b): + + helper_result = _prepare_array_func_args(a, b) + ret_cf = _divide_cosmo_factor(helper_result["cfs"][1], helper_result["cfs"][0]) + res = unyt_linalg_solve(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.linalg.tensorsolve) +def linalg_tensorsolve(a, b, axes=None): + + helper_result = _prepare_array_func_args(a, b, axes=axes) + ret_cf = _divide_cosmo_factor(helper_result["cfs"][1], helper_result["cfs"][0]) + res = unyt_linalg_tensorsolve(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.linalg.eig) +def linalg_eig(a): + + helper_result = _prepare_array_func_args(a) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) + ress = unyt_linalg_eig(*helper_result["args"], **helper_result["kwargs"]) + return (_return_helper(ress[0], helper_result, ret_cf), ress[1]) + + +@implements(np.linalg.eigh) +def linalg_eigh(a, UPLO="L"): + + helper_result = _prepare_array_func_args(a, UPLO=UPLO) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) + ress = unyt_linalg_eigh(*helper_result["args"], **helper_result["kwargs"]) + return (_return_helper(ress[0], helper_result, ret_cf), ress[1]) + + +implements(np.linalg.eigvals)( + _default_unary_wrapper(unyt_linalg_eigvals, _preserve_cosmo_factor) +) +implements(np.linalg.eigvalsh)( + _default_unary_wrapper(unyt_linalg_eigvalsh, _preserve_cosmo_factor) +) + + +@implements(np.savetxt) +def savetxt( + fname, + X, + fmt="%.18e", + delimiter=" ", + newline="\n", + header="", + footer="", + comments="# ", + encoding=None, +): + + warnings.warn( + "numpy.savetxt does not preserve units or cosmo_array information, " + "and will only save the raw numerical data from the cosmo_array object.\n" + "If this is the intended behaviour, call `numpy.savetxt(file, arr.d)` " + "to silence this warning.\n", + stacklevel=4, + ) + helper_result = _prepare_array_func_args( + fname, + X, + fmt=fmt, + delimiter=delimiter, + newline=newline, + header=header, + footer=footer, + comments=comments, + encoding=encoding, + ) + with warnings.catch_warnings(): + warnings.filterwarnings( + action="ignore", + category=UserWarning, + message="numpy.savetxt does not preserve units", + ) + unyt_savetxt(*helper_result["args"], **helper_result["kwargs"]) + return + + +@implements(np.apply_over_axes) +def apply_over_axes(func, a, axes): + res = func(a, axes[0]) + if len(axes) > 1: + # this function is recursive by nature, + # here we intentionally do not call the base _implementation + return np.apply_over_axes(func, res, axes[1:]) + else: + return res + + +@implements(np.fill_diagonal) +def fill_diagonal(a, val, wrap=False): + + helper_result = _prepare_array_func_args(a, val, wrap=wrap) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][1]) + # must pass a directly here because it's modified in-place + comoving = getattr(a, "comoving", None) + if comoving: + val.convert_to_comoving() + elif comoving is False: + val.convert_to_physical() + unyt_fill_diagonal(a, val, **helper_result["kwargs"]) + + +implements(np.isin)(_default_comparison_wrapper(unyt_isin)) + + +@implements(np.place) +def place(arr, mask, vals): + + helper_result = _prepare_array_func_args(arr, mask, vals) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) + # must pass arr directly here because it's modified in-place + if isinstance(vals, objects.cosmo_array): + comoving = getattr(arr, "comoving", None) + if comoving: + vals.convert_to_comoving() + elif comoving is False: + vals.convert_to_physical() + unyt_place(arr, mask, vals) + + +@implements(np.put) +def put(a, ind, v, mode="raise"): + + helper_result = _prepare_array_func_args(a, ind, v, mode=mode) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) + # must pass arr directly here because it's modified in-place + if isinstance(v, objects.cosmo_array): + comoving = getattr(a, "comoving", None) + if comoving: + v.convert_to_comoving() + elif comoving is False: + v.convert_to_physical() + unyt_put(a, ind, v, **helper_result["kwargs"]) + + +@implements(np.put_along_axis) +def put_along_axis(arr, indices, values, axis): + + helper_result = _prepare_array_func_args(arr, indices, values, axis) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) + # must pass arr directly here because it's modified in-place + if isinstance(values, objects.cosmo_array): + comoving = getattr(arr, "comoving", None) + if comoving: + values.convert_to_comoving() + elif comoving is False: + values.convert_to_physical() + unyt_put_along_axis(arr, indices, values, axis) + + +@implements(np.putmask) +def putmask(a, mask, values): + + helper_result = _prepare_array_func_args(a, mask, values) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) + # must pass arr directly here because it's modified in-place + if isinstance(values, objects.cosmo_array): + comoving = getattr(a, "comoving", None) + if comoving: + values.convert_to_comoving() + elif comoving is False: + values.convert_to_physical() + unyt_putmask(a, mask, values) + + +implements(np.searchsorted)( + _default_binary_wrapper(unyt_searchsorted, _return_without_cosmo_factor) +) + + +@implements(np.select) +def select(condlist, choicelist, default=0): + + helper_result = _prepare_array_func_args(condlist, choicelist, default=default) + helper_result_choicelist = _prepare_array_func_args(*choicelist) + ret_cf = _preserve_cosmo_factor(*helper_result_choicelist["cfs"]) + res = unyt_select( + helper_result["args"][0], + helper_result_choicelist["args"], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result, ret_cf) + + +implements(np.setdiff1d)( + _default_binary_wrapper(unyt_setdiff1d, _preserve_cosmo_factor) +) + + +@implements(np.sinc) +def sinc(x): + + # unyt just casts to array and calls the numpy implementation + # so let's just hand off to them + return unyt_sinc(x) + + +@implements(np.clip) +def clip( + a, + a_min=np._NoValue, + a_max=np._NoValue, + out=None, + *, + min=np._NoValue, + max=np._NoValue, + **kwargs, +): + + # can't work out how to properly handle min and max, + # just leave them in kwargs I guess (might be a numpy version conflict?) + helper_result = _prepare_array_func_args( + a, a_min=a_min, a_max=a_max, out=out, **kwargs + ) + ret_cf = _preserve_cosmo_factor( + helper_result["cfs"][0], + helper_result["kw_cfs"]["a_min"], + helper_result["kw_cfs"]["a_max"], + ) + res = unyt_clip( + helper_result["args"][0], + helper_result["kwargs"]["a_min"], + helper_result["kwargs"]["a_max"], + out=helper_result["kwargs"]["out"], + **kwargs, + ) + return _return_helper(res, helper_result, ret_cf, out=out) + + +@implements(np.where) +def where(condition, *args): + + helper_result = _prepare_array_func_args(condition, *args) + if len(args) == 0: # just condition + ret_cf = _return_without_cosmo_factor(helper_result["cfs"][0]) + res = unyt_where(*helper_result["args"], **helper_result["kwargs"]) + elif len(args) < 2: + # error message borrowed from numpy 1.24.1 + raise ValueError("either both or neither of x and y should be given") + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][1], helper_result["cfs"][2]) + res = unyt_where(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +implements(np.triu)(_default_unary_wrapper(unyt_triu, _preserve_cosmo_factor)) +implements(np.tril)(_default_unary_wrapper(unyt_tril, _preserve_cosmo_factor)) + + +@implements(np.einsum) +def einsum( + subscripts, + *operands, + out=None, + dtype=None, + order="K", + casting="safe", + optimize=False, +): + + helper_result = _prepare_array_func_args( + subscripts, + operands, + out=out, + dtype=dtype, + order=order, + casting=casting, + optimize=optimize, + ) + helper_result_operands = _prepare_array_func_args(*operands) + ret_cf = _preserve_cosmo_factor(*helper_result_operands["cfs"]) + res = unyt_einsum( + helper_result["args"][0], + *helper_result_operands["args"], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result_operands, ret_cf, out=out) + + +implements(np.convolve)(_default_binary_wrapper(unyt_convolve, _multiply_cosmo_factor)) +implements(np.correlate)( + _default_binary_wrapper(unyt_correlate, _multiply_cosmo_factor) +) +implements(np.tensordot)( + _default_binary_wrapper(unyt_tensordot, _multiply_cosmo_factor) +) + + +@implements(np.unwrap) +def unwrap(p, discont=None, axis=-1, *, period=6.283_185_307_179_586): + + helper_result = _prepare_array_func_args( + p, discont=discont, axis=axis, period=period + ) + ret_cf = _preserve_cosmo_factor( + helper_result["cfs"][0], + helper_result["kw_cfs"]["discont"], + helper_result["kw_cfs"]["period"], + ) + res = unyt_unwrap(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.interp) +def interp(x, xp, fp, left=None, right=None, period=None): + + helper_result = _prepare_array_func_args( + x, xp, fp, left=left, right=right, period=period + ) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][2]) + res = unyt_interp(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.array_repr) +def array_repr(arr, max_line_width=None, precision=None, suppress_small=None): + + helper_result = _prepare_array_func_args( + arr, + max_line_width=max_line_width, + precision=precision, + suppress_small=suppress_small, + ) + rep = unyt_array_repr(*helper_result["args"], **helper_result["kwargs"])[:-1] + if hasattr(arr, "comoving"): + rep += f", comoving='{arr.comoving}'" + if hasattr(arr, "cosmo_factor"): + rep += f", cosmo_factor='{arr.cosmo_factor}'" + if hasattr(arr, "valid_transform"): + rep += f", valid_transform='{arr.valid_transform}'" + rep += ")" + return rep + + +implements(np.linalg.outer)( + _default_binary_wrapper(unyt_linalg_outer, _multiply_cosmo_factor) +) + + +@implements(np.trapezoid) +def trapezoid(y, x=None, dx=1.0, axis=-1): + + helper_result = _prepare_array_func_args(y, x=x, dx=dx, axis=axis) + if x is None: + ret_cf = _multiply_cosmo_factor( + helper_result["cfs"][0], helper_result["kw_cfs"]["dx"] + ) + else: + ret_cf = _multiply_cosmo_factor( + helper_result["cfs"][0], helper_result["kw_cfs"]["x"] + ) + res = unyt_trapezoid(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +implements(np.in1d)(_default_comparison_wrapper(unyt_in1d)) +implements(np.take)(_default_unary_wrapper(unyt_take, _preserve_cosmo_factor)) + +# Now we wrap functions that unyt does not handle explicitly: + +implements(np.average)( + _propagate_cosmo_array_attributes_to_result(np.average._implementation) +) +implements(np.max)(_propagate_cosmo_array_attributes_to_result(np.max._implementation)) +implements(np.min)(_propagate_cosmo_array_attributes_to_result(np.min._implementation)) +implements(np.mean)( + _propagate_cosmo_array_attributes_to_result(np.mean._implementation) +) +implements(np.median)( + _propagate_cosmo_array_attributes_to_result(np.median._implementation) +) +implements(np.sort)( + _propagate_cosmo_array_attributes_to_result(np.sort._implementation) +) +implements(np.sum)(_propagate_cosmo_array_attributes_to_result(np.sum._implementation)) +implements(np.partition)( + _propagate_cosmo_array_attributes_to_result(np.partition._implementation) +) + + +@implements(np.meshgrid) +def meshgrid(*xi, **kwargs): + # meshgrid is a unique case: arguments never interact with each other, so we don't + # want to use our _prepare_array_func_args helper (that will try to coerce to + # compatible comoving, cosmo_factor). + # However we can't just use _propagate_cosmo_array_attributes_to_result because we + # need to iterate over arguments. + res = np.meshgrid._implementation(*xi, **kwargs) + return tuple( + _copy_cosmo_array_attributes_if_present(x, r) for (x, r) in zip(xi, res) + ) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index a8cd8ad8..7382f6d1 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -123,7 +123,7 @@ def read_ranges_from_file_unchunked( already_read = 0 handle_multidim = handle.ndim > 1 - for (read_start, read_end) in ranges: + for read_start, read_end in ranges: if read_end == read_start: continue diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index d246c1a1..0ff3878f 100644 --- a/swiftsimio/conversions.py +++ b/swiftsimio/conversions.py @@ -4,12 +4,11 @@ """ from swiftsimio.optional_packages import ASTROPY_AVAILABLE -import unyt +from swiftsimio.objects import cosmo_quantity if ASTROPY_AVAILABLE: from astropy.cosmology import w0waCDM from astropy.cosmology.core import Cosmology - import astropy.version import astropy.constants as const import astropy.units as astropy_units import numpy as np @@ -46,12 +45,14 @@ def swift_neutrinos_to_astropy(N_eff, N_ur, M_nu_eV, deg_nu): raise AttributeError( "SWIFTsimIO uses astropy, which cannot handle this cosmological model." ) - if not int(N_eff) == deg_nu.astype(int).sum() + int(N_ur): + if not int(N_eff) == deg_nu.astype(int).sum() + int(np.squeeze(N_ur)): raise AttributeError( - "SWIFTsimIO uses astropy, which cannot handle this cosmological model." + "SWIFTSimIO uses astropy, which cannot handle this cosmological model." ) ap_m_nu = [[m] * int(d) for m, d in zip(M_nu_eV, deg_nu)] # replicate - ap_m_nu = sum(ap_m_nu, []) + [0.0] * int(N_ur) # flatten + add massless + ap_m_nu = sum(ap_m_nu, []) + [0.0] * int( + np.squeeze(N_ur) + ) # flatten + add massless ap_m_nu = np.array(ap_m_nu) * astropy_units.eV return ap_m_nu @@ -74,7 +75,7 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: correct parameters. """ - H0 = unyt.unyt_quantity(cosmo["H0 [internal units]"][0], units=1.0 / units.time) + H0 = cosmo_quantity(cosmo["H0 [internal units]"][0], units=1.0 / units.time) Omega_b = cosmo["Omega_b"][0] Omega_lambda = cosmo["Omega_lambda"][0] diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index eb05a248..44de628b 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -5,17 +5,15 @@ import warnings -import unyt import h5py import numpy as np from swiftsimio.metadata.objects import SWIFTMetadata -from swiftsimio.objects import InvalidSnapshot +from swiftsimio.objects import InvalidSnapshot, cosmo_array, cosmo_quantity from swiftsimio.accelerated import ranges_from_array -from typing import Dict class SWIFTMask(object): @@ -229,11 +227,21 @@ def _unpack_cell_metadata(self): self.counts[key] = counts[sort] # Also need to sort centers in the same way - self.centers = unyt.unyt_array(centers_handle[:][sort], units=self.units.length) + self.centers = cosmo_array( + centers_handle[:][sort], + units=self.units.length, + comoving=True, + scale_factor=self.metadata.scale_factor, + scale_exponent=1, + ) # Note that we cannot assume that these are cubic, unfortunately. - self.cell_size = unyt.unyt_array( - metadata_handle.attrs["size"], units=self.units.length + self.cell_size = cosmo_array( + metadata_handle.attrs["size"], + units=self.units.length, + comoving=True, + scale_factor=self.metadata.scale_factor, + scale_exponent=1, ) return @@ -242,8 +250,8 @@ def constrain_mask( self, group_name: str, quantity: str, - lower: unyt.array.unyt_quantity, - upper: unyt.array.unyt_quantity, + lower: cosmo_quantity, + upper: cosmo_quantity, ): """ Constrains the mask further for a given particle type, and bounds a @@ -263,10 +271,10 @@ def constrain_mask( quantity : str quantity being constrained - lower : unyt.array.unyt_quantity + lower : ~swiftsimio.objects.cosmo_quantity constraint lower bound - upper : unyt.array.unyt_quantity + upper : ~swiftsimio.objects.cosmo_quantity constraint upper bound See Also @@ -300,13 +308,32 @@ def constrain_mask( handle = handle_dict[quantity] + physical_dict = { + k: v + for k, v in zip(group_metadata.field_names, group_metadata.field_physicals) + } + + physical = physical_dict[quantity] + + cosmologies_dict = { + k: v + for k, v in zip( + group_metadata.field_names, group_metadata.field_cosmologies + ) + } + + cosmology_factor = cosmologies_dict[quantity] + # Load in the relevant data. with h5py.File(self.metadata.filename, "r") as file: # Surprisingly this is faster than just using the boolean # indexing because h5py has slow indexing routines. - data = unyt.unyt_array( - np.take(file[handle], np.where(current_mask)[0], axis=0), units=unit + data = cosmo_array( + np.take(file[handle], np.where(current_mask)[0], axis=0), + units=unit, + comoving=not physical, + cosmo_factor=cosmology_factor, ) new_mask = np.logical_and.reduce([data > lower, data <= upper]) @@ -521,8 +548,8 @@ def constrain_index(self, index: int): if not self.metadata.homogeneous_arrays: raise RuntimeError( - "Cannot constrain to a single row in a non-homogeneous array; you currently " - f"are using a {self.metadata.output_type} file" + "Cannot constrain to a single row in a non-homogeneous array; you " + f"currently are using a {self.metadata.output_type} file" ) if not self.spatial_only: @@ -548,8 +575,8 @@ def constrain_indices(self, indices: list[int]): if not self.metadata.homogeneous_arrays: raise RuntimeError( - "Cannot constrain to a single row in a non-homogeneous array; you currently " - f"are using a {self.metadata.output_type} file" + "Cannot constrain to a single row in a non-homogeneous array; you " + f"currently are using a {self.metadata.output_type} file" ) if self.spatial_only: diff --git a/swiftsimio/metadata/metadata/metadata_fields.py b/swiftsimio/metadata/metadata/metadata_fields.py index a032e039..3358c6c1 100644 --- a/swiftsimio/metadata/metadata/metadata_fields.py +++ b/swiftsimio/metadata/metadata/metadata_fields.py @@ -2,7 +2,7 @@ Contains the description of the metadata fields in the SWIFT snapshots. """ -from ..objects import cosmo_factor, a +from ..objects import cosmo_factor metadata_fields_to_read = { "Code": "code", @@ -77,7 +77,7 @@ def generate_cosmo_args_header_unpack_arrays(scale_factor) -> dict: # should not be cosmo_array'd). cosmo_args = { "boxsize": dict( - cosmo_factor=cosmo_factor(a ** 1, scale_factor=scale_factor), + cosmo_factor=cosmo_factor.create(scale_factor, 1), comoving=True, # if it's not, then a=1 and it doesn't matter valid_transform=True, ) diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index df60c643..df46cd2f 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -12,7 +12,7 @@ import h5py from swiftsimio.conversions import swift_cosmology_to_astropy from swiftsimio import metadata -from swiftsimio.objects import cosmo_array, cosmo_factor, a +from swiftsimio.objects import cosmo_array, cosmo_factor from abc import ABC, abstractmethod import re @@ -610,12 +610,8 @@ def get_units(unit_attribute): # We should probably warn the user here... pass - # Deal with case where we _really_ have a dimensionless quantity. Comparing with - # 1.0 doesn't work, beacause in these cases unyt reverts to a floating point - # comparison. - try: - units.units - except AttributeError: + # Deal with case where we _really_ have a dimensionless quantity. + if not hasattr(units, "units"): units = None return units @@ -702,9 +698,7 @@ def get_cosmo(dataset): # Can't load, 'graceful' fallback. cosmo_exponent = 0.0 - a_factor_this_dataset = a ** cosmo_exponent - - return cosmo_factor(a_factor_this_dataset, current_scale_factor) + return cosmo_factor.create(current_scale_factor, cosmo_exponent) self.field_cosmologies = [ get_cosmo(self.metadata.handle[x]) for x in self.field_paths @@ -847,17 +841,8 @@ def handle(self): Property that gets the file handle, which can be shared with other objects for efficiency reasons. """ - if isinstance(self._handle, h5py.File): - # Can be open or closed, let's test. - try: - file = self._handle.file - - return self._handle - except ValueError: - # This will be the case if there is no active file handle - pass - - self._handle = h5py.File(self.filename, "r") + if not self._handle: # if self._handle is None, or if file closed (h5py #1363) + self._handle = h5py.File(self.filename, "r") return self._handle @@ -901,8 +886,8 @@ def __del__(self): def metadata_discriminator(filename: str, units: SWIFTUnits) -> "SWIFTMetadata": """ - Discriminates between the different types of metadata objects read from SWIFT-compatile - files. + Discriminates between the different types of metadata objects read from + SWIFT-compatible files. Parameters ---------- diff --git a/swiftsimio/metadata/soap/soap_types.py b/swiftsimio/metadata/soap/soap_types.py index cecb2730..5caa12c2 100644 --- a/swiftsimio/metadata/soap/soap_types.py +++ b/swiftsimio/metadata/soap/soap_types.py @@ -2,6 +2,7 @@ Includes the fancy names. """ + # Describes the conversion of hdf5 groups to names def get_soap_name_underscore(group: str) -> str: soap_name_underscores = { diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index edae2fd7..3538012e 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -1,18 +1,20 @@ """ -Contains global objects, e.g. the superclass version of the -unyt_array that we use, called cosmo_array. +Contains classes for our custom :class:`~swiftsimio.objects.cosmo_array`, +:class:`~swiftsimio.objects.cosmo_quantity` and +:class:`~swiftsimio.objects.cosmo_factor` objects for cosmology-aware +arrays, extending the functionality of the :class:`~unyt.array.unyt_array`. + +For developers, see also :mod:`swiftsimio._array_functions` containing +helpers, wrappers and implementations that enable most :mod:`numpy` and +:mod:`unyt` functions to work with our cosmology-aware arrays. """ -import warnings - import unyt -from unyt import unyt_array -from unyt.array import multiple_output_operators - -try: - from unyt.array import POWER_MAPPING -except ImportError: - raise ImportError("unyt >=2.9.0 required") +from unyt import unyt_array, unyt_quantity +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 +from collections.abc import Collection import sympy import numpy as np @@ -95,304 +97,123 @@ isnat, heaviside, matmul, + vecdot, +) +from numpy._core.umath import _ones_like, clip +from ._array_functions import ( + _propagate_cosmo_array_attributes_to_result, + _ensure_result_is_cosmo_array_or_quantity, + _sqrt_cosmo_factor, + _multiply_cosmo_factor, + _preserve_cosmo_factor, + _power_cosmo_factor, + _square_cosmo_factor, + _cbrt_cosmo_factor, + _divide_cosmo_factor, + _reciprocal_cosmo_factor, + _passthrough_cosmo_factor, + _return_without_cosmo_factor, + _arctan2_cosmo_factor, + _comparison_cosmo_factor, + _prepare_array_func_args, + _default_binary_wrapper, ) -from numpy._core.umath import _ones_like try: - from numpy._core.umath import clip + import pint except ImportError: - clip = None + pass # only for type hinting + +try: + import astropy.units +except ImportError: + pass # only for type hinting # The scale factor! a = sympy.symbols("a") -class InvalidConversionError(Exception): - def __init__(self, message="Could not convert to comoving coordinates"): - self.message = message - - -def _propagate_cosmo_array_attributes(func): - def wrapped(self, *args, **kwargs): - ret = func(self, *args, **kwargs) - if not type(ret) is cosmo_array: - return ret - if hasattr(self, "cosmo_factor"): - ret.cosmo_factor = self.cosmo_factor - if hasattr(self, "comoving"): - ret.comoving = self.comoving - if hasattr(self, "valid_transform"): - ret.valid_transform = self.valid_transform - return ret - - return wrapped - - -def _sqrt_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor( - ca_cf, (False, None), power=0.5 - ) # ufunc sqrt not supported - - -def _multiply_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - if (cf1 is None) and (cf2 is None): - # neither has cosmo_factor information: - return None - elif not ca1 and ca2: - # one is not a cosmo_array, allow e.g. multiplication by constants: - return cf2 - elif ca1 and not ca2: - # two is not a cosmo_array, allow e.g. multiplication by constants: - return cf1 - elif (ca1 and ca2) and ((cf1 is None) or (cf2 is None)): - # both cosmo_array but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors ({cf1} and {cf2})," - f" discarding cosmo_factor in return value.", - RuntimeWarning, - ) - return None - elif (ca1 and ca2) and ((cf1 is not None) and (cf2 is not None)): - # both cosmo_array and both with cosmo_factor: - return cf1 * cf2 # cosmo_factor.__mul__ raises if scale factors differ - else: - raise RuntimeError("Unexpected state, please report this error on github.") - - -def _preserve_cosmo_factor(ca_cf1, ca_cf2=None, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - if ca_cf2 is None: - # single argument, return promptly - return cf1 - elif (cf1 is None) and (cf2 is None): - # neither has cosmo_factor information: - return None - elif ca1 and not ca2: - # only one is cosmo_array - return cf1 - elif ca2 and not ca1: - # only one is cosmo_array - return cf2 - elif (ca1 and ca2) and (cf1 is None and cf2 is not None): - # both cosmo_array, but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - return cf2 - elif (ca1 and ca2) and (cf1 is not None and cf2 is None): - # both cosmo_array, but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf1}) for all arguments.", - RuntimeWarning, - ) - return cf1 - elif (ca1 and ca2) and (cf1 != cf2): - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." - ) - elif (ca1 and ca2) and (cf1 == cf2): - return cf1 # or cf2, they're equal - else: - raise RuntimeError("Unexpected state, please report this error on github.") - - -def _power_cosmo_factor(ca_cf1, ca_cf2, inputs=None, power=None): - if inputs is not None and power is not None: - raise ValueError - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - power = inputs[1] if inputs else power - if hasattr(power, "units"): - if not power.units.is_dimensionless: - raise ValueError("Exponent must be dimensionless.") - elif power.units is not unyt.dimensionless: - power = power.to_value(unyt.dimensionless) - # else power.units is unyt.dimensionless, do nothing - if ca2 and cf2.a_factor != 1.0: - raise ValueError("Exponent has scaling with scale factor != 1.") - if cf1 is None: - return None - return np.power(cf1, power) - - -def _square_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=2) - - -def _cbrt_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=1.0 / 3.0) - - -def _divide_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - return _multiply_cosmo_factor( - (ca1, cf1), (ca2, _reciprocal_cosmo_factor((ca2, cf2))) - ) - - -def _reciprocal_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=-1) +def _verify_valid_transform_validity(obj: "cosmo_array") -> None: + """ + Checks that ``comoving`` and ``valid_transform`` attributes are compatible. + Comoving arrays must be able to transform, while arrays that don't transform must + be physical. This function raises if this is not the case. -def _passthrough_cosmo_factor(ca_cf, ca_cf2=None, **kwargs): - ca, cf = ca_cf - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - if ca_cf2 is None: - # no second argument, return promptly - return cf - elif (cf2 is not None) and cf != cf2: - # if both have cosmo_factor information and it differs this is an error - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." - ) - else: - # passthrough is for e.g. ufuncs with a second dimensionless argument, - # so ok if cf2 is None and cf1 is not - return cf + Parameters + ---------- + obj : swiftsimio.objects.cosmo_array + The array whose validity is to be checked. + Raises + ------ + AssertionError + When an invalid combination of ``comoving`` and ``valid_transform`` is found. + """ + if not obj.valid_transform: + assert ( + not obj.comoving + ), "Cosmo arrays without a valid transform to comoving units must be physical" + if obj.comoving: + assert ( + obj.valid_transform + ), "Comoving cosmo_arrays must be able to be transformed to physical" -def _return_without_cosmo_factor(ca_cf, ca_cf2=None, inputs=None, zero_comparison=None): - ca, cf = ca_cf - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - if ca_cf2 is None: - # no second argument - pass - elif ca and not ca2: - # one is not a cosmo_array, warn on e.g. comparison to constants: - if not zero_comparison: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing" - f" assuming provided cosmo_factor ({cf}) for all arguments.", - RuntimeWarning, - ) - elif not ca and ca2: - # two is not a cosmo_array, warn on e.g. comparison to constants: - if not zero_comparison: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing" - f" assuming provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - elif (ca and ca2) and (cf is not None and cf2 is None): - # one has no cosmo_factor information, warn: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf}) for all arguments.", - RuntimeWarning, - ) - elif (ca and ca2) and (cf is None and cf2 is not None): - # two has no cosmo_factor information, warn: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - elif (cf is not None) and (cf2 is not None) and (cf != cf2): - # both have cosmo_factor, don't match: - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." - ) - elif ((cf is not None) and (cf2 is not None) and (cf == cf2)) or ( - (cf is None) and (cf2 is None) - ): - # both have cosmo_factor, and they match, or neither has cosmo_factor: - pass - else: - raise RuntimeError("Unexpected state, please report this error on github.") - # return without cosmo_factor - return None +class InvalidConversionError(Exception): + """ + Raised when converting from comoving from physical to comoving is not allowed. -def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - if (cf1 is None) and (cf2 is None): - return None - if cf1 is None and cf2 is not None: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - if cf1 is not None and cf2 is None: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf1}) for all arguments.", - RuntimeWarning, - ) - if (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." - ) - return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) + Parameters + ---------- + message : str, optional + Message to print in case of invalid conversion. + """ + def __init__( + self, message: str = "Could not convert to comoving coordinates." + ) -> None: + """ + Constructor for warning of invalid conversion. -def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - try: - iter(inputs[0]) - except TypeError: - if ca1: - input1_iszero = not inputs[0].value and inputs[0] is not False - else: - input1_iszero = not inputs[0] and inputs[0] is not False - else: - if ca1: - input1_iszero = not inputs[0].value.any() - else: - input1_iszero = not inputs[0].any() - try: - iter(inputs[1]) - except IndexError: - input2_iszero = None - except TypeError: - if ca2: - input2_iszero = not inputs[1].value and inputs[1] is not False - else: - input2_iszero = not inputs[1] and inputs[1] is not False - else: - if ca2: - input2_iszero = not inputs[1].value.any() - else: - input2_iszero = not inputs[1].any() - zero_comparison = input1_iszero or input2_iszero - return _return_without_cosmo_factor( - ca_cf1, ca_cf2=ca_cf2, inputs=inputs, zero_comparison=zero_comparison - ) + Parameters + ---------- + message : str, optional + Message to print in case of invalid conversion. + """ + self.message = message class InvalidScaleFactor(Exception): """ Raised when a scale factor is invalid, such as when adding two cosmo_factors with inconsistent scale factors. + + Parameters + ---------- + message : str, optional + Message to print in case of invalid scale factor. """ - def __init__(self, message=None, *args): + def __init__(self, message: str = None, *args) -> None: """ - Constructor for warning of invalid scale factor + Constructor for warning of invalid scale factor. Parameters ---------- - message : str, optional - Message to print in case of invalid scale factor + Message to print in case of invalid scale factor. """ self.message = message def __str__(self): """ - Print warning message of invalid scale factor + Print warning message for invalid scale factor. + + Returns + ------- + out : str + The error message. """ return f"InvalidScaleFactor: {self.message}" @@ -401,116 +222,196 @@ class InvalidSnapshot(Exception): """ Generated when a snapshot is invalid (e.g. you are trying to partially load a sub-snapshot). + + Parameters + ---------- + message : str, optional + Message to print in case of invalid snapshot. """ - def __init__(self, message=None, *args): + def __init__(self, message: str = None, *args) -> None: """ Constructor for warning of invalid snapshot Parameters ---------- - message : str, optional Message to print in case of invalid snapshot """ self.message = message - def __str__(self): + def __str__(self) -> str: """ Print warning message of invalid snapshot """ return f"InvalidSnapshot: {self.message}" -class cosmo_factor: +class cosmo_factor(object): """ Cosmology factor class for storing and computing conversion between comoving and physical coordinates. This takes the expected exponent of the array that can be parsed - by sympy, and the current value of the cosmological scale factor a. + by :mod:`sympy`, and the current value of the cosmological scale factor ``a``. This should be given as the conversion from comoving to physical, i.e. + :math:`r = a^f \times r` where :math:`a` is the scale factor, + :math:`r` is a physical quantity and :math`r'` a comoving quantity. + + Parameters + ---------- + expr : sympy.Expr + Expression used to convert between comoving and physical coordinates. + scale_factor : float + The scale factor (a). + + Attributes + ---------- + expr : sympy.Expr + Expression used to convert between comoving and physical coordinates. - r = cosmo_factor * r' with r in physical and r' comoving + scale_factor : float + The scale factor (a). Examples -------- + Mass density transforms as :math:`a^3`. To set up a ``cosmo_factor``, supposing + a current ``scale_factor=0.97``, we import the scale factor ``a`` and initialize + as: - Typically this would make cosmo_factor = a for the conversion between - comoving positions r' and physical co-ordinates r. + :: - To do this, use the a imported from objects multiplied as you'd like: + from swiftsimio.objects import a # the scale factor (a sympy symbol object) + density_cosmo_factor = cosmo_factor(a**3, scale_factor=0.97) - ``density_cosmo_factor = cosmo_factor(a**3, scale_factor=0.97)`` + :class:`~swiftsimio.objects.cosmo_factor` supports arithmetic, for example: + :: + + >>> cosmo_factor(a**2, scale_factor=0.5) * cosmo_factor(a**-1, scale_factor=0.5) + cosmo_factor(expr=a, scale_factor=0.5) + + See Also + -------- + swiftsimio.objects.cosmo_factor.create """ - def __init__(self, expr, scale_factor): + def __init__(self, expr: sympy.Expr, scale_factor: float) -> None: """ - Constructor for cosmology factor class + Constructor for cosmology factor class. Parameters ---------- - expr : sympy.expr - expression used to convert between comoving and physical coordinates + Expression used to convert between comoving and physical coordinates. scale_factor : float - the scale factor of the simulation data + The scale factor (a). + + See Also + -------- + swiftsimio.objects.cosmo_factor.create """ self.expr = expr self.scale_factor = scale_factor pass - def __str__(self): + @classmethod + def create(cls, scale_factor: float, exponent: numeric_type) -> "cosmo_factor": """ - Print exponent and current scale factor + Create a :class:`~swiftsimio.objects.cosmo_factor` from a scale factor and + exponent. + + Parameters + ---------- + scale_factor : :obj:`float` + The scale factor. + + exponent : :obj:`int` or :obj:`float` + The exponent defining the scaling with the scale factor. + + Examples + -------- + :: + + >>> cosmo_factor.create(0.5, 2) + cosmo_factor(expr=a**2, scale_factor=0.5) + """ + + obj = cls(a ** exponent, scale_factor) + + return obj + + def __str__(self) -> str: + """ + Print exponent and current scale factor. Returns ------- - - str - string to print exponent and current scale factor + out : str + String with exponent and current scale factor. """ return str(self.expr) + f" at a={self.scale_factor}" @property - def a_factor(self): + def a_factor(self) -> float: """ - The a-factor for the unit. + The multiplicative factor for conversion from comoving to physical. - e.g. for density this is 1 / a**3. + For example, for density this is :math:`a^{-3}`. Returns ------- - - float - the a-factor for given unit + out : float + The multiplicative factor for conversion from comoving to physical. """ + if (self.expr is None) or (self.scale_factor is None): + return None return float(self.expr.subs(a, self.scale_factor)) @property - def redshift(self): + def redshift(self) -> float: """ - Compute the redshift from the scale factor. + The redshift computed from the scale factor. + + Returns the redshift :math:`z = \\frac{1}{a} - 1`, where :math:`a` is the scale + factor. Returns ------- + out : float + The redshift. + """ + if self.scale_factor is None: + return None + return (1.0 / self.scale_factor) - 1.0 - float - redshift from the given scale factor + def __add__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Add two :class:`~swiftsimio.objects.cosmo_factor`s. - Notes - ----- + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to add to this one. - Returns the redshift - ..math:: z = \\frac{1}{a} - 1, - where :math: `a` is the scale factor - """ - return (1.0 / self.scale_factor) - 1.0 + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The sum of the two :class:`~swiftsimio.objects.cosmo_factor`s. - def __add__(self, b): + Raises + ------ + ValueError + If the object to be summed is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only add cosmo_factor to another cosmo_factor.") if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to add two cosmo_factors with different scale factors " @@ -525,7 +426,34 @@ def __add__(self, b): return cosmo_factor(expr=self.expr, scale_factor=self.scale_factor) - def __sub__(self, b): + def __sub__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Subtract two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to subtract from this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The difference of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be subtracted is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError( + "Can only subtract cosmo_factor from another cosmo_factor." + ) if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to subtract two cosmo_factors with different scale factors " @@ -540,96 +468,568 @@ def __sub__(self, b): return cosmo_factor(expr=self.expr, scale_factor=self.scale_factor) - def __mul__(self, b): + def __mul__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Multiply two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to multiply this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be multiplied is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError( + "Can only multiply cosmo_factor with another cosmo_factor." + ) if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to multiply two cosmo_factors with different scale factors " f"{self.scale_factor} and {b.scale_factor}" ) + if ((self.expr is None) and (b.expr is not None)) or ( + (self.expr is not None) and (b.expr is None) + ): + raise InvalidScaleFactor( + "Attempting to multiply an initialized cosmo_factor with an " + f"uninitialized cosmo_factor {self} and {b}." + ) + if (self.expr is None) and (b.expr is None): + # let's be permissive and allow two uninitialized cosmo_factors through + return cosmo_factor(expr=None, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr * b.expr, scale_factor=self.scale_factor) - def __truediv__(self, b): + def __truediv__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Divide two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to divide this one by. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The quotient of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to divide by is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only divide cosmo_factor with another cosmo_factor.") if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to divide two cosmo_factors with different scale factors " f"{self.scale_factor} and {b.scale_factor}" ) + if ((self.expr is None) and (b.expr is not None)) or ( + (self.expr is not None) and (b.expr is None) + ): + raise InvalidScaleFactor( + "Attempting to divide an initialized cosmo_factor with an " + f"uninitialized cosmo_factor {self} and {b}." + ) + if (self.expr is None) and (b.expr is None): + # let's be permissive and allow two uninitialized cosmo_factors through + return cosmo_factor(expr=None, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr / b.expr, scale_factor=self.scale_factor) - def __radd__(self, b): + def __radd__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Add two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to add to this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The sum of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be summed is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return self.__add__(b) - def __rsub__(self, b): + def __rsub__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Subtract two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to subtract from this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The difference of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be subtracted is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return self.__sub__(b) - def __rmul__(self, b): + def __rmul__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Multiply two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to multiply this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be multiplied is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return self.__mul__(b) - def __rtruediv__(self, b): + def __rtruediv__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Divide two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to divide this one by. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The quotient of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to divide by is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return b.__truediv__(self) - def __pow__(self, p): - return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) + def __pow__(self, p: float) -> "cosmo_factor": + """ + Raise this :class:`~swiftsimio.objects.cosmo_factor` to an exponent. + + Parameters + ---------- + p : float + The exponent by which to raise this :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The exponentiated :class:`~swiftsimio.objects.cosmo_factor`s. + """ + if self.expr is None: + return cosmo_factor(expr=None, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) + + def __lt__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") + return self.a_factor < b.a_factor + + def __gt__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") + return self.a_factor > b.a_factor + + def __le__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") + return self.a_factor <= b.a_factor + + def __ge__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") + return self.a_factor >= b.a_factor + + def __eq__(self, b: "cosmo_factor") -> bool: + """ + Compare the expressions and values of two + :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. Notice that unlike ``__gt__``, + ``__ge__``, ``__lt__`` and ``__le__``, (in)equality comparisons check that the + expression is (un)equal as well as the value. This is so that e.g. ``a**1`` and + ``a**2``, both with ``scale_factor=1.0`` are not equal (both have + ``a_factor==1``). + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") + return (self.scale_factor == b.scale_factor) and (self.a_factor == b.a_factor) + + def __ne__(self, b: "cosmo_factor") -> bool: + """ + Compare the expressions and values of two + :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. Notice that unlike ``__gt__``, + ``__ge__``, ``__lt__`` and ``__le__``, (in)equality comparisons check that the + expression is (un)equal as well as the value. This is so that e.g. ``a**1`` and + ``a**2``, both with ``scale_factor=1.0`` are not equal (both have + ``a_factor==1``). + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + return not self.__eq__(b) + + def __repr__(self) -> str: + """ + Get a string representation of the scaling with the scale factor. + + Returns + ------- + out : str + String representation of the scaling with the scale factor. + """ + return f"cosmo_factor(expr={self.expr}, scale_factor={self.scale_factor})" - def __lt__(self, b): - return self.a_factor < b.a_factor - def __gt__(self, b): - return self.a_factor > b.a_factor +NULL_CF = cosmo_factor(None, None) # helps avoid name collisions with kwargs below - def __le__(self, b): - return self.a_factor <= b.a_factor - def __ge__(self, b): - return self.a_factor >= b.a_factor +def _parse_cosmo_factor_args( + cf: cosmo_factor = None, + scale_factor: float = None, + scale_exponent: numeric_type = None, +) -> cosmo_factor: + """ + Decide what provided cosmology information to use, or raise an error. - def __eq__(self, b): - # Doesn't handle some corner cases, e.g. cosmo_factor(a ** 1, scale_factor=1) - # is considered equal to cosmo_factor(a ** 2, scale_factor=1) because - # 1 ** 1 == 1 ** 2. Should check self.expr vs b.expr with sympy? - return (self.scale_factor == b.scale_factor) and (self.a_factor == b.a_factor) + If both a ``cosmo_factor`` and a (``scale_factor``, ``scale_exponent``) pair are + given then this is an error. If only one of ``scale_factor`` and ``scale_exponent`` + is given this is an error. Otherwise we construct the + :class:`~swiftsimio.objects.cosmo_factor`, unless it's going to be a ``NULL_CF`` with + the information we have - in that case we return ``None`` to give a chance for it + to be filled in elsewhere before assuming the ``NULL_CF`` default. - def __ne__(self, b): - return not self.__eq__(b) + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` passed as an explicit argument. + scale_factor : numeric_type + The scale factor passed as a kwarg. + scale_exponent : float + The exponent for the scale factor to convert to/from comoving passed as a kwarg. + + Returns + ------- + out : cosmo_factor or None + The :class:`~swiftsimio.objects.cosmo_factor` to use, or ``None``. + + Raises + ------ + ValueError + If multiple values or incomplete information for the desired + :class:`~swiftsimio.objects.cosmo_factor` are provided. + """ + if cf is None and scale_factor is None and scale_exponent is None: + # we can return promptly + return None + if cf is not None: + if scale_factor is not None or scale_exponent is not None: + if cosmo_factor.create(scale_factor, scale_exponent) != cf: + raise ValueError( + "Provide either `cosmo_factor` or (`scale_factor` and " + "`scale_exponent`, not both (perhaps there was a `cosmo_factor` " + "attached to the input array or scalar?)." + ) + else: + # the duplicate information matches so let's allow it + return cf + else: + return cf + else: + if (scale_factor is not None and scale_exponent is None) or ( + scale_factor is None and scale_exponent is not None + ): + raise ValueError( + "Provide values for both `scale_factor` and `scale_exponent`." + ) + if scale_factor is None and scale_exponent is None: + return NULL_CF + else: + return cosmo_factor.create(scale_factor, scale_exponent) class cosmo_array(unyt_array): """ Cosmology array class. - This inherits from the unyt.unyt_array, and adds - four variables: compression, cosmo_factor, comoving, and valid_transform. - Data is assumed to be comoving when passed to the object but you - can override this by setting the latter flag to be False. + This inherits from the :class:`~unyt.array.unyt_array`, and adds + four attributes: ``compression``, ``cosmo_factor``, ``comoving``, and + ``valid_transform``. + + .. note:: + + :class:`~swiftsimio.objects.cosmo_array` and the related + :class:`~swiftsimio.objects.cosmo_quantity` are now intended to support all + :mod:`numpy` functions, propagating units (thanks to :mod:`unyt`) and + cosmology information. There are a large number of functions, and a very + large number of possible parameter combinations, so some corner cases may + have been missed in testing. Please report any issues on github, they are + usually easy to fix for future use! Currently :mod:`scipy` functions are + not supported (although some might "just work"). Requests to fully support + specific functions can also be submitted as github issues. Parameters ---------- - - unyt_array : unyt.unyt_array - the inherited unyt_array + input_array : np.ndarray, unyt.array.unyt_array or iterable + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, powers must be specified using + python syntax (``cm**3``, not ``cm^3``). + registry : unyt.unit_registry.UnitRegistry, optional + The registry to create units from. If ``units`` is already associated + with a unit registry and this is specified, this will be used instead of the + registry associated with the unit object. + dtype : np.dtype or str, optional + The dtype of the array data. Defaults to the dtype of the input data, or, if + none is found, uses ``np.float64``. + bypass_validation : bool, optional + If ``True``, all input validation is skipped. Using this option may produce + corrupted or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. Defaults to ``False``. + name : str, optional + The name of the array. Defaults to ``None``. This attribute does not propagate + through mathematical operations, but is preserved under indexing and unit + conversions. + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. + comoving : bool + Flag to indicate whether using comoving coordinates. + valid_transform : bool + Flag to indicate whether this array can be converted to comoving. If ``False``, + then ``comoving`` must be ``False``. + compression : string + Description of the compression filters that were applied to that array in the + hdf5 file. Attributes ---------- - comoving : bool - if True then the array is in comoving co-ordinates, and if - False then it is in physical units. + If ``True`` then the array is in comoving coordinates, if``False`` then it is in + physical units. - cosmo_factor : float - Object to store conversion data between comoving and physical coordinates + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. compression : string String describing any compression that was applied to this array in the hdf5 file. valid_transform: bool - if True then the array can be converted from physical to comoving units - + If ``True`` then the array can be converted from physical to comoving units. + + Notes + ----- + This class will generally try to make sense of input and initialize an array-like + object consistent with the input, and warn or raise if this cannot be done + consistently. However, the way that :class:`~unyt.array.unyt_array` handles input + imposes some limits to this. In particular, nested non-numpy containers given in + input are not traversed recursively, but only one level deep. This means that + while with this input the attributes are detected by the new array correctly: + + :: + + >>> from swiftsimio.objects import cosmo_array, cosmo_factor + >>> x = cosmo_array( + ... np.arange(3), + ... u.kpc, + ... comoving=True, + ... scale_factor=1.0, + ... scale_exponent=1, + ... ) + >>> cosmo_array([x, x]) + cosmo_array([[0, 1, 2], + [0, 1, 2]], 'kpc', comoving='True', cosmo_factor='a at a=1.0', + valid_transform='True') + + with this input they are lost: + + :: + + >>> cosmo_array([[x, x],[x, x]]) + cosmo_array([[[0, 1, 2],[0, 1, 2]],[[0, 1, 2],[0, 1, 2]]], + '(dimensionless)', comoving='None', cosmo_factor='None at a=None', + valid_transform='True') + + See Also + -------- + swiftsimio.objects.cosmo_quantity """ - # TODO: _cosmo_factor_ufunc_registry = { add: _preserve_cosmo_factor, subtract: _preserve_cosmo_factor, @@ -718,65 +1118,71 @@ class cosmo_array(unyt_array): _ones_like: _preserve_cosmo_factor, matmul: _multiply_cosmo_factor, clip: _passthrough_cosmo_factor, + vecdot: _multiply_cosmo_factor, } def __new__( cls, - input_array, - units=None, - registry=None, - dtype=None, - bypass_validation=False, - input_units=None, - name=None, - cosmo_factor=None, - comoving=None, - valid_transform=True, - compression=None, - ): + input_array: Iterable, + units: Union[str, unyt.unit_object.Unit, "astropy.units.core.Unit"] = None, + *, + registry: unyt.unit_registry.UnitRegistry = None, + dtype: Union[np.dtype, str] = None, + bypass_validation: bool = False, + name: str = None, + cosmo_factor: cosmo_factor = None, + scale_factor: Optional[float] = None, + scale_exponent: Optional[float] = None, + comoving: bool = None, + valid_transform: bool = True, + compression: str = None, + ) -> "cosmo_array": """ - Essentially a copy of the __new__ constructor. + Closely inspired by the :meth:`unyt.array.unyt_array.__new__` constructor. Parameters ---------- - input_array : iterable - A tuple, list, or array to attach units to - units : str, unyt.unit_symbols or astropy.unit, optional - The units of the array. Powers must be specified using python syntax - (cm**3, not cm^3). + input_array : np.ndarray, unyt.array.unyt_array or iterable + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, powers must be specified using + python syntax (``cm**3``, not ``cm^3``). registry : unyt.unit_registry.UnitRegistry, optional - The registry to create units from. If input_units is already associated with a - unit registry and this is specified, this will be used instead of the registry - associated with the unit object. + The registry to create units from. If ``units`` is already associated + with a unit registry and this is specified, this will be used instead of the + registry associated with the unit object. dtype : np.dtype or str, optional The dtype of the array data. Defaults to the dtype of the input data, or, if - none is found, uses np.float64 + none is found, uses ``np.float64``. bypass_validation : bool, optional - If True, all input validation is skipped. Using this option may produce - corrupted, invalid units or array data, but can lead to significant speedups - in the input validation logic adds significant overhead. If set, input_units - must be a valid unit object. Defaults to False. - input_units : str, optional - deprecated in favour of units option + If ``True``, all input validation is skipped. Using this option may produce + corrupted or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. Defaults to ``False``. name : str, optional - The name of the array. Defaults to None. This attribute does not propagate + The name of the array. Defaults to ``None``. This attribute does not propagate through mathematical operations, but is preserved under indexing and unit conversions. - cosmo_factor : cosmo_factor - cosmo_factor object to store conversion data between comoving and physical - coordinates + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. + scale_factor : float + The scale factor associated to the data. Also provide a value for + ``scale_exponent``. + scale_exponent : int or float + The exponent for the scale factor giving the scaling for conversion to/from + comoving units. Also provide a value for ``scale_factor``. comoving : bool - flag to indicate whether using comoving coordinates + Flag to indicate whether using comoving coordinates. valid_transform : bool - flag to indicate whether this array can be converted to comoving + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. compression : string - description of the compression filters that were applied to that array in the - hdf5 file + Description of the compression filters that were applied to that array in the + hdf5 file. """ - cosmo_factor: cosmo_factor + if bypass_validation is True: - try: obj = super().__new__( cls, input_array, @@ -786,81 +1192,153 @@ def __new__( bypass_validation=bypass_validation, name=name, ) - except TypeError: - # Older versions of unyt (before input_units was deprecated) - obj = super().__new__( - cls, - input_array, - units=units, - registry=registry, - dtype=dtype, - bypass_validation=bypass_validation, - input_units=input_units, - name=name, + + # dtype, units, registry & name are handled by unyt + obj.comoving = comoving + obj.cosmo_factor = cosmo_factor if cosmo_factor is not None else NULL_CF + if scale_factor is not None: # ambiguity, but this is `bypass_validation` + obj.cosmo_factor = cosmo_factor.create(scale_factor, scale_exponent) + obj.valid_transform = valid_transform + obj.compression = compression + + return obj + + if isinstance(input_array, cosmo_array): + + obj = input_array.view(cls) + + # do cosmo_factor first since it can be used in comoving/physical conversion: + cosmo_factor = ( + input_array.cosmo_factor + if input_array.cosmo_factor != NULL_CF + else None ) - except TypeError: - # Even older versions of unyt (before name was added) - obj = super().__new__( - cls, - input_array, - units=units, - registry=registry, - dtype=dtype, - bypass_validation=bypass_validation, - input_units=input_units, + cosmo_factor = _parse_cosmo_factor_args( + cf=cosmo_factor, + scale_factor=scale_factor, + scale_exponent=scale_exponent, + ) + if cosmo_factor is not None: + obj.cosmo_factor = cosmo_factor + # else is already copied from input_array + + if comoving is True: + obj.convert_to_comoving() + elif comoving is False: + obj.convert_to_physical() + # else is already copied from input_array + + # only overwrite valid_transform after transforming so that invalid + # transformations raise: + obj.valid_transform = valid_transform + _verify_valid_transform_validity(obj) + + obj.compression = ( + compression if compression is not None else obj.compression + ) + + return obj + + elif isinstance(input_array, np.ndarray) and input_array.dtype != object: + + # guard np.ndarray so it doesn't get caught by _iterable in next case + + # ndarray with object dtype goes to next case to properly handle e.g. + # ndarrays containing cosmo_quantities + + cosmo_factor = _parse_cosmo_factor_args( + cf=cosmo_factor, + scale_factor=scale_factor, + scale_exponent=scale_exponent, + ) + + elif _iterable(input_array) and input_array: + # if _prepare_array_func_args finds cosmo_array input it will convert to: + default_cm = comoving if comoving is not None else True + + # coerce any cosmo_array inputs to consistency: + helper_result = _prepare_array_func_args( + *input_array, _default_cm=default_cm + ) + + input_array = helper_result["args"] + + # default to comoving, cosmo_factor and compression given as kwargs + comoving = helper_result["comoving"] if comoving is None else comoving + cosmo_factor = _parse_cosmo_factor_args( + cf=cosmo_factor, + scale_factor=scale_factor, + scale_exponent=scale_exponent, + ) + cosmo_factor = ( + _preserve_cosmo_factor(*helper_result["cfs"]) + if cosmo_factor is None + else cosmo_factor ) + compression = ( + helper_result["compression"] if compression is None else compression + ) + # valid_transform has a non-None default, so we have to decide to always + # respect it + + obj = super().__new__( + cls, + input_array, + units=units, + registry=registry, + dtype=dtype, + bypass_validation=bypass_validation, + name=name, + ) if isinstance(obj, unyt_array) and not isinstance(obj, cls): obj = obj.view(cls) - obj.cosmo_factor = cosmo_factor + # attach our attributes: obj.comoving = comoving + # unyt allows creating a unyt_array from e.g. arrays with heterogenous units + # (it probably shouldn't...), so we don't recurse deeply and therefore + # can't guarantee that cosmo_factor isn't None at this point, guard with default + obj.cosmo_factor = cosmo_factor if cosmo_factor is not None else NULL_CF obj.compression = compression obj.valid_transform = valid_transform - if not obj.valid_transform: - assert ( - not obj.comoving - ), "Cosmo arrays without a valid transform to comoving units must be physical" - if obj.comoving: - assert ( - obj.valid_transform - ), "Comoving Cosmo arrays must be able to be transformed to physical" + _verify_valid_transform_validity(obj) return obj - def __array_finalize__(self, obj): + def __array_finalize__(self, obj: "cosmo_array") -> None: super().__array_finalize__(obj) if obj is None: return - self.cosmo_factor = getattr(obj, "cosmo_factor", None) + self.cosmo_factor = getattr(obj, "cosmo_factor", NULL_CF) self.comoving = getattr(obj, "comoving", None) self.compression = getattr(obj, "compression", None) self.valid_transform = getattr(obj, "valid_transform", True) - def __str__(self): + def __str__(self) -> str: if self.comoving: comoving_str = "(Comoving)" + elif self.comoving is None: + comoving_str = "(Physical/comoving not set)" else: comoving_str = "(Physical)" return super().__str__() + " " + comoving_str - def __repr__(self): - if self.comoving: - comoving_str = ", comoving=True)" - elif self.comoving is None: - comoving_str = ")" - else: - comoving_str = ", comoving=False)" - # Remove final parenthesis and append comoving flag - return super().__repr__()[:-1] + comoving_str + def __repr__(self) -> str: + return super().__repr__() - def __reduce__(self): + def __reduce__(self) -> tuple: """ - Pickle reduction method + Pickle reduction method. - Here we add an extra element at the start of the unyt_array state - tuple to store the cosmology info. + Here we add an extra element at the start of the :class:`~unyt.array.unyt_array` + state tuple to store the cosmology info. + + Returns + ------- + out : tuple + The state ready for pickling. """ np_ret = super(cosmo_array, self).__reduce__() obj_state = np_ret[2] @@ -870,53 +1348,56 @@ def __reduce__(self): new_ret = np_ret[:2] + cosmo_state + np_ret[3:] return new_ret - def __setstate__(self, state): + def __setstate__(self, state: Tuple) -> None: """ - Pickle setstate method + Pickle setstate method. Here we extract the extra cosmology info we added to the object - state and pass the rest to unyt_array.__setstate__. + state and pass the rest to :meth:`unyt.array.unyt_array.__setstate__`. + + Parameters + ---------- + state : tuple + A :obj:`tuple` containing the extra state information. """ super(cosmo_array, self).__setstate__(state[1:]) self.cosmo_factor, self.comoving, self.valid_transform = state[0] # Wrap functions that return copies of cosmo_arrays so that our # attributes get passed through: - __getitem__ = _propagate_cosmo_array_attributes(unyt_array.__getitem__) - __copy__ = _propagate_cosmo_array_attributes(unyt_array.__copy__) - __deepcopy__ = _propagate_cosmo_array_attributes(unyt_array.__deepcopy__) - in_cgs = _propagate_cosmo_array_attributes(unyt_array.in_cgs) - astype = _propagate_cosmo_array_attributes(unyt_array.astype) - in_units = _propagate_cosmo_array_attributes(unyt_array.in_units) - byteswap = _propagate_cosmo_array_attributes(unyt_array.byteswap) - compress = _propagate_cosmo_array_attributes(unyt_array.compress) - diagonal = _propagate_cosmo_array_attributes(unyt_array.diagonal) - flatten = _propagate_cosmo_array_attributes(unyt_array.flatten) - ravel = _propagate_cosmo_array_attributes(unyt_array.ravel) - repeat = _propagate_cosmo_array_attributes(unyt_array.repeat) - reshape = _propagate_cosmo_array_attributes(unyt_array.reshape) - swapaxes = _propagate_cosmo_array_attributes(unyt_array.swapaxes) - take = _propagate_cosmo_array_attributes(unyt_array.take) - transpose = _propagate_cosmo_array_attributes(unyt_array.transpose) - view = _propagate_cosmo_array_attributes(unyt_array.view) - - # Also wrap some array "attributes": - - @property - def T(self): - return self.transpose() # transpose is wrapped above. - - @property - def ua(self): - return _propagate_cosmo_array_attributes(np.ones_like)(self) + astype = _propagate_cosmo_array_attributes_to_result(unyt_array.astype) + in_units = _propagate_cosmo_array_attributes_to_result(unyt_array.in_units) + byteswap = _propagate_cosmo_array_attributes_to_result(unyt_array.byteswap) + compress = _propagate_cosmo_array_attributes_to_result(unyt_array.compress) + diagonal = _propagate_cosmo_array_attributes_to_result(unyt_array.diagonal) + flatten = _propagate_cosmo_array_attributes_to_result(unyt_array.flatten) + ravel = _propagate_cosmo_array_attributes_to_result(unyt_array.ravel) + repeat = _propagate_cosmo_array_attributes_to_result(unyt_array.repeat) + swapaxes = _propagate_cosmo_array_attributes_to_result(unyt_array.swapaxes) + transpose = _propagate_cosmo_array_attributes_to_result(unyt_array.transpose) + view = _propagate_cosmo_array_attributes_to_result(unyt_array.view) + __copy__ = _propagate_cosmo_array_attributes_to_result(unyt_array.__copy__) + __deepcopy__ = _propagate_cosmo_array_attributes_to_result(unyt_array.__deepcopy__) + in_cgs = _propagate_cosmo_array_attributes_to_result(unyt_array.in_cgs) + take = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_array.take) + ) + reshape = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_array.reshape) + ) + __getitem__ = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_array.__getitem__) + ) + dot = _default_binary_wrapper(unyt_array.dot, _multiply_cosmo_factor) - @property - def unit_array(self): - return _propagate_cosmo_array_attributes(np.ones_like)(self) + # Also wrap some array "properties": + T = property(_propagate_cosmo_array_attributes_to_result(unyt_array.transpose)) + ua = property(_propagate_cosmo_array_attributes_to_result(np.ones_like)) + unit_array = property(_propagate_cosmo_array_attributes_to_result(np.ones_like)) def convert_to_comoving(self) -> None: """ - Convert the internal data to be in comoving units. + Convert the internal data in-place to be in comoving units. """ if self.comoving: return @@ -930,7 +1411,7 @@ def convert_to_comoving(self) -> None: def convert_to_physical(self) -> None: """ - Convert the internal data to be in physical units. + Convert the internal data in-place to be in physical units. """ if self.comoving is None: raise InvalidConversionError @@ -943,28 +1424,28 @@ def convert_to_physical(self) -> None: values *= self.cosmo_factor.a_factor self.comoving = False - def to_physical(self): + def to_physical(self) -> "cosmo_array": """ Creates a copy of the data in physical units. Returns ------- - cosmo_array - copy of cosmo_array in physical units + out : swiftsimio.objects.cosmo_array + Copy of this array in physical units. """ copied_data = self.in_units(self.units, cosmo_factor=self.cosmo_factor) copied_data.convert_to_physical() return copied_data - def to_comoving(self): + def to_comoving(self) -> "cosmo_array": """ Creates a copy of the data in comoving units. Returns ------- - cosmo_array - copy of cosmo_array in comoving units + out : swiftsimio.objects.cosmo_array + Copy of this array in comoving units """ if not self.valid_transform: raise InvalidConversionError @@ -973,60 +1454,82 @@ def to_comoving(self): return copied_data - def compatible_with_comoving(self): + def compatible_with_comoving(self) -> bool: """ - Is this cosmo_array compatible with a comoving cosmo_array? + Is this :class:`~swiftsimio.objects.cosmo_array` compatible with a comoving + :class:`~swiftsimio.objects.cosmo_array`? - This is the case if the cosmo_array is comoving, or if the scale factor - exponent is 0 (cosmo_factor.a_factor() == 1) + This is the case if the :class:`~swiftsimio.objects.cosmo_array` is comoving, or + if the scale factor exponent is 0, or the scale factor is 1 + (either case satisfies ``cosmo_factor.a_factor() == 1``). + + Returns + ------- + out : bool + ``True`` if compatible, ``False`` otherwise. """ return self.comoving or (self.cosmo_factor.a_factor == 1.0) - def compatible_with_physical(self): + def compatible_with_physical(self) -> bool: """ - Is this cosmo_array compatible with a physical cosmo_array? + Is this :class:`~swiftsimio.objects.cosmo_array` compatible with a physical + :class:`~swiftsimio.objects.cosmo_array`? + + This is the case if the :class:`~swiftsimio.objects.cosmo_array` is physical, or + if the scale factor exponent is 0, or the scale factor is 1 + (either case satisfies ``cosmo_factor.a_factor() == 1``). - This is the case if the cosmo_array is physical, or if the scale factor - exponent is 0 (cosmo_factor.a_factor == 1) + Returns + ------- + out : bool + ``True`` if compatible, ``False`` otherwise. """ return (not self.comoving) or (self.cosmo_factor.a_factor == 1.0) @classmethod def from_astropy( cls, - arr, - unit_registry=None, - comoving=None, - cosmo_factor=None, - compression=None, - valid_transform=True, - ): + arr: "astropy.units.quantity.Quantity", + unit_registry: unyt.unit_registry.UnitRegistry = None, + comoving: bool = None, + cosmo_factor: cosmo_factor = cosmo_factor(None, None), + compression: str = None, + valid_transform: bool = True, + ) -> "cosmo_array": """ - Convert an AstroPy "Quantity" to a cosmo_array. + Convert an :class:`astropy.units.quantity.Quantity` to a + :class:`~swiftsimio.objects.cosmo_array`. Parameters ---------- - arr: AstroPy Quantity - The Quantity to convert from. - unit_registry: yt UnitRegistry, optional - A yt unit registry to use in the conversion. If one is not supplied, the + arr: astropy.units.quantity.Quantity + The quantity to convert from. + unit_registry : unyt.unit_registry.UnitRegistry, optional + A unyt registry to use in the conversion. If one is not supplied, the default one will be used. comoving : bool - if True then the array is in comoving co-ordinates, and if False then it is in - physical units. - cosmo_factor : float - Object to store conversion data between comoving and physical coordinates + Flag to indicate whether using comoving coordinates. + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. compression : string - String describing any compression that was applied to this array in the hdf5 - file. + Description of the compression filters that were applied to that array in the + hdf5 file. valid_transform : bool - flag to indicate whether this array can be converted to comoving + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + A cosmology-aware array. Example ------- - >>> from astropy.units import kpc - >>> cosmo_array.from_astropy([1, 2, 3] * kpc) - cosmo_array([1., 2., 3.], 'kpc') + :: + + >>> from astropy.units import kpc + >>> cosmo_array.from_astropy([1, 2, 3] * kpc) + cosmo_array([1., 2., 3.], 'kpc') """ obj = super().from_astropy(arr, unit_registry=unit_registry).view(cls) @@ -1040,46 +1543,54 @@ def from_astropy( @classmethod def from_pint( cls, - arr, - unit_registry=None, - comoving=None, - cosmo_factor=None, - compression=None, - valid_transform=True, - ): + arr: "pint.registry.Quantity", + unit_registry: unyt.unit_registry.UnitRegistry = None, + comoving: bool = None, + cosmo_factor: cosmo_factor = cosmo_factor(None, None), + compression: str = None, + valid_transform: bool = True, + ) -> "cosmo_array": """ - Convert a Pint "Quantity" to a cosmo_array. + Convert a :class:`pint.registry.Quantity` to a + :class:`~swiftsimio.objects.cosmo_array`. Parameters ---------- - arr : Pint Quantity - The Quantity to convert from. - unit_registry : yt UnitRegistry, optional - A yt unit registry to use in the conversion. If one is not - supplied, the default one will be used. + arr: pint.registry.Quantity + The quantity to convert from. + unit_registry : unyt.unit_registry.UnitRegistry, optional + A unyt registry to use in the conversion. If one is not supplied, the + default one will be used. comoving : bool - if True then the array is in comoving co-ordinates, and if False then it is in - physical units. - cosmo_factor : float - Object to store conversion data between comoving and physical coordinates + Flag to indicate whether using comoving coordinates. + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. compression : string - String describing any compression that was applied to this array in the hdf5 - file. + Description of the compression filters that were applied to that array in the + hdf5 file. valid_transform : bool - flag to indicate whether this array can be converted to comoving + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + A cosmology-aware array. Examples -------- - >>> from pint import UnitRegistry - >>> import numpy as np - >>> ureg = UnitRegistry() - >>> a = np.arange(4) - >>> b = ureg.Quantity(a, "erg/cm**3") - >>> b - - >>> c = cosmo_array.from_pint(b) - >>> c - cosmo_array([0, 1, 2, 3], 'erg/cm**3') + :: + + >>> from pint import UnitRegistry + >>> import numpy as np + >>> ureg = UnitRegistry() + >>> a = np.arange(4) + >>> b = ureg.Quantity(a, "erg/cm**3") + >>> b + + >>> c = cosmo_array.from_pint(b) + >>> c + cosmo_array([0, 1, 2, 3], 'erg/cm**3') """ obj = super().from_pint(arr, unit_registry=unit_registry).view(cls) obj.comoving = comoving @@ -1089,41 +1600,41 @@ def from_pint( return obj - # TODO: - def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): - cms = [ - (hasattr(inp, "comoving"), getattr(inp, "comoving", None)) for inp in inputs - ] - cfs = [ - (hasattr(inp, "cosmo_factor"), getattr(inp, "cosmo_factor", None)) - for inp in inputs - ] - comps = [ - (hasattr(inp, "compression"), getattr(inp, "compression", None)) - for inp in inputs - ] - - # if we're here at least one input must be a cosmo_array - if all([cm[1] for cm in cms if cm[0]]): - # all cosmo_array inputs are comoving - ret_cm = True - elif not any([cm[1] for cm in cms if cm[0]]): - # all cosmo_array inputs are physical - ret_cm = False - else: - # mix of comoving and physical inputs - inputs = [ - inp.to_comoving() if cm[0] and not cm[1] else inp - for inp, cm in zip(inputs, cms) - ] - ret_cm = True - - if len(set(comps)) == 1: - # all compressions identical, preserve it - ret_comp = comps[0] - else: - # mixed compressions, strip it off - ret_comp = None + def __array_ufunc__( + self, ufunc: np.ufunc, method: str, *inputs, **kwargs + ) -> object: + """ + Handles :mod:`numpy` ufunc calls on :class:`~swiftsimio.objects.cosmo_array` + input. + + :mod:`numpy` facilitates wrapping array classes by handing off to this function + when a function of :class:`numpy.ufunc` type is called with arguments from an + inheriting array class. Since we inherit from :class:`~unyt.array.unyt_array`, + we let :mod:`unyt` handle what to do with the units and take care of processing + the cosmology information via our helper functions. + + Parameters + ---------- + ufunc : numpy.ufunc + The numpy function being called. + + method : str, optional + Some ufuncs have methods accessed as attributes, such as ``"reduce"``. + If using such a method, this argument receives its name. + + inputs : tuple + Arguments to the ufunc. + + kwargs : dict + Keyword arguments to the ufunc. + + Returns + ------- + out : object + The result of the ufunc call, with our cosmology attribute processing applied. + """ + helper_result = _prepare_array_func_args(*inputs, **kwargs) + cfs = helper_result["cfs"] # make sure we evaluate the cosmo_factor_ufunc_registry function: # might raise/warn even if we're not returning a cosmo_array @@ -1131,48 +1642,422 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): power_map = POWER_MAPPING[ufunc] if "axis" in kwargs and kwargs["axis"] is not None: ret_cf = _power_cosmo_factor( - cfs[0], - (False, None), - power=power_map(inputs[0].shape[kwargs["axis"]]), + cfs[0], None, power=power_map(inputs[0].shape[kwargs["axis"]]) ) else: ret_cf = _power_cosmo_factor( - cfs[0], (False, None), power=power_map(inputs[0].size) + cfs[0], None, power=power_map(inputs[0].size) ) + elif ( + ufunc in (logical_and, logical_or, logical_xor, logical_not) + and method == "reduce" + ): + ret_cf = _return_without_cosmo_factor(cfs[0]) else: ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs) - ret = super().__array_ufunc__(ufunc, method, *inputs, **kwargs) + ret = _ensure_result_is_cosmo_array_or_quantity(super().__array_ufunc__)( + ufunc, method, *inputs, **kwargs + ) # if we get a tuple we have multiple return values to deal with - # if unyt returns a bare ndarray, do the same - # otherwise we create a view and attach our attributes if isinstance(ret, tuple): - ret = tuple( - r.view(type(self)) if isinstance(r, unyt_array) else r for r in ret - ) for r in ret: - if isinstance(r, type(self)): - r.comoving = ret_cm + if isinstance(r, cosmo_array): # also recognizes cosmo_quantity + r.comoving = helper_result["comoving"] r.cosmo_factor = ret_cf - r.compression = ret_comp - if isinstance(ret, unyt_array): - ret = ret.view(type(self)) - ret.comoving = ret_cm + r.compression = helper_result["compression"] + elif isinstance(ret, cosmo_array): # also recognizes cosmo_quantity + ret.comoving = helper_result["comoving"] ret.cosmo_factor = ret_cf - ret.compression = ret_comp + ret.compression = helper_result["compression"] if "out" in kwargs: out = kwargs.pop("out") if ufunc not in multiple_output_operators: out = out[0] - if isinstance(out, cosmo_array): - out.comoving = ret_cm + if isinstance(out, cosmo_array): # also recognizes cosmo_quantity + out.comoving = helper_result["comoving"] out.cosmo_factor = ret_cf - out.compression = ret_comp + out.compression = helper_result["compression"] else: for o in out: - if isinstance(o, type(self)): - o.comoving = ret_cm + if isinstance(o, cosmo_array): # also recognizes cosmo_quantity + o.comoving = helper_result["comoving"] o.cosmo_factor = ret_cf - o.compression = ret_comp + o.compression = helper_result["compression"] + + return ret + def __array_function__( + self, func: Callable, types: Collection, args: tuple, kwargs: dict + ): + """ + Handles :mod:`numpy` function calls on :class:`~swiftsimio.objects.cosmo_array` + input. + + :mod:`numpy` facilitates wrapping array classes by handing off to this function + when a numpy-defined function is called with arguments from an + inheriting array class. Since we inherit from :class:`~unyt.array.unyt_array`, + we let :mod:`unyt` handle what to do with the units and take care of processing + the cosmology information via our helper functions. + + Parameters + ---------- + func : Callable + The numpy function being called. + + types : collections.abc.Collection + A collection of unique argument types from the original :mod:`numpy` function + call that implement ``__array_function__``. + + args : tuple + Arguments to the functions. + + kwargs : dict + Keyword arguments to the function. + + Returns + ------- + out : object + The result of the ufunc call, with our cosmology attribute processing applied. + """ + # Follow NEP 18 guidelines + # https://numpy.org/neps/nep-0018-array-function-protocol.html + from ._array_functions import _HANDLED_FUNCTIONS + from unyt._array_functions import ( + _HANDLED_FUNCTIONS as _UNYT_HANDLED_FUNCTIONS, + _UNSUPPORTED_FUNCTIONS as _UNYT_UNSUPPORTED_FUNCTIONS, + ) + + # Let's claim to support everything supported by unyt. + # If we can't do this in future, follow their pattern of + # defining out own _UNSUPPORTED_FUNCTIONS in a _array_functions.py file + _UNSUPPORTED_FUNCTIONS = _UNYT_UNSUPPORTED_FUNCTIONS + + if func in _UNSUPPORTED_FUNCTIONS: + # following NEP 18, return NotImplemented as a sentinel value + # which will lead to raising a TypeError, while + # leaving other arguments a chance to take the lead + return NotImplemented + + if not all(issubclass(t, cosmo_array) or t is np.ndarray for t in types): + # Note: this allows subclasses that don't override + # __array_function__ to handle cosmo_array objects + return NotImplemented + + if func in _HANDLED_FUNCTIONS: + function_to_invoke = _HANDLED_FUNCTIONS[func] + elif func in _UNYT_HANDLED_FUNCTIONS: + function_to_invoke = _UNYT_HANDLED_FUNCTIONS[func] + else: + # default to numpy's private implementation + function_to_invoke = func._implementation + return function_to_invoke(*args, **kwargs) + + def __mul__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Multiply this :class:`~swiftsimio.objects.cosmo_array`. + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to multiply with this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the multiplication. + """ + if isinstance(b, unyt.unit_object.Unit): + retval = self.__copy__() + retval.units = retval.units * b + return retval + else: + return super().__mul__(b) + + def __rmul__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Multiply this :class:`~swiftsimio.objects.cosmo_array` (as the right argument). + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + .. note:: + + This function is never called when `b` is a :class:`unyt.unit_object.Unit` + because :mod:`unyt` handles the operation. This results in a silent demotion + to a :class:`unyt.array.unyt_array`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to multiply with this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the multiplication. + """ + if isinstance(b, unyt.unit_object.Unit): + return self.__mul__(b) + else: + return super().__rmul__(b) + + def __truediv__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Divide this :class:`~swiftsimio.objects.cosmo_array`. + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to divide this one by. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the division. + """ + if isinstance(b, unyt.unit_object.Unit): + return self.__mul__(1 / b) + else: + return super().__truediv__(b) + + def __rtruediv__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Divide this :class:`~swiftsimio.objects.cosmo_array` (as the right argument). + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + .. note:: + + This function is never called when `b` is a :class:`unyt.unit_object.Unit` + because :mod:`unyt` handles the operation. This results in a silent demotion + to a :class:`unyt.array.unyt_array`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to divide by this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the division. + """ + if isinstance(b, unyt.unit_object.Unit): + return (1 / self).__mul__(b) + else: + return super().__rtruediv__(b) + + +class cosmo_quantity(cosmo_array, unyt_quantity): + """ + Cosmology scalar class. + + This inherits from both the :class:`~swiftsimio.objects.cosmo_array` and the + :class:`~unyt.array.unyt_quantity`, and has the same four attributes as + :class:`~swiftsimio.objects.cosmo_array`: ``compression``, ``cosmo_factor``, + ``comoving``, and ``valid_transform``. + + Like :class:`unyt.array.unyt_quantity`, it is intended to hold a scalar value. + Values of this type will be returned by :mod:`numpy` functions that return + scalar values. + + Other than containing a scalar, functionality is identical to + :class:`~swiftsimio.objects.cosmo_array`. Refer to that class's documentation. + + Parameters + ---------- + input_scalar : float or unyt.array.unyt_quantity + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, powers must be specified using + python syntax (``cm**3``, not ``cm^3``). + registry : unyt.unit_registry.UnitRegistry, optional + The registry to create units from. If ``units`` is already associated + with a unit registry and this is specified, this will be used instead of the + registry associated with the unit object. + dtype : np.dtype or str, optional + The dtype of the array data. Defaults to the dtype of the input data, or, if + none is found, uses ``np.float64``. + bypass_validation : bool, optional + If ``True``, all input validation is skipped. Using this option may produce + corrupted or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. Defaults to ``False``. + name : str, optional + The name of the array. Defaults to ``None``. This attribute does not propagate + through mathematical operations, but is preserved under indexing and unit + conversions. + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. + comoving : bool + Flag to indicate whether using comoving coordinates. + valid_transform : bool + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. + compression : string + Description of the compression filters that were applied to that array in the + hdf5 file. + + Attributes + ---------- + comoving : bool + if True then the array is in comoving co-ordinates, and if + False then it is in physical units. + + cosmo_factor : float + Object to store conversion data between comoving and physical coordinates + + compression : string + String describing any compression that was applied to this array in the + hdf5 file. + + valid_transform: bool + if True then the array can be converted from physical to comoving units + """ + + def __new__( + cls, + input_scalar: numeric_type, + units: Optional[ + Union[str, unyt.unit_object.Unit, "astropy.units.core.Unit"] + ] = None, + *, + registry: Optional[unyt.unit_registry.UnitRegistry] = None, + dtype: Optional[Union[np.dtype, str]] = None, + bypass_validation: bool = False, + name: Optional[str] = None, + cosmo_factor: Optional[cosmo_factor] = None, + scale_factor: Optional[float] = None, + scale_exponent: Optional[float] = None, + comoving: Optional[bool] = None, + valid_transform: bool = True, + compression: Optional[str] = None, + ) -> "cosmo_quantity": + """ + Closely inspired by the :meth:`unyt.array.unyt_quantity.__new__` constructor. + + Parameters + ---------- + input_scalar : float or unyt.array.unyt_quantity + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, powers must be specified using + python syntax (``cm**3``, not ``cm^3``). + registry : unyt.unit_registry.UnitRegistry, optional + The registry to create units from. If ``units`` is already associated + with a unit registry and this is specified, this will be used instead of the + registry associated with the unit object. + dtype : np.dtype or str, optional + The dtype of the array data. Defaults to the dtype of the input data, or, if + none is found, uses ``np.float64``. + bypass_validation : bool, optional + If ``True``, all input validation is skipped. Using this option may produce + corrupted or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. Defaults to ``False``. + name : str, optional + The name of the array. Defaults to ``None``. This attribute does not propagate + through mathematical operations, but is preserved under indexing and unit + conversions. + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. + The same information can be provided using the ``scale_factor`` and + ``scale_exponent`` arguments, instead. + scale_factor : float + The scale factor associated to the data. Also provide a value for + ``scale_exponent``. + scale_exponent : int or float + The exponent for the scale factor giving the scaling for conversion to/from + comoving units. Also provide a value for ``scale_factor``. + comoving : bool + Flag to indicate whether using comoving coordinates. + valid_transform : bool + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. + compression : string + Description of the compression filters that were applied to that array in the + hdf5 file. + """ + if bypass_validation is True: + ret = super().__new__( + cls, + np.asarray(input_scalar), + units=units, + registry=registry, + dtype=dtype, + bypass_validation=bypass_validation, + name=name, + cosmo_factor=cosmo_factor, + scale_factor=scale_factor, + scale_exponent=scale_exponent, + comoving=comoving, + valid_transform=valid_transform, + compression=compression, + ) + + if not isinstance(input_scalar, (numeric_type, np.number, np.ndarray)): + raise RuntimeError("cosmo_quantity values must be numeric") + + # Use values from kwargs, if None use values from input_scalar + units = getattr(input_scalar, "units", None) if units is None else units + name = getattr(input_scalar, "name", None) if name is None else name + cosmo_factor = ( + getattr(input_scalar, "cosmo_factor", None) + if cosmo_factor is None + else cosmo_factor + ) + comoving = ( + getattr(input_scalar, "comoving", None) if comoving is None else comoving + ) + valid_transform = ( + getattr(input_scalar, "valid_transform", None) + if valid_transform is None + else valid_transform + ) + compression = ( + getattr(input_scalar, "compression", None) + if compression is None + else compression + ) + ret = super().__new__( + cls, + np.asarray(input_scalar), + units=units, + registry=registry, + dtype=dtype, + bypass_validation=bypass_validation, + name=name, + cosmo_factor=cosmo_factor, + scale_factor=scale_factor, + scale_exponent=scale_exponent, + comoving=comoving, + valid_transform=valid_transform, + compression=compression, + ) + if ret.size > 1: + raise RuntimeError("cosmo_quantity instances must be scalars") return ret + + __round__ = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_quantity.__round__) + ) diff --git a/swiftsimio/reader.py b/swiftsimio/reader.py index f83db4aa..07aee825 100644 --- a/swiftsimio/reader.py +++ b/swiftsimio/reader.py @@ -172,8 +172,8 @@ def getter(self): self, f"_{name}", cosmo_array( - # Only use column data if array is multidimensional, otherwise - # we will crash here + # Only use column data if array is multidimensional, + # otherwise we will crash here ( handle[field][:, columns] if handle[field].ndim > 1 @@ -366,7 +366,10 @@ def __init__(self, field_path: str, named_columns: List[str], name: str): return def __str__(self): - return f'Named columns instance with {self.named_columns} available for "{self.name}"' + return ( + f"Named columns instance with {self.named_columns} available " + f'for "{self.name}"' + ) def __repr__(self): return self.__str__() diff --git a/swiftsimio/snapshot_writer.py b/swiftsimio/snapshot_writer.py index 2d9a4206..9315ed64 100644 --- a/swiftsimio/snapshot_writer.py +++ b/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. @@ -162,11 +163,7 @@ def check_consistent(self) -> bool: return True - def generate_smoothing_lengths( - self, - boxsize: Union[List[unyt.unyt_quantity], unyt.unyt_quantity], - dimension: int, - ): + def generate_smoothing_lengths(self, boxsize: cosmo_array, dimension: int): """ Automatically generates the smoothing lengths as 2 * the mean interparticle separation. @@ -175,7 +172,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 @@ -506,7 +503,7 @@ class SWIFTSnapshotWriter(object): def __init__( self, unit_system: Union[unyt.UnitSystem, str], - box_size: Union[list, unyt.unyt_quantity], + boxsize: cosmo_array, dimension=3, compress=True, extra_header: Union[None, dict] = None, @@ -522,7 +519,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 @@ -545,13 +542,13 @@ def __init__( # Validate the boxsize and convert to our units. try: - for x in box_size: + for x in boxsize: x.convert_to_base(self.unit_system) - self.box_size = box_size + self.boxsize = boxsize except TypeError: # This is just a single number (i.e. uniform in all dimensions) - box_size.convert_to_base(self.unit_system) - self.box_size = box_size + boxsize.convert_to_base(self.unit_system) + self.boxsize = boxsize self.dimension = dimension self.compress = compress @@ -637,7 +634,7 @@ def _write_metadata(self, handle: h5py.File, names_to_write: List): mass_table[_ptype_str_to_int(number)] = getattr(self, name).masses[0] attrs = { - "BoxSize": self.box_size, + "BoxSize": self.boxsize, "NumPart_Total": number_of_particles, "NumPart_Total_HighWord": [0] * 6, "Flag_Entropy_ICs": 0, diff --git a/swiftsimio/swiftsnap.py b/swiftsimio/swiftsnap.py index 3c586460..ac74b601 100755 --- a/swiftsimio/swiftsnap.py +++ b/swiftsimio/swiftsnap.py @@ -64,7 +64,6 @@ def decode(bytestring: bytes) -> str: def swiftsnap(): import swiftsimio as sw from swiftsimio.metadata.objects import metadata_discriminator - import unyt from swiftsimio.metadata.particle import particle_name_underscores from textwrap import wrap diff --git a/swiftsimio/visualisation/__init__.py b/swiftsimio/visualisation/__init__.py index 05928f1b..b496dede 100644 --- a/swiftsimio/visualisation/__init__.py +++ b/swiftsimio/visualisation/__init__.py @@ -2,7 +2,8 @@ Visualisation sub-module for swiftismio. """ -from .projection import scatter, project_gas, project_gas_pixel_grid -from .slice import slice_scatter as slice -from .slice import slice_gas, slice_gas_pixel_grid +from .projection import project_gas +from .slice import slice_gas +from .volume_render import render_gas +from .ray_trace import panel_gas from .smoothing_length import generate_smoothing_lengths diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py new file mode 100644 index 00000000..f71b0a8e --- /dev/null +++ b/swiftsimio/visualisation/_vistools.py @@ -0,0 +1,139 @@ +import numpy as np +from warnings import warn +from swiftsimio.objects import cosmo_array +from swiftsimio._array_functions import _copy_cosmo_array_attributes_if_present + + +def _get_projection_field(data, field_name): + return ( + getattr(data, field_name) + if field_name is not None + else np.ones_like(data.particle_ids) + ) + + +def _get_region_info(data, region, z_slice=None, require_cubic=False, periodic=True): + boxsize = data.metadata.boxsize + if region is not None: + region = cosmo_array(region) + if data.coordinates.comoving: + boxsize.convert_to_comoving() + if region is not None: + region.convert_to_comoving() + elif data.coordinates.comoving is False: # compare to False in case None + boxsize.convert_to_physical() + if region is not None: + region.convert_to_physical() + z_slice_included = z_slice is not None + if not z_slice_included: + z_slice = np.zeros_like(boxsize[0]) + box_x, box_y, box_z = boxsize + if region is not None: + x_min, x_max, y_min, y_max = region[:4] + if len(region) == 6: + z_slice_included = True + z_min, z_max = region[4:] + else: + z_min, z_max = np.zeros_like(box_z), box_z + else: + x_min, x_max = np.zeros_like(box_x), box_x + y_min, y_max = np.zeros_like(box_y), box_y + z_min, z_max = np.zeros_like(box_z), box_z + + if z_slice_included and (z_slice > box_z) or (z_slice < np.zeros_like(box_z)): + raise ValueError("Please enter a slice value inside the box.") + + x_range = x_max - x_min + y_range = y_max - y_min + z_range = z_max - z_min + max_range = np.r_[x_range, y_range].max() + + if require_cubic and not ( + np.isclose(x_range, y_range) and np.isclose(x_range, z_range) + ): + raise AttributeError( + "Projection code is currently not able to handle non-cubic images." + ) + + periodic_box_x, periodic_box_y, periodic_box_z = ( + boxsize / max_range if periodic else np.zeros(3) + ) + + return { + "x_min": x_min, + "x_max": x_max, + "y_min": y_min, + "y_max": y_max, + "z_min": z_min, + "z_max": z_max, + "x_range": x_range, + "y_range": y_range, + "z_range": z_range, + "max_range": max_range, + "z_slice_included": z_slice_included, + "periodic_box_x": periodic_box_x, + "periodic_box_y": periodic_box_y, + "periodic_box_z": periodic_box_z, + } + + +def _get_rotated_coordinates(data, rotation_matrix, rotation_center): + if rotation_center is not None: + if data.coordinates.comoving: + rotation_center = rotation_center.to_comoving() + elif data.coordinates.comoving is False: + rotation_center = rotation_center.to_physical() + # Rotate co-ordinates as required + x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) + + x += rotation_center[0] + y += rotation_center[1] + z += rotation_center[2] + else: + x, y, z = data.coordinates.T + return x, y, z + + +def backend_restore_cosmo_and_units(backend_func, norm=1.0): + def wrapper(*args, **kwargs): + comoving = getattr(kwargs["m"], "comoving", None) + if comoving is True: + if kwargs["x"].comoving is False or kwargs["y"].comoving is False: + warn( + "Projecting a comoving quantity with physical input for coordinates. " + "Converting coordinate grid to comoving." + ) + kwargs["x"].convert_to_comoving() + kwargs["y"].convert_to_comoving() + if kwargs["h"].comoving is False: + warn( + "Projecting a comoving quantity with physical input for smoothing " + "lengths. Converting smoothing lengths to comoving." + ) + kwargs["h"].convert_to_comoving() + norm.convert_to_comoving() + elif comoving is False: # don't use else in case None + if kwargs["x"].comoving or kwargs["y"].comoving: + warn( + "Projecting a physical quantity with comoving input for coordinates. " + "Converting coordinate grid to physical." + ) + kwargs["x"].convert_to_physical() + kwargs["y"].convert_to_physical() + if kwargs["h"].comoving: + warn( + "Projecting a physical quantity with comoving input for smoothing " + "lengths. Converting smoothing lengths to physical." + ) + kwargs["h"].convert_to_physical() + norm.convert_to_physical() + return ( + _copy_cosmo_array_attributes_if_present( + kwargs["m"], + backend_func(*args, **kwargs).view(cosmo_array), + copy_units=True, + ) + / norm + ) + + return wrapper diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index cdd56f17..42a870eb 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -2,14 +2,13 @@ Tools for creating power spectra from SWIFT data. """ -from numpy import float32, float64, int32, zeros, ndarray +from numpy import float32, float64, int32, zeros, ndarray, zeros_like import numpy as np import scipy.fft -import unyt from swiftsimio.optional_packages import tqdm from swiftsimio.accelerated import jit, NUM_THREADS, prange -from swiftsimio import cosmo_array +from swiftsimio import cosmo_array, cosmo_quantity from swiftsimio.reader import __SWIFTGroupDataset from typing import Optional, Dict, Tuple @@ -206,18 +205,20 @@ def render_to_deposit( if positions.comoving: if not quantity.compatible_with_comoving(): raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' + f'Physical quantity "{project}" is not compatible with comoving ' + "coordinates!" ) else: if not quantity.compatible_with_physical(): raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' + f'Comoving quantity "{project}" is not compatible with physical ' + "coordinates!" ) # Get the box size - box_size = data.metadata.boxsize + boxsize = data.metadata.boxsize - if not box_size.units == positions.units: + if not boxsize.units == positions.units: raise AttributeError("Box size and positions have different units!") # Deposit the particles @@ -228,9 +229,9 @@ def render_to_deposit( m=quantity.v, res=resolution, fold=folding, - box_x=box_size[0].v, - box_y=box_size[1].v, - box_z=box_size[2].v, + box_x=boxsize[0].v, + box_y=boxsize[1].v, + box_z=boxsize[2].v, ) if parallel: @@ -256,10 +257,10 @@ def render_to_deposit( def folded_depositions_to_power_spectrum( depositions: Dict[int, cosmo_array], - box_size: cosmo_array, + boxsize: cosmo_array, number_of_wavenumber_bins: int, cross_depositions: Optional[Dict[int, cosmo_array]] = None, - wavenumber_range: Optional[Tuple[unyt.unyt_quantity]] = None, + wavenumber_range: Optional[Tuple[cosmo_quantity]] = None, log_wavenumber_bins: bool = True, workers: Optional[int] = None, minimal_sample_modes: Optional[int] = 0, @@ -267,7 +268,7 @@ def folded_depositions_to_power_spectrum( track_progress: bool = False, transition: str = "simple", shot_noise_norm: Optional[float] = None, -) -> Tuple[unyt.unyt_array]: +) -> Tuple[cosmo_array]: """ Convert some folded depositions to power spectra. @@ -278,13 +279,13 @@ def folded_depositions_to_power_spectrum( Dictionary of depositions, where the key is the base folding. So that would be 0 for no folding, 1 for a half-box-size folding, 2 for a quarter, etc. The 'real' folding is 2 ** depositions.keys(). - box_size: cosmo_array + boxsize: cosmo_array The box size of the deposition, from the dataset. number_of_wavenumber_bins: int The number of bins to use in the power spectrum. cross_depositions: Optional[dict[int, cosmo_array]] An optional dictionary of cross-depositions, where the key is the folding. - wavenumber_range: Optional[tuple[unyt.unyt_quantity]] + wavenumber_range: Optional[tuple[cosmo_quantity]] The range of wavenumbers to use. Officially optional, but is required for now. log_wavenumber_bins: bool @@ -313,13 +314,13 @@ def folded_depositions_to_power_spectrum( Returns ------- - wavenumber_bins: unyt.unyt_array[float32] + wavenumber_bins: cosmo_array[float32] The wavenumber bins. - wavenumber_centers: unyt.unyt_array[float32] + wavenumber_centers: cosmo_array[float32] The centers of the wavenumber bins. - power_spectrum: unyt.unyt_array[float32] + power_spectrum: cosmo_array[float32] The power spectrum. folding_tracker: np.array @@ -336,32 +337,28 @@ def folded_depositions_to_power_spectrum( raise NotImplementedError if log_wavenumber_bins: - wavenumber_bins = unyt.unyt_array( - np.logspace( - np.log10(min(wavenumber_range).v), - np.log10(max(wavenumber_range).v), - number_of_wavenumber_bins + 1, - ), - wavenumber_range[0].units, - name="Wavenumber bins", + wavenumber_bins = np.geomspace( + np.min(cosmo_array(wavenumber_range)), + np.max(cosmo_array(wavenumber_range)), + number_of_wavenumber_bins + 1, ) + else: - wavenumber_bins = unyt.unyt_array( - np.linspace( - min(wavenumber_range).v, - max(wavenumber_range).v, - number_of_wavenumber_bins + 1, - ), - wavenumber_range[0].units, - name="Wavenumber bins", + wavenumber_bins = np.linspace( + np.min(cosmo_array(wavenumber_range)), + np.max(cosmo_array(wavenumber_range)), + number_of_wavenumber_bins + 1, ) + wavenumber_bins.name = "Wavenumber bins" wavenumber_centers = 0.5 * (wavenumber_bins[1:] + wavenumber_bins[:-1]) wavenumber_centers.name = r"Wavenumbers $k$" - box_volume = np.prod(box_size) - power_spectrum = unyt.unyt_array( + box_volume = np.prod(boxsize) + power_spectrum = cosmo_array( np.zeros(number_of_wavenumber_bins), units=box_volume.units, + comoving=box_volume.comoving, + cosmo_factor=box_volume.cosmo_factor ** -1, name="Power spectrum $P(k)$", ) folding_tracker = np.ones(number_of_wavenumber_bins, dtype=float) @@ -382,7 +379,7 @@ def folded_depositions_to_power_spectrum( folded_counts, ) = deposition_to_power_spectrum( deposition=depositions[folding], - box_size=box_size, + boxsize=boxsize, folding=folding, wavenumber_bins=wavenumber_bins, workers=workers, @@ -399,14 +396,14 @@ def folded_depositions_to_power_spectrum( if folding != final_folding: cutoff_wavenumber = ( - 2.0 ** folding * np.min(depositions[folding].shape) / np.min(box_size) + 2.0 ** folding * np.min(depositions[folding].shape) / np.min(boxsize) ) if cutoff_above_wavenumber_fraction is not None: maximally_sampled_wavenumber = np.max( folded_wavenumber_centers[use_bins] ) - cutoff_wavenumber = min( + cutoff_wavenumber = np.min( cutoff_above_wavenumber_fraction * maximally_sampled_wavenumber, cutoff_above_wavenumber_fraction * cutoff_wavenumber, ) @@ -474,31 +471,31 @@ def folded_depositions_to_power_spectrum( def deposition_to_power_spectrum( - deposition: unyt.unyt_array, - box_size: cosmo_array, + deposition: cosmo_array, + boxsize: cosmo_array, folding: int = 0, - cross_deposition: Optional[unyt.unyt_array] = None, - wavenumber_bins: Optional[unyt.unyt_array] = None, + cross_deposition: Optional[cosmo_array] = None, + wavenumber_bins: Optional[cosmo_array] = None, workers: Optional[int] = None, shot_noise_norm: Optional[float] = None, -) -> Tuple[unyt.unyt_array]: +) -> Tuple[cosmo_array]: """ Convert a deposition to a power spectrum, by default using a linear binning strategy. Parameters ---------- - deposition: unyt.unyt_array[float32, float32, float32] + deposition: ~swiftsimio.objects.cosmo_array[float32, float32, float32] The deposition to convert to a power spectrum. - box_size: cosmo_array + boxsize: ~swiftsimio.objects.cosmo_array The box size of the deposition, from the dataset. folding: int The folding number (i.e. box-size is divided by 2^folding) that was used here. - cross_deposition: unyt.unyt_array[float32, float32, float32] + cross_deposition: ~swiftsimio.objects.cosmo_array[float32, float32, float32] An optional second deposition to cross-correlate with the first. If not provided, we assume you want an auto-spectrum. - wavenumber_bins: unyt.unyt_array[float32], optional + wavenumber_bins: ~swiftsimio.objects.cosmo_array[float32], optional Optionally you can provide the specific bins that you would like to use. workers: Optional[int] The number of threads to use. @@ -509,12 +506,12 @@ def deposition_to_power_spectrum( Returns ------- - wavenumber_centers: unyt.unyt_array[float32] + wavenumber_centers: ~swiftsimio.objects.cosmo_array[float32] The k-values of the power spectrum, with units. These are the real bin centers, calculated from the mean value of k that was used in the binning process. - power_spectrum: unyt.unyt_array[float32] + power_spectrum: ~swiftsimio.objects.cosmo_array[float32] The power spectrum, with units. binned_counts: np.array[int32] @@ -540,18 +537,18 @@ def deposition_to_power_spectrum( folding = 2.0 ** folding - box_size_folded = box_size[0] / folding + boxsize_folded = boxsize[0] / folding npix = deposition.shape[0] - mean_deposition = np.mean(deposition.v) - overdensity = (deposition.v - mean_deposition) / mean_deposition + mean_deposition = np.mean(deposition) + overdensity = (deposition - mean_deposition) / mean_deposition fft = scipy.fft.fftn(overdensity / np.prod(deposition.shape), workers=workers) if cross_deposition is not None: - mean_cross_deposition = np.mean(cross_deposition.v) + mean_cross_deposition = np.mean(cross_deposition) cross_overdensity = ( - cross_deposition.v - mean_cross_deposition + cross_deposition - mean_cross_deposition ) / mean_cross_deposition conj_fft = scipy.fft.fftn( @@ -560,10 +557,10 @@ def deposition_to_power_spectrum( else: conj_fft = fft.conj() - fourier_amplitudes = (fft * conj_fft).real * box_size_folded ** 3 + fourier_amplitudes = (fft * conj_fft).real * boxsize_folded ** 3 # Calculate k-value spacing (centered FFT) - dk = 2 * np.pi / (box_size_folded) + dk = 2 * np.pi / (boxsize_folded) # Create k-values array (adjust range based on your needs) kfreq = np.fft.fftfreq(npix, d=1 / dk) * npix @@ -578,7 +575,7 @@ def deposition_to_power_spectrum( else: kbins = wavenumber_bins - binned_amplitudes = np.histogram(knrm, bins=kbins, weights=fourier_amplitudes.v)[0] + binned_amplitudes = np.histogram(knrm, bins=kbins, weights=fourier_amplitudes)[0] binned_counts = np.histogram(knrm, bins=kbins)[0] # Also compute the 'real' average wavenumber point contributing to this bin. binned_wavenumbers = np.histogram(knrm, bins=kbins, weights=knrm)[0] @@ -595,18 +592,15 @@ def deposition_to_power_spectrum( binned_amplitudes *= folding ** 3 # Correct units and names - wavenumbers = unyt.unyt_array( - binned_wavenumbers / divisor, units=knrm.units, name="Wavenumber $k$" - ) + wavenumbers = binned_wavenumbers / divisor + wavenumbers.name = "Wavenumber $k$" shot_noise = ( - (box_size[0] ** 3 / shot_noise_norm) if shot_noise_norm is not None else 0.0 - ) - - power_spectrum = ( - unyt.unyt_array((binned_amplitudes) / divisor, units=fourier_amplitudes.units) - - shot_noise + (boxsize[0] ** 3 / shot_noise_norm) + if shot_noise_norm is not None + else zeros_like(boxsize[0] ** 3) # copy cosmo properties ) + power_spectrum = (binned_amplitudes / divisor) - shot_noise power_spectrum.name = "Power Spectrum $P(k)$" diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index fa0738b3..a615e035 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -3,53 +3,28 @@ """ 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_, -) -from unyt import unyt_array, unyt_quantity, exceptions +import numpy as np 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.smoothing_length import backends_get_hsml +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_info, + _get_rotated_coordinates, + backend_restore_cosmo_and_units, ) -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, resolution: int, project: Union[str, None] = "masses", - region: Union[None, unyt_array] = None, - mask: Union[None, array] = None, - rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, 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, cosmo_array] = None, parallel: bool = False, backend: str = "fast", periodic: bool = True, @@ -68,9 +43,6 @@ def project_pixel_grid( data: __SWIFTGroupDataset The SWIFT dataset that you wish to visualise (get this from ``load``) - boxsize: unyt_array - The box-size of the simulation. - resolution: int The resolution of the image. All images returned are square, ``res`` by ``res``, pixel grids. @@ -79,9 +51,10 @@ def project_pixel_grid( Variable to project to get the weighted density of. By default, this is mass. If you would like to mass-weight any other variable, you can always create it as ``data.gas.my_variable = data.gas.other_variable - * data.gas.masses``. + * data.gas.masses``. The result is comoving if this is comoving, else + it is physical. - 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,8 +89,9 @@ def project_pixel_grid( Returns ------- - image: unyt_array + image: cosmo_array Projected image with units of project / length^2, of size ``res`` x ``res``. + Comoving if ``project`` data are comoving, else physical. Notes @@ -130,254 +104,49 @@ def project_pixel_grid( array if you want it to be visualised the 'right way up'. """ - if rotation_center is not None: - try: - if rotation_center.units == data.coordinates.units: - pass - else: - raise exceptions.InvalidUnitOperation( - "Units of coordinates and rotation center must agree" - ) - except AttributeError: - raise exceptions.InvalidUnitOperation( - "Ensure that rotation_center is a unyt array with the same units as coordinates" - ) - - number_of_particles = data.coordinates.shape[0] - - if project is None: - m = ones(number_of_particles, dtype=float32) - else: - m = getattr(data, project) - if data.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value - - # This provides a default 'slice it all' mask. - if mask is None: - mask = s_[:] - - box_x, box_y, box_z = boxsize - - # Set the limits of the image. - z_slice_included = False - - if region is not None: - x_min, x_max, y_min, y_max = region[:4] - - if len(region) == 6: - z_slice_included = True - z_min, z_max = region[4:] - else: - z_min = unyt_quantity(0.0, units=box_z.units) - z_max = box_z - else: - x_min = unyt_quantity(0.0, units=box_x.units) - x_max = box_x - y_min = unyt_quantity(0.0, units=box_y.units) - y_max = box_y - - x_range = x_max - x_min - y_range = y_max - y_min - - # Deal with non-cubic boxes: - # we always use the maximum of x_range and y_range to normalise the coordinates - # empty pixels in the resulting square image are trimmed afterwards - max_range = max(x_range, y_range) - - try: - try: - hsml = data.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.smoothing_length - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - f"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!" - ) - except AttributeError: - # No hsml present. If they are using the 'histogram' backend, we - # should just mock them to be anything as it doesn't matter. - if backend == "histogram": - hsml = empty_like(m) - else: - raise AttributeError - - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = matmul(rotation_matrix, (data.coordinates - rotation_center).T) - - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - - else: - x, y, z = data.coordinates.T - - if z_slice_included: - combined_mask = logical_and(mask, logical_and(z <= z_max, z >= z_min)).astype( - bool - ) - else: - combined_mask = mask - - if periodic: - periodic_box_x = box_x / max_range - periodic_box_y = box_y / max_range - else: - periodic_box_x = 0.0 - periodic_box_y = 0.0 - - common_arguments = dict( - x=(x[combined_mask] - x_min) / max_range, - y=(y[combined_mask] - y_min) / max_range, - m=m[combined_mask], - h=hsml[combined_mask] / max_range, + m = _get_projection_field(data, project) + region_info = _get_region_info(data, region) + hsml = backends_get_hsml["sph" if backend != "histogram" else "histogram"](data) + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) + mask = mask if mask is not None else np.s_[...] + if not region_info["z_slice_included"]: + mask = np.logical_and( + mask, np.logical_and(z <= region_info["z_max"], z >= region_info["z_min"]) + ).astype(bool) + + kwargs = dict( + x=(x[mask] - region_info["x_min"]) / region_info["max_range"], + y=(y[mask] - region_info["y_min"]) / region_info["max_range"], + m=m[mask], + h=hsml[mask] / region_info["max_range"], res=resolution, - box_x=periodic_box_x, - box_y=periodic_box_y, + box_x=region_info["periodic_box_x"], + box_y=region_info["periodic_box_y"], ) - - if parallel: - image = backends_parallel[backend](**common_arguments) - else: - image = backends[backend](**common_arguments) + norm = region_info["x_range"] * region_info["y_range"] + backend_func = (backends_parallel if parallel else backends)[backend] + image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs) # determine the effective number of pixels for each dimension - xres = int(ceil(resolution * (x_range / max_range))) - yres = int(ceil(resolution * (y_range / max_range))) + xres = int( + np.ceil(resolution * (region_info["x_range"] / region_info["max_range"])) + ) + yres = int( + np.ceil(resolution * (region_info["y_range"] / region_info["max_range"])) + ) # trim the image to remove empty pixels return image[:xres, :yres] -def project_gas_pixel_grid( - data: SWIFTDataset, - resolution: int, - project: Union[str, None] = "masses", - region: Union[None, unyt_array] = None, - mask: Union[None, array] = None, - rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - parallel: bool = False, - backend: str = "fast", - periodic: bool = True, -): - r""" - Creates a 2D projection of a SWIFT dataset, projected by the "project" - variable (e.g. if project is Temperature, we return: \bar{T} = \sum_j T_j - W_{ij}). - - This function is the same as ``project_gas`` but does not include units. - - Default projection variable is mass. If it is None, then we don't - weight with anything, providing a number density image. - - Parameters - ---------- - - data: SWIFTDataset - The SWIFT dataset that you wish to visualise (get this from ``load``) - - resolution: int - The resolution of the image. All images returned are square, ``res`` - by ``res``, pixel grids. - - project: str, optional - Variable to project to get the weighted density of. By default, this - is mass. If you would like to mass-weight any other variable, you can - always create it as ``data.gas.my_variable = data.gas.other_variable - * data.gas.masses``. - - region: unyt_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: - ``[x_min, x_max, y_min, y_max, {z_min, z_max}]`` - - mask: np.array, optional - Allows only a sub-set of the particles in data to be visualised. Useful - in cases where you have read data out of a ``velociraptor`` catalogue, - or if you only want to visualise e.g. star forming particles. This boolean - mask is applied just before visualisation. - - rotation_center: np.array, optional - Center of the rotation. If you are trying to rotate around a galaxy, this - should be the most bound particle. - - rotation_matrix: np.array, optional - Rotation matrix (3x3) that describes the rotation of the box around - ``rotation_center``. In the default case, this provides a projection - along the z axis. - - parallel: bool, optional - Defaults to ``False``, whether or not to create the image in parallel. - The parallel version of this function uses significantly more memory. - - backend: str, optional - Backend to use. See documentation for details. Defaults to 'fast'. - - periodic: bool, optional - Account for periodic boundary conditions for the simulation box? - Defaults to ``True``. - - Returns - ------- - - image: np.array - Projected image with dimensions of project / length^2, of size - ``res`` x ``res``. - - - Notes - ----- - - + Particles outside of this range are still considered if their smoothing - lengths overlap with the range. - + The returned array has x as the first component and y as the second component, - which is the opposite to what ``imshow`` requires. You should transpose the - array if you want it to be visualised the 'right way up'. - """ - - image = project_pixel_grid( - data=data.gas, - boxsize=data.metadata.boxsize, - resolution=resolution, - project=project, - mask=mask, - parallel=parallel, - region=region, - rotation_matrix=rotation_matrix, - rotation_center=rotation_center, - backend=backend, - periodic=periodic, - ) - - return image - - def project_gas( data: SWIFTDataset, resolution: int, project: Union[str, None] = "masses", - region: Union[None, unyt_array] = None, - mask: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - rotation_matrix: Union[None, array] = None, + region: Union[None, cosmo_array] = None, + mask: Union[None, np.array] = None, + rotation_center: Union[None, cosmo_array] = None, + rotation_matrix: Union[None, np.array] = None, parallel: bool = False, backend: str = "fast", periodic: bool = True, @@ -404,9 +173,10 @@ def project_gas( Variable to project to get the weighted density of. By default, this is mass. If you would like to mass-weight any other variable, you can always create it as ``data.gas.my_variable = data.gas.other_variable - * data.gas.masses``. + * data.gas.masses``. The result is comoving if this is comoving, else + it is physical. - 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,9 +212,9 @@ def project_gas( Returns ------- - image: unyt_array - Projected image with units of project / length^2, of size ``res`` x - ``res``. + image: cosmo_array + Projected image with units of project / length^2, of size ``res`` x ``res``. + Comoving if ``project`` data are comoving, else physical. Notes @@ -457,8 +227,8 @@ def project_gas( array if you want it to be visualised the 'right way up'. """ - image = project_gas_pixel_grid( - data=data, + return project_pixel_grid( + data=data.gas, resolution=resolution, project=project, mask=mask, @@ -469,31 +239,3 @@ def project_gas( backend=backend, periodic=periodic, ) - - if region is not None: - 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) - # 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) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - 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 - else: - new_cosmo_factor = coord_cosmo_factor ** (-2) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) diff --git a/swiftsimio/visualisation/projection_backends/fast.py b/swiftsimio/visualisation/projection_backends/fast.py index 4555e2cb..96121146 100644 --- a/swiftsimio/visualisation/projection_backends/fast.py +++ b/swiftsimio/visualisation/projection_backends/fast.py @@ -2,14 +2,9 @@ Fast backend. This uses float32 precision and no special cases. -""" - -""" The original smoothing code. This provides no renormalisation. """ - -from typing import Union from math import sqrt from numpy import float64, float32, int32, zeros, ndarray @@ -17,10 +12,7 @@ from swiftsimio.visualisation.projection_backends.kernels import ( kernel_single_precision as kernel, ) -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_constant, - kernel_gamma, -) +from swiftsimio.visualisation.projection_backends.kernels import kernel_gamma @jit(nopython=True, fastmath=True) @@ -162,8 +154,9 @@ def scatter( particle_cell_x + cells_spanned + 1, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float32(cell_x) + 0.5) * pixel_width - float32( x_pos ) diff --git a/swiftsimio/visualisation/projection_backends/histogram.py b/swiftsimio/visualisation/projection_backends/histogram.py index cb82a540..871da53d 100644 --- a/swiftsimio/visualisation/projection_backends/histogram.py +++ b/swiftsimio/visualisation/projection_backends/histogram.py @@ -4,24 +4,20 @@ Uses double precision. """ - -from typing import Union -from math import sqrt, ceil -from numpy import float32, float64, int32, zeros, ndarray - +import numpy as np from swiftsimio.accelerated import jit, NUM_THREADS, prange @jit(nopython=True, fastmath=True) def scatter( - x: float64, - y: float64, - m: float32, - h: float32, + x: np.float64, + y: np.float64, + m: np.float32, + h: np.float32, res: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, -) -> ndarray: + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, +) -> np.ndarray: """ Creates a weighted scatter plot @@ -32,34 +28,34 @@ def scatter( Parameters ---------- - x : np.array[float64] + x : np.array[np.float64] array of x-positions of the particles. Must be bounded by [0, 1]. - y : np.array[float64] + y : np.array[np.float64] array of y-positions of the particles. Must be bounded by [0, 1]. - m : np.array[float32] + m : np.array[np.float32] array of masses (or otherwise weights) of the particles - h : np.array[float32] + h : np.array[np.float32] array of smoothing lengths of the particles res : int the number of pixels along one axis, i.e. this returns a square of res * res. - box_x: float64 + box_x: np.float64 box size in x, in the same rescaled length units as x and y. Used for periodic wrapping. - box_y: float64 + box_y: np.float64 box size in y, in the same rescaled length units as x and y. Used for periodic wrapping. Returns ------- - np.array[float32, float32, float32] + np.array[np.float32, np.float32, np.float32] pixel grid of quantity See Also @@ -75,13 +71,13 @@ def scatter( floats and integers is also an improvement over using the numba ones. """ # Output array for our image - image = zeros((res, res), dtype=float64) - maximal_array_index = int32(res) - 1 + image = np.zeros((res, res), dtype=np.float64) + maximal_array_index = np.int32(res) - 1 # Change that integer to a float, we know that our x, y are bounded # by [0, 1]. # We need this for combining with the x_pos and y_pos variables. - float_res = float64(res) + float_res = np.float64(res) # Pre-calculate this constant for use with the above inverse_cell_area = float_res * float_res @@ -107,8 +103,8 @@ def scatter( # Calculate the cell that this particle; use the 64 bit version of the # resolution as this is the same type as the positions - particle_cell_x = int32(float_res * x_pos) - particle_cell_y = int32(float_res * y_pos) + particle_cell_x = np.int32(float_res * x_pos) + particle_cell_y = np.int32(float_res * y_pos) if not ( particle_cell_x < 0 @@ -117,7 +113,7 @@ def scatter( or particle_cell_y >= maximal_array_index ): image[particle_cell_x, particle_cell_y] += ( - float64(mass) * inverse_cell_area + np.float64(mass) * inverse_cell_area ) return image @@ -125,14 +121,14 @@ def scatter( @jit(nopython=True, fastmath=True, parallel=True) def scatter_parallel( - x: float64, - y: float64, - m: float32, - h: float32, + x: np.float64, + y: np.float64, + m: np.float32, + h: np.float32, res: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, -) -> ndarray: + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, +) -> np.ndarray: """ Parallel implementation of scatter @@ -143,34 +139,34 @@ def scatter_parallel( Parameters ---------- - x : np.array[float64] + x : np.array[np.float64] array of x-positions of the particles. Must be bounded by [0, 1]. - y : np.array[float64] + y : np.array[np.float64] array of y-positions of the particles. Must be bounded by [0, 1]. - m : np.array[float32] + m : np.array[np.float32] array of masses (or otherwise weights) of the particles - h : np.array[float32] + h : np.array[np.float32] array of smoothing lengths of the particles res : int the number of pixels along one axis, i.e. this returns a square of res * res. - box_x: float64 + box_x: np.float64 box size in x, in the same rescaled length units as x and y. Used for periodic wrapping. - box_y: float64 + box_y: np.float64 box size in y, in the same rescaled length units as x and y. Used for periodic wrapping. Returns ------- - np.array[float32, float32, float32] + np.array[np.float32, np.float32, np.float32] pixel grid of quantity See Also @@ -192,7 +188,7 @@ def scatter_parallel( number_of_particles = x.size core_particles = number_of_particles // NUM_THREADS - output = zeros((res, res), dtype=float64) + output = np.zeros((res, res), dtype=np.float64) for thread in prange(NUM_THREADS): # Left edge is easy, just start at 0 and go to 'final' diff --git a/swiftsimio/visualisation/projection_backends/kernels.py b/swiftsimio/visualisation/projection_backends/kernels.py index dc62cfb1..bc69990d 100644 --- a/swiftsimio/visualisation/projection_backends/kernels.py +++ b/swiftsimio/visualisation/projection_backends/kernels.py @@ -2,19 +2,16 @@ Projection kernels. """ -from typing import Union -from math import sqrt -from numpy import float64, float32, int32 - -from swiftsimio.accelerated import jit, NUM_THREADS, prange +import numpy as np +from swiftsimio.accelerated import jit # Taken from Dehnen & Aly 2012 -kernel_gamma = float32(1.897367) -kernel_constant = float32(7.0 / 3.14159) +kernel_gamma = np.float32(1.897367) +kernel_constant = np.float32(7.0 / 3.14159) @jit("float32(float32, float32)", nopython=True, fastmath=True) -def kernel_single_precision(r: float32, H: float32): +def kernel_single_precision(r: np.float32, H: np.float32): """ Single precision kernel implementation for swiftsimio. @@ -23,16 +20,16 @@ def kernel_single_precision(r: float32, H: float32): Parameters ---------- - r : float32 + r : np.float32 radius used in kernel computation - H : float32 + H : np.float32 kernel width (i.e. radius of compact support for the kernel) Returns ------- - float32 + np.float32 Contribution to the density by the particle See Also @@ -45,7 +42,7 @@ def kernel_single_precision(r: float32, H: float32): .. [1] Dehnen W., Aly H., 2012, MNRAS, 425, 1068 """ - kernel_constant = float32(2.22817109) + kernel_constant = np.float32(2.22817109) inverse_H = 1.0 / H ratio = r * inverse_H @@ -65,7 +62,7 @@ def kernel_single_precision(r: float32, H: float32): @jit("float64(float64, float64)", nopython=True, fastmath=True) -def kernel_double_precision(r: float64, H: float64): +def kernel_double_precision(r: np.float64, H: np.float64): """ Single precision kernel implementation for swiftsimio. @@ -74,15 +71,15 @@ def kernel_double_precision(r: float64, H: float64): Parameters ---------- - r : float32 + r : np.float32 radius used in kernel computation - H : float32 + H : np.float32 kernel width (i.e. radius of compact support for the kernel) Returns ------- - float32 + np.float32 Contribution to the density by the particle See Also @@ -95,7 +92,7 @@ def kernel_double_precision(r: float64, H: float64): .. [2] Dehnen W., Aly H., 2012, MNRAS, 425, 1068 """ - kernel_constant = float64(2.22817109) + kernel_constant = np.float64(2.22817109) inverse_H = 1.0 / H ratio = r * inverse_H diff --git a/swiftsimio/visualisation/projection_backends/reference.py b/swiftsimio/visualisation/projection_backends/reference.py index dfdd59a0..affe96c5 100644 --- a/swiftsimio/visualisation/projection_backends/reference.py +++ b/swiftsimio/visualisation/projection_backends/reference.py @@ -5,12 +5,9 @@ Uses double precision. """ - -from typing import Union -from math import sqrt, ceil +from math import sqrt from numpy import float32, float64, int32, zeros, ndarray -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.visualisation.projection_backends.kernels import ( kernel_double_precision as kernel, @@ -152,8 +149,9 @@ def scatter( particle_cell_x + cells_spanned + 1, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float64(cell_x) + 0.5) * pixel_width - float64( x_pos ) diff --git a/swiftsimio/visualisation/projection_backends/renormalised.py b/swiftsimio/visualisation/projection_backends/renormalised.py index 295cab0c..4ae6bb60 100644 --- a/swiftsimio/visualisation/projection_backends/renormalised.py +++ b/swiftsimio/visualisation/projection_backends/renormalised.py @@ -4,16 +4,11 @@ This version of the function is the same as `fast` but provides an explicit renormalisation of each kernel such that the mass is conserved up to floating point precision. -""" - -""" The original smoothing code. This provides basic renormalisation -of the kernel on each call. +of the kernel on each call. """ - -from typing import Union from math import sqrt from numpy import float64, float32, int32, zeros, ndarray @@ -22,10 +17,7 @@ from swiftsimio.visualisation.projection_backends.kernels import ( kernel_single_precision as kernel, ) -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_constant, - kernel_gamma, -) +from swiftsimio.visualisation.projection_backends.kernels import kernel_gamma @jit(nopython=True, fastmath=True) @@ -167,8 +159,9 @@ def scatter( particle_cell_x - cells_spanned + 1, particle_cell_x + cells_spanned, ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float32(cell_x) + 0.5) * pixel_width - float32( x_pos ) @@ -199,8 +192,9 @@ def scatter( particle_cell_x + cells_spanned + 1, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float32(cell_x) + 0.5) * pixel_width - float32( x_pos ) diff --git a/swiftsimio/visualisation/projection_backends/subsampled.py b/swiftsimio/visualisation/projection_backends/subsampled.py index 15ee5432..fc95f08a 100644 --- a/swiftsimio/visualisation/projection_backends/subsampled.py +++ b/swiftsimio/visualisation/projection_backends/subsampled.py @@ -5,20 +5,14 @@ scales uses subsampling. Uses double precision. -""" - -""" The original smoothing code. This provides a paranoid supersampling of the kernel. """ - -from typing import Union from math import sqrt, ceil from numpy import float32, float64, int32, zeros, ndarray -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.visualisation.projection_backends.kernels import ( kernel_double_precision as kernel, @@ -282,9 +276,10 @@ def scatter( ), ): float_cell_y = float64(cell_y) - # Now we subsample the pixels to get a more accurate determination - # of the kernel weight. We take the mean of the kernel evaluations - # within a given pixel and apply this as the true 'kernel evaluation'. + # Now we subsample the pixels to get a more accurate + # determination of the kernel weight. We take the mean of the + # kernel evaluations within a given pixel and apply this as + # the true 'kernel evaluation'. kernel_eval = float64(0.0) for subsample_x in range(0, subsample_factor): diff --git a/swiftsimio/visualisation/projection_backends/subsampled_extreme.py b/swiftsimio/visualisation/projection_backends/subsampled_extreme.py index cfd0a2f3..56d0c375 100644 --- a/swiftsimio/visualisation/projection_backends/subsampled_extreme.py +++ b/swiftsimio/visualisation/projection_backends/subsampled_extreme.py @@ -5,20 +5,14 @@ scales uses subsampling. Uses double precision. -""" - -""" The original smoothing code. This provides a paranoid supersampling of the kernel. """ - -from typing import Union from math import sqrt, ceil from numpy import float32, float64, int32, zeros, ndarray -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.visualisation.projection_backends.kernels import ( kernel_double_precision as kernel, @@ -284,9 +278,10 @@ def scatter( ), ): float_cell_y = float64(cell_y) - # Now we subsample the pixels to get a more accurate determination - # of the kernel weight. We take the mean of the kernel evaluations - # within a given pixel and apply this as the true 'kernel evaluation'. + # Now we subsample the pixels to get a more accurate + # determination of the kernel weight. We take the mean of the + # kernel evaluations within a given pixel and apply this as + # the true 'kernel evaluation'. kernel_eval = float64(0.0) for subsample_x in range(0, subsample_factor): diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 61960109..55bc2a1c 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -2,328 +2,52 @@ Ray tracing module for visualisation. """ -# There should be three implementations here: -# - An example case that uses 'screening' with multiple panes. -# - An example case that builds a 3D mesh and uses that -# - An example case that uses a 'real' algorithm (even if it has to -# be single-threaded) - from typing import Union import numpy as np -import math from swiftsimio.objects import cosmo_array from swiftsimio.reader import __SWIFTGroupDataset, SWIFTDataset -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_gamma, - kernel_double_precision as kernel, -) - -from swiftsimio.accelerated import jit, prange, NUM_THREADS -import unyt - - -@jit(nopython=True, fastmath=True) -def core_panels( - x: np.float64, - y: np.float64, - z: np.float64, - h: np.float32, - m: np.float32, - res: int, - panels: int, - min_z: np.float64, - max_z: np.float64, -) -> np.array: - """ - Creates a 2D array of the projected density of particles in a 3D volume using the - 'renormalised' strategy, with multiple panels across the z-range. - - Parameters - ---------- - - x: np.array[np.float64] - The x-coordinates of the particles. - - y: np.array[np.float64] - The y-coordinates of the particles. - - z: np.array[np.float64] - The z-coordinates of the particles. - - h: np.array[np.float32] - The smoothing lengths of the particles. - - m: np.array[np.float32] - The masses of the particles. - - res: int - The resolution of the output array. - - panels: int - The number of panels to use in the z-direction. - - min_z: np.float64 - The minimum z-coordinate of the volume. - - max_z: np.float64 - The maximum z-coordinate of the volume. - - Returns - ------- - - A 3D array of shape (res, res, panels) containing the projected density in each pixel. - """ - output = np.zeros((res, res, panels)) - maximal_array_index = res - 1 - - number_of_particles = len(x) - float_res = float(res) - pixel_width = 1.0 / float_res - - assert len(y) == number_of_particles - assert len(z) == number_of_particles - assert len(h) == number_of_particles - assert len(m) == number_of_particles - - z_per_panel = (max_z - min_z) / panels - - inverse_cell_area = float_res * float_res - - for i in range(number_of_particles): - panel = int(z[i] / z_per_panel) - - if panel < 0 or panel >= panels: - continue - - particle_cell_x = int(float_res * x[i]) - particle_cell_y = int(float_res * y[i]) - - kernel_width = kernel_gamma * h[i] - cells_spanned = int(1.0 + kernel_width * float_res) - - if ( - particle_cell_x + cells_spanned < 0 - or particle_cell_x - cells_spanned > maximal_array_index - or particle_cell_y + cells_spanned < 0 - or particle_cell_y - cells_spanned > maximal_array_index - ): - # Can happily skip this particle - continue - - if cells_spanned <= 1: - if ( - particle_cell_x >= 0 - and particle_cell_x <= maximal_array_index - and particle_cell_y >= 0 - and particle_cell_y <= maximal_array_index - ): - output[particle_cell_x, particle_cell_y, panel] += ( - m[i] * inverse_cell_area - ) - continue - - normalisation = 0.0 - - for cell_x in range( - particle_cell_x - cells_spanned, particle_cell_x + cells_spanned + 1 - ): - distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] - distance_x_2 = distance_x * distance_x - - for cell_y in range( - particle_cell_y - cells_spanned, particle_cell_y + cells_spanned + 1 - ): - distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] - distance_y_2 = distance_y * distance_y - - r = math.sqrt(distance_x_2 + distance_y_2) - - normalisation += kernel(r, kernel_width) - - # Now have the normalisation - normalisation = m[i] * inverse_cell_area / normalisation - - for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we'll segfault - max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the array bounds, - # otherwise we'll segfault (oops). - min(particle_cell_x + cells_spanned + 1, maximal_array_index + 1), - ): - distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] - distance_x_2 = distance_x * distance_x - - for cell_y in range( - max(0, particle_cell_y - cells_spanned), - min(particle_cell_y + cells_spanned + 1, maximal_array_index + 1), - ): - distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] - distance_y_2 = distance_y * distance_y - - r = math.sqrt(distance_x_2 + distance_y_2) - - output[cell_x, cell_y, panel] += kernel(r, kernel_width) * normalisation - - return output - - -@jit(nopython=True, fastmath=True) -def core_panels_parallel( - x: np.float64, - y: np.float64, - z: np.float64, - h: np.float32, - m: np.float32, - res: int, - panels: int, - min_z: np.float64, - max_z: np.float64, -): - # Same as scatter, but executes in parallel! This is actually trivial, - # we just make NUM_THREADS images and add them together at the end. - - number_of_particles = x.size - core_particles = number_of_particles // NUM_THREADS - - output = np.zeros((res, res, panels), dtype=np.float32) - - for thread in prange(NUM_THREADS): - # Left edge is easy, just start at 0 and go to 'final' - left_edge = thread * core_particles - - # Right edge is harder in case of left over particles... - right_edge = thread + 1 - if right_edge == NUM_THREADS: - right_edge = number_of_particles - else: - right_edge *= core_particles - - output += core_panels( - x[left_edge:right_edge], - y[left_edge:right_edge], - z[left_edge:right_edge], - h[left_edge:right_edge], - m[left_edge:right_edge], - res, - panels, - min_z, - max_z, - ) - - return output +from swiftsimio.visualisation.ray_trace_backends import backends +from swiftsimio.visualisation.smoothing_length import backends_get_hsml +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_info, + _get_rotated_coordinates, + backend_restore_cosmo_and_units, +) def panel_pixel_grid( data: __SWIFTGroupDataset, - boxsize: unyt.unyt_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: - if rotation_center is not None: - try: - if rotation_center.units == data.coordinates.units: - pass - else: - raise unyt.exceptions.InvalidUnitOperation( - "Units of coordinates and rotation center must agree" - ) - except AttributeError: - raise unyt.exceptions.InvalidUnitOperation( - "Ensure that rotation_center is a unyt array with the same units as coordinates" - ) - - number_of_particles = data.coordinates.shape[0] - - if project is None: - m = np.ones(number_of_particles, dtype=np.float32) - else: - m = getattr(data, project) - if data.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value - - # This provides a default 'slice it all' mask. - if mask is None: - mask = np.s_[:] - - box_x, box_y, box_z = boxsize - - # Set the limits of the image. - if region is not None: - x_min, x_max, y_min, y_max = region[:4] - - if len(region) == 6: - z_min, z_max = region[4:] - else: - z_min = unyt.unyt_quantity(0.0, units=box_z.units) - z_max = box_z - else: - x_min = unyt.unyt_quantity(0.0, units=box_x.units) - x_max = box_x - y_min = unyt.unyt_quantity(0.0, units=box_y.units) - y_max = box_y - z_min = unyt.unyt_quantity(0.0, units=box_z.units) - z_max = box_z - - x_range = x_max - x_min - y_range = y_max - y_min - - # Deal with non-cubic boxes: - # we always use the maximum of x_range and y_range to normalise the coordinates - # empty pixels in the resulting square image are trimmed afterwards - max_range = max(x_range, y_range) - - try: - hsml = data.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.smoothing_length - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - "Physical smoothing length is not compatible with comoving coordinates!" - ) - else: - if not hsml.compatible_with_physical(): - raise AttributeError( - "Comoving smoothing length is not compatible with physical coordinates!" - ) - - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) - - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - else: - x, y, z = data.coordinates.T + rotation_center: Union[None, cosmo_array] = None, +) -> cosmo_array: - return core_panels( - x=x[mask] / max_range, - y=y[mask] / max_range, + m = _get_projection_field(data, project) + region_info = _get_region_info(data, region) + hsml = backends_get_hsml["sph"](data) + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) + mask = np.s_[...] if mask is None else mask + + # There's a parallel version of core_panels but it seems + # that it's never used anywhere. + norm = region_info["x_range"] * region_info["y_range"] + return backend_restore_cosmo_and_units(backends["core_panels"], norm=norm)( + x=x[mask] / region_info["max_range"], + y=y[mask] / region_info["max_range"], z=z[mask], - h=hsml[mask] / max_range, + h=hsml[mask] / region_info["max_range"], m=m[mask], res=resolution, panels=panels, - min_z=z_min, - max_z=z_max, + min_z=region_info["z_min"], + max_z=region_info["z_max"], ) @@ -332,14 +56,13 @@ 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( + return panel_pixel_grid( data=data.gas, - boxsize=data.metadata.boxsize, resolution=resolution, panels=panels, project=project, @@ -348,146 +71,3 @@ def panel_gas( rotation_matrix=rotation_matrix, rotation_center=rotation_center, ) - - if region is not None: - 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) - # 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) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - 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 - else: - new_cosmo_factor = coord_cosmo_factor ** (-2) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) - - -# --- Functions that actually perform the 'ray tracing'. - - -def transfer_function(value, width, center): - """ - A simple gaussian transfer function centered around a specific value. - """ - return ( - 1 - / (width * np.sqrt(2.0 * np.pi)) - * np.exp(-0.5 * ((value - center) / width) ** 2) - ) - - -@jit(fastmath=True, nopython=True) -def integrate_ray_numba_specific( - input: np.array, red: float, green: float, blue: float, center: float, width: float -): - """ - Given a ray, integrate the transfer function along it - """ - - value = np.array([0.0, 0.0, 0.0], dtype=np.float32) - color = np.array([red, green, blue], dtype=np.float32) - - for i in input: - value += ( - color - * 1 - / (width * np.sqrt(2.0 * np.pi)) - * np.exp(-0.5 * ((i - center) / width) ** 2) - ) - - return value / len(input) - - -@jit(fastmath=True, nopython=True) -def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): - """ - Given a ray, integrate the transfer function along it - """ - - value = np.float32(0.0) - - for i in input: - value *= 0.99 - value += ( - 1 - / (width * np.sqrt(2.0 * np.pi)) - * np.exp(-0.5 * ((i - center) / width) ** 2) - ) - - return np.float32(value / len(input)) - - -# #%% -# data = np.load("voxel_1024.npy") -# # %% -# log_data = np.log10(data) -# # %% -# transfer = lambda x: transfer_function(x, np.mean(log_data), np.std(log_data) * 0.5) -# # %% -# color = np.array([1.0, 0.0, 0.0], dtype=np.float32) -# #%% -# from tqdm import tqdm -# # %% -# @numba.njit(fastmath=True) -# def make_grid(color, center, width): -# output = np.zeros((len(log_data), len(log_data[0])), dtype=np.float32) -# for x in numba.prange(len(log_data)): -# for y in range(len(log_data)): -# data = log_data[x, y] - -# value = np.float32(0.0) - -# for index, i in enumerate(data): -# factor = index / len(data) - -# if factor > 0.5: -# factor = 1.0 - factor - -# value += 1 / (width * np.sqrt(2.0 * np.pi)) * np.exp(-0.5 * ((i - center) / width) ** 2) - -# output[x, y] = value - -# return output - - -# # %% -# import matplotlib.pyplot as plt -# import swiftascmaps -# # %% -# std = np.std(log_data) -# width = 0.05 -# centers = [np.mean(log_data) + x * std for x in [0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5]] -# # %% -# colors = plt.get_cmap("swift.nineteen_eighty_nine")((np.linspace(0, 1, len(centers))))[:, :3] - -# grids = [ -# make_grid(color, center, width) for color, center in zip(colors, centers) -# ] -# #%% - -# #%% -# make_image = lambda x, y: np.array([x * y[0], x * y[1], x * y[2]]).T -# images = [make_image(grid / np.max(grid), color) for color, grid in zip(colors, grids)] -# # %% -# combined_image = sum(images) -# plt.imsave("test.png", combined_image / np.max(combined_image)) -# # %% -# for id, image in zip(centers, images): -# plt.imsave(f"test{id}.png", image) -# # %% diff --git a/swiftsimio/visualisation/ray_trace_backends/__init__.py b/swiftsimio/visualisation/ray_trace_backends/__init__.py new file mode 100644 index 00000000..d14af439 --- /dev/null +++ b/swiftsimio/visualisation/ray_trace_backends/__init__.py @@ -0,0 +1,12 @@ +""" +Backends for ray tracing +""" + +from swiftsimio.visualisation.ray_trace_backends.core_panels import ( + core_panels, + core_panels_parallel, +) + +backends = {"core_panels": core_panels} + +backends_parallel = {"core_panels": core_panels_parallel} diff --git a/swiftsimio/visualisation/ray_trace_backends/core_panels.py b/swiftsimio/visualisation/ray_trace_backends/core_panels.py new file mode 100644 index 00000000..9919aa18 --- /dev/null +++ b/swiftsimio/visualisation/ray_trace_backends/core_panels.py @@ -0,0 +1,332 @@ +""" +Ray tracing module for visualisation. +""" + +# There should be three implementations here: +# - An example case that uses 'screening' with multiple panes. +# - An example case that builds a 3D mesh and uses that +# - An example case that uses a 'real' algorithm (even if it has to +# be single-threaded) + +import numpy as np +import math + +from swiftsimio.visualisation.projection_backends.kernels import ( + kernel_gamma, + kernel_double_precision as kernel, +) + +from swiftsimio.accelerated import jit, prange, NUM_THREADS + + +@jit(nopython=True, fastmath=True) +def core_panels( + x: np.float64, + y: np.float64, + z: np.float64, + h: np.float32, + m: np.float32, + res: int, + panels: int, + min_z: np.float64, + max_z: np.float64, +) -> np.array: + """ + Creates a 2D array of the projected density of particles in a 3D volume using the + 'renormalised' strategy, with multiple panels across the z-range. + + Parameters + ---------- + + x: np.array[np.float64] + The x-coordinates of the particles. + + y: np.array[np.float64] + The y-coordinates of the particles. + + z: np.array[np.float64] + The z-coordinates of the particles. + + h: np.array[np.float32] + The smoothing lengths of the particles. + + m: np.array[np.float32] + The masses of the particles. + + res: int + The resolution of the output array. + + panels: int + The number of panels to use in the z-direction. + + min_z: np.float64 + The minimum z-coordinate of the volume. + + max_z: np.float64 + The maximum z-coordinate of the volume. + + Returns + ------- + + A 3D array of shape (res, res, panels) containing the projected density in each pixel. + """ + output = np.zeros((res, res, panels)) + maximal_array_index = res - 1 + + number_of_particles = len(x) + float_res = float(res) + pixel_width = 1.0 / float_res + + assert len(y) == number_of_particles + assert len(z) == number_of_particles + assert len(h) == number_of_particles + assert len(m) == number_of_particles + + z_per_panel = (max_z - min_z) / panels + + inverse_cell_area = float_res * float_res + + for i in range(number_of_particles): + panel = int(z[i] / z_per_panel) + + if panel < 0 or panel >= panels: + continue + + particle_cell_x = int(float_res * x[i]) + particle_cell_y = int(float_res * y[i]) + + kernel_width = kernel_gamma * h[i] + cells_spanned = int(1.0 + kernel_width * float_res) + + if ( + particle_cell_x + cells_spanned < 0 + or particle_cell_x - cells_spanned > maximal_array_index + or particle_cell_y + cells_spanned < 0 + or particle_cell_y - cells_spanned > maximal_array_index + ): + # Can happily skip this particle + continue + + if cells_spanned <= 1: + if ( + particle_cell_x >= 0 + and particle_cell_x <= maximal_array_index + and particle_cell_y >= 0 + and particle_cell_y <= maximal_array_index + ): + output[particle_cell_x, particle_cell_y, panel] += ( + m[i] * inverse_cell_area + ) + continue + + normalisation = 0.0 + + for cell_x in range( + particle_cell_x - cells_spanned, particle_cell_x + cells_spanned + 1 + ): + distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] + distance_x_2 = distance_x * distance_x + + for cell_y in range( + particle_cell_y - cells_spanned, particle_cell_y + cells_spanned + 1 + ): + distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] + distance_y_2 = distance_y * distance_y + + r = math.sqrt(distance_x_2 + distance_y_2) + + normalisation += kernel(r, kernel_width) + + # Now have the normalisation + normalisation = m[i] * inverse_cell_area / normalisation + + for cell_x in range( + # Ensure that the lowest x value is 0, otherwise we'll segfault + max(0, particle_cell_x - cells_spanned), + # Ensure that the highest x value lies within the array bounds, + # otherwise we'll segfault (oops). + min(particle_cell_x + cells_spanned + 1, maximal_array_index + 1), + ): + distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] + distance_x_2 = distance_x * distance_x + + for cell_y in range( + max(0, particle_cell_y - cells_spanned), + min(particle_cell_y + cells_spanned + 1, maximal_array_index + 1), + ): + distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] + distance_y_2 = distance_y * distance_y + + r = math.sqrt(distance_x_2 + distance_y_2) + + output[cell_x, cell_y, panel] += kernel(r, kernel_width) * normalisation + + return output + + +@jit(nopython=True, fastmath=True) +def core_panels_parallel( + x: np.float64, + y: np.float64, + z: np.float64, + h: np.float32, + m: np.float32, + res: int, + panels: int, + min_z: np.float64, + max_z: np.float64, +): + # Same as scatter, but executes in parallel! This is actually trivial, + # we just make NUM_THREADS images and add them together at the end. + + number_of_particles = x.size + core_particles = number_of_particles // NUM_THREADS + + output = np.zeros((res, res, panels), dtype=np.float32) + + for thread in prange(NUM_THREADS): + # Left edge is easy, just start at 0 and go to 'final' + left_edge = thread * core_particles + + # Right edge is harder in case of left over particles... + right_edge = thread + 1 + + if right_edge == NUM_THREADS: + right_edge = number_of_particles + else: + right_edge *= core_particles + + output += core_panels( + x[left_edge:right_edge], + y[left_edge:right_edge], + z[left_edge:right_edge], + h[left_edge:right_edge], + m[left_edge:right_edge], + res, + panels, + min_z, + max_z, + ) + + return output + + +# --- Functions that actually perform the 'ray tracing'. + + +def transfer_function(value, width, center): + """ + A simple gaussian transfer function centered around a specific value. + """ + return ( + 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((value - center) / width) ** 2) + ) + + +@jit(fastmath=True, nopython=True) +def integrate_ray_numba_specific( + input: np.array, red: float, green: float, blue: float, center: float, width: float +): + """ + Given a ray, integrate the transfer function along it + """ + + value = np.array([0.0, 0.0, 0.0], dtype=np.float32) + color = np.array([red, green, blue], dtype=np.float32) + + for i in input: + value += ( + color + * 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((i - center) / width) ** 2) + ) + + return value / len(input) + + +@jit(fastmath=True, nopython=True) +def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): + """ + Given a ray, integrate the transfer function along it + """ + + value = np.float32(0.0) + + for i in input: + value *= 0.99 + value += ( + 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((i - center) / width) ** 2) + ) + + return np.float32(value / len(input)) + + +# #%% +# data = np.load("voxel_1024.npy") +# # %% +# log_data = np.log10(data) +# # %% +# transfer = lambda x: transfer_function(x, np.mean(log_data), np.std(log_data) * 0.5) +# # %% +# color = np.array([1.0, 0.0, 0.0], dtype=np.float32) +# #%% +# from tqdm import tqdm +# # %% +# @numba.njit(fastmath=True) +# def make_grid(color, center, width): +# output = np.zeros((len(log_data), len(log_data[0])), dtype=np.float32) +# for x in numba.prange(len(log_data)): +# for y in range(len(log_data)): +# data = log_data[x, y] + +# value = np.float32(0.0) + +# for index, i in enumerate(data): +# factor = index / len(data) + +# if factor > 0.5: +# factor = 1.0 - factor + +# value += ( +# 1 +# / (width * np.sqrt(2.0 * np.pi)) +# * np.exp(-0.5 * ((i - center) / width) ** 2) +# ) + +# output[x, y] = value + +# return output + + +# # %% +# import matplotlib.pyplot as plt +# import swiftascmaps +# # %% +# std = np.std(log_data) +# width = 0.05 +# centers = [np.mean(log_data) + x * std for x in [0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5]] +# # %% +# colors = plt.get_cmap("swift.nineteen_eighty_nine")((np.linspace(0, 1, len(centers))))[ +# :, :3 +# ] + +# grids = [ +# make_grid(color, center, width) for color, center in zip(colors, centers) +# ] +# #%% + +# #%% +# make_image = lambda x, y: np.array([x * y[0], x * y[1], x * y[2]]).T +# images = [make_image(grid / np.max(grid), color) for color, grid in zip(colors, grids)] +# # %% +# combined_image = sum(images) +# plt.imsave("test.png", combined_image / np.max(combined_image)) +# # %% +# for id, image in zip(centers, images): +# plt.imsave(f"test{id}.png", image) +# # %% diff --git a/swiftsimio/visualisation/rotation.py b/swiftsimio/visualisation/rotation.py index c9d01c2f..5fab663d 100644 --- a/swiftsimio/visualisation/rotation.py +++ b/swiftsimio/visualisation/rotation.py @@ -2,10 +2,9 @@ Rotation matrix calculation routines. """ -from swiftsimio.accelerated import jit -from numpy import float64, array, matrix, cross, identity, dot, matmul +from numpy import float64, array, cross, identity, dot, matmul from numpy.linalg import norm, inv -from math import sin, cos, sqrt, acos +from math import sin, acos def rotation_matrix_from_vector(vector: float64, axis: str = "z") -> array: diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index 0240311e..25ad1dd0 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -3,31 +3,27 @@ """ from typing import Union, Optional -from numpy import float32, array, ones, matmul -from unyt import unyt_array, unyt_quantity -from swiftsimio import SWIFTDataset, cosmo_array +import numpy as np +from swiftsimio import SWIFTDataset, cosmo_array, cosmo_quantity from swiftsimio.visualisation.slice_backends import backends, backends_parallel from swiftsimio.visualisation.smoothing_length import backends_get_hsml - -from swiftsimio.visualisation.slice_backends.sph import ( - kernel, - kernel_constant, - kernel_gamma, +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_info, + _get_rotated_coordinates, + backend_restore_cosmo_and_units, ) -slice_scatter = backends["sph"] -slice_scatter_parallel = backends_parallel["sph"] - -def slice_gas_pixel_grid( +def slice_gas( data: SWIFTDataset, resolution: int, - z_slice: Optional[unyt_quantity] = None, + z_slice: Optional[cosmo_quantity] = None, project: Union[str, None] = "masses", parallel: bool = False, - rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, + rotation_matrix: Union[None, np.array] = None, + rotation_center: Union[None, cosmo_array] = None, + region: Union[None, cosmo_array] = None, backend: str = "sph", periodic: bool = True, ): @@ -41,33 +37,34 @@ def slice_gas_pixel_grid( Dataset from which slice is extracted resolution : int - Specifies size of return array + Specifies size of return np.array - z_slice : unyt_quantity + z_slice : cosmo_quantity Specifies the location along the z-axis where the slice is to be extracted, relative to the rotation center or the origin of the box if no rotation center is provided. If the perspective is rotated this value refers to the location along the rotated z-axis. project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles + Data field to be projected. Default is mass. If ``None`` then simply + count number of particles. The result is comoving if this is comoving, + else it is physical. parallel : bool used to determine if we will create the image in parallel. This defaults to False, but can speed up the creation of large images significantly at the cost of increased memory usage. - rotation_matrix: np.array, optional + rotation_matrix: np.np.array, optional Rotation matrix (3x3) that describes the rotation of the box around ``rotation_center``. In the default case, this provides a slice perpendicular to the z axis. - rotation_center: np.array, optional + rotation_center: np.np.array, optional Center of the rotation. If you are trying to rotate around a galaxy, this should be the most bound particle. - region : unyt_array, optional + region : cosmo_array, optional 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, and take the form: @@ -87,245 +84,47 @@ def slice_gas_pixel_grid( Returns ------- - ndarray of float32 - Creates a `resolution` x `resolution` array and returns it, - without appropriate units. + image : cosmo_array + Slice image with units of project / length^2, of size ``res`` x ``res``. + Comoving if ``project`` data are comoving, else physical. See Also -------- render_gas_voxel_grid : Creates a 3D voxel grid from a SWIFT dataset """ + data = data.gas + z_slice = np.zeros_like(data.metadata.boxsize[0]) if z_slice is None else z_slice - if z_slice is None: - z_slice = 0.0 * data.gas.coordinates.units - - number_of_gas_particles = data.gas.coordinates.shape[0] - - if project is None: - m = ones(number_of_gas_particles, dtype=float32) - else: - m = getattr(data.gas, project) - if data.gas.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value - - box_x, box_y, box_z = data.metadata.boxsize - - if z_slice > box_z or z_slice < (0 * box_z): - raise ValueError("Please enter a slice value inside the box.") - - # Set the limits of the image. - if region is not None: - x_min, x_max, y_min, y_max = region - else: - x_min = (0 * box_x).to(box_x.units) - x_max = box_x - y_min = (0 * box_y).to(box_y.units) - y_max = box_y - - x_range = x_max - x_min - y_range = y_max - y_min - - # Deal with non-cubic boxes: - # we always use the maximum of x_range and y_range to normalise the coordinates - # empty pixels in the resulting square image are trimmed afterwards - max_range = max(x_range, y_range) - - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = matmul(rotation_matrix, (data.gas.coordinates - rotation_center).T) - - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - - z_center = rotation_center[2] - - else: - x, y, z = data.gas.coordinates.T - - z_center = 0 * box_z - + m = _get_projection_field(data, project) + region_info = _get_region_info(data, region, z_slice=z_slice, periodic=periodic) hsml = backends_get_hsml[backend](data) - if data.gas.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - f"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!" - ) - - if periodic: - periodic_box_x = box_x / max_range - periodic_box_y = box_y / max_range - periodic_box_z = box_z / max_range - else: - periodic_box_x = 0.0 - periodic_box_y = 0.0 - periodic_box_z = 0.0 + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) + z_center = ( + rotation_center[2] + if rotation_center is not None + else np.zeros_like(data.metadata.boxsize[2]) + ) # determine the effective number of pixels for each dimension - xres = int(resolution * x_range / max_range) - yres = int(resolution * y_range / max_range) + xres = int(resolution * region_info["x_range"] / region_info["max_range"]) + yres = int(resolution * region_info["y_range"] / region_info["max_range"]) - common_parameters = dict( - x=(x - x_min) / max_range, - y=(y - y_min) / max_range, - z=z / max_range, + kwargs = dict( + x=(x - region_info["x_min"]) / region_info["max_range"], + y=(y - region_info["y_min"]) / region_info["max_range"], + z=z / region_info["max_range"], m=m, - h=hsml / max_range, - z_slice=(z_center + z_slice) / max_range, + h=hsml / region_info["max_range"], + z_slice=(z_center + z_slice) / region_info["max_range"], xres=xres, yres=yres, - box_x=periodic_box_x, - box_y=periodic_box_y, - box_z=periodic_box_z, + box_x=region_info["periodic_box_x"], + box_y=region_info["periodic_box_y"], + box_z=region_info["periodic_box_z"], ) - - if parallel: - image = backends_parallel[backend](**common_parameters) - else: - image = backends[backend](**common_parameters) + norm = region_info["x_range"] * region_info["y_range"] * region_info["z_range"] + backend_func = (backends_parallel if parallel else backends)[backend] + image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs) return image - - -def slice_gas( - data: SWIFTDataset, - resolution: int, - z_slice: Optional[unyt_quantity] = None, - project: Union[str, None] = "masses", - parallel: bool = False, - rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, - backend: str = "sph", - periodic: bool = True, -): - """ - Creates a 2D slice of a SWIFT dataset, weighted by data field - - Parameters - ---------- - data : SWIFTDataset - Dataset from which slice is extracted - - resolution : int - Specifies size of return array - - z_slice : unyt_quantity - Specifies the location along the z-axis where the slice is to be - extracted, relative to the rotation center or the origin of the box - if no rotation center is provided. If the perspective is rotated - this value refers to the location along the rotated z-axis. - - project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles - - parallel : bool, optional - used to determine if we will create the image in parallel. This - defaults to False, but can speed up the creation of large images - significantly at the cost of increased memory usage. - - rotation_matrix: np.array, optional - Rotation matrix (3x3) that describes the rotation of the box around - ``rotation_center``. In the default case, this provides a slice - perpendicular to the z axis. - - rotation_center: np.array, optional - Center of the rotation. If you are trying to rotate around a galaxy, this - should be the most bound particle. - - region : array, optional - 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, and take the form: - - [x_min, x_max, y_min, y_max] - - Particles outside of this range are still considered if their - smoothing lengths overlap with the range. - - backend : str, optional - Backend to use. Choices are "sph" for interpolation using kernel weights or - "nearest_neighbours" for nearest neighbour interpolation. - - periodic : bool, optional - Account for periodic boundaries for the simulation box? - Default is ``True``. - - Returns - ------- - ndarray of float32 - a `resolution` x `resolution` array of the contribution - of the projected data field to the voxel grid from all of the particles - - See Also - -------- - slice_gas_pixel grid : Creates a 2D slice of a SWIFT dataset - render_gas : Creates a 3D voxel grid of a SWIFT dataset with appropriate units - - Notes - ----- - This is a wrapper function for slice_gas_pixel_grid ensuring that output units are - appropriate - """ - - if z_slice is None: - z_slice = 0.0 * data.gas.coordinates.units - - image = slice_gas_pixel_grid( - data, - resolution, - z_slice, - project, - parallel, - rotation_matrix, - rotation_center, - region, - backend, - periodic, - ) - - if region is not None: - x_range = region[1] - region[0] - y_range = region[3] - region[2] - max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 3) - # 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 * data.metadata.boxsize.units) - ) - else: - max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 3) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) - - 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 ** 3 - else: - new_cosmo_factor = coord_cosmo_factor ** (-3) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) diff --git a/swiftsimio/visualisation/smoothing_length/__init__.py b/swiftsimio/visualisation/smoothing_length/__init__.py index 7cccb0fe..2417ab74 100644 --- a/swiftsimio/visualisation/smoothing_length/__init__.py +++ b/swiftsimio/visualisation/smoothing_length/__init__.py @@ -1,8 +1,10 @@ +import numpy as np from .sph import get_hsml as sph_get_hsml from .nearest_neighbours import get_hsml as nearest_neighbours_get_hsml from .generate import generate_smoothing_lengths backends_get_hsml = { + "histogram": lambda m: np.empty_like(m), "sph": sph_get_hsml, "nearest_neighbours": nearest_neighbours_get_hsml, } diff --git a/swiftsimio/visualisation/smoothing_length/generate.py b/swiftsimio/visualisation/smoothing_length/generate.py index dd271f60..e151fc0b 100644 --- a/swiftsimio/visualisation/smoothing_length/generate.py +++ b/swiftsimio/visualisation/smoothing_length/generate.py @@ -3,19 +3,17 @@ that do not usually carry a smoothing length field (e.g. dark matter). """ -from typing import Union -from unyt import unyt_array -from numpy import empty, float32 - -from swiftsimio import SWIFTDataset, cosmo_array -from swiftsimio.visualisation.projection_backends.kernels import kernel_gamma +import numpy as np +from swiftsimio import cosmo_array from swiftsimio.optional_packages import KDTree, TREE_AVAILABLE +from swiftsimio._array_functions import _propagate_cosmo_array_attributes_to_result +@_propagate_cosmo_array_attributes_to_result # copies attrs of first arg to result def generate_smoothing_lengths( - coordinates: Union[unyt_array, cosmo_array], - boxsize: Union[unyt_array, cosmo_array], - kernel_gamma: float32, + coordinates: cosmo_array, + boxsize: cosmo_array, + kernel_gamma: np.float32, neighbours=32, speedup_fac=2, dimension=3, @@ -25,11 +23,11 @@ def generate_smoothing_lengths( Parameters ---------- - coordinates : unyt_array or cosmo_array + coordinates : cosmo_array a cosmo_array that gives the co-ordinates of all particles - boxsize : unyt_array or cosmo_array + boxsize : cosmo_array the size of the box (3D) - kernel_gamma : float32 + kernel_gamma : np.float32 the kernel gamma of the kernel being used neighbours : int, optional the number of neighbours to encompass @@ -46,8 +44,8 @@ def generate_smoothing_lengths( Returns ------- - smoothing lengths : unyt_array - an unyt array of smoothing lengths. + smoothing lengths : cosmo_array + a cosmo_array of smoothing lengths. """ if not TREE_AVAILABLE: @@ -59,7 +57,7 @@ def generate_smoothing_lengths( tree = KDTree(coordinates.value, boxsize=boxsize.to(coordinates.units).value) - smoothing_lengths = empty(number_of_parts, dtype=float32) + smoothing_lengths = np.empty(number_of_parts, dtype=np.float32) smoothing_lengths[-1] = -0.1 # Include speedup_fac stuff here: @@ -103,15 +101,8 @@ def generate_smoothing_lengths( smoothing_lengths[starting_index:ending_index] = d[:, -1] - if isinstance(coordinates, cosmo_array): - return cosmo_array( - smoothing_lengths * (hsml_correction_fac_speedup / kernel_gamma), - units=coordinates.units, - comoving=coordinates.comoving, - cosmo_factor=coordinates.cosmo_factor, - ) - else: - return unyt_array( - smoothing_lengths * (hsml_correction_fac_speedup / kernel_gamma), - units=coordinates.units, - ) + return type(coordinates)( + smoothing_lengths + * (hsml_correction_fac_speedup / kernel_gamma) + * coordinates.units + ) diff --git a/swiftsimio/visualisation/smoothing_length/sph.py b/swiftsimio/visualisation/smoothing_length/sph.py index 29d09446..d761ad46 100644 --- a/swiftsimio/visualisation/smoothing_length/sph.py +++ b/swiftsimio/visualisation/smoothing_length/sph.py @@ -14,9 +14,8 @@ def get_hsml(data: SWIFTDataset) -> cosmo_array: ------- The extracted smoothing lengths. """ - try: - hsml = data.gas.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.gas.smoothing_length - return hsml + return ( + data.smoothing_lengths + if hasattr(data, "smoothing_lengths") + else data.smoothing_length # backwards compatibility + ) diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index fbb9e591..8662acdb 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -5,591 +5,30 @@ from typing import List, Literal, Tuple, Union from math import sqrt, exp, pi -from numpy import ( - float64, - float32, - int32, - zeros, - array, - ndarray, - ones, - isclose, - linspace, - matmul, - max as np_max, -) -from unyt import unyt_array +import numpy as np from swiftsimio import SWIFTDataset, cosmo_array +from swiftsimio.accelerated import jit -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.optional_packages import plt -from .slice import kernel, kernel_gamma - - -@jit(nopython=True, fastmath=True) -def scatter( - x: float64, - y: float64, - z: float64, - m: float32, - h: float32, - res: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, - box_z: float64 = 0.0, -) -> ndarray: - """ - Creates a weighted voxel grid - - Computes contributions to a voxel grid from particles with positions - (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. - This includes periodic boundary effects. - - Parameters - ---------- - - x : np.array[float64] - array of x-positions of the particles. Must be bounded by [0, 1]. - - y : np.array[float64] - array of y-positions of the particles. Must be bounded by [0, 1]. - - z : np.array[float64] - array of z-positions of the particles. Must be bounded by [0, 1]. - - m : np.array[float32] - array of masses (or otherwise weights) of the particles - - h : np.array[float32] - array of smoothing lengths of the particles - - res : int - the number of voxels along one axis, i.e. this returns a cube - of res * res * res. - - box_x: float64 - box size in x, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_y: float64 - box size in y, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_z: float64 - box size in z, in the same rescaled length units as x, y and z. - Used for periodic wrapping - - Returns - ------- - - np.array[float32, float32, float32] - voxel grid of quantity - - See Also - -------- - - scatter_parallel : Parallel implementation of this function - slice_scatter : Create scatter plot of a slice of data - slice_scatter_parallel : Create scatter plot of a slice of data in parallel - - Notes - ----- - - Explicitly defining the types in this function allows - for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba ones. - """ - # Output array for our image - image = zeros((res, res, res), dtype=float32) - maximal_array_index = int32(res) - 1 - - # Change that integer to a float, we know that our x, y are bounded - # by [0, 1]. - float_res = float32(res) - pixel_width = 1.0 / float_res - - # We need this for combining with the x_pos and y_pos variables. - float_res_64 = float64(res) - - # If the kernel width is smaller than this, we drop to just PIC method - drop_to_single_cell = pixel_width * 0.5 - - # Pre-calculate this constant for use with the above - inverse_cell_volume = float_res * float_res * float_res - - if box_x == 0.0: - xshift_min = 0 - xshift_max = 1 - else: - xshift_min = -1 - xshift_max = 2 - if box_y == 0.0: - yshift_min = 0 - yshift_max = 1 - else: - yshift_min = -1 - yshift_max = 2 - if box_z == 0.0: - zshift_min = 0 - zshift_max = 1 - else: - zshift_min = -1 - zshift_max = 2 - - for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( - x, y, z, m, h - ): - # loop over periodic copies of the particle - for xshift in range(xshift_min, xshift_max): - for yshift in range(yshift_min, yshift_max): - for zshift in range(zshift_min, zshift_max): - x_pos = x_pos_original + xshift * box_x - y_pos = y_pos_original + yshift * box_y - z_pos = z_pos_original + zshift * box_z - - # Calculate the cell that this particle; use the 64 bit version of the - # resolution as this is the same type as the positions - particle_cell_x = int32(float_res_64 * x_pos) - particle_cell_y = int32(float_res_64 * y_pos) - particle_cell_z = int32(float_res_64 * z_pos) - - # SWIFT stores hsml as the FWHM. - kernel_width = kernel_gamma * hsml - - # The number of cells that this kernel spans - cells_spanned = int32(1.0 + kernel_width * float_res) - - if ( - particle_cell_x + cells_spanned < 0 - or particle_cell_x - cells_spanned > maximal_array_index - or particle_cell_y + cells_spanned < 0 - or particle_cell_y - cells_spanned > maximal_array_index - or particle_cell_z + cells_spanned < 0 - or particle_cell_z - cells_spanned > maximal_array_index - ): - # Can happily skip this particle - continue - - if kernel_width < drop_to_single_cell: - # Easygame, gg - if ( - particle_cell_x >= 0 - and particle_cell_x <= maximal_array_index - and particle_cell_y >= 0 - and particle_cell_y <= maximal_array_index - and particle_cell_z >= 0 - and particle_cell_z <= maximal_array_index - ): - image[ - particle_cell_x, particle_cell_y, particle_cell_z - ] += (mass * inverse_cell_volume) - else: - # Now we loop over the square of cells that the kernel lives in - for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we'll segfault - max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the array bounds, - # otherwise we'll segfault (oops). - min( - particle_cell_x + cells_spanned, maximal_array_index + 1 - ), - ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre - distance_x = ( - float32(cell_x) + 0.5 - ) * pixel_width - float32(x_pos) - distance_x_2 = distance_x * distance_x - for cell_y in range( - max(0, particle_cell_y - cells_spanned), - min( - particle_cell_y + cells_spanned, - maximal_array_index + 1, - ), - ): - distance_y = ( - float32(cell_y) + 0.5 - ) * pixel_width - float32(y_pos) - distance_y_2 = distance_y * distance_y - for cell_z in range( - max(0, particle_cell_z - cells_spanned), - min( - particle_cell_z + cells_spanned, - maximal_array_index + 1, - ), - ): - distance_z = ( - float32(cell_z) + 0.5 - ) * pixel_width - float32(z_pos) - distance_z_2 = distance_z * distance_z - - r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) - - kernel_eval = kernel(r, kernel_width) - - image[cell_x, cell_y, cell_z] += mass * kernel_eval - - return image - - -@jit(nopython=True, fastmath=True) -def scatter_limited_z( - x: float64, - y: float64, - z: float64, - m: float32, - h: float32, - res: int, - res_ratio_z: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, - box_z: float64 = 0.0, -) -> ndarray: - """ - Creates a weighted voxel grid - - Computes contributions to a voxel grid from particles with positions - (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. - This includes periodic boundary effects. - - Parameters - ---------- - - x : np.array[float64] - array of x-positions of the particles. Must be bounded by [0, 1]. - - y : np.array[float64] - array of y-positions of the particles. Must be bounded by [0, 1]. - - z : np.array[float64] - array of z-positions of the particles. Must be bounded by [0, 1]. - - m : np.array[float32] - array of masses (or otherwise weights) of the particles - - h : np.array[float32] - array of smoothing lengths of the particles - - res : int - the number of voxels along one axis, i.e. this returns a cube - of res * res * res. - - res_ratio_z: int - the number of voxels along the x and y axes relative to the z - axis. If this is, for instance, 8, and the res is 128, then the - output array will be 128 x 128 x 16. - - box_x: float64 - box size in x, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_y: float64 - box size in y, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_z: float64 - box size in z, in the same rescaled length units as x, y and z. - Used for periodic wrapping - - Returns - ------- - - np.array[float32, float32, float32] - voxel grid of quantity - - See Also - -------- - - scatter_parallel : Parallel implementation of this function - slice_scatter : Create scatter plot of a slice of data - slice_scatter_parallel : Create scatter plot of a slice of data in parallel - - Notes - ----- - - Explicitly defining the types in this function allows - for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba ones. - """ - # Output array for our image - res_z = res // res_ratio_z - image = zeros((res, res, res_z), dtype=float32) - maximal_array_index = int32(res) - 1 - maximal_array_index_z = int32(res_z) - 1 - - # Change that integer to a float, we know that our x, y are bounded - # by [0, 1]. - float_res = float32(res) - float_res_z = float32(res_z) - pixel_width = 1.0 / float_res - pixel_width_z = 1.0 / float_res_z - - # We need this for combining with the x_pos and y_pos variables. - float_res_64 = float64(res) - float_res_z_64 = float64(res_z) - - # If the kernel width is smaller than this, we drop to just PIC method - drop_to_single_cell = pixel_width * 0.5 - drop_to_single_cell_z = pixel_width_z * 0.5 - - # Pre-calculate this constant for use with the above - inverse_cell_volume = float_res * float_res * float_res_z - - if box_x == 0.0: - xshift_min = 0 - xshift_max = 1 - else: - xshift_min = -1 - xshift_max = 2 - if box_y == 0.0: - yshift_min = 0 - yshift_max = 1 - else: - yshift_min = -1 - yshift_max = 2 - if box_z == 0.0: - zshift_min = 0 - zshift_max = 1 - else: - zshift_min = -1 - zshift_max = 2 - - for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( - x, y, z, m, h - ): - # loop over periodic copies of the particle - for xshift in range(xshift_min, xshift_max): - for yshift in range(yshift_min, yshift_max): - for zshift in range(zshift_min, zshift_max): - x_pos = x_pos_original + xshift * box_x - y_pos = y_pos_original + yshift * box_y - z_pos = z_pos_original + zshift * box_z - - # Calculate the cell that this particle; use the 64 bit version of the - # resolution as this is the same type as the positions - particle_cell_x = int32(float_res_64 * x_pos) - particle_cell_y = int32(float_res_64 * y_pos) - particle_cell_z = int32(float_res_z_64 * z_pos) - - # SWIFT stores hsml as the FWHM. - kernel_width = kernel_gamma * hsml - - # The number of cells that this kernel spans - cells_spanned = int32(1.0 + kernel_width * float_res) - cells_spanned_z = int32(1.0 + kernel_width * float_res_z) - - if ( - particle_cell_x + cells_spanned < 0 - or particle_cell_x - cells_spanned > maximal_array_index - or particle_cell_y + cells_spanned < 0 - or particle_cell_y - cells_spanned > maximal_array_index - or particle_cell_z + cells_spanned_z < 0 - or particle_cell_z - cells_spanned_z > maximal_array_index_z - ): - # Can happily skip this particle - continue - - if ( - kernel_width < drop_to_single_cell - or kernel_width < drop_to_single_cell_z - ): - # Easygame, gg - if ( - particle_cell_x >= 0 - and particle_cell_x <= maximal_array_index - and particle_cell_y >= 0 - and particle_cell_y <= maximal_array_index - and particle_cell_z >= 0 - and particle_cell_z <= maximal_array_index_z - ): - image[ - particle_cell_x, particle_cell_y, particle_cell_z - ] += (mass * inverse_cell_volume) - else: - # Now we loop over the square of cells that the kernel lives in - for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we'll segfault - max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the array bounds, - # otherwise we'll segfault (oops). - min( - particle_cell_x + cells_spanned, maximal_array_index + 1 - ), - ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre - distance_x = ( - float32(cell_x) + 0.5 - ) * pixel_width - float32(x_pos) - distance_x_2 = distance_x * distance_x - for cell_y in range( - max(0, particle_cell_y - cells_spanned), - min( - particle_cell_y + cells_spanned, - maximal_array_index + 1, - ), - ): - distance_y = ( - float32(cell_y) + 0.5 - ) * pixel_width - float32(y_pos) - distance_y_2 = distance_y * distance_y - for cell_z in range( - max(0, particle_cell_z - cells_spanned_z), - min( - particle_cell_z + cells_spanned_z, - maximal_array_index_z + 1, - ), - ): - distance_z = ( - float32(cell_z) + 0.5 - ) * pixel_width_z - float32(z_pos) - distance_z_2 = distance_z * distance_z - - r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) - - kernel_eval = kernel(r, kernel_width) - - image[cell_x, cell_y, cell_z] += mass * kernel_eval - - return image - - -@jit(nopython=True, fastmath=True, parallel=True) -def scatter_parallel( - x: float64, - y: float64, - z: float64, - m: float32, - h: float32, - res: int, - res_ratio_z: int = 1, - box_x: float64 = 0.0, - box_y: float64 = 0.0, - box_z: float64 = 0.0, -) -> ndarray: - """ - Parallel implementation of scatter - - Compute contributions to a voxel grid from particles with positions - (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. - This ignores boundary effects. - - Parameters - ---------- - x : array of float64 - array of x-positions of the particles. Must be bounded by [0, 1]. - - y : array of float64 - array of y-positions of the particles. Must be bounded by [0, 1]. - - z : array of float64 - array of z-positions of the particles. Must be bounded by [0, 1]. - - m : array of float32 - array of masses (or otherwise weights) of the particles - - h : array of float32 - array of smoothing lengths of the particles - - res : int - the number of voxels along one axis, i.e. this returns a cube - of res * res * res. - - res_ratio_z: int - the number of voxels along the x and y axes relative to the z - - box_x: float64 - box size in x, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_y: float64 - box size in y, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_z: float64 - box size in z, in the same rescaled length units as x, y and z. - Used for periodic wrapping - - Returns - ------- - - ndarray of float32 - voxel grid of quantity - - See Also - -------- - - scatter : Create voxel grid of quantity - slice_scatter : Create scatter plot of a slice of data - slice_scatter_parallel : Create scatter plot of a slice of data in parallel - - Notes - ----- - - Explicitly defining the types in this function allows - for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba ones. - - """ - # Same as scatter, but executes in parallel! This is actually trivial, - # we just make NUM_THREADS images and add them together at the end. - - number_of_particles = x.size - core_particles = number_of_particles // NUM_THREADS - - output = zeros((res, res, res), dtype=float32) - - for thread in prange(NUM_THREADS): - # Left edge is easy, just start at 0 and go to 'final' - left_edge = thread * core_particles - - # Right edge is harder in case of left over particles... - right_edge = thread + 1 - - if right_edge == NUM_THREADS: - right_edge = number_of_particles - else: - right_edge *= core_particles - - # using kwargs is unsupported in numba - if res_ratio_z == 1: - output += scatter( - x=x[left_edge:right_edge], - y=y[left_edge:right_edge], - z=z[left_edge:right_edge], - m=m[left_edge:right_edge], - h=h[left_edge:right_edge], - res=res, - box_x=box_x, - box_y=box_y, - box_z=box_z, - ) - else: - output += scatter_limited_z( - x=x[left_edge:right_edge], - y=y[left_edge:right_edge], - z=z[left_edge:right_edge], - m=m[left_edge:right_edge], - h=h[left_edge:right_edge], - res=res, - res_ratio_z=res_ratio_z, - box_x=box_x, - box_y=box_y, - box_z=box_z, - ) - - return output +from swiftsimio.visualisation.smoothing_length import backends_get_hsml +from swiftsimio.visualisation.volume_render_backends import backends, backends_parallel +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_info, + _get_rotated_coordinates, + backend_restore_cosmo_and_units, +) -def render_gas_voxel_grid( +def render_gas( data: SWIFTDataset, resolution: int, project: Union[str, None] = "masses", parallel: bool = False, - rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, + rotation_matrix: Union[None, np.array] = None, + rotation_center: Union[None, cosmo_array] = None, + region: Union[None, cosmo_array] = None, periodic: bool = True, ): """ @@ -602,27 +41,28 @@ def render_gas_voxel_grid( Dataset from which slice is extracted resolution : int - Specifies size of return array + Specifies size of return np.array project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles + Data field to be projected. Default is ``"mass"``. If ``None`` then simply + count number of particles. The result is comoving if this is comoving, else + it is physical. parallel : bool used to determine if we will create the image in parallel. This defaults to False, but can speed up the creation of large images significantly at the cost of increased memory usage. - rotation_matrix: np.array, optional + rotation_matrix: np.np.array, optional Rotation matrix (3x3) that describes the rotation of the box around ``rotation_center``. In the default case, this provides a volume render viewed along the z axis. - rotation_center: np.array, optional + rotation_center: cosmo_array, optional Center of the rotation. If you are trying to rotate around a galaxy, this should be the most bound particle. - region : unyt_array, optional + region : cosmo_array, optional determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges, and front and back edges) if it is not None. It should have a @@ -639,231 +79,44 @@ def render_gas_voxel_grid( Returns ------- - ndarray of float32 - Creates a `resolution` x `resolution` x `resolution` array and - returns it, without appropriate units. + cosmo_array + Voxel grid with units of project / length^3, of size ``resolution`` x + ``resolution`` x ``resolution``. Comoving if ``project`` data are + comoving, else physical. See Also -------- slice_gas_pixel_grid : Creates a 2D slice of a SWIFT dataset """ + data = data.gas - number_of_gas_particles = data.gas.particle_ids.size - - if project is None: - m = ones(number_of_gas_particles, dtype=float32) - else: - m = getattr(data.gas, project) - if data.gas.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value - - box_x, box_y, box_z = data.metadata.boxsize - - # Set the limits of the image. - if region is not None: - x_min, x_max, y_min, y_max, z_min, z_max = region - else: - x_min = (0 * box_x).to(box_x.units) - x_max = box_x - y_min = (0 * box_y).to(box_y.units) - y_max = box_y - z_min = (0 * box_z).to(box_z.units) - z_max = box_z - - x_range = x_max - x_min - y_range = y_max - y_min - z_range = z_max - z_min - - # Test that we've got a cubic box - if not ( - isclose(x_range.value, y_range.value) and isclose(x_range.value, z_range.value) - ): - raise AttributeError( - "Projection code is currently not able to handle non-cubic images" - ) - - # Let's just hope that the box is square otherwise we're probably SOL - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = matmul(rotation_matrix, (data.gas.coordinates - rotation_center).T) - - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - - else: - x, y, z = data.gas.coordinates.T - - try: - hsml = data.gas.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.gas.smoothing_length - if data.gas.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - f"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!" - ) - - if periodic: - periodic_box_x = box_x / x_range - periodic_box_y = box_y / y_range - periodic_box_z = box_z / z_range - else: - periodic_box_x = 0.0 - periodic_box_y = 0.0 - periodic_box_z = 0.0 - - arguments = dict( - x=(x - x_min) / x_range, - y=(y - y_min) / y_range, - z=(z - z_min) / z_range, + m = _get_projection_field(data, project) + region_info = _get_region_info(data, region, require_cubic=True) + hsml = backends_get_hsml["sph"](data) + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) + + kwargs = dict( + x=(x - region_info["x_min"]) / region_info["x_range"], + y=(y - region_info["y_min"]) / region_info["y_range"], + z=(z - region_info["z_min"]) / region_info["z_range"], m=m, - h=hsml / x_range, + h=hsml / region_info["x_range"], # cubic so x_range == y_range == z_range res=resolution, - box_x=periodic_box_x, - box_y=periodic_box_y, - box_z=periodic_box_z, + box_x=region_info["periodic_box_x"], + box_y=region_info["periodic_box_y"], + box_z=region_info["periodic_box_z"], ) - - if parallel: - image = scatter_parallel(**arguments) - else: - image = scatter(**arguments) + norm = region_info["x_range"] * region_info["y_range"] * region_info["z_range"] + backend_func = (backends_parallel if parallel else backends)["scatter"] + image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs) return image -def render_gas( - data: SWIFTDataset, - resolution: int, - project: Union[str, None] = "masses", - parallel: bool = False, - rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, - periodic: bool = True, -): - """ - Creates a 3D voxel grid of a SWIFT dataset, weighted by data field - - Parameters - ---------- - - data : SWIFTDataset - Dataset from which slice is extracted - - resolution : int - Specifies size of return array - - project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles - - parallel : bool - used to determine if we will create the image in parallel. This - defaults to False, but can speed up the creation of large images - significantly at the cost of increased memory usage. - - rotation_matrix: np.array, optional - Rotation matrix (3x3) that describes the rotation of the box around - ``rotation_center``. In the default case, this provides a volume render - viewed along the z axis. - - rotation_center: np.array, optional - Center of the rotation. If you are trying to rotate around a galaxy, this - should be the most bound particle. - - region : unyt_array, optional - determines where the image will be created - (this corresponds to the left and right-hand edges, and top and bottom - edges, and front and back edges) if it is not None. It should have a - length of six, and take the form: - ``[x_min, x_max, y_min, y_max, z_min, z_max]`` - Particles outside of this range are still considered if their - smoothing lengths overlap with the range. - - periodic : bool, optional - Account for periodic boundaries for the simulation box? - Default is ``True``. - - Returns - ------- - - ndarray of float32 - a `resolution` x `resolution` x `resolution` array of the contribution - of the projected data field to the voxel grid from all of the particles - - See Also - -------- - - slice_gas : Creates a 2D slice of a SWIFT dataset with appropriate units - render_gas_voxel_grid : Creates a 3D voxel grid of a SWIFT dataset - - Notes - ----- - - This is a wrapper function for slice_gas_pixel_grid ensuring that output - units are appropriate - """ - - image = render_gas_voxel_grid( - data, - resolution, - project, - parallel, - rotation_matrix, - rotation_center, - region=region, - periodic=periodic, - ) - - if region is not None: - x_range = region[1] - region[0] - y_range = region[3] - region[2] - z_range = region[5] - region[4] - units = 1.0 / (x_range * y_range * z_range) - units.convert_to_units(1.0 / (x_range.units * y_range.units * z_range.units)) - else: - units = 1.0 / ( - data.metadata.boxsize[0] - * data.metadata.boxsize[1] - * data.metadata.boxsize[2] - ) - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) - - 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 ** 3 - else: - new_cosmo_factor = coord_cosmo_factor ** (-3) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) - - @jit(nopython=True, fastmath=True) def render_voxel_to_array(data, center, width): - output = zeros((data.shape[0], data.shape[1])) + output = np.zeros((data.shape[0], data.shape[1])) for i in range(data.shape[0]): for j in range(data.shape[1]): @@ -880,24 +133,24 @@ def render_voxel_to_array(data, center, width): def visualise_render( - render: ndarray, + render: np.ndarray, centers: List[float], widths: Union[List[float], float], cmap: str = "viridis", return_type: Literal["all", "lighten", "add"] = "lighten", norm: Union[List["plt.Normalize"], "plt.Normalize", None] = None, -) -> Tuple[Union[List[ndarray], ndarray], List["plt.Normalize"]]: +) -> Tuple[Union[List[np.ndarray], np.ndarray], List["plt.Normalize"]]: """ Visualises a render with multiple centers and widths. Parameters ---------- - render : np.array + render : np.np.array The render to visualise. You should scale this appropriately before using this function (e.g. use a logarithmic transform!) - and pass in the 'value' array, not the original cosmo array or - unyt array. + and pass in the 'value' np.array, not the original cosmo_array or + unyt_array. centers : list[float] The centers of your rendering functions @@ -907,8 +160,8 @@ def visualise_render( will have the same width. cmap : str - The colormap to use for the rendering functions. - + The colormap to use for the rendering functions. + return_type : Literal["all", "lighten", "add"] The type of return. If "all", all images are returned. If "lighten", the maximum of all images is returned. If "add", the sum of all images @@ -921,7 +174,7 @@ def visualise_render( Returns ------- - list[np.array] | np.array + list[np.np.array] | np.np.array The images of the rendering functions. If return_type is "all", this will be a list of images. If return_type is "lighten" or "add", this will be a single image. @@ -938,7 +191,7 @@ def visualise_render( elif not isinstance(norm, list): norm = [norm] * len(centers) - colors = plt.get_cmap(cmap)(linspace(0, 1, len(centers)))[:, :3] + colors = plt.get_cmap(cmap)(np.linspace(0, 1, len(centers)))[:, :3] images = [ n(render_voxel_to_array(render, center, width)) @@ -946,7 +199,7 @@ def visualise_render( ] images = [ - array([color[0] * x, color[1] * x, color[2] * x]).T + np.array([color[0] * x, color[1] * x, color[2] * x]).T for color, x in zip(colors, images) ] @@ -954,7 +207,7 @@ def visualise_render( return images, norm if return_type == "lighten": - return np_max(images, axis=0), norm + return np.max(images, axis=0), norm if return_type == "add": return sum(images), norm @@ -974,11 +227,11 @@ def visualise_render_options( centers : list[float] The centers of your rendering functions - + widths : list[float] | float The widths of your rendering functions. If a single float, all functions will have the same width. - + cmap : str The colormap to use for the rendering functions. @@ -993,10 +246,10 @@ def visualise_render_options( if isinstance(widths, float): widths = [widths] * len(centers) - colors = plt.get_cmap(cmap)(linspace(0, 1, len(centers)))[:, :3] + colors = plt.get_cmap(cmap)(np.linspace(0, 1, len(centers)))[:, :3] for center, width, color in zip(centers, widths, colors): - xs = linspace(center - 5.0 * width, center + 5.0 * width, 100) + xs = np.linspace(center - 5.0 * width, center + 5.0 * width, 100) ys = [ exp(-0.5 * ((center - x) / width) ** 2) / (width * sqrt(2.0 * pi)) for x in xs diff --git a/swiftsimio/visualisation/volume_render_backends/__init__.py b/swiftsimio/visualisation/volume_render_backends/__init__.py new file mode 100644 index 00000000..f8108b5f --- /dev/null +++ b/swiftsimio/visualisation/volume_render_backends/__init__.py @@ -0,0 +1,12 @@ +""" +Backends for volume rendering +""" + +from swiftsimio.visualisation.volume_render_backends.scatter import ( + scatter, + scatter_parallel, +) + +backends = {"scatter": scatter} + +backends_parallel = {"scatter": scatter_parallel} diff --git a/swiftsimio/visualisation/volume_render_backends/scatter.py b/swiftsimio/visualisation/volume_render_backends/scatter.py new file mode 100644 index 00000000..d52218bd --- /dev/null +++ b/swiftsimio/visualisation/volume_render_backends/scatter.py @@ -0,0 +1,568 @@ +""" +Basic volume render for SPH data. This takes the 3D positions +of the particles and projects them onto a grid. +""" + +from math import sqrt +import numpy as np + +from swiftsimio.accelerated import jit, NUM_THREADS, prange + +from swiftsimio.visualisation.slice_backends.sph import kernel, kernel_gamma + + +@jit(nopython=True, fastmath=True) +def scatter( + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, + res: int, + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: + """ + Creates a weighted voxel grid + + Computes contributions to a voxel grid from particles with positions + (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. + This includes periodic boundary effects. + + Parameters + ---------- + + x : np.np.array[np.float64] + np.array of x-positions of the particles. Must be bounded by [0, 1]. + + y : np.np.array[np.float64] + np.array of y-positions of the particles. Must be bounded by [0, 1]. + + z : np.np.array[np.float64] + np.array of z-positions of the particles. Must be bounded by [0, 1]. + + m : np.np.array[np.float32] + np.array of masses (or otherwise weights) of the particles + + h : np.np.array[np.float32] + np.array of smoothing lengths of the particles + + res : int + the number of voxels along one axis, i.e. this returns a cube + of res * res * res. + + box_x: np.float64 + box size in x, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_y: np.float64 + box size in y, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_z: np.float64 + box size in z, in the same rescaled length units as x, y and z. + Used for periodic wrapping + + Returns + ------- + + np.np.array[np.float32, np.float32, np.float32] + voxel grid of quantity + + See Also + -------- + + scatter_parallel : Parallel implementation of this function + slice_scatter : Create scatter plot of a slice of data + slice_scatter_parallel : Create scatter plot of a slice of data in parallel + + Notes + ----- + + Explicitly defining the types in this function allows + for a 25-50% performance improvement. In our testing, using numpy + floats and integers is also an improvement over using the numba np.ones. + """ + # Output np.array for our image + image = np.zeros((res, res, res), dtype=np.float32) + maximal_array_index = np.int32(res) - 1 + + # Change that integer to a float, we know that our x, y are bounded + # by [0, 1]. + float_res = np.float32(res) + pixel_width = 1.0 / float_res + + # We need this for combining with the x_pos and y_pos variables. + float_res_64 = np.float64(res) + + # If the kernel width is smaller than this, we drop to just PIC method + drop_to_single_cell = pixel_width * 0.5 + + # Pre-calculate this constant for use with the above + inverse_cell_volume = float_res * float_res * float_res + + if box_x == 0.0: + xshift_min = 0 + xshift_max = 1 + else: + xshift_min = -1 + xshift_max = 2 + if box_y == 0.0: + yshift_min = 0 + yshift_max = 1 + else: + yshift_min = -1 + yshift_max = 2 + if box_z == 0.0: + zshift_min = 0 + zshift_max = 1 + else: + zshift_min = -1 + zshift_max = 2 + + for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( + x, y, z, m, h + ): + # loop over periodic copies of the particle + for xshift in range(xshift_min, xshift_max): + for yshift in range(yshift_min, yshift_max): + for zshift in range(zshift_min, zshift_max): + x_pos = x_pos_original + xshift * box_x + y_pos = y_pos_original + yshift * box_y + z_pos = z_pos_original + zshift * box_z + + # Calculate the cell that this particle; use the 64 bit version of the + # resolution as this is the same type as the positions + particle_cell_x = np.int32(float_res_64 * x_pos) + particle_cell_y = np.int32(float_res_64 * y_pos) + particle_cell_z = np.int32(float_res_64 * z_pos) + + # SWIFT stores hsml as the FWHM. + kernel_width = kernel_gamma * hsml + + # The number of cells that this kernel spans + cells_spanned = np.int32(1.0 + kernel_width * float_res) + + if ( + particle_cell_x + cells_spanned < 0 + or particle_cell_x - cells_spanned > maximal_array_index + or particle_cell_y + cells_spanned < 0 + or particle_cell_y - cells_spanned > maximal_array_index + or particle_cell_z + cells_spanned < 0 + or particle_cell_z - cells_spanned > maximal_array_index + ): + # Can happily skip this particle + continue + + if kernel_width < drop_to_single_cell: + # Easygame, gg + if ( + particle_cell_x >= 0 + and particle_cell_x <= maximal_array_index + and particle_cell_y >= 0 + and particle_cell_y <= maximal_array_index + and particle_cell_z >= 0 + and particle_cell_z <= maximal_array_index + ): + image[ + particle_cell_x, particle_cell_y, particle_cell_z + ] += (mass * inverse_cell_volume) + else: + # Now we loop over the square of cells that the kernel lives in + for cell_x in range( + # Ensure that the lowest x value is 0, otherwise we segfault + max(0, particle_cell_x - cells_spanned), + # Ensure that the highest x value lies within the np.array + # bounds, otherwise we'll segfault (oops). + min( + particle_cell_x + cells_spanned, maximal_array_index + 1 + ), + ): + # The distance in x to our new favourite cell -- remember that + # our x, y are all in a box of [0, 1]; calculate the distance + # to the cell centre + distance_x = ( + np.float32(cell_x) + 0.5 + ) * pixel_width - np.float32(x_pos) + distance_x_2 = distance_x * distance_x + for cell_y in range( + max(0, particle_cell_y - cells_spanned), + min( + particle_cell_y + cells_spanned, + maximal_array_index + 1, + ), + ): + distance_y = ( + np.float32(cell_y) + 0.5 + ) * pixel_width - np.float32(y_pos) + distance_y_2 = distance_y * distance_y + for cell_z in range( + max(0, particle_cell_z - cells_spanned), + min( + particle_cell_z + cells_spanned, + maximal_array_index + 1, + ), + ): + distance_z = ( + np.float32(cell_z) + 0.5 + ) * pixel_width - np.float32(z_pos) + distance_z_2 = distance_z * distance_z + + r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) + + kernel_eval = kernel(r, kernel_width) + + image[cell_x, cell_y, cell_z] += mass * kernel_eval + + return image + + +@jit(nopython=True, fastmath=True) +def scatter_limited_z( + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, + res: int, + res_ratio_z: int, + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: + """ + Creates a weighted voxel grid + + Computes contributions to a voxel grid from particles with positions + (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. + This includes periodic boundary effects. + + Parameters + ---------- + + x : np.np.array[np.float64] + np.array of x-positions of the particles. Must be bounded by [0, 1]. + + y : np.np.array[np.float64] + np.array of y-positions of the particles. Must be bounded by [0, 1]. + + z : np.np.array[np.float64] + np.array of z-positions of the particles. Must be bounded by [0, 1]. + + m : np.np.array[np.float32] + np.array of masses (or otherwise weights) of the particles + + h : np.np.array[np.float32] + np.array of smoothing lengths of the particles + + res : int + the number of voxels along one axis, i.e. this returns a cube + of res * res * res. + + res_ratio_z: int + the number of voxels along the x and y axes relative to the z + axis. If this is, for instance, 8, and the res is 128, then the + output np.array will be 128 x 128 x 16. + + box_x: np.float64 + box size in x, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_y: np.float64 + box size in y, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_z: np.float64 + box size in z, in the same rescaled length units as x, y and z. + Used for periodic wrapping + + Returns + ------- + + np.np.array[np.float32, np.float32, np.float32] + voxel grid of quantity + + See Also + -------- + + scatter_parallel : Parallel implementation of this function + slice_scatter : Create scatter plot of a slice of data + slice_scatter_parallel : Create scatter plot of a slice of data in parallel + + Notes + ----- + + Explicitly defining the types in this function allows + for a 25-50% performance improvement. In our testing, using numpy + floats and integers is also an improvement over using the numba np.ones. + """ + # Output np.array for our image + res_z = res // res_ratio_z + image = np.zeros((res, res, res_z), dtype=np.float32) + maximal_array_index = np.int32(res) - 1 + maximal_array_index_z = np.int32(res_z) - 1 + + # Change that integer to a float, we know that our x, y are bounded + # by [0, 1]. + float_res = np.float32(res) + float_res_z = np.float32(res_z) + pixel_width = 1.0 / float_res + pixel_width_z = 1.0 / float_res_z + + # We need this for combining with the x_pos and y_pos variables. + float_res_64 = np.float64(res) + float_res_z_64 = np.float64(res_z) + + # If the kernel width is smaller than this, we drop to just PIC method + drop_to_single_cell = pixel_width * 0.5 + drop_to_single_cell_z = pixel_width_z * 0.5 + + # Pre-calculate this constant for use with the above + inverse_cell_volume = float_res * float_res * float_res_z + + if box_x == 0.0: + xshift_min = 0 + xshift_max = 1 + else: + xshift_min = -1 + xshift_max = 2 + if box_y == 0.0: + yshift_min = 0 + yshift_max = 1 + else: + yshift_min = -1 + yshift_max = 2 + if box_z == 0.0: + zshift_min = 0 + zshift_max = 1 + else: + zshift_min = -1 + zshift_max = 2 + + for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( + x, y, z, m, h + ): + # loop over periodic copies of the particle + for xshift in range(xshift_min, xshift_max): + for yshift in range(yshift_min, yshift_max): + for zshift in range(zshift_min, zshift_max): + x_pos = x_pos_original + xshift * box_x + y_pos = y_pos_original + yshift * box_y + z_pos = z_pos_original + zshift * box_z + + # Calculate the cell that this particle; use the 64 bit version of the + # resolution as this is the same type as the positions + particle_cell_x = np.int32(float_res_64 * x_pos) + particle_cell_y = np.int32(float_res_64 * y_pos) + particle_cell_z = np.int32(float_res_z_64 * z_pos) + + # SWIFT stores hsml as the FWHM. + kernel_width = kernel_gamma * hsml + + # The number of cells that this kernel spans + cells_spanned = np.int32(1.0 + kernel_width * float_res) + cells_spanned_z = np.int32(1.0 + kernel_width * float_res_z) + + if ( + particle_cell_x + cells_spanned < 0 + or particle_cell_x - cells_spanned > maximal_array_index + or particle_cell_y + cells_spanned < 0 + or particle_cell_y - cells_spanned > maximal_array_index + or particle_cell_z + cells_spanned_z < 0 + or particle_cell_z - cells_spanned_z > maximal_array_index_z + ): + # Can happily skip this particle + continue + + if ( + kernel_width < drop_to_single_cell + or kernel_width < drop_to_single_cell_z + ): + # Easygame, gg + if ( + particle_cell_x >= 0 + and particle_cell_x <= maximal_array_index + and particle_cell_y >= 0 + and particle_cell_y <= maximal_array_index + and particle_cell_z >= 0 + and particle_cell_z <= maximal_array_index_z + ): + image[ + particle_cell_x, particle_cell_y, particle_cell_z + ] += (mass * inverse_cell_volume) + else: + # Now we loop over the square of cells that the kernel lives in + for cell_x in range( + # Ensure that the lowest x value is 0, otherwise we segfault + max(0, particle_cell_x - cells_spanned), + # Ensure that the highest x value lies within the np.array + # bounds, otherwise we'll segfault (oops). + min( + particle_cell_x + cells_spanned, maximal_array_index + 1 + ), + ): + # The distance in x to our new favourite cell -- remember that + # our x, y are all in a box of [0, 1]; calculate the distance + # to the cell centre + distance_x = ( + np.float32(cell_x) + 0.5 + ) * pixel_width - np.float32(x_pos) + distance_x_2 = distance_x * distance_x + for cell_y in range( + max(0, particle_cell_y - cells_spanned), + min( + particle_cell_y + cells_spanned, + maximal_array_index + 1, + ), + ): + distance_y = ( + np.float32(cell_y) + 0.5 + ) * pixel_width - np.float32(y_pos) + distance_y_2 = distance_y * distance_y + for cell_z in range( + max(0, particle_cell_z - cells_spanned_z), + min( + particle_cell_z + cells_spanned_z, + maximal_array_index_z + 1, + ), + ): + distance_z = ( + np.float32(cell_z) + 0.5 + ) * pixel_width_z - np.float32(z_pos) + distance_z_2 = distance_z * distance_z + + r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) + + kernel_eval = kernel(r, kernel_width) + + image[cell_x, cell_y, cell_z] += mass * kernel_eval + + return image + + +@jit(nopython=True, fastmath=True, parallel=True) +def scatter_parallel( + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, + res: int, + res_ratio_z: int = 1, + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: + """ + Parallel implementation of scatter + + Compute contributions to a voxel grid from particles with positions + (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. + This ignores boundary effects. + + Parameters + ---------- + x : np.array of np.float64 + np.array of x-positions of the particles. Must be bounded by [0, 1]. + + y : np.array of np.float64 + np.array of y-positions of the particles. Must be bounded by [0, 1]. + + z : np.array of np.float64 + np.array of z-positions of the particles. Must be bounded by [0, 1]. + + m : np.array of np.float32 + np.array of masses (or otherwise weights) of the particles + + h : np.array of np.float32 + np.array of smoothing lengths of the particles + + res : int + the number of voxels along one axis, i.e. this returns a cube + of res * res * res. + + res_ratio_z: int + the number of voxels along the x and y axes relative to the z + + box_x: np.float64 + box size in x, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_y: np.float64 + box size in y, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_z: np.float64 + box size in z, in the same rescaled length units as x, y and z. + Used for periodic wrapping + + Returns + ------- + + np.ndarray of np.float32 + voxel grid of quantity + + See Also + -------- + + scatter : Create voxel grid of quantity + slice_scatter : Create scatter plot of a slice of data + slice_scatter_parallel : Create scatter plot of a slice of data in parallel + + Notes + ----- + + Explicitly defining the types in this function allows + for a 25-50% performance improvement. In our testing, using numpy + floats and integers is also an improvement over using the numba np.ones. + + """ + # Same as scatter, but executes in parallel! This is actually trivial, + # we just make NUM_THREADS images and add them together at the end. + + number_of_particles = x.size + core_particles = number_of_particles // NUM_THREADS + + output = np.zeros((res, res, res), dtype=np.float32) + + for thread in prange(NUM_THREADS): + # Left edge is easy, just start at 0 and go to 'final' + left_edge = thread * core_particles + + # Right edge is harder in case of left over particles... + right_edge = thread + 1 + + if right_edge == NUM_THREADS: + right_edge = number_of_particles + else: + right_edge *= core_particles + + # using kwargs is unsupported in numba + if res_ratio_z == 1: + output += scatter( + x=x[left_edge:right_edge], + y=y[left_edge:right_edge], + z=z[left_edge:right_edge], + m=m[left_edge:right_edge], + h=h[left_edge:right_edge], + res=res, + box_x=box_x, + box_y=box_y, + box_z=box_z, + ) + else: + output += scatter_limited_z( + x=x[left_edge:right_edge], + y=y[left_edge:right_edge], + z=z[left_edge:right_edge], + m=m[left_edge:right_edge], + h=h[left_edge:right_edge], + res=res, + res_ratio_z=res_ratio_z, + box_x=box_x, + box_y=box_y, + box_z=box_z, + ) + + return output diff --git a/tests/basic_test.py b/tests/basic_test.py index e125da1f..adb4e7fc 100644 --- a/tests/basic_test.py +++ b/tests/basic_test.py @@ -4,7 +4,6 @@ from swiftsimio import load from swiftsimio import Writer -from swiftsimio.units import cosmo_units import unyt import numpy as np @@ -53,7 +52,7 @@ def test_load(): """ x = load("test.hdf5") - density = x.gas.internal_energy - coordinates = x.gas.coordinates + x.gas.internal_energy + x.gas.coordinates os.remove("test.hdf5") diff --git a/tests/helper.py b/tests/helper.py index f2d6907d..58b0afd8 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -6,8 +6,8 @@ import os import h5py from swiftsimio.subset_writer import find_links, write_metadata -from swiftsimio import mask, cosmo_array, load -from numpy import mean +from swiftsimio import mask, cosmo_array +from numpy import mean, zeros_like webstorage_location = "http://virgodb.cosma.dur.ac.uk/swift-webstorage/IOExamples/" test_data_location = "test_data/" @@ -82,7 +82,7 @@ def create_single_particle_dataset(filename: str, output_name: str): # Create a dummy mask in order to write metadata data_mask = mask(filename) boxsize = data_mask.metadata.boxsize - region = [[0, b] for b in boxsize] + region = [[zeros_like(b), b] for b in boxsize] data_mask.constrain_spatial(region) # Write the metadata diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 0b1b4d90..3f87702d 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -2,18 +2,125 @@ Tests the initialisation of a cosmo_array. """ +import pytest +import os +import warnings import numpy as np import unyt as u -from swiftsimio.objects import cosmo_array, cosmo_factor, a from copy import copy, deepcopy +from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a + +savetxt_file = "saved_array.txt" + + +def getfunc(fname): + """ + Helper for our tests: get the function handle from a name (possibly with attribute + access). + """ + func = np + for attr in fname.split("."): + func = getattr(func, attr) + return func + + +def ca(x, unit=u.Mpc): + """ + Helper for our tests: turn an array into a cosmo_array. + """ + return cosmo_array(x, unit, comoving=False, scale_factor=0.5, scale_exponent=1) + + +def cq(x, unit=u.Mpc): + """ + Helper for our tests: turn a scalar into a cosmo_quantity. + """ + return cosmo_quantity(x, unit, comoving=False, scale_factor=0.5, scale_exponent=1) + + +def arg_to_ua(arg): + """ + Helper for our tests: recursively convert cosmo_* in an argument (possibly an + iterable) to their unyt_* equivalents. + """ + if type(arg) in (list, tuple): + return type(arg)([arg_to_ua(a) for a in arg]) + else: + return to_ua(arg) + + +def to_ua(x): + """ + Helper for our tests: turn a cosmo_* object into its unyt_* equivalent. + """ + return u.unyt_array(x) if hasattr(x, "comoving") else x + + +def check_result(x_c, x_u, ignore_values=False): + """ + Helper for our tests: check that a result with cosmo input matches what we + expected based on the result with unyt input. + + We check: + - that the type of the result makes sense, recursing if needed. + - that the value of the result matches (unless ignore_values=False). + - that the units match. + """ + if x_u is None: + assert x_c is None + return + if isinstance(x_u, str): + assert isinstance(x_c, str) + return + if isinstance(x_u, type) or isinstance(x_u, np.dtype): + assert x_u == x_c + return + if type(x_u) in (list, tuple): + assert type(x_u) is type(x_c) + assert len(x_u) == len(x_c) + for x_c_i, x_u_i in zip(x_c, x_u): + check_result(x_c_i, x_u_i) + return + # careful, unyt_quantity is a subclass of unyt_array: + if isinstance(x_u, u.unyt_quantity): + assert isinstance(x_c, cosmo_quantity) + elif isinstance(x_u, u.unyt_array): + assert isinstance(x_c, cosmo_array) and not isinstance(x_c, cosmo_quantity) + else: + assert not isinstance(x_c, cosmo_array) + if not ignore_values: + assert np.allclose(x_c, x_u) + return + assert x_c.units == x_u.units + if not ignore_values: + assert np.allclose(x_c.to_value(x_c.units), x_u.to_value(x_u.units)) + if isinstance(x_c, cosmo_array): # includes cosmo_quantity + assert x_c.comoving is False + if x_c.units != u.dimensionless: + assert x_c.cosmo_factor is not None + return class TestCosmoArrayInit: + """ + Test different ways of initializing a cosmo_array. + """ + def test_init_from_ndarray(self): + """ + Check initializing from a bare numpy array. + """ + arr = cosmo_array( + np.ones(5), units=u.Mpc, scale_factor=1.0, scale_exponent=1, comoving=False + ) + assert hasattr(arr, "cosmo_factor") + assert hasattr(arr, "comoving") + assert isinstance(arr, cosmo_array) + # also with a cosmo_factor argument instead of scale_factor & scale_exponent arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1.0), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -21,10 +128,24 @@ def test_init_from_ndarray(self): assert isinstance(arr, cosmo_array) def test_init_from_list(self): + """ + Check initializing from a list of values. + """ arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + scale_factor=1.0, + scale_exponent=1, + comoving=False, + ) + assert hasattr(arr, "cosmo_factor") + assert hasattr(arr, "comoving") + assert isinstance(arr, cosmo_array) + # also with a cosmo_factor argument instead of scale_factor & scale_exponent + arr = cosmo_array( + [1, 1, 1, 1, 1], + units=u.Mpc, + cosmo_factor=cosmo_factor(a ** 1, 1.0), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -32,9 +153,22 @@ def test_init_from_list(self): assert isinstance(arr, cosmo_array) def test_init_from_unyt_array(self): + """ + Check initializing from a unyt_array. + """ arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a ** 1, 1), + scale_factor=1.0, + scale_exponent=1, + comoving=False, + ) + assert hasattr(arr, "cosmo_factor") + assert hasattr(arr, "comoving") + assert isinstance(arr, cosmo_array) + # also with a cosmo_factor argument instead of scale_factor & scale_exponent + arr = cosmo_array( + u.unyt_array(np.ones(5), units=u.Mpc), + cosmo_factor=cosmo_factor(a ** 1, 1.0), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -42,17 +176,844 @@ def test_init_from_unyt_array(self): assert isinstance(arr, cosmo_array) def test_init_from_list_of_unyt_arrays(self): + """ + Check initializing from a list of unyt_array's. + + Note that unyt won't recurse deeper than one level on inputs, so we don't test + deeper than one level of lists. This behaviour is documented in cosmo_array. + """ arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a ** 1, 1), + scale_factor=1.0, + scale_exponent=1, comoving=False, ) assert hasattr(arr, "cosmo_factor") assert hasattr(arr, "comoving") assert isinstance(arr, cosmo_array) + # also with a cosmo_factor argument instead of scale_factor & scale_exponent + arr = cosmo_array( + [u.unyt_array(1, units=u.Mpc) for _ in range(5)], + cosmo_factor=cosmo_factor(a ** 1, 1.0), + comoving=False, + ) + assert hasattr(arr, "cosmo_factor") + assert hasattr(arr, "comoving") + assert isinstance(arr, cosmo_array) + + def test_init_from_list_of_cosmo_arrays(self): + """ + Check initializing from a list of cosmo_array's. + + Note that unyt won't recurse deeper than one level on inputs, so we don't test + deeper than one level of lists. This behaviour is documented in cosmo_array. + """ + arr = cosmo_array( + [ + cosmo_array( + [1], units=u.Mpc, comoving=False, scale_factor=1.0, scale_exponent=1 + ) + for _ in range(5) + ] + ) + assert isinstance(arr, cosmo_array) + assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor( + a ** 1, 1 + ) + assert hasattr(arr, "comoving") and arr.comoving is False + # also with a cosmo_factor argument instead of scale_factor & scale_exponent + arr = cosmo_array( + [ + cosmo_array( + [1], + units=u.Mpc, + comoving=False, + cosmo_factor=cosmo_factor(a ** 1, 1.0), + ) + for _ in range(5) + ] + ) + assert isinstance(arr, cosmo_array) + assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor( + a ** 1, 1 + ) + assert hasattr(arr, "comoving") and arr.comoving is False + + def test_expected_init_failures(self): + for cls, inp in ((cosmo_array, [1]), (cosmo_quantity, 1)): + # we refuse both cosmo_factor and scale_factor/scale_exponent provided: + with pytest.raises(ValueError): + cls( + inp, + units=u.Mpc, + comoving=False, + cosmo_factor=cosmo_factor.create(1.0, 1), + scale_factor=0.5, + scale_exponent=1, + ) + # unless they match, that's fine: + cls( + inp, + units=u.Mpc, + comoving=False, + cosmo_factor=cosmo_factor.create(1.0, 1), + scale_factor=1.0, + scale_exponent=1, + ) + # we refuse scale_factor with missing scale_exponent and vice-versa: + with pytest.raises(ValueError): + cls(inp, units=u.Mpc, comoving=False, scale_factor=1.0) + with pytest.raises(ValueError): + cls(inp, units=u.Mpc, comoving=False, scale_exponent=1) + # we refuse overriding an input cosmo_array's information: + with pytest.raises(ValueError): + cls( + cls( + inp, + units=u.Mpc, + comoving=False, + cosmo_factor=cosmo_factor.create(1.0, 1), + ), + units=u.Mpc, + comoving=False, + scale_factor=0.5, + scale_exponent=1, + ) + # unless it matches, that's fine: + cls( + cls( + inp, + units=u.Mpc, + comoving=False, + cosmo_factor=cosmo_factor.create(1.0, 1), + ), + units=u.Mpc, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + ) + + +class TestNumpyFunctions: + """ + Check that numpy functions recognize our cosmo classes as input and handle them + correctly. + """ + + def test_explicitly_handled_funcs(self): + """ + Make sure we at least handle everything that unyt does, and anything that + 'just worked' for unyt but that we need to handle by hand. + + We don't try to be exhaustive here, but at give some basic input to every function + that we expect to be able to take cosmo input. We then use our helpers defined + above to convert the inputs to unyt equivalents and call the numpy function on + both cosmo and unyt input. Then we use our helpers to check the results for + consistency. For instnace, if with unyt input we got back a unyt_array, we + should expect a cosmo_array. + + We are not currently explicitly testing that the results of any specific function + are numerically what we expected them to be (seems like overkill), nor that the + cosmo_factor's make sense given the input. The latter would be a useful addition, + but I can't think of a sensible way to implement this besides writing in the + expectation for every output value of every function by hand. + + As long as no functions outright crash, the test will report the list of functions + that we should have covered that we didn't cover in tests, and/or the list of + functions whose output values were not what we expected based on running them with + unyt input. Otherwise we just get a stack trace of the first function that + crashed. + """ + from unyt._array_functions import _HANDLED_FUNCTIONS + from unyt.tests.test_array_functions import NOOP_FUNCTIONS + + functions_to_check = { + # FUNCTIONS UNYT HANDLES EXPLICITLY: + "array2string": (ca(np.arange(3)),), + "dot": (ca(np.arange(3)), ca(np.arange(3))), + "vdot": (ca(np.arange(3)), ca(np.arange(3))), + "inner": (ca(np.arange(3)), ca(np.arange(3))), + "outer": (ca(np.arange(3)), ca(np.arange(3))), + "kron": (ca(np.arange(3)), ca(np.arange(3))), + "histogram_bin_edges": (ca(np.arange(3)),), + "linalg.inv": (ca(np.eye(3)),), + "linalg.tensorinv": (ca(np.eye(9).reshape((3, 3, 3, 3))),), + "linalg.pinv": (ca(np.eye(3)),), + "linalg.svd": (ca(np.eye(3)),), + "histogram": (ca(np.arange(3)),), + "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), + "histogramdd": (ca(np.arange(3)).reshape((1, 3)),), + "concatenate": (ca(np.eye(3)),), + "cross": (ca(np.arange(3)), ca(np.arange(3))), + "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), + "union1d": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.norm": (ca(np.arange(3)),), + "vstack": (ca(np.arange(3)),), + "hstack": (ca(np.arange(3)),), + "dstack": (ca(np.arange(3)),), + "column_stack": (ca(np.arange(3)),), + "stack": (ca(np.arange(3)),), + "around": (ca(np.arange(3)),), + "block": ([[ca(np.arange(3))], [ca(np.arange(3))]],), + "fft.fft": (ca(np.arange(3)),), + "fft.fft2": (ca(np.eye(3)),), + "fft.fftn": (ca(np.arange(3)),), + "fft.hfft": (ca(np.arange(3)),), + "fft.rfft": (ca(np.arange(3)),), + "fft.rfft2": (ca(np.eye(3)),), + "fft.rfftn": (ca(np.arange(3)),), + "fft.ifft": (ca(np.arange(3)),), + "fft.ifft2": (ca(np.eye(3)),), + "fft.ifftn": (ca(np.arange(3)),), + "fft.ihfft": (ca(np.arange(3)),), + "fft.irfft": (ca(np.arange(3)),), + "fft.irfft2": (ca(np.eye(3)),), + "fft.irfftn": (ca(np.arange(3)),), + "fft.fftshift": (ca(np.arange(3)),), + "fft.ifftshift": (ca(np.arange(3)),), + "sort_complex": (ca(np.arange(3)),), + "isclose": (ca(np.arange(3)), ca(np.arange(3))), + "allclose": (ca(np.arange(3)), ca(np.arange(3))), + "array_equal": (ca(np.arange(3)), ca(np.arange(3))), + "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), + "linspace": (cq(1), cq(2)), + "logspace": (cq(1, unit=u.dimensionless), cq(2, unit=u.dimensionless)), + "geomspace": (cq(1), cq(1)), + "copyto": (ca(np.arange(3)), ca(np.arange(3))), + "prod": (ca(np.arange(3)),), + "var": (ca(np.arange(3)),), + "trace": (ca(np.eye(3)),), + "percentile": (ca(np.arange(3)), 30), + "quantile": (ca(np.arange(3)), 0.3), + "nanpercentile": (ca(np.arange(3)), 30), + "nanquantile": (ca(np.arange(3)), 0.3), + "linalg.det": (ca(np.eye(3)),), + "diff": (ca(np.arange(3)),), + "ediff1d": (ca(np.arange(3)),), + "ptp": (ca(np.arange(3)),), + "cumprod": (ca(np.arange(3)),), + "pad": (ca(np.arange(3)), 3), + "choose": (np.arange(3), ca(np.eye(3))), + "insert": (ca(np.arange(3)), 1, cq(1)), + "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.tensorsolve": ( + ca(np.eye(24).reshape((6, 4, 2, 3, 4))), + ca(np.ones((6, 4))), + ), + "linalg.eig": (ca(np.eye(3)),), + "linalg.eigh": (ca(np.eye(3)),), + "linalg.eigvals": (ca(np.eye(3)),), + "linalg.eigvalsh": (ca(np.eye(3)),), + "savetxt": (savetxt_file, ca(np.arange(3))), + "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), + "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), + "isin": (ca(np.arange(3)), ca(np.arange(3))), + "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), + "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), + "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), + "select": ( + [np.arange(3) < 1, np.arange(3) > 1], + [ca(np.arange(3)), ca(np.arange(3))], + cq(1), + ), + "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))), + "sinc": (ca(np.arange(3)),), + "clip": (ca(np.arange(3)), cq(1), cq(2)), + "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + "triu": (ca(np.ones((3, 3))),), + "tril": (ca(np.ones((3, 3))),), + "einsum": ("ii->i", ca(np.eye(3))), + "convolve": (ca(np.arange(3)), ca(np.arange(3))), + "correlate": (ca(np.arange(3)), ca(np.arange(3))), + "tensordot": (ca(np.eye(3)), ca(np.eye(3))), + "unwrap": (ca(np.arange(3)),), + "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + "array_repr": (ca(np.arange(3)),), + "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), + "trapezoid": (ca(np.arange(3)),), + "in1d": (ca(np.arange(3)), ca(np.arange(3))), # np deprecated + "take": (ca(np.arange(3)), np.arange(3)), + # FUNCTIONS THAT UNYT DOESN'T HANDLE EXPLICITLY (THEY "JUST WORK"): + "all": (ca(np.arange(3)),), + "amax": (ca(np.arange(3)),), # implemented via max + "amin": (ca(np.arange(3)),), # implemented via min + "angle": (cq(complex(1, 1)),), + "any": (ca(np.arange(3)),), + "append": (ca(np.arange(3)), cq(1)), + "apply_along_axis": (lambda x: x, 0, ca(np.eye(3))), + "argmax": (ca(np.arange(3)),), # implemented via max + "argmin": (ca(np.arange(3)),), # implemented via min + "argpartition": (ca(np.arange(3)), 1), # implemented via partition + "argsort": (ca(np.arange(3)),), # implemented via sort + "argwhere": (ca(np.arange(3)),), + "array_str": (ca(np.arange(3)),), + "atleast_1d": (ca(np.arange(3)),), + "atleast_2d": (ca(np.arange(3)),), + "atleast_3d": (ca(np.arange(3)),), + "average": (ca(np.arange(3)),), + "can_cast": (ca(np.arange(3)), np.float64), + "common_type": (ca(np.arange(3)), ca(np.arange(3))), + "result_type": (ca(np.ones(3)), ca(np.ones(3))), + "iscomplex": (ca(np.arange(3)),), + "iscomplexobj": (ca(np.arange(3)),), + "isreal": (ca(np.arange(3)),), + "isrealobj": (ca(np.arange(3)),), + "nan_to_num": (ca(np.arange(3)),), + "nanargmax": (ca(np.arange(3)),), # implemented via max + "nanargmin": (ca(np.arange(3)),), # implemented via min + "nanmax": (ca(np.arange(3)),), # implemented via max + "nanmean": (ca(np.arange(3)),), # implemented via mean + "nanmedian": (ca(np.arange(3)),), # implemented via median + "nanmin": (ca(np.arange(3)),), # implemented via min + "trim_zeros": (ca(np.arange(3)),), + "max": (ca(np.arange(3)),), + "mean": (ca(np.arange(3)),), + "median": (ca(np.arange(3)),), + "min": (ca(np.arange(3)),), + "ndim": (ca(np.arange(3)),), + "shape": (ca(np.arange(3)),), + "size": (ca(np.arange(3)),), + "sort": (ca(np.arange(3)),), + "sum": (ca(np.arange(3)),), + "repeat": (ca(np.arange(3)), 2), + "tile": (ca(np.arange(3)), 2), + "shares_memory": (ca(np.arange(3)), ca(np.arange(3))), + "nonzero": (ca(np.arange(3)),), + "count_nonzero": (ca(np.arange(3)),), + "flatnonzero": (ca(np.arange(3)),), + "isneginf": (ca(np.arange(3)),), + "isposinf": (ca(np.arange(3)),), + "empty_like": (ca(np.arange(3)),), + "full_like": (ca(np.arange(3)), cq(1)), + "ones_like": (ca(np.arange(3)),), + "zeros_like": (ca(np.arange(3)),), + "copy": (ca(np.arange(3)),), + "meshgrid": (ca(np.arange(3)), ca(np.arange(3))), + "transpose": (ca(np.eye(3)),), + "reshape": (ca(np.arange(3)), (3,)), + "resize": (ca(np.arange(3)), 6), + "roll": (ca(np.arange(3)), 1), + "rollaxis": (ca(np.arange(3)), 0), + "rot90": (ca(np.eye(3)),), + "expand_dims": (ca(np.arange(3)), 0), + "squeeze": (ca(np.arange(3)),), + "flip": (ca(np.eye(3)),), + "fliplr": (ca(np.eye(3)),), + "flipud": (ca(np.eye(3)),), + "delete": (ca(np.arange(3)), 0), + "partition": (ca(np.arange(3)), 1), + "broadcast_to": (ca(np.arange(3)), 3), + "broadcast_arrays": (ca(np.arange(3)),), + "split": (ca(np.arange(3)), 1), + "array_split": (ca(np.arange(3)), 1), + "dsplit": (ca(np.arange(27)).reshape(3, 3, 3), 1), + "hsplit": (ca(np.arange(3)), 1), + "vsplit": (ca(np.eye(3)), 1), + "swapaxes": (ca(np.eye(3)), 0, 1), + "moveaxis": (ca(np.eye(3)), 0, 1), + "nansum": (ca(np.arange(3)),), # implemented via sum + "std": (ca(np.arange(3)),), + "nanstd": (ca(np.arange(3)),), + "nanvar": (ca(np.arange(3)),), + "nanprod": (ca(np.arange(3)),), + "diag": (ca(np.eye(3)),), + "diag_indices_from": (ca(np.eye(3)),), + "diagflat": (ca(np.eye(3)),), + "diagonal": (ca(np.eye(3)),), + "ravel": (ca(np.arange(3)),), + "ravel_multi_index": (np.eye(2, dtype=int), (2, 2)), + "unravel_index": (np.arange(3), (3,)), + "fix": (ca(np.arange(3)),), + "round": (ca(np.arange(3)),), # implemented via around + "may_share_memory": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.matrix_power": (ca(np.eye(3)), 2), + "linalg.cholesky": (ca(np.eye(3)),), + "linalg.multi_dot": ((ca(np.eye(3)), ca(np.eye(3))),), + "linalg.matrix_rank": (ca(np.eye(3)),), + "linalg.qr": (ca(np.eye(3)),), + "linalg.slogdet": (ca(np.eye(3)),), + "linalg.cond": (ca(np.eye(3)),), + "gradient": (ca(np.arange(3)),), + "cumsum": (ca(np.arange(3)),), + "nancumsum": (ca(np.arange(3)),), + "nancumprod": (ca(np.arange(3)),), + "bincount": (ca(np.arange(3)),), + "unique": (ca(np.arange(3)),), + "min_scalar_type": (ca(np.arange(3)),), + "extract": (0, ca(np.arange(3))), + "setxor1d": (ca(np.arange(3)), ca(np.arange(3))), + "lexsort": (ca(np.arange(3)),), + "digitize": (ca(np.arange(3)), ca(np.arange(3))), + "tril_indices_from": (ca(np.eye(3)),), + "triu_indices_from": (ca(np.eye(3)),), + "imag": (ca(np.arange(3)),), + "real": (ca(np.arange(3)),), + "real_if_close": (ca(np.arange(3)),), + "einsum_path": ("ij,jk->ik", ca(np.eye(3)), ca(np.eye(3))), + "cov": (ca(np.arange(3)),), + "corrcoef": (ca(np.arange(3)),), + "compress": (np.zeros(3), ca(np.arange(3))), + "take_along_axis": (ca(np.arange(3)), np.ones(3, dtype=int), 0), + "linalg.cross": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.diagonal": (ca(np.eye(3)),), + "linalg.matmul": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.matrix_norm": (ca(np.eye(3)),), + "linalg.matrix_transpose": (ca(np.eye(3)),), + "linalg.svdvals": (ca(np.eye(3)),), + "linalg.tensordot": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.trace": (ca(np.eye(3)),), + "linalg.vecdot": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.vector_norm": (ca(np.arange(3)),), + "astype": (ca(np.arange(3)), float), + "matrix_transpose": (ca(np.eye(3)),), + "unique_all": (ca(np.arange(3)),), + "unique_counts": (ca(np.arange(3)),), + "unique_inverse": (ca(np.arange(3)),), + "unique_values": (ca(np.arange(3)),), + "cumulative_sum": (ca(np.arange(3)),), + "cumulative_prod": (ca(np.arange(3)),), + "unstack": (ca(np.arange(3)),), + } + functions_checked = list() + bad_funcs = dict() + for fname, args in functions_to_check.items(): + # ----- this is to be removed ------ + # ---- see test_block_is_broken ---- + if fname == "block": + # we skip this function due to issue in unyt with unreleased fix + functions_checked.append(np.block) + continue + # ---------------------------------- + ua_args = list() + for arg in args: + ua_args.append(arg_to_ua(arg)) + func = getfunc(fname) + try: + with warnings.catch_warnings(): + if "savetxt" in fname: + warnings.filterwarnings( + action="ignore", + category=UserWarning, + message="numpy.savetxt does not preserve units", + ) + try: + ua_result = func(*ua_args) + except: + print(f"Crashed in {fname} with unyt input.") + raise + except u.exceptions.UnytError: + raises_unyt_error = True + else: + raises_unyt_error = False + if "savetxt" in fname and os.path.isfile(savetxt_file): + os.remove(savetxt_file) + functions_checked.append(func) + if raises_unyt_error: + with pytest.raises(u.exceptions.UnytError): + result = func(*args) + continue + with warnings.catch_warnings(): + if "savetxt" in fname: + warnings.filterwarnings( + action="ignore", + category=UserWarning, + message="numpy.savetxt does not preserve units or cosmo", + ) + if "unwrap" in fname: + # haven't bothered to pass a cosmo_quantity for period + warnings.filterwarnings( + action="ignore", + category=RuntimeWarning, + message="Mixing arguments with and without cosmo_factors", + ) + try: + result = func(*args) + except: + print(f"Crashed in {fname} with cosmo input.") + raise + if fname.split(".")[-1] in ( + "fill_diagonal", + "copyto", + "place", + "put", + "put_along_axis", + "putmask", + ): + # treat inplace modified values for relevant functions as result + result = args[0] + ua_result = ua_args[0] + if "savetxt" in fname and os.path.isfile(savetxt_file): + os.remove(savetxt_file) + ignore_values = fname in {"empty_like"} # empty_like has arbitrary data + try: + check_result(result, ua_result, ignore_values=ignore_values) + except AssertionError: + bad_funcs["np." + fname] = result, ua_result + if len(bad_funcs) > 0: + raise AssertionError( + "Some functions did not return expected types " + "(obtained, obtained with unyt input): " + str(bad_funcs) + ) + unchecked_functions = [ + f + for f in set(_HANDLED_FUNCTIONS) | NOOP_FUNCTIONS + if f not in functions_checked + ] + try: + assert len(unchecked_functions) == 0 + except AssertionError: + raise AssertionError( + "Did not check functions", + [ + (".".join((f.__module__, f.__name__)).replace("numpy", "np")) + for f in unchecked_functions + ], + ) + + @pytest.mark.xfail + def test_block_is_broken(self): + """ + There is an issue in unyt affecting np.block and fixed in + https://github.com/yt-project/unyt/pull/571 + + When this fix is released: + - This test will unexpectedly pass (instead of xfailing). + - Remove lines flagged with a comment in `test_explicitly_handled_funcs`. + - Remove this test. + """ + assert isinstance( + np.block([[ca(np.arange(3))], [ca(np.arange(3))]]), cosmo_array + ) + + # the combinations of units and cosmo_factors is nonsense but it's just for testing... + @pytest.mark.parametrize( + "func_args", + ( + ( + np.histogram, + ( + cosmo_array( + [1, 2, 3], + u.m, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + ), + ), + ), + ( + np.histogram2d, + ( + cosmo_array( + [1, 2, 3], + u.m, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + ), + cosmo_array( + [1, 2, 3], + u.K, + comoving=False, + scale_factor=1.0, + scale_exponent=2, + ), + ), + ), + ( + np.histogramdd, + ( + [ + cosmo_array( + [1, 2, 3], + u.m, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + ), + cosmo_array( + [1, 2, 3], + u.K, + comoving=False, + scale_factor=1.0, + scale_exponent=2, + ), + cosmo_array( + [1, 2, 3], + u.kg, + comoving=False, + scale_factor=1.0, + scale_exponent=3, + ), + ], + ), + ), + ), + ) + @pytest.mark.parametrize( + "weights", + ( + None, + cosmo_array( + [1, 2, 3], u.s, comoving=False, scale_factor=1.0, scale_exponent=1 + ), + np.array([1, 2, 3]), + ), + ) + @pytest.mark.parametrize("bins_type", ("int", "np", "ca")) + @pytest.mark.parametrize("density", (None, True)) + def test_histograms(self, func_args, weights, bins_type, density): + """ + Test that histograms give sensible output. + + Histograms are tricky with possible density and weights arguments, and the way + that attributes need validation and propagation between the bins and values. + They are also commonly used. They therefore need a bespoke test. + """ + func, args = func_args + bins = { + "int": 10, + "np": [np.linspace(0, 5, 11)] * 3, + "ca": [ + cosmo_array( + np.linspace(0, 5, 11), + u.kpc, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + ), + cosmo_array( + np.linspace(0, 5, 11), + u.K, + comoving=False, + scale_factor=1.0, + scale_exponent=2, + ), + cosmo_array( + np.linspace(0, 5, 11), + u.Msun, + comoving=False, + scale_factor=1.0, + scale_exponent=3, + ), + ], + }[bins_type] + bins = ( + bins[ + { + np.histogram: np.s_[0], + np.histogram2d: np.s_[:2], + np.histogramdd: np.s_[:], + }[func] + ] + if bins_type in ("np", "ca") + else bins + ) + result = func(*args, bins=bins, density=density, weights=weights) + ua_args = tuple( + ( + to_ua(arg) + if not isinstance(arg, tuple) + else tuple(to_ua(item) for item in arg) + ) + for arg in args + ) + ua_bins = ( + to_ua(bins) + if not isinstance(bins, tuple) + else tuple(to_ua(item) for item in bins) + ) + ua_result = func( + *ua_args, bins=ua_bins, density=density, weights=to_ua(weights) + ) + if isinstance(ua_result, tuple): + assert isinstance(result, tuple) + assert len(result) == len(ua_result) + for r, ua_r in zip(result, ua_result): + check_result(r, ua_r) + else: + check_result(result, ua_result) + if not density and not isinstance(weights, cosmo_array): + assert not isinstance(result[0], cosmo_array) + else: + assert result[0].comoving is False + if density and not isinstance(weights, cosmo_array): + assert ( + result[0].cosmo_factor + == { + np.histogram: cosmo_factor(a ** -1, 1.0), + np.histogram2d: cosmo_factor(a ** -3, 1.0), + np.histogramdd: cosmo_factor(a ** -6, 1.0), + }[func] + ) + elif density and isinstance(weights, cosmo_array): + assert result[0].comoving is False + assert ( + result[0].cosmo_factor + == { + np.histogram: cosmo_factor(a ** 0, 1.0), + np.histogram2d: cosmo_factor(a ** -2, 1.0), + np.histogramdd: cosmo_factor(a ** -5, 1.0), + }[func] + ) + elif not density and isinstance(weights, cosmo_array): + assert result[0].comoving is False + assert ( + result[0].cosmo_factor + == { + np.histogram: cosmo_factor(a ** 1, 1.0), + np.histogram2d: cosmo_factor(a ** 1, 1.0), + np.histogramdd: cosmo_factor(a ** 1, 1.0), + }[func] + ) + ret_bins = { + np.histogram: [result[1]], + np.histogram2d: result[1:], + np.histogramdd: result[1], + }[func] + for b, expt_cf in zip( + ret_bins, + ( + [ + cosmo_factor(a ** 1, 1.0), + cosmo_factor(a ** 2, 1.0), + cosmo_factor(a ** 3, 1.0), + ] + ), + ): + assert b.comoving is False + assert b.cosmo_factor == expt_cf + + def test_getitem(self): + """ + Make sure that we don't degrade to an ndarray on slicing. + """ + assert isinstance(ca(np.arange(3))[0], cosmo_quantity) + + def test_reshape_to_scalar(self): + """ + Make sure that we convert to a cosmo_quantity when we reshape to a scalar. + """ + assert isinstance(ca(np.ones(1)).reshape(tuple()), cosmo_quantity) + + def test_iter(self): + """ + Make sure that we get cosmo_quantity's when iterating over a cosmo_array. + """ + for cq in ca(np.arange(3)): + assert isinstance(cq, cosmo_quantity) + + def test_dot(self): + """ + Make sure that we get a cosmo_array when we use array attribute dot. + """ + res = ca(np.arange(3)).dot(ca(np.arange(3))) + assert isinstance(res, cosmo_quantity) + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5) + assert res.valid_transform is True + + +class TestCosmoQuantity: + """ + Test that the cosmo_quantity class works as desired, mostly around issues converting + back and forth with cosmo_array. + """ + + @pytest.mark.parametrize( + "func, args", + [ + ("astype", (float,)), + ("in_units", (u.m,)), + ("byteswap", tuple()), + ("compress", ([True],)), + ("flatten", tuple()), + ("ravel", tuple()), + ("repeat", (1,)), + ("reshape", (1,)), + ("take", ([0],)), + ("transpose", tuple()), + ("view", tuple()), + ], + ) + def test_propagation_func(self, func, args): + """ + Test that functions that are supposed to propagate our attributes do so. + """ + cq = cosmo_quantity( + 1, + u.m, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + valid_transform=True, + ) + res = getattr(cq, func)(*args) + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.valid_transform is True + + def test_round(self): + """ + Test that attributes propagate through the round builtin. + """ + cq = cosmo_quantity( + 1.03, + u.m, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + valid_transform=True, + ) + res = round(cq) + assert res.value == 1.0 + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.valid_transform is True + + def test_scalar_return_func(self): + """ + Make sure that default-wrapped functions that take a cosmo_array and return a + scalar convert to a cosmo_quantity. + """ + ca = cosmo_array( + np.arange(3), + u.m, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + valid_transform=True, + ) + res = np.min(ca) + assert isinstance(res, cosmo_quantity) + + @pytest.mark.parametrize("prop", ["T", "ua", "unit_array"]) + def test_propagation_props(self, prop): + """ + Test that properties propagate our attributes as intended. + """ + cq = cosmo_quantity( + 1, + u.m, + comoving=False, + scale_factor=1.0, + scale_exponent=1, + valid_transform=True, + ) + res = getattr(cq, prop) + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.valid_transform is True class TestCosmoArrayCopy: + """ + Tests of explicit (deep)copying of cosmo_array. + """ + def test_copy(self): """ Check that when we copy a cosmo_array it preserves its values and attributes. @@ -60,7 +1021,8 @@ def test_copy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + scale_factor=1.0, + scale_exponent=1, comoving=False, ) copy_arr = copy(arr) @@ -76,7 +1038,8 @@ def test_deepcopy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + scale_factor=1.0, + scale_exponent=1, comoving=False, ) copy_arr = deepcopy(arr) @@ -92,7 +1055,8 @@ def test_to_cgs(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + scale_factor=1.0, + scale_exponent=1, comoving=False, ) cgs_arr = arr.in_cgs() @@ -100,3 +1064,68 @@ def test_to_cgs(self): assert cgs_arr.units == u.cm assert cgs_arr.cosmo_factor == arr.cosmo_factor assert cgs_arr.comoving == arr.comoving + + +class TestMultiplicationByUnyt: + def test_multiplication_by_unyt(self): + """ + We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as + for `cosmo_array(...) * u.Mpc`. + + Right-sided multiplication & division can't be supported without upstream + changes in unyt, see `test_rmultiplication_by_unyt`. + """ + ca = cosmo_array( + np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1 + ) + # required so that can test right-sided division with the same assertions: + assert np.allclose(ca.to_value(ca.units), 1) + # the reference result: + multiplied_by_quantity = ca * (1 * u.Mpc) # parentheses very important here + # get the same result twice through left-sided multiplication and division: + lmultiplied_by_unyt = ca * u.Mpc + ldivided_by_unyt = ca / u.Mpc ** -1 + + for multiplied_by_unyt in (lmultiplied_by_unyt, ldivided_by_unyt): + assert isinstance(multiplied_by_quantity, cosmo_array) + assert isinstance(multiplied_by_unyt, cosmo_array) + assert np.allclose( + multiplied_by_unyt.to_value(multiplied_by_quantity.units), + multiplied_by_quantity.to_value(multiplied_by_quantity.units), + ) + + @pytest.mark.xfail + def test_rmultiplication_by_unyt(self): + """ + We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as + for `cosmo_array(...) * u.Mpc`. + + But unyt will call it's own __mul__ before we get a chance to use our __rmul__ + when the cosmo_array is the right-hand argument. + + We can't handle this case without upstream changes in unyt, so this test is marked + to xfail. + + If this is fixed in the future this test will pass and can be merged with + `test_multiplication_by_unyt` to tidy up. + + See https://github.com/yt-project/unyt/pull/572 + """ + ca = cosmo_array( + np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1 + ) + # required so that can test right-sided division with the same assertions: + assert np.allclose(ca.to_value(ca.units), 1) + # the reference result: + multiplied_by_quantity = ca * (1 * u.Mpc) # parentheses very important here + # get 2x the same result through right-sided multiplication and division: + rmultiplied_by_unyt = u.Mpc * ca + rdivided_by_unyt = u.Mpc ** 2 / ca + + for multiplied_by_unyt in (rmultiplied_by_unyt, rdivided_by_unyt): + assert isinstance(multiplied_by_quantity, cosmo_array) + assert isinstance(multiplied_by_unyt, cosmo_array) + assert np.allclose( + multiplied_by_unyt.to_value(multiplied_by_quantity.units), + multiplied_by_quantity.to_value(multiplied_by_quantity.units), + ) diff --git a/tests/test_cosmo_array_attrs.py b/tests/test_cosmo_array_ufuncs.py similarity index 77% rename from tests/test_cosmo_array_attrs.py rename to tests/test_cosmo_array_ufuncs.py index bfa625d2..d4b1562c 100644 --- a/tests/test_cosmo_array_attrs.py +++ b/tests/test_cosmo_array_ufuncs.py @@ -1,19 +1,32 @@ """ -Tests that functions returning copies of cosmo_array - preserve the comoving and cosmo_factor attributes. +Tests that ufuncs handling cosmo_array's properly handle our extra attributes. """ import pytest import numpy as np import unyt as u -from swiftsimio.objects import cosmo_array, cosmo_factor, a, multiple_output_operators +from swiftsimio.objects import ( + cosmo_array, + cosmo_factor, + a, + multiple_output_operators, + InvalidScaleFactor, +) class TestCopyFuncs: + """ + Test ufuncs that copy arrays. + """ + @pytest.mark.parametrize( ("func"), ["byteswap", "diagonal", "flatten", "ravel", "transpose", "view"] ) def test_argless_copyfuncs(self, func): + """ + Make sure that our attributes are preserved through copying functions that + take no arguments. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -24,6 +37,9 @@ def test_argless_copyfuncs(self, func): assert hasattr(getattr(arr, func)(), "comoving") def test_astype(self): + """ + Make sure that our attributes are preserved through astype. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -35,6 +51,9 @@ def test_astype(self): assert hasattr(res, "comoving") def test_in_units(self): + """ + Make sure that our attributes are preserved through in_units (from unyt). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -46,6 +65,9 @@ def test_in_units(self): assert hasattr(res, "comoving") def test_compress(self): + """ + Make sure that our attributes are preserved through compress. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -57,6 +79,9 @@ def test_compress(self): assert hasattr(res, "comoving") def test_repeat(self): + """ + Make sure that our attributes are preserved through repeat. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -68,6 +93,9 @@ def test_repeat(self): assert hasattr(res, "comoving") def test_T(self): + """ + Make sure that our attributes are preserved through transpose (T). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -79,6 +107,9 @@ def test_T(self): assert hasattr(res, "comoving") def test_ua(self): + """ + Make sure that our attributes are preserved through ua (from unyt). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -90,6 +121,9 @@ def test_ua(self): assert hasattr(res, "comoving") def test_unit_array(self): + """ + Make sure that our attributes are preserved through unit_array (from unyt). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -101,6 +135,10 @@ def test_unit_array(self): assert hasattr(res, "comoving") def test_compatibility(self): + """ + Check that the compatible_with_comoving and compatible_with_physical functions + give correct compatibility checks. + """ # comoving array at high redshift arr = cosmo_array( np.ones((10, 10)), @@ -158,19 +196,37 @@ def test_compatibility(self): class TestCheckUfuncCoverage: + """ + Check that we've wrapped all functions that unyt wraps. + """ + def test_multi_output_coverage(self): + """ + Compare our list of multi_output_operators with unyt's to make sure we cover + everything. + """ assert set(multiple_output_operators.keys()) == set( (np.modf, np.frexp, np.divmod) ) def test_ufunc_coverage(self): + """ + Compare our list of ufuncs with unyt's to make sure we cover everything. + """ assert set(u.unyt_array._ufunc_registry.keys()) == set( cosmo_array._ufunc_registry.keys() ) class TestCosmoArrayUfuncs: + """ + Test some example functions using each of our wrappers for correct output. + """ + def test_preserving_ufunc(self): + """ + Tests of the _preserve_cosmo_factor wrapper. + """ # 1 argument inp = cosmo_array( [2], @@ -183,25 +239,27 @@ def test_preserving_ufunc(self): assert res.comoving is False assert res.cosmo_factor == inp.cosmo_factor # 2 argument, no cosmo_factors - inp = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp = cosmo_array([2], u.kpc, comoving=False) res = inp + inp assert res.to_value(u.kpc) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # 2 argument, one is not cosmo_array inp1 = u.unyt_array([2], u.kpc) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - res = inp1 + inp2 + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.warns(RuntimeWarning, match="Mixing arguments"): + res = inp1 + inp2 assert res.to_value(u.kpc) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # 2 argument, two is not cosmo_array - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = u.unyt_array([2], u.kpc) - res = inp1 + inp2 + with pytest.warns(RuntimeWarning, match="Mixing arguments"): + res = inp1 + inp2 assert res.to_value(u.kpc) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # 2 argument, only one has cosmo_factor inp1 = cosmo_array( [2], @@ -209,25 +267,23 @@ def test_preserving_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 + inp2 - assert res.to_value(u.kpc) == 4 - assert res.comoving is False - assert res.cosmo_factor == inp1.cosmo_factor + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + inp1 + inp2 # 2 argument, only two has cosmo_factor - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = cosmo_array( [2], u.kpc, comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 + inp2 - assert res.to_value(u.kpc) == 4 - assert res.comoving is False - assert res.cosmo_factor == inp2.cosmo_factor + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + inp1 + inp2 # 2 argument, mismatched cosmo_factors inp1 = cosmo_array( [2], @@ -242,9 +298,9 @@ def test_preserving_ufunc(self): cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( - ValueError, match="Ufunc arguments have cosmo_factors that differ" + ValueError, match="Arguments have cosmo_factors that differ" ): - res = inp1 + inp2 + inp1 + inp2 # 2 argument, matched cosmo_factors inp = cosmo_array( [2], @@ -258,21 +314,24 @@ def test_preserving_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor def test_multiplying_ufunc(self): + """ + Tests of the _multiply_cosmo_factor wrapper. + """ # no cosmo_factors - inp = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp = cosmo_array([2], u.kpc, comoving=False) res = inp * inp assert res.to_value(u.kpc ** 2) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # one is not cosmo_array inp1 = 2 - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp2 = cosmo_array([2], u.kpc, comoving=False) res = inp1 * inp2 assert res.to_value(u.kpc) == 4 assert res.comoving is False assert res.cosmo_factor == inp2.cosmo_factor # two is not cosmo_array - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = 2 res = inp1 * inp2 assert res.to_value(u.kpc) == 4 @@ -285,25 +344,19 @@ def test_multiplying_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 * inp2 - assert res.to_value(u.kpc ** 2) == 4 - assert res.comoving is False - assert res.cosmo_factor is None + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): + inp1 * inp2 # only two has cosmo_factor - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = cosmo_array( [2], u.kpc, comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 * inp2 - assert res.to_value(u.kpc ** 2) == 4 - assert res.comoving is False - assert res.cosmo_factor is None + with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): + inp1 * inp2 # cosmo_factors both present inp = cosmo_array( [2], @@ -317,6 +370,9 @@ def test_multiplying_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_dividing_ufunc(self): + """ + Tests of the _divide_cosmo_factor wrapper. + """ inp = cosmo_array( [2.0], u.kpc, @@ -329,6 +385,9 @@ def test_dividing_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 0 def test_return_without_ufunc(self): + """ + Tests of the _return_without_cosmo_factor wrapper. + """ # 1 argument inp = cosmo_array( [1], @@ -363,21 +422,21 @@ def test_return_without_ufunc(self): cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( - ValueError, match="Ufunc arguments have cosmo_factors that differ" + ValueError, match="Arguments have cosmo_factors that differ" ): - res = np.logaddexp(inp1, inp2) + np.logaddexp(inp1, inp2) # 2 arguments, one missing comso_factor - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = cosmo_array( [2], u.kpc, comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = np.logaddexp(inp1, inp2) - assert res == np.logaddexp(2, 2) - assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + np.logaddexp(inp1, inp2) # 2 arguments, two missing comso_factor inp1 = cosmo_array( [2], @@ -385,9 +444,11 @@ def test_return_without_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = np.logaddexp(inp1, inp2) + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + np.logaddexp(inp1, inp2) assert res == np.logaddexp(2, 2) assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) # 2 arguments, one not cosmo_array @@ -398,7 +459,7 @@ def test_return_without_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = np.logaddexp(inp1, inp2) assert res == np.logaddexp(2, 2) assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) @@ -410,12 +471,15 @@ def test_return_without_ufunc(self): cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = u.unyt_array([2], u.kpc) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = np.logaddexp(inp1, inp2) assert res == np.logaddexp(2, 2) assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) def test_sqrt_ufunc(self): + """ + Tests of the _sqrt_cosmo_factor wrapper. + """ inp = cosmo_array( [4], u.kpc, @@ -428,6 +492,9 @@ def test_sqrt_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 0.5 def test_square_ufunc(self): + """ + Tests of the _square_cosmo_factor wrapper. + """ inp = cosmo_array( [2], u.kpc, @@ -440,6 +507,9 @@ def test_square_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_cbrt_ufunc(self): + """ + Tests of the _cbrt_cosmo_factor wrapper. + """ inp = cosmo_array( [8], u.kpc, @@ -452,6 +522,9 @@ def test_cbrt_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** (1.0 / 3.0) def test_reciprocal_ufunc(self): + """ + Tests of the _reciprocal_cosmo_factor wrapper. + """ inp = cosmo_array( [2.0], u.kpc, @@ -464,6 +537,9 @@ def test_reciprocal_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_passthrough_ufunc(self): + """ + Tests of the _passthrough_cosmo_factor wrapper. + """ # 1 argument inp = cosmo_array( [2], @@ -499,10 +575,15 @@ def test_passthrough_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) - with pytest.raises(ValueError): - res = np.copysign(inp1, inp2) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + np.copysign(inp1, inp2) def test_arctan2_ufunc(self): + """ + Tests of the _arctan2_cosmo_factor wrapper. + """ inp = cosmo_array( [2], u.kpc, @@ -515,6 +596,9 @@ def test_arctan2_ufunc(self): assert res.cosmo_factor.a_factor == 1 # also ensures cosmo_factor present def test_comparison_ufunc(self): + """ + Tests of the _comparison_cosmo_factor wrapper. + """ inp1 = cosmo_array( [1], u.kpc, @@ -532,19 +616,35 @@ def test_comparison_ufunc(self): assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) def test_out_arg(self): + """ + Test that our helpers can handle functions with an ``out`` kwarg. + """ inp = cosmo_array( [1], u.kpc, comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - out = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) + out = cosmo_array([np.nan], u.dimensionless, comoving=True) np.abs(inp, out=out) assert out.to_value(u.kpc) == np.abs(inp.to_value(u.kpc)) assert out.comoving is False assert out.cosmo_factor == inp.cosmo_factor + inp = cosmo_array( + [1], + u.kpc, + comoving=False, + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + ) + # make sure we can also pass a non-cosmo type for out without crashing + out = np.array([np.nan]) + np.abs(inp, out=out) + assert out == np.abs(inp.to_value(u.kpc)) def test_reduce_multiply(self): + """ + Test that we can handle the reduce method for the multiply ufunc. + """ inp = cosmo_array( [[1, 2], [3, 4]], u.kpc, @@ -557,6 +657,9 @@ def test_reduce_multiply(self): assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_reduce_divide(self): + """ + Test that we can handle the reduce method for the divide ufunc. + """ inp = cosmo_array( [[1.0, 2.0], [1.0, 4.0], [1.0, 1.0]], u.kpc, @@ -569,6 +672,9 @@ def test_reduce_divide(self): assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_reduce_other(self): + """ + Test that we can handle other ufuncs with a reduce method. + """ inp = cosmo_array( [[1.0, 2.0], [1.0, 2.0]], u.kpc, @@ -581,6 +687,9 @@ def test_reduce_other(self): assert res.cosmo_factor == inp.cosmo_factor def test_multi_output(self): + """ + Test that we can handle functions with multiple return values. + """ # with passthrough inp = cosmo_array( [2.5], @@ -609,6 +718,10 @@ def test_multi_output(self): assert isinstance(res2, np.ndarray) and not isinstance(res2, u.unyt_array) def test_multi_output_with_out_arg(self): + """ + Test that we can handle multiple return values in conjunction with an ``out`` + kwarg. + """ # with two out arrays inp = cosmo_array( [2.5], @@ -616,8 +729,8 @@ def test_multi_output_with_out_arg(self): comoving=False, cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) - out1 = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) - out2 = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) + out1 = cosmo_array([np.nan], u.dimensionless, comoving=True) + out2 = cosmo_array([np.nan], u.dimensionless, comoving=True) np.modf(inp, out=(out1, out2)) assert out1.to_value(u.kpc) == 0.5 assert out2.to_value(u.kpc) == 2.0 @@ -627,6 +740,10 @@ def test_multi_output_with_out_arg(self): assert out2.cosmo_factor == inp.cosmo_factor def test_comparison_with_zero(self): + """ + Test that we don't produce warnings for dangerous comparisons on comparison with + zero. + """ inp1 = cosmo_array( [1, 1, 1], u.kpc, @@ -643,7 +760,7 @@ def test_comparison_with_zero(self): cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = 0.5 - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = inp1 > inp2 assert res.all() inp1 = cosmo_array( @@ -666,7 +783,7 @@ def test_comparison_with_zero(self): cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = np.ones(3) * u.kpc - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): assert (inp1 == inp2).all() inp1 = cosmo_array( [1, 1, 1], diff --git a/tests/test_data.py b/tests/test_data.py index 232864d2..69df9a85 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -5,8 +5,7 @@ be read in. """ -import pytest - +import numpy as np from tests.helper import requires from swiftsimio import load, mask @@ -37,8 +36,8 @@ def test_cosmology_metadata(filename): @requires("cosmological_volume.hdf5") def test_time_metadata(filename): """ - This tests the time metadata and also tests the ability to include two items at once from - the same header attribute. + This tests the time metadata and also tests the ability to include two items at once + from the same header attribute. """ data = load(filename) @@ -153,7 +152,11 @@ def test_cell_metadata_is_valid(filename): mask_region = mask(filename) # Because we sort by offset if we are using the metadata we # must re-order the data to be in the correct order - mask_region.constrain_spatial([[0 * b, b] for b in mask_region.metadata.boxsize]) + mask_region.constrain_spatial( + cosmo_array( + [np.zeros_like(mask_region.metadata.boxsize), mask_region.metadata.boxsize] + ).T + ) data = load(filename, mask=mask_region) cell_size = mask_region.cell_size.to(data.gas.coordinates.units) @@ -161,13 +164,6 @@ def test_cell_metadata_is_valid(filename): offsets = mask_region.offsets["gas"] counts = mask_region.counts["gas"] - # can be removed when issue #128 resolved: - boxsize = cosmo_array( - boxsize, - comoving=True, - cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), - ) - start_offset = offsets stop_offset = offsets + counts @@ -187,11 +183,8 @@ def test_cell_metadata_is_valid(filename): continue # Give it a little wiggle room. - # Mask_region provides unyt_array, not cosmo_array, anticipate warnings. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert max <= upper * 1.05 - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert min > lower * 0.95 + assert max <= upper * 1.05 + assert min > lower * 0.95 @requires("cosmological_volume_dithered.hdf5") @@ -206,7 +199,11 @@ def test_dithered_cell_metadata_is_valid(filename): mask_region = mask(filename) # Because we sort by offset if we are using the metadata we # must re-order the data to be in the correct order - mask_region.constrain_spatial([[0 * b, b] for b in mask_region.metadata.boxsize]) + mask_region.constrain_spatial( + cosmo_array( + [np.zeros_like(mask_region.metadata.boxsize), mask_region.metadata.boxsize] + ).T + ) data = load(filename, mask=mask_region) cell_size = mask_region.cell_size.to(data.dark_matter.coordinates.units) @@ -240,12 +237,8 @@ def test_dithered_cell_metadata_is_valid(filename): continue # Give it a little wiggle room - # Mask_region provides unyt_array, not cosmo_array, anticipate warnings. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert max <= upper * 1.05 - # Mask_region provides unyt_array, not cosmo_array, anticipate warnings. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert min > lower * 0.95 + assert max <= upper * 1.05 + assert min > lower * 0.95 @requires("cosmological_volume.hdf5") @@ -274,27 +267,22 @@ def test_reading_select_region_metadata(filename): selected_coordinates = selected_data.gas.coordinates # Now need to repeat teh selection by hand: - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) - ] - ) + + subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) + ] + ) # We also need to repeat for the thing we just selected; the cells only give # us an _approximate_ selection! - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - selected_subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip( - selected_data.gas.coordinates.T, restrict - ) - ] - ) + selected_subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(selected_data.gas.coordinates.T, restrict) + ] + ) hand_selected_coordinates = full_data.gas.coordinates[subset_mask] @@ -331,27 +319,21 @@ def test_reading_select_region_metadata_not_spatial_only(filename): selected_coordinates = selected_data.gas.coordinates # Now need to repeat the selection by hand: - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) - ] - ) + subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) + ] + ) # We also need to repeat for the thing we just selected; the cells only give # us an _approximate_ selection! - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - selected_subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip( - selected_data.gas.coordinates.T, restrict - ) - ] - ) + selected_subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(selected_data.gas.coordinates.T, restrict) + ] + ) hand_selected_coordinates = full_data.gas.coordinates[subset_mask] diff --git a/tests/test_extraparts.py b/tests/test_extraparts.py index 796d4cc3..1a2f9db0 100644 --- a/tests/test_extraparts.py +++ b/tests/test_extraparts.py @@ -1,6 +1,7 @@ """ Test for extra particle types. """ + from swiftsimio import load, metadata from swiftsimio import Writer from swiftsimio.units import cosmo_units diff --git a/tests/test_mask.py b/tests/test_mask.py index 1a75b498..ae071ab4 100644 --- a/tests/test_mask.py +++ b/tests/test_mask.py @@ -6,7 +6,8 @@ from swiftsimio import load, mask import numpy as np -from unyt import unyt_array as array, dimensionless +from unyt import dimensionless +from swiftsimio import cosmo_array @requires("cosmological_volume.hdf5") @@ -22,8 +23,8 @@ def test_reading_select_region_spatial(filename): mask_region = mask(filename, spatial_only=True) mask_region_nospatial = mask(filename, spatial_only=False) - restrict = array( - [[0.0, 0.0, 0.0] * full_data.metadata.boxsize, full_data.metadata.boxsize * 0.5] + restrict = cosmo_array( + [np.zeros_like(full_data.metadata.boxsize), full_data.metadata.boxsize * 0.5] ).T mask_region.constrain_spatial(restrict=restrict) @@ -54,11 +55,8 @@ def test_reading_select_region_half_box(filename): # Mask off the lower bottom corner of the volume. mask_region = mask(filename, spatial_only=True) - restrict = array( - [ - [0.0, 0.0, 0.0] * full_data.metadata.boxsize, - full_data.metadata.boxsize * 0.49, - ] + restrict = cosmo_array( + [np.zeros_like(full_data.metadata.boxsize), full_data.metadata.boxsize * 0.49] ).T mask_region.constrain_spatial(restrict=restrict) diff --git a/tests/test_physical_conversion.py b/tests/test_physical_conversion.py index 26b72d29..2ade790f 100644 --- a/tests/test_physical_conversion.py +++ b/tests/test_physical_conversion.py @@ -10,7 +10,12 @@ def test_convert(filename): """ data = load(filename) coords = data.gas.coordinates + units = coords.units coords_physical = coords.to_physical() - assert array_equal(coords * data.metadata.a, coords_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) + ) return diff --git a/tests/test_read_ic.py b/tests/test_read_ic.py index 4db743ac..d1f39f5c 100644 --- a/tests/test_read_ic.py +++ b/tests/test_read_ic.py @@ -1,6 +1,7 @@ from swiftsimio import load from swiftsimio import Writer from swiftsimio.units import cosmo_units +from swiftsimio import cosmo_array import unyt import numpy as np @@ -70,7 +71,9 @@ def test_reading_ic_units(simple_snapshot_data, field): data = load(test_filename) - assert unyt.array.allclose_units( + assert isinstance(getattr(data.gas, field), cosmo_array) + # np.allclose checks unit consistency + assert np.allclose( getattr(data.gas, field), getattr(simple_snapshot_data.gas, field), rtol=1.0e-4 ) return diff --git a/tests/test_rotate_visualisations.py b/tests/test_rotate_visualisations.py index 694f6213..1e73aec5 100644 --- a/tests/test_rotate_visualisations.py +++ b/tests/test_rotate_visualisations.py @@ -6,7 +6,6 @@ from swiftsimio.visualisation.rotation import rotation_matrix_from_vector from numpy import array_equal from os import remove -import pytest @requires("cosmological_volume.hdf5") @@ -29,7 +28,6 @@ def test_project(filename): centre = data.gas.coordinates[0] rotate_vec = [0.5, 0.5, 0.5] matrix = rotation_matrix_from_vector(rotate_vec, axis="z") - boxsize = data.metadata.boxsize unrotated = project_gas(data, resolution=1024, project="masses", parallel=True) @@ -67,7 +65,6 @@ def test_slice(filename): centre = data.gas.coordinates[0] rotate_vec = [0.5, 0.5, 0.5] matrix = rotation_matrix_from_vector(rotate_vec, axis="z") - boxsize = data.metadata.boxsize slice_z = centre[2] @@ -114,7 +111,6 @@ def test_render(filename): centre = data.gas.coordinates[0] rotate_vec = [0.5, 0.5, 0.5] matrix = rotation_matrix_from_vector(rotate_vec, axis="z") - boxsize = data.metadata.boxsize unrotated = render_gas(data, resolution=256, project="masses", parallel=True) diff --git a/tests/test_soap.py b/tests/test_soap.py index 66ec93fa..89401eb8 100644 --- a/tests/test_soap.py +++ b/tests/test_soap.py @@ -4,13 +4,12 @@ from tests.helper import requires -from swiftsimio import load, mask -import unyt +from swiftsimio import load, mask, cosmo_quantity @requires("soap_example.hdf5") def test_soap_can_load(filename): - data = load(filename) + load(filename) return @@ -43,8 +42,20 @@ def test_soap_can_mask_non_spatial(filename): def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): this_mask = mask(filename, spatial_only=False) - lower = unyt.unyt_quantity(1e5, "Msun") - upper = unyt.unyt_quantity(1e13, "Msun") + lower = cosmo_quantity( + 1e5, + "Msun", + comoving=True, + scale_factor=this_mask.metadata.scale_factor, + scale_exponent=0, + ) + upper = cosmo_quantity( + 1e13, + "Msun", + comoving=True, + scale_factor=this_mask.metadata.scale_factor, + scale_exponent=0, + ) this_mask.constrain_mask( "spherical_overdensity_200_mean", "total_mass", lower, upper ) @@ -60,9 +71,7 @@ def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): masses2 = data2.spherical_overdensity_200_mean.total_mass # Manually mask - custom_mask = ( - unyt.unyt_array(masses2.to_value(masses2.units), masses2.units) >= lower - ) & (unyt.unyt_array(masses2.to_value(masses2.units), masses2.units) <= upper) + custom_mask = (masses2 >= lower) & (masses2 <= upper) assert len(masses2[custom_mask]) == len(masses) diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index 99a91790..3059fcc4 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -1,37 +1,33 @@ import pytest -from swiftsimio import load -from swiftsimio.visualisation import scatter, slice, volume_render -from swiftsimio.visualisation.projection import ( - scatter_parallel, - project_gas, - project_pixel_grid, +from swiftsimio import load, mask +from swiftsimio.visualisation.projection import project_gas, project_pixel_grid +from swiftsimio.visualisation.slice import slice_gas +from swiftsimio.visualisation.volume_render import render_gas +from swiftsimio.visualisation.ray_trace import panel_gas + +from swiftsimio.visualisation.slice_backends import ( + backends as slice_backends, + backends_parallel as slice_backends_parallel, ) -from swiftsimio.visualisation.slice import ( - slice_scatter, - slice_scatter_parallel, - slice_gas, +from swiftsimio.visualisation.volume_render_backends import ( + backends as volume_render_backends, + backends_parallel as volume_render_backends_parallel, ) -from swiftsimio.visualisation.volume_render import render_gas -from swiftsimio.visualisation.volume_render import scatter as volume_scatter +from swiftsimio.visualisation.projection_backends import ( + backends as projection_backends, + backends_parallel as projection_backends_parallel, +) + from swiftsimio.visualisation.power_spectrum import ( deposit, deposition_to_power_spectrum, render_to_deposit, folded_depositions_to_power_spectrum, ) -from swiftsimio.visualisation.ray_trace import panel_gas -from swiftsimio.visualisation.smoothing_length import generate_smoothing_lengths -from swiftsimio.visualisation.projection_backends import ( - backends as projection_backends, - backends_parallel as projection_backends_parallel, -) -from swiftsimio.visualisation.slice_backends import ( - backends as slice_backends, - backends_parallel as slice_backends_parallel, -) from swiftsimio.visualisation.smoothing_length import generate_smoothing_lengths from swiftsimio.optional_packages import CudaSupportError, CUDA_AVAILABLE -from swiftsimio.objects import cosmo_array, a + +from swiftsimio.objects import cosmo_array, cosmo_quantity, a from unyt.array import unyt_array import unyt @@ -54,13 +50,13 @@ def test_scatter(save=False): for backend in projection_backends.keys(): try: image = projection_backends[backend]( - np.array([0.0, 1.0, 1.0, -0.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([1.0, 1.0, 1.0, 1.0]), - np.array([0.2, 0.2, 0.2, 0.000002]), - 256, - 1.0, - 1.0, + x=np.array([0.0, 1.0, 1.0, -0.000_001]), + y=np.array([0.0, 0.0, 1.0, 1.000_001]), + m=np.array([1.0, 1.0, 1.0, 1.0]), + h=np.array([0.2, 0.2, 0.2, 0.000_002]), + res=256, + box_x=1.0, + box_y=1.0, ) except CudaSupportError: if CUDA_AVAILABLE: @@ -75,7 +71,7 @@ def test_scatter(save=False): def test_scatter_mass_conservation(): - np.random.seed(971263) + np.random.seed(971_263) # Width of 0.8 centered on 0.5, 0.5. x = 0.8 * np.random.rand(100) + 0.1 y = 0.8 * np.random.rand(100) + 0.1 @@ -86,11 +82,12 @@ def test_scatter_mass_conservation(): total_mass = np.sum(m) for resolution in resolutions: - image = scatter(x, y, m, h, resolution, 1.0, 1.0) + scatter = projection_backends["fast"] + image = scatter(x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0) mass_in_image = image.sum() / (resolution ** 2) # Check mass conservation to 5% - assert np.isclose(mass_in_image, total_mass, 0.05) + assert np.isclose(mass_in_image.view(np.ndarray), total_mass, 0.05) return @@ -113,9 +110,25 @@ def test_scatter_parallel(save=False): hsml = np.random.rand(number_of_parts).astype(np.float32) * h_max masses = np.ones(number_of_parts, dtype=np.float32) - image = scatter(coordinates[0], coordinates[1], masses, hsml, resolution, 1.0, 1.0) + scatter = projection_backends["fast"] + scatter_parallel = projection_backends_parallel["fast"] + image = scatter( + x=coordinates[0], + y=coordinates[1], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, + ) image_par = scatter_parallel( - coordinates[0], coordinates[1], masses, hsml, resolution, 1.0, 1.0 + x=coordinates[0], + y=coordinates[1], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, ) if save: @@ -127,17 +140,19 @@ def test_scatter_parallel(save=False): def test_slice(save=False): + slice = slice_backends["sph"] image = slice( - np.array([0.0, 1.0, 1.0, -0.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([1.0, 1.0, 1.0, 1.0]), - np.array([0.2, 0.2, 0.2, 0.000002]), - 0.99, - 256, - 1.0, - 1.0, - 1.0, + x=np.array([0.0, 1.0, 1.0, -0.000_001]), + y=np.array([0.0, 0.0, 1.0, 1.000_001]), + z=np.array([0.0, 0.0, 1.0, 1.000_001]), + m=np.array([1.0, 1.0, 1.0, 1.0]), + h=np.array([0.2, 0.2, 0.2, 0.000_002]), + z_slice=0.99, + xres=256, + yres=256, + box_x=1.0, + box_y=1.0, + box_z=1.0, ) if save: @@ -167,30 +182,30 @@ def test_slice_parallel(save=False): for backend in slice_backends.keys(): image = slice_backends[backend]( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - z_slice, - resolution, - resolution, - 1.0, - 1.0, - 1.0, + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + z_slice=z_slice, + xres=resolution, + yres=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, ) image_par = slice_backends_parallel[backend]( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - z_slice, - resolution, - resolution, - 1.0, - 1.0, - 1.0, + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + z_slice=z_slice, + xres=resolution, + yres=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, ) assert np.isclose(image, image_par).all() @@ -203,16 +218,17 @@ def test_slice_parallel(save=False): def test_volume_render(): # render image - volume_render.scatter( - np.array([0.0, 1.0, 1.0, -0.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([1.0, 1.0, 1.0, 1.0]), - np.array([0.2, 0.2, 0.2, 0.000002]), - 64, - 1.0, - 1.0, - 1.0, + scatter = volume_render_backends["scatter"] + scatter( + x=np.array([0.0, 1.0, 1.0, -0.000_001]), + y=np.array([0.0, 0.0, 1.0, 1.000_001]), + z=np.array([0.0, 0.0, 1.0, 1.000_001]), + m=np.array([1.0, 1.0, 1.0, 1.0]), + h=np.array([0.2, 0.2, 0.2, 0.000_002]), + res=64, + box_x=1.0, + box_y=1.0, + box_z=1.0, ) return @@ -231,28 +247,29 @@ def test_volume_parallel(): hsml = np.random.rand(number_of_parts).astype(np.float32) * h_max masses = np.ones(number_of_parts, dtype=np.float32) - image = volume_render.scatter( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - resolution, - 1.0, - 1.0, - 1.0, - ) - image_par = volume_render.scatter_parallel( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - resolution, - 1, - 1.0, - 1.0, - 1.0, + scatter = volume_render_backends["scatter"] + image = scatter( + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + ) + scatter_parallel = volume_render_backends_parallel["scatter"] + image_par = scatter_parallel( + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, ) assert np.isclose(image, image_par).all() @@ -274,7 +291,7 @@ def test_selection_render(filename): project_gas(data, 256, parallel=True, region=[0 * bs, 0.001 * bs] * 2) # render non-square project_gas( - data, 256, parallel=True, region=[0 * bs, 0.00 * bs, 0.25 * bs, 0.75 * bs] + data, 256, parallel=True, region=[0 * bs, 0.50 * bs, 0.25 * bs, 0.75 * bs] ) # Slicing @@ -317,20 +334,38 @@ def test_render_outside_region(): h = 10 ** np.random.rand(number_of_parts) - 1.0 h[h > 0.5] = 0.05 m = np.ones_like(h) - projection_backends["histogram"](x, y, m, h, resolution, 1.0, 1.0) + projection_backends["histogram"]( + x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0 + ) for backend in projection_backends.keys(): try: - projection_backends[backend](x, y, m, h, resolution, 1.0, 1.0) + projection_backends[backend]( + x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0 + ) except CudaSupportError: if CUDA_AVAILABLE: raise ImportError("Optional loading of the CUDA module is broken") else: continue - slice_scatter_parallel(x, y, z, m, h, 0.2, resolution, 1.0, 1.0, 1.0) + slice_backends_parallel["sph"]( + x=x, + y=y, + z=z, + m=m, + h=h, + z_slice=0.2, + xres=resolution, + yres=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + ) - volume_render.scatter_parallel(x, y, z, m, h, resolution, 1, 1.0, 1.0, 1.0) + volume_render_backends_parallel["scatter"]( + x=x, y=y, z=z, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, box_z=1.0 + ) @requires("cosmological_volume.hdf5") @@ -339,40 +374,77 @@ def test_comoving_versus_physical(filename): Test what happens if you try to mix up physical and comoving quantities. """ + # this test is pretty slow if we don't mask out some particles + m = mask(filename) + boxsize = m.metadata.boxsize + m.constrain_spatial([[0.0 * b, 0.2 * b] for b in boxsize]) + region = [ + 0.0 * boxsize[0], + 0.2 * boxsize[0], + 0.0 * boxsize[1], + 0.2 * boxsize[1], + 0.0 * boxsize[2], + 0.2 * boxsize[2], + ] for func, aexp in [(project_gas, -2.0), (slice_gas, -3.0), (render_gas, -3.0)]: # normal case: everything comoving - data = load(filename) + data = load(filename, mask=m) # we force the default (project="masses") to check the cosmo_factor # conversion in this case - img = func(data, resolution=256, project=None) - assert img.comoving + img = func(data, resolution=64, project=None, region=region) + assert data.gas.masses.comoving and img.comoving assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0 - img = func(data, resolution=256, project="densities") - assert img.comoving + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving and img.comoving assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 - # try to mix comoving coordinates with a physical variable + # try to mix comoving coordinates with a physical variable: + # the coordinates should convert to physical internally and warn data.gas.densities.convert_to_physical() - with pytest.raises(AttributeError, match="not compatible with comoving"): - img = func(data, resolution=256, project="densities") - # convert coordinates to physical (but not smoothing lengths) + with pytest.warns( + UserWarning, match="Converting smoothing lengths to physical." + ): + with pytest.warns( + UserWarning, match="Converting coordinate grid to physical." + ): + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving is False and img.comoving is False + assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 + # convert coordinates to physical (but not smoothing lengths): + # the coordinates (copy) should convert back to comoving to match the masses data.gas.coordinates.convert_to_physical() - with pytest.raises(AttributeError, match=""): - img = func(data, resolution=256, project="masses") + with pytest.warns(UserWarning, match="Converting coordinate grid to comoving."): + img = func(data, resolution=64, project="masses", region=region) + assert data.gas.masses.comoving and img.comoving + assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0 # also convert smoothing lengths to physical + # everything should still convert back to comoving to match masses data.gas.smoothing_lengths.convert_to_physical() - # masses are always compatible with either - img = func(data, resolution=256, project="masses") - # check that we get a physical result - assert not img.comoving + with pytest.warns( + UserWarning, match="Converting smoothing lengths to comoving." + ): + with pytest.warns( + UserWarning, match="Converting coordinate grid to comoving." + ): + img = func(data, resolution=64, project="masses", region=region) + assert data.gas.masses.comoving and img.comoving assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0 - # densities are still compatible with physical - img = func(data, resolution=256, project="densities") - assert not img.comoving + # densities are physical, make sure this works with physical coordinates and + # smoothing lengths + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving is False and img.comoving is False assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 - # now try again with comoving densities + # now try again with comoving densities, should work and give a comoving img + # with internal conversions to comoving data.gas.densities.convert_to_comoving() - with pytest.raises(AttributeError, match="not compatible with physical"): - img = func(data, resolution=256, project="densities") + with pytest.warns( + UserWarning, match="Converting smoothing lengths to comoving." + ): + with pytest.warns( + UserWarning, match="Converting coordinate grid to comoving." + ): + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving and img.comoving + assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 @requires("cosmological_volume.hdf5") @@ -381,17 +453,12 @@ def test_nongas_smoothing_lengths(filename): Test that the visualisation tools to calculate smoothing lengths give usable results. """ - # If project_pixel_grid runs without error the smoothing lengths seem usable. + # If project_gas runs without error the smoothing lengths seem usable. data = load(filename) data.dark_matter.smoothing_length = generate_smoothing_lengths( data.dark_matter.coordinates, data.metadata.boxsize, kernel_gamma=1.8 ) - project_pixel_grid( - data.dark_matter, - boxsize=data.metadata.boxsize, - resolution=256, - project="masses", - ) + project_pixel_grid(data.dark_matter, resolution=256, project="masses") assert isinstance(data.dark_matter.smoothing_length, cosmo_array) # We should also be able to use a unyt_array (rather than cosmo_array) as input, @@ -536,7 +603,8 @@ def test_periodic_boundary_wrapping(): assert (image1 == image2).all() # test the volume rendering scatter function - image1 = volume_render.scatter( + scatter = volume_render_backends["scatter"] + image1 = scatter( x=coordinates_periodic[:, 0], y=coordinates_periodic[:, 1], z=coordinates_periodic[:, 2], @@ -547,8 +615,7 @@ def test_periodic_boundary_wrapping(): box_y=boxsize, box_z=boxsize, ) - - image2 = volume_render.scatter( + image2 = scatter( x=coordinates_non_periodic[:, 0], y=coordinates_non_periodic[:, 1], z=coordinates_non_periodic[:, 2], @@ -581,20 +648,20 @@ def test_volume_render_and_unfolded_deposit(): # 1.0 implies no folding deposition = deposit(x, y, z, m, res, 1.0, boxsize, boxsize, boxsize) - # Need to norm for the volume render - volume = volume_scatter( - x / boxsize, - y / boxsize, - z / boxsize, - m, - h / boxsize, - res, - boxsize, - boxsize, - boxsize, + # Need to norm coords and box for the volume render + volume = volume_render_backends["scatter"]( + x=x / boxsize, + y=y / boxsize, + z=z / boxsize, + m=m, + h=h / boxsize, + res=res, + box_x=boxsize, + box_y=boxsize, + box_z=boxsize, ) - assert np.allclose(deposition, volume) + assert np.allclose(deposition, volume.view(np.ndarray)) def test_folding_deposit(): @@ -606,7 +673,6 @@ def test_folding_deposit(): y = np.array([100, 200]) z = np.array([100, 200]) m = np.array([1, 1]) - h = np.array([1e-10, 1e-10]) res = 256 boxsize = 1.0 * res @@ -634,12 +700,20 @@ def test_volume_render_and_unfolded_deposit_with_units(filename): # Volume render the particles volume = render_gas(data, npix, parallel=False).to_physical() - mean_density_deposit = (np.sum(deposition) / npix ** 3).to("Msun / kpc**3").v - mean_density_volume = (np.sum(volume) / npix ** 3).to("Msun / kpc**3").v + mean_density_deposit = ( + (np.sum(deposition) / npix ** 3) + .to_comoving() + .to_value(unyt.solMass / unyt.kpc ** 3) + ) + mean_density_volume = ( + (np.sum(volume) / npix ** 3) + .to_comoving() + .to_value(unyt.solMass / unyt.kpc ** 3) + ) mean_density_calculated = ( - (np.sum(data.gas.masses) / (data.metadata.boxsize[0] * data.metadata.a) ** 3) - .to("Msun / kpc**3") - .v + (np.sum(data.gas.masses) / np.prod(data.metadata.boxsize)) + .to_comoving() + .to_value(unyt.solMass / unyt.kpc ** 3) ) assert np.isclose(mean_density_deposit, mean_density_calculated) @@ -658,13 +732,23 @@ def test_dark_matter_power_spectrum(filename, save=False): # Collate a bunch of raw depositions folds = {} - min_k = 1e-2 / unyt.Mpc - max_k = 1e2 / unyt.Mpc - - bins = unyt.unyt_array( - np.logspace(np.log10(min_k.v), np.log10(max_k.v), 32), units=min_k.units + min_k = cosmo_quantity( + 1e-2, + unyt.Mpc ** -1, + comoving=True, + scale_factor=data.metadata.scale_factor, + scale_exponent=-1, + ) + max_k = cosmo_quantity( + 1e2, + unyt.Mpc ** -1, + comoving=True, + scale_factor=data.metadata.scale_factor, + scale_exponent=-1, ) + bins = np.geomspace(min_k, max_k, 32) + output = {} for npix in [32, 128]: # Deposit the particles @@ -703,7 +787,7 @@ def test_dark_matter_power_spectrum(filename, save=False): _, all_centers, all_ps, folding_tracker = folded_depositions_to_power_spectrum( depositions=folds, - box_size=data.metadata.boxsize, + boxsize=data.metadata.boxsize, number_of_wavenumber_bins=32, cross_depositions=None, wavenumber_range=(min_k, max_k),