1 #!/usr/bin/env python
   2 import sys
   3 import numpy as npy
   4 import pylab as P
   5 from scipy.weave import inline, converters, blitz
   6 from scipy.testing import measure
   7 # Blitz conversion is terrific, but sometimes you don't have fixed array sizes
   8 # in your problem. Fortunately numpy iterators still make writing inline
   9 # weave code very, very simple. 
  10 
  11 def multi_iter_example():
  12     # This is a very simple example of multi dimensional iterators, and
  13     # their power to "broadcast" arrays of compatible shapes. It shows that
  14     # the very same code that is entirely ignorant of dimensionality can
  15     # achieve completely different computations based on the rules of
  16     # broadcasting.
  17 
  18     # it is important to know that the weave array conversion of "a"
  19     # gives you access in C++ to:
  20     # py_a -- PyObject *
  21     # a_array -- PyArrayObject *
  22     # a -- py_array->data
  23     
  24     a = npy.ones((4,4), npy.float64)
  25     # for the sake of driving home the "dynamic code" approach...
  26     dtype2ctype = {
  27         npy.dtype(npy.float64): 'double',
  28         npy.dtype(npy.float32): 'float',
  29         npy.dtype(npy.int32): 'int',
  30         npy.dtype(npy.int16): 'short',
  31     }
  32     dt = dtype2ctype.get(a.dtype)
  33     
  34     # this code does a = a*b inplace, broadcasting b to fit the shape of a
  35     code = \
  36 """
  37 %s *p1, *p2;
  38 PyObject *itr;
  39 itr = PyArray_MultiIterNew(2, a_array, b_array);
  40 while(PyArray_MultiIter_NOTDONE(itr)) {
  41   p1 = (%s *) PyArray_MultiIter_DATA(itr, 0);
  42   p2 = (%s *) PyArray_MultiIter_DATA(itr, 1);
  43   *p1 = (*p1) * (*p2);
  44   PyArray_MultiIter_NEXT(itr);
  45 }
  46 """ % (dt, dt, dt)
  47 
  48     b = npy.arange(4, dtype=a.dtype)
  49     print '\n         A                  B     '
  50     print a, b
  51     # this reshaping is redundant, it would be the default broadcast
  52     b.shape = (1,4)
  53     inline(code, ['a', 'b'])
  54     print "\ninline version of a*b[None,:],"
  55     print a
  56     a = npy.ones((4,4), npy.float64)
  57     b = npy.arange(4, dtype=a.dtype)
  58     b.shape = (4,1)
  59     inline(code, ['a', 'b'])
  60     print "\ninline version of a*b[:,None],"
  61     print a
  62 
  63 def data_casting_test():
  64     # In my MR application, raw data is stored as a file with one or more
  65     # (block-hdr, block-data) pairs. Block data is one or more
  66     # rows of Npt complex samples in big-endian integer pairs (real, imag).
  67     #
  68     # At the block level, I encounter three different raw data layouts--
  69     # 1) one plane, or slice: Y rows by 2*Npt samples
  70     # 2) one volume: Z slices * Y rows by 2*Npt samples
  71     # 3) one row sliced across the z-axis: Z slices by 2*Npt samples
  72     #
  73     # The task is to tease out one volume at a time from any given layout,
  74     # and cast the integer precision data into a complex64 array.
  75     # Given that contiguity is not guaranteed, and the number of dimensions
  76     # can vary, Numpy iterators are useful to provide a single code that can
  77     # carry out the conversion.
  78     #
  79     # Other solutions include:
  80     # 1) working entirely with the string data from file.read() with string
  81     #    manipulations (simulated below).
  82     # 2) letting numpy handle automatic byteorder/dtype conversion
  83     
  84     nsl, nline, npt = (20,64,64)
  85     hdr_dt = npy.dtype('>V28')
  86     # example 1: a block is one slice of complex samples in short integer pairs
  87     blk_dt1 = npy.dtype(('>i2', nline*npt*2))
  88     dat_dt = npy.dtype({'names': ['hdr', 'data'], 'formats': [hdr_dt, blk_dt1]})
  89     # create an empty volume-- nsl contiguous blocks
  90     vol = npy.empty((nsl,), dat_dt)
  91     t = time_casting(vol[:]['data'])
  92     P.plot(100*t/t.max(), 'b--', label='vol=20 contiguous blocks')
  93     P.plot(100*t/t.max(), 'bo')
  94     # example 2: a block is one entire volume
  95     blk_dt2 = npy.dtype(('>i2', nsl*nline*npt*2))
  96     dat_dt = npy.dtype({'names': ['hdr', 'data'], 'formats': [hdr_dt, blk_dt2]})
  97     # create an empty volume-- 1 block
  98     vol = npy.empty((1,), dat_dt)
  99     t = time_casting(vol[0]['data'])
 100     P.plot(100*t/t.max(), 'g--', label='vol=1 contiguous block')
 101     P.plot(100*t/t.max(), 'go')    
 102     # example 3: a block slices across the z dimension, long integer precision
 103     # ALSO--a given volume is sliced discontiguously
 104     blk_dt3 = npy.dtype(('>i4', nsl*npt*2))
 105     dat_dt = npy.dtype({'names': ['hdr', 'data'], 'formats': [hdr_dt, blk_dt3]})
 106     # a real data set has volumes interleaved, so create two volumes here
 107     vols = npy.empty((2*nline,), dat_dt)
 108     # and work on casting the first volume
 109     t = time_casting(vols[0::2]['data'])
 110     P.plot(100*t/t.max(), 'r--', label='vol=64 discontiguous blocks')
 111     P.plot(100*t/t.max(), 'ro')    
 112     P.xticks([0,1,2], ('strings', 'numpy auto', 'inline'))
 113     P.gca().set_xlim((-0.25, 2.25))
 114     P.gca().set_ylim((0, 110))
 115     P.gca().set_ylabel(r"% of slowest time")
 116     P.legend(loc=8)
 117     P.title('Casting raw file data to an MR volume')
 118     P.show()
 119     
 120 
 121 def time_casting(int_data):
 122     nblk = 1 if len(int_data.shape) < 2 else int_data.shape[0]
 123     bias = (npy.random.rand(nblk) + \
 124             1j*npy.random.rand(nblk)).astype(npy.complex64)
 125     dstr = int_data.tostring()
 126     dt = npy.int16 if int_data.dtype.itemsize == 2 else npy.int32
 127     fshape = list(int_data.shape)
 128     fshape[-1] = fshape[-1]/2
 129     float_data = npy.empty(fshape, npy.complex64)
 130     # method 1: string conversion
 131     float_data.shape = (npy.product(fshape),)
 132     tstr = measure("float_data[:] = complex_fromstring(dstr, dt)", times=25)
 133     float_data.shape = fshape
 134     print "to-/from- string: ", tstr, "shape=",float_data.shape
 135 
 136     # method 2: numpy dtype magic
 137     sl = [None, slice(None)] if len(fshape)<2 else [slice(None)]*len(fshape)
 138     # need to loop since int_data need not be contiguous
 139     tnpy = measure("""
 140 for fline, iline, b in zip(float_data[sl], int_data[sl], bias):
 141     cast_to_complex_npy(fline, iline, bias=b)""", times=25)
 142     print"numpy automagic: ", tnpy
 143 
 144     # method 3: plain inline brute force!
 145     twv = measure("cast_to_complex(float_data, int_data, bias=bias)",
 146                   times=25)
 147     print"inline casting: ", twv
 148     return npy.array([tstr, tnpy, twv], npy.float64)
 149     
 150 def complex_fromstring(data, numtype):
 151     if sys.byteorder == "little":
 152         return npy.fromstring(
 153             npy.fromstring(data,numtype).byteswap().astype(npy.float32).tostring(),
 154             npy.complex64)
 155     else:
 156         return npy.fromstring(
 157 	    npy.fromstring(data,numtype).astype(npy.float32).tostring(),
 158             npy.complex64)
 159 
 160 def cast_to_complex(cplx_float, cplx_integer, bias=None):
 161     if cplx_integer.dtype.itemsize == 4:
 162         replacements = tuple(["l", "long", "SWAPLONG", "l"]*2)
 163     else:
 164         replacements = tuple(["s", "short", "SWAPSHORT", "s"]*2)
 165     if sys.byteorder == "big":
 166         replacements[-2] = replacements[-6] = "NOP"
 167 
 168     cast_code = """
 169     #define SWAPSHORT(x) ((short) ((x >> 8) | (x << 8)) )
 170     #define SWAPLONG(x) ((long) ((x >> 24) | (x << 24) | ((x & 0x00ff0000) >> 8) | ((x & 0x0000ff00) << 8)) )
 171     #define NOP(x) x
 172     
 173     unsigned short *s;
 174     unsigned long *l;
 175     float repart, impart;
 176     PyObject *itr;
 177     itr = PyArray_IterNew(py_cplx_integer);
 178     while(PyArray_ITER_NOTDONE(itr)) {
 179 
 180       // get real part
 181       %s = (unsigned %s *) PyArray_ITER_DATA(itr);
 182       repart = %s(*%s);
 183       PyArray_ITER_NEXT(itr);
 184       // get imag part
 185       %s = (unsigned %s *) PyArray_ITER_DATA(itr);
 186       impart = %s(*%s);
 187       PyArray_ITER_NEXT(itr);
 188       *(cplx_float++) = std::complex<float>(repart, impart);
 189 
 190     }
 191     """ % replacements
 192     
 193     inline(cast_code, ['cplx_float', 'cplx_integer'])
 194     if bias is not None:
 195         if len(cplx_float.shape) > 1:
 196             bsl = [slice(None)]*(len(cplx_float.shape)-1) + [None]
 197         else:
 198             bsl = slice(None)
 199         npy.subtract(cplx_float, bias[bsl], cplx_float)
 200 
 201 def cast_to_complex_npy(cplx_float, cplx_integer, bias=None):
 202     cplx_float.real[:] = cplx_integer[0::2]
 203     cplx_float.imag[:] = cplx_integer[1::2]
 204     if bias is not None:
 205         npy.subtract(cplx_float, bias, cplx_float)
 206 
 207 if __name__=="__main__":
 208     data_casting_test()
 209     multi_iter_example()