Matplotlib 3D surface plot from 2D pandas dataframe

Question

I have a pandas dataframe containing four columns of measurement data. I'd like to create a 3D surface plot with the row index as X, the column index as Y, and the data as Z. (The data in each column is a series of discrete measurements output from a test that steps through all values of X for each category Y)

import pandas as pd
import numpy as np


df = pd.DataFrame(np.random.randn(5, 4), columns=['A', 'B', 'C', 'D'])
print(df)

   A         B         C         D
0  0.791692 -0.945571  0.183304  2.039369
1 -0.474666  1.117902 -0.483240  0.137620
2  1.448765  0.228217  0.294523  0.728543
3 -0.196164  0.898117 -1.770550  1.259608
4  0.646730 -0.366295 -0.893671 -0.745815

I tried converting the df into a numpy grid using np.meshgrid as below but not sure I really understand what is required, or if I can use the df indices in this way.

import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

x = df.columns
y = df.index
X,Y = np.meshgrid(x,y)
Z = df
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, Z)

I've read through the matplotlib 3D tutorial and related answers here, but am still stuck. Would be very grateful if someone could point me in the right direction please.

Solution 1

The general strategy you pursue is fine. The only error you have is that you create a meshgrid from a list of strings. Of course maplotlib cannot plot strings.

You can therfore create an array of the same length as the number of columns in your dataframe and plug that into the meshgrid.

x = np.arange(len(df.columns))

Solution 2

just take of the columns names (['A', 'B', 'C', 'D']) and it should work.

you can later change the ticks of the axis for ['A', 'B', 'C', 'D'].

import pandas as pd
import numpy as np


df = pd.DataFrame(np.random.randn(5, 4))

import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

x = df.columns
y = df.index
X,Y = np.meshgrid(x,y)
Z = df
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, Z)

Comments

[Violet]

Doh! Had convinced myself the approach was completely wrong, hence overlooked the obvious. Thank you very much!