Drawing a surface 3D plot using "plotnine" library

Question

Question : Using the python library 'plotnine', can we draw an interactive 3D surface plot?

Backup Explanations

1. What I'd like to do is, under python environment, creating an interactive 3D plot with R plot grammars like we do with ggplot2 library in R. It's because I have hard time remembering grammars of matplotlib and other libraries like seaborn.

2. An interactive 3D plot means a 3D plot that you can zoom in, zoom out, and scroll up and down, etc.

3. It seems like only Java supported plotting libraries scuh as bokeh or plotly can create interactive 3D plots. But I want to create it with the library 'plotnine' because the library supports ggplot-like grammar, which is easy to remember.

4. For example, can I draw a 3D surface plot like the one below with the library 'plotnine'?

   import plotly.plotly as py
   import plotly.graph_objs as go
   import pandas as pd
   
   # Read data from a csv
   z_data =
   pd.read_csv('https://raw.githubusercontent.com/plotly/datasets/
   master/api_docs/mt_bruno_elevation.csv')
   
    data = [
           go.Surface(
           z=z_data.as_matrix()
           )]
    layout = go.Layout(
    title='Mt Bruno Elevation',
    autosize=False,
    width=500,
    height=500,
    margin=dict(
    l=65,
    r=50,
    b=65,
    t=90
      )
    )
    fig = go.Figure(data=data, layout=layout)
    py.iplot(fig, filename='elevations-3d-surface')
   

The codes above make a figure like below.

You can check out the complete interactive 3D surface plot in this link

p.s. If i can draw an interactive 3D plot with ggplot-like grammar, it does not have to be the 'plotnine' library that we should use.

Thank you for your time for reading this question!

Answer 1

It is possible, if you are willing to expand plotnine a bit, and caveats apply. The final code is as simple as:

(
    ggplot_3d(mt_bruno_long)
    + aes(x='x', y='y', z='height')
    + geom_polygon_3d(size=0.01)
    + theme_minimal()
)

And the result:

First, you need to transform your data into long format:

z_data = pd.read_csv('https://raw.githubusercontent.com/plotly/datasets/master/api_docs/mt_bruno_elevation.csv', index_col=0)
z = z_data.values
nrows, ncols = z.shape
x, y = np.linspace(0, 1, nrows), np.linspace(0, 1, ncols)
x, y = np.meshgrid(x, y)
mt_bruno_long = pd.DataFrame({'x': x.flatten(), 'y': y.flatten(), 'height': z.flatten()})

Then, we need to create equivalents for ggplot and geom_polygon with awareness of the third dimension:

from plotnine import ggplot, geom_polygon
from plotnine.utils import to_rgba, SIZE_FACTOR


class ggplot_3d(ggplot):
    def _create_figure(self):
        figure = plt.figure()
        axs = [plt.axes(projection='3d')]
        
        figure._themeable = {}
        self.figure = figure
        self.axs = axs
        return figure, axs
    
    def _draw_labels(self):
        ax = self.axs[0]
        ax.set_xlabel(self.layout.xlabel(self.labels))
        ax.set_ylabel(self.layout.ylabel(self.labels))
        ax.set_zlabel(self.labels['z'])


class geom_polygon_3d(geom_polygon):
    REQUIRED_AES = {'x', 'y', 'z'}

    @staticmethod
    def draw_group(data, panel_params, coord, ax, **params):
        data = coord.transform(data, panel_params, munch=True)
        data['size'] *= SIZE_FACTOR

        grouper = data.groupby('group', sort=False)
        for i, (group, df) in enumerate(grouper):
            fill = to_rgba(df['fill'], df['alpha'])
            polyc = ax.plot_trisurf(
                df['x'].values,
                df['y'].values,
                df['z'].values,
                facecolors=fill if any(fill) else 'none',
                edgecolors=df['color'] if any(df['color']) else 'none',
                linestyles=df['linetype'],
                linewidths=df['size'],
                zorder=params['zorder'],
                rasterized=params['raster'],
            )
            # workaround for https://github.com/matplotlib/matplotlib/issues/9535
            if len(set(fill)) == 1:
                polyc.set_facecolors(fill[0])

For interactivity you can use any matplotlib backend of your liking, I went with ipympl (pip install ipympl and then %matplotlib widget in a jupyter notebook cell).

The caveats are:

• while shading works nice, plot_trisurf does not handle facecolors well (there is a PR to fix it here)
• you may want to add a parameter allowing to disable shading, see matplotlib 3D shading examples
• faceting, flipping axes etc will not work without further fiddling - this could however be addressed in the future as discussed in this plotnine issue about bringing 3D plots to plotnine.

Edit: In case if the dataset becomes unavailable, here is a self-contained example based on matplotlib's documentation:

import numpy as np

n_radii = 8
n_angles = 36

radii = np.linspace(0.125, 1.0, n_radii)
angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)[..., np.newaxis]

x = np.append(0, (radii*np.cos(angles)).flatten())
y = np.append(0, (radii*np.sin(angles)).flatten())

z = np.sin(-x*y)
df = pd.DataFrame(dict(x=x,y=y,z=z))

(
    ggplot_3d(df)
    + aes(x='x', y='y', z='z')
    + geom_polygon_3d(size=0.01)
    + theme_minimal()
)