Draw skewed image on skewed axis.

imshow_skewed.py

__all__ = ["imshow_skewed"]

from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.transforms import Affine2D
from mpl_toolkits.axisartist import floating_axes
from mpl_toolkits.axisartist.grid_finder import DictFormatter, FixedLocator


def imshow_skewed(
    data: np.ndarray,
    figsize: Tuple[float, float] = [10, 10],
    aspect: float = 2,
    angle: float = 45,
    fig_angle: float = 0,
    **kwargs,
) -> Tuple[plt.Figure, plt.Axes]:
    assert data.ndim == 2
    fig, ax = plt.subplots(figsize=figsize)
    fig.delaxes(ax)
    transform = Affine2D().scale(1, aspect).skew_deg(0, angle).rotate_deg(fig_angle)
    grid_helper = floating_axes.GridHelperCurveLinear(
        transform,
        extremes=(-0.5, data.shape[1] - 0.5, -0.5, data.shape[0] - 0.5),
        grid_locator1=FixedLocator([]),
        grid_locator2=FixedLocator([]),
        tick_formatter1=DictFormatter({}),
        tick_formatter2=DictFormatter({}),
    )
    _ax = floating_axes.FloatingSubplot(fig, 111, grid_helper=grid_helper)
    fig.add_subplot(_ax)
    ax = _ax.get_aux_axes(transform)
    grid_helper.grid_finder.grid_locator1._nbins = data.shape[0]
    grid_helper.grid_finder.grid_locator2._nbins = data.shape[1]

    ax.imshow(data, **kwargs)
    return fig, ax