In [12]:
# Using the 'where' command in python is somewhat non-intuitive, 
# especially if coming from IDL. As an example, I'll first set up
# a 2D array
import numpy as np
a = np.arange(12).reshape(3,4)
a


Out[12]:
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

In [13]:
# Now find all the even elements
np.where((a%2)==0)


Out[13]:
(array([0, 0, 1, 1, 2, 2]), array([0, 2, 0, 2, 0, 2]))

In [14]:
# The return value is a tuple. Each array within the tuple is
# the index value along the respective axis. Note that the arrays
# are 1D, and once you index with them the return value is also
# 1D, as you would expect since the where clause can choose arbitrary
# elements out of the array
a[np.where((a%2)==0)]


Out[14]:
array([ 0,  2,  4,  6,  8, 10])

In [15]:
# A hangup if you're used to IDL comes with chaining together
# where statements
ii = np.where((a%2)==0)
jj = np.where(a[ii]>6)
a[ii[jj]]


---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-15-85e9dd30ca21> in <module>()
      3 ii = np.where((a%2)==0)
      4 jj = np.where(a[ii]>6)
----> 5 a[ii[jj]]

TypeError: tuple indices must be integers, not tuple

In [16]:
# This happens because both ii and jj are tuples, not arrays.
# You can work around it like this:
a[ii][jj]


Out[16]:
array([ 8, 10])

In [17]:
# But in general I think that using boolean arrays makes for more
# readable code than using where, especially compared to the
# 'complement' keyword in IDL.
a[((a%2)==0) & (a>6)]


Out[17]:
array([ 8, 10])

In [18]:
a[((a%2)==0) & ~(a>6)]


Out[18]:
array([0, 2, 4, 6])

In [19]:
# Note another difference with IDL: the logical operators
# (and, or, not) operate on *objects*. Numpy arrays are 
# ndarray objects, so operating on them with logical
# operators is undefined.
((a%2)==0) and (a>6)


---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-19-a4e70d079b3b> in <module>()
      3 # ndarray objects, so operating on them with logical
      4 # operators is undefined.
----> 5 ((a%2)==0) and (a>6)

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

In [20]:
# Using bitwise operators (&, |, ~, etc.) is defined, 
# since that means operating over each element of the array.
((a%2)==0) & (a>6)


Out[20]:
array([[False, False, False, False],
       [False, False, False, False],
       [ True, False,  True, False]], dtype=bool)

In [21]:
# Working with images (2D arrays)
# as an example, a quick implementation of optimal spectral extraction (Horne 1986)
# this method extracts a 1D spectrum from a 2D image by weighting each pixel
# using the object profile (i.e., the PSF)
# In this example I'll create a simple spectrum that is a sine wave function of
# wavelength, and has a gaussian profile on the spatial axis
#
# First, set up the index arrays (x is spectral axis, y is spatial axis)
x = np.linspace(0,20,100)
y = np.linspace(0,1,20)
X,Y = np.meshgrid(x,y)
X.shape


Out[21]:
(20, 100)

In [22]:
X[0,:5] # this is like wavelength


Out[22]:
array([ 0.        ,  0.2020202 ,  0.4040404 ,  0.60606061,  0.80808081])

In [23]:
Y[:,0] # this is like slit position


Out[23]:
array([ 0.        ,  0.05263158,  0.10526316,  0.15789474,  0.21052632,
        0.26315789,  0.31578947,  0.36842105,  0.42105263,  0.47368421,
        0.52631579,  0.57894737,  0.63157895,  0.68421053,  0.73684211,
        0.78947368,  0.84210526,  0.89473684,  0.94736842,  1.        ])

In [24]:
# Now create the spectrum. The gaussian width of the PSF is 0.1 pix
sigY = 0.1
spec2d = (np.sin(X)/2+0.5)* \
  (1./(np.sqrt(2*np.pi)*sigY))*np.exp(-(Y-0.5)**2/(2*sigY**2))

In [25]:
import matplotlib.pyplot as plt
plt.imshow(spec2d,origin='lower',interpolation='nearest')


Out[25]:
<matplotlib.image.AxesImage at 0x10acb2250>

In [26]:
# add some gaussian noise to make it like real data (which is never gaussian)
spec2d[:] += np.random.normal(scale=1,size=spec2d.shape)
plt.imshow(spec2d,origin='lower',interpolation='nearest')


Out[26]:
<matplotlib.image.AxesImage at 0x1059f57d0>

In [27]:
# now define the object profile, making the assumption that it is perfectly gaussian
# with parameters we already know, and that it doesn't vary across the image, so that
# it is simply a 1D function of slit position
P = (1./(np.sqrt(2*np.pi)*sigY))*np.exp(-(y-0.5)**2/(2*sigY**2))
P.shape


Out[27]:
(20,)

In [28]:
# now to do optimal extraction, we need to multiply the spectrum by the profile, which
# gives us the gaussian weighting of each pixel we need. but it doesn't work, because
# the array shapes don't match
spec1d = P*spec2d


---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-28-466a83da54e0> in <module>()
      2 # gives us the gaussian weighting of each pixel we need. but it doesn't work, because
      3 # the array shapes don't match
----> 4 spec1d = P*spec2d

ValueError: operands could not be broadcast together with shapes (20,) (20,100) 

In [29]:
# but you can just add an axis to the 1D profile to make it a 2D array, i.e., a
# column vector, so that the multiplication is propagated along the spectral axis
# (the x-axis)
P = P[:,np.newaxis]
P.shape


Out[29]:
(20, 1)

In [30]:
# lastly define the variance array, where we've assumed every pixel has the same noise
V = 1.0**2

In [31]:
# and finally do optimal extraction
spec1d = np.sum(P*spec2d/V,axis=0) / \
           np.sum(P**2/V,axis=0)
spec1d.shape


Out[31]:
(100,)

In [32]:
# now compare the result to boxcar extraction, where each pixel gets equal weight
# and you just sum along the spatial axis
plt.plot(spec2d.mean(axis=0),c='blue')  # boxcar extraction
plt.plot(spec1d,c='red')                # the optimal extraction result


Out[32]:
[<matplotlib.lines.Line2D at 0x10acca210>]

In [33]:
# the red line is less wiggly, so I guess it worked!

In [33]:


In [ ]: