Plotting

ERLabPy provides a number of plotting functions to help visualize data and create publication quality figures.

Importing

The key module to plotting is erlab.plotting, which contains various plotting functions. To import it, use the following code:

import matplotlib.pyplot as plt
import erlab.plotting as eplt

First, let us generate some example data from a simple tight binding model of graphene. A rigid shift of 200 meV has been applied so that the Dirac cone is visible.

from erlab.io.exampledata import generate_data

dat = generate_data(bandshift=-0.2, seed=1).T

dat

<xarray.DataArray (eV: 300, ky: 250, kx: 250)> Size: 150MB
3.17 3.71 2.969 4.094 4.071 2.702 ... 0.02378 0.1156 0.1156 0.0235 0.0007165
Coordinates:
  * kx       (kx) float64 2kB -0.89 -0.8829 -0.8757 ... 0.8757 0.8829 0.89
  * ky       (ky) float64 2kB -0.89 -0.8829 -0.8757 ... 0.8757 0.8829 0.89
  * eV       (eV) float64 2kB -0.45 -0.4482 -0.4464 ... 0.08639 0.08819 0.09

xarray.DataArray

• eV: 300
• ky: 250
• kx: 250

• 3.17 3.71 2.969 4.094 4.071 ... 0.02378 0.1156 0.1156 0.0235 0.0007165

  array([[[3.16988000e+00, 3.70981552e+00, 2.96864224e+00, ...,
           4.77441778e+00, 2.68780734e+00, 3.58014595e+00],
          [3.11472181e+00, 4.99662991e+00, 4.33036264e+00, ...,
           4.14316680e+00, 2.51227236e+00, 2.86280882e+00],
          [4.20805285e+00, 5.11327317e+00, 4.61750090e+00, ...,
           3.48265432e+00, 3.11414840e+00, 2.25097081e+00],
          ...,
          [6.78462033e+00, 5.72291886e+00, 4.72469251e+00, ...,
           2.79661715e+00, 3.92313953e+00, 3.66758190e+00],
          [4.35915511e+00, 4.90442132e+00, 3.74273234e+00, ...,
           2.95530884e+00, 5.28886388e+00, 5.24115485e+00],
          [4.33768324e+00, 3.47113544e+00, 2.55882867e+00, ...,
           2.93477333e+00, 5.33625560e+00, 7.91872470e+00]],
  
         [[3.55132863e+00, 3.38371741e+00, 2.84406026e+00, ...,
           3.69613211e+00, 3.93494449e+00, 4.19966501e+00],
          [2.47746428e+00, 3.71616202e+00, 4.81859481e+00, ...,
           3.46830901e+00, 3.92669279e+00, 4.37095141e+00],
          [3.35680809e+00, 3.58224642e+00, 4.64347774e+00, ...,
           2.94637355e+00, 4.10706201e+00, 3.75000791e+00],
  ...
          [4.21180462e-02, 8.10836798e-02, 9.76421093e-02, ...,
           6.93990816e-03, 5.24098590e-03, 2.27311413e-02],
          [5.10658278e-03, 1.91810304e-02, 2.79012144e-02, ...,
           9.28219325e-02, 3.72252576e-02, 9.14126278e-02],
          [7.44505241e-05, 6.53480462e-04, 2.34407666e-02, ...,
           3.70198064e-01, 9.58749828e-02, 2.44888315e-02]],
  
         [[3.53356451e-03, 2.27695745e-01, 9.13149099e-01, ...,
           7.00876446e-02, 2.32873892e-01, 2.78236113e-01],
          [7.21992044e-04, 4.65028451e-02, 1.86489496e-01, ...,
           1.05797035e-01, 4.13523732e-01, 1.49476809e-01],
          [7.21312930e-04, 2.48403636e-02, 9.43059870e-02, ...,
           4.29708001e-02, 1.72646211e-01, 6.79042692e-02],
          ...,
          [2.06386786e-01, 3.89175368e-01, 1.33661244e-01, ...,
           2.32977441e-02, 5.66623178e-04, 3.55179865e-04],
          [2.48401658e-02, 9.19184247e-02, 5.08321046e-02, ...,
           1.09886317e-01, 6.37623966e-03, 1.48810240e-03],
          [3.61495409e-04, 3.18070344e-03, 1.14273325e-01, ...,
           1.15613895e-01, 2.35049226e-02, 7.16525028e-04]]])

• Coordinates: (3)
  • kx
    (kx)
    float64
    -0.89 -0.8829 ... 0.8829 0.89

    array([-0.89    , -0.882851, -0.875703, ...,  0.875703,  0.882851,  0.89    ])

  • ky
    (ky)
    float64
    -0.89 -0.8829 ... 0.8829 0.89

    array([-0.89    , -0.882851, -0.875703, ...,  0.875703,  0.882851,  0.89    ])

  • eV
    (eV)
    float64
    -0.45 -0.4482 ... 0.08819 0.09

    array([-0.45    , -0.448194, -0.446388, ...,  0.086388,  0.088194,  0.09    ])

• Indexes: (3)
  • kx
    PandasIndex

    PandasIndex(Index([              -0.89,   -0.88285140562249,   -0.87570281124498,
           -0.8685542168674699, -0.8614056224899599, -0.8542570281124499,
           -0.8471084337349397, -0.8399598393574297, -0.8328112449799197,
           -0.8256626506024096,
           ...
            0.8256626506024095,  0.8328112449799195,  0.8399598393574296,
            0.8471084337349396,  0.8542570281124499,  0.8614056224899599,
            0.8685542168674699,    0.87570281124498,    0.88285140562249,
                          0.89],
          dtype='float64', name='kx', length=250))

  • ky
    PandasIndex

    PandasIndex(Index([              -0.89,   -0.88285140562249,   -0.87570281124498,
           -0.8685542168674699, -0.8614056224899599, -0.8542570281124499,
           -0.8471084337349397, -0.8399598393574297, -0.8328112449799197,
           -0.8256626506024096,
           ...
            0.8256626506024095,  0.8328112449799195,  0.8399598393574296,
            0.8471084337349396,  0.8542570281124499,  0.8614056224899599,
            0.8685542168674699,    0.87570281124498,    0.88285140562249,
                          0.89],
          dtype='float64', name='ky', length=250))

  • eV
    PandasIndex

    PandasIndex(Index([               -0.45,  -0.4481939799331104, -0.44638795986622076,
           -0.44458193979933114, -0.44277591973244146, -0.44096989966555183,
            -0.4391638795986622,  -0.4373578595317726, -0.43555183946488296,
           -0.43374581939799334,
           ...
            0.07374581939799335,  0.07555183946488303,   0.0773578595317726,
            0.07916387959866228,  0.08096989966555185,  0.08277591973244153,
             0.0845819397993311,  0.08638795986622078,  0.08819397993311046,
                           0.09],
          dtype='float64', name='eV', length=300))

• Attributes: (0)

We can see that the generated data is a three-dimensional xarray.DataArray. Now, let’s extract a cut along \(k_y = 0.3\).

cut = dat.qsel(ky=0.3)
cut

<xarray.DataArray (eV: 300, kx: 250)> Size: 600kB
3.52 2.469 3.126 3.968 2.827 2.746 ... 0.1384 0.05118 0.09941 0.1574 0.5755
Coordinates:
  * eV       (eV) float64 2kB -0.45 -0.4482 -0.4464 ... 0.08639 0.08819 0.09
  * kx       (kx) float64 2kB -0.89 -0.8829 -0.8757 ... 0.8757 0.8829 0.89
    ky       float64 8B 0.2967

xarray.DataArray

• eV: 300
• kx: 250

• 3.52 2.469 3.126 3.968 2.827 ... 0.1384 0.05118 0.09941 0.1574 0.5755

  array([[3.51996103, 2.46860004, 3.12627143, ..., 3.63247043, 3.19189828,
          2.56384799],
         [1.99131562, 2.44728588, 3.13024957, ..., 3.35094604, 2.36664741,
          1.60045821],
         [1.8102282 , 2.42394519, 3.25970425, ..., 2.7023653 , 1.66755731,
          1.63550534],
         ...,
         [0.00580997, 0.09148465, 0.2993771 , ..., 0.01552326, 0.03751128,
          0.01112317],
         [0.02342888, 0.10960098, 0.09828399, ..., 0.03836185, 0.10246751,
          0.13451497],
         [0.02775587, 0.1098784 , 0.02903789, ..., 0.09941059, 0.15744388,
          0.57552933]])

• Coordinates: (3)
  • eV
    (eV)
    float64
    -0.45 -0.4482 ... 0.08819 0.09

    array([-0.45    , -0.448194, -0.446388, ...,  0.086388,  0.088194,  0.09    ])

  • kx
    (kx)
    float64
    -0.89 -0.8829 ... 0.8829 0.89

    array([-0.89    , -0.882851, -0.875703, ...,  0.875703,  0.882851,  0.89    ])

  • ky
    ()
    float64
    0.2967

    array(0.29666667)

• Indexes: (2)
  • eV
    PandasIndex

    PandasIndex(Index([               -0.45,  -0.4481939799331104, -0.44638795986622076,
           -0.44458193979933114, -0.44277591973244146, -0.44096989966555183,
            -0.4391638795986622,  -0.4373578595317726, -0.43555183946488296,
           -0.43374581939799334,
           ...
            0.07374581939799335,  0.07555183946488303,   0.0773578595317726,
            0.07916387959866228,  0.08096989966555185,  0.08277591973244153,
             0.0845819397993311,  0.08638795986622078,  0.08819397993311046,
                           0.09],
          dtype='float64', name='eV', length=300))

  • kx
    PandasIndex

    PandasIndex(Index([              -0.89,   -0.88285140562249,   -0.87570281124498,
           -0.8685542168674699, -0.8614056224899599, -0.8542570281124499,
           -0.8471084337349397, -0.8399598393574297, -0.8328112449799197,
           -0.8256626506024096,
           ...
            0.8256626506024095,  0.8328112449799195,  0.8399598393574296,
            0.8471084337349396,  0.8542570281124499,  0.8614056224899599,
            0.8685542168674699,    0.87570281124498,    0.88285140562249,
                          0.89],
          dtype='float64', name='kx', length=250))

• Attributes: (0)

Plotting 2D data

The fastest way to plot a 2D array like this is to use plot_array. Each axis is automatically labeled.

eplt.plot_array(cut)

<matplotlib.image.AxesImage at 0x7f9e68008650>

plot_array takes many arguments that can customize the look of your plot. The following is an example of some of the functionality provided. For all arguments, see the API reference.

eplt.plot_array(
    cut, cmap="Greys", gamma=0.5, colorbar=True, colorbar_kw=dict(width=10, ticks=[])
)

<matplotlib.image.AxesImage at 0x7f9e5fff3f50>

plot_array can also be accessed (for 2D data) through xarray.DataArray.qplot().

cut.qplot(cmap="Greys", gamma=0.5)

<matplotlib.image.AxesImage at 0x7f9e5f7cd0a0>

Next, let’s add some annotations! The following code adds a line indicating the Fermi level, labels high symmetry points, and adds a colorbar. Here, unlike the previous example, the colorbar was added after plotting. Like this, adding elements separately instead of using keyword arguments can make the code more readable in complex plots.

eplt.plot_array(cut, cmap="Greys", gamma=0.5)

eplt.fermiline()
eplt.mark_points([-0.525, 0.525], ["K", "K"], fontsize=10, pad=(0, 10))
eplt.nice_colorbar(width=10, ticks=[])

<matplotlib.colorbar.Colorbar at 0x7f9e605a0950>

Slices

What if we want to plot multiple slices at once? We should create subplots to place the slices. plt.subplots is very useful in managing multiple axes and figures. If you are unfamiliar with the syntax, visit the relevant matplotlib documentation.

Suppose we want to plot constant energy surfaces at specific binding energies, say, at [-0.4, -0.2, 0.0]. We could create three subplots and iterate over the axes.

energies = [-0.4, -0.2, 0.0]

fig, axs = plt.subplots(1, 3, layout="compressed", sharey=True)
for energy, ax in zip(energies, axs):
    const_energy_surface = dat.qsel(eV=energy)
    eplt.plot_array(const_energy_surface, ax=ax, gamma=0.5, aspect="equal")

Here, we plotted each constant energy surface with plot_array.

To remove the duplicated y axis labels and add some annotations, we can use clean_labels and label_subplot_properties:

fig, axs = plt.subplots(1, 3, layout="compressed", sharey=True)
for energy, ax in zip(energies, axs):
    const_energy_surface = dat.qsel(eV=energy)
    eplt.plot_array(const_energy_surface, ax=ax, gamma=0.5, aspect="equal")

eplt.clean_labels(axs)  # removes shared y labels
eplt.label_subplot_properties(axs, values={"Eb": energies})  # annotates energy

Not bad. However, when it gets to multiple slices along multiple datasets, it gets cumbersome.

Luckily, ERLabPy provides a function that automates the subplot creation, slicing, and annotation for you: plot_slices, which reduces the same code to a one-liner.

fig, axs = eplt.plot_slices([dat], eV=[-0.4, -0.2, 0.0], gamma=0.5, axis="image")

We can also plot the data integrated over an energy window, in this case with a width of 200 meV by adding the eV_width argument:

fig, axs = eplt.plot_slices(
    [dat], eV=[-0.4, -0.2, 0.0], eV_width=0.2, gamma=0.5, axis="image"
)

Cuts along constant \(k_y\) can be plotted analogously.

fig, axs = eplt.plot_slices([dat], ky=[0.0, 0.1, 0.3], gamma=0.5, figsize=(6, 2))

Here, we notice that the first two plots slices through regions with less spectral weight, so the color across the three subplots are not on the same scale. This may be misleading in some occasions where intensity across different slices are important. Luckily, we have a function that can unify the color limits across multiple axes.

Note

The same effect can be achieved by passing on same_limits=True to plot_slices.

fig, axs = eplt.plot_slices([dat], ky=[0.0, 0.1, 0.3], gamma=0.5, figsize=(6, 2))
eplt.unify_clim(axs)

We can also choose a reference axis to get the color limits from.

fig, axs = eplt.plot_slices([dat], ky=[0.0, 0.1, 0.3], gamma=0.5, figsize=(6, 2))
eplt.unify_clim(axs, target=axs.flat[1])

What if we want to plot constant energy surfaces and cuts in the same figure? We can create the subplots first and then utilize the axes argument of plot_slices.

fig, axs = plt.subplots(2, 3, layout="compressed", sharex=True, sharey="row")
eplt.plot_slices([dat], eV=[-0.4, -0.2, 0.0], gamma=0.5, axes=axs[0, :], axis="image")
eplt.plot_slices([dat], ky=[0.0, 0.1, 0.3], gamma=0.5, axes=axs[1, :])
eplt.clean_labels(axs)

2D colormaps

2D colormaps are a method to visualize two data with a single image by mapping one of the data to the lightness of the color and the other to the hue. This is useful when visualizing dichroic or spin-resolved ARPES data[2].

Let us begin with the simulated constant energy contours of Graphene, 0.3 eV below and above the Fermi level.

dat0, dat1 = generate_data(
    shape=(250, 250, 2), Erange=(-0.3, 0.3), temp=0.0, seed=1, count=1e6
).T

_, axs = eplt.plot_slices(
    [dat0, dat1],
    order="F",
    subplot_kw={"layout": "compressed", "sharey": "row"},
    axis="scaled",
    label=True,
)
# eplt.label_subplot_properties(axs, values=dict(Eb=[-0.3, 0.3]))

Suppose we want to visualize the sum and the normalized difference between the two. The simplest way is to plot them side by side.

dat_sum = dat0 + dat1
dat_ndiff = (dat0 - dat1) / dat_sum

eplt.plot_slices(
    [dat_sum, dat_ndiff],
    order="F",
    subplot_kw={"layout": "compressed", "sharey": "row"},
    cmap=["viridis", "bwr"],
    axis="scaled",
)
eplt.proportional_colorbar()

<matplotlib.colorbar.Colorbar at 0x7f9e5ca7ebd0>

The difference array is noisy for small values of the sum. We can plot using a 2D colomap, where dat_ndiff is mapped to the color along the colormap and dat_sum is mapped to the lightness of the colormap.

eplt.plot_array_2d(dat_sum, dat_ndiff)

(<matplotlib.image.AxesImage at 0x7f9e5c93d8b0>,
 <matplotlib.colorbar.Colorbar at 0x7f9e5d888830>)

The color normalization for each axis can be set independently with lnorm and cnorm. The appearance of the colorbar axes can be customized with the returned Colorbar object.

_, cb = eplt.plot_array_2d(
    dat_sum,
    dat_ndiff,
    lnorm=eplt.InversePowerNorm(0.5),
    cnorm=eplt.CenteredInversePowerNorm(0.7, vcenter=0.0, halfrange=1.0),
)
cb.ax.set_xticks(cb.ax.get_xlim())
cb.ax.set_xticklabels(["Min", "Max"])

[Text(0.0, 0, 'Min'), Text(41.33577630506462, 0, 'Max')]

Styling figures

You can control the look and feel of matplotlib figures with style sheets and rcParams. In addition to the options provided by matplotlib, ERLabPy provides some style sheets that are listed below. Note that style sheets that change the default font requires the font to be installed on the system. To see how each one looks, try running the code provided by matplotlib.

Style Name	Description
khan	Personal preferences of the package author.
fira	Changes the default font to Fira Sans.
firalight	Changes the default font to Fira Sans Light.
times	Changes the default font to Times New Roman.
nature	Changes the default font to Arial, and tweaks some aspects such as padding and default figure size.

with plt.style.context(["nature"]):
    eplt.plot_array(cut, cmap="Greys", gamma=0.5)

Tips

• In the python ecosystem, there are some libraries that provide great colormaps, such as cmasher, cmocean, colorcet, and cmcrameri.

• Although matplotlib is a powerful library, it is heavy and slow, and better suited for static plots. For interactive plots, libraries such as Plotly or Bokeh are popular.

  The hvplot library is a high-level plotting library that provides a simple interface to Bokeh, Plotly, and Matplotlib. It is particularly useful for interactive plots and can be used with xarray objects. Here are some examples that uses the Bokeh backend:

import hvplot.xarray

cut.hvplot(x="kx", y="eV", cmap="Greys", aspect=1.5)

dat.hvplot(x="kx", y="ky", cmap="Greys", aspect="equal", widget_location="bottom")

Note

If you are viewing this documentation online, the slider above will not work. To see the interactive plot, you can run the notebook locally after installing hvplot.

For more information, see the hvplot documentation.