test_nifty_mpi_data.py 77.9 KB
Newer Older
1
2
# -*- coding: utf-8 -*-

Ultima's avatar
Ultima committed
3
from numpy.testing import assert_equal,\
4
5
    assert_almost_equal,\
    assert_raises
6

7
8
9
10
11
12
from nose_parameterized import parameterized
import unittest

import itertools
import os
import numpy as np
Ultima's avatar
Ultima committed
13
14
15
import warnings
import tempfile

16
import nifty
theos's avatar
theos committed
17
18
from nifty.d2o import distributed_data_object,\
                      STRATEGIES
19

20
21
from distutils.version import LooseVersion as lv

22
FOUND = {}
23
try:
Ultima's avatar
Ultima committed
24
    import h5py
25
    FOUND['h5py'] = True
Ultima's avatar
Ultima committed
26
    FOUND['h5py_parallel'] = h5py.get_config().mpi
27
except(ImportError):
28
    FOUND['h5py'] = False
Ultima's avatar
Ultima committed
29
    FOUND['h5py_parallel'] = False
30
31
32
33

try:
    from mpi4py import MPI
    FOUND['MPI'] = True
34
except(ImportError):
35
36
    import mpi_dummy as MPI
    FOUND['MPI'] = False
Ultima's avatar
Ultima committed
37
38


39
40
41
42
43
44
45
46
47
48
###############################################################################

comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size

###############################################################################

np.random.seed(123)

49
###############################################################################
50

51
# all_datatypes = [np.bool_, np.int16, np.uint16, np.uint32, np.int32, np.int_,
52
53
#             np.int, np.int64, np.uint64, np.float32, np.float_, np.float,
#             np.float64, np.float128, np.complex64, np.complex_,
54
#             np.complex, np.complex128]
55
all_datatypes = [np.dtype('bool'), np.dtype('int16'), np.dtype('uint16'),
56
57
58
59
60
                 np.dtype('uint32'), np.dtype('int32'), np.dtype('int'),
                 np.dtype('int64'), np.dtype('uint'), np.dtype('uint64'),
                 np.dtype('float32'), np.dtype(
                     'float64'), np.dtype('float128'),
                 np.dtype('complex64'), np.dtype('complex128')]
61
62

###############################################################################
63

Ultima's avatar
Ultima committed
64
65
66
67
all_distribution_strategies = STRATEGIES['all']
global_distribution_strategies = STRATEGIES['global']
local_distribution_strategies = STRATEGIES['local']
hdf5_distribution_strategies = STRATEGIES['hdf5']
68
69

###############################################################################
70

71
72
73
74
binary_non_inplace_operators = ['__add__', '__radd__', '__sub__', '__rsub__',
                                '__div__', '__truediv__', '__rdiv__',
                                '__rtruediv__', '__floordiv__',
                                '__rfloordiv__', '__mul__', '__rmul__',
75
76
                                '__pow__', '__rpow__']
binary_inplace_operators = ['__iadd__', '__isub__', '__idiv__', '__itruediv__',
77
78
79
                            '__ifloordiv__', '__imul__', '__ipow__']
comparison_operators = ['__ne__', '__lt__', '__le__', '__eq__', '__ge__',
                        '__gt__', ]
80

81
82
###############################################################################

83
hdf5_test_paths = [  # ('hdf5_init_test.hdf5', None),
84
    ('hdf5_init_test.hdf5', os.path.join(os.path.dirname(nifty.__file__),
85
86
87
88
                                         'test/hdf5_init_test.hdf5')),
    ('hdf5_init_test.hdf5',
     os.path.join(os.path.dirname(nifty.__file__),
                  'test//hdf5_test_folder/hdf5_init_test.hdf5'))]
89
90

###############################################################################
91

92

93
def custom_name_func(testcase_func, param_num, param):
94
    return "%s_%s" % (
95
96
97
        testcase_func.__name__,
        parameterized.to_safe_name("_".join(str(x) for x in param.args)),
    )
98
99


Ultima's avatar
Ultima committed
100
101
102
###############################################################################
###############################################################################

103
def generate_data(global_shape, dtype, distribution_strategy,
104
                  strictly_positive=False):
Ultima's avatar
Ultima committed
105
106
    if distribution_strategy in global_distribution_strategies:
        a = np.arange(np.prod(global_shape))
107
        a -= np.prod(global_shape) // 2
Ultima's avatar
Ultima committed
108

109
        a = a * (1 + 1j)
Ultima's avatar
Ultima committed
110
111
        if strictly_positive:
            a = abs(a)
112

Ultima's avatar
Ultima committed
113
114
115
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            a = a.reshape(global_shape).astype(dtype)
116
117
        obj = distributed_data_object(
                                a, distribution_strategy=distribution_strategy)
Ultima's avatar
Ultima committed
118
        global_a = a
119

Ultima's avatar
Ultima committed
120
121
    elif distribution_strategy in local_distribution_strategies:
        local_shape = list(global_shape)
122
        if rank % 2 == 1:
Ultima's avatar
Ultima committed
123
124
            local_shape[0] = 0
        else:
125
126
127
            local_shape[0] = global_shape[0] // np.ceil(size / 2.)
            number_of_extras = global_shape[
                0] - local_shape[0] * np.ceil(size / 2.)
Ultima's avatar
Ultima committed
128
129
130
            if number_of_extras > rank:
                local_shape[0] += 1
        local_shape = tuple(local_shape)
131

Ultima's avatar
Ultima committed
132
        a = np.arange(np.prod(local_shape))
133
        a -= np.prod(local_shape) // 2
134

135
        a = a * (1 + 1j)
Ultima's avatar
Ultima committed
136
137
138
139
140
141
        if strictly_positive:
            a = abs(a)

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            a = a.reshape(local_shape).astype(dtype)
142
143
144
        a *= (rank + 1)
        obj = distributed_data_object(
                    local_data=a, distribution_strategy=distribution_strategy)
Ultima's avatar
Ultima committed
145
146
        a_list = comm.allgather(a)
        global_a = np.concatenate(a_list)
147

Ultima's avatar
Ultima committed
148
    return (global_a, obj)
149
150


151
152
153
###############################################################################
###############################################################################

Ultima's avatar
Ultima committed
154

155
156
class Test_Globaltype_Initialization(unittest.TestCase):

157
    @parameterized.expand(
158
159
160
161
162
        itertools.product([(1,), (7,), (78, 11), (256, 256)],
                          all_datatypes,
                          global_distribution_strategies,
                          [True, False]),
        testcase_func_name=custom_name_func)
163
    def test_successful_init_via_global_shape_and_dtype(self,
164
                                                        global_shape,
165
166
167
                                                        dtype,
                                                        distribution_strategy,
                                                        hermitian):
168
169
170
171
172
        obj = distributed_data_object(
                                  global_shape=global_shape,
                                  dtype=dtype,
                                  distribution_strategy=distribution_strategy,
                                  hermitian=hermitian)
173

174
175
176
177
178
        assert_equal(obj.dtype, dtype)
        assert_equal(obj.shape, global_shape)
        assert_equal(obj.distribution_strategy, distribution_strategy)
        assert_equal(obj.hermitian, hermitian)
        assert_equal(obj.data.dtype, np.dtype(dtype))
179
180


181
###############################################################################
182

183
    @parameterized.expand(
184
185
186
187
        itertools.product([(1,), (7,), (77, 11), (256, 256)],
                          all_datatypes,
                          global_distribution_strategies),
        testcase_func_name=custom_name_func)
188
    def test_successful_init_via_global_data(self,
189
190
191
                                             global_shape,
                                             dtype,
                                             distribution_strategy):
192

193
194
195
196
        a = (np.random.rand(*global_shape) * 100 - 50).astype(dtype)
        obj = distributed_data_object(
                                  global_data=a,
                                  distribution_strategy=distribution_strategy)
197
198
199
200
201
        assert_equal(obj.dtype, np.dtype(dtype))
        assert_equal(obj.shape, global_shape)
        assert_equal(obj.distribution_strategy, distribution_strategy)
        assert_equal(obj.data.dtype, np.dtype(dtype))

202
203
204
###############################################################################

    @parameterized.expand(
205
206
207
208
209
        itertools.product([(1,), (7,), (77, 11), (256, 256)],
                          ['tuple', 'list'],
                          all_datatypes,
                          global_distribution_strategies),
        testcase_func_name=custom_name_func)
210
    def test_successful_init_via_tuple_and_list(self,
211
212
213
214
                                                global_shape,
                                                global_data_type,
                                                dtype,
                                                distribution_strategy):
215

216
        a = (np.random.rand(*global_shape) * 100 - 50).astype(dtype)
217
218
219
220
        if global_data_type == 'list':
            a = a.tolist()
        elif global_data_type == 'tuple':
            a = tuple(a.tolist())
221

222
223
224
        obj = distributed_data_object(
                                global_data=a,
                                distribution_strategy=distribution_strategy)
225
226
        assert_equal(obj.shape, global_shape)
        assert_equal(obj.distribution_strategy, distribution_strategy)
227

228
###############################################################################
229

230
    @parameterized.expand(itertools.product([
231
232
233
234
235
236
237
238
        [1, (13, 7), np.dtype('float64'),
         (13, 7), np.dtype('float64')],
        [np.array([1]), (13, 7), np.dtype('float64'),
         (1,), np.dtype('float64')],
        [np.array([[1., 2.], [3., 4.]]), (13, 7), np.dtype('int'),
         (2, 2), np.dtype('int')],
        [None, (10, 10), None,
         (10, 10), np.dtype('float64')],
239
240
241
242
    ], global_distribution_strategies),
        testcase_func_name=custom_name_func)
    def test_special_init_cases(self,
                                (global_data,
243
244
245
246
                                 global_shape,
                                 dtype,
                                 expected_shape,
                                 expected_dtype),
247
                                distribution_strategy):
248
249
250
251
252
        obj = distributed_data_object(
                                  global_data=global_data,
                                  global_shape=global_shape,
                                  dtype=dtype,
                                  distribution_strategy=distribution_strategy)
253
254
        assert_equal(obj.shape, expected_shape)
        assert_equal(obj.dtype, expected_dtype)
255

256
###############################################################################
257

258
    if FOUND['h5py'] == True:
259
        @parameterized.expand(itertools.product(hdf5_test_paths,
260
261
                                                hdf5_distribution_strategies),
                              testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
262
        def test_hdf5_init(self, (alias, path), distribution_strategy):
263
264
265
266
267
268
            obj = distributed_data_object(
                                  global_data=1.,
                                  global_shape=(12, 6),
                                  alias=alias,
                                  path=path,
                                  distribution_strategy=distribution_strategy)
269
            assert_equal(obj.dtype, np.complex128)
270
            assert_equal(obj.shape, (13, 7))
271
272

###############################################################################
273

274
275
    @parameterized.expand(
        itertools.product(
276
277
278
279
280
281
282
283
            [(None, None, None, None, None),
             (None, None, np.int_, None, None),
                (None, (), np.dtype('int'), None, None),
                (1, None, None, None, None),
             (None, None, None, np.array([1, 2, 3]), (3,)),
                (None, None, np.int_, None, (3,))],
            global_distribution_strategies),
        testcase_func_name=custom_name_func)
284
    def test_failed_init_on_unsufficient_parameters(self,
285
286
287
288
289
                                                    (global_data,
                                                     global_shape,
                                                     dtype,
                                                     local_data, local_shape),
                                                    distribution_strategy):
290
291
        assert_raises(ValueError,
                      lambda: distributed_data_object(
292
293
294
295
296
297
                          global_data=global_data,
                          global_shape=global_shape,
                          dtype=dtype,
                          local_data=local_data,
                          local_shape=local_shape,
                          distribution_strategy=distribution_strategy))
298

Ultima's avatar
Ultima committed
299
300
###############################################################################

301
    if size > 1:
Ultima's avatar
Ultima committed
302
303
        @parameterized.expand(
            itertools.product(
304
305
306
307
308
309
310
311
312
313
314
315
316
                [(None, (10, rank, 10), np.dtype('int'), None, None),
                 (None, (2, 2),
                  np.dtype('int') if (rank == 0) else np.dtype('float'),
                  None, None), ],
                global_distribution_strategies),
            testcase_func_name=custom_name_func)
        def test_failed_init_unsufficient_params_mpi(self,
                                                     (global_data,
                                                      global_shape,
                                                      dtype,
                                                      local_data,
                                                      local_shape),
                                                     distribution_strategy):
Ultima's avatar
Ultima committed
317
318
            assert_raises(ValueError,
                          lambda: distributed_data_object(
319
320
321
322
323
324
                              global_data=global_data,
                              global_shape=global_shape,
                              dtype=dtype,
                              local_data=local_data,
                              local_shape=local_shape,
                              distribution_strategy=distribution_strategy))
Ultima's avatar
Ultima committed
325
326


327
328
329
###############################################################################

    @parameterized.expand(
330
331
332
        itertools.product([(0,), (1, 0), (0, 1), (25, 0, 10), (0, 0)],
                          global_distribution_strategies),
        testcase_func_name=custom_name_func)
333
    def test_init_with_zero_type_shape(self, global_shape,
334
                                       distribution_strategy):
335
336
337
338
        obj = distributed_data_object(
                                  global_shape=global_shape,
                                  dtype=np.int,
                                  distribution_strategy=distribution_strategy)
339
        assert_equal(obj.shape, global_shape)
340
341
342
343


###############################################################################
###############################################################################
344

345
346
class Test_Localtype_Initialization(unittest.TestCase):

347
    ##########################################################################
348
    @parameterized.expand(
349
350
351
352
353
354
        itertools.product([(1,), (7,), (78, 11), (256, 256)],
                          [False, True],
                          all_datatypes,
                          local_distribution_strategies,
                          [True, False]),
        testcase_func_name=custom_name_func)
355
    def test_successful_init_via_local_shape_and_dtype(self,
356
357
358
359
360
                                                       local_shape,
                                                       different_shapes,
                                                       dtype,
                                                       distribution_strategy,
                                                       hermitian):
361

362
        if different_shapes is True:
363
            expected_global_shape = np.array(local_shape)
364
            expected_global_shape[0] *= size * (size - 1) / 2
365
            expected_global_shape = tuple(expected_global_shape)
366
367
368
369
370
371
372
373
            local_shape = list(local_shape)
            local_shape[0] *= rank
            local_shape = tuple(local_shape)
        else:
            expected_global_shape = np.array(local_shape)
            expected_global_shape[0] *= size
            expected_global_shape = tuple(expected_global_shape)

374
375
376
377
378
        obj = distributed_data_object(
                                  local_shape=local_shape,
                                  dtype=dtype,
                                  distribution_strategy=distribution_strategy,
                                  hermitian=hermitian)
379
380
381
382
383
384

        assert_equal(obj.dtype, dtype)
        assert_equal(obj.shape, expected_global_shape)
        assert_equal(obj.distribution_strategy, distribution_strategy)
        assert_equal(obj.hermitian, hermitian)
        assert_equal(obj.data.dtype, np.dtype(dtype))
385
386


387
###############################################################################
388

389
    @parameterized.expand(
390
391
392
393
394
        itertools.product([(1,), (7,), (77, 11), (256, 256)],
                          [False, True],
                          all_datatypes,
                          local_distribution_strategies),
        testcase_func_name=custom_name_func)
395
    def test_successful_init_via_local_data(self,
396
397
398
399
400
                                            local_shape,
                                            different_shapes,
                                            dtype,
                                            distribution_strategy):
        if different_shapes is True:
401
            expected_global_shape = np.array(local_shape)
402
            expected_global_shape[0] *= size * (size - 1) / 2
403
404
405
406
407
408
409
410
            expected_global_shape = tuple(expected_global_shape)
            local_shape = list(local_shape)
            local_shape[0] *= rank
            local_shape = tuple(local_shape)
        else:
            expected_global_shape = np.array(local_shape)
            expected_global_shape[0] *= size
            expected_global_shape = tuple(expected_global_shape)
411

412
413
414
415
        a = (np.random.rand(*local_shape) * 100 - 50).astype(dtype)
        obj = distributed_data_object(
                                  local_data=a,
                                  distribution_strategy=distribution_strategy)
416

417
418
419
420
421
422
423
424
        assert_equal(obj.dtype, np.dtype(dtype))
        assert_equal(obj.shape, expected_global_shape)
        assert_equal(obj.distribution_strategy, distribution_strategy)
        assert_equal(obj.data.dtype, np.dtype(dtype))

###############################################################################

    @parameterized.expand(
Ultima's avatar
Ultima committed
425
        itertools.product([(1,), (7,), (77, 11)],
426
427
428
429
                          ['tuple', 'list'],
                          all_datatypes,
                          local_distribution_strategies),
        testcase_func_name=custom_name_func)
430
    def test_successful_init_via_tuple_and_list(self,
431
432
433
434
                                                local_shape,
                                                local_data_type,
                                                dtype,
                                                distribution_strategy):
435

436
        a = (np.random.rand(*local_shape) * 100).astype(dtype)
437
438
439
440
        if local_data_type == 'list':
            a = a.tolist()
        elif local_data_type == 'tuple':
            a = tuple(a.tolist())
441
442
443
        obj = distributed_data_object(
                                local_data=a,
                                distribution_strategy=distribution_strategy)
444
445
446
447
448
449
450
451
452
453

        expected_global_shape = np.array(local_shape)
        expected_global_shape[0] *= size
        expected_global_shape = tuple(expected_global_shape)

        assert_equal(obj.shape, expected_global_shape)
        assert_equal(obj.distribution_strategy, distribution_strategy)


###############################################################################
454

455
    @parameterized.expand(itertools.product([
456
457
458
459
        [1, (13, 7), np.float64, (13 * size, 7), np.float64],
        [np.array([1]), (13, 7), np.float64, (1 * size,), np.float64],
        [np.array([[1., 2.], [3., 4.]]), (13, 7),
         np.int, (2 * size, 2), np.int]
460
461
462
463
    ], local_distribution_strategies),
        testcase_func_name=custom_name_func)
    def test_special_init_cases(self,
                                (local_data,
464
465
466
467
                                 local_shape,
                                 dtype,
                                 expected_shape,
                                 expected_dtype),
468
                                distribution_strategy):
469
470
471
472
473
        obj = distributed_data_object(
                                  local_data=local_data,
                                  local_shape=local_shape,
                                  dtype=dtype,
                                  distribution_strategy=distribution_strategy)
474

475
476
        assert_equal(obj.shape, expected_shape)
        assert_equal(obj.dtype, expected_dtype)
Ultima's avatar
Ultima committed
477

478
###############################################################################
Ultima's avatar
Ultima committed
479
480

    def test_special_init_from_d2o_cases(self):
481
        global_shape = (8, 8)
482
        dtype = np.dtype('int')
Ultima's avatar
Ultima committed
483
        (a, obj) = generate_data(global_shape, dtype, 'equal')
484
485
486
487
        # Given dtype overwrites the one from data
        p = distributed_data_object(global_data=obj,
                                    dtype=np.dtype('float'),
                                    distribution_strategy='freeform')
Ultima's avatar
Ultima committed
488
        assert_equal(p.dtype, np.dtype('float'))
489
490
491
492
        # Global d2o overwrites local data
        p = distributed_data_object(global_data=obj,
                                    local_data=np.array([1, 2, 3]),
                                    distribution_strategy='freeform')
Ultima's avatar
Ultima committed
493
        assert_equal(obj.get_full_data(), p.get_full_data())
494
495
496
497
        # Global d2o overwrites local shapes
        p = distributed_data_object(global_data=obj,
                                    local_shape=(4, 4),
                                    distribution_strategy='freeform')
Ultima's avatar
Ultima committed
498
499
500
501
        assert_equal(obj.get_full_data(), p.get_full_data())


###############################################################################
502

Ultima's avatar
Ultima committed
503
504
    @parameterized.expand(
        itertools.product(
505
506
507
508
509
510
511
512
            [(None, None, None, None, None),
             (None, (8, 8), None, None, None),
                (None, None, np.int_, None, None),
                (1, None, None, None, None),
             (np.array([1, 2, 3]), (3,), None, None, None),
                (None, (4, 4), np.int_, None, None)],
            local_distribution_strategies),
        testcase_func_name=custom_name_func)
513
    def test_failed_init_on_unsufficient_parameters(self,
514
515
516
517
518
519
                                                    (global_data,
                                                     global_shape,
                                                     dtype,
                                                     local_data,
                                                     local_shape),
                                                    distribution_strategy):
Ultima's avatar
Ultima committed
520
521
        assert_raises(ValueError,
                      lambda: distributed_data_object(
522
523
524
525
526
527
                          global_data=global_data,
                          global_shape=global_shape,
                          dtype=dtype,
                          local_data=local_data,
                          local_shape=local_shape,
                          distribution_strategy=distribution_strategy))
Ultima's avatar
Ultima committed
528
529

###############################################################################
530
531

    if size > 1:
Ultima's avatar
Ultima committed
532
533
        @parameterized.expand(
            itertools.product(
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
                [(None, None, np.dtype('int'), None, (2, rank, 3)),
                 (None, None, None, np.arange(2 * rank).reshape((2, rank)),
                  None),
                 (None, None,
                  np.dtype('int') if (rank == 0) else np.dtype('float'),
                  None, (2, 2))],
                local_distribution_strategies),
            testcase_func_name=custom_name_func)
        def test_failed_init_unsufficient_params_mpi(self,
                                                     (global_data,
                                                      global_shape,
                                                      dtype,
                                                      local_data,
                                                      local_shape),
                                                     distribution_strategy):
Ultima's avatar
Ultima committed
549
550
            assert_raises(ValueError,
                          lambda: distributed_data_object(
551
552
553
554
555
556
                              global_data=global_data,
                              global_shape=global_shape,
                              dtype=dtype,
                              local_data=local_data,
                              local_shape=local_shape,
                              distribution_strategy=distribution_strategy))
Ultima's avatar
Ultima committed
557

558
##########################################################################
559

Ultima's avatar
Ultima committed
560
    @parameterized.expand(
561
562
563
        itertools.product([(0,), (1, 0), (0, 1), (25, 0, 10), (0, 0)],
                          local_distribution_strategies),
        testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
564
565
    def test_init_with_zero_type_shape(self, local_shape,
                                       distribution_strategy):
566

567
568
569
570
        obj = distributed_data_object(
                                   local_shape=local_shape,
                                   dtype=np.int,
                                   distribution_strategy=distribution_strategy)
571

Ultima's avatar
Ultima committed
572
573
        expected_global_shape = np.array(local_shape)
        expected_global_shape[0] *= size
574
        expected_global_shape = tuple(expected_global_shape)
Ultima's avatar
Ultima committed
575
576
577
578
579
580
581
582

        assert_equal(obj.shape, expected_global_shape)


###############################################################################
###############################################################################

class Test_init_from_existing_d2o(unittest.TestCase):
583

Ultima's avatar
Ultima committed
584
    @parameterized.expand(
585
586
        itertools.product(all_distribution_strategies,
                          all_distribution_strategies),
Ultima's avatar
Ultima committed
587
588
        testcase_func_name=custom_name_func)
    def test_simple_init_from_existing_d2o(self, old_strat, new_strat):
589
        global_shape = (8, 8)
Ultima's avatar
Ultima committed
590
591
        dtype = np.dtype('float')
        (a, obj) = generate_data(global_shape, dtype, old_strat)
592
593
        p = distributed_data_object(global_data=obj,
                                    distribution_strategy=new_strat)
Ultima's avatar
Ultima committed
594
        if old_strat == 'not' and new_strat in local_distribution_strategies:
595
            old_blown_up = np.concatenate([obj.get_full_data(), ] * size)
Ultima's avatar
Ultima committed
596
597
598
599
600
601
602
            assert_equal(old_blown_up, p.get_full_data())
        else:
            assert_equal(obj.get_full_data(), p.get_full_data())
        assert_equal(obj.distribution_strategy, old_strat)
        assert_equal(p.distribution_strategy, new_strat)


603
604
###############################################################################
###############################################################################
605

606
class Test_set_get_full_and_local_data(unittest.TestCase):
607

608
    @parameterized.expand(
609
610
611
612
613
614
615
616
        itertools.product([(1,), (7,), (2, 7), (77, 11), (256, 256)],
                          all_datatypes,
                          all_distribution_strategies),
        testcase_func_name=custom_name_func)
    def test_init_data_and_get_full_and_local_data(self,
                                                   global_shape,
                                                   dtype,
                                                   distribution_strategy):
Ultima's avatar
Ultima committed
617
        (a, obj) = generate_data(global_shape, dtype,
618
                                 distribution_strategy)
619
620
        assert_equal(obj.get_full_data(), a)

Ultima's avatar
Ultima committed
621
622
###############################################################################

623
    if FOUND['h5py']:
Ultima's avatar
Ultima committed
624
        @parameterized.expand(
625
626
            itertools.product(hdf5_test_paths, hdf5_distribution_strategies),
            testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
627
        def test_loading_hdf5_file(self, (alias, path), distribution_strategy):
628
            a = np.arange(13 * 7).reshape((13, 7)).astype(np.float)
629
            b = a[::-1, ::-1]
630
631
632
633
634
            a = a + 1j * b
            obj = distributed_data_object(
                                  alias=alias,
                                  path=path,
                                  distribution_strategy=distribution_strategy)
635
636
            assert_equal(obj.get_full_data(), a)

637
###############################################################################
638
639

    @parameterized.expand(
640
641
642
643
        itertools.product([(1,), (7,), (2, 7), (77, 11), (256, 256)],
                          all_datatypes,
                          all_distribution_strategies),
        testcase_func_name=custom_name_func)
644
    def test_set_get_full_data(self,
645
646
647
                               global_shape,
                               dtype,
                               distribution_strategy):
Ultima's avatar
Ultima committed
648
        (a, obj) = generate_data(global_shape, dtype,
649
650
                                 distribution_strategy)
        new_a = np.array(5 * a, copy=True, dtype=dtype)
Ultima's avatar
Ultima committed
651
652
653
654
        obj.set_full_data(new_a)
        assert_equal(obj.get_full_data(), new_a)

###############################################################################
655

656
    @parameterized.expand(
657
658
659
660
        itertools.product([(1,), (7,), (2, 7), (77, 11), (256, 256)],
                          all_datatypes,
                          all_distribution_strategies),
        testcase_func_name=custom_name_func)
661
    def test_get_set_local_data(self,
662
663
664
                                global_shape,
                                dtype,
                                distribution_strategy):
Ultima's avatar
Ultima committed
665
        (a, obj) = generate_data(global_shape, dtype,
666
                                 distribution_strategy)
667
        b = obj.get_local_data()
668
669
        c = (np.random.random(b.shape) * 100).astype(np.dtype(dtype))
        obj.set_local_data(data=c)
670
671
        assert_equal(obj.get_local_data(), c)

Ultima's avatar
Ultima committed
672

673
674
##########################################################################
##########################################################################
675

676
class Test_slicing_get_set_data(unittest.TestCase):
677

678
    @parameterized.expand(
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
        itertools.product(
            [(4, 4, 4)],  # (20,21), (256,256)],
            [np.dtype('uint'), np.dtype('int'), np.dtype('float')],
            # all_datatypes,
            all_distribution_strategies,
            [slice(None, None, None),
             slice(5, 18),
             slice(5, 18, 4),
             slice(18, 5, -3),
             slice(6, 14),
             slice(6, 14, 4),
             slice(14, 6, -4),
             slice(5, None),
             slice(5, None, 3),
             slice(None, 5, -2),
             slice(None, 10),
             slice(None, 10, 3),
             slice(10, None, -2),
             slice(None, None, 3),
             slice(None, None, -3),
             slice(2, 2),
             slice(2, 2, 4),
             slice(2, 2, -1),
             slice(2, 2, -4),
             slice(-5, -2),
             slice(-1000, 300),
             slice(5, 300),
             slice(1000, -300, -1),
             slice(1000, -300, -3),
             (3,),
             (2, slice(5, 18)),
             (slice(None), 2),
             (slice(5, 18), slice(18, 5, -1))],
theos's avatar
theos committed
712
            all_distribution_strategies + ['np']),
713
        testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
714
    def test_get_set_slicing_data(self,
715
716
717
718
719
                                  global_shape,
                                  dtype,
                                  distribution_strategy,
                                  slice_tuple,
                                  from_distribution_strategy):
Ultima's avatar
Ultima committed
720
721

        (a, obj) = generate_data(global_shape, dtype,
722
                                 distribution_strategy)
723
724
725

        assert_equal(obj[slice_tuple].get_full_data(), a[slice_tuple])

Ultima's avatar
Ultima committed
726
        if from_distribution_strategy == 'np':
727
            b = a.copy() * 100
Ultima's avatar
Ultima committed
728
            obj[slice_tuple] = b[slice_tuple]
729
            a[slice_tuple] = b[slice_tuple]
Ultima's avatar
Ultima committed
730
731
        else:
            (b, p) = generate_data(global_shape, dtype,
732
                                   from_distribution_strategy)
Ultima's avatar
Ultima committed
733
734
735
            b *= 100
            p *= 100
            obj[slice_tuple] = p[slice_tuple]
736
737
738
739
            a[slice_tuple] = b[slice_tuple]

        assert_equal(obj.get_full_data(), a)

Ultima's avatar
Ultima committed
740
741
742
        a[slice_tuple] = 111
        obj[slice_tuple] = 111

743
744
745
        assert_equal(obj.get_full_data(), a)


Ultima's avatar
Ultima committed
746
747
###############################################################################

Ultima's avatar
Ultima committed
748
749
    @parameterized.expand(all_distribution_strategies,
                          testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
750
751
    def test_get_single_value_from_d2o(self, distribution_strategy):
        (a, obj) = generate_data((4,), np.dtype('float'),
752
                                 distribution_strategy)
Ultima's avatar
Ultima committed
753
754
        assert_equal(obj[0], a[0])

Ultima's avatar
Ultima committed
755

756
###############################################################################
757

Ultima's avatar
Ultima committed
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
    @parameterized.expand(
        itertools.product(all_distribution_strategies,
                          all_distribution_strategies),
        testcase_func_name=custom_name_func)
    def test_single_row_from_d2o(self, distribution_strategy1,
                                 distribution_strategy2):
        (a, obj) = generate_data((8, 8), np.dtype('float'),
                                 distribution_strategy1)
        (b, p) = generate_data((8,), np.dtype('float'),
                               distribution_strategy2)
        a[4] = b
        obj[4] = p
        assert_equal(obj.get_full_data(), a)


###############################################################################
###############################################################################
775

Ultima's avatar
Ultima committed
776
class Test_boolean_get_set_data(unittest.TestCase):
777

Ultima's avatar
Ultima committed
778
    @parameterized.expand(
779
780
781
782
783
784
        itertools.product(
            [(4, 4)],  # (20,21)],
            all_datatypes,
            all_distribution_strategies,
            all_distribution_strategies),
        testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
785
    def test_boolean_get_data(self,
786
787
788
789
                              global_shape,
                              dtype,
                              distribution_strategy_1,
                              distribution_strategy_2):
Ultima's avatar
Ultima committed
790
791

        (a, obj) = generate_data(global_shape, dtype,
792
                                 distribution_strategy_1)
793

Ultima's avatar
Ultima committed
794
        b = ((a < -5) + (a > 4))
795

796
797
        p = obj.copy_empty(dtype=np.dtype('bool'),
                           distribution_strategy=distribution_strategy_2)
Ultima's avatar
Ultima committed
798
        p[:] = b
799
800
801
802

        assert_equal(obj[b].get_full_data(), a[b])
        assert_equal(obj[p].get_full_data(), a[b])

Ultima's avatar
Ultima committed
803
###############################################################################
804

Ultima's avatar
Ultima committed
805
    @parameterized.expand(
806
807
808
809
810
        itertools.product(
            [(4, 4), (20, 21)],
            all_datatypes,
            all_distribution_strategies),
        testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
811
    def test_boolean_set_data(self,
812
813
814
                              global_shape,
                              dtype,
                              distribution_strategy):
Ultima's avatar
Ultima committed
815
816
817

        (a, obj_mem) = generate_data(global_shape, dtype,
                                     distribution_strategy)
818

Ultima's avatar
Ultima committed
819
820
        p = ((obj_mem < -5) + (obj_mem > 4))
        b = ((a < -5) + (a > 4))
821
822
        c = a[b] * 100
        q = obj_mem[p] * 100
823

Ultima's avatar
Ultima committed
824
        a[b] = c
825

Ultima's avatar
Ultima committed
826
827
        obj = obj_mem.copy()
        obj[b] = c
828
        assert_equal(obj.get_full_data(), a)
Ultima's avatar
Ultima committed
829

830
        obj = obj_mem.copy()
Ultima's avatar
Ultima committed
831
        obj[b] = q
832
        assert_equal(obj.get_full_data(), a)
Ultima's avatar
Ultima committed
833

834
        obj = obj_mem.copy()
Ultima's avatar
Ultima committed
835
        obj[p] = q
836
        assert_equal(obj.get_full_data(), a)
Ultima's avatar
Ultima committed
837
838
839
840
841

###############################################################################

    @parameterized.expand(all_distribution_strategies)
    def test_boolean_set_data_of_a_scalar(self, distribution_strategy):
842
        global_shape = (8, 8)
Ultima's avatar
Ultima committed
843
844
        dtype = np.dtype('float')
        (a, obj) = generate_data(global_shape, dtype,
845
                                 distribution_strategy)
846
847
848

        b = a.copy()
        b[a < 3] = 14
Ultima's avatar
Ultima committed
849
850
851
852

        p = obj.copy()
        p[obj < 3] = 14
        assert_equal(p.get_full_data(), b)
853

Ultima's avatar
Ultima committed
854
855
856
        p = obj.copy()
        p[a < 3] = 14
        assert_equal(p.get_full_data(), b)
857
858


Ultima's avatar
Ultima committed
859
860
###############################################################################
###############################################################################
861

Ultima's avatar
Ultima committed
862
class Test_list_get_set_data(unittest.TestCase):
863

Ultima's avatar
Ultima committed
864
    @parameterized.expand(
865
866
867
868
869
870
        itertools.product(
            [(4, 4, 4)],
            all_datatypes,
            all_distribution_strategies,
            all_distribution_strategies),
        testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
871
    def test_list_get_data(self,
872
873
874
875
                           global_shape,
                           dtype,
                           distribution_strategy_1,
                           distribution_strategy_2):
Ultima's avatar
Ultima committed
876
877

        (a, obj) = generate_data(global_shape, dtype,
878
                                 distribution_strategy_1)
879
        w = np.where(a > 28)
880
        p = obj.copy(distribution_strategy=distribution_strategy_2)
881
        wo = (p > 28).where()
882

Ultima's avatar
Ultima committed
883
884
        assert_equal(obj[w].get_full_data(), a[w])
        assert_equal(obj[wo].get_full_data(), a[w])
885

Ultima's avatar
Ultima committed
886
        w0 = 3
887
888
889
        w1 = np.array([0, 1, 3])
        w1o = distributed_data_object(global_data=w1,
                                      distribution_strategy='equal')
890

Ultima's avatar
Ultima committed
891
        w = [w0, w1, w1]
892
        wo = [w0, w1, w1o]
Ultima's avatar
Ultima committed
893
894
895
896
        assert_equal(obj[w].get_full_data(), a[w])
        assert_equal(obj[wo].get_full_data(), a[w])

        w = [w1, w0, w1]
897
        wo = [w1, w0, w1o]
Ultima's avatar
Ultima committed
898
899
900
901
        assert_equal(obj[w].get_full_data(), a[w])
        assert_equal(obj[wo].get_full_data(), a[w])

        w = [w1, w0, w1]
902
        wo = [w1o, w0, w1]
Ultima's avatar
Ultima committed
903
904
905
        assert_equal(obj[w].get_full_data(), a[w])
        assert_equal(obj[wo].get_full_data(), a[w])

906

907
#############################################################################
908

Ultima's avatar
Ultima committed
909
    @parameterized.expand(
910
911
912
913
914
        itertools.product(
            [(4, 4), (20, 21)],
            all_datatypes,
            all_distribution_strategies),
        testcase_func_name=custom_name_func)
Ultima's avatar
Ultima committed
915
    def test_list_set_data(self,
916
917
918
                           global_shape,
                           dtype,
                           distribution_strategy):
Ultima's avatar
Ultima committed
919
920
921

        (a, obj_mem) = generate_data(global_shape, dtype,
                                     distribution_strategy)
922

Ultima's avatar
Ultima committed
923
924
        p = ((obj_mem < -5) + (obj_mem > 4)).where()
        b = np.where((a < -5) + (a > 4))
925
926
        c = a[b] * 100
        q = obj_mem[p] * 100
927

Ultima's avatar
Ultima committed
928
        a[b] = c
929

Ultima's avatar
Ultima committed
930
931
        obj = obj_mem.copy()
        obj[b] = c
932
        assert_equal(obj.get_full_data(), a)
Ultima's avatar
Ultima committed
933

934
        obj = obj_mem.copy()
Ultima's avatar
Ultima committed
935
        obj[b] = q
936
        assert_equal(obj.get_full_data(), a)
Ultima's avatar
Ultima committed
937

938
        obj = obj_mem.copy()
Ultima's avatar
Ultima committed
939
        obj[p] = q
940
        assert_equal(obj.get_full_data(), a)
Ultima's avatar
Ultima committed
941
942
943
944
945

###############################################################################

    @parameterized.expand(all_distribution_strategies)
    def test_list_set_data_of_a_scalar(self, distribution_strategy):
946
        global_shape = (8, 8)
Ultima's avatar
Ultima committed
947
948
        dtype = np.dtype('float')
        (a, obj) = generate_data(global_shape, dtype,
949
                                 distribution_strategy)
950
951
952

        b = a.copy()
        b[np.where(a < 3)] = 14
Ultima's avatar
Ultima committed
953
954
955
956

        p = obj.copy()
        p[(obj < 3).where()] = 14
        assert_equal(p.get_full_data(), b)
957

Ultima's avatar
Ultima committed
958
959
960
961
962
963
964
965
966
        p = obj.copy()
        p[np.where(a < 3)] = 14
        assert_equal(p.get_full_data(), b)


###############################################################################
###############################################################################

class Test_non_local_keys(unittest.TestCase):
967

968
    @parameterized.expand(all_distribution_strategies)
Ultima's avatar
Ultima committed
969
    def test_get_from_non_local_keys(self, distribution_strategy):