From 000425a5987552a8f611f32a67056582f1f39b5e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Fri, 18 Feb 2022 19:44:02 +0100 Subject: [PATCH] ENH: add support for VelocityCallback and MagFieldCallback in spherical coordinates --- yt/fields/vector_operations.py | 64 +++++++++++++++++++- yt/frontends/stream/fields.py | 10 +++- yt/visualization/plot_modifications.py | 39 ++++++++++++- yt/visualization/plot_window.py | 4 +- yt/visualization/tests/test_callbacks.py | 74 +++++++++++++++++------- 5 files changed, 164 insertions(+), 27 deletions(-) diff --git a/yt/fields/vector_operations.py b/yt/fields/vector_operations.py index 0b678be0469..ef3bb03f295 100644 --- a/yt/fields/vector_operations.py +++ b/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, diff --git a/yt/frontends/stream/fields.py b/yt/frontends/stream/fields.py index 1f14e39cfb6..e26d51d5e78 100644 --- a/yt/frontends/stream/fields.py +++ b/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): diff --git a/yt/visualization/plot_modifications.py b/yt/visualization/plot_modifications.py index bc087d9dcf2..49ddc3ea0f5 100644 --- a/yt/visualization/plot_modifications.py +++ b/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__( diff --git a/yt/visualization/plot_window.py b/yt/visualization/plot_window.py index 9c9c666047f..553ba039b2a 100644 --- a/yt/visualization/plot_window.py +++ b/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(): diff --git a/yt/visualization/tests/test_callbacks.py b/yt/visualization/tests/test_callbacks.py index f4b2040d5cb..ebfe9c9644e 100644 --- a/yt/visualization/tests/test_callbacks.py +++ b/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)