ENH: add support for VelocityCallback and MagFieldCallBack in spherical coordinates

yt/fields/vector_operations.py

@@ -195,6 +195,8 @@ def create_vector_fields(registry, basename, field_units, ftype="gas", slice_inf
         validators=[ValidateParameter(f"bulk_{basename}")],
     )
 
+    axis_names = registry.ds.coordinates.axis_order
+
     if not is_curvilinear(registry.ds.geometry):
 
         # The following fields are invalid for curvilinear geometries
@@ -553,7 +555,7 @@ def _cartesian_z(field, data):
                     data[(ftype, "theta")]
                 ) - data[(ftype, f"{basename}_theta")] * np.sin(data[(ftype, "theta")])
 
-        if registry.ds.dimensionality == 3:
+        if registry.ds.dimensionality >= 2:
             registry.add_field(
                 (ftype, f"{basename}_cartesian_z"),
                 sampling_type="local",
@@ -562,6 +564,66 @@ def _cartesian_z(field, data):
                 display_field=True,
             )
 
+    if registry.ds.geometry == "spherical":
+
+        def _cylindrical_radius_component(field, data):
+            return (
+                np.sin(data[(ftype, "theta")]) * data[(ftype, f"{basename}_r")]
+                + np.cos(data[(ftype, "theta")]) * data[(ftype, f"{basename}_theta")]
+            )
+
+        registry.add_field(
+            (ftype, f"{basename}_cylindrical_radius"),
+            sampling_type="local",
+            function=_cylindrical_radius_component,
+            units=field_units,
+            display_field=True,
+        )
+
+        registry.alias(
+            (ftype, f"{basename}_cylindrical_z"),
+            (ftype, f"{basename}_cartesian_z"),
+        )
+
+        # define vector components appropriate for 'theta'-normal plots.
+        # The projection plane is called 'conic plane' in the code base as well as docs.
+        # Contrary to 'poloidal' and 'toroidal', this isn't a widely spread
+        # naming convention, but here it is exposed to users as part of dedicated
+        # field names, so it needs to be stable.
+        def _conic_x(field, data):
+            rax = axis_names.index("r")
+            pax = axis_names.index("phi")
+            bc = data.get_field_parameter(f"bulk_{basename}")
+            return np.cos(data[ftype, "phi"]) * (
+                data[ftype, f"{basename}_r"] - bc[rax]
+            ) - np.sin(data[ftype, "phi"]) * (data[ftype, f"{basename}_phi"] - bc[pax])
+
+        def _conic_y(field, data):
+            rax = axis_names.index("r")
+            pax = axis_names.index("phi")
+            bc = data.get_field_parameter(f"bulk_{basename}")
+            return np.sin(data[(ftype, "phi")]) * (
+                data[(ftype, f"{basename}_r")] - bc[rax]
+            ) + np.cos(data[(ftype, "phi")]) * (
+                data[(ftype, f"{basename}_phi")] - bc[pax]
+            )
+
+        if registry.ds.dimensionality == 3:
+            registry.add_field(
+                (ftype, f"{basename}_conic_x"),
+                sampling_type="local",
+                function=_conic_x,
+                units=field_units,
+                display_field=True,
+            )
+            registry.add_field(
+                (ftype, f"{basename}_conic_y"),
+                sampling_type="local",
+                function=_conic_y,
+                units=field_units,
+                display_field=True,
+            )
+
 
 def create_averaged_field(
     registry,

yt/frontends/stream/fields.py

@@ -19,6 +19,12 @@ class StreamFieldInfo(FieldInfoContainer):
         ("magnetic_field_x", ("gauss", [], None)),
         ("magnetic_field_y", ("gauss", [], None)),
         ("magnetic_field_z", ("gauss", [], None)),
+        ("velocity_r", ("code_length/code_time", ["velocity_r"], None)),
+        ("velocity_theta", ("code_length/code_time", ["velocity_theta"], None)),
+        ("velocity_phi", ("code_length/code_time", ["velocity_phi"], None)),
+        ("magnetic_field_r", ("gauss", [], None)),
+        ("magnetic_field_theta", ("gauss", [], None)),
+        ("magnetic_field_phi", ("gauss", [], None)),
         (
             "radiation_acceleration_x",
             ("code_length/code_time**2", ["radiation_acceleration_x"], None),
@@ -89,7 +95,9 @@ def setup_fluid_fields(self):
             if units != "":
                 self.add_output_field(field, sampling_type="cell", units=units)
         setup_magnetic_field_aliases(
-            self, "stream", [f"magnetic_field_{ax}" for ax in "xyz"]
+            self,
+            "stream",
+            [f"magnetic_field_{ax}" for ax in self.ds.coordinates.axis_order],
         )
 
     def add_output_field(self, name, sampling_type, **kwargs):

yt/visualization/plot_modifications.py

@@ -406,7 +406,13 @@ class VelocityCallback(PlotCallback):
     """
 
     _type_name = "velocity"
-    _supported_geometries = ("cartesian", "spectral_cube", "polar", "cylindrical")
+    _supported_geometries = (
+        "cartesian",
+        "spectral_cube",
+        "polar",
+        "cylindrical",
+        "spherical",
+    )
     _incompatible_plot_types = ("OffAxisProjection", "Particle")
 
     def __init__(
@@ -475,6 +481,17 @@ def __call__(self, plot):
                 # should convert r-theta plane to x-y plane
                 xv = (ftype, "velocity_cartesian_x")
                 yv = (ftype, "velocity_cartesian_y")
+            elif plot.data.ds.geometry == "spherical":
+                if axis_names[plot.data.axis] == "phi":
+                    xv = (ftype, "velocity_cylindrical_radius")
+                    yv = (ftype, "velocity_cylindrical_z")
+                elif axis_names[plot.data.axis] == "theta":
+                    xv = (ftype, "velocity_conic_x")
+                    yv = (ftype, "velocity_conic_y")
+                else:
+                    raise NotImplementedError(
+                        f"annotate_velocity is missing support for normal={axis_names[plot.data.axis]!r}"
+                    )
             else:
                 # for other cases (even for cylindrical geometry),
                 # orthogonal planes are generically Cartesian
@@ -511,7 +528,13 @@ class MagFieldCallback(PlotCallback):
     """
 
     _type_name = "magnetic_field"
-    _supported_geometries = ("cartesian", "spectral_cube", "polar", "cylindrical")
+    _supported_geometries = (
+        "cartesian",
+        "spectral_cube",
+        "polar",
+        "cylindrical",
+        "spherical",
+    )
     _incompatible_plot_types = ("OffAxisProjection", "Particle")
 
     def __init__(
@@ -571,6 +594,17 @@ def __call__(self, plot):
                 # should convert r-theta plane to x-y plane
                 xv = (ftype, "magnetic_field_cartesian_x")
                 yv = (ftype, "magnetic_field_cartesian_y")
+            elif plot.data.ds.geometry == "spherical":
+                if axis_names[plot.data.axis] == "phi":
+                    xv = (ftype, "magnetic_field_cylindrical_radius")
+                    yv = (ftype, "magnetic_field_cylindrical_z")
+                elif axis_names[plot.data.axis] == "theta":
+                    xv = (ftype, "magnetic_field_conic_x")
+                    yv = (ftype, "magnetic_field_conic_y")
+                else:
+                    raise NotImplementedError(
+                        f"annotate_magnetic_field is missing support for normal={axis_names[plot.data.axis]!r}"
+                    )
             else:
                 # for other cases (even for cylindrical geometry),
                 # orthogonal planes are generically Cartesian
@@ -694,6 +728,7 @@ class QuiverCallback(BaseQuiverCallback):
         "spectral_cube",
         "polar",
         "cylindrical",
+        "spherical",
     )
 
     def __init__(

yt/visualization/plot_window.py

@@ -1433,8 +1433,8 @@ def run_callbacks(self):
                 )
                 try:
                     callback(cbw)
-                except YTDataTypeUnsupported as e:
-                    raise e
+                except (NotImplementedError, YTDataTypeUnsupported):
+                    raise
                 except Exception as e:
                     raise YTPlotCallbackError(callback._type_name) from e
             for key in self.frb.keys():

yt/visualization/tests/test_callbacks.py

@@ -485,16 +485,23 @@ def test_velocity_callback():
             slc.annotate_velocity()
             assert_fname(slc.save(prefix)[0])
 
+
+def test_velocity_callback_spherical():
+    ds = fake_amr_ds(
+        fields=("density", "velocity_r", "velocity_theta", "velocity_phi"),
+        units=("g/cm**3", "cm/s", "cm/s", "cm/s"),
+        geometry="spherical",
+    )
+
     with _cleanup_fname() as prefix:
-        ds = fake_amr_ds(
-            fields=("density", "velocity_r", "velocity_theta", "velocity_phi"),
-            units=("g/cm**3", "cm/s", "cm/s", "cm/s"),
-            geometry="spherical",
-        )
-        p = ProjectionPlot(ds, "r", ("gas", "density"))
-        assert_raises(
-            YTDataTypeUnsupported, p.annotate_velocity, factor=40, normalize=True
-        )
+        p = ProjectionPlot(ds, "phi", ("stream", "density"))
+        p.annotate_velocity(factor=40, normalize=True)
+        assert_fname(p.save(prefix)[0])
+
+    with _cleanup_fname() as prefix:
+        p = ProjectionPlot(ds, "r", ("stream", "density"))
+        p.annotate_velocity(factor=40, normalize=True)
+        assert_raises(NotImplementedError, p.save, prefix)
 
 
 @requires_file(cyl_2d)
@@ -568,9 +575,23 @@ def test_magnetic_callback():
             ),
             geometry="spherical",
         )
+        p = ProjectionPlot(ds, "phi", ("gas", "density"))
+        p.annotate_magnetic_field(
+            factor=8, scale=0.5, scale_units="inches", normalize=True
+        )
+        assert_fname(slc.save(prefix)[0])
+
+        p = ProjectionPlot(ds, "theta", ("gas", "density"))
+        p.annotate_magnetic_field(
+            factor=8, scale=0.5, scale_units="inches", normalize=True
+        )
+        assert_fname(slc.save(prefix)[0])
+
         p = ProjectionPlot(ds, "r", ("gas", "density"))
-        kwargs = dict(factor=8, scale=0.5, scale_units="inches", normalize=True)
-        assert_raises(YTDataTypeUnsupported, p.annotate_magnetic_field, **kwargs)
+        p.annotate_magnetic_field(
+            factor=8, scale=0.5, scale_units="inches", normalize=True
+        )
+        assert_raises(NotImplementedError, p.save, prefix)
 
 
 @requires_file(cyl_2d)
@@ -628,26 +649,37 @@ def test_quiver_callback():
         slc.annotate_quiver(("gas", "velocity_r"), ("gas", "velocity_z"))
         assert_fname(slc.save(prefix)[0])
 
+
+def test_quiver_callback_spherical():
+    ds = fake_amr_ds(
+        fields=("density", "velocity_r", "velocity_theta", "velocity_phi"),
+        units=("g/cm**3", "cm/s", "cm/s", "cm/s"),
+        geometry="spherical",
+    )
+
     with _cleanup_fname() as prefix:
-        ds = fake_amr_ds(
-            fields=("density", "velocity_x", "velocity_theta", "velocity_phi"),
-            units=("g/cm**3", "cm/s", "cm/s", "cm/s"),
-            geometry="spherical",
+        p = ProjectionPlot(ds, "phi", ("gas", "density"))
+        p.annotate_quiver(
+            ("gas", "velocity_cylindrical_radius"),
+            ("gas", "velocity_cylindrical_z"),
+            factor=8,
+            scale=0.5,
+            scale_units="inches",
+            normalize=True,
         )
+        assert_fname(p.save(prefix)[0])
+
+    with _cleanup_fname() as prefix:
         p = ProjectionPlot(ds, "r", ("gas", "density"))
-        args = (
+        p.annotate_quiver(
             ("gas", "velocity_theta"),
             ("gas", "velocity_phi"),
-        )
-        kwargs = dict(
             factor=8,
             scale=0.5,
             scale_units="inches",
             normalize=True,
-            bv_x=0.5 * u.cm / u.s,
-            bv_y=0.5 * u.cm / u.s,
         )
-        assert_raises(YTDataTypeUnsupported, p.annotate_quiver, *args, **kwargs)
+        assert_fname(p.save(prefix)[0])
 
 
 @requires_file(cyl_2d)