def test_mask():
    length = 100
    ramps = [np.linspace(0, 4 * np.pi, length),
             np.linspace(0, 8 * np.pi, length),
             np.linspace(0, 6 * np.pi, length)]
    image = np.vstack(ramps)
    mask_1d = np.ones((length,), dtype=np.bool)
    mask_1d[0] = mask_1d[-1] = False
    for i in range(len(ramps)):
        # mask all ramps but the i'th one
        mask = np.zeros(image.shape, dtype=np.bool)
        mask |= mask_1d.reshape(1, -1)
        mask[i, :] = False   # unmask i'th ramp
        image_wrapped = np.ma.array(np.angle(np.exp(1j * image)), mask=mask)
        image_unwrapped = unwrap_phase(image_wrapped)
        image_unwrapped -= image_unwrapped[0, 0]    # remove phase shift
        # The end of the unwrapped array should have value equal to the
        # endpoint of the unmasked ramp
        assert_array_almost_equal_nulp(image_unwrapped[:, -1], image[i, -1])
        assert_(np.ma.isMaskedArray(image_unwrapped))

        # Same tests, but forcing use of the 3D unwrapper by reshaping
        with expected_warnings(['length 1 dimension']):
            shape = (1,) + image_wrapped.shape
            image_wrapped_3d = image_wrapped.reshape(shape)
            image_unwrapped_3d = unwrap_phase(image_wrapped_3d)
            # remove phase shift
            image_unwrapped_3d -= image_unwrapped_3d[0, 0, 0]
        assert_array_almost_equal_nulp(image_unwrapped_3d[:, :, -1],
                                       image[i, -1])

def check_wrap_around(ndim, axis):
    # create a ramp, but with the last pixel along axis equalling the first
    elements = 100
    ramp = np.linspace(0, 12 * np.pi, elements)
    ramp[-1] = ramp[0]
    image = ramp.reshape(tuple([elements if n == axis else 1
                                for n in range(ndim)]))
    image_wrapped = np.angle(np.exp(1j * image))

    index_first = tuple([0] * ndim)
    index_last = tuple([-1 if n == axis else 0 for n in range(ndim)])
    # unwrap the image without wrap around
    with warnings.catch_warnings():
        # We do not want warnings about length 1 dimensions
        warnings.simplefilter("ignore")
        image_unwrap_no_wrap_around = unwrap_phase(image_wrapped, seed=0)
    print('endpoints without wrap_around:',
          image_unwrap_no_wrap_around[index_first],
          image_unwrap_no_wrap_around[index_last])
    # without wrap around, the endpoints of the image should differ
    assert_(abs(image_unwrap_no_wrap_around[index_first] -
                image_unwrap_no_wrap_around[index_last]) > np.pi)
    # unwrap the image with wrap around
    wrap_around = [n == axis for n in range(ndim)]
    with warnings.catch_warnings():
        # We do not want warnings about length 1 dimensions
        warnings.simplefilter("ignore")
        image_unwrap_wrap_around = unwrap_phase(image_wrapped, wrap_around,
                                                seed=0)
    print('endpoints with wrap_around:',
          image_unwrap_wrap_around[index_first],
          image_unwrap_wrap_around[index_last])
    # with wrap around, the endpoints of the image should be equal
    assert_almost_equal(image_unwrap_wrap_around[index_first],
                        image_unwrap_wrap_around[index_last])

def test_unwrap_1d():
    image = np.linspace(0, 10 * np.pi, 100)
    check_unwrap(image)
    # Masked arrays are not allowed in 1D
    with testing.raises(ValueError):
        check_unwrap(image, True)
    # wrap_around is not allowed in 1D
    with testing.raises(ValueError):
        unwrap_phase(image, True, seed=0)

def test_unwrap_3d_all_masked():
    # all elements masked
    image = np.ma.zeros((10, 10, 10))
    image[:] = np.ma.masked
    unwrap = unwrap_phase(image)
    assert_(np.ma.isMaskedArray(unwrap))
    assert_(np.all(unwrap.mask))

    # 1 unmasked element, still zero edges
    image = np.ma.zeros((10, 10, 10))
    image[:] = np.ma.masked
    image[0, 0, 0] = 0
    unwrap = unwrap_phase(image)
    assert_(np.ma.isMaskedArray(unwrap))
    assert_(np.sum(unwrap.mask) == 999)   # all but one masked
    assert_(unwrap[0, 0, 0] == 0)

def test_invalid_input():
    with testing.raises(ValueError):
        unwrap_phase(np.zeros([]))
    with testing.raises(ValueError):
        unwrap_phase(np.zeros((1, 1, 1, 1)))
    with testing.raises(ValueError):
        unwrap_phase(np.zeros((1, 1)), 3 * [False])
    with testing.raises(ValueError):
        unwrap_phase(np.zeros((1, 1)), 'False')

def check_unwrap(image, mask=None):
    image_wrapped = np.angle(np.exp(1j * image))
    if mask is not None:
        print('Testing a masked image')
        image = np.ma.array(image, mask=mask)
        image_wrapped = np.ma.array(image_wrapped, mask=mask)
    image_unwrapped = unwrap_phase(image_wrapped, seed=0)
    assert_phase_almost_equal(image_unwrapped, image)

def test_unwrap_2d_all_masked():
    # Segmentation fault when image is masked array with a all elements masked
    # GitHub issue #1347
    # all elements masked
    image = np.ma.zeros((10, 10))
    image[:] = np.ma.masked
    unwrap = unwrap_phase(image)
    assert_(np.ma.isMaskedArray(unwrap))
    assert_(np.all(unwrap.mask))

    # 1 unmasked element, still zero edges
    image = np.ma.zeros((10, 10))
    image[:] = np.ma.masked
    image[0, 0] = 0
    unwrap = unwrap_phase(image)
    assert_(np.ma.isMaskedArray(unwrap))
    assert_(np.sum(unwrap.mask) == 99)    # all but one masked
    assert_(unwrap[0, 0] == 0)

def load_folder_pupil(folder, flag = 'mod', radius = 49 ):
    # load and plot all the pupil functions in the chosen folder.
    pupil_list = glob.glob(folder + '*'+ flag+'*.npy')
    pupil_list.sort(key = os.path.getmtime)
    for pname in pupil_list:
        session_name = os.path.basename(pname).split('.')[0]
        pupil_name = session_name + '_pupil'
        print(session_name)
        raw_pupil = unwrap_phase(np.load(pname))/(2*np.pi) # convert to wavelength
        pupil = pupil_crop(raw_pupil,mask = True, rad = radius)
        figp = pupil_showcross(pupil)
        figp.savefig(folder+pupil_name)
        plt.close()

def single_2Dgrating_analyses(img, img_ref=None, harmonicPeriod=None,
                              unwrapFlag=True, plotFlag=True, verbose=False):
    """
    Function to process the data of single 2D grating Talbot imaging. It
    wraps other functions in order to make all the process transparent

    """

    # Obtain Harmonic images
    h_img = single_grating_harmonic_images(img, harmonicPeriod,
                                           plotFlag=plotFlag,
                                           verbose=verbose)

    if img_ref is not None:  # relative wavefront

        h_img_ref = single_grating_harmonic_images(img_ref, harmonicPeriod,
                                                   plotFlag=plotFlag,
                                                   verbose=verbose)

        int00 = np.abs(h_img[0])/np.abs(h_img_ref[0])
        int01 = np.abs(h_img[1])/np.abs(h_img_ref[1])
        int10 = np.abs(h_img[2])/np.abs(h_img_ref[2])

        if unwrapFlag is True:

            arg01 = (unwrap_phase(np.angle(h_img[1]), seed=72673) -
                     unwrap_phase(np.angle(h_img_ref[1]), seed=72673))

            arg10 = (unwrap_phase(np.angle(h_img[2]), seed=72673) -
                     unwrap_phase(np.angle(h_img_ref[2]), seed=72673))

        else:
            arg01 = np.angle(h_img[1]) - np.angle(h_img_ref[1])
            arg10 = np.angle(h_img[2]) - np.angle(h_img_ref[2])

    else:  # absolute wavefront

        int00 = np.abs(h_img[0])
        int01 = np.abs(h_img[1])
        int10 = np.abs(h_img[2])

        if unwrapFlag is True:

            arg01 = unwrap_phase(np.angle(h_img[1]), seed=72673)
            arg10 = unwrap_phase(np.angle(h_img[2]), seed=72673)
        else:
            arg01 = np.angle(h_img[1])
            arg10 = np.angle(h_img[2])

    if unwrapFlag is True:  # remove pi jump
        arg01 -= int(np.round(np.mean(arg01/np.pi)))*np.pi
        arg10 -= int(np.round(np.mean(arg10/np.pi)))*np.pi

    darkField01 = int01/int00
    darkField10 = int10/int00

    return [int00, int01, int10,
            darkField01, darkField10,
            arg01, arg10]

def wave_retrieval(data,r_data=None,filt=None,STEPS=False,s=0.1,ARGS=False):
  # go to fourier space
  f_data = np.fft.fft2(data)
  # get the k-vector (lots of finagling here for fft shift crap)
  if r_data is None:
    K = get_wave(f_data)
    kx,ky = np.fft.fftfreq(f_data.shape[0]),np.fft.fftfreq(f_data.shape[1])
    k = kx[K[0]],ky[K[1]]
    nx,ny = data.shape
    Y,X = np.meshgrid(np.arange(nx),np.arange(ny),sparse=False,indexing='xy')
  #  phplot.imageshow(data)
  #  phplot.imageshow(np.real(r_data))
  #  phplot.imageshow(np.real(r_data*np.exp(+1j*2.*np.pi*(k[0]*X+k[1]*Y))))
    # reference beam
    r_data = np.exp(+1j*2.*np.pi*(k[0]*X+k[1]*Y))


  # filter k-space with a gaussian around data of interest
  if filt is None:
    try:
      gf_data = f_gaussian(f_data,K)*f_data
    except UnboundLocalError:
      K = get_wave(f_data)
    finally:
      gf_data = f_gaussian(f_data,K)*f_data
      filt = np.ones(f_data.shape)
      filt = f_gaussian(filt,K)*filt
#      band = get_band(data,K)
#      filt = f_bandpass(filt,band)
#      filt = hilbert_transform(filt,'x')
  else:
    gf_data = filt*f_data
  # check it out!!
#  phplot.imageshow(phplot.dB_data(gauss_filter))
  # go back to xy space
  g_data = np.fft.ifft2(gf_data)

  o_data = g_data*r_data  # retrieve the wrapped phase from the digital reference beam
  wrapped_phase = np.arctan2(o_data.imag,o_data.real)
  # unwrap the phase (may be unneccesary)
  unwrapped_phase = unwrap_phase(wrapped_phase)
  phase = unwrapped_phase
  #phase = wrapped_phase

  if STEPS:
    return data,f_data,gauss_filter,wrapped_phase,unwrapped_phase,plane,phase

  if ARGS:
    return -phase,r_data,filt

  return -phase

def compute_local_phase_field(camp):
    """ Compute local phase of each cell
    """
    width, height, depth = camp.shape

    # compute thetas
    tau = compute_tau(camp)

    theta = np.empty((width, height)).tolist()
    for j in range(width):
        for i in range(height):
            cur = compute_phase_variable(camp[i, j], tau)
            theta[i][j] = np.array(cur)
    theta = np.rollaxis(np.array(theta), 2, 0)

    # fix phase jumps greater than PI
    for i in range(len(theta)):
        theta[i] = unwrap_phase(theta[i])

    return theta

def parse_enz_solution(cfg, betak):
    if betak.ndim == 2:
        gpf = cfg.cpsf.czern.eval_grid(betak[:, 0])
    else:
        gpf = cfg.cpsf.czern.eval_grid(betak)
    wrph = np.arctan2(gpf.imag, gpf.real)
    wrph = np.mod(wrph, 2*np.pi) - np.pi

    ut1 = time()
    unph = unwrap_phase(wrph.reshape(
        (cfg.fit_L, cfg.fit_K), order='F')).ravel(order='F')
    ut2 = time()

    ft1 = time()
    alpha_hat = cfg.phase_fit.fit(unph)
    ft2 = time()

    alpha_hat[0] = 0.0

    return alpha_hat, ft2 - ft1, ut2 - ut1, wrph, unph

def hilbert_retrieval(data,plane=None,filt=None,STEPS=False,direction='x',ARGS=False):
  # fourier transform data and filter
  window_f = np.hanning
  w2 = np.sqrt(np.outer(window_f(data.shape[0]),window_f(data.shape[1])))
  f_data = np.fft.fft2(data)
  if filt is None:
    band = get_band(f_data)
    filt_f_data = f_bandpass(f_data,band)
    hilbert_f_data = hilbert_transform(filt_f_data,direction)
    filt = np.ones(f_data.shape)
    filt = f_bandpass(filt,band)
    filt = hilbert_transform(filt,direction)
  else:
    hilbert_f_data = f_data *filt

  # hilbert-transform
#  phplot.imageshow(phplot.dB_data(filt_f_data))
  hilbert_data = np.fft.ifft2(hilbert_f_data)

  # retrieve the wrapped phase
  wrapped_phase = np.arctan2(hilbert_data.imag,hilbert_data.real)

  # unwrap the phase
  unwrapped_phase = unwrap_phase(wrapped_phase)

  # flatten the data (if required - some methods can use the same level for all data)
  if plane is None:
    plane = diff_level(unwrapped_phase)
  phase = unwrapped_phase - plane

  if STEPS:
    return data,f_data,filt_f_data,hilbert_f_data,hilbert_data,wrapped_phase,unwrapped_phase,plane,phase

  if ARGS:
    return -phase,plane,filt
  return -phase

 def unwrap(self):
     unwrapped = unwrap_phase(self._PF.phase)
     return unwrapped

"""

import numpy as np
from matplotlib import pyplot as plt
from skimage import data, img_as_float, color, exposure
from skimage.restoration import unwrap_phase


# Load an image as a floating-point grayscale
image = color.rgb2gray(img_as_float(data.chelsea()))
# Scale the image to [0, 4*pi]
image = exposure.rescale_intensity(image, out_range=(0, 4 * np.pi))
# Create a phase-wrapped image in the interval [-pi, pi)
image_wrapped = np.angle(np.exp(1j * image))
# Perform phase unwrapping
image_unwrapped = unwrap_phase(image_wrapped)

fig, ax = plt.subplots(2, 2)
ax1, ax2, ax3, ax4 = ax.ravel()

fig.colorbar(ax1.imshow(image, cmap='gray', vmin=0, vmax=4 * np.pi), ax=ax1)
ax1.set_title('Original')

fig.colorbar(ax2.imshow(image_wrapped, cmap='gray', vmin=-np.pi, vmax=np.pi), ax=ax2)
ax2.set_title('Wrapped phase')

fig.colorbar(ax3.imshow(image_unwrapped, cmap='gray'), ax=ax3)
ax3.set_title('After phase unwrapping')

fig.colorbar(ax4.imshow(image_unwrapped - image, cmap='gray'), ax=ax4)
ax4.set_title('Unwrapped minus original')

def diagnosticPlot(solver):
    """
    Returns a matplotlib figure
    """
    ur = pint.UnitRegistry()
    MASK_CUTOFF = 1.0
    # N_PHASE_TICKS = 6
    NUM_COLOURS = 100
    # DISP_FRACTION = 1/10
    KS_FRACTION = 1/4
    ur = UnitRegistry()
    x, y = solver.grid.getSpatialGrid(scaled=False)
    x *= 1e6
    y *= 1e6

    x_ax = x[0, :]
    y_ax = y[:, 0]
    psi = solver.psi.get()
    # n = solver.n.get()
    P = solver.Pdt.get()
    P /= solver.dt

    N = psi.shape[0]
    time = solver.times
    energy = solver.energy
    number = solver.number
    pumpSlice = P[N/2, :]
    PthScaled = solver.Pth / np.max(P)
    pumpSlice /= np.max(P)

    rMiddle = abs(x[N/2, np.argmax(pumpSlice)])
    density = solver.getDensity(scaled=True)
    pumpCircle = plt.Circle((0.0, 0.0), radius=rMiddle, color='w',
                            linestyle='dashed', fill=False)

    fig, axes = plt.subplots(2, 3)
    time = solver.times
    energy = solver.energy
    number = solver.number

    density = solver.getDensity(scaled=True)
    pumpCircle = plt.Circle((0.0, 0.0), radius=rMiddle, color='w',
                            linestyle='dashed', fill=False)

    fig, axes = plt.subplots(2, 3)

    fig, axes = plt.subplots(2, 3)
    p0 = axes[0, 0].contourf(x, y, density, NUM_COLOURS)

    # RS density
    axes[0, 0].set_title("Density $\mu m ^{-2}$")
    # axes[0, 0].set_aspect('equal', 'box')
    axes[0, 0].set_aspect(1./axes[0, 0].get_data_ratio())
    axes[0, 0].add_artist(pumpCircle)
    plt.colorbar(p0, ax=axes[0, 0], fraction=0.046, pad=0.04)

    # KS density
    ksDensity = np.absolute(fft.fftshift(fft.fft2(psi)))
    ksLowerLim = int(N/2 - KS_FRACTION * N / 2)
    ksUpperLim = int(N/2 + KS_FRACTION * N / 2)
    p1 = axes[1, 0].contourf(x_ax[ksLowerLim:ksUpperLim],
                             y_ax[ksLowerLim:ksUpperLim],
                             ksDensity[ksLowerLim:ksUpperLim,
                                       ksLowerLim:ksUpperLim] /
                             np.max(ksDensity), NUM_COLOURS)
    axes[1, 0].set_title("KS Density")
    axes[1, 0].set_aspect('equal')
    plt.colorbar(p1, ax=axes[1, 0], fraction=0.046, pad=0.04)

    # Phase
    # p2 = axes[0, 1].contourf(x, y, np.angle(psi) + np.pi)
    phase = unwrap_phase(np.angle(psi))
    # Mask the values that are out the damping region are ignored
    mask = solver.damping.get() < MASK_CUTOFF
    phase = np.ma.array(phase, mask=mask) / np.pi
    # Set the minumum phase to zero
    phase -= np.min(phase)
    # maxPhase = np.ceil(np.max(phase) / np.pi)
    # minPhase = np.floor(np.min(phase) / np.pi)
    # step = np.ceil((maxPhase - minPhase) / N_PHASE_TICKS)
    # ticks = np.pi * np.arange(minPhase, maxPhase + 1, step=step)
    # labels = ["$%d \pi$" % i for i in np.arange(minPhase, maxPhase + 1)]
    p2 = axes[0, 1].contourf(x, y, phase, NUM_COLOURS)
    axes[0, 1].set_title("Phase")
    # axes[0, 1].set_aspect('equal', 'box')
    axes[0, 1].set_aspect(1./axes[0, 1].get_data_ratio())
    pumpCircle = plt.Circle((0.0, 0.0), radius=rMiddle, color='w',
                            linestyle='dashed', fill=False)
    axes[0, 1].add_artist(pumpCircle)
    cbar2 = plt.colorbar(p2, ax=axes[0, 1], fraction=0.046, pad=0.04)
    labels = [t.get_text() + '$\pi$' for t in cbar2.ax.get_yticklabels()]
    # cbar2 = plt.colorbar(p2, ax=axes[0, 1],
    #                      ticks=ticks,
    #                      fraction=0.046, pad=0.04)
    cbar2.ax.set_yticklabels(labels)
    # cbar2 = plt.colorbar(p2, ax=axes[0, 1],
    #                      ticks=np.pi * np.array([0, 1.0, 2.0]),
    #                      fraction=0.046, pad=0.04)
    # cbar2.ax.set_yticklabels(['0', '$\pi$', '$2\pi$'])

#.........这里部分代码省略.........

templateX, templateY, templateWidth  = locate_glowplug(img, approxTemplate, debug=False)
template = get_autolite_template(templateWidth+templateSlop, img_h/2)

## perform analysis
if not os.path.exists('results'):
    os.makedirs('results')

frameNums = map(str,frameNums)
for i in frameNums:
    imageFilename = ('cropped/' + imageFileBase + i + '_crop.jpg')

    orig = WTP.imagefile2dat(imageFilename)

    wrapped = WTP.performWTP(wavelet_type = wavelet_type, ridge_alg = ridge_alg, 
                        starting_scale = starting_scale, ending_scale=ending_scale, use_FFT=use_FFT, Morlet_sigma=Morlet_sigma)
    unwrapped = unwrap_phase(wrapped)

    phase = unbias_image(unwrapped)

    masked = remove_glowplug(phase, templateX-templateSlop/2, templateY, template)

    x, y, centered  = center_image_x(masked, templateX + templateWidth/2, templateY, dx)

    r, f_abel, deriv, smoothed  = perform_abel_discontinuous(x, centered)

    delta_n = f_abel*lam / (2*np.pi)
    rho= (n_air - delta_n - 1)/K_air

    fileroot = ('results/' + imageFileBase + i + '_crop')
    np.savetxt(fileroot+"_orig.csv", orig, delimiter=",")
    np.savetxt(fileroot+"_wrapped.csv", wrapped, delimiter=",")

 def complex(self, new):
     self._complex = new
     self._amplitude = abs(new)
     self._phase = unwrap_phase(np.angle(new))

    wpu.plot_slide_colorbar(intensity,
                            title='Intensity',
                            xlabel=r'x [$\mu m$]',
                            ylabel=r'y [$\mu m$]',
                            extent=wpu.extent_func(dpc_1d, pixelSize)*1e6)

    # %% Dark Field

    wpu.plot_slide_colorbar(dk_field, title='Dark Field',
                            xlabel=r'x [$\mu m$]',
                            ylabel=r'y [$\mu m$]',
                            extent=wpu.extent_func(dpc_1d, pixelSize)*1e6)

    # %% DPC

    dpc_1d = unwrap_phase(dpc_1d)
    wpu.plot_slide_colorbar(dpc_1d/np.pi/2.0,
                            title=r'DPC [$\pi rad$]',
                            xlabel=r'x [$\mu m$]',
                            ylabel=r'y [$\mu m$]',
                            extent=wpu.extent_func(dpc_1d, pixelSize)*1e6)

    # %%
    xx, yy = wpu.realcoordmatrix(dpc_1d.shape[1], pixelSize,
                                 dpc_1d.shape[0], pixelSize)
    wpu.plot_profile(xx*1e3, yy*1e3, dpc_1d/np.pi/2.0,
                     xlabel='[mm]', ylabel='[mm]')

    # %% chi2

    plt.figure()

def test_unwrap_2d_compressed_mask():
    # ValueError when image is masked array with a compressed mask (no masked
    # elments).  GitHub issue #1346
    image = np.ma.zeros((10, 10))
    unwrap = unwrap_phase(image)
    assert_(np.all(unwrap == 0))