1 from numpy import *
   2 
   3 
   4 def interpn_nearest( z, targetcoords, bincoords=None ):
   5         '''Interpolate some data, z, located at bincoords in the target locations targetcoords
   6         using nearest neighbor interpolation.
   7                 z - is the data array
   8                 targetcoords - are the coordinates where we want to evaluate the data.
   9                         It must have a format so that z[targetcoords] would make sense
  10                         if targetcoords where integers.
  11                         mgrid and ogrid can be used to get such a format.
  12                 bincoords - is a list of arrays. Each array should containt.
  13         '''
  14         coords = interpn_check_data(z,targetcoords,bincoords)
  15         indices = [ c.round().astype('i') for c in coords ]
  16         return z[indices]
  17 
  18 
  19 def interpn_linear( z, targetcoords, bincoords=None ):
  20         '''Interpolate some data, z, located at bincoords in the target locations targetcoords
  21         using linear interpolation.
  22                 z - is the data array
  23                 targetcoords - are the coordinates where we want to evaluate the data.
  24                         It must have a format so that z[targetcoords] would make sense
  25                         if targetcoords where integers.
  26                         mgrid and ogrid can be used to get such a format.
  27                 bincoords - is a list of arrays. Each array should containt.
  28         '''
  29         coords = interpn_check_data(z,targetcoords,bincoords)
  30 
  31         indices = [ x.astype('i') for x in coords ]
  32         weights = [ x-i for x,i in zip(coords,indices) ]
  33 
  34         res = z[indices]*0
  35 
  36         for selector in ndindex( *z.ndim*(2,) ):
  37                 weight = 1
  38                 for w,s in zip(weights,selector):
  39                         if s: weight = weight*w
  40                         else: weight = weight*(1-w)
  41                 value = z[ [ i+s for i,s in zip(indices,selector) ] ]
  42                 res += weight*value
  43 
  44         return res
  45 
  46 
  47 
  48 def rebin( a, newshape ):
  49         '''Rebin an array to a new shape, without interpolation.
  50         This can be usefull if the array type is not a vector space.
  51         '''
  52         assert a.ndim == len(newshape)
  53 
  54         slices = [ slice(0,old, float(old)/new) for old,new in zip(a.shape,newshape) ]
  55         coordinates = ogrid[slices]
  56         indices = [ c.astype('i') for c in coordinates ]
  57         return a[indices]
  58 
  59 
  60 
  61 #################################
  62 ## more internal functions
  63 
  64 def array_coordinates( targetx, binx ):
  65         '''Computes the coordinates in an array reference.
  66                 targetx are the coordinates of the points that we want to get in the array reference.
  67                 binx    are the coordinates of the array bins
  68         Example:
  69                 >>> array_coordinates( [2, 10], [1, 3, 13] )
  70                 [0.5, 1.7]
  71         '''
  72         tol = 1e-12
  73         binx = asarray(binx)
  74         nx = clip( targetx, binx.min()*(1+tol), binx.max()*(1-tol) )
  75 
  76         i = searchsorted( binx, nx )      # index of the biggest smaller bin
  77 
  78         bx = empty((len(binx)+2,), 'd')   # avoid probles at the border
  79         bx[0] = binx[0] - 1
  80         bx[1:-1] = binx
  81         bx[-1] = binx[-1] + 1
  82 
  83         return i-1 + ( nx - bx[i] )/( bx[i+1] - bx[i] )
  84 
  85 
  86 def interpn_check_data( z, targetcoords, bincoords ):
  87         '''Check the validity of the input data for intepn_x functions
  88         and returns the target coordinates in the array reference
  89         '''
  90         dim = z.ndim
  91         if bincoords:
  92                 if len(bincoords) != dim:
  93                         raise ValueError, 'bincoords shape mismatch.'
  94                 for i in range(dim):
  95                         if prod(bincoords[i].shape) != z.shape[i]:
  96                                 raise ValueError, 'bincoords shape mismatch.'
  97 
  98                 coords = [ array_coordinates(targetcoords[i],bincoords[i].ravel()) for i in range(dim) ]
  99         else:
 100                 coords = targetcoords
 101         return coords
 102 
 103 
 104 def interp2_linear( z, tx, ty, binx=None, biny=None  ):
 105         '''Toy function just like interpn_linear in 2 dimensions.
 106         This function exists just to help the understanding and mantaining of interpn_linear.
 107         '''
 108         if not binx is None: tx = array_coordinates( tx,binx )
 109         if not biny is None: ty = array_coordinates( ty,biny )
 110 
 111         ix = tx.astype('i')
 112         iy = ty.astype('i')
 113         wx = tx-ix
 114         wy = ty-iy
 115 
 116         return    (1-wy)*(1-wx) * z[iy  ,ix  ] \
                + (1-wy)*   wx  * z[iy  ,ix+1] \
                +    wy *(1-wx) * z[iy+1,ix  ] \
                +    wy *   wx  * z[iy+1,ix+1]
 117 
 118 
 119 
 120 
 121 
 122 
 123 ###########################################
 124 ## some visual testing
 125 ##
 126 import pylab
 127 
 128 def test():
 129         x,y = ogrid[ -1:1:10j, -1:1:10j ]
 130         z = sin( x**2 + y**2 )
 131 
 132         binx = (x,y)
 133         tx = ogrid[ -2:2:100j, -2:2:100j ]
 134 
 135         pylab.subplot(221)
 136         pylab.title('original')
 137         pylab.imshow(z)
 138         pylab.subplot(223)
 139         pylab.title('interpn_nearest')
 140         pylab.imshow( interpn_nearest( z, tx, binx ) )
 141         pylab.subplot(222)
 142         pylab.title('interpn_linear')
 143         pylab.imshow( interpn_linear( z, tx, binx ) )
 144         pylab.subplot(224)
 145         pylab.title('interp2_linear')
 146         pylab.imshow( interp2_linear( z, tx[0],tx[1], x.ravel(),y.ravel() ) )
 147         pylab.show()