aeolus.plot.pv — aeolus 0.4.9+9.g5a71939 documentation

# Source code for aeolus.plot.pv

    """Prepare data in spherical coordinates for pyvista plotting."""
    import numpy as np

    from pyvista import grid_from_sph_coords, transform_vectors_sph_to_cart

    from ..coord import _cell_bounds
    from ..exceptions import _warn
    from ..model import um

    __all__ = ("grid_for_scalar_cube_sph", "grid_for_vector_cubes_sph")

    def grid_for_scalar_cube_sph(cube, z_scale=1, z_offset=0, grid=None, label="scalar3d", model=um):
        """
        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.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        pyvista.StructuredGrid
           PyVista grid with data in cell_arrays.
        """
        if grid is None:
            lons = _cell_bounds(cube.coord(model.x).points)
            lats = 90.0 - _cell_bounds(cube.coord(model.y).points)
            if cube.ndim == 3:
                # TODO check if _cell_bounds should be here
                levels = _cell_bounds(cube.coord(model.z).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}")
        grid.cell_arrays[label] = np.array(cube.data).swapaxes(-2, -1).ravel("C")
        return grid

    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",
        model=um,
    ):
        """
        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.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        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(model.x).points
        lats = 90.0 - u.coord(model.y).points
        levels = z_scale * u.coord(model.z).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