docs/source/tutorial_jupyter.ipynb
Warning: This notebook needs a running kernel to be fully interactive, please run it locally or on mybinder.
</div>Vaex can process about 1 billion rows per second, and in combination with the Jupyter notebook, this allows for interactive exporation of large datasets.
The vaex-jupyter package contains the building blocks to interactively define an N-dimensional grid, which is then used for visualizations.
We start by defining the building blocks (vaex.jupyter.model.Axis, vaex.jupyter.model.DataArray and vaex.jupyter.view.DataArray) used to define and visualize our N-dimensional grid.
Let us first import the relevant packages, and open the example DataFrame:
import vaex
import vaex.jupyter.model as vjm
import numpy as np
import matplotlib.pyplot as plt
df = vaex.example()
df
We want to build a 2 dimensinoal grid with the number counts in each bin. To do this, we first define two axis objects:
E_axis = vjm.Axis(df=df, expression=df.E, shape=140)
Lz_axis = vjm.Axis(df=df, expression=df.Lz, shape=100)
Lz_axis
When we inspect the Lz_axis object we see that the min, max, and bin centers are all None. This is because Vaex calculates them in the background, so the kernel stays interactive, meaning you can continue working in the notebook. We can ask Vaex to wait until all background calculations are done. Note that for billions of rows, this can take over a second.
await vaex.jupyter.gather() # wait until Vaex is done with all background computation
Lz_axis # now min and max are computed, and bin_centers is set
Note that the Axis is a traitlets HasTrait object, similar to all ipywidget objects. This means that we can link all of its properties to an ipywidget and thus creating interactivity. We can also use observe to listen to any changes to our model.
Now that we have defined our two axes, we can create a vaex.jupyter.model.DataArray (model) together with a vaex.jupyter.view.DataArray (view).
A convenient way to do this, is to use the widget accessor data_array method, which creates both, links them together and will return a view for us.
The returned view is an ipywidget object, which becomes a visual element in the Jupyter notebook when displayed.
data_array_widget = df.widget.data_array(axes=[Lz_axis, E_axis], selection=[None, 'default'])
data_array_widget # being the last expression in the cell, Jupyter will 'display' the widget
Note: If you see this notebook on readthedocs, you will see the selection coordinate already has [None, 'default'], because cells below have already been executed and have updated this widget. If you run this notebook yourself (say on mybinder), you will see after executing the above cell, the selection will have [None] as its only value.
From the specification of the axes and the selections, Vaex computes a 3d histogram, the first dimension being the selections. Interally this is simply a numpy array, but we wrap it in an xarray DataArray object. An xarray DataArray object can be seen as a labeled Nd array, i.e. a numpy array with extra metadata to make it fully self-describing.
Notice that in the above code cell, we specified the selection argument with a list containing two elements in this case, None and 'default'. The None selection simply shows all the data, while the default refers to any selection made without explicitly naming it. Even though the later has not been defined at this point, we can still pre-emptively include it, in case we want to modify it later.
The most important properties of the data_array are printed out below:
# NOTE: since the computations are done in the background, data_array_widget.model.grid is initially None.
# We can ask vaex-jupyter to wait till all executions are done using:
await vaex.jupyter.gather()
# get a reference to the xarray DataArray object
data_array = data_array_widget.model.grid
print(f"type:", type(data_array))
print("dims:", data_array.dims)
print("data:", data_array.data)
print("coords:", data_array.coords)
print("Lz's data:", data_array.coords['Lz'].data)
print("Lz's attrs:", data_array.coords['Lz'].attrs)
print("And displaying the xarray DataArray:")
display(data_array) # this is what the vaex.jupyter.view.DataArray uses
Note that data_array.coords['Lz'].data is the same as Lz_axis.bin_centers and data_array.coords['Lz'].attrs contains the same min/max as the Lz_axis.
Also, we see that displaying the xarray.DataArray object (data_array_view.model.grid) gives us the same output as the data_array_view above. There is a big difference however. If we change a selection:
df.select(df.x > 0)
and scroll back we see that the data_array_view widget has updated itself, and now contains two selections! This is a very powerful feature, that allows us to make interactive visualizations.
To make interactive plots we can pass a custom display_function to the data_array_widget. This will override the default notebook behaviour which is a call to display(data_array_widget). In the following example we create a function that displays a matplotlib figure:
# NOTE: da is short for 'data array'
def plot2d(da):
plt.figure(figsize=(8, 8))
ar = da.data[1] # take the numpy data, and select take the selection
print(f'imshow of a numpy array of shape: {ar.shape}')
plt.imshow(np.log1p(ar.T), origin='lower')
df.widget.data_array(axes=[Lz_axis, E_axis], display_function=plot2d, selection=[None, True])
In the above figure, we choose index 1 along the selection axis, which referes to the 'default' selection. Choosing an index of 0 would correspond to the None selection, and all the data would be displayed. If we now change the selection, the figure will update itself:
df.select(df.id < 10)
As xarray's DataArray is fully self describing, we can improve the plot by using the dimension names for labeling, and setting the extent of the figure's axes.
Note that we don't need any information from the Axis objects created above, and in fact, we should not use them, since they may not be in sync with the xarray DataArray object. Later on, we will create a widget that will edit the Axis' expression.
Our improved visualization with proper axes and labeling:
def plot2d_with_labels(da):
plt.figure(figsize=(8, 8))
grid = da.data # take the numpy data
dim_x = da.dims[0]
dim_y = da.dims[1]
plt.title(f'{dim_y} vs {dim_x} - shape: {grid.shape}')
extent = [
da.coords[dim_x].attrs['min'], da.coords[dim_x].attrs['max'],
da.coords[dim_y].attrs['min'], da.coords[dim_y].attrs['max']
]
plt.imshow(np.log1p(grid.T), origin='lower', extent=extent, aspect='auto')
plt.xlabel(da.dims[0])
plt.ylabel(da.dims[1])
da_plot_view_nicer = df.widget.data_array(axes=[Lz_axis, E_axis], display_function=plot2d_with_labels)
da_plot_view_nicer
We can also create more sophisticated plots, for example one where we show all of the selections. Note that we can pre-emptively expect a selection and define it later:
def plot2d_with_selections(da):
grid = da.data
# Create 1 row and #selections of columns of matplotlib axes
fig, axgrid = plt.subplots(1, grid.shape[0], sharey=True, squeeze=False)
for selection_index, ax in enumerate(axgrid[0]):
ax.imshow(np.log1p(grid[selection_index].T), origin='lower')
df.widget.data_array(axes=[Lz_axis, E_axis], display_function=plot2d_with_selections,
selection=[None, 'default', 'rest'])
Modifying a selection will update the figure.
df.select(df.id < 10) # select 10 objects
df.select(df.id >= 10, name='rest') # and the rest
Another advantage of using xarray is its excellent plotting capabilities. It handles a lot of the boring stuff like axis labeling, and also provides a nice interface for slicing the data even more.
Let us introduce another axis, FeH (fun fact: FeH is a property of stars that tells us how much iron relative to hydrogen is contained in them, an idicator of their origin):
FeH_axis = vjm.Axis(df=df, expression='FeH', min=-3, max=1, shape=5)
da_view = df.widget.data_array(axes=[E_axis, Lz_axis, FeH_axis], selection=[None, 'default'])
da_view
We can see that we now have a 4 dimensional grid, which we would like to visualize.
And xarray's plot make our life much easier:
def plot_with_xarray(da):
da_log = np.log1p(da) # Note that an xarray DataArray is like a numpy array
da_log.plot(x='Lz', y='E', col='FeH', row='selection', cmap='viridis')
plot_view = df.widget.data_array([E_axis, Lz_axis, FeH_axis], display_function=plot_with_xarray,
selection=[None, 'default', 'rest'])
plot_view
We only have to tell xarray which axis it should map to which 'aesthetic', speaking in Grammar of Graphics terms.
Although we can change the selection in the notebook (e.g. df.select(df.id > 20)), if we create a dashboard (using Voila) we cannot execute arbitrary code. Vaex-jupyter also comes with many widgets, and one of them is a selection_expression widget:
selection_widget = df.widget.selection_expression()
selection_widget
The counter_selection creates a widget which keeps track of the number of rows in a selection. In this case we ask it to be 'lazy', which means that it will not cause extra passes over the data, but will ride along if some user action triggers a calculation.
await vaex.jupyter.gather()
w = df.widget.counter_selection('default', lazy=True)
w
Let us create new axis objects using the same expressions as before, but give them more general names (x_axis and y_axis), because we want to change the expressions interactively.
x_axis = vjm.Axis(df=df, expression=df.Lz)
y_axis = vjm.Axis(df=df, expression=df.E)
da_xy_view = df.widget.data_array(axes=[x_axis, y_axis], display_function=plot2d_with_labels, shape=180)
da_xy_view
Again, we can change the expressions of the axes programmatically:
# wait for the previous plot to finish
await vaex.jupyter.gather()
# Change both the x and y axis
x_axis.expression = np.log(df.x**2)
y_axis.expression = df.y
# Note that both assignment will create 1 computation in the background (minimal amount of passes over the data)
await vaex.jupyter.gather()
# vaex computed the new min/max, and the xarray DataArray
# x_axis.min, x_axis.max, da_xy_view.model.grid
But, if we want to create a dashboard with Voila, we need to have a widget that controls them:
x_widget = df.widget.expression(x_axis.expression, label='X axis')
x_widget
This widget will allow us to edit an expression, which will be validated by Vaex. How do we 'link' the value of the widget to the axis expression? Because both the Axis as well as the x_widget are HasTrait objects, we can link their traits together:
from ipywidgets import link
link((x_widget, 'value'), (x_axis, 'expression'))
Since this operation is so common, we can also directly pass the Axis object, and Vaex will set up the linking for us:
y_widget = df.widget.expression(y_axis, label='X axis')
# vaex now does this for us, much shorter
# link((y_widget, 'value'), (y_axis, 'expression'))
y_widget
await vaex.jupyter.gather() # lets wait again till all calculations are finished
If you are familiar with the ipyvuetify components, you can combine them to create very pretty widgets. Vaex-jupyter comes with a nice container:
from vaex.jupyter.widgets import ContainerCard
ContainerCard(title='My plot',
subtitle="using vaex-jupyter",
main=da_xy_view,
controls=[x_widget, y_widget], show_controls=True)
We can directly assign a Vaex expression to the x_axis.expression, or to x_widget.value since they are linked.
y_axis.expression = df.vx
So far we have been using interactive widgets to control the axes in the view. The figure itself however was not interactive, and we could not have panned or zoomed for example.
Vaex has a few builtin visualizations, most notably a heatmap and histogram using bqplot:
df = vaex.example() # we create the dataframe again, to leave all the plots above 'alone'
heatmap_xy = df.widget.heatmap(df.x, df.y, selection=[None, True])
heatmap_xy
Note that we passed expressions, and not axis objects. Vaex recognizes this and will create the axis objects for you. You can access them from the model:
heatmap_xy.model.x
The heatmap itself is again a widget. Thus we can combine it with other widgets to create a more sophisticated interface.
x_widget = df.widget.expression(heatmap_xy.model.x, label='X axis')
y_widget = df.widget.expression(heatmap_xy.model.y, label='X axis')
ContainerCard(title='My plot',
subtitle="using vaex-jupyter and bqplot",
main=heatmap_xy,
controls=[x_widget, y_widget, selection_widget],
show_controls=True,
card_props={'style': 'min-width: 800px;'})
By switching the tool in the toolbar (click <i aria-hidden="true" class="v-icon notranslate material-icons theme--light">pan_tool</i>, or changing it programmmatically in the next cell), we can zoom in. The plot's axis bounds are directly synched to the axis object (the x_min is linked to the x_axis min, etc). Thus a zoom action causes the axis objects to be changed, which will trigger a recomputation.
heatmap_xy.tool = 'pan-zoom' # we can also do this programmatically.
Since we can access the Axis objects, we can also programmatically change the heatmap. Note that both the expression widget, the plot axis label and the heatmap it self is updated. Everything is linked together!
heatmap_xy.model.x.expression = np.log10(df.x**2)
await vaex.jupyter.gather() # and we wait before we continue
Another visualization based on bqplot is the interactive histogram. In the example below, we show all the data, but the selection interaction will affect/set the 'default' selection.
histogram_Lz = df.widget.histogram(df.Lz, selection_interact='default')
histogram_Lz.tool = 'select-x'
histogram_Lz
# You can graphically select a particular region, in this case we do it programmatically
# for reproducability of this notebook
histogram_Lz.plot.figure.interaction.selected = [1200, 1300]
This shows an interesting structure in the heatmap above
The primary goal of Vaex-Jupyter is to provide users with a framework to create dashboard and new visualizations. Over time more visualizations will go into the vaex-jupyter package, but giving you the option to create new ones is more important. To help you create new visualization, we have examples on how to create your own:
If you want to create your own visualization on this framework, check out these examples:
The examples can also be found at the Examples page.