Numpy Array Slicing using a polygon in Matplotlib
astronomy, matplotlib, numpy, polygon, slice
Solution
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)
Problem
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.