For an application I am developing, I need to call ax.scatter then change the colors of the points after the fact. However, if the input x/y values contain NaN values, then setting the colors after the fact doesn't work properly. The following example:

import matplotlib
matplotlib.use('Agg')

import numpy as np
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, aspect='auto')

x = np.array([1, 2, 3])
y = np.array([2, np.nan, 4])
c = np.array([3, 4, 5])

s = ax.scatter(x, y, c=c, vmin=3, vmax=5, s=320, zorder=1, edgecolor='none')

s = ax.scatter(x, y, s=160, zorder=2, edgecolor='none')
s.set_array(c)
s.set_clim(3, 5)

s = ax.scatter([], [], s=40, zorder=3, edgecolor='none')
s.set_offsets(np.vstack([x, y]).transpose())
s.set_array(c)
s.set_clim(3, 5)

fig.savefig('test.png')

produces the following output:

The top right point should be all yellow. However it isn't, because in the scatter call, invalid values are removed, so when I try and set the colors with set_array, the color values are offset from the x/y values. In the final case, things work since the 'bad' values are masked from neither the offsets nor the colors.

Is there a way to achieve something along the lines of the middle example (setting the colors after the fact)