def get_image(n):
"""
Gets the image for snapshot n, and also returns the associated
SWIFT metadata object.
"""
filename = f"{snapshot_name}_{n:04d}.hdf5"
data = load(filename)
boxsize = data.metadata.boxsize[0].value
output = np.zeros((resolution, resolution * 4), dtype=float)
x, y, z = data.gas.coordinates.value.T
# This is an oblong box but we can only make squares!
for box, box_edges in enumerate([[0.0, 1.1], [0.9, 2.1], [1.9, 3.1],
[2.9, 4.0]]):
mask = np.logical_and(x >= box_edges[0], x <= box_edges[1])
masked_x = x[mask] - np.float64(box)
masked_y = y[mask]
masked_z = z[mask]
hsml = data.gas.smoothing_length.value[mask]
if plot == "density":
mass = data.gas.masses.value[mask]
image = slice(
x=masked_y,
y=masked_x,
z=masked_z,
m=mass,
h=hsml,
z_slice=0.5,
res=resolution,
)
else:
quantity = getattr(data.gas, plot).value[mask]
# Need to divide out the particle density for non-projected density quantities
image = scatter(
x=masked_y,
y=masked_x,
z=masked_z,
m=quantity,
h=hsml,
z_slice=0.5,
res=resolution,
) / scatter(
x=masked_y,
y=masked_x,
z=masked_z,
m=np.ones_like(quantity),
h=hsml,
z_slice=0.5,
res=resolution,
)
output[:, box * resolution:(box + 1) * resolution] = image
return output, data.metadata
def test_slice(save=False):
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,
)
if save:
imsave("test_image_creation.png", image)
return
def test_slice_parallel(save=False):
"""
Asserts that we create the same image with the parallel version of the code
as with the serial version.
"""
number_of_parts = 1000
h_max = np.float32(0.05)
z_slice = 0.5
resolution = 256
coordinates = (
np.random.rand(3 * number_of_parts)
.reshape((3, number_of_parts))
.astype(np.float64)
)
hsml = np.random.rand(number_of_parts).astype(np.float32) * h_max
masses = np.ones(number_of_parts, dtype=np.float32)
image = slice(
coordinates[0],
coordinates[1],
coordinates[2],
masses,
hsml,
z_slice,
resolution,
)
image_par = slice_scatter_parallel(
coordinates[0],
coordinates[1],
coordinates[2],
masses,
hsml,
z_slice,
resolution,
)
if save:
imsave("test_image_creation.png", image)
assert np.isclose(image, image_par).all()
return