Source code for aeolus.plot.pv

"""Prepare data in spherical coordinates for pyvista plotting."""
from warnings import warn

import numpy as np

from pyvista import grid_from_sph_coords, transform_vectors_sph_to_cart

from ..exceptions import AeolusWarning
from ..grid import _cell_bounds


__all__ = (
    "grid_for_scalar_cube_sph",
    "grid_for_vector_cubes_sph",
)


[docs]def grid_for_scalar_cube_sph(cube, z_scale=1, z_offset=0, grid=None, label="scalar3d"): """ Create a `pyvista` grid for an `iris` cube (2D or 3D) in spherical coordinates. Parameters ---------- cube: iris.cube.Cube 2D or 3D cube with (longitude, latitude, [level_height]) coordinates. z_scale: float, optional Scaling factor for the vertical level dimension. z_offset: float, optional Scaling offset for the vertical level dimension. grid: pyvista.StructuredGrid, optional If given, add data to the existing grid, otherwise create a new one from the input cube's coordinates. label: str, optional Label for the data within the grid. Returns ------- pyvista.StructuredGrid PyVista grid with data in cell_arrays. """ if grid is None: lons = _cell_bounds(cube.coord("longitude").points) lats = 90.0 - _cell_bounds(cube.coord("latitude").points) if cube.ndim == 3: # TODO check if _cell_bounds should be here levels = _cell_bounds(cube.coord("level_height").points) else: levels = np.array([0]) levels = z_scale * levels + z_offset grid = grid_from_sph_coords(lons, lats, levels) # Add data arras to grid if label in grid.cell_arrays: warn(f"Label '{label}' exists in {grid}", AeolusWarning) grid.cell_arrays[label] = np.array(cube.data).swapaxes(-2, -1).ravel("C") return grid
[docs]def grid_for_vector_cubes_sph( u, v, w, vector_scale=1, vertical_wind_scale=1, z_scale=1, z_offset=0, xstride=1, ystride=1, grid=None, label="vector3d", ): """ Take wind vectors in spherical coordinates and create a `pyvista` grid for them. Parameters ---------- u: iris.cube.Cube 2D or 3D cube of x-wind component (zonal wind). v: iris.cube.Cube 2D or 3D cube of y-wind component (meridional wind). w: iris.cube.Cube 2D or 3D cube of z-wind component (vertical wind). vector_scale: float, optional Scaling factor for vectors. vertical_wind_scale: float, optional Scaling factor for the vertical wind component (for better visibility). z_scale: float, optional Scaling factor for the vertical level dimension. z_offset: float, optional Scaling offset for the vertical level dimension. xstride: float, optional Stride along the longitude axis. ystride: float, optional Stride along the latitude axis. grid: pyvista.StructuredGrid, optional If given, add data to the existing grid, otherwise create a new one from the input cube's coordinates. label: str, optional Label for the data within the grid. Returns ------- pyvista.StructuredGrid PyVista grid with vector data in point_arrays. """ assert (u.shape == v.shape) and ( u.shape == w.shape ), "Wind components should have the same array size!" lons = u.coord("longitude").points lats = 90.0 - u.coord("latitude").points levels = z_scale * u.coord("level_height").points + z_offset # Stride vectors lons_s = lons[::xstride] lats_s = lats[::ystride] inv_axes = [*range(u.ndim - 1, -1, -1)] u_sdata = u.data[..., ::ystride, ::xstride].transpose(inv_axes) v_sdata = v.data[..., ::ystride, ::xstride].transpose(inv_axes) w_sdata = w.data[..., ::ystride, ::xstride].transpose(inv_axes) # Rescale vertical wind for better viz w_sdata *= vertical_wind_scale # Create a stack of vectors transformed from spherical to cartesian coordinates vectors = np.stack( [ i.transpose(inv_axes).swapaxes(u.ndim - 2, u.ndim - 1).ravel("C") for i in transform_vectors_sph_to_cart( lons_s, lats_s, levels, u_sdata, -v_sdata, w_sdata ) ], axis=1, ) # Scale vectors vectors *= vector_scale if grid is None: grid = grid_from_sph_coords(lons_s, lats_s, levels) # Add vectors to the grid grid.point_arrays[label] = vectors return grid