[1]:
import matplotlib.pyplot as plt
import numpy as np

%config InlineBackend.figure_formats = ['svg']

Plotting with xarray

Overall architecture

We use the xarray package, which provides some convenience commands to create plots. In this notebook, we use matplotlib as the plotting back end.

As our plots need to get more sophisticated and tuned, we start using more and more matplotlib commands. We show towards the end of the section how to draw the xarray data with matplotlib using only matplotlib commands for the plots. At the end, we make a reference to the HoloViews tutorial, which is especially suited for multidimensional and dynamic plots.

Getting the xarray

For now, we load a particular data file.

[2]:
from postopus import Run
[3]:
cd ../../tests/data/benzene/
/builds/octopus-code/postopus/tests/data/benzene
/usr/local/lib/python3.9/dist-packages/IPython/core/magics/osm.py:417: UserWarning: using dhist requires you to install the `pickleshare` library.
  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]
[4]:
!octopus > out_gs.log 2>&1
[5]:
run = Run(".")

The data set contains the electron density of the Benzene molecule, sampled on a 3d grid.

To make processing and plotting easier, we ask Postopus for the xarray representation of the data:

[6]:
density = run.default.scf.density
xa = density.get_converged(source="cube")

Calling density.get() without passing a parameter for step, automatically the converged field from the “static” folder is loaded. xarray objects contain the data values, but also the coordinates for which these values are defined:

[7]:
xa.coords
[7]:
Coordinates:
  * step     (step) int64 1
  * x        (x) float64 -13.04 -12.47 -11.91 -11.34 ... 11.34 11.91 12.47 13.04
  * y        (y) float64 -13.61 -13.04 -12.47 -11.91 ... 11.91 12.47 13.04 13.61
  * z        (z) float64 -9.071 -8.504 -7.937 -7.37 ... 7.37 7.937 8.504 9.071
[8]:
xa.values.shape
[8]:
(1, 47, 49, 33)

Because the xarray object knows the coordinates (and the names of the coordinates, such as x, y and z), it carries the metadata that is required to search and plot the data. This can make the plotting of the data easier. It also helps selecting subsets of the data (which may be plotted subsequently).

Here are some examples.

Selecting a slice by index

We know we have 33 sampling points in the z-direction. We expect the benzene molecule to be in the middle of the sampled region, so we can ask for index \(i_z = 33/2 \approx 16\) to retrieve density data points in the x-y plane where the index for z is 16:

[9]:
s0 = xa.isel(z=16)

s0 is another xarray object. (It provides a view on the xarray, so that it does not cost additional memory to create it.)

Creating a plot quickly (.plot())

The quickest way to visualise any xarray object is to call the plot method on it:

[10]:
s0.plot()
[10]:
<matplotlib.collections.QuadMesh at 0x7fbd4819b9d0>
../_images/notebooks_xarray-plots1_17_1.svg

We get a visualisation of the data that has correct coordinate labels in the x and y direction, and the values of the z-coordinate at which we sample the data are shown as a title. The colourbar shows the colourscale for the density.

[TODO: The units of the density will be shown in the future]

There are things one may wish to improve in this plot, for example having an equal aspect ratio (so that a circle in the data would look like a circle when plotted, not like an ellipse). We will see below how this can be achieved.

Here, the lesson is: A quick visualisation of the data can be achieved using the plot method on an xarray object. (Note: the output of the plot method will depend on the dimension of the data.)

Inverting the order of x and y

As we can see from the image above the y-axis is presented on the horizontal axis and the x-axis is presented on the vertical axis. This is the xarray plotting standard, which uses the second dimension for the horizontal axis. In natural sciences, we commonly use the x dimension (the first one) as the horizontal dimension and y (the second dimension) as the vertical one. If we want to follow this convention, we need to specify the x-axis explicitly with plot(x="x"):

[11]:
s0.plot(x="x")
[11]:
<matplotlib.collections.QuadMesh at 0x7fbd25ee1df0>
../_images/notebooks_xarray-plots1_20_1.svg

For simplicity, we will not add the (x="x") to every plot of this tutorial, but it can be added where required. We will keep xarray’s standard and use plot() as it is.

Selecting a slice by coordinate

In the example above, we have chosen index 33 in the xa.isel(z=33) command to access the plane at \(z\approx 0\).

As xarray knows the coordinate values that are attached to our data points, we can also ask it to find the plane for which z=0. As the data point that is meant to represent \(z=0\) is actually located at \(z=1.588\cdot 10^{-6}\), we need to tell xarray that we will accept the nearest data point to \(z=0\):

[12]:
xa = xa.squeeze()  # only one step, squeeze this dimension.
[13]:
s1 = xa.sel(z=0, method="nearest")
[14]:
s1.plot()
[14]:
<matplotlib.collections.QuadMesh at 0x7fbd25dc8eb0>
../_images/notebooks_xarray-plots1_25_1.svg

Fine-tuning: no grid lines

To get rid of the grid lines in the plot, we can use plot.imshow() instead of plot() on the xarray object:

[15]:
s1.plot.imshow()
[15]:
<matplotlib.image.AxesImage at 0x7fbd25ea4af0>
../_images/notebooks_xarray-plots1_27_1.svg

[The grid lines only seem to appear with the svg plotting backend.]

As a second example, let’s slice the data in the \(x\approx 0\) plane:

[16]:
s2 = xa.sel(x=0, method="nearest")
s2.plot.imshow()
[16]:
<matplotlib.image.AxesImage at 0x7fbd25de22b0>
../_images/notebooks_xarray-plots1_29_1.svg

Fine-tuning: equal aspect ratio

To draw each pixel as a square, we need to (in matplotlib terminology) set the aspect ratio to equal.

Once we move from a quick plot of the data to a more refined visualisation, it is worth using the (somewhat standard) matplotlib figure and axes template:

  1. Using fig, ax = plt.subplots() we create a figure (fig) and an axes (ax) object. The axes object defines the axes which will span our plot. (The name function is called subplots because we can create an array of axes objects with it - even though in this example we just want one, and that is why we do not pass any arguments to subplots).

  2. Pass the ax object to the xarray drawing routine.

  3. Modify the ax (and if we want the fig object) afterwards to fine-tune the plot.

Here is our full example:

[17]:
fig, ax = plt.subplots()
s1.plot.imshow(ax=ax)
ax.set_aspect("equal")
../_images/notebooks_xarray-plots1_31_0.svg

Selecting sub-sets of data

In our first example, we select data points in the xarray xa where the field value is greater than 0.001 (xa>0.001), and remove (drop=True) all other points.

We then slice the remaining data at \(z\approx 0\), and show a plot:

[18]:
s2 = xa.where(xa > 0.001, drop=True).sel(z=0, method="nearest")
s2.plot()
[18]:
<matplotlib.collections.QuadMesh at 0x7fbd22240e50>
../_images/notebooks_xarray-plots1_33_1.svg

We can also use data selection to focus on features of interest. For example, to only display data for which $ -3 \le `x :nbsphinx-math:le 3`$ and \(-4 \le x \le 4\), and to plot a slice of the data at \(z\approx 0\), we can use:

[19]:
s3 = xa.where((abs(xa.x) <= 3) & (abs(xa.y) <= 4), drop=True).sel(z=0, method="nearest")
s3.plot()
[19]:
<matplotlib.collections.QuadMesh at 0x7fbd221b3520>
../_images/notebooks_xarray-plots1_35_1.svg

Plot directly with matplotlib

Here is an example that re-creates the plots from above using only the data from the xarray object, and plain matplotlib plotting commands otherwise.

In the example below, the extend argument tells the imshow command what the x-axis and y-axis values represent.

[20]:
fig, ax = plt.subplots()

# figure out the data (d) and coordinates (x, y):
d = s3
x = s3.coords["x"].values
y = s3.coords["y"].values

# create the actual image, with `magma` colour map
p = ax.imshow(d, extent=[x.min(), x.max(), y.min(), y.max()], aspect="auto")
ax.set_xlabel("x")
ax.set_ylabel("y")

# add colour bar
fig.colorbar(p)
[20]:
<matplotlib.colorbar.Colorbar at 0x7fbd25b00970>
../_images/notebooks_xarray-plots1_37_1.svg

As a second example, we draw some contour lines.

[21]:
fig, ax = plt.subplots()

# contour needs x and y coordinates on matrices X and Y with the same dimensions as the data
Y, X = np.meshgrid(x, y, indexing="ij")

p2 = ax.contour(X, Y, d, levels=10)
fig.colorbar(p2)
[21]:
<matplotlib.colorbar.Colorbar at 0x7fbd21fe8580>
../_images/notebooks_xarray-plots1_39_1.svg

We can also combine different representations in the same axis:

[22]:
fig, ax = plt.subplots(figsize=(10, 8))

z = xa.sel(z=0, method="nearest")
x = xa.coords["x"].values
y = xa.coords["y"].values
p = ax.imshow(
    z,
    vmin=z.min(),
    vmax=z.max(),
    extent=[x.min(), x.max(), y.min(), y.max()],
    interpolation="nearest",
    aspect="equal",
    alpha=1,
)
fig.colorbar(p)

Y, X = np.meshgrid(x, y, indexing="ij")
p2 = ax.contour(X, Y, z, cmap="magma", levels=5)
fig.colorbar(p2)
[22]:
<matplotlib.colorbar.Colorbar at 0x7fbd21dfb2e0>
../_images/notebooks_xarray-plots1_41_1.svg

Changing the colourmap

See matplotlib documentation for the available colourmaps

[23]:
fig, ax = plt.subplots()
s3.plot.imshow(ax=ax, cmap="magma")
ax.set_aspect("equal")
../_images/notebooks_xarray-plots1_44_0.svg

Saving a figure

[24]:
fig, ax = plt.subplots()
s3.plot.imshow(ax=ax, cmap="cividis")
ax.set_aspect("equal")

fig.savefig("test.pdf")
../_images/notebooks_xarray-plots1_46_0.svg

The fig.savefig command creates a pdf file with the plot on disk in the current working directory. Other formats can be chosen through the file extension, as is usual for matplotlib’s savefig command.

Surface plot

There are probably better ways to create such plots, but here is an example for completeness.

[25]:
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator
import numpy as np

fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(10, 10))

# Make data.

z = s3
x = z.coords["x"].values
y = z.coords["y"].values

Y, X = np.meshgrid(x, y, indexing="ij")
Z = z.values

# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False)

# Customize the z axis.
ax.set_zlim(-1, 1)
ax.zaxis.set_major_locator(LinearLocator(11))

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=7)

plt.show()
../_images/notebooks_xarray-plots1_49_0.svg

Plotting with HoloViews

For dynamic/multidimensional plots: s. the holoviews tutorial