fix(monitor): fix another issue where `color` is not transposed
@@ -854,10 +854,12 @@ class StreamPlotMonitor2D(BaseMonitor):
def _plot_streamlines(self, ax, us, vs, norms, cb_idx, is_grad=False):
ax.clear()
if self.mask is not None:
+ # FIXME if mask covers all points in the meshgrid, the following ValueError will be raised
+ # "Need at least one array to concatenate"
us[~self.mask] = np.nan
vs[~self.mask] = np.nan
ax.pcolor(*self._pcolor_args, shading='auto', cmap='Purples')
- kwargs = dict(color=norms)
+ kwargs = dict(color=norms.transpose())
kwargs.update(self.stream_kwargs)
stream = ax.streamplot(self.xs_plot[:, 0], self.ys_plot[0, :], us.transpose(), vs.transpose(), **kwargs)
if self.cbs[cb_idx] is not None:
