"""Create image slices through a volume."""
import numpy as np
from swiftsimio import SWIFTDataset, cosmo_array, cosmo_quantity
from swiftsimio.reader import __SWIFTGroupDataset
from swiftsimio.visualisation.slice_backends import backends, backends_parallel
from swiftsimio.visualisation.smoothing_length import backends_get_hsml
from swiftsimio.visualisation._vistools import (
_get_projection_field,
_get_region_info,
_get_rotated_and_wrapped_coordinates,
backend_strip_and_restore_cosmo_and_units,
)
[docs]
def slice_pixel_grid(
data: __SWIFTGroupDataset,
resolution: int,
z_slice: cosmo_quantity | None = None,
project: str | None = "masses",
parallel: bool = False,
rotation_matrix: np.ndarray | None = None,
rotation_center: cosmo_array | None = None,
region: cosmo_array | None = None,
backend: str = "sph",
periodic: bool = True,
mask: np.ndarray | None = None,
) -> cosmo_array:
"""
Create a data field-weighted 2D slice through a particle dataset as a pixel grid.
Parameters
----------
data : __SWIFTGroupDataset
Particle dataset to slice (e.g. ``data.gas``, ``data.dark_matter``).
resolution : int
Specifies size of return np.array.
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. 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.ndarray, 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 : cosmo_array, optional
Center of the rotation. If you are trying to rotate around a galaxy, this
should be the most bound particle.
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:
[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" (default) 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``.
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.
Returns
-------
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
--------
project_pixel_grid
Creates a 2D projection of a particle dataset.
render_voxel_grid
Creates a 3D voxel grid from a particle dataset.
slice_gas
Convenience wrapper for slicing gas particles.
"""
z_slice = np.zeros_like(data.metadata.boxsize[0]) if z_slice is None else z_slice
m = _get_projection_field(data, project)
region_info = _get_region_info(data, region, periodic=periodic)
hsml = backends_get_hsml[backend](data)
x, y, z = _get_rotated_and_wrapped_coordinates(
data, rotation_matrix, rotation_center, periodic
)
z_center = (
rotation_center[2]
if rotation_center is not None
else np.zeros_like(data.metadata.boxsize[2])
)
mask = mask if mask is not None else np.s_[...]
# determine the effective number of pixels for each dimension
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"]))
normed_x = (x[mask] - region_info["x_min"]) / region_info["max_range"]
normed_y = (y[mask] - region_info["y_min"]) / region_info["max_range"]
normed_z = z[mask] / region_info["max_range"]
normed_z_slice = (z_slice + z_center) / region_info["max_range"]
if periodic:
# place everything in the region inside [0, 1], the backend will tile as needed
normed_x %= region_info["periodic_box_x"]
normed_y %= region_info["periodic_box_y"]
normed_z %= region_info["periodic_box_z"]
normed_z_slice %= region_info["periodic_box_z"]
# Apply the mask to the other arrays
m = m[mask]
hsml = hsml[mask]
kwargs = dict(
x=normed_x,
y=normed_y,
z=normed_z,
m=m,
h=hsml / region_info["max_range"],
z_slice=normed_z_slice,
xres=xres,
yres=yres,
box_x=region_info["periodic_box_x"],
box_y=region_info["periodic_box_y"],
box_z=region_info["periodic_box_z"],
)
norm = region_info["max_range"] ** 3
backend_func = (backends_parallel if parallel else backends)[backend]
image = backend_strip_and_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs)
return image
[docs]
def slice_gas(
data: SWIFTDataset,
resolution: int,
z_slice: cosmo_quantity | None = None,
project: str | None = "masses",
parallel: bool = False,
rotation_matrix: np.ndarray | None = None,
rotation_center: cosmo_array | None = None,
region: cosmo_array | None = None,
backend: str = "sph",
periodic: bool = True,
mask: np.ndarray | None = None,
) -> cosmo_array:
"""
Create a data field-weighted 2D slice through the gas of a SWIFT dataset as a pixel grid.
Parameters
----------
data : SWIFTDataset
Dataset from which slice is extracted.
resolution : int
Specifies size of return np.array.
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. 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.ndarray, 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 : cosmo_array, optional
Center of the rotation. If you are trying to rotate around a galaxy, this
should be the most bound particle.
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:
[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" (default) 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``.
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.
Returns
-------
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
--------
slice_pixel_grid
Slices any particle type, not just gas.
render_voxel_grid
Creates a 3D voxel grid from a particle dataset.
"""
return slice_pixel_grid(
data=data.gas,
resolution=resolution,
z_slice=z_slice,
project=project,
parallel=parallel,
rotation_matrix=rotation_matrix,
rotation_center=rotation_center,
region=region,
backend=backend,
periodic=periodic,
mask=mask,
)