Draw an image and a stem plot in 3d with matpltolib

I'm trying to make a 3d plot with python that shows an image at a certain level and certain specific points as a stem plot (basically the image is a sort of 2d map and the stem plots represent the value of a certain parameter at the point in the map)

Current code :

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(-1, 1, 101)
y = np.linspace(-1, 1, 101)

X, Y = np.meshgrid(x, y)
z = np.exp(-(X**2) - Y**2)

x0 = [-0.5, 0, 0.5]
y0 = [-0.5, 0, 0.5]
z0 = [0.5, 1, 2]

fig, ax = plt.subplots(subplot_kw=dict(projection="3d"))
ax.contourf(
    x,
    y,
    z,
    zdir="z",
    offset=-0.0,
)
ps = ax.stem(x0, y0, z0)
for p in ps:
    p.set_zorder(20)

plt.show()

I added the set_zorder() part to try to put the stem plot forward, however the resulting plot looks like this:

Where the lines are hidden behind the 2d image while I would like it to look like the stem where originating on the 2d plot.

Random image I've found on the internet to explicate what I'm trying to achieve (with bars instead of stem plots because I couldn't find one with stem plots). Many thanks.

You can use the computed_zorder kwarg to control this. By setting that to False, a 3D axes will respect the zorder values that you set. The default is True, which tries to work out what they should be.

See the docs here.

The code to change is your example is:

fig, ax = plt.subplots(subplot_kw=dict(projection="3d", computed_zorder=False))

which yields: