ENH: use Lagrangian tesselation for smoothing N-body plotting

yt/utilities/lib/image_utilities.pyx

@@ -7,28 +7,132 @@ Utilities for images
 
 import numpy as np
 
+cimport cython
 cimport numpy as np
+from cython.parallel cimport prange
 
 from yt.utilities.lib.fp_utils cimport iclip
 
 
+@cython.boundscheck(False)
+cdef inline void _add_point_to_greyscale_image(
+ np.float64_t[:, :] buffer,
+ np.uint8_t[:, :] buffer_mask,
+ np.float64_t px,
+ np.float64_t py,
+ np.float64_t pv,
+ int xs,
+ int ys
+) nogil:
+ cdef int i, j
+ j = <int> (xs * px)
+ i = <int> (ys * py)
+ buffer[i, j] += pv
+ buffer_mask[i, j] = 1
+
+
 def add_points_to_greyscale_image(
  np.ndarray[np.float64_t, ndim=2] buffer,
  np.ndarray[np.uint8_t, ndim=2] buffer_mask,
  np.ndarray[np.float64_t, ndim=1] px,
  np.ndarray[np.float64_t, ndim=1] py,
  np.ndarray[np.float64_t, ndim=1] pv):
- cdef int i, j, pi
+ cdef int pi
  cdef int npx = px.shape[0]
  cdef int xs = buffer.shape[0]
  cdef int ys = buffer.shape[1]
+
+ cdef np.float64_t[:, :] buffer_view = buffer
+ cdef np.uint8_t[:, :] buffer_mask_view = buffer_mask
+
  for pi in range(npx):
- j = <int> (xs * px[pi])
- i = <int> (ys * py[pi])
- buffer[i, j] += pv[pi]
- buffer_mask[i, j] = 1
+ _add_point_to_greyscale_image(buffer_view, buffer_mask_view, px[pi], py[pi], pv[pi], xs, ys)
  return
 
+
+@cython.boundscheck(False)
+def add_points_to_greyscale_image_with_lagrangian_tesselation(
+ np.ndarray[np.float64_t, ndim=2] buffer,
+ np.ndarray[np.uint8_t, ndim=2] buffer_mask,
+ np.ndarray[np.float64_t, ndim=4] p3d,
+ np.ndarray[np.float64_t, ndim=3] pv,
+ int split=100,
+):
+ cdef int xs = buffer.shape[0]
+ cdef int ys = buffer.shape[1]
+ cdef int Ngrid = p3d.shape[0] - 1
+
+ cdef np.float64_t[:, :] buffer_view = buffer
+ cdef np.uint8_t[:, :] buffer_mask_view = buffer_mask
+
+ cdef np.ndarray[np.uint8_t, ndim=2] vert = np.array(
+ (
+ (0, 0, 0),
+ (1, 0, 0),
+ (1, 1, 0),
+ (0, 1, 0),
+ (0, 0, 1),
+ (1, 0, 1),
+ (1, 1, 1),
+ (0, 1, 1),
+ ),
+ dtype=np.uint8
+ )
+ cdef np.ndarray[np.uint8_t, ndim=2] conn = np.array(
+ (
+ (4, 0, 7, 1),
+ (1, 0, 3, 7),
+ (5, 1, 4, 7),
+ (2, 3, 7, 1),
+ (1, 5, 6, 7),
+ (2, 6, 1, 7),
+ ),
+ dtype=np.uint8
+ )
+
+ cdef np.float64_t[3] orig
+ cdef np.float64_t[3] a
+ cdef np.float64_t[3] b
+ cdef np.float64_t[3] c
+
+ cdef np.float64_t[3] newp
+
+ cdef int i, j, k, m, s, idim
+
+ cdef np.float64_t f = 1.0 / split
+
+ cdef np.uint8_t[:, :] off
+ cdef np.float64_t[:, :] ro
+
+ for m in range(len(conn)):
+ ro = np.random.random(size=(split, 3)) # random offsets
+ off = vert[conn[m]]
+
+ for i in range(Ngrid):
+ for j in range(Ngrid):
+ for k in range(Ngrid):
+ for idim in range(2): # stop at 2, because z is not used
+ orig[idim] = p3d[i + off[3][0], j + off[3][1], k + off[3][2], idim]
+ a[idim] = p3d[i + off[0][0], j + off[0][1], k + off[0][2], idim] - orig[idim]
+ b[idim] = p3d[i + off[1][0], j + off[1][1], k + off[1][2], idim] - orig[idim]
+ c[idim] = p3d[i + off[2][0], j + off[2][1], k + off[2][2], idim] - orig[idim]
+
+ for s in prange(split, nogil=True):
+ for idim in range(2): # stop at 2, because z is not used
+ newp[idim] = orig[idim] + (
+ ro[s, 0] * a[idim] +
+ ro[s, 1] * b[idim] +
+ ro[s, 2] * c[idim]
+ )
+ while newp[idim] < 0.0:
+ newp[idim] += 1.0
+ while newp[idim] >= 1.0:
+ newp[idim] -= 1.0
+
+ _add_point_to_greyscale_image(
+ buffer_view, buffer_mask_view, newp[0], newp[1], f, xs, ys
+ )
+
 def add_points_to_image(
  np.ndarray[np.uint8_t, ndim=3] buffer,
  np.ndarray[np.float64_t, ndim=1] px,

yt/visualization/fixed_resolution.py

@@ -13,6 +13,7 @@
 from yt.utilities.lib.api import ( # type: ignore
  CICDeposit_2,
  add_points_to_greyscale_image,
+ add_points_to_greyscale_image_with_lagrangian_tesselation,
 )
 from yt.utilities.lib.pixelization_routines import pixelize_cylinder
 from yt.utilities.math_utils import compute_stddev_image
@@ -661,6 +662,58 @@ def _sq_field(field, data, item: FieldKey):
  return ia
 
 
+def _get_tet_interpolated_particles(Ngrid, p3d, split=1):
+ """A fast function to create particles inside tetrahedra
+ with the same linear map that made the tet."""
+ vert = np.array(
+ (
+ (0, 0, 0),
+ (1, 0, 0),
+ (1, 1, 0),
+ (0, 0, 1),
+ (0, 0, 1),
+ (1, 0, 1),
+ (1, 1, 1),
+ (0, 1, 1),
+ )
+ )
+ conn = np.array(
+ (
+ (4, 0, 7, 1),
+ (1, 0, 7, 3),
+ (5, 1, 7, 4),
+ (2, 7, 1, 3),
+ (1, 5, 7, 6),
+ (2, 6, 1, 7),
+ )
+ )
+ Ntetpp = len(conn)
+ Np = Ngrid * Ngrid * Ngrid
+ newp = np.zeros((Ntetpp, split, Np, 3))
+ ro = np.random.random(size=(Ntetpp, split, 3)) # random offsets
+ for m in range(Ntetpp): # 6 tets
+ off = vert[conn[m]]
+ sx, sy, sz = (
+ [slice(off[i][j], Ngrid + off[i][j]) for i in range(4)] for j in range(3)
+ )
+
+ orig = p3d[sx[3], sy[3], sz[3], :]
+
+ b = (p3d[sx[1], sy[1], sz[1], :] - orig).reshape((Np, 3))
+ c = (p3d[sx[2], sy[2], sz[2], :] - orig).reshape((Np, 3))
+ a = (p3d[sx[0], sy[0], sz[0], :] - orig).reshape((Np, 3))
+
+ oo = orig.reshape(Np, 3)[None, :, :]
+
+ for i in range(3):
+ newp[m, :, :, i] = oo[:, :, i] + (
+ ro[m, :, None, 0] * a[None, :, i]
+ + ro[m, :, None, 1] * b[None, :, i]
+ + ro[m, :, None, 2] * c[None, :, i]
+ )
+ return newp
+
+
 class ParticleImageBuffer(FixedResolutionBuffer):
  """
 
@@ -689,9 +742,11 @@ def __init__(
  axis = self.axis
  self.xax = self.ds.coordinates.x_axis[axis]
  self.yax = self.ds.coordinates.y_axis[axis]
+ self.zax = list({0, 1, 2} - {self.xax, self.yax})[0]
  ax_field_template = "particle_position_%s"
  self.x_field = ax_field_template % self.ds.coordinates.axis_name[self.xax]
  self.y_field = ax_field_template % self.ds.coordinates.axis_name[self.yax]
+ self.z_field = ax_field_template % self.ds.coordinates.axis_name[self.zax]
 
  def __getitem__(self, item):
  if item in self.data:
@@ -763,6 +818,20 @@ def __getitem__(self, item):
  x_bin_edges,
  y_bin_edges,
  )
+ elif deposition == "lagrangian_tesselation":
+ z_data = self.data_source.dd[ftype, self.z_field]
+ dz = z_data.in_units("code_length").d
+ # TODO: handle periodicity
+ pz = dz / (bounds[1] - bounds[0])
+ Nsplit = 100
+ order = np.argsort(self.data_source.dd[ftype, "particle_index"])
+ p3d = np.stack([px[order], py[order], pz[order]], axis=1).reshape(
+ 64, 64, 64, 3
+ )
+
+ add_points_to_greyscale_image_with_lagrangian_tesselation(
+ buff, buff_mask, p3d, splat_vals[None, None, :], Nsplit
+ )
  else:
  raise ValueError(f"Received unknown deposition method '{deposition}'")
 
@@ -785,7 +854,7 @@ def __getitem__(self, item):
  ia = ImageArray(buff, units=units, info=info)
 
  # divide by the weight_field, if needed
- if weight_field is not None:
+ if weight_field is not None and deposition != "lagrangian_tesselation":
  weight_buff = np.zeros(self.buff_size)
  weight_buff_mask = np.zeros(self.buff_size, dtype="uint8")
  if deposition == "ngp":