From 807aede13bf4eeffb0af593a6030f65feca7be02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sat, 20 Nov 2021 19:08:22 +0100 Subject: [PATCH] ENH: make plot callback with 'factor' argument accept 2-tuples to distinguish x/y resolutions --- yt/visualization/plot_modifications.py | 73 +++++++++++++++++------- yt/visualization/tests/test_callbacks.py | 2 +- 2 files changed, 53 insertions(+), 22 deletions(-) diff --git a/yt/visualization/plot_modifications.py b/yt/visualization/plot_modifications.py index 23aaa5091ac..f1f7d409ec8 100644 --- a/yt/visualization/plot_modifications.py +++ b/yt/visualization/plot_modifications.py @@ -1,7 +1,8 @@ import re import warnings from functools import wraps -from numbers import Number +from numbers import Integral, Number +from typing import Tuple, Union import matplotlib import numpy as np @@ -46,6 +47,26 @@ def _check_geometry(self, plot): return _check_geometry +def _validate_factor_tuple( + factor: Union[Tuple[Integral, Integral], Integral] +) -> Tuple[int, int]: + if ( + is_sequence(factor) + and len(factor) == 2 + and all(isinstance(_, Integral) for _ in factor) + ): + # - checking for "is_sequence" allows lists, numpy arrays and other containers + # - checking for Integral type allows numpy integer types + # in any case we return a with strict typing + return tuple(int(_) for _ in factor) + elif isinstance(factor, Integral): + return (int(factor), int(factor)) + else: + raise TypeError( + f"Expected a single, or a pair of integers, received {factor!r}" + ) + + class PlotCallback: # _supported_geometries is set by subclasses of PlotCallback to a tuple of # strings corresponding to the names of the geometries that a callback @@ -307,10 +328,15 @@ class VelocityCallback(PlotCallback): _supported_geometries = ("cartesian", "spectral_cube", "polar", "cylindrical") def __init__( - self, factor=16, scale=None, scale_units=None, normalize=False, plot_args=None + self, + factor: Union[Tuple[int, int], int] = 16, + scale=None, + scale_units=None, + normalize=False, + plot_args=None, ): PlotCallback.__init__(self) - self.factor = factor + self.factor = _validate_factor_tuple(factor) self.scale = scale self.scale_units = scale_units self.normalize = normalize @@ -397,10 +423,15 @@ class MagFieldCallback(PlotCallback): _supported_geometries = ("cartesian", "spectral_cube", "polar", "cylindrical") def __init__( - self, factor=16, scale=None, scale_units=None, normalize=False, plot_args=None + self, + factor: Union[Tuple[int, int], int] = 16, + scale=None, + scale_units=None, + normalize=False, + plot_args=None, ): PlotCallback.__init__(self) - self.factor = factor + self.factor = _validate_factor_tuple(factor) self.scale = scale self.scale_units = scale_units self.normalize = normalize @@ -477,7 +508,7 @@ def __init__( self, field_x, field_y, - factor=16, + factor: Union[Tuple[int, int], int] = 16, scale=None, scale_units=None, normalize=False, @@ -490,7 +521,7 @@ def __init__( self.field_y = field_y self.bv_x = bv_x self.bv_y = bv_y - self.factor = factor + self.factor = _validate_factor_tuple(factor) self.scale = scale self.scale_units = scale_units self.normalize = normalize @@ -529,8 +560,8 @@ def _transformed_field(field, data): # We are feeding this size into the pixelizer, where it will properly # set it in reverse order - nx = plot.image._A.shape[1] // self.factor - ny = plot.image._A.shape[0] // self.factor + nx = plot.image._A.shape[1] // self.factor[0] + ny = plot.image._A.shape[0] // self.factor[1] pixX = plot.data.ds.coordinates.pixelize( plot.data.axis, plot.data, @@ -587,7 +618,7 @@ def __init__( self, field, ncont=5, - factor=4, + factor: Union[Tuple[int, int], int] = 4, clim=None, plot_args=None, label=False, @@ -601,7 +632,7 @@ def __init__( def_text_args = {"colors": "w"} self.ncont = ncont self.field = field - self.factor = factor + self.factor = _validate_factor_tuple(factor) self.clim = clim self.take_log = take_log if plot_args is None: @@ -639,8 +670,8 @@ def __call__(self, plot): # We want xi, yi in plot coordinates xi, yi = np.mgrid[ - xx0 : xx1 : numPoints_x / (self.factor * 1j), - yy0 : yy1 : numPoints_y / (self.factor * 1j), + xx0 : xx1 : numPoints_x / (self.factor[0] * 1j), + yy0 : yy1 : numPoints_y / (self.factor[1] * 1j), ] data = self.data_source or plot.data @@ -684,7 +715,7 @@ def __call__(self, plot): triangulation = Triangulation(x, y) zi = LinearTriInterpolator(triangulation, z)(xi, yi) elif plot._type_name == "OffAxisProjection": - zi = plot.frb[self.field][:: self.factor, :: self.factor].transpose() + zi = plot.frb[self.field][:: self.factor[0], :: self.factor[1]].transpose() if self.take_log is None: field = data._determine_fields([self.field])[0] @@ -907,7 +938,7 @@ def __init__( self, field_x, field_y, - factor=16, + factor: Union[Tuple[int, int], int] = 16, density=1, field_color=None, display_threshold=None, @@ -918,7 +949,7 @@ def __init__( self.field_x = field_x self.field_y = field_y self.field_color = field_color - self.factor = factor + self.factor = _validate_factor_tuple(factor) self.dens = density self.display_threshold = display_threshold if plot_args is None: @@ -931,8 +962,8 @@ def __call__(self, plot): # We are feeding this size into the pixelizer, where it will properly # set it in reverse order - nx = plot.image._A.shape[1] // self.factor - ny = plot.image._A.shape[0] // self.factor + nx = plot.image._A.shape[1] // self.factor[0] + ny = plot.image._A.shape[0] // self.factor[1] pixX = plot.data.ds.coordinates.pixelize( plot.data.axis, plot.data, self.field_x, bounds, (nx, ny) ) @@ -1124,7 +1155,7 @@ def __init__( PlotCallback.__init__(self) self.field_x = field_x self.field_y = field_y - self.factor = factor + self.factor = _validate_factor_tuple(factor) self.scale = scale self.scale_units = scale_units self.normalize = normalize @@ -1135,8 +1166,8 @@ def __init__( def __call__(self, plot): x0, x1, y0, y1 = self._physical_bounds(plot) xx0, xx1, yy0, yy1 = self._plot_bounds(plot) - nx = plot.image._A.shape[1] // self.factor - ny = plot.image._A.shape[0] // self.factor + nx = plot.image._A.shape[1] // self.factor[0] + ny = plot.image._A.shape[0] // self.factor[1] indices = np.argsort(plot.data["index", "dx"])[::-1].astype(np.int_) pixX = np.zeros((ny, nx), dtype="f8") diff --git a/yt/visualization/tests/test_callbacks.py b/yt/visualization/tests/test_callbacks.py index ce7040c7776..7abc505d383 100644 --- a/yt/visualization/tests/test_callbacks.py +++ b/yt/visualization/tests/test_callbacks.py @@ -641,7 +641,7 @@ def test_contour_callback(): slc.annotate_contour( ("gas", "plasma_beta"), ncont=2, - factor=7.0, + factor=7, take_log=False, clim=(1.0e-1, 1.0e1), label=True,