Plots a scalar 4d field using 2d slices.

plot_field.py

#
# Copyright(c) Daniel Knüttel 2024
#

#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

import matplotlib.pyplot as plt
import numpy as np

def get_slice(axis_A, axis_B, axis_C, axis_D, cut_C, cut_D):
    slc = [None, None, None, None]
    slc[axis_A] = slice(None, None, None)
    slc[axis_B] = slice(None, None, None)
    slc[axis_C] = cut_C
    slc[axis_D] = cut_D
    return tuple(slc)


def plot_scalar_field(field_np, axis_A=0, axis_B=1, axis_C=2, axis_D=3
                      , ax=None, cmap="afmhot", fig=None, add_hist=False
                      , use_abs=False, tiny_axis_label_color="black"
                      , vmin=None, vmax=None, omit_colorbar=False
                      , grid_color=None):
    """
    Plots a scalar 4d (``x,y,z,t``) field using ``matplotlib.pyplot.imshow``. First publicly used in [1]_.
    Cite as [2]_.

    The field is sliced along three of the coordinates and these slices are put next to each other.
    See the ascii art below as an explanation.

    The parameters ``axis_A``, ``axis_B``, ``axis_C``, ``axis_D`` control the ordering of the slicing.
    ``ax`` and ``fig`` can be used to plot using the given axis and figure.
    If ``add_hist is True``, a histogram is added to the plot.
    ``cmap`` controls the color map; ``vmin`` and ``vmax`` control the scaling.
    ``grid_color`` and ``tiny_axis_label`` can be used to adjust the colors that are used
    inside the overlay.

    Slicing visualization using a 3d cube::

                     +======================================+  < i = 0
                    /::::::::::::::::::::::::::::::::::::::/|
                   /::::::::::::::::::::::::::::::::::::::/ +  < i = 1
                  /::::::::::::::::::::::::::::::::::::::/ /|
                 /::::::::::::::::::::::::::::::::::::::/ / +  < ...
                /::::::::::::::::::::::::::::::::::::::/ / /|
               /::::::::::::::::::::::::::::::::::::::/ / / +
              /::::::::::::::::::::::::::::::::::::::/ / / /|
             /::::::::::::::::::::::::::::::::::::::/ / / / +
            +======================================+ / / / /|
            |                                      |/ / / / +
            +--------------------------------------+ / / / /|
            |                                      |/ / / / +                               
            +--------------------------------------+ / / / /|
            |                                      |/ / / / +
            +--------------------------------------+ / / / /|
            |                                      |/ / / / +
            +--------------------------------------+ / / / /|
            |                                      |/ / / / +
            +--------------------------------------+ / / / /
            |                                      |/ / / /
            +--------------------------------------+ / / /
            |                                      |/ / /
            +--------------------------------------+ / /
            |                                      |/ /
            +--------------------------------------+ /
            |                                      |/
            +======================================+                           
            ^^^
            |||
            ||...
            |j = 1
            j = 0

            
             i = 0                                  | i = 1                                  | ...
                                                    |                                        |                                        
            +--------------------------------------+|+--------------------------------------+|+--------------------------------------+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +......................................+|+......................................+|+......................................+
            +--------------------------------------+|+--------------------------------------+|+--------------------------------------+
            ^^^                                      ^^^                                      ^^^   
            |||                                      |||                                      |||
            ||...                                    ||...                                    ||...
            |j = 1                                   |j = 1                                   |j = 1
            j = 0                                    j = 0                                    j = 0
    
    
        

    
    .. [1]: 10.22323/1.453.0037
    .. [2]: 10.5281/zenodo.14051252
    """

    if ax is None:
        ax = plt.gca()
    if fig is None:
        fig = plt.gcf()
    
    shape = field_np.shape
    plotable = np.copy(field_np).reshape(( shape[axis_A]*shape[axis_C], shape[axis_B]*shape[axis_D]))
    
    axes_labels = ["x", "y", "z", "t"]
    
    for i in range(shape[axis_C]):
        for j in range(shape[axis_D]):
            plotable[i*shape[axis_A]: (i+1)*shape[axis_A]
                    , j*shape[axis_B]: (j+1)*shape[axis_B]] \
                            = field_np[get_slice(axis_A, axis_B, axis_C, axis_D, i, j)]

    # See https://stackoverflow.com/questions/23090791/matplotlib-colorbar-not-working-due-to-garbage-collection
    if not use_abs:
        mappable = ax.imshow(plotable.real, cmap=cmap, vmin=vmin, vmax=vmax)
    else:
        mappable = ax.imshow(np.abs(plotable), cmap=cmap, vmin=vmin, vmax=vmax)
    if not omit_colorbar:
        fig.colorbar(mappable, ax=ax, shrink=0.6)
    
    ax.set_yticks(np.arange(0, shape[axis_A]*shape[axis_C], shape[axis_A]) - 0.5, np.arange(0, shape[axis_C], 1)) 
    ax.set_xticks(np.arange(0, shape[axis_B]*shape[axis_D], shape[axis_B]) - 0.5, np.arange(0, shape[axis_D], 1)) 
    ax.set_xlabel(f"${axes_labels[axis_D]}$")
    ax.set_ylabel(f"${axes_labels[axis_C]}$")
    
    ax.arrow(0.5, 0.5, 0, shape[axis_B] // 2, width=0.3, edgecolor=tiny_axis_label_color, facecolor=tiny_axis_label_color)
    ax.arrow(0.5, 0.5, shape[axis_A] // 2, 0, width=0.3, edgecolor=tiny_axis_label_color, facecolor=tiny_axis_label_color)
    ax.text(0.5 + shape[axis_B] // 2 + 2, 1.5, f"${axes_labels[axis_B]}$", color=tiny_axis_label_color)
    ax.text(1.5, 0.5 + shape[axis_A] // 2 + 2, f"${axes_labels[axis_A]}$", color=tiny_axis_label_color)

    if grid_color is None:
        ax.grid(visible=True, axis="both")
    else:
        ax.grid(visible=True, axis="both", color=grid_color)
    
    if add_hist:
        def mul(itbl):
            res = 1
            for i in itbl:
                res *= i
            return res

        inset = ax.inset_axes([0.47, 0.47, 0.47, 0.47])
        inset.patch.set_alpha(0.01)
        inset.grid(False)
        inset.hist(plotable.reshape(mul(shape)).real, density=True, histtype="step")