def fftconvolve(in1, in2, mode="full"):
    """Convolve two N-dimensional arrays using FFT. See convolve.
    
    """
    s1 = array(in1.shape)
    s2 = array(in2.shape)
    if (s1.dtype.char in ['D','F']) or (s2.dtype.char in ['D', 'F']):
        cmplx=1
    else: cmplx=0
    size = s1+s2-1
    IN1 = fftn(in1,size)
    IN1 *= fftn(in2,size)
    ret = ifftn(IN1)
    del IN1
    if not cmplx:
        ret = real(ret)
    if mode == "full":
        return ret
    elif mode == "same":
        if product(s1,axis=0) > product(s2,axis=0):
            osize = s1
        else:
            osize = s2
        return _centered(ret,osize)
    elif mode == "valid":
        return _centered(ret,abs(s2-s1)+1)

def fftconvolve(in1, in2, in3=None, mode="full"):
    """Convolve two N-dimensional arrays using FFT. See convolve.

    copied from scipy, but here used to try out inverse filter
    doesn't work or I can't get it to work
    """
    s1 = array(in1.shape)
    s2 = array(in2.shape)
    complex_result = (np.issubdtype(in1.dtype, np.complex) or
                      np.issubdtype(in2.dtype, np.complex))
    size = s1+s2-1

    # Always use 2**n-sized FFT
    fsize = 2**np.ceil(np.log2(size))
    IN1 = fftn(in1,fsize)
    #IN1 *= fftn(in2,fsize)
    IN1 /= fftn(in2,fsize)  # use inverse filter
    # note the inverse is elementwise not matrix inverse
    # is this correct, NO  doesn't seem to work
    fslice = tuple([slice(0, int(sz)) for sz in size])
    ret = ifftn(IN1)[fslice].copy()
    del IN1
    if not complex_result:
        ret = ret.real
    if mode == "full":
        return ret
    elif mode == "same":
        if product(s1,axis=0) > product(s2,axis=0):
            osize = s1
        else:
            osize = s2
        return _centered(ret,osize)
    elif mode == "valid":
        return _centered(ret,abs(s2-s1)+1)

    def bench_random(self):
        from numpy.fft import fftn as numpy_fftn
        print()
        print('    Multi-dimensional Fast Fourier Transform')
        print('===================================================')
        print('          |    real input     |   complex input    ')
        print('---------------------------------------------------')
        print('   size   |  scipy  |  numpy  |  scipy  |  numpy ')
        print('---------------------------------------------------')
        for size,repeat in [((100,100),100),((1000,100),7),
                            ((256,256),10),
                            ((512,512),3),
                            ]:
            print('%9s' % ('%sx%s'%size), end=' ')
            sys.stdout.flush()

            for x in [random(size).astype(double),
                      random(size).astype(cdouble)+random(size).astype(cdouble)*1j
                      ]:
                y = fftn(x)
                #if size > 500: y = fftn(x)
                #else: y = direct_dft(x)
                assert_array_almost_equal(fftn(x),y)
                print('|%8.2f' % measure('fftn(x)',repeat), end=' ')
                sys.stdout.flush()

                assert_array_almost_equal(numpy_fftn(x),y)
                print('|%8.2f' % measure('numpy_fftn(x)',repeat), end=' ')
                sys.stdout.flush()

            print(' (secs for %s calls)' % (repeat))

        sys.stdout.flush()

 def test_shape_argument(self):
     small_x = [[1,2,3],[4,5,6]]
     large_x1 = [[1,2,3,0],[4,5,6,0],[0,0,0,0],[0,0,0,0]]
     y = fftn(small_x,shape=(4,4))
     assert_array_almost_equal (y,fftn(large_x1))
     y = fftn(small_x,shape=(3,4))
     assert_array_almost_equal (y,fftn(large_x1[:-1]))

def fftconvolve3(in1, in2=None, in3=None, mode="full"):
    """Convolve two N-dimensional arrays using FFT. See convolve.

    for use with arma  (old version: in1=num in2=den in3=data

    * better for consistency with other functions in1=data in2=num in3=den
    * note in2 and in3 need to have consistent dimension/shape
      since I'm using max of in2, in3 shapes and not the sum

    copied from scipy.signal.signaltools, but here used to try out inverse
    filter doesn't work or I can't get it to work

    2010-10-23
    looks ok to me for 1d,
    from results below with padded data array (fftp)
    but it doesn't work for multidimensional inverse filter (fftn)
    original signal.fftconvolve also uses fftn
    """
    if (in2 is None) and (in3 is None):
        raise ValueError('at least one of in2 and in3 needs to be given')
    s1 = np.array(in1.shape)
    if not in2 is None:
        s2 = np.array(in2.shape)
    else:
        s2 = 0
    if not in3 is None:
        s3 = np.array(in3.shape)
        s2 = max(s2, s3) # try this looks reasonable for ARMA
        #s2 = s3


    complex_result = (np.issubdtype(in1.dtype, np.complex) or
                      np.issubdtype(in2.dtype, np.complex))
    size = s1+s2-1

    # Always use 2**n-sized FFT
    fsize = 2**np.ceil(np.log2(size))
    #convolve shorter ones first, not sure if it matters
    if not in2 is None:
        IN1 = fft.fftn(in2, fsize)
    if not in3 is None:
        IN1 /= fft.fftn(in3, fsize)  # use inverse filter
    # note the inverse is elementwise not matrix inverse
    # is this correct, NO  doesn't seem to work for VARMA
    IN1 *= fft.fftn(in1, fsize)
    fslice = tuple([slice(0, int(sz)) for sz in size])
    ret = fft.ifftn(IN1)[fslice].copy()
    del IN1
    if not complex_result:
        ret = ret.real
    if mode == "full":
        return ret
    elif mode == "same":
        if np.product(s1,axis=0) > np.product(s2,axis=0):
            osize = s1
        else:
            osize = s2
        return trim_centered(ret,osize)
    elif mode == "valid":
        return trim_centered(ret,abs(s2-s1)+1)

    def psf_calc(self, psf, kz, data_size):
        '''Pre calculate OTFs etc ...'''
        g = psf;

        self.height = data_size[0]
        self.width  = data_size[1]
        self.depth  = data_size[2]

        (x,y,z) = mgrid[-floor(self.height/2.0):(ceil(self.height/2.0)), -floor(self.width/2.0):(ceil(self.width/2.0)), -floor(self.depth/2.0):(ceil(self.depth/2.0))]
        
       
        gs = shape(g);
        g = g[int(floor((gs[0] - self.height)/2)):int(self.height + floor((gs[0] - self.height)/2)), int(floor((gs[1] - self.width)/2)):int(self.width + floor((gs[1] - self.width)/2)), int(floor((gs[2] - self.depth)/2)):int(self.depth + floor((gs[2] - self.depth)/2))]
	
        g = abs(ifftshift(ifftn(abs(fftn(g)))));
        g = (g/sum(sum(sum(g))));
	
        self.g = g;
        
        self.H = cast['f'](fftn(g));
        self.Ht = cast['f'](ifftn(g));

        tk = 2*kz*z
        
        t = g*exp(1j*tk)
        self.He = cast['F'](fftn(t));    
        self.Het = cast['F'](ifftn(t));

        tk = 2*tk
        
        t = g*exp(1j*tk)
        self.He2 = cast['F'](fftn(t));    
        self.He2t = cast['F'](ifftn(t));

def fftconvolve(in1, in2, mode="full"):
    """Convolve two N-dimensional arrays using FFT. See convolve.

    """
    s1 = array(in1.shape)
    s2 = array(in2.shape)
    complex_result = (np.issubdtype(in1.dtype, np.complex) or
                      np.issubdtype(in2.dtype, np.complex))
    size = s1 + s2 - 1

    # Always use 2**n-sized FFT
    fsize = 2 ** np.ceil(np.log2(size))
    IN1 = fftn(in1, fsize)
    IN1 *= fftn(in2, fsize)
    fslice = tuple([slice(0, int(sz)) for sz in size])
    ret = ifftn(IN1)[fslice].copy()
    del IN1
    if not complex_result:
        ret = ret.real
    if mode == "full":
        return ret
    elif mode == "same":
        if product(s1, axis=0) > product(s2, axis=0):
            osize = s1
        else:
            osize = s2
        return _centered(ret, osize)
    elif mode == "valid":
        return _centered(ret, abs(s2 - s1) + 1)

def fftconvolve(in1, in2, mode="full"):
    """Convolve two N-dimensional arrays using FFT. See convolve.

    """
    s1 = array(in1.shape)
    s2 = array(in2.shape)
    complex_result = (np.issubdtype(in1.dtype, np.complex) or
                      np.issubdtype(in2.dtype, np.complex))
    size = s1+s2-1
    IN1 = fftn(in1,size)
    IN1 *= fftn(in2,size)
    ret = ifftn(IN1)
    del IN1
    if not complex_result:
        ret = ret.real
    if mode == "full":
        return ret
    elif mode == "same":
        if product(s1,axis=0) > product(s2,axis=0):
            osize = s1
        else:
            osize = s2
        return _centered(ret,osize)
    elif mode == "valid":
        return _centered(ret,abs(s2-s1)+1)

def fft_correlation(in1, in2, normalize=False):
    """Correlation of two N-dimensional arrays using FFT.

    Adapted from scipy's fftconvolve.

    Parameters
    ----------
    in1, in2 : array
    normalize: bool
        If True performs phase correlation

    """
    s1 = np.array(in1.shape)
    s2 = np.array(in2.shape)
    size = s1 + s2 - 1
    # Use 2**n-sized FFT
    fsize = 2 ** np.ceil(np.log2(size))
    IN1 = fftn(in1, fsize)
    IN1 *= fftn(in2, fsize).conjugate()
    if normalize is True:
        ret = ifftn(np.nan_to_num(IN1 / np.absolute(IN1))).real.copy()
    else:
        ret = ifftn(IN1).real.copy()
    del IN1
    return ret

    def test_size_accuracy_large(self, size):
        x = np.random.rand(size, 3) + 1j*np.random.rand(size, 3)
        y1 = fftn(x.real.astype(np.float32))
        y2 = fftn(x.real.astype(np.float64)).astype(np.complex64)

        assert_equal(y1.dtype, np.complex64)
        assert_array_almost_equal_nulp(y1, y2, 2000)

    def test_float16_input_large(self, size):
        x = np.random.rand(size, 3) + 1j*np.random.rand(size, 3)
        y1 = fftn(x.real.astype(np.float16))
        y2 = fftn(x.real.astype(np.float64)).astype(np.complex64)

        assert_equal(y1.dtype, np.complex64)
        assert_array_almost_equal_nulp(y1, y2, 2e6)

def shear_fft2d(array,x,time,type='rho',rm_shear=0):
    #array of size (nx,ny,nz,3) if type='v'
    #array of size (nx,ny,nz) if type='rho'

    q=1.5e0 ; Omega=1.e-3
    Lx=1. ; Ly=2.*np.arcsin(1.e0) ; Lz=1.
    twopi=4.*np.arcsin(1.e0)

    if rank(array)==3:
        nx,ny,nz=shape(array)
    if rank(array)==4:
        nx,ny,nz,ndim=shape(array)

    # Remove background velocity shear if needed...
    if (rm_shear==1):
        array[:,:,:,1]=remove_shear(array[:,:,:,1],x)

    # Unshear data in real space...
    array=unshear(array,x,time,type=type)

    # Compute FFT...
    fft_array=array
    if (type=='rho'):
        for k in range(nz):
            fft_array[:,:,k]=fftn(array[:,:,k]-np.mean(array[:,:,k]))
    else:
        for k in range(nz):
            fft_array[:,:,k,0]=fftn(array[:,:,k,0])
            fft_array[:,:,k,1]=fftn(array[:,:,k,1])
            fft_array[:,:,k,2]=fftn(array[:,:,k,2])

    return fft_array

def customfftconvolve(in1, in2, mode="full", types=('','')):
  """ Pretty much the same as original fftconvolve, but supports
      having operands as fft already 
  """

  in1 = asarray(in1)
  in2 = asarray(in2)

  if in1.ndim == in2.ndim == 0:  # scalar inputs
    return in1 * in2
  elif not in1.ndim == in2.ndim:
    raise ValueError("in1 and in2 should have the same dimensionality")
  elif in1.size == 0 or in2.size == 0:  # empty arrays
    return array([])

  s1 = array(in1.shape)
  s2 = array(in2.shape)
  complex_result = False
  #complex_result = (np.issubdtype(in1.dtype, np.complex) or
  #                  np.issubdtype(in2.dtype, np.complex))
  shape = s1 + s2 - 1
  
  if mode == "valid":
    _check_valid_mode_shapes(s1, s2)

  # Speed up FFT by padding to optimal size for FFTPACK
  fshape = [_next_regular(int(d)) for d in shape]
  fslice = tuple([slice(0, int(sz)) for sz in shape])

  if not complex_result:
    if types[0] == 'fft':
      fin1 = in1#_unfold_fft(in1, fshape)
    else:
      fin1 = rfftn(in1, fshape)

    if types[1] == 'fft':
      fin2 = in2#_unfold_fft(in2, fshape)
    else:
      fin2 = rfftn(in2, fshape)
    ret = irfftn(fin1 * fin2, fshape)[fslice].copy()
  else:
    if types[0] == 'fft':
      fin1 = _unfold_fft(in1, fshape)
    else:
      fin1 = fftn(in1, fshape)
    if types[1] == 'fft':
      fin2 = _unfold_fft(in2, fshape)
    else:
      fin2 = fftn(in2, fshape)
    ret = ifftn(fin1 * fin2)[fslice].copy()

  if mode == "full":
    return ret
  elif mode == "same":
    return _centered(ret, s1)
  elif mode == "valid":
    return _centered(ret, s1 - s2 + 1)
  else:
    raise ValueError("Acceptable mode flags are 'valid',"
                     " 'same', or 'full'.")

 def test_definition_float16(self):
     x = [[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]]
     y = fftn(np.array(x, np.float16))
     assert_equal(y.dtype, np.complex64)
     y_r = np.array(fftn(x), np.complex64)
     assert_array_almost_equal_nulp(y, y_r)

    def test_definition(self):
        x = [[1,2,3],[4,5,6],[7,8,9]]
        y = fftn(np.array(x, np.float32))
        if not y.dtype == np.complex64:
            raise ValueError("double precision output with single precision")

        y_r = np.array(fftn(x), np.complex64)
        assert_array_almost_equal_nulp(y, y_r)

def bgtensor(img, lsigma, rho=0.2):
    eps = 1e-12
    fimg = fftn(img, overwrite_x=True)

    for s in lsigma:
        jvbuffer = bgkern3(kerlen=math.ceil(s)*6+1, sigma=s, rho=rho)
        jvbuffer = fftn(jvbuffer, shape=fimg.shape, overwrite_x=True) * fimg
        fimg = ifftn(jvbuffer, overwrite_x=True)
        yield hessian3(np.real(fimg))

    def test_invalid_sizes(self):
        with assert_raises(ValueError,
                           match="invalid number of data points"
                           r" \(\[1 0\]\) specified"):
            fftn([[]])

        with assert_raises(ValueError,
                           match="invalid number of data points"
                           r" \(\[ 4 -3\]\) specified"):
            fftn([[1, 1], [2, 2]], (4, -3))

def _fftconvolve(in1, in2, mode="full", axis=None):
    """ Convolve two N-dimensional arrays using FFT. See convolve.

    This is a fix of scipy.signal.fftconvolve, adding an axis argument and
    importing locally the stuff only needed for this function
    
    """
    #Locally import stuff only required for this:
    from scipy.fftpack import fftn, fft, ifftn, ifft
    from scipy.signal.signaltools import _centered
    from numpy import array, product


    s1 = array(in1.shape)
    s2 = array(in2.shape)
    complex_result = (np.issubdtype(in1.dtype, np.complex) or
                      np.issubdtype(in2.dtype, np.complex))

    if axis is None:
        size = s1+s2-1
        fslice = tuple([slice(0, int(sz)) for sz in size])
    else:
        equal_shapes = s1==s2
        # allow equal_shapes[axis] to be False
        equal_shapes[axis] = True
        assert equal_shapes.all(), 'Shape mismatch on non-convolving axes'
        size = s1[axis]+s2[axis]-1
        fslice = [slice(l) for l in s1]
        fslice[axis] = slice(0, int(size))
        fslice = tuple(fslice)

    # Always use 2**n-sized FFT
    fsize = 2**np.ceil(np.log2(size))
    if axis is None:
        IN1 = fftn(in1,fsize)
        IN1 *= fftn(in2,fsize)
        ret = ifftn(IN1)[fslice].copy()
    else:
        IN1 = fft(in1,fsize,axis=axis)
        IN1 *= fft(in2,fsize,axis=axis)
        ret = ifft(IN1,axis=axis)[fslice].copy()
    if not complex_result:
        del IN1
        ret = ret.real
    if mode == "full":
        return ret
    elif mode == "same":
        if product(s1,axis=0) > product(s2,axis=0):
            osize = s1
        else:
            osize = s2
        return _centered(ret,osize)
    elif mode == "valid":
        return _centered(ret,abs(s2-s1)+1)

    def test_shape_axes_argument(self):
        small_x = [[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9]]
        large_x1 = array([[1, 2, 3, 0],
                          [4, 5, 6, 0],
                          [7, 8, 9, 0],
                          [0, 0, 0, 0]])
        y = fftn(small_x, shape=(4, 4), axes=(-2, -1))
        assert_array_almost_equal(y, fftn(large_x1))
        y = fftn(small_x, shape=(4, 4), axes=(-1, -2))

        assert_array_almost_equal(y, swapaxes(
            fftn(swapaxes(large_x1, -1, -2)), -1, -2))

def GetXi2d(array,x,time,type,periodic):
    #Return correlation function

    ciso2=1.e-6

    # Get dims
    if rank(array)==3:
        nx,ny,nz=shape(array)
    if rank(array)==4:
        nx,ny,nz,ndim=shape(array)

    if ((type=='rho') or (type=='B')):
        rm_shear=0
    else:
        rm_shear=1

    # y-array
    Ly=2.*np.arcsin(1.0) ; dy=Ly/ny
    y=dy/2.+arange(ny)*dy

    # Perform FFTs
    if periodic:
        if rank(array)==3:
            fft_array=fftn(array)
        if rank(array)==4:
            fft_array=array
            fft_array[:,:,:,0]=fftn(array[:,:,:,0])
            fft_array[:,:,:,1]=fftn(array[:,:,:,1])
            fft_array[:,:,:,2]=fftn(array[:,:,:,2])
    else:
        fft_array=shear_fft2d(array,x,time,type=type,rm_shear=rm_shear)

    # Compute correlation function
    if (type=='rho'):
        xi=fft_array
        for k in range(nz):
            xi[:,:,k]=abs(ifftn(fft_array[:,:,k]*conjugate(fft_array[:,:,k])))/nx/ny
    else:
        xi=fft_array[:,:,:,0]
        for k in range(nz):
            xi[:,:,k]=          abs(ifftn(fft_array[:,:,k,1]*conjugate(fft_array[:,:,k,1])))/nx/ny
#            xi[:,:,k]=          abs(ifftn(fft_array[:,:,k,0]*conjugate(fft_array[:,:,k,0])))/nx/ny
#            xi[:,:,k]=xi[:,:,k]+abs(ifftn(fft_array[:,:,k,1]*conjugate(fft_array[:,:,k,1])))/nx/ny
#            xi[:,:,k]=xi[:,:,k]+abs(ifftn(fft_array[:,:,k,2]*conjugate(fft_array[:,:,k,2])))/nx/ny
            xi[:,:,k]=xi[:,:,k]/ciso2
    xi=np.roll(np.roll(xi,nx/2,0),ny/2,1)
    if not(periodic):
        xi=unshear(xi,x,time,type='rho',direct=-1.e0)

    return xi