import pytest
import numpy as np

import random

import pyop.operators as operators
import pyop

num_tests = 100

array_max_size = 5
dimensions_max = 4


###############
#  FFT Tests  #
###############

def testFftInputErrors():

    with pytest.raises(ValueError):
        operators.fft((8, 8), (8,))

    with pytest.raises(ValueError):
        operators.fft((8, 8), (8, -8))

    with pytest.raises(ValueError):
        operators.fft((8, -8), (8, 8))


def testFftRandom():
    for _ in range(num_tests):
        d = random.randint(1, dimensions_max + 1)

        arr = np.random.rand(*tuple(
            random.randint(1, array_max_size + 1) for _ in range(d)))

        order = random.choice(('C', 'F'))

        s = tuple(random.randint(1, array_max_size*2) for _ in range(d))

        F = operators.fft(arr.shape, s = s, order = order)

        pyop.adjointTest(F)

        np.testing.assert_allclose(
            np.reshape(F._forward(np.ravel(arr, order)), s, order),
            np.fft.fftn(arr, s = s))


def testIfftInputErrors():

    with pytest.raises(ValueError):
        operators.ifft((8, 8), (8,))

    with pytest.raises(ValueError):
        operators.ifft((8, 8), (8, -8))

    with pytest.raises(ValueError):
        operators.ifft((8, -8), (8, 8))


def testIfftRandom():
    for _ in range(num_tests):
        d = random.randint(1, dimensions_max + 1)

        arr = np.random.rand(*tuple(
            random.randint(1, array_max_size + 1) for _ in range(d)))

        order = random.choice(('C', 'F'))

        s = tuple(random.randint(1, array_max_size*2) for _ in range(d))

        F = operators.ifft(arr.shape, s = s, order = order)

        pyop.adjointTest(F)

        np.testing.assert_allclose(
            np.reshape(F._forward(np.ravel(arr, order)), s, order),
            np.fft.ifftn(arr, s = s))

@@ -1811,24 +1811,1896 @@
 .fft.ifftn(arr, s = s))%0A
+%0A%0A#####################%0A#  FFT Shift Tests  #%0A#####################%0A%0Adef randomAxes(ndim):%0A    return tuple(random.randint(0, ndim)%0A        for _ in range(random.randint(0, ndim + 2)))%0A%0Adef testFftshiftInputErrors():%0A%0A    with pytest.raises(ValueError):%0A        operators.fftshift((8, -8), (1,))%0A%0A    with pytest.raises(ValueError):%0A        operators.fftshift((8, 8), (-1,))%0A%0A    with pytest.raises(ValueError):%0A        operators.fftshift((8, 8), (2, ))%0A%0A%0Adef testFftshiftRandom():%0A    for _ in range(num_tests):%0A        d = random.randint(1, dimensions_max + 1)%0A%0A        arr = np.random.rand(*tuple(%0A            random.randint(1, array_max_size + 1) for _ in range(d)))%0A%0A        order = random.choice(('C', 'F'))%0A%0A        axes = random.choice((None, randomAxes(arr.ndim - 1)))%0A%0A        F = operators.fftshift(arr.shape, axes, order = order)%0A%0A        pyop.adjointTest(F)%0A%0A        np.testing.assert_allclose(%0A            np.reshape(F._forward(np.ravel(arr, order)), arr.shape, order),%0A            np.fft.fftshift(arr, axes))%0A%0A%0Adef testIfftshiftInputErrors():%0A%0A    with pytest.raises(ValueError):%0A        operators.ifftshift((8, -8), (1,))%0A%0A    with pytest.raises(ValueError):%0A        operators.ifftshift((8, 8), (-1,))%0A%0A    with pytest.raises(ValueError):%0A        operators.ifftshift((8, 8), (2, ))%0A%0A%0Adef testIfftshiftRandom():%0A    for _ in range(num_tests):%0A        d = random.randint(1, dimensions_max + 1)%0A%0A        arr = np.random.rand(*tuple(%0A            random.randint(1, array_max_size + 1) for _ in range(d)))%0A%0A        order = random.choice(('C', 'F'))%0A%0A        axes = random.choice((None, randomAxes(arr.ndim - 1)))%0A%0A        F = operators.ifftshift(arr.shape, axes, order = order)%0A%0A        pyop.adjointTest(F)%0A%0A        np.testing.assert_allclose(%0A            np.reshape(F._forward(np.ravel(arr, order)), arr.shape, order),%0A            np.fft.ifftshift(arr, axes))%0A%0A

Add tests fftshift, ifftshift.