[SOLVED] matplotlib waterfall plot with surfaces shows black artifacts on border of plot

I have a script to write a heatmap (or contour map) inside an arbitrary closed shape. A bounding box of grid points is created and a mask is used to clip any points outside of the shape. However, if I want to create a stacked plot of these maps (to show changes along a fourth dimension), the grid points are painted black.

Here is the script for a single plot:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.path import Path
from scipy.interpolate import griddata

arbitrary_shape_points = [
    (0.0, 0.5), (0.1, 0.6), (0.3, 0.55), (0.5, 0.4), (0.6, 0.2),  # Top curve
    (0.65, 0.0), (0.6, -0.2), (0.5, -0.4), (0.3, -0.55), (0.1, -0.6), # Bottom curve right
    (0.0, -0.5), (-0.1, -0.6), (-0.3, -0.55), (-0.5, -0.4), (-0.6, -0.2), # Bottom curve left
    (-0.65, 0.0), (-0.6, 0.2), (-0.5, 0.4), (-0.3, 0.55), (-0.1, 0.6), # Top curve left
    (0.0, 0.5) # Closing point
]

shape_path = Path(arbitrary_shape_points)

np.random.seed(42)
num_points = 100
x_data = np.random.uniform(-0.7, 0.7, num_points) 
y_data = np.random.uniform(-0.7, 0.7, num_points) 
z_data = np.pi*np.sin(np.pi * x_data) + np.exp(-np.pi*y_data)

# Bounding box 
shape_xmin = min([p[0] for p in arbitrary_shape_points[:-1]])
shape_ymin = min([p[1] for p in arbitrary_shape_points[:-1]])
shape_xmax = max([p[0] for p in arbitrary_shape_points[:-1]])
shape_ymax = max([p[1] for p in arbitrary_shape_points[:-1]])

# Grid
grid_resolution = 500
x_grid = np.linspace(shape_xmin, shape_xmax, grid_resolution)
y_grid = np.linspace(shape_ymin, shape_ymax, grid_resolution)
xx, yy = np.meshgrid(x_grid, y_grid)

# Interpolate data
interpolation_method = 'cubic'
zz_interpolated = griddata((x_data, y_data), z_data, (xx, yy), method=interpolation_method)


# Mask the grid outside the shape
grid_points = np.column_stack((xx.flatten(), yy.flatten()))
mask = shape_path.contains_points(grid_points).reshape(xx.shape)
zz_masked = np.where(mask, zz_interpolated, np.nan)

# Plot Heatmap using pcolormesh
plt.figure(figsize=(8, 6))

heatmap = plt.pcolormesh(xx, yy, zz_masked, cmap='viridis', shading='auto') 
plt.colorbar(heatmap, label='Z Value')

plt.legend()

plt.title('Heatmap within Arbitrary Shape') 
plt.xlabel('X')
plt.ylabel('Y')

x_shape_original, y_shape_original = zip(*arbitrary_shape_points)
plt.plot(x_shape_original, y_shape_original, 'k-', linewidth=1)

# create whitespace around the bounding box
whitespace_factor = 1.2
plt.xlim(shape_xmin * whitespace_factor, shape_xmax * whitespace_factor)
plt.ylim(shape_ymin * whitespace_factor, shape_ymax * whitespace_factor)
plt.gca().set_aspect('equal', adjustable='box')
plt.show()

The result is this:

This is the stacked plot script:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.path import Path
from scipy.interpolate import griddata

arbitrary_shape_points = [
    (0.0, 0.5), (0.1, 0.6), (0.3, 0.55), (0.5, 0.4), (0.6, 0.2),  # Top curve
    (0.65, 0.0), (0.6, -0.2), (0.5, -0.4), (0.3, -0.55), (0.1, -0.6), # Bottom curve right
    (0.0, -0.5), (-0.1, -0.6), (-0.3, -0.55), (-0.5, -0.4), (-0.6, -0.2), # Bottom curve left
    (-0.65, 0.0), (-0.6, 0.2), (-0.5, 0.4), (-0.3, 0.55), (-0.1, 0.6), # Top curve left
    (0.0, 0.5) # Closing point
]
shape_path = Path(arbitrary_shape_points)

np.random.seed(42)
num_points = 100
x_data = np.random.uniform(-0.7, 0.7, num_points)
y_data = np.random.uniform(-0.7, 0.7, num_points)
fourth_dimension_values = np.linspace(0, 1, 5) 

shape_xmin = min([p[0] for p in arbitrary_shape_points[:-1]])
shape_ymin = min([p[1] for p in arbitrary_shape_points[:-1]])
shape_xmax = max([p[0] for p in arbitrary_shape_points[:-1]])
shape_ymax = max([p[1] for p in arbitrary_shape_points[:-1]])

grid_resolution = 100
x_grid = np.linspace(shape_xmin, shape_xmax, grid_resolution)
y_grid = np.linspace(shape_ymin, shape_ymax, grid_resolution)
xx, yy = np.meshgrid(x_grid, y_grid)
grid_points = np.column_stack((xx.flatten(), yy.flatten()))
mask = shape_path.contains_points(grid_points).reshape(xx.shape)

interpolation_method = 'cubic'

fig = plt.figure(figsize=(10, 8), facecolor=(0, 0, 0, 0)) 
ax = fig.add_subplot(111, projection='3d', facecolor=(0, 0, 0, 0)) 


z_offset = 0
z_step = 1

for i, fd_value in enumerate(fourth_dimension_values):

    z_data = np.pi*np.sin(np.pi * x_data) + np.exp(-np.pi*y_data) + fd_value

    zz_interpolated = griddata((x_data, y_data), z_data, (xx, yy), method=interpolation_method)

    # Mask the grid outside the shape
    zz_masked = np.where(mask, zz_interpolated, np.nan)

    # Prepare Z values for 3D plot - constant Z for each slice, offset along z-axis
    z_surface = np.full_like(xx, z_offset + i * z_step)

    z_min_slice = np.nanmin(zz_masked)
    z_max_slice = np.nanmax(zz_masked)

    if z_max_slice > z_min_slice:
        zz_normalized = (zz_masked - z_min_slice) / (z_max_slice - z_min_slice)
    else:
        zz_normalized = np.zeros_like(zz_masked)

    # Create facecolors from normalized data
    facecolors_heatmap = plt.cm.viridis(zz_normalized)

    # Make masked areas fully transparent 
    facecolors_heatmap[np.isnan(zz_masked)] = [0, 0, 0, 0] 

    # Plot each heatmap slice as a surface
    surf = ax.plot_surface(xx, yy, z_surface, facecolors=facecolors_heatmap, linewidth=0, antialiased=False, shade=False, alpha=0.8, rstride=1, cstride=1)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Fourth Dimension Index') 
ax.set_title('Stacked Heatmaps along Fourth Dimension') 

ax.view_init(elev=30, azim=-45) 

ax.set_box_aspect([np.diff(ax.get_xlim())[0], np.diff(ax.get_ylim())[0], np.diff(ax.get_zlim())[0]*0.5]) 

plt.show()

The result is this:

I am not sure how to stop the black border from appearing. I tried setting those points to be transparent but that does not do anything.