From 1799bdf3559d072e5af3869d5c73b36ad2e8ac83 Mon Sep 17 00:00:00 2001 From: Chris Havlin Date: Wed, 25 Oct 2023 13:13:21 -0500 Subject: [PATCH 1/4] initial implementation --- yt/utilities/lib/pixelization_routines.pyx | 301 ++++++++++++++++-- yt/visualization/fixed_resolution.py | 1 + yt/visualization/plot_window.py | 7 + .../volume_rendering/off_axis_projection.py | 10 + 4 files changed, 301 insertions(+), 18 deletions(-) diff --git a/yt/utilities/lib/pixelization_routines.pyx b/yt/utilities/lib/pixelization_routines.pyx index 72197abdb3c..a053c558d51 100644 --- a/yt/utilities/lib/pixelization_routines.pyx +++ b/yt/utilities/lib/pixelization_routines.pyx @@ -1284,6 +1284,191 @@ def pixelize_sph_kernel_projection( return mask +@cython.initializedcheck(False) +@cython.boundscheck(False) +@cython.wraparound(False) +@cython.cdivision(True) +def pixelize_sph_kernel_projection_with_depth_check( + np.float64_t[:, :] buff, + np.uint8_t[:, :] mask, + any_float[:] posx, + any_float[:] posy, + any_float[:] posz, + any_float[:] hsml, + any_float[:] pmass, + any_float[:] pdens, + any_float[:] quantity_to_smooth, + bounds, + kernel_name="cubic", + weight_field=None, + int check_period=1, + period=None): + + cdef np.intp_t xsize, ysize + cdef np.float64_t x_min, x_max, y_min, y_max, z_min, z_max, prefactor_j + cdef np.int64_t xi, yi, x0, x1, y0, y1, xxi, yyi + cdef np.float64_t q_ij2, posx_diff, posy_diff, ih_j2 + cdef np.float64_t x, y, dx, dy, idx, idy, h_j2, px, py, pz + cdef np.float64_t period_x = 0, period_y = 0, period_z = 0 + cdef int i, j, ii, jj, kk + cdef np.float64_t[:] _weight_field + cdef int * xiter + cdef int * yiter + cdef int * ziter + cdef np.float64_t * xiterv + cdef np.float64_t * yiterv + cdef np.float64_t * ziterv + + if weight_field is not None: + _weight_field = weight_field + + if period is not None: + period_x = period[0] + period_y = period[1] + period_z = period[2] + + # we find the x and y range over which we have pixels and we find how many + # pixels we have in each dimension + xsize, ysize = buff.shape[0], buff.shape[1] + x_min = bounds[0] + x_max = bounds[1] + y_min = bounds[2] + y_max = bounds[3] + + # we also want to skip particles outside the rotated z bounds + z_min = bounds[4] + z_max = bounds[5] + + dx = (x_max - x_min) / xsize + dy = (y_max - y_min) / ysize + + idx = 1.0/dx + idy = 1.0/dy + + if kernel_name not in kernel_tables: + kernel_tables[kernel_name] = SPHKernelInterpolationTable(kernel_name) + cdef SPHKernelInterpolationTable itab = kernel_tables[kernel_name] + + with nogil, parallel(): + # loop through every particle + # NOTE: this loop can be quite time consuming. see description in + # pixelize_sph_kernel_projection. + + local_buff = malloc(sizeof(np.float64_t) * xsize * ysize) + xiterv = malloc(sizeof(np.float64_t) * 2) + yiterv = malloc(sizeof(np.float64_t) * 2) + ziterv = malloc(sizeof(np.float64_t) * 2) + xiter = malloc(sizeof(int) * 2) + yiter = malloc(sizeof(int) * 2) + ziter = malloc(sizeof(int) * 2) + xiter[0] = yiter[0] = ziter[0] = 0 + xiterv[0] = yiterv[0] = ziterv[0] = 0.0 + for i in range(xsize * ysize): + local_buff[i] = 0.0 + + for j in prange(0, posx.shape[0], schedule="dynamic"): + if j % 100000 == 0: + with gil: + PyErr_CheckSignals() + + xiter[1] = yiter[1] = ziter[1] = 999 + + if check_period == 1: + if posx[j] - hsml[j] < x_min: + xiter[1] = +1 + xiterv[1] = period_x + elif posx[j] + hsml[j] > x_max: + xiter[1] = -1 + xiterv[1] = -period_x + if posy[j] - hsml[j] < y_min: + yiter[1] = +1 + yiterv[1] = period_y + elif posy[j] + hsml[j] > y_max: + yiter[1] = -1 + yiterv[1] = -period_y + if posz[j] - hsml[j] < z_min: + ziter[1] = +1 + ziterv[1] = period_z + elif posz[j] + hsml[j] > z_max: + ziter[1] = -1 + ziterv[1] = -period_z + + # we set the smoothing length squared with lower limit of the pixel + h_j2 = fmax(hsml[j]*hsml[j], dx*dy) + ih_j2 = 1.0/h_j2 + + prefactor_j = pmass[j] / pdens[j] / hsml[j]**2 * quantity_to_smooth[j] + if weight_field is not None: + prefactor_j *= _weight_field[j] + + for ii in range(2): + if xiter[ii] == 999: continue + px = posx[j] + xiterv[ii] + if (px + hsml[j] < x_min) or (px - hsml[j] > x_max): continue + + + for jj in range(2): + if yiter[jj] == 999: continue + py = posy[j] + yiterv[jj] + if (py + hsml[j] < y_min) or (py - hsml[j] > y_max): continue + + for kk in range(2): + # discard if z is outside bounds + if ziter[kk] == 999: continue + pz = posz[j] + ziterv[kk] + if (pz + hsml[j] < z_min) or (pz - hsml[j] > z_max): continue + + # here we find the pixels which this particle contributes to + x0 = ((px - hsml[j] - x_min)*idx) + x1 = ((px + hsml[j] - x_min)*idx) + x0 = iclip(x0-1, 0, xsize) + x1 = iclip(x1+1, 0, xsize) + + y0 = ((py - hsml[j] - y_min)*idy) + y1 = ((py + hsml[j] - y_min)*idy) + y0 = iclip(y0-1, 0, ysize) + y1 = iclip(y1+1, 0, ysize) + + # found pixels we deposit on, loop through those pixels + for xi in range(x0, x1): + # we use the centre of the pixel to calculate contribution + x = (xi + 0.5) * dx + x_min + + posx_diff = px - x + posx_diff = posx_diff * posx_diff + + if posx_diff > h_j2: continue + + for yi in range(y0, y1): + y = (yi + 0.5) * dy + y_min + + posy_diff = py - y + posy_diff = posy_diff * posy_diff + if posy_diff > h_j2: continue + + q_ij2 = (posx_diff + posy_diff) * ih_j2 + if q_ij2 >= 1: + continue + + # see equation 32 of the SPLASH paper + # now we just use the kernel projection + local_buff[xi + yi*xsize] += prefactor_j * itab.interpolate(q_ij2) + mask[xi, yi] = 1 + + with gil: + for xxi in range(xsize): + for yyi in range(ysize): + buff[xxi, yyi] += local_buff[xxi + yyi*xsize] + free(local_buff) + free(xiterv) + free(yiterv) + free(ziterv) + free(xiter) + free(yiter) + free(ziter) + + return mask + @cython.boundscheck(False) @cython.wraparound(False) def interpolate_sph_positions_gather(np.float64_t[:] buff, @@ -1910,6 +2095,61 @@ def rotate_particle_coord(np.float64_t[:] px, return px_rotated, py_rotated, rot_bounds_x0, rot_bounds_x1, rot_bounds_y0, rot_bounds_y1 +@cython.boundscheck(False) +@cython.wraparound(False) +def rotate_particle_coord_with_z(np.float64_t[:] px, + np.float64_t[:] py, + np.float64_t[:] pz, + center, + width, + normal_vector, + north_vector): + # same as rotate_particle_coord but returns rotated z coords and bounds. + # This is a separate function so that we can avoid allocating the z-related + # arrays when they will not be used. + cdef int num_particles = np.size(px) + cdef np.float64_t[:] z_axis = np.array([0., 0., 1.], dtype="float64") + cdef np.float64_t[:] y_axis = np.array([0., 1., 0.], dtype="float64") + cdef np.float64_t[:, :] normal_rotation_matrix + cdef np.float64_t[:] transformed_north_vector + cdef np.float64_t[:, :] north_rotation_matrix + cdef np.float64_t[:, :] rotation_matrix + + normal_rotation_matrix = get_rotation_matrix(normal_vector, z_axis) + transformed_north_vector = np.matmul(normal_rotation_matrix, north_vector) + north_rotation_matrix = get_rotation_matrix(transformed_north_vector, y_axis) + rotation_matrix = np.matmul(north_rotation_matrix, normal_rotation_matrix) + + cdef np.float64_t[:] px_rotated = np.empty(num_particles, dtype="float64") + cdef np.float64_t[:] py_rotated = np.empty(num_particles, dtype="float64") + cdef np.float64_t[:] pz_rotated = np.empty(num_particles, dtype="float64") + cdef np.float64_t[:] coordinate_matrix = np.empty(3, dtype="float64") + cdef np.float64_t[:] rotated_coordinates + cdef np.float64_t[:] rotated_center + rotated_center = rotation_matmul( + rotation_matrix, np.array([center[0], center[1], center[2]])) + + # set up the rotated bounds + cdef np.float64_t rot_bounds_x0 = rotated_center[0] - width[0] / 2 + cdef np.float64_t rot_bounds_x1 = rotated_center[0] + width[0] / 2 + cdef np.float64_t rot_bounds_y0 = rotated_center[1] - width[1] / 2 + cdef np.float64_t rot_bounds_y1 = rotated_center[1] + width[1] / 2 + cdef np.float64_t rot_bounds_z0 = rotated_center[2] - width[2] / 2 + cdef np.float64_t rot_bounds_z1 = rotated_center[2] + width[2] / 2 + + for i in range(num_particles): + coordinate_matrix[0] = px[i] + coordinate_matrix[1] = py[i] + coordinate_matrix[2] = pz[i] + rotated_coordinates = rotation_matmul( + rotation_matrix, coordinate_matrix) + px_rotated[i] = rotated_coordinates[0] + py_rotated[i] = rotated_coordinates[1] + pz_rotated[i] = rotated_coordinates[2] + + return px_rotated, py_rotated, pz_rotated, rot_bounds_x0, rot_bounds_x1, rot_bounds_y0, rot_bounds_y1, rot_bounds_z0, rot_bounds_z1 + + @cython.boundscheck(False) @cython.wraparound(False) def off_axis_projection_SPH(np.float64_t[:] px, @@ -1926,28 +2166,53 @@ def off_axis_projection_SPH(np.float64_t[:] px, np.uint8_t[:, :] mask, normal_vector, north_vector, - weight_field=None): + weight_field=None, + depth_set=False): # Do nothing in event of a 0 normal vector if np.allclose(normal_vector, 0.): return - px_rotated, py_rotated, \ - rot_bounds_x0, rot_bounds_x1, \ - rot_bounds_y0, rot_bounds_y1 = rotate_particle_coord(px, py, pz, - center, width, normal_vector, north_vector) - - pixelize_sph_kernel_projection(projection_array, - mask, - px_rotated, - py_rotated, - smoothing_lengths, - particle_masses, - particle_densities, - quantity_to_smooth, - [rot_bounds_x0, rot_bounds_x1, - rot_bounds_y0, rot_bounds_y1], - weight_field=weight_field, - check_period=0) + if depth_set: + px_rotated, py_rotated, pz_rotated, \ + rot_bounds_x0, rot_bounds_x1, \ + rot_bounds_y0, rot_bounds_y1, \ + rot_bounds_z0, rot_bounds_z1 = rotate_particle_coord_with_z(px, py, pz, + center, width, normal_vector, north_vector) + + pixelize_sph_kernel_projection_with_depth_check(projection_array, + mask, + px_rotated, + py_rotated, + pz_rotated, + smoothing_lengths, + particle_masses, + particle_densities, + quantity_to_smooth, + [rot_bounds_x0, rot_bounds_x1, + rot_bounds_y0, rot_bounds_y1, + rot_bounds_z0, rot_bounds_z1], + weight_field=weight_field, + check_period=0) + else: + px_rotated, py_rotated, \ + rot_bounds_x0, rot_bounds_x1, \ + rot_bounds_y0, rot_bounds_y1 = rotate_particle_coord(px, py, pz, + center, width, normal_vector, north_vector) + + pixelize_sph_kernel_projection(projection_array, + mask, + px_rotated, + py_rotated, + smoothing_lengths, + particle_masses, + particle_densities, + quantity_to_smooth, + [rot_bounds_x0, rot_bounds_x1, + rot_bounds_y0, rot_bounds_y1], + weight_field=weight_field, + check_period=0) + + @cython.boundscheck(False) diff --git a/yt/visualization/fixed_resolution.py b/yt/visualization/fixed_resolution.py index 7e1123cdf6f..dfd49de0faa 100644 --- a/yt/visualization/fixed_resolution.py +++ b/yt/visualization/fixed_resolution.py @@ -649,6 +649,7 @@ def _generate_image_and_mask(self, item) -> None: interpolated=dd.interpolated, north_vector=dd.north_vector, method=dd.method, + depth_set=dd.depth_set, ) if self.data_source.moment == 2: diff --git a/yt/visualization/plot_window.py b/yt/visualization/plot_window.py index 412cc28f8ad..8f3ecf10bf2 100644 --- a/yt/visualization/plot_window.py +++ b/yt/visualization/plot_window.py @@ -2275,6 +2275,7 @@ def __init__( data_source=None, *, moment=1, + depth_set=False, ): validate_moment(moment, weight) self.center = center @@ -2301,6 +2302,11 @@ def __init__( self.orienter = Orientation(normal_vector, north_vector=north_vector) self.moment = moment + # note: the depth_set bool indicates whether we need to limit particles + # by the rotated z-bounds. If using all particles in the rotated x-y + # plane (depth_set is False), then we can avoid some computation. + self.depth_set = depth_set + def _determine_fields(self, *args): return self.dd._determine_fields(*args) @@ -2463,6 +2469,7 @@ def __init__( method=method, data_source=data_source, moment=moment, + depth_set=depth is not None, ) validate_mesh_fields(OffAxisProj, fields) diff --git a/yt/visualization/volume_rendering/off_axis_projection.py b/yt/visualization/volume_rendering/off_axis_projection.py index 381cde35031..6ae5953cdbb 100644 --- a/yt/visualization/volume_rendering/off_axis_projection.py +++ b/yt/visualization/volume_rendering/off_axis_projection.py @@ -29,6 +29,8 @@ def off_axis_projection( north_vector=None, num_threads=1, method="integrate", + *, + depth_set=False, ): r"""Project through a dataset, off-axis, and return the image plane. @@ -94,6 +96,10 @@ def off_axis_projection( This should only be used for uniform resolution grid datasets, as other datasets may result in unphysical images. or camera movements. + depth_set : bool + If True, will check the rotated z bounds and limit particles included. + If False (the default), will use all particles in the rotated x-y plane. + Returns ------- image : array @@ -210,6 +216,7 @@ def off_axis_projection( ounits = finfo.output_units bounds = [x_min, x_max, y_min, y_max, z_min, z_max] + print(f"depth_set is {depth_set}") if weight is None: for chunk in data_source.chunks([], "io"): off_axis_projection_SPH( @@ -227,6 +234,7 @@ def off_axis_projection( mask, normal_vector, north, + depth_set=depth_set, ) # Assure that the path length unit is in the default length units @@ -266,6 +274,7 @@ def off_axis_projection( normal_vector, north, weight_field=chunk[weight].in_units(wounits), + depth_set=depth_set, ) for chunk in data_source.chunks([], "io"): @@ -284,6 +293,7 @@ def off_axis_projection( mask, normal_vector, north, + depth_set=depth_set, ) normalization_2d_utility(buf, weight_buff) From 118182fb4007ec79fd2dabe96eac5b06f5e792b8 Mon Sep 17 00:00:00 2001 From: Chris Havlin Date: Wed, 25 Oct 2023 13:30:45 -0500 Subject: [PATCH 2/4] rm errant print --- yt/visualization/volume_rendering/off_axis_projection.py | 1 - 1 file changed, 1 deletion(-) diff --git a/yt/visualization/volume_rendering/off_axis_projection.py b/yt/visualization/volume_rendering/off_axis_projection.py index 6ae5953cdbb..90cf2178e59 100644 --- a/yt/visualization/volume_rendering/off_axis_projection.py +++ b/yt/visualization/volume_rendering/off_axis_projection.py @@ -216,7 +216,6 @@ def off_axis_projection( ounits = finfo.output_units bounds = [x_min, x_max, y_min, y_max, z_min, z_max] - print(f"depth_set is {depth_set}") if weight is None: for chunk in data_source.chunks([], "io"): off_axis_projection_SPH( From 4dfe236ebee9414edeadc8af0a10443f6ad99cb2 Mon Sep 17 00:00:00 2001 From: chrishavlin Date: Tue, 31 Oct 2023 15:35:43 -0500 Subject: [PATCH 3/4] add a test --- .../tests/test_offaxisprojection.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/yt/visualization/tests/test_offaxisprojection.py b/yt/visualization/tests/test_offaxisprojection.py index 20950393ec5..41d0f3dd17f 100644 --- a/yt/visualization/tests/test_offaxisprojection.py +++ b/yt/visualization/tests/test_offaxisprojection.py @@ -12,6 +12,7 @@ assert_rel_equal, fake_octree_ds, fake_random_ds, + fake_sph_grid_ds, ) from yt.visualization.api import OffAxisProjectionPlot, OffAxisSlicePlot from yt.visualization.image_writer import write_projection @@ -246,3 +247,22 @@ def _vlos_sq(field, data): p2.frb["gas", "velocity_los"], 10, ) + + +def test_off_axis_sph_depth(): + ds = fake_sph_grid_ds() + p1 = OffAxisProjectionPlot( + ds, + [1, 1, 1], + ("io", "density"), + buff_size=(50, 50), + ) + p2 = OffAxisProjectionPlot( + ds, + [1, 1, 1], + ("io", "density"), + buff_size=(50, 50), + depth=(0.1, "cm"), + ) + # checks that fewer particles are picked up + assert p2.frb["io", "density"].sum() < p1.frb["io", "density"].sum() From 26a2cc6b645ef97d1383c83b393fedab810a6852 Mon Sep 17 00:00:00 2001 From: chavlin Date: Thu, 2 Nov 2023 11:08:09 -0500 Subject: [PATCH 4/4] fix the no-depth case --- yt/visualization/plot_window.py | 8 +++++++- yt/visualization/tests/test_offaxisprojection.py | 1 + 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/yt/visualization/plot_window.py b/yt/visualization/plot_window.py index 8f3ecf10bf2..ca30e530798 100644 --- a/yt/visualization/plot_window.py +++ b/yt/visualization/plot_window.py @@ -2446,6 +2446,12 @@ def __init__( f"currently supported geometries: {self._supported_geometries!r}" ) + depth_set = depth is not None + if not depth_set: + # need a valid depth for calculating bounds, but those + # bounds will not be used to limit particles. + depth = (1, "1") + (bounds, center_rot) = get_oblique_window_parameters( normal, center, width, ds, depth=depth ) @@ -2469,7 +2475,7 @@ def __init__( method=method, data_source=data_source, moment=moment, - depth_set=depth is not None, + depth_set=depth_set, ) validate_mesh_fields(OffAxisProj, fields) diff --git a/yt/visualization/tests/test_offaxisprojection.py b/yt/visualization/tests/test_offaxisprojection.py index 41d0f3dd17f..2d286bc531a 100644 --- a/yt/visualization/tests/test_offaxisprojection.py +++ b/yt/visualization/tests/test_offaxisprojection.py @@ -256,6 +256,7 @@ def test_off_axis_sph_depth(): [1, 1, 1], ("io", "density"), buff_size=(50, 50), + depth=None, ) p2 = OffAxisProjectionPlot( ds,