From f55ccee083ed01b34e8ef6493b73674017db2578 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 21 Aug 2025 11:37:21 +0000 Subject: [PATCH 1/3] Refactor: Separate classic and quiver logic The files _plot.py and _animate.py were becoming too large and contained mixed logic for classic (scalar) and quiver (vector) plots and animations. This refactoring separates these concerns into two new files: - `mapflow/_classic.py`: Contains `PlotModel` and `Animation` for scalar data. - `mapflow/_quiver.py`: Contains `plot_da_quiver` and a new `QuiverAnimation` class for vector data. The original `_plot.py` and `_animate.py` files have been removed, and the package's `__init__.py` has been updated to expose the same public API. This change improves code organization and makes the library easier to maintain and extend. --- mapflow/__init__.py | 14 +- mapflow/{_animate.py => _classic.py} | 611 ++++++++++++++------------- mapflow/_misc.py | 33 ++ mapflow/_plot.py | 401 ------------------ mapflow/_quiver.py | 383 +++++++++++++++++ 5 files changed, 735 insertions(+), 707 deletions(-) rename mapflow/{_animate.py => _classic.py} (52%) delete mode 100644 mapflow/_plot.py create mode 100644 mapflow/_quiver.py diff --git a/mapflow/__init__.py b/mapflow/__init__.py index 42ddc7c..ed8ab13 100644 --- a/mapflow/__init__.py +++ b/mapflow/__init__.py @@ -1,11 +1,19 @@ from importlib.metadata import version -from ._animate import Animation, animate, animate_quiver +from ._classic import Animation, PlotModel, animate, plot_da +from ._quiver import QuiverAnimation, animate_quiver, plot_da_quiver from ._misc import check_ffmpeg -from ._plot import PlotModel, plot_da, plot_da_quiver check_ffmpeg() -__all__ = ["Animation", "PlotModel", "animate", "plot_da", "plot_da_quiver", "animate_quiver"] +__all__ = [ + "Animation", + "PlotModel", + "animate", + "plot_da", + "plot_da_quiver", + "animate_quiver", + "QuiverAnimation", +] __version__ = version("mapflow") diff --git a/mapflow/_animate.py b/mapflow/_classic.py similarity index 52% rename from mapflow/_animate.py rename to mapflow/_classic.py index 7ddb2fa..e25529a 100644 --- a/mapflow/_animate.py +++ b/mapflow/_classic.py @@ -1,18 +1,312 @@ -import subprocess +from copy import copy from multiprocessing import Pool from os import cpu_count from pathlib import Path from tempfile import TemporaryDirectory +import subprocess import geopandas as gpd import matplotlib.pyplot as plt import numpy as np import xarray as xr +from matplotlib.collections import PatchCollection +from matplotlib.colors import LogNorm, Normalize +from matplotlib.patches import Polygon as PolygonPatch from pyproj import CRS +from shapely.geometry import MultiPolygon from tqdm.auto import tqdm -from ._misc import TIME_NAME_CANDIDATES, X_NAME_CANDIDATES, Y_NAME_CANDIDATES -from ._plot import PlotModel +from ._misc import ( + TIME_NAME_CANDIDATES, + X_NAME_CANDIDATES, + Y_NAME_CANDIDATES, + check_da, + guess_coord_name, + process_crs, +) + + +class PlotModel: + """A class for plotting 2D data with geographic borders. Useful for multiple + plots of the same geographic domain, as it pre-computes geographic borders. + + Args: + x, y: Coordinates for the plot. + crs: Coordinate Reference System. Defaults to 4326. + borders (gpd.GeoDataFrame | gpd.GeoSeries | None): Custom borders to use. + If None, defaults to world borders from a packaged GeoPackage. + + .. code-block:: python + + import xarray as xr + from mapflow import PlotModel + + ds = xr.tutorial.open_dataset("era5-2mt-2019-03-uk.grib") + da = ds["t2m"].isel(time=0) + + p = PlotModel(x=da.longitude, y=da.latitude) + p(da) + + """ + + def __init__(self, x, y, crs=4326, borders=None): + self.x = np.asarray_chkfinite(x) + self.y = np.asarray_chkfinite(y) + if self.x.ndim != self.y.ndim: + raise ValueError("x and y must have the same dimensionality (both 1D or both 2D)") + + self.crs = CRS.from_user_input(crs) + if self.crs.is_geographic: + self.aspect = 1 / np.cos((self.y.mean() * np.pi / 180)) + else: + self.aspect = 1 + if self.x.ndim == 1: + self.dx = abs(self.x[1] - self.x[0]) + self.dy = abs(self.y[1] - self.y[0]) + else: + self.dx = np.diff(self.x, axis=1).max() + self.dy = np.diff(self.y, axis=0).max() + bbox = ( + self.x.min() - 10 * self.dx, + self.y.min() - 10 * self.dy, + self.x.max() + 10 * self.dx, + self.y.max() + 10 * self.dy, + ) + + if borders is None: + borders_ = gpd.read_file(Path(__file__).parent / "_static" / "world.gpkg") + elif isinstance(borders, (gpd.GeoDataFrame, gpd.GeoSeries)): + borders_ = borders + else: + raise TypeError("borders must be a geopandas GeoDataFrame, GeoSeries, or None.") + borders_ = borders_.to_crs(self.crs).clip(bbox) + self.borders = self._shp_to_patches(borders_) + + @staticmethod + def _shp_to_patches(gdf): + patches = [] + for poly in gdf.geometry.values: + if isinstance(poly, MultiPolygon): + for polygon in poly.geoms: + patches.append(PolygonPatch(polygon.exterior.coords)) + else: + patches.append(PolygonPatch(poly.exterior.coords)) + return PatchCollection(patches, facecolor="none", linewidth=0.5, edgecolor="k") + + @staticmethod + def _log_norm(data, vmin, vmax, qmin, qmax): + """Generates a logarithmic normalization.""" + positive_data = data[data > 0] + if len(positive_data) == 0: + return Normalize(vmin=1e-1, vmax=1e0) + + vmin = np.nanpercentile(positive_data, q=qmin) if vmin is None else vmin + vmax = np.nanpercentile(positive_data, q=qmax) if vmax is None else vmax + + if vmin <= 0 or vmax <= 0: + raise ValueError(f"Normalization range for log scale must be positive. Got vmin={vmin}, vmax={vmax}") + return LogNorm(vmin=vmin, vmax=vmax) + + @staticmethod + def _norm(data, vmin, vmax, qmin, qmax, norm, log): + """Generates a normalization based on the specified parameters. + + Args: + data (array-like): Data to normalize. + vmin (float): Minimum value for normalization. + vmax (float): Maximum value for normalization. + qmin (float): Minimum quantile for normalization (0-100). + qmax (float): Maximum quantile for normalization (0-100). + norm (matplotlib.colors.Normalize): Custom normalization object. + log (bool): Indicates if a logarithmic scale should be used. + + Returns: + matplotlib.colors.Normalize: Normalization object. + + Raises: + ValueError: If qmin/qmax are not between 0-100 or if log=True with no positive values. + """ + # Validate quantile ranges + if not (0 <= qmin <= 100): + raise ValueError(f"qmin must be between 0 and 100, got {qmin}") + if not (0 <= qmax <= 100): + raise ValueError(f"qmax must be between 0 and 100, got {qmax}") + if qmin >= qmax: + raise ValueError(f"qmin must be less than qmax, got {qmin} and {qmax}") + + if norm is not None: + return norm + + if log: + return PlotModel._log_norm(data, vmin, vmax, qmin, qmax) + + vmin = np.nanpercentile(data, q=qmin) if vmin is None else vmin + vmax = np.nanpercentile(data, q=qmax) if vmax is None else vmax + return Normalize(vmin=vmin, vmax=vmax) + + def _process_data(self, data): + if not isinstance(data, np.ndarray): + data = np.asarray(data) + data = np.squeeze(data) + if data.ndim != 2: + raise ValueError("Data must be a 2D array.") + + if self.x.ndim == 1: + if data.shape[0] != self.y.size or data.shape[1] != self.x.size: + raise ValueError("Data shape does not match x and y dimensions.") + else: + if data.shape != self.x.shape: + raise ValueError("Data shape does not match x and y dimensions.") + return data + + def __call__( + self, + data, + figsize=None, + qmin=0.01, + qmax=99.9, + vmin=None, + vmax=None, + log=False, + cmap="jet", + norm=None, + shading="nearest", + shrink=0.4, + label=None, + title=None, + show=True, + ): + """ + Plots a 2D data array using imshow or pcolormesh. + + This method handles the actual plotting of a single frame. It applies + normalization, colormaps, adds a colorbar, overlays borders, sets the + aspect ratio, title, and optionally displays the plot. + + Args: + data (np.ndarray): 2D array of data to plot. + figsize (tuple[float, float], optional): Figure size (width, height) + in inches. Defaults to None (matplotlib's default). + qmin (float, optional): Minimum quantile for color normalization if + vmin is not set. Defaults to 0.01. + qmax (float, optional): Maximum quantile for color normalization if + vmax is not set. Defaults to 99.9. + vmin (float, optional): Minimum value for color normalization. + Overrides qmin. Defaults to None. + vmax (float, optional): Maximum value for color normalization. + Overrides qmax. Defaults to None. + log (bool, optional): Whether to use a logarithmic color scale. + Defaults to False. + cmap (str, optional): Colormap to use. Defaults to "jet". + norm (matplotlib.colors.Normalize, optional): Custom normalization object. + Overrides vmin, vmax, qmin, qmax, log. Defaults to None. + shading (str, optional): Shading method for pcolormesh. + Defaults to "nearest". + shrink (float, optional): Factor by which to shrink the colorbar. + Defaults to 0.4. + label (str, optional): Label for the colorbar. Defaults to None. + title (str, optional): Title for the plot. Defaults to None. + show (bool, optional): Whether to display the plot using `plt.show()`. + Defaults to True. + """ + data = self._process_data(data) + norm = self._norm(data, vmin, vmax, qmin, qmax, norm, log=log) + plt.figure(figsize=figsize) + if (self.x.ndim == 1) and (self.y.ndim == 1): + plt.imshow( + X=data, + cmap=cmap, + norm=norm, + origin="lower", + extent=( + self.x.min() - self.dx / 2, + self.x.max() + self.dx / 2, + self.y.min() - self.dy / 2, + self.y.max() + self.dy / 2, + ), + interpolation=shading, + ) + else: + plt.pcolormesh( + self.x, + self.y, + data, + cmap=cmap, + norm=norm, + shading=shading, + rasterized=True, + ) + plt.colorbar(shrink=shrink, label=label) + plt.xlim(self.x.min() - self.dx / 2, self.x.max() + self.dx / 2) + plt.ylim(self.y.min() - self.dy / 2, self.y.max() + self.dy / 2) + plt.gca().add_collection(copy(self.borders)) + plt.gca().set_aspect(self.aspect) + plt.title(title) + plt.gca().axis("off") + plt.tight_layout() + if show: + plt.show() + + +def plot_da(da: xr.DataArray, x_name=None, y_name=None, crs=4326, **kwargs): + """Convenience function for quick plotting of an xarray DataArray using PlotModel. + + This is a simplified wrapper around the `PlotModel` class that handles: + - Automatic coordinate detection + - CRS processing + - Data sorting and longitude wrapping (for geographic CRS) + - Single-call plotting + + For better performance when making multiple plots of the same geographic domain, + consider using `PlotModel` directly, which pre-computes geographic borders and + can be reused for multiple plots. + + Args: + da: xarray DataArray with 2D data to plot. Must have appropriate coordinates. + x_name: Name of the x-coordinate dimension. If None, will attempt to guess. + y_name: Name of the y-coordinate dimension. If None, will attempt to guess. + crs: Coordinate Reference System. Can be: + - EPSG code (e.g., 4326 for WGS84) + - PROJ string + - pyproj.CRS object + - If the DataArray has a 'crs' attribute, that will be used by default + **kwargs: Additional arguments passed to PlotModel.__call__(), including: + - figsize: Tuple (width, height) in inches + - qmin/qmax: Quantile ranges for color scaling (0-100) + - vmin/vmax: Explicit value ranges for color scaling + - log: Whether to use logarithmic color scale + - cmap: Colormap name + - norm: Custom normalization + - shading: Color shading method + - shrink: Colorbar shrink factor + - label: Colorbar label + - title: Plot title + - show: Whether to display the plot + + Example: + .. code-block:: python + + import xarray as xr + from mapflow import plot_da + + ds = xr.tutorial.open_dataset("era5-2mt-2019-03-uk.grib") + plot_da(da=ds['t2m'].isel(time=0)) + + See Also: + PlotModel: The underlying plotting class used by this function + """ + actual_x_name = guess_coord_name(da.coords, X_NAME_CANDIDATES, x_name, "x") + actual_y_name = guess_coord_name(da.coords, Y_NAME_CANDIDATES, y_name, "y") + + if da[actual_x_name].ndim == 1 and da[actual_y_name].ndim == 1: + da = da.sortby(actual_x_name).sortby(actual_y_name) + crs_ = process_crs(da, crs) + if crs_.is_geographic: + da[actual_x_name] = xr.where(da[actual_x_name] > 180, da[actual_x_name] - 360, da[actual_x_name]) + + p = PlotModel(x=da[actual_x_name].values, y=da[actual_y_name].values, crs=crs_) + data = p._process_data(da.values) + p(data, **kwargs) class Animation: @@ -40,7 +334,6 @@ class Animation: animation(da, "animation.mp4") """ - def __init__(self, x, y, crs=4326, verbose=0, borders=None): self.plot = PlotModel(x=x, y=y, crs=crs, borders=borders) self.verbose = verbose @@ -60,6 +353,8 @@ def upsample(data, ratio=5): @staticmethod def _process_title(title, upsample_ratio): + if title is None: + return if isinstance(title, str): return [title] * upsample_ratio elif isinstance(title, (list, tuple)): @@ -139,52 +434,6 @@ def __call__( timeout=timeout, ) - def quiver( - self, - u, - v, - path, - subsample: int = 1, - figsize: tuple = None, - title=None, - fps: int = 24, - upsample_ratio: int = 2, - cmap="jet", - qmin=0.01, - qmax=99.9, - vmin=None, - vmax=None, - norm=None, - log=False, - label=None, - dpi=180, - n_jobs=None, - timeout="auto", - **kwargs, - ): - self._animate( - data=(u, v), - path=path, - frame_generator=self._generate_quiver_frame, - figsize=figsize, - title=title, - fps=fps, - upsample_ratio=upsample_ratio, - cmap=cmap, - norm=norm, - label=label, - dpi=dpi, - n_jobs=n_jobs, - timeout=timeout, - subsample=subsample, - qmin=qmin, - qmax=qmax, - vmin=vmin, - vmax=vmax, - log=log, - **kwargs, - ) - def _animate( self, data, @@ -200,65 +449,29 @@ def _animate( dpi=180, n_jobs=None, timeout="auto", - **kwargs, ): titles = self._process_title(title, upsample_ratio) - - # Upsample data - if isinstance(data, tuple): # For quiver - u, v = data - magnitude = np.sqrt(u**2 + v**2) - norm = self.plot._norm( - magnitude.values, - vmin=kwargs.get("vmin"), - vmax=kwargs.get("vmax"), - qmin=kwargs.get("qmin", 0.01), - qmax=kwargs.get("qmax", 99.9), - norm=kwargs.get("norm"), - log=kwargs.get("log", False), - ) - kwargs["x_name"] = _guess_coord_name(u.coords, X_NAME_CANDIDATES, kwargs.get("x_name"), "x") - kwargs["y_name"] = _guess_coord_name(u.coords, Y_NAME_CANDIDATES, kwargs.get("y_name"), "y") - u = self.upsample(u.values, ratio=upsample_ratio) - v = self.upsample(v.values, ratio=upsample_ratio) - data = (u, v) - data_len = len(u) - else: # For single field - norm = self.plot._norm( - data, - vmin=kwargs.get("vmin"), - vmax=kwargs.get("vmax"), - qmin=kwargs.get("qmin", 0.01), - qmax=kwargs.get("qmax", 99.9), - norm=kwargs.get("norm"), - log=kwargs.get("log", False), - ) - data = self.upsample(data, ratio=upsample_ratio) - data_len = len(data) + data = self.upsample(data, ratio=upsample_ratio) + data_len = len(data) with TemporaryDirectory() as tempdir: frame_paths = [Path(tempdir) / f"frame_{k:08d}.png" for k in range(data_len)] args = [] for k in range(data_len): - if isinstance(data, tuple): - frame_data = (data[0][k], data[1][k]) - else: - frame_data = data[k] - + frame_data = data[k] arg_tuple = ( frame_data, frame_paths[k], figsize, - titles[k], + titles[k] if titles and k < len(titles) else None, cmap, norm, label, dpi, - kwargs, + {}, # No kwargs ) args.append(arg_tuple) - # Generate frames in parallel n_jobs = int(2 / 3 * cpu_count()) if n_jobs is None else n_jobs with Pool(processes=n_jobs) as pool: list( @@ -290,50 +503,6 @@ def _generate_frame(self, args): plt.clf() plt.close() - def _generate_quiver_frame(self, args): - """Generates a quiver frame and saves it as a PNG.""" - (u_frame, v_frame), frame_path, figsize, title, cmap, norm, label, dpi, kwargs = args - x_name = kwargs.get("x_name") - y_name = kwargs.get("y_name") - coords = {y_name: self.plot.y, x_name: self.plot.x} - dims = (y_name, x_name) - u_da = xr.DataArray(u_frame, coords=coords, dims=dims) - v_da = xr.DataArray(v_frame, coords=coords, dims=dims) - magnitude = np.sqrt(u_frame**2 + v_frame**2) - self.plot( - data=magnitude, - figsize=figsize, - title=title, - show=False, - cmap=cmap, - norm=norm, - label=label, - ) - subsample = kwargs.get("subsample", 1) - if subsample > 1: - u_sub = u_da.isel({y_name: slice(None, None, subsample), x_name: slice(None, None, subsample)}) - v_sub = v_da.isel({y_name: slice(None, None, subsample), x_name: slice(None, None, subsample)}) - x = u_sub[x_name].values - y = u_sub[y_name].values - u_sub = u_sub.values - v_sub = v_sub.values - else: - x = u_da[x_name].values - y = u_da[y_name].values - u_sub = u_da.values - v_sub = v_da.values - - if u_da[x_name].ndim == 1: - x, y = np.meshgrid(x, y) - - arrows_kwgs = kwargs.get("arrows_kwgs") - if arrows_kwgs is None: - arrows_kwgs = {} - plt.quiver(x, y, u_sub, v_sub, **arrows_kwgs) - plt.savefig(frame_path, dpi=dpi, bbox_inches="tight", pad_inches=0.05) - plt.clf() - plt.close() - @staticmethod def _build_ffmpeg_cmd(tempdir, path, fps): path = Path(path) @@ -370,7 +539,7 @@ def _create_video(tempdir, path, fps, timeout): stderr=subprocess.PIPE, timeout=timeout, ) - if result.stdout: # Only print if there's output + if result.stdout: print(result.stdout) except subprocess.CalledProcessError as e: print(f"Error during video creation: {e}") @@ -383,69 +552,6 @@ def _create_video(tempdir, path, fps, timeout): raise -def process_crs(da, crs): - if crs is None: - if "spatial_ref" in da.coords: - crs = da.spatial_ref.attrs.get("crs_wkt", 4326) - else: - crs = 4326 - ret = CRS.from_user_input(crs) - return ret - - -def check_da(da: xr.DataArray, time_name, x_name, y_name, crs): - if not isinstance(da, xr.DataArray): - raise TypeError(f"Expected xarray.DataArray, got {type(da)}") - for dim in (x_name, y_name, time_name): - if dim not in da.coords: - raise ValueError(f"Dimension '{dim}' not found in DataArray coordinates: {da.dims}") - crs_ = process_crs(da, crs) - if crs_.is_geographic: - da[x_name] = xr.where(da[x_name] > 180, da[x_name] - 360, da[x_name]) - - # For non-rectilinear grids (2D coordinates), sorting by spatial dimensions is not possible. - if da[x_name].ndim == 1 and da[y_name].ndim == 1: - da = da.sortby(x_name).sortby(y_name) - - da = da.sortby(time_name).squeeze() - - if da.ndim != 3: - raise ValueError( - f"DataArray must have 3 dimensions ({time_name}, {y_name}, {x_name}), got {da.ndim} dimensions." - ) - - # Ensure time is the first dimension - if da[x_name].ndim == 1 and da[y_name].ndim == 1: - da = da.transpose(time_name, y_name, x_name) - elif list(da.dims)[0] != time_name: - current_dims = list(da.dims) - current_dims.remove(time_name) - new_order = [time_name] + current_dims - da = da.transpose(*new_order) - return da, crs_ - - -def _guess_coord_name(da_coords, candidates, provided_name, coord_type_for_error): - """ - Guesses the coordinate name if not provided. - Iterates through da_coords, compares lowercased names with candidates. - """ - if provided_name is not None: - return provided_name - - for coord_name_key in da_coords: - # Convert coord_name_key to string before lower() in case it's not already a string - coord_name_str = str(coord_name_key).lower() - if coord_name_str in candidates: - return str(coord_name_key) # Return original case name - - raise ValueError( - f"Could not automatically detect {coord_type_for_error}-coordinate. " - f"Please specify '{coord_type_for_error}_name' from available coordinates: {list(da_coords.keys())}. " - f"Tried to guess from candidates: {candidates}." - ) - - def animate( da: xr.DataArray, path: str, @@ -506,10 +612,9 @@ def animate( animate(da=ds['t2m'].isel(time=slice(120)), path='animation.mp4') """ - # Guess coordinate names if not provided - actual_time_name = _guess_coord_name(da.coords, TIME_NAME_CANDIDATES, time_name, "time") - actual_x_name = _guess_coord_name(da.coords, X_NAME_CANDIDATES, x_name, "x") - actual_y_name = _guess_coord_name(da.coords, Y_NAME_CANDIDATES, y_name, "y") + actual_time_name = guess_coord_name(da.coords, TIME_NAME_CANDIDATES, time_name, "time") + actual_x_name = guess_coord_name(da.coords, X_NAME_CANDIDATES, x_name, "x") + actual_y_name = guess_coord_name(da.coords, Y_NAME_CANDIDATES, y_name, "y") da, crs_ = check_da(da, actual_time_name, actual_x_name, actual_y_name, crs) @@ -534,103 +639,3 @@ def animate( label=unit, **kwargs, ) - - -def animate_quiver( - u: xr.DataArray, - v: xr.DataArray, - path: str, - time_name: str = None, - x_name: str = None, - y_name: str = None, - crs=None, - field_name: str = None, - borders: gpd.GeoDataFrame | gpd.GeoSeries | None = None, - verbose: int = 0, - subsample: int = 1, - arrows_kwgs: dict = None, - **kwargs, -): - """ - Creates a quiver animation from two xarray DataArrays. - - Args: - u (xr.DataArray): Input DataArray for the U-component with at least time, x, and y dimensions. - v (xr.DataArray): Input DataArray for the V-component with at least time, x, and y dimensions. - path (str): Output path for the video file. Supported formats are avi, mov and mp4. - time_name (str, optional): Name of the time coordinate in `da`. If None, - it's guessed from ['time', 't', 'times']. Defaults to None. - x_name (str, optional): Name of the x-coordinate (e.g., longitude) in `da`. - If None, it's guessed from ['x', 'lon', 'longitude']. Defaults to None. - y_name (str, optional): Name of the y-coordinate (e.g., latitude) in `da`. - If None, it's guessed from ['y', 'lat', 'latitude']. Defaults to None. - crs (int | str | CRS, optional): Coordinate Reference System of the data. - Defaults to 4326 (WGS84). - borders (gpd.GeoDataFrame | gpd.GeoSeries | None, optional): - Custom borders to use for plotting. If None, defaults to - world borders. Defaults to None. - verbose (int, optional): Verbosity level for the Animation class. - Defaults to 0. - subsample (int, optional): The subsampling factor for the quiver arrows. - For example, a value of 10 will plot one arrow for every 10 grid points. - Defaults to 1. - arrows_kwgs (dict, optional): Additional keyword arguments passed to - `matplotlib.pyplot.quiver`. Defaults to None. - **kwargs: Additional keyword arguments passed to the Animation class, including: - - cmap (str): Colormap for the plot. Defaults to "jet". - - norm (matplotlib.colors.Normalize): Custom normalization object. - - log (bool): Use logarithmic color scale. Defaults to False. - - qmin (float): Minimum quantile for color normalization. Defaults to 0.01. - - qmax (float): Maximum quantile for color normalization. Defaults to 99.9. - - vmin (float): Minimum value for color normalization. Overrides qmin. - - vmax (float): Maximum value for color normalization. Overrides qmax. - - time_format (str): Strftime format for time in titles. Defaults to "%Y-%m-%dT%H". - - upsample_ratio (int): Factor to upsample data temporally. Defaults to 4. - - fps (int): Frames per second for the video. Defaults to 24. - - n_jobs (int): Number of parallel jobs for frame generation. - - dpi (int): Dots per inch for the saved frames. Defaults to 180. - - timeout (str | int): Timeout for video creation. Defaults to 'auto'. - - Example: - .. code-block:: python - - import xarray as xr - from mapflow import animate_quiver - - ds = xr.tutorial.load_dataset("air_temperature_gradient") - animate_quiver(u=ds["dTdx"], v=ds["dTdy"], path='animation.mkv', subsample=3) - """ - actual_time_name = _guess_coord_name(u.coords, TIME_NAME_CANDIDATES, time_name, "time") - actual_x_name = _guess_coord_name(u.coords, X_NAME_CANDIDATES, x_name, "x") - actual_y_name = _guess_coord_name(u.coords, Y_NAME_CANDIDATES, y_name, "y") - - u, crs_ = check_da(u, actual_time_name, actual_x_name, actual_y_name, crs) - v, _ = check_da(v, actual_time_name, actual_x_name, actual_y_name, crs) - - animation = Animation( - x=u[actual_x_name].values, - y=u[actual_y_name].values, - crs=crs_, - verbose=verbose, - borders=borders, - ) - output_path = Path(path) - output_path.parent.mkdir(exist_ok=True, parents=True) - unit = u.attrs.get("unit", None) or u.attrs.get("units", None) - time_format = kwargs.get("time_format", "%Y-%m-%dT%H") - time = u[actual_time_name].dt.strftime(time_format).values - if field_name is None: - titles = [f"{t}" for t in time] - else: - titles = [f"{field_name} · {t}" for t in time] - - animation.quiver( - u=u, - v=v, - path=output_path, - title=titles, - label=unit, - subsample=subsample, - arrows_kwgs=arrows_kwgs, - **kwargs, - ) diff --git a/mapflow/_misc.py b/mapflow/_misc.py index 6401bb8..46e1540 100644 --- a/mapflow/_misc.py +++ b/mapflow/_misc.py @@ -49,3 +49,36 @@ def process_crs(da, crs): else: crs = 4326 return CRS.from_user_input(crs) + + +def check_da(da, time_name, x_name, y_name, crs): + import xarray as xr + if not isinstance(da, xr.DataArray): + raise TypeError(f"Expected xarray.DataArray, got {type(da)}") + for dim in (x_name, y_name, time_name): + if dim not in da.coords: + raise ValueError(f"Dimension '{dim}' not found in DataArray coordinates: {da.dims}") + crs_ = process_crs(da, crs) + if crs_.is_geographic: + da[x_name] = xr.where(da[x_name] > 180, da[x_name] - 360, da[x_name]) + + # For non-rectilinear grids (2D coordinates), sorting by spatial dimensions is not possible. + if da[x_name].ndim == 1 and da[y_name].ndim == 1: + da = da.sortby(x_name).sortby(y_name) + + da = da.sortby(time_name).squeeze() + + if da.ndim != 3: + raise ValueError( + f"DataArray must have 3 dimensions ({time_name}, {y_name}, {x_name}), got {da.ndim} dimensions." + ) + + # Ensure time is the first dimension + if da[x_name].ndim == 1 and da[y_name].ndim == 1: + da = da.transpose(time_name, y_name, x_name) + elif list(da.dims)[0] != time_name: + current_dims = list(da.dims) + current_dims.remove(time_name) + new_order = [time_name] + current_dims + da = da.transpose(*new_order) + return da, crs_ diff --git a/mapflow/_plot.py b/mapflow/_plot.py deleted file mode 100644 index f99bc0c..0000000 --- a/mapflow/_plot.py +++ /dev/null @@ -1,401 +0,0 @@ -from copy import copy -from pathlib import Path - -import geopandas as gpd -import matplotlib.pyplot as plt -import numpy as np -import xarray as xr -from matplotlib.collections import PatchCollection -from matplotlib.colors import LogNorm, Normalize -from matplotlib.patches import Polygon as PolygonPatch -from pyproj import CRS -from shapely.geometry import MultiPolygon - -from ._misc import X_NAME_CANDIDATES, Y_NAME_CANDIDATES, guess_coord_name, process_crs - - -class PlotModel: - """A class for plotting 2D data with geographic borders. Useful for multiple - plots of the same geographic domain, as it pre-computes geographic borders. - - Args: - x, y: Coordinates for the plot. - crs: Coordinate Reference System. Defaults to 4326. - borders (gpd.GeoDataFrame | gpd.GeoSeries | None): Custom borders to use. - If None, defaults to world borders from a packaged GeoPackage. - - .. code-block:: python - - import xarray as xr - from mapflow import PlotModel - - ds = xr.tutorial.open_dataset("era5-2mt-2019-03-uk.grib") - da = ds["t2m"].isel(time=0) - - p = PlotModel(x=da.longitude, y=da.latitude) - p(da) - - """ - - def __init__(self, x, y, crs=4326, borders=None): - self.x = np.asarray_chkfinite(x) - self.y = np.asarray_chkfinite(y) - if self.x.ndim != self.y.ndim: - raise ValueError("x and y must have the same dimensionality (both 1D or both 2D)") - - self.crs = CRS.from_user_input(crs) - if self.crs.is_geographic: - self.aspect = 1 / np.cos((self.y.mean() * np.pi / 180)) - else: - self.aspect = 1 - if self.x.ndim == 1: - self.dx = abs(self.x[1] - self.x[0]) - self.dy = abs(self.y[1] - self.y[0]) - else: - self.dx = np.diff(self.x, axis=1).max() - self.dy = np.diff(self.y, axis=0).max() - bbox = ( - self.x.min() - 10 * self.dx, - self.y.min() - 10 * self.dy, - self.x.max() + 10 * self.dx, - self.y.max() + 10 * self.dy, - ) - - if borders is None: - borders_ = gpd.read_file(Path(__file__).parent / "_static" / "world.gpkg") - elif isinstance(borders, (gpd.GeoDataFrame, gpd.GeoSeries)): - borders_ = borders - else: - raise TypeError("borders must be a geopandas GeoDataFrame, GeoSeries, or None.") - borders_ = borders_.to_crs(self.crs).clip(bbox) - self.borders = self._shp_to_patches(borders_) - - @staticmethod - def _shp_to_patches(gdf): - patches = [] - for poly in gdf.geometry.values: - if isinstance(poly, MultiPolygon): - for polygon in poly.geoms: - patches.append(PolygonPatch(polygon.exterior.coords)) - else: - patches.append(PolygonPatch(poly.exterior.coords)) - return PatchCollection(patches, facecolor="none", linewidth=0.5, edgecolor="k") - - @staticmethod - def _log_norm(data, vmin, vmax, qmin, qmax): - """Generates a logarithmic normalization.""" - positive_data = data[data > 0] - if len(positive_data) == 0: - return Normalize(vmin=1e-1, vmax=1e0) - - vmin = np.nanpercentile(positive_data, q=qmin) if vmin is None else vmin - vmax = np.nanpercentile(positive_data, q=qmax) if vmax is None else vmax - - if vmin <= 0 or vmax <= 0: - raise ValueError(f"Normalization range for log scale must be positive. Got vmin={vmin}, vmax={vmax}") - return LogNorm(vmin=vmin, vmax=vmax) - - @staticmethod - def _norm(data, vmin, vmax, qmin, qmax, norm, log): - """Generates a normalization based on the specified parameters. - - Args: - data (array-like): Data to normalize. - vmin (float): Minimum value for normalization. - vmax (float): Maximum value for normalization. - qmin (float): Minimum quantile for normalization (0-100). - qmax (float): Maximum quantile for normalization (0-100). - norm (matplotlib.colors.Normalize): Custom normalization object. - log (bool): Indicates if a logarithmic scale should be used. - - Returns: - matplotlib.colors.Normalize: Normalization object. - - Raises: - ValueError: If qmin/qmax are not between 0-100 or if log=True with no positive values. - """ - # Validate quantile ranges - if not (0 <= qmin <= 100): - raise ValueError(f"qmin must be between 0 and 100, got {qmin}") - if not (0 <= qmax <= 100): - raise ValueError(f"qmax must be between 0 and 100, got {qmax}") - if qmin >= qmax: - raise ValueError(f"qmin must be less than qmax, got {qmin} and {qmax}") - - if norm is not None: - return norm - - if log: - return PlotModel._log_norm(data, vmin, vmax, qmin, qmax) - - vmin = np.nanpercentile(data, q=qmin) if vmin is None else vmin - vmax = np.nanpercentile(data, q=qmax) if vmax is None else vmax - return Normalize(vmin=vmin, vmax=vmax) - - def _process_data(self, data): - if not isinstance(data, np.ndarray): - data = np.asarray(data) - data = np.squeeze(data) - if data.ndim != 2: - raise ValueError("Data must be a 2D array.") - - if self.x.ndim == 1: - if data.shape[0] != self.y.size or data.shape[1] != self.x.size: - raise ValueError("Data shape does not match x and y dimensions.") - else: - if data.shape != self.x.shape: - raise ValueError("Data shape does not match x and y dimensions.") - return data - - def __call__( - self, - data, - figsize=None, - qmin=0.01, - qmax=99.9, - vmin=None, - vmax=None, - log=False, - cmap="jet", - norm=None, - shading="nearest", - shrink=0.4, - label=None, - title=None, - show=True, - ): - """ - Plots a 2D data array using imshow or pcolormesh. - - This method handles the actual plotting of a single frame. It applies - normalization, colormaps, adds a colorbar, overlays borders, sets the - aspect ratio, title, and optionally displays the plot. - - Args: - data (np.ndarray): 2D array of data to plot. - figsize (tuple[float, float], optional): Figure size (width, height) - in inches. Defaults to None (matplotlib's default). - qmin (float, optional): Minimum quantile for color normalization if - vmin is not set. Defaults to 0.01. - qmax (float, optional): Maximum quantile for color normalization if - vmax is not set. Defaults to 99.9. - vmin (float, optional): Minimum value for color normalization. - Overrides qmin. Defaults to None. - vmax (float, optional): Maximum value for color normalization. - Overrides qmax. Defaults to None. - log (bool, optional): Whether to use a logarithmic color scale. - Defaults to False. - cmap (str, optional): Colormap to use. Defaults to "jet". - norm (matplotlib.colors.Normalize, optional): Custom normalization object. - Overrides vmin, vmax, qmin, qmax, log. Defaults to None. - shading (str, optional): Shading method for pcolormesh. - Defaults to "nearest". - shrink (float, optional): Factor by which to shrink the colorbar. - Defaults to 0.4. - label (str, optional): Label for the colorbar. Defaults to None. - title (str, optional): Title for the plot. Defaults to None. - show (bool, optional): Whether to display the plot using `plt.show()`. - Defaults to True. - """ - data = self._process_data(data) - norm = self._norm(data, vmin, vmax, qmin, qmax, norm, log=log) - plt.figure(figsize=figsize) - if (self.x.ndim == 1) and (self.y.ndim == 1): - plt.imshow( - X=data, - cmap=cmap, - norm=norm, - origin="lower", - extent=( - self.x.min() - self.dx / 2, - self.x.max() + self.dx / 2, - self.y.min() - self.dy / 2, - self.y.max() + self.dy / 2, - ), - interpolation=shading, - ) - else: - plt.pcolormesh( - self.x, - self.y, - data, - cmap=cmap, - norm=norm, - shading=shading, - rasterized=True, - ) - plt.colorbar(shrink=shrink, label=label) - plt.xlim(self.x.min() - self.dx / 2, self.x.max() + self.dx / 2) - plt.ylim(self.y.min() - self.dy / 2, self.y.max() + self.dy / 2) - plt.gca().add_collection(copy(self.borders)) - plt.gca().set_aspect(self.aspect) - plt.title(title) - plt.gca().axis("off") - plt.tight_layout() - if show: - plt.show() - - -def plot_da(da: xr.DataArray, x_name=None, y_name=None, crs=4326, **kwargs): - """Convenience function for quick plotting of an xarray DataArray using PlotModel. - - This is a simplified wrapper around the `PlotModel` class that handles: - - Automatic coordinate detection - - CRS processing - - Data sorting and longitude wrapping (for geographic CRS) - - Single-call plotting - - For better performance when making multiple plots of the same geographic domain, - consider using `PlotModel` directly, which pre-computes geographic borders and - can be reused for multiple plots. - - Args: - da: xarray DataArray with 2D data to plot. Must have appropriate coordinates. - x_name: Name of the x-coordinate dimension. If None, will attempt to guess. - y_name: Name of the y-coordinate dimension. If None, will attempt to guess. - crs: Coordinate Reference System. Can be: - - EPSG code (e.g., 4326 for WGS84) - - PROJ string - - pyproj.CRS object - - If the DataArray has a 'crs' attribute, that will be used by default - **kwargs: Additional arguments passed to PlotModel.__call__(), including: - - figsize: Tuple (width, height) in inches - - qmin/qmax: Quantile ranges for color scaling (0-100) - - vmin/vmax: Explicit value ranges for color scaling - - log: Whether to use logarithmic color scale - - cmap: Colormap name - - norm: Custom normalization - - shading: Color shading method - - shrink: Colorbar shrink factor - - label: Colorbar label - - title: Plot title - - show: Whether to display the plot - - Example: - .. code-block:: python - - import xarray as xr - from mapflow import plot_da - - ds = xr.tutorial.open_dataset("era5-2mt-2019-03-uk.grib") - plot_da(da=ds['t2m'].isel(time=0)) - - See Also: - PlotModel: The underlying plotting class used by this function - """ - actual_x_name = guess_coord_name(da.coords, X_NAME_CANDIDATES, x_name, "x") - actual_y_name = guess_coord_name(da.coords, Y_NAME_CANDIDATES, y_name, "y") - - if da[actual_x_name].ndim == 1 and da[actual_y_name].ndim == 1: - da = da.sortby(actual_x_name).sortby(actual_y_name) - crs_ = process_crs(da, crs) - if crs_.is_geographic: - da[actual_x_name] = xr.where(da[actual_x_name] > 180, da[actual_x_name] - 360, da[actual_x_name]) - - p = PlotModel(x=da[actual_x_name].values, y=da[actual_y_name].values, crs=crs_) - data = p._process_data(da.values) - p(data, **kwargs) - - -def plot_da_quiver( - u, - v, - x_name=None, - y_name=None, - crs=4326, - subsample: int = 1, - show=True, - arrows_kwgs: dict = None, - **kwargs, -): - """ - Plots a quiver plot from two xarray DataArrays, representing the U and V - components of a vector field. - - The magnitude of the vector field is represented by a color mesh, and the - direction is shown with quiver arrows. - - Args: - u (xr.DataArray): DataArray for the U-component of the vector field. - v (xr.DataArray): DataArray for the V-component of the vector field. - x_name (str, optional): Name of the x-coordinate dimension. - If None, will attempt to guess. - y_name (str, optional): Name of the y-coordinate dimension. - If None, will attempt to guess. - crs: Coordinate Reference System. Can be: - - EPSG code (e.g., 4326 for WGS84) - - PROJ string - - pyproj.CRS object - - If the DataArray has a 'crs' attribute, that will be used by default - subsample (int, optional): The subsampling factor for the quiver arrows. - For example, a value of 10 will plot one arrow for every 10 grid points. - Defaults to 1. - show: Whether to display the plot - arrows_kwgs (dict, optional): Additional keyword arguments passed to - `matplotlib.pyplot.quiver`. Defaults to None. - **kwargs: Additional arguments passed to PlotModel.__call__(), including: - - figsize: Tuple (width, height) in inches - - qmin/qmax: Quantile ranges for color scaling (0-100) - - vmin/vmax: Explicit value ranges for color scaling - - log: Whether to use logarithmic color scale - - cmap: Colormap name - - norm: Custom normalization - - shading: Color shading method - - shrink: Colorbar shrink factor - - label: Colorbar label - - title: Plot title - - Example: - .. code-block:: python - - import xarray as xr - from mapflow import plot_da_quiver - - ds = xr.tutorial.load_dataset("air_temperature_gradient").isel(time=0) - plot_da_quiver(u=ds["dTdx"], v=ds["dTdy"], subsample=4) - - See Also: - PlotModel: The underlying plotting class used by this function. - """ - actual_x_name = guess_coord_name(u.coords, X_NAME_CANDIDATES, x_name, "x") - actual_y_name = guess_coord_name(u.coords, Y_NAME_CANDIDATES, y_name, "y") - - if u[actual_x_name].ndim == 1 and u[actual_y_name].ndim == 1: - u = u.sortby(actual_x_name).sortby(actual_y_name) - v = v.sortby(actual_x_name).sortby(actual_y_name) - - crs_ = process_crs(u, crs) - if crs_.is_geographic: - u[actual_x_name] = xr.where(u[actual_x_name] > 180, u[actual_x_name] - 360, u[actual_x_name]) - v[actual_x_name] = xr.where(v[actual_x_name] > 180, v[actual_x_name] - 360, v[actual_x_name]) - - magnitude = np.sqrt(u**2 + v**2) - p = PlotModel(x=u[actual_x_name].values, y=u[actual_y_name].values, crs=crs_) - data = p._process_data(magnitude.values) - p(data, show=False, **kwargs) - - if subsample > 1: - u_subsampled = u.isel( - {actual_y_name: slice(None, None, subsample), actual_x_name: slice(None, None, subsample)} - ) - v_subsampled = v.isel( - {actual_y_name: slice(None, None, subsample), actual_x_name: slice(None, None, subsample)} - ) - x = u_subsampled[actual_x_name].values - y = u_subsampled[actual_y_name].values - u_subsampled = u_subsampled.values - v_subsampled = v_subsampled.values - else: - x = u[actual_x_name].values - y = u[actual_y_name].values - u_subsampled = u.values - v_subsampled = v.values - - if u[actual_x_name].ndim == 1: - x, y = np.meshgrid(x, y) - - if arrows_kwgs is None: - arrows_kwgs = {} - plt.quiver(x, y, u_subsampled, v_subsampled, **arrows_kwgs) - if show: - plt.show() diff --git a/mapflow/_quiver.py b/mapflow/_quiver.py new file mode 100644 index 0000000..14c34af --- /dev/null +++ b/mapflow/_quiver.py @@ -0,0 +1,383 @@ +from multiprocessing import Pool +from os import cpu_count +from pathlib import Path +from tempfile import TemporaryDirectory + +import geopandas as gpd +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr +from tqdm.auto import tqdm + +from ._classic import Animation, PlotModel +from ._misc import ( + TIME_NAME_CANDIDATES, + X_NAME_CANDIDATES, + Y_NAME_CANDIDATES, + check_da, + guess_coord_name, + process_crs, +) + + +def plot_da_quiver( + u, + v, + x_name=None, + y_name=None, + crs=4326, + subsample: int = 1, + show=True, + arrows_kwgs: dict = None, + **kwargs, +): + """ + Plots a quiver plot from two xarray DataArrays, representing the U and V + components of a vector field. + + The magnitude of the vector field is represented by a color mesh, and the + direction is shown with quiver arrows. + + Args: + u (xr.DataArray): DataArray for the U-component of the vector field. + v (xr.DataArray): DataArray for the V-component of the vector field. + x_name (str, optional): Name of the x-coordinate dimension. + If None, will attempt to guess. + y_name (str, optional): Name of the y-coordinate dimension. + If None, will attempt to guess. + crs: Coordinate Reference System. Can be: + - EPSG code (e.g., 4326 for WGS84) + - PROJ string + - pyproj.CRS object + - If the DataArray has a 'crs' attribute, that will be used by default + subsample (int, optional): The subsampling factor for the quiver arrows. + For example, a value of 10 will plot one arrow for every 10 grid points. + Defaults to 1. + show: Whether to display the plot + arrows_kwgs (dict, optional): Additional keyword arguments passed to + `matplotlib.pyplot.quiver`. Defaults to None. + **kwargs: Additional arguments passed to PlotModel.__call__(), including: + - figsize: Tuple (width, height) in inches + - qmin/qmax: Quantile ranges for color scaling (0-100) + - vmin/vmax: Explicit value ranges for color scaling + - log: Whether to use logarithmic color scale + - cmap: Colormap name + - norm: Custom normalization + - shading: Color shading method + - shrink: Colorbar shrink factor + - label: Colorbar label + - title: Plot title + + Example: + .. code-block:: python + + import xarray as xr + from mapflow import plot_da_quiver + + ds = xr.tutorial.load_dataset("air_temperature_gradient").isel(time=0) + plot_da_quiver(u=ds["dTdx"], v=ds["dTdy"], subsample=4) + + See Also: + PlotModel: The underlying plotting class used by this function. + """ + actual_x_name = guess_coord_name(u.coords, X_NAME_CANDIDATES, x_name, "x") + actual_y_name = guess_coord_name(u.coords, Y_NAME_CANDIDATES, y_name, "y") + + if u[actual_x_name].ndim == 1 and u[actual_y_name].ndim == 1: + u = u.sortby(actual_x_name).sortby(actual_y_name) + v = v.sortby(actual_x_name).sortby(actual_y_name) + + crs_ = process_crs(u, crs) + if crs_.is_geographic: + u[actual_x_name] = xr.where(u[actual_x_name] > 180, u[actual_x_name] - 360, u[actual_x_name]) + v[actual_x_name] = xr.where(v[actual_x_name] > 180, v[actual_x_name] - 360, v[actual_x_name]) + + magnitude = np.sqrt(u**2 + v**2) + p = PlotModel(x=u[actual_x_name].values, y=u[actual_y_name].values, crs=crs_) + data = p._process_data(magnitude.values) + p(data, show=False, **kwargs) + + if subsample > 1: + u_subsampled = u.isel( + {actual_y_name: slice(None, None, subsample), actual_x_name: slice(None, None, subsample)} + ) + v_subsampled = v.isel( + {actual_y_name: slice(None, None, subsample), actual_x_name: slice(None, None, subsample)} + ) + x = u_subsampled[actual_x_name].values + y = u_subsampled[actual_y_name].values + u_subsampled = u_subsampled.values + v_subsampled = v_subsampled.values + else: + x = u[actual_x_name].values + y = u[actual_y_name].values + u_subsampled = u.values + v_subsampled = v.values + + if u[actual_x_name].ndim == 1: + x, y = np.meshgrid(x, y) + + if arrows_kwgs is None: + arrows_kwgs = {} + plt.quiver(x, y, u_subsampled, v_subsampled, **arrows_kwgs) + if show: + plt.show() + + +class QuiverAnimation(Animation): + def quiver( + self, + u, + v, + path, + subsample: int = 1, + figsize: tuple = None, + title=None, + fps: int = 24, + upsample_ratio: int = 2, + cmap="jet", + qmin=0.01, + qmax=99.9, + vmin=None, + vmax=None, + norm=None, + log=False, + label=None, + dpi=180, + n_jobs=None, + timeout="auto", + **kwargs, + ): + magnitude = np.sqrt(u**2 + v**2) + norm = self.plot._norm( + magnitude.values, + vmin=vmin, + vmax=vmax, + qmin=qmin, + qmax=qmax, + norm=norm, + log=log, + ) + self._animate( + data=(u, v), + path=path, + frame_generator=self._generate_quiver_frame, + figsize=figsize, + title=title, + fps=fps, + upsample_ratio=upsample_ratio, + cmap=cmap, + norm=norm, + label=label, + dpi=dpi, + n_jobs=n_jobs, + timeout=timeout, + subsample=subsample, + **kwargs, + ) + + def _animate( + self, + data, + path, + frame_generator, + figsize: tuple = None, + title=None, + fps: int = 24, + upsample_ratio: int = 2, + cmap="jet", + norm=None, + label=None, + dpi=180, + n_jobs=None, + timeout="auto", + **kwargs, + ): + titles = self._process_title(title, upsample_ratio) + + u, v = data + u_upsampled = self.upsample(u.values, ratio=upsample_ratio) + v_upsampled = self.upsample(v.values, ratio=upsample_ratio) + data = (u_upsampled, v_upsampled) + data_len = len(u_upsampled) + + with TemporaryDirectory() as tempdir: + frame_paths = [Path(tempdir) / f"frame_{k:08d}.png" for k in range(data_len)] + args = [] + for k in range(data_len): + frame_data = (data[0][k], data[1][k]) + arg_tuple = ( + frame_data, + frame_paths[k], + figsize, + titles[k] if titles and k < len(titles) else None, + cmap, + norm, + label, + dpi, + kwargs, + ) + args.append(arg_tuple) + + n_jobs = int(2 / 3 * cpu_count()) if n_jobs is None else n_jobs + with Pool(processes=n_jobs) as pool: + list( + tqdm( + pool.imap(frame_generator, args), + total=data_len, + disable=(not self.verbose), + desc="Frames generation", + leave=False, + ) + ) + + timeout = max(20, 0.1 * data_len) if timeout == "auto" else timeout + self._create_video(tempdir, path, fps, timeout=timeout) + + def _generate_quiver_frame(self, args): + """Generates a quiver frame and saves it as a PNG.""" + (u_frame, v_frame), frame_path, figsize, title, cmap, norm, label, dpi, kwargs = args + x_name = kwargs.get("x_name") + y_name = kwargs.get("y_name") + coords = {y_name: self.plot.y, x_name: self.plot.x} + dims = (y_name, x_name) + u_da = xr.DataArray(u_frame, coords=coords, dims=dims) + v_da = xr.DataArray(v_frame, coords=coords, dims=dims) + magnitude = np.sqrt(u_frame**2 + v_frame**2) + self.plot( + data=magnitude, + figsize=figsize, + title=title, + show=False, + cmap=cmap, + norm=norm, + label=label, + ) + subsample = kwargs.get("subsample", 1) + if subsample > 1: + u_sub = u_da.isel({y_name: slice(None, None, subsample), x_name: slice(None, None, subsample)}) + v_sub = v_da.isel({y_name: slice(None, None, subsample), x_name: slice(None, None, subsample)}) + x = u_sub[x_name].values + y = u_sub[y_name].values + u_sub = u_sub.values + v_sub = v_sub.values + else: + x = u_da[x_name].values + y = u_da[y_name].values + u_sub = u_da.values + v_sub = v_da.values + + if u_da[x_name].ndim == 1: + x, y = np.meshgrid(x, y) + + arrows_kwgs = kwargs.get("arrows_kwgs") + if arrows_kwgs is None: + arrows_kwgs = {} + plt.quiver(x, y, u_sub, v_sub, **arrows_kwgs) + plt.savefig(frame_path, dpi=dpi, bbox_inches="tight", pad_inches=0.05) + plt.clf() + plt.close() + + +def animate_quiver( + u: xr.DataArray, + v: xr.DataArray, + path: str, + time_name: str = None, + x_name: str = None, + y_name: str = None, + crs=None, + field_name: str = None, + borders: gpd.GeoDataFrame | gpd.GeoSeries | None = None, + verbose: int = 0, + subsample: int = 1, + arrows_kwgs: dict = None, + **kwargs, +): + """ + Creates a quiver animation from two xarray DataArrays. + + Args: + u (xr.DataArray): Input DataArray for the U-component with at least time, x, and y dimensions. + v (xr.DataArray): Input DataArray for the V-component with at least time, x, and y dimensions. + path (str): Output path for the video file. Supported formats are avi, mov and mp4. + time_name (str, optional): Name of the time coordinate in `da`. If None, + it's guessed from ['time', 't', 'times']. Defaults to None. + x_name (str, optional): Name of the x-coordinate (e.g., longitude) in `da`. + If None, it's guessed from ['x', 'lon', 'longitude']. Defaults to None. + y_name (str, optional): Name of the y-coordinate (e.g., latitude) in `da`. + If None, it's guessed from ['y', 'lat', 'latitude']. Defaults to None. + crs (int | str | CRS, optional): Coordinate Reference System of the data. + Defaults to 4326 (WGS84). + borders (gpd.GeoDataFrame | gpd.GeoSeries | None, optional): + Custom borders to use for plotting. If None, defaults to + world borders. Defaults to None. + verbose (int, optional): Verbosity level for the Animation class. + Defaults to 0. + subsample (int, optional): The subsampling factor for the quiver arrows. + For example, a value of 10 will plot one arrow for every 10 grid points. + Defaults to 1. + arrows_kwgs (dict, optional): Additional keyword arguments passed to + `matplotlib.pyplot.quiver`. Defaults to None. + **kwargs: Additional keyword arguments passed to the Animation class, including: + - cmap (str): Colormap for the plot. Defaults to "jet". + - norm (matplotlib.colors.Normalize): Custom normalization object. + - log (bool): Use logarithmic color scale. Defaults to False. + - qmin (float): Minimum quantile for color normalization. Defaults to 0.01. + - qmax (float): Maximum quantile for color normalization. Defaults to 99.9. + - vmin (float): Minimum value for color normalization. Overrides qmin. + - vmax (float): Maximum value for color normalization. Overrides qmax. + - time_format (str): Strftime format for time in titles. Defaults to "%Y-%m-%dT%H". + - upsample_ratio (int): Factor to upsample data temporally. Defaults to 4. + - fps (int): Frames per second for the video. Defaults to 24. + - n_jobs (int): Number of parallel jobs for frame generation. + - dpi (int): Dots per inch for the saved frames. Defaults to 180. + - timeout (str | int): Timeout for video creation. Defaults to 'auto'. + + Example: + .. code-block:: python + + import xarray as xr + from mapflow import animate_quiver + + ds = xr.tutorial.load_dataset("air_temperature_gradient") + animate_quiver(u=ds["dTdx"], v=ds["dTdy"], path='animation.mkv', subsample=3) + """ + actual_time_name = guess_coord_name(u.coords, TIME_NAME_CANDIDATES, time_name, "time") + actual_x_name = guess_coord_name(u.coords, X_NAME_CANDIDATES, x_name, "x") + actual_y_name = guess_coord_name(u.coords, Y_NAME_CANDIDATES, y_name, "y") + + u, crs_ = check_da(u, actual_time_name, actual_x_name, actual_y_name, crs) + v, _ = check_da(v, actual_time_name, actual_x_name, actual_y_name, crs) + + animation = QuiverAnimation( + x=u[actual_x_name].values, + y=u[actual_y_name].values, + crs=crs_, + verbose=verbose, + borders=borders, + ) + output_path = Path(path) + output_path.parent.mkdir(exist_ok=True, parents=True) + unit = u.attrs.get("unit", None) or u.attrs.get("units", None) + time_format = kwargs.get("time_format", "%Y-%m-%dT%H") + time = u[actual_time_name].dt.strftime(time_format).values + if field_name is None: + titles = [f"{t}" for t in time] + else: + titles = [f"{field_name} · {t}" for t in time] + + quiver_kwargs = kwargs.copy() + quiver_kwargs["x_name"] = actual_x_name + quiver_kwargs["y_name"] = actual_y_name + + animation.quiver( + u=u, + v=v, + path=output_path, + title=titles, + label=unit, + subsample=subsample, + arrows_kwgs=arrows_kwgs, + **quiver_kwargs, + ) From c46d31799a0b7dd4cfb8cb3451c1fe114e209a47 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 21 Aug 2025 11:39:46 +0000 Subject: [PATCH 2/3] Format code with Ruff --- mapflow/_classic.py | 1 + mapflow/_misc.py | 1 + 2 files changed, 2 insertions(+) diff --git a/mapflow/_classic.py b/mapflow/_classic.py index e25529a..ccf308a 100644 --- a/mapflow/_classic.py +++ b/mapflow/_classic.py @@ -334,6 +334,7 @@ class Animation: animation(da, "animation.mp4") """ + def __init__(self, x, y, crs=4326, verbose=0, borders=None): self.plot = PlotModel(x=x, y=y, crs=crs, borders=borders) self.verbose = verbose diff --git a/mapflow/_misc.py b/mapflow/_misc.py index 46e1540..4c7e25e 100644 --- a/mapflow/_misc.py +++ b/mapflow/_misc.py @@ -53,6 +53,7 @@ def process_crs(da, crs): def check_da(da, time_name, x_name, y_name, crs): import xarray as xr + if not isinstance(da, xr.DataArray): raise TypeError(f"Expected xarray.DataArray, got {type(da)}") for dim in (x_name, y_name, time_name): From c4731ed41c05342cf9637da48c933bf2f4d9e55b Mon Sep 17 00:00:00 2001 From: Cyril Date: Thu, 21 Aug 2025 11:43:03 +0000 Subject: [PATCH 3/3] import --- mapflow/_misc.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/mapflow/_misc.py b/mapflow/_misc.py index 4c7e25e..83fbf39 100644 --- a/mapflow/_misc.py +++ b/mapflow/_misc.py @@ -1,6 +1,6 @@ import subprocess import warnings - +import xarray as xr from pyproj import CRS X_NAME_CANDIDATES = ("x", "lon", "longitude") @@ -52,8 +52,6 @@ def process_crs(da, crs): def check_da(da, time_name, x_name, y_name, crs): - import xarray as xr - if not isinstance(da, xr.DataArray): raise TypeError(f"Expected xarray.DataArray, got {type(da)}") for dim in (x_name, y_name, time_name):