Plotting#

[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 ../octopus_data/benzene/

/builds/octopus-code/postopus/docs/octopus_data/benzene

[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]:

xa = run.default.scf.density(source="cube")

Calling density() provides the data accross all iterations in output_iter/scf.*/density.cube as well as static/density.cube. The iteration step is part of the coordinates of the given DataArray:

[7]:

xa.coords

[7]:

Coordinates:
  * step     (step) int64 32B 5 10 15 16
  * x        (x) float64 376B -13.04 -12.47 -11.91 -11.34 ... 11.91 12.47 13.04
  * y        (y) float64 392B -13.61 -13.04 -12.47 -11.91 ... 12.47 13.04 13.61
  * z        (z) float64 264B -9.071 -8.504 -7.937 -7.37 ... 7.937 8.504 9.071

[8]:

xa.values.shape

[8]:

(4, 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#

First, we only want to work with the data provided in static/ which is always the last iteration:

[9]:

xa = xa.isel(step=-1)
xa

[9]:

<xarray.DataArray 'density' (x: 47, y: 49, z: 33)> Size: 608kB
[75999 values with dtype=float64]
Coordinates:
    step     int64 8B 16
  * x        (x) float64 376B -13.04 -12.47 -11.91 -11.34 ... 11.91 12.47 13.04
  * y        (y) float64 392B -13.61 -13.04 -12.47 -11.91 ... 12.47 13.04 13.61
  * z        (z) float64 264B -9.071 -8.504 -7.937 -7.37 ... 7.937 8.504 9.071
Attributes:
    units:    au

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:

[10]:

s0 = xa.isel(z=16)
s0

[10]:

<xarray.DataArray 'density' (x: 47, y: 49)> Size: 18kB
[2303 values with dtype=float64]
Coordinates:
    step     int64 8B 16
  * x        (x) float64 376B -13.04 -12.47 -11.91 -11.34 ... 11.91 12.47 13.04
  * y        (y) float64 392B -13.61 -13.04 -12.47 -11.91 ... 12.47 13.04 13.61
    z        float64 8B 3e-06
Attributes:
    units:    au

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:

[11]:

s0.plot();
../_images/notebooks_xarray-plots1_20_0.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.

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"):

[12]:

s0.plot(x="x");
../_images/notebooks_xarray-plots1_23_0.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 16 in the xa.isel(z=16) 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\):

[13]:

s1 = xa.sel(z=0, method="nearest")

[14]:

s1.plot();
../_images/notebooks_xarray-plots1_27_0.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();
../_images/notebooks_xarray-plots1_29_0.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();
../_images/notebooks_xarray-plots1_31_0.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_33_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();
../_images/notebooks_xarray-plots1_35_0.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 y \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();
../_images/notebooks_xarray-plots1_37_0.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);
../_images/notebooks_xarray-plots1_39_0.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);
../_images/notebooks_xarray-plots1_41_0.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);
../_images/notebooks_xarray-plots1_43_0.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_46_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_48_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_51_0.svg

Plotting with HoloViews#

For dynamic/multidimensional plots: s. the holoviews tutorial