Numpy Array Slicing using a polygon in Matplotlib

Question

This seems like a fairly straightforward problem, but I'm new to Python and I'm struggling to resolve it. I've got a scatter plot / heatmap generated from two numpy arrays (about 25,000 pieces of information). The y-axis is taken directly from an array and the x-axis is generated from a simple subtraction operation on two arrays.

What I need to do now is slice up the data so that I can work with a selection that falls within certain parameters on the plot. For example, I need to extract all the points that fall within the parallelogram:

I'm able to cut out a rectangle using simple inequalities (see indexing idx_c, idx_h and idx, below) but I really need a way to select the points using a more complex geometry. It looks like this slicing can be done by specifying the vertices of the polygon. This is about the closest I can find to a solution, but I can't figure out how to implement it:

http://matplotlib.org/api/nxutils_api.html#matplotlib.nxutils.points_inside_poly

Ideally, I really need something akin to the indexing below, i.e. something like colorjh[idx]. Ultimately I'll have to plot different quantities (for example, colorjh[idx] vs colorhk[idx]), so the indexing needs to be transferable to all the arrays in the dataset (lots of arrays). Maybe that's obvious, but I would imagine there are solutions that might not be as flexible. In other words, I'll use this plot to select the points I'm interested in, and then I'll need those indices to work for other arrays from the same table.

Here's the code I'm working with:

import numpy as np
from numpy import ndarray
import matplotlib.pyplot as plt
import matplotlib
import atpy
from pylab import *

twomass = atpy.Table()

twomass.read('/IRSA_downloads/2MASS_GCbox1.tbl')

hmag = list([twomass['h_m']])
jmag = list([twomass['j_m']])
kmag = list([twomass['k_m']])

hmag = np.array(hmag)
jmag = np.array(jmag)
kmag = np.array(kmag)

colorjh = np.array(jmag - hmag)
colorhk = np.array(hmag - kmag)

idx_c = (colorjh > -1.01) & (colorjh < 6)  #manipulate x-axis slicing here here
idx_h = (hmag > 0) & (hmag < 17.01)        #manipulate y-axis slicing here
idx = idx_c & idx_h

# heatmap below
heatmap, xedges, yedges = np.histogram2d(hmag[idx], colorjh[idx], bins=200)
extent = [yedges[0], yedges[-1], xedges[-1], xedges[0]]
plt.clf()
plt.imshow(heatmap, extent=extent, aspect=0.65)

plt.xlabel('Color(J-H)', fontsize=15)           #adjust axis labels here
plt.ylabel('Magnitude (H)', fontsize=15)

plt.gca().invert_yaxis()       #I put this in to recover familiar axis orientation

plt.legend(loc=2)
plt.title('CMD for Galactic Center (2MASS)', fontsize=20)
plt.grid(True)
colorbar()

plt.show()

Like I say, I'm new to Python, so the less jargon-y the explanation the more likely I'll be able to implement it. Thanks for any help y'all can provide.

Answer 1

a = np.random.randint(0,10,(100,100))

x = np.linspace(-1,5.5,100) # tried to mimic your data boundaries
y = np.linspace(8,16,100)
xx, yy = np.meshgrid(x,y)

m = np.all([yy > xx**2, yy < 10* xx, xx < 4, yy > 9], axis = 0)