mplot3d: contourf offset, limits and ticks

Question

I'm trying to get a nice contourf plot under a mplot3d surface. I'd want it to appear on the floor of the 3d axis cube with a little offset from my data lower limits. Right now I'm doing something like this:

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

N = 50

fig = plt.figure()
ax = fig.gca(projection='3d')

surface = np.zeros((N,N))

# gaussian
for x in np.arange(N, dtype=float):
    for y in np.arange(N, dtype=float):
        sigma = 0.2
        xn = (x - N/2)/N
        yn = (y - N/2)/N
        r = np.sqrt(xn**2.0 + yn**2.0)
        surface[x,y] = np.exp((-r**2.0)/(2.0*sigma**2.0))

# mesh grid NxN points in [0,1]
gx, gy = np.meshgrid(np.linspace(0,1,N),np.linspace(0,1,N))

ax.plot_surface(gx, gy, surface, rstride=2, cstride=2, cmap=mpl.cm.Spectral)

# extend z axis limit to make room for contourf
ax.set_zlim3d(np.min(surface) - 0.5, np.max(surface))

# contour on the floor
levels = np.linspace(np.min(surface), np.max(surface), 20)
ax.contourf(gx, gy, surface, levels=levels,
            offset=(np.min(surface) - 0.5), cmap=mpl.cm.Spectral)

That plots this image

This looks fine but it adds a couple of ticks where I extend the zaxis under the data minimum. I'd like to not show any tick under the minimum but still extend the zaxis to offset the contourf plot.

Any idea? How do I hide or not draw at all the red circled ticks?

Answer 1

It was easier than I thought, many thanks to @gboffi for pointing me to the correct apis.

s_min = np.min(surface)
s_max = np.max(surface)

# filter out extra ticks that exceed data limits
ax.set_zticks(filter(lambda x: s_min <= x <= s_max, ax.get_zticks()))