ENH: add support for VelocityCallback and MagFieldCallBack in spherical coordinates

[neutrinoceros]

PR Summary

minimal demo

import numpy as np
import yt
from unyt import cm, s

shape = (32, 128, 1)
a = np.sin(np.linspace(-np.pi/2, np.pi/2, 128))
b = np.ones((32, 128))
c = np.zeros((32, 128))

data = {
    ("gas", "velocity_r"): np.reshape(a * b, shape) * cm/s,
    ("gas", "velocity_theta"): np.reshape(a * c, shape) * cm/s,
}

ds = yt.load_uniform_grid(
    data,
    shape,
    bbox=np.array([[0.0, 5.0], [np.pi / 2 - 0.5, np.pi / 2 + 0.5], [-0.1, +0.1]]),
    geometry="spherical",
)

p = yt.SlicePlot(ds, "phi", "velocity_r")
p.annotate_velocity()
p.save("/tmp/")

image actualized on 2022-08-01

older version

note that the quiver plot is drawn well outside of the actual data domain. I think this effect is due to the current implementation of the pixelizer routine (pixelize_cylinder), combined with the "low resolution" buffer associated with the quiver plot for readability. I consider this a bug, but it's out of scope here.
Edit: this is solved in #3818

[neutrinoceros]

While I'm at it, I'll add support for "conic" projections too (theta-normal), illustrated here with a purely azimuthal magnetic field

import numpy as np
import yt
from yt.units import G

shape = (32, 128, 128)
a = np.sin(np.linspace(-np.pi / 2, np.pi / 2, 128))
b = np.ones((32, 128, 128))
c = np.zeros((32, 128, 128))

data = {
    ("gas", "magnetic_field_r"): np.reshape(a * c, shape) * G,
    ("gas", "magnetic_field_theta"): np.reshape(a * c, shape) * G,
    ("gas", "magnetic_field_phi"): np.reshape(a * b, shape) * G,
}

ds = yt.load_uniform_grid(
    data,
    shape,
    bbox=np.array([[1, 5], [np.pi / 2 - 0.5, np.pi / 2 + 0.5], [0, +np.pi]]),
    geometry="spherical",
)

p = yt.SlicePlot(ds, "theta", "magnetic_field_phi")
p.annotate_magnetic_field()
p.save("/tmp/")

image actualized on 2022-08-01

older version

[neutrinoceros]

Tests exploding everywhere, switching to draft for now.

[matthewturk]

I think your diagnosis -- the buffer being low-resolution -- is the root cause, but perhaps it can be fixed by more carefully adjusting the arguments to the pixelization routine to query at the specific pixel centers.

[matthewturk]

I did not realize it depended on yt.mods. Oof! Also, I could be wrong about the mag field.

[neutrinoceros]

I think your diagnosis -- the buffer being low-resolution -- is the root cause, but perhaps it can be fixed by more carefully adjusting the arguments to the pixelization routine to query at the specific pixel centers.

How about the proposed fix in #3818 ?

[matthewturk]

LGTM

[neutrinoceros]

I need to revisit this. Switching to draft in the mean time. For now I'll just resolve a conflict and add a comment as requested by Matt.

[neutrinoceros]

refreshing CI

[neutrinoceros]

rebased to resolve conflicts and refresh CI

[neutrinoceros]

this is now rebased on top of #4021. I also cleaned up a couple redundant definitions as previously discussed, and fixed a couple outstanding bugs.

[neutrinoceros]

@yt-fido test this please

[neutrinoceros]

@matthewturk can we merge this ?