Making an `imshow` for Bokeh that roughly does what I expect.

imshow.py

from bokeh.plotting import figure
import inspect

def retrieve_name(var):
    """
    Gets the name of var. Does it from the out most frame inner-wards.
    :param var: variable to get name from.
    :return: string
    """
    for fi in reversed(inspect.stack()):
        names = [var_name for var_name, var_val in fi.frame.f_locals.items() if var_val is var]
        if len(names) > 0:
            return names[0]

def imshow(X, pixel_size=50, palette='Viridis256'):
    """
    Show an array as an image
    :param X: the array to visualise
    :param pixel_size: roughly how big each pixel is
    :param palette: the Bokeh palette to use
    :return: a `bokeh.Figure` object (use `bokeh.io.show` to render)
    """
    p = figure(
        title=f"Array plot of {retrieve_name(X)}", 
        width=X.shape[1] * pixel_size, 
        height=X.shape[0] * pixel_size
    )
    p.image(
        image=[X], 
        x=[0], y=[0], 
        dw=[X.shape[1]], dh=[X.shape[0]], 
        palette=palette
    )
    return p

Example usage

import jax.numpy as np
from jax import random

from mdmd.core import compute_dmd

key = random.key(314159)

length = 20
b = np.array(random.uniform(key, shape=[length]), dtype=np.float32)
lam = random.uniform(key, shape=[length])

# Build snapshot matrix of shape (n, t)
n, t, T = b.shape[0], 20, 2
X = np.column_stack([(lam ** k) * (1 + np.sin(k / T)) * b for k in range(t)])

# Show a plot of X
show(imshow(X, pixel_size=25))

gives this output