unwrap: accept (and return if given) an masked array as input argument

unwrap2D/test_unwrap.py

@@ -4,9 +4,10 @@ from unwrap2D import unwrap2D
 import numpy as np
 from numpy import outer, arange, ones, abs, empty, power, indices
 
+import numpy.ma as ma
 
 def test_unwrap2D():
-    nx, ny = 10, 10
+    nx, ny = 32, 32
     x = np.arange(nx)
     y = np.arange(ny)
     x.shape = (1,-1)
@@ -14,10 +15,15 @@ def test_unwrap2D():
 
     z = np.exp(1j*x*0.2*np.pi) * np.exp(1j*y*0.1*np.pi)
     phi_w = np.angle(z)
+    phi = unwrap2D(phi_w)
+
     mask = 0*np.ones((nx, ny), dtype = np.uint8)
-    mask[4:6, 5:7] = 1
-    phi = unwrap2D(phi_w.astype(np.float32), mask)
-    return phi_w/(np.pi*2), np.asarray(phi)/(np.pi*2)
+    mask[4:16, 4:16] = 1
+    phi_w_ma = ma.array(phi_w, dtype = np.float32, mask = mask)
+    phi_ma = unwrap2D(phi_w_ma)
+
+    return (phi_w/(np.pi*2), phi/(np.pi*2),
+            phi_w_ma/(np.pi*2), phi_ma/(np.pi*2),)
     
 
 # class test_unwrap(TestCase):
@@ -46,8 +52,15 @@ def test_unwrap2D():
 if __name__=="__main__":
     #NumpyTest().run()
     import matplotlib.pyplot as plt
-    p,p2 = test_unwrap2D()
+    p1,p2,p3,p4 = test_unwrap2D()
     plt.clf()
+    plt.subplot(221)
+    plt.imshow(p1,interpolation = 'nearest')
+    plt.subplot(222)
     plt.imshow(p2, interpolation = 'nearest')
+    plt.subplot(223)
+    plt.imshow(p3, interpolation = 'nearest')
+    plt.subplot(224)
+    plt.imshow(p4, interpolation = 'nearest')
     plt.draw()
     plt.show()

unwrap2D/unwrap2D.pyx

@@ -1,12 +1,31 @@
 import numpy as np
 cimport numpy as np
 
-cdef extern int unwrap(float* WrappedImage, float* UnwrappedImage, unsigned char* input_mask, 
+import numpy.ma
+
+cdef extern int unwrap(float* wrapped_image, 
+                       float* unwrapped_image, 
+                       unsigned char* input_mask, 
                        int image_width, int image_height,
                        int wrap_around_x, int wrap_around_y)
 
-def unwrap2D(float[:,::1] array, unsigned char[:,::1] mask,
-             wrap_around_x = False, wrap_around_y = False):
+def unwrap2D(input, wrap_around_x = False, wrap_around_y = False):
+    
+    masked_array = numpy.ma.asarray(input, dtype = np.float32)
+    unwrapped_array = _unwrap2D(masked_array.data, 
+                                numpy.ma.getmaskarray(masked_array).astype(np.uint8),
+                                wrap_around_x, wrap_around_y)
+    if numpy.ma.isarray(input):
+        return numpy.ma.array(unwrapped_array, mask = input.mask)
+    else:
+        return unwrapped_array
+
+    #TODO: set_fill to minimum value
+    
+
+cdef _unwrap2D(float[:,::1] array, 
+               unsigned char[:,::1] mask,
+               wrap_around_x, wrap_around_y):
     cdef float[:,::1] unwrapped_array = np.empty_like(array)
     cdef int h = array.shape[0]
     cdef int w = array.shape[1]