Utils

ginjax.utils

setup_plot(figsize: tuple[int, int] = (8, 6)) -> matplotlib.axes.Axes

Create a figure of figsize and return the axes.

Parameters:

Name	Type	Description	Default
figsize	tuple[int, int]	the specified figure size, (width,height)	(8, 6)

Returns:

Type	Description
Axes	the new axes

Source code in ginjax/utils.py

def setup_plot(figsize: tuple[int, int] = (8, 6)) -> matplotlib.axes.Axes:
    """
    Create a figure of figsize and return the axes.

    args:
        figsize: the specified figure size, (width,height)

    returns:
        the new axes
    """
    return plt.figure(figsize=figsize).gca()

nobox(ax: matplotlib.axes.Axes) -> None

Turn axes and xticks,yticks off.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis to edit	required

Source code in ginjax/utils.py

def nobox(ax: matplotlib.axes.Axes) -> None:
    """
    Turn axes and xticks,yticks off.

    args:
        ax: the axis to edit
    """
    ax.set_xticks([])
    ax.set_yticks([])
    ax.axis("off")

finish_plot(ax: matplotlib.axes.Axes, title: str, xs: np.ndarray, ys: np.ndarray, D: int) -> None

Set title, axis limits, equal aspect ratio, and no box.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis to edit	required
title	str	title of the plot	required
xs	ndarray	x dimensions	required
ys	ndarray	y dimensions	required
D	int	dimension of the space	required

Source code in ginjax/utils.py

def finish_plot(
    ax: matplotlib.axes.Axes, title: str, xs: np.ndarray, ys: np.ndarray, D: int
) -> None:
    """
    Set title, axis limits, equal aspect ratio, and no box.

    args:
        ax: the axis to edit
        title: title of the plot
        xs: x dimensions
        ys: y dimensions
        D: dimension of the space
    """
    ax.set_title(title)
    if D == 2:
        ax.set_xlim(np.min(xs) - 0.55, np.max(xs) + 0.55)
        ax.set_ylim(np.min(ys) - 0.55, np.max(ys) + 0.55)
    if D == 3:
        ax.set_xlim(np.min(xs) - 0.75, np.max(xs) + 0.75)
        ax.set_ylim(np.min(ys) - 0.75, np.max(ys) + 0.75)
    ax.set_aspect("equal")
    nobox(ax)

plot_boxes(ax: matplotlib.axes.Axes, xs: np.ndarray, ys: np.ndarray) -> None

Plot the boxes for a geometric image.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis to plot on	required
xs	ndarray	pixels in the x direction	required
ys	ndarray	pixels in the y direction	required

Source code in ginjax/utils.py

def plot_boxes(ax: matplotlib.axes.Axes, xs: np.ndarray, ys: np.ndarray) -> None:
    """
    Plot the boxes for a geometric image.

    args:
        ax: the axis to plot on
        xs: pixels in the x direction
        ys: pixels in the y direction
    """
    ax.plot(
        xs[None] + np.array([-0.5, -0.5, 0.5, 0.5, -0.5]).reshape((5, 1)),
        ys[None] + np.array([-0.5, 0.5, 0.5, -0.5, -0.5]).reshape((5, 1)),
        "k-",
        lw=0.5,
        zorder=10,
    )

fill_boxes(ax: matplotlib.axes.Axes, xs: np.ndarray, ys: np.ndarray, ws: np.ndarray, vmin: float, vmax: float, cmap: Union[matplotlib.colors.Colormap, str], zorder: int = -100, colorbar: bool = False, alpha: float = 1.0) -> None

Fill boxes with color according to the ws values.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis we are plotting on	required
xs	ndarray	the x coordinates of the pixels	required
ys	ndarray	the y coordinates of the pixels	required
ws	ndarray	the values to determine the fill color	required
vmin	float	min value used for color	required
vmax	float	max value used for color	required
cmap	Union[Colormap, str]	the color map	required
zorder	int	whether to put the color in front or behind other elements	-100
colorbar	bool	whether to include the colorbar	False
alpha	float	the opacity of the box fill	1.0

Source code in ginjax/utils.py

def fill_boxes(
    ax: matplotlib.axes.Axes,
    xs: np.ndarray,
    ys: np.ndarray,
    ws: np.ndarray,
    vmin: float,
    vmax: float,
    cmap: Union[matplotlib.colors.Colormap, str],
    zorder: int = -100,
    colorbar: bool = False,
    alpha: float = 1.0,
) -> None:
    """
    Fill boxes with color according to the ws values.

    args:
        ax: the axis we are plotting on
        xs: the x coordinates of the pixels
        ys: the y coordinates of the pixels
        ws: the values to determine the fill color
        vmin: min value used for color
        vmax: max value used for color
        cmap: the color map
        zorder: whether to put the color in front or behind other elements
        colorbar: whether to include the colorbar
        alpha: the opacity of the box fill
    """
    plotted_img = ax.imshow(
        ws.reshape((np.max(xs) + 1, np.max(ys) + 1)).T,
        vmin=vmin,
        vmax=vmax,
        cmap=cmap,
        alpha=alpha,
        zorder=zorder,
    )
    if colorbar:
        plt.colorbar(plotted_img, ax=ax)

plot_scalars(ax: matplotlib.axes.Axes, spatial_dims: tuple[int, ...], xs: np.ndarray, ys: np.ndarray, ws: np.ndarray, boxes: bool = True, fill: bool = True, symbols: bool = True, vmin: float = -2.0, vmax: float = 2.0, cmap: matplotlib.colors.Colormap | str | None = 'BrBG', colorbar: bool = False) -> None

Plot scalars from a scalar image.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis we are plotting on	required
spatial_dims	tuple[int, ...]	the spatial dimensions of the image	required
xs	ndarray	the x coordinates of the pixels	required
ys	ndarray	the y coordinates of the pixels	required
ws	ndarray	the scalar pixel values	required
boxes	bool	whether to plot the surrounding boxes	True
fill	bool	whether to plot the color fill	True
symbols	bool	whether to plot symbol representation of the scalar values	True
vmin	float	min value used for color	-2.0
vmax	float	max value used for color	2.0
cmap	Colormap | str | None	the color map	'BrBG'
colorbar	bool	whether to include the colorbar	False

Source code in ginjax/utils.py

def plot_scalars(
    ax: matplotlib.axes.Axes,
    spatial_dims: tuple[int, ...],
    xs: np.ndarray,
    ys: np.ndarray,
    ws: np.ndarray,
    boxes: bool = True,
    fill: bool = True,
    symbols: bool = True,
    vmin: float = -2.0,
    vmax: float = 2.0,
    cmap: matplotlib.colors.Colormap | str | None = "BrBG",
    colorbar: bool = False,
) -> None:
    """
    Plot scalars from a scalar image.

    args:
        ax: the axis we are plotting on
        spatial_dims: the spatial dimensions of the image
        xs: the x coordinates of the pixels
        ys: the y coordinates of the pixels
        ws: the scalar pixel values
        boxes: whether to plot the surrounding boxes
        fill: whether to plot the color fill
        symbols: whether to plot symbol representation of the scalar values
        vmin: min value used for color
        vmax: max value used for color
        cmap: the color map
        colorbar: whether to include the colorbar
    """
    if boxes:
        plot_boxes(ax, xs, ys)
    if fill:
        if cmap is None:
            cmap = "BrBG"

        fill_boxes(ax, xs, ys, ws, vmin, vmax, cmap, colorbar=colorbar)
    if symbols:
        height = ax.get_window_extent().height
        ss = (5 * height / spatial_dims[0]) * np.abs(ws)
        ax.scatter(xs[ws > TINY], ys[ws > TINY], marker="+", c="k", s=ss[ws > TINY], zorder=100)
        ax.scatter(
            xs[ws < -TINY],
            ys[ws < -TINY],
            marker="_",
            c="k",
            s=ss[ws < -TINY],
            zorder=100,
        )

plot_vectors(ax: matplotlib.axes.Axes, xs: np.ndarray, ys: np.ndarray, ws: np.ndarray, boxes: bool = True, fill: bool = True, vmin: float = 0.0, vmax: float = 2.0, cmap: matplotlib.colors.Colormap | str | None = 'PuRd', scaling: float = 0.33) -> None

Plot vectors from a vector image.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis we are plotting on	required
xs	ndarray	the x coordinates of the pixels	required
ys	ndarray	the y coordinates of the pixels	required
ws	ndarray	the scalar pixel values	required
boxes	bool	whether to plot the surrounding boxes	True
fill	bool	whether to plot the color fill	True
vmin	float	min value used for color	0.0
vmax	float	max value used for color	2.0
cmap	Colormap | str | None	the color map	'PuRd'
scaling	float	how much to scale the vectors	0.33

Source code in ginjax/utils.py

def plot_vectors(
    ax: matplotlib.axes.Axes,
    xs: np.ndarray,
    ys: np.ndarray,
    ws: np.ndarray,
    boxes: bool = True,
    fill: bool = True,
    vmin: float = 0.0,
    vmax: float = 2.0,
    cmap: matplotlib.colors.Colormap | str | None = "PuRd",
    scaling: float = 0.33,
) -> None:
    """
    Plot vectors from a vector image.

    args:
        ax: the axis we are plotting on
        xs: the x coordinates of the pixels
        ys: the y coordinates of the pixels
        ws: the scalar pixel values
        boxes: whether to plot the surrounding boxes
        fill: whether to plot the color fill
        vmin: min value used for color
        vmax: max value used for color
        cmap: the color map
        scaling: how much to scale the vectors
    """
    if boxes:
        plot_boxes(ax, xs, ys)
    if fill:
        if cmap is None:
            cmap = "PuRd"

        fill_boxes(ax, xs, ys, np.linalg.norm(ws, axis=-1), vmin, vmax, cmap, alpha=0.25)

    normws = np.linalg.norm(ws, axis=1)

    xs = xs[normws > TINY]
    ys = ys[normws > TINY]
    ws = ws[normws > TINY]

    for x, y, w, normw in zip(xs, ys, ws, normws[normws > TINY]):
        ax.arrow(
            x - scaling * w[0],
            y - scaling * w[1],
            2 * scaling * w[0],
            2 * scaling * w[1],
            length_includes_head=True,
            head_width=0.24 * scaling * normw,
            head_length=0.72 * scaling * normw,
            color="k",
            zorder=100,
        )

plot_tensors(ax: matplotlib.axes.Axes, xs: np.ndarray, ys: np.ndarray, Ts: np.ndarray, boxes: bool = True, arrow_scaling: float = 0.33) -> None

Plot a tensor image.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis to plot on	required
xs	ndarray	the x coordinates of the pixels	required
ys	ndarray	the y coordinates of the pixels	required
ws		the tensor values at the pixels	required
boxes	bool	whether to also plot the boxes	True

Source code in ginjax/utils.py

def plot_tensors(
    ax: matplotlib.axes.Axes,
    xs: np.ndarray,
    ys: np.ndarray,
    Ts: np.ndarray,
    boxes: bool = True,
    arrow_scaling: float = 0.33,
) -> None:
    """
    Plot a tensor image.

    args:
        ax: the axis to plot on
        xs: the x coordinates of the pixels
        ys: the y coordinates of the pixels
        ws: the tensor values at the pixels
        boxes: whether to also plot the boxes
    """
    if boxes:
        plot_boxes(ax, xs, ys)

    # untested for D=3
    T_shape = Ts.shape[1:]

    if T_shape != (2, 2):
        print(f"plot_tensors: Can only plot D=2 tensor images, got {T_shape} tensor.")
        return  # don't want plotting ever to break

    Ts_trace = np.einsum("...ii", Ts) / 2
    Ts_antisym = ((Ts - np.transpose(Ts, (0, 2, 1))) / 2)[:, 0, 1]
    Ts_sym = ((Ts + np.transpose(Ts, (0, 2, 1))) / 2) - Ts_trace[:, None, None] * np.eye(
        T_shape[0]
    )[None]

    # -3,3 is the scalar filters default
    fill_boxes(ax, xs, ys, Ts_trace, -3.0, 3.0, cmap="BrBG")

    patches = []
    for x, y, T_trace, T_antisym, T_sym in zip(xs, ys, Ts_trace, Ts_antisym, Ts_sym):
        if np.abs(T_antisym) > TINY:
            patches.append(
                Circle(
                    (x, y),
                    radius=0.25,
                    color=matplotlib.colormaps["bwr"](T_antisym),
                    zorder=0,
                    alpha=0.25,
                )
            )

        arrow_val = T_sym[:, 0]
        arrow_norm = np.linalg.norm(arrow_val)
        if arrow_norm > TINY:
            patches.append(
                Arrow(
                    x - arrow_scaling * arrow_val[0],
                    y - arrow_scaling * arrow_val[1],
                    2 * arrow_scaling * arrow_val[0],
                    2 * arrow_scaling * arrow_val[1],
                    width=float(0.5 * arrow_norm),
                    color="k",
                    zorder=100,
                )
            )

    ax.add_collection(PatchCollection(patches, match_original=True))

plot_nothing(ax: matplotlib.axes.Axes) -> None

Set the title to empty string and print no boxes.

Parameters:

Name	Type	Description	Default
ax	Axes	the axis to plot nothing on	required

Source code in ginjax/utils.py

def plot_nothing(ax: matplotlib.axes.Axes) -> None:
    """
    Set the title to empty string and print no boxes.

    args:
        ax: the axis to plot nothing on
    """
    ax.set_title(" ")
    nobox(ax)

plot_grid(images: list, names: list[str], n_cols: int, **kwargs: bool) -> matplotlib.figure.Figure

Plot a grid of GeometricImages. The grid will have columns equal to n_cols, and the number of rows are calculated automatically. If there are extra spots in the grid, plot_nothing is called on those spaces.

Parameters:

Name	Type	Description	Default
images	list	images to plot in the grid	required
names	list[str]	names of each image to plot as the title	required
n_cols	int	number of columns in the gride	required
kwargs	bool	keyword arguments passed along to plot	{}

Returns:

Type	Description
Figure	the figure plotted

Source code in ginjax/utils.py

def plot_grid(
    images: list, names: list[str], n_cols: int, **kwargs: bool
) -> matplotlib.figure.Figure:
    """
    Plot a grid of GeometricImages. The grid will have columns equal to n_cols, and the number of
    rows are calculated automatically. If there are extra spots in the grid, plot_nothing is called
    on those spaces.

    args:
        images: images to plot in the grid
        names: names of each image to plot as the title
        n_cols: number of columns in the gride
        kwargs: keyword arguments passed along to plot

    returns:
        the figure plotted
    """
    n_rows = max(1, np.ceil(len(images) / n_cols).astype(int))
    assert len(images) <= n_cols * n_rows
    bar = 8.0  # figure width in inches?
    fig, axes = plt.subplots(
        n_rows,
        n_cols,
        figsize=(bar, 1.15 * bar * n_rows / n_cols),  # magic
        squeeze=False,
    )
    axes = axes.flatten()
    plt.subplots_adjust(
        left=0.001 / n_cols,
        right=1 - 0.001 / n_cols,
        wspace=0.2 / n_cols,
        bottom=0.001 / n_rows,
        top=1 - 0.001 / n_rows - 0.1 / n_rows,
        hspace=0.2 / n_rows,
    )

    for img, name, axis in zip(images, names, axes):
        img.plot(ax=axis, title=name, **kwargs)

    for axis in axes[len(images) :]:
        plot_nothing(axis)

    return fig

power(img: jax.Array) -> tuple[jax.Array, jax.Array, jax.Array]

Compute the power of image From: https://bertvandenbroucke.netlify.app/2019/05/24/computing-a-power-spectrum-in-python/

Parameters:

Name	Type	Description	Default
img	Array	scalar image data, shape (batch,channel,spatial)	required

Returns:

Type	Description
tuple[Array, Array, Array]	tuple of k, P, N

Source code in ginjax/utils.py

def power(img: jax.Array) -> tuple[jax.Array, jax.Array, jax.Array]:
    """
    Compute the power of image
    From: https://bertvandenbroucke.netlify.app/2019/05/24/computing-a-power-spectrum-in-python/

    args:
        img: scalar image data, shape (batch,channel,spatial)

    returns:
        tuple of k, P, N
    """
    # images are assumed to be scalar images
    spatial_dims = img.shape[2:]

    # some assumptions about the image
    assert len(spatial_dims) == 2  # assert the D=2
    assert spatial_dims[0] == spatial_dims[1]  # the image is square

    kmax = min(s for s in spatial_dims) // 2
    even = spatial_dims[0] % 2 == 0  # are the image sides even?

    img = jnp.fft.fftn(img, s=spatial_dims)  # fourier transform
    P = img.real**2 + img.imag**2
    P = jnp.sum(
        jnp.mean(P, axis=0), axis=0
    )  # mean over batch, then sum over channels. Shape (spatial,)

    kfreq = jnp.fft.fftfreq(spatial_dims[0]) * spatial_dims[0]
    kfreq2D = jnp.meshgrid(kfreq, kfreq)
    k = jnp.linalg.norm(jnp.stack(kfreq2D, axis=0), axis=0)

    N = np.full_like(P, 2, dtype=jnp.int32)
    N[..., 0] = 1
    if even:
        N[..., -1] = 1

    k = k.flatten()
    P = P.flatten()
    N = N.flatten()

    kbin = jnp.ceil(k).astype(jnp.int32)
    k = jnp.bincount(kbin, weights=k * N)
    P = jnp.bincount(kbin, weights=P * N)
    N = jnp.bincount(kbin, weights=N).round().astype(jnp.int32)

    # drop k=0 mode and cut at kmax (smallest Nyquist)
    k = k[1 : 1 + kmax]
    P = P[1 : 1 + kmax]
    N = N[1 : 1 + kmax]

    k /= N
    P /= N

    return k, P, N

plot_power(fields: list[jax.Array], labels: Optional[list[str]], ax: matplotlib.axes.Axes, title: str = '', xlabel: str = 'unnormalized wavenumber', ylabel: str = 'unnormalized power', hide_ticks: bool = False) -> None

Plot the power spectrum of each image onto the same plot.

Parameters:

Name	Type	Description	Default
fields	list[Array]	list of fields to plot the power spectrum, shape (batch,channel,spatial)	required
labels	Optional[list[str]]	label for each field	required
ax	Axes	the axis to plot on	required
title	str	title of the plot	''
xlabel	str	the x axis label	'unnormalized wavenumber'
ylabel	str	the y axis label	'unnormalized power'
hide_ticks	bool	whether to hide the ticks	False

Source code in ginjax/utils.py

def plot_power(
    fields: list[jax.Array],
    labels: Optional[list[str]],
    ax: matplotlib.axes.Axes,
    title: str = "",
    xlabel: str = "unnormalized wavenumber",
    ylabel: str = "unnormalized power",
    hide_ticks: bool = False,
) -> None:
    """
    Plot the power spectrum of each image onto the same plot.

    args:
        fields: list of fields to plot the power spectrum, shape (batch,channel,spatial)
        labels: label for each field
        ax: the axis to plot on
        title: title of the plot
        xlabel: the x axis label
        ylabel: the y axis label
        hide_ticks: whether to hide the ticks
    """

    ks, Ps = [], []
    for field in fields:
        k, P, _ = power(field)
        ks.append(k)
        Ps.append(P)

    used_labels = labels if labels else [""] * len(ks)
    for k, P, l in zip(ks, Ps, used_labels):
        ax.loglog(k, P, label=l, alpha=0.7)

    if labels:
        ax.legend(fontsize=36)

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)

    if hide_ticks:
        ax.tick_params(
            axis="both",
            which="both",
            bottom=False,
            left=False,
            labelbottom=False,
            labelleft=False,
        )