15

I am trying to add legend to a surface plot but unable to do so. Here is the code.

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import random

def fun(x, y):
  return 0.063*x**2 + 0.0628*x*y - 0.15015876*x + 96.1659*y**2 - 74.05284306*y  +      14.319143466051


fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
x = y = np.arange(-1.0, 1.0, 0.05)
X, Y = np.meshgrid(x, y)
zs = np.array([fun(x,y) for x,y in zip(np.ravel(X), np.ravel(Y))])
Z = zs.reshape(X.shape)

ax.plot_surface(X, Y, Z)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.plot(color='red',label='Lyapunov function on XY plane',linewidth=4)  # Adding legend

plt.show()

Kindly help. Thanks in advance.

Benoît Latinier
  • 2,062
  • 2
  • 24
  • 36
Chikorita Rai
  • 851
  • 5
  • 13
  • 17
  • What do you mean when you say "I am unable to do so"? Did you run into an error? Are you just unsure of how to do it? – gcarvelli Dec 12 '14 at 17:53
  • The way I am trying to include the legend is not giving the correct answer. So, yes, I do not know how to do it. – Chikorita Rai Dec 12 '14 at 17:56
  • `matplotlib` has a pretty cool tutorial on custom legends [here](http://matplotlib.org/users/legend_guide.html). There are a bunch of code examples that will help you write yours. – gcarvelli Dec 12 '14 at 18:00

2 Answers2

19

It is not trivial to make a legend in a 3D axis. You can use the following hack:

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib as mpl
import random

def fun(x, y):
  return 0.063*x**2 + 0.0628*x*y - 0.15015876*x + 96.1659*y**2 - 74.05284306*y  +      14.319143466051


fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
x = y = np.arange(-1.0, 1.0, 0.05)
X, Y = np.meshgrid(x, y)
zs = np.array([fun(x,y) for x,y in zip(np.ravel(X), np.ravel(Y))])
Z = zs.reshape(X.shape)

ax.plot_surface(X, Y, Z)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
fake2Dline = mpl.lines.Line2D([0],[0], linestyle="none", c='b', marker = 'o')
ax.legend([fake2Dline], ['Lyapunov function on XY plane'], numpoints = 1)
plt.show()

enter image description here

I would say a title is more appropriate than a legend in this case.

elyase
  • 39,479
  • 12
  • 112
  • 119
  • 1
    why do you need `fake2Dline = mpl.lines.Line2D([0],[0], linestyle="none", c='b', marker = 'o')` for it to work? – Charlie Parker Sep 08 '17 at 16:37
  • do you mind explaining what your code is doing and why the solution is as it is, it makes no sense to me and its hard to adapt to my own code without some explanation on the rationale. – Charlie Parker Sep 08 '17 at 16:43
  • also, there must be something missing in your instructions because my surface does not match the color of ur `fake2Dline`, it has its own separate color. – Charlie Parker Sep 08 '17 at 19:21
  • for those looking for explanation, he adds a fake line to the plot and add that line's legend instead of surface's legend in the plot. – alok.m Aug 03 '22 at 04:45
5

According to this question, the issue is ongoing, and there is a relatively simple workaround. You can manually set the two missing attributes that would allow legend to automatically create the patch for you:

surf = ax.plot_surface(X, Y, Z, label='Lyapunov function on XY plane')
surf._edgecolors2d = surf._edgecolor3d
surf._facecolors2d = surf._facecolor3d

ax.legend()

The attribute names on the right hand side of the assignment are surf._edgecolors3d and surf.facecolors3d for matplotlib < v3.3.3.

Mad Physicist
  • 107,652
  • 25
  • 181
  • 264