tvregdiff/test_tvregdiff.py at master · stur86/tvregdiff · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
"""
Unit tests for tvregdiff.py which test if this Python
implementation of TVREGDiff replicates the MATLAB demo
scipt outputs from Rick Chartrand's webpage:
https://sites.google.com/site/dnartrahckcir/home/tvdiff-code
"""

import os
import unittest
import numpy as np
from numpy.testing import assert_allclose
from tvregdiff import TVRegDiff
try:
    import matplotlib.pyplot as plt
except ImportError:
    plt = False

# Functions to test on
functions = {
    'abs': [np.abs, np.sign],
    'sigmoid': [lambda x: np.exp(x)/(1+np.exp(x)),
                lambda x: np.exp(x)/(1+np.exp(x))-np.exp(2*x)/(1+np.exp(x))**2]
}


def rms(x):
    return np.average(x**2)**0.5


class TVDiffTest(unittest.TestCase):

    def setUp(self):
        self.data_dir = 'test_data'
        self.longMessage = True

    def test_small(self):

        n = 50
        x = np.linspace(-5, 5, n)
        dx = x[1]-x[0]

        for fname, (f, df) in functions.items():

            data = f(x) + (np.random.random(n)-0.5)*0.05
            targ = df(x)

            u1a = TVRegDiff(data, 1, 0.2, dx=dx, plotflag=False,
                            precondflag=False)
            u1s = TVRegDiff(data, 1, 0.2, dx=dx, plotflag=False,
                            diffkernel='sq')

            # Loose tolerance for this
            self.assertLess(rms(u1a-targ), 0.3,
                            '[function: {0}]'.format(fname))
            self.assertLess(rms(u1s-targ), 0.3,
                            '[function: {0}]'.format(fname))

            # Tigher for more iterations
            u10a = TVRegDiff(data, 10, 0.2, dx=dx, plotflag=False,
                             precondflag=False)

            self.assertLess(rms(u10a-targ), 0.2,
                            '[function: {0}]'.format(fname))

            if plt:
                plt.title('scale = small')
                plt.plot(x, targ)
                plt.plot(x, u1a, label='u1a')
                plt.plot(x, u1s, label='u1s')
                plt.plot(x, u10a, label='u10a')
                plt.legend()
                plt.show()

    def test_large(self):

        n = 1000
        x = np.linspace(-5, 5, n)
        dx = x[1]-x[0]

        for fname, (f, df) in functions.items():

            data = f(x) + (np.random.random(n)-0.5)*0.05
            targ = df(x)

            u1a = TVRegDiff(data, 1, 0.2, dx=dx, plotflag=False,
                            scale='large')
            u1s = TVRegDiff(data, 1, 0.2, dx=dx, plotflag=False,
                            diffkernel='sq', scale='large')

            self.assertLess(rms(u1a-targ), 0.3,
                            '[function: {0}]'.format(fname))
            self.assertLess(rms(u1s-targ), 0.3,
                            '[function: {0}]'.format(fname))

            # Tigher for more iterations
            u10a = TVRegDiff(data, 10, 0.2, dx=dx, plotflag=False,
                             scale='large')

            self.assertLess(rms(u10a-targ), 0.2,
                            '[function: {0}]'.format(fname))

            if plt:
                plt.title('scale = large')
                plt.plot(x, targ)
                plt.plot(x, u1a, label='u1a')
                plt.plot(x, u1s, label='u1s')
                plt.plot(x, u10a, label='u10a')
                plt.legend()
                plt.show()


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