Use stamping approach to deposit off-axis cells

[cphyc]

PR Summary

On main, off-axis projections are very slow due to the need to cast rays through the domain. This is particularly the case for octree-datasets to to the hot-loop for grid, mask in data_source.blocks:. Indeed, this iterates over all octs in the domain, one-by-one.

While one solution might be to move this chunk to a Cython-accelerated code, it will remain somewhat slow no matter what.

In this PR, I take a completely different approach. In the first part of the code, I compute the projected depth of a cube as seen from the camera position (i.e. how long is the segment traversing the cube at different pixel locations). One way of seeing this is that we now have a precomputed kernel, which we can apply to each cell at a scale proportional to the cell size.

This yields good performances: ~2s on an i9 for an 800x800 simulation with 10,000,000 cells. Interestingly, this could also be used to offload the rendering to the GPU.

While the code is relatively stable, I still get slight differences between axis-aligned projections and this method, which suggests that there are some minor details that I got wrong (see below).

Comparisons

Using the code

import yt

ds = yt.load_sample("output_00080")
p = yt.ProjectionPlot(
    ds,
    "y",
    "density",
    weight_field="density"
)
p.set_zlim(("gas", "density"), 6e-31, 2e-26)
p.save("/tmp/projection.png")
p = yt.OffAxisProjectionPlot(
    ds,
    [0, 1, 0],
    "density",
    north_vector=[1, 0, 0],
    width=(1, "unitary"),
    center=[0.5, 0.5, 0.5],
    buff_size=(800, 800),
    weight_field="density",
)
p.set_zlim(("gas", "density"), 6e-31, 2e-26)
p.set_axes_unit("Mpc")
p.save("/tmp/new.png")

I obtain (as of #f8d38de) very similar performances as the on-axis projection (!!!)

	Result from ProjectionPlot	Result from new off-axis projection
Image
Timing (s)	1.45	2.14 (8 cores), 3.13 (1 core)

[chrishavlin]

No insight into why you're getting different results between the on-axis projection and the off-axis when aligned with the axis, but here's an example that results in very different behavior. The standard on axis projection returns what I would expect:

fld = ("ramses", "z")
wt = ("index", "ones")
p = yt.ProjectionPlot(
    ds, 
    "y",
    fld,
    weight_field=wt,
    buff_size=(800,800),
)
p.set_axes_unit("Mpc")
p.set_unit(fld, "Mpc")
p.save('old.png')

But the new off axis produces a constant value of 1?

p = yt.OffAxisProjectionPlot(
    ds, 
    [0, 1, 0], 
    fld,
    north_vector=[1, 0, 0], 
    width=(1, "unitary"),
    center=[0.5, 0.5, 0.5],
    buff_size=(800, 800),
    weight_field=wt,
)
p.set_axes_unit("Mpc")
p.set_unit(fld, "Mpc")

mx = p.plots[fld].image.get_array().max()
mn = p.plots[fld].image.get_array().min()
print((mn, mx))
# (1.0, 1.0)
p.save('new.png')

Been reading and stepping through the code, I kinda think it's related to how the weighting field is handled? maybe units related?

[cphyc]

I have finally updated the code. Here are some examples

Example 1: One of the galaxy in output_00080

import yt

ds = yt.load("output_00080")
ad = ds.all_data()
center = ds.arr(ad.argmax("density"))
width = ds.quan(200, "kpc")
zlim = (6e-30, 2e-24)
data_source = ds.sphere(center, width * 2)

t0 = time()
p = yt.ProjectionPlot(
    ds,
    "x",
    "density",
    buff_size=(800, 800),
    weight_field="density",
    center=center,
    width=width,
    data_source=data_source
)
pref = p
p.set_zlim(("gas", "density"), pref.frb["gas", "density"].min(), pref.frb["gas", "density"].max())
p.save("/tmp/projection.png")

p = yt.OffAxisProjectionPlot(
    ds,
    [1, 0, 0],
    "density",
    north_vector=[0, 0, 1],
    buff_size=(800, 800),
    weight_field="density",
    center=center,
    width=width,
    data_source=data_source
)
p.set_zlim(("gas", "density"), pref.frb["gas", "density"].min(), pref.frb["gas", "density"].max())
p.set_axes_unit("Mpc")
p.save("/tmp/new.png")

On axis projection	Off axis projection	Off axis projection (from [1, 2, 3])
0.19s	1.3s (on 8 cores), 5.8s (on 1 core)	2.25s (on 8 cores), 12s (on 1 core)

Example 2: plot of the z component

This is @chrishavlin example from https://github.com/yt-project/yt/pull/4741#issuecomment-1813024283.

On axis projection	Off axis projection

Performances

This new implementation is particularly efficient when we have lots of cells to project. For example, to do an off-axis projection of the entirety of the dataset output_00080 I obtain:

Method	Version	Time (s)	OMP_NUM_THREADS	Speedup w.r.t old
Axis projection		0.5	-	-
Off-axis projection along x	new	2.7	1	4.5
Off-axis projection along [1, 2, 3]		3.8	1	3.8
Off-axis projection along x		1.1	8	11
Off-axis projection along [1, 2, 3]		0.89	8	13
Off-axis projection along x		1.1	16	11
Off-axis projection along [1, 2, 3]		0.9	16	16
Off-axis projection along x	old	12.1	-	1
Off-axis projection along [1, 2, 3]		14.5	-	1

[neutrinoceros]

I would have liked to review this but unfortunately I'm approaching review request bankruptcy so realistically it's not going to happen. I think it would be fine to merge it with Chris' approval. For release notes though, should this be in new features, perf improvements, minor API change, all three ?

[cphyc]

I added labels (improvement and new feature). The API doesn't change, it's a drop-in replacement for the existing off-axis renderer.

[neutrinoceros]

Let me rephrase that: is there any observable change besides improved performance ?

[cphyc]

Yes; the method does not yield exactly the same images as the previous implementation (although they should be extremely close).

[chrishavlin]

I didn't mention it in my review, but while reviewing i did actually run some more comparison tests between old and new methods with a couple datasets and some different cutting planes. the differences were on the order of 1e-40 relative difference in all the cases I looked at. it wasn't exhaustive testing by any means, but I think it's pretty safe to say that it's unlikely for folks to observe any difference by eye but may see differences when comparing actual pixel values. Certainly worth giving a heads up in the release notes though!

[neutrinoceros]

thank you both ! Feel free to merge away then !