1from numpy.testing import *
2import numpy
3
4from algopy.utpm import *
5from algopy.utpm.algorithms import *
6
7# explicitly import some of the helpers that have underscores
8from algopy.utpm.algorithms import _plus_const
9from algopy.utpm.algorithms import _taylor_polynomials_of_ode_solutions
10
11
12class Test_Helper_Functions(TestCase):
13    def test_vdot(self):
14        (D,P,N,M) = 4,3,2,5
15        A = numpy.array([ i for i in range(D*P*N*M)],dtype=float)
16        A = A.reshape((D,P,N,M))
17        B = A.transpose((0,1,3,2)).copy()
18
19        R  = vdot(A[0],B[0])
20        R2 = numpy.zeros((P,N,N))
21        for p in range(P):
22            R2[p,:,:] = numpy.dot(A[0,p],B[0,p])
23
24        S  = vdot(A,B)
25        S2 = numpy.zeros((D,P,N,N))
26        for d in range(D):
27            for p in range(P):
28                S2[d,p,:,:] = numpy.dot(A[d,p],B[d,p])
29
30        assert_array_almost_equal(R,R2)
31        assert_array_almost_equal(S,S2)
32
33    def test_triple_truncated_dot(self):
34        D,P,N,M = 3,1,1,1
35        A = numpy.random.rand(D,P,N,M)
36        B = numpy.random.rand(D,P,N,M)
37        C = numpy.random.rand(D,P,N,M)
38
39        S = A[0]*B[1]*C[1] + A[1]*B[0]*C[1] + A[1]*B[1]*C[0]
40        R = truncated_triple_dot(A,B,C,2)
41
42        assert_array_almost_equal(R,S)
43
44        D,P,N,M = 4,1,1,1
45        A = numpy.random.rand(D,P,N,M)
46        B = numpy.random.rand(D,P,N,M)
47        C = numpy.random.rand(D,P,N,M)
48
49        S = A[0]*B[1]*C[2] + A[0]*B[2]*C[1] + \
50            A[1]*B[0]*C[2] + A[1]*B[1]*C[1] + A[1]*B[2]*C[0] +\
51            A[2]*B[1]*C[0] + A[2]*B[0]*C[1]
52        R = truncated_triple_dot(A,B,C, 3)
53
54        assert_array_almost_equal(R,S)
55
56
57    def test_broadcast_arrays_shape(self):
58        D,P = 3,4
59
60        y_shp = (D,P)
61        x_shp = (D,P)
62        z_shp1 = broadcast_arrays_shape(x_shp, y_shp)
63        z_shp2 = broadcast_arrays_shape(y_shp, x_shp)
64        assert_array_equal((D,P), z_shp1)
65        assert_array_equal((D,P), z_shp2)
66
67        y_shp = (D,P,1)
68        x_shp = (D,P)
69        z_shp1 = broadcast_arrays_shape(x_shp, y_shp)
70        z_shp2 = broadcast_arrays_shape(y_shp, x_shp)
71        assert_array_equal((D,P,1), z_shp1)
72        assert_array_equal((D,P,1), z_shp2)
73
74        y_shp = (D,P,1,2,3)
75        x_shp = (D,P)
76        z_shp1 = broadcast_arrays_shape(x_shp, y_shp)
77        z_shp2 = broadcast_arrays_shape(y_shp, x_shp)
78        assert_array_equal((D,P,1,2,3), z_shp1)
79        assert_array_equal((D,P,1,2,3), z_shp2)
80
81        y_shp = (D,P,1,2,3)
82        x_shp = (D,P,3,1,1)
83        z_shp1 = broadcast_arrays_shape(x_shp, y_shp)
84        z_shp2 = broadcast_arrays_shape(y_shp, x_shp)
85        assert_array_equal((D,P,3,2,3), z_shp1)
86        assert_array_equal((D,P,3,2,3), z_shp2)
87
88        y_shp = (D,P,7, 1,2,3)
89        x_shp = (D,P,3,1,1)
90        z_shp1 = broadcast_arrays_shape(x_shp, y_shp)
91        z_shp2 = broadcast_arrays_shape(y_shp, x_shp)
92        assert_array_equal((D,P,7,3,2,3), z_shp1)
93        assert_array_equal((D,P,7,3,2,3), z_shp2)
94
95        y_shp = (D,P,7, 1,2,3)
96        x_shp = (D,P,3,1,1)
97        z_shp1 = broadcast_arrays_shape(x_shp, y_shp)
98        z_shp2 = broadcast_arrays_shape(y_shp, x_shp)
99        assert_array_equal((D,P,7,3,2,3), z_shp1)
100        assert_array_equal((D,P,7,3,2,3), z_shp2)
101
102
103class Test_taylor_polynomials_of_ode_solutions(TestCase):
104
105    def test_log(self):
106
107        for shape in (
108                (2, 3),
109                (4, 3, 2, 5),
110                ):
111
112            # sample some positive numbers for taking logs
113            x = UTPM(numpy.exp(numpy.random.randn(*shape)))
114
115            # construct the u and v arrays
116            u_data = x.data.copy()
117            v_data = numpy.empty_like(u_data)
118            v_data[0, ...] = numpy.log(u_data[0])
119
120            # construct values like in Table (13.2) of "Evaluating Derivatives"
121            a_data = numpy.zeros_like(u_data)
122            b_data = u_data.copy()
123            c_data = _plus_const(numpy.zeros_like(u_data), 1)
124
125            # fill the rest of the v_data
126            _taylor_polynomials_of_ode_solutions(
127                a_data, b_data, c_data,
128                u_data, v_data)
129
130            # compare the v_data array to the UTPM log data
131            assert_allclose(v_data, UTPM.log(x).data)
132
133    def test_log1p(self):
134
135        for shape in (
136                (2, 3),
137                (4, 3, 2, 5),
138                ):
139
140            # sample some positive numbers for taking logs
141            x = UTPM(numpy.exp(numpy.random.randn(*shape)))
142
143            # construct the u and v arrays
144            u_data = x.data.copy()
145            v_data = numpy.empty_like(u_data)
146            v_data[0, ...] = numpy.log1p(u_data[0])
147
148            # construct values like in Table (13.2) of "Evaluating Derivatives"
149            a_data = numpy.zeros_like(u_data)
150            b_data = _plus_const(u_data.copy(), 1)
151            c_data = _plus_const(numpy.zeros_like(u_data), 1)
152
153            # fill the rest of the v_data
154            _taylor_polynomials_of_ode_solutions(
155                a_data, b_data, c_data,
156                u_data, v_data)
157
158            # compare the v_data array to the UTPM log1p data
159            assert_allclose(v_data, UTPM.log1p(x).data)
160
161    def test_exp(self):
162
163        for shape in (
164                (2, 3),
165                (4, 3, 2, 5),
166                ):
167
168            # sample some random numbers
169            x = UTPM(numpy.random.randn(*shape))
170
171            # construct the u and v arrays
172            u_data = x.data.copy()
173            v_data = numpy.empty_like(u_data)
174            v_data[0, ...] = numpy.exp(u_data[0])
175
176            # construct values like in Table (13.2) of "Evaluating Derivatives"
177            a_data = _plus_const(numpy.zeros_like(u_data), 1)
178            b_data = _plus_const(numpy.zeros_like(u_data), 1)
179            c_data = numpy.zeros_like(u_data)
180
181            # fill the rest of the v_data
182            _taylor_polynomials_of_ode_solutions(
183                a_data, b_data, c_data,
184                u_data, v_data)
185
186            # compare the v_data array to the UTPM exp data
187            assert_allclose(v_data, UTPM.exp(x).data)
188
189    def test_expm1(self):
190
191        for shape in (
192                (2, 3),
193                (4, 3, 2, 5),
194                ):
195
196            # sample some random numbers
197            x = UTPM(numpy.random.randn(*shape))
198
199            # construct the u and v arrays
200            u_data = x.data.copy()
201            v_data = numpy.empty_like(u_data)
202            v_data[0, ...] = numpy.expm1(u_data[0])
203
204            # construct values like in Table (13.2) of "Evaluating Derivatives"
205            a_data = _plus_const(numpy.zeros_like(u_data), 1)
206            b_data = _plus_const(numpy.zeros_like(u_data), 1)
207            c_data = _plus_const(numpy.zeros_like(u_data), 1)
208
209            # fill the rest of the v_data
210            _taylor_polynomials_of_ode_solutions(
211                a_data, b_data, c_data,
212                u_data, v_data)
213
214            # compare the v_data array to the UTPM expm1 data
215            assert_allclose(v_data, UTPM.expm1(x).data)
216
217    def test_power(self):
218
219        # define a constant real exponent
220        r = 1.23
221
222        for shape in (
223                (2, 3),
224                (4, 3, 2, 5),
225                ):
226
227            # sample some positive numbers for taking fractional real powers
228            x = UTPM(numpy.exp(numpy.random.randn(*shape)))
229
230            # construct the u and v arrays
231            u_data = x.data.copy()
232            v_data = numpy.empty_like(u_data)
233            v_data[0, ...] = numpy.power(u_data[0], r)
234
235            # construct values like in Table (13.2) of "Evaluating Derivatives"
236            a_data = _plus_const(numpy.zeros_like(u_data), r)
237            b_data = u_data.copy()
238            c_data = numpy.zeros_like(u_data)
239
240            # fill the rest of the v_data
241            _taylor_polynomials_of_ode_solutions(
242                a_data, b_data, c_data,
243                u_data, v_data)
244
245            # compare the v_data array to the UTPM power data
246            assert_allclose(v_data, (x ** r).data)
247
248    def test_dawsn(self):
249
250        for shape in (
251                (2, 3),
252                (4, 3, 2, 5),
253                ):
254
255            # sample some random numbers
256            x = UTPM(numpy.random.randn(*shape))
257
258            # construct the u and v arrays
259            u_data = x.data.copy()
260            v_data = numpy.empty_like(u_data)
261            v_data[0, ...] = scipy.special.dawsn(u_data[0])
262
263            # construct values like in Table (13.2) of "Evaluating Derivatives"
264            a_data = -2 * u_data.copy()
265            b_data = _plus_const(numpy.zeros_like(u_data), 1)
266            c_data = _plus_const(numpy.zeros_like(u_data), 1)
267
268            # fill the rest of the v_data
269            _taylor_polynomials_of_ode_solutions(
270                a_data, b_data, c_data,
271                u_data, v_data)
272
273            # compare the v_data array to the UTPM dawsn data
274            assert_allclose(v_data, UTPM.dawsn(x).data)
275
276
277
278class Test_pushforward_class_functions(TestCase):
279    """
280    Test the push forward class functions that operate directly on data.
281    """
282
283    def test__itruediv(self):
284        X_data = 2 * numpy.random.rand(2,2,2,2)
285        Z_data = 3 * numpy.random.rand(2,2,2,2)
286        Z2_data = Z_data.copy()
287
288        UTPM._itruediv(Z_data, X_data)
289
290        X = UTPM(X_data)
291        Z = UTPM(Z2_data)
292
293        Z/=X
294
295        assert_array_almost_equal(Z_data, Z.data)
296
297    def test__div(self):
298        X_data = 2 * numpy.random.rand(2,2,2,2)
299        Y_data = 3 * numpy.random.rand(2,2,2,2)
300        Z_data = numpy.zeros((2,2,2,2))
301
302        X = UTPM(X_data)
303        Y = UTPM(Y_data)
304
305        Z = X/Y
306
307        UTPM._truediv(X_data, Y_data, out = Z_data)
308
309        assert_array_almost_equal(Z_data, Z.data)
310
311    def test__transpose(self):
312        D,P,M,N = 2,3,4,5
313        X_data = numpy.random.rand(D,P,N,M)
314        Y_data = UTPM._transpose(X_data)
315
316        assert_array_almost_equal(X_data.transpose((0,1,3,2)), Y_data)
317
318
319class Test_aliasing(TestCase):
320
321    def test_mul_aliasing(self):
322        D, P, M, N = 5, 4, 3, 2
323        x = numpy.random.randn(D, P, M, N)
324        y1 = numpy.random.randn(D, P, M, N)
325        y2 = numpy.empty_like(x)
326        UTPM._mul(x, y1, out=y2)
327        UTPM._mul(x, y1, out=y1)
328        assert_allclose(y1, y2)
329
330    def test_div_aliasing(self):
331        D, P, M, N = 5, 4, 3, 2
332        x = numpy.random.randn(D, P, M, N)
333        y1 = numpy.random.randn(D, P, M, N)
334        y2 = numpy.empty_like(x)
335        UTPM._truediv(x, y1, out=y2)
336        UTPM._truediv(x, y1, out=y1)
337        assert_allclose(y1, y2)
338
339    def test_square_aliasing(self):
340        D, P, M, N = 5, 4, 3, 2
341        x = numpy.random.randn(D, P, M, N)
342        y = numpy.empty_like(x)
343        UTPM._square(x, out=y)
344        UTPM._square(x, out=x)
345        assert_allclose(x, y)
346
347    def test_sign_aliasing(self):
348        D, P, M, N = 5, 4, 3, 2
349        x = numpy.random.randn(D, P, M, N)
350        y = numpy.empty_like(x)
351        UTPM._sign(x, out=y)
352        UTPM._sign(x, out=x)
353        assert_allclose(x, y)
354
355    def test_reciprocal_aliasing(self):
356        D, P, M, N = 5, 4, 3, 2
357        x = numpy.random.randn(D, P, M, N)
358        y = numpy.empty_like(x)
359        UTPM._reciprocal(x, out=y)
360        UTPM._reciprocal(x, out=x)
361        assert_allclose(x, y)
362
363    def test_negative_aliasing(self):
364        D, P, M, N = 5, 4, 3, 2
365        x = numpy.random.randn(D, P, M, N)
366        y = numpy.empty_like(x)
367        UTPM._negative(x, out=y)
368        UTPM._negative(x, out=x)
369        assert_allclose(x, y)
370
371    def test_absolute_aliasing(self):
372        D, P, M, N = 5, 4, 3, 2
373        x = numpy.random.randn(D, P, M, N)
374        y = numpy.empty_like(x)
375        UTPM._absolute(x, out=y)
376        UTPM._absolute(x, out=x)
377        assert_allclose(x, y)
378
379    def test_exp_aliasing(self):
380        D, P, M, N = 5, 4, 3, 2
381        x = numpy.random.randn(D, P, M, N)
382        y = numpy.empty_like(x)
383        UTPM._exp(x, out=y)
384        UTPM._exp(x, out=x)
385        assert_allclose(x, y)
386
387    def test_log_aliasing(self):
388        D, P, M, N = 5, 4, 3, 2
389        x = numpy.exp(numpy.random.randn(D, P, M, N))
390        y = numpy.empty_like(x)
391        UTPM._log(x, out=y)
392        UTPM._log(x, out=x)
393        assert_allclose(x, y)
394
395    def test_sqrt_aliasing(self):
396        D, P, M, N = 5, 4, 3, 2
397        x = numpy.exp(numpy.random.randn(D, P, M, N))
398        y = numpy.empty_like(x)
399        UTPM._sqrt(x, out=y)
400        UTPM._sqrt(x, out=x)
401        assert_allclose(x, y)
402
403
404if __name__ == "__main__":
405    run_module_suite()
406