import unittest
import numpy as np
from dmdtools import *

class TestDMD(unittest.TestCase):
    def check_2d(self, f, **kwargs):
        eigvals = [0.5, 3.0]
        A = np.diag(eigvals)
        x1 = [1,2]
        x2 = [-1,1]
        X = np.array([x1, x2]).T
        Y = np.dot(A, X)
        modes, evals = f(X, Y, **kwargs)
        np.testing.assert_array_almost_equal(np.sort(evals), eigvals)
        # TODO: test modes as well

    def test_2d(self):
        self.check_2d(dmd)

    def test_2d_kernel_none(self):
        self.check_2d(kdmd, kernel=None)

    def test_2d_kernel_0(self):
        self.check_2d(kdmd, kernel=0)

# This test will fail; regularization will change the eigenvalues
#    def test_2d_kernel_5(self):
#        # TODO: fix this test.  Probably shouldn't recover the eigenvalues of A
#        self.check_2d(kdmd, kernel=5)

    def check_2d_random_data(self, n_data, f, **kwargs):

        eigvals = [0.9, 0.5]
        eigvecs = np.array([[1., 0.], [0., 1.]])
        A = np.diag(eigvals)

        np.random.seed(0)  # set the seed to 0 for reproducability
        X = np.random.randn(2, n_data)
        Y = A.dot(X)

        modes, evals = f(X, Y, **kwargs)

        # Check that we have an eigenvalue near 1
        self.assertAlmostEqual(min(abs(evals - 1)), 0)

        # Check that we have an eigenvalue near the true values
        for eigval in eigvals:
            self.assertAlmostEqual(min(abs(evals - eigval)), 0)

        # Check that all modes BUT the ones associated
        # with the system eigenvalues are 0
        for ii in xrange(0, len(evals)):
            try:  # Either the mode is 0 in norm
                np.testing.assert_array_almost_equal(modes[:, ii], np.zeros(2))
            except AssertionError:  # OR we match an eigval/eigvec
                # Find which system eigenvalue we best match
                index = np.argmin(abs(eigvals - evals[ii]))

                # Check that the "best fit" eigenvalue is a good match
                self.assertAlmostEqual(eigvals[index], evals[ii])

                # AND, once normalize, we match an eigenvector
                np.testing.assert_array_almost_equal(
                    abs(modes[:, ii]/np.linalg.norm(modes[:, ii])),
                    eigvecs[:, index])

    def test_2d_kernel_powers_random(self):
        n_data = 100
        for kernel_value in xrange(1, 5):
            self.check_2d_random_data(n_data, kdmd, kernel=kernel_value)




if __name__ == "__main__":
    unittest.main()

@@ -628,228 +628,8 @@
 0)%0A%0A
-# This test will fail; regularization will change the eigenvalues%0A#    def test_2d_kernel_5(self):%0A#        # TODO: fix this test.  Probably shouldn't recover the eigenvalues of A%0A#        self.check_2d(kdmd, kernel=5)%0A%0A
     
@@ -1339,25 +1339,24 @@
   for ii in 
-x
 range(0, len
@@ -2139,17 +2139,16 @@
 alue in 
-x
 range(1,
@@ -2226,17 +2226,16 @@
 value)%0A%0A
-%0A
 %0A%0Aif __n

Change xrange to range for python3 compatibility