nifty_mpi_data.py 108 KB
Newer Older
ultimanet's avatar
ultimanet committed
1
# -*- coding: utf-8 -*-
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# NIFTY (Numerical Information Field Theory) has been developed at the
# Max-Planck-Institute for Astrophysics.
#
# Copyright (C) 2015 Max-Planck-Society
#
# Author: Theo Steininger
# Project homepage: <http://www.mpa-garching.mpg.de/ift/nifty/>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
22
23


ultimanet's avatar
ultimanet committed
24
import numpy as np
25
from weakref import WeakValueDictionary as weakdict
ultimanet's avatar
ultimanet committed
26

27
28
29
from keepers import about,\
                    global_configuration as gc,\
                    global_dependency_injector as gdi
ultimanet's avatar
ultimanet committed
30

31
32
33
MPI = gdi[gc['mpi_module']]
h5py = gdi.get('h5py')
pyfftw = gdi.get('pyfftw')
ultimanet's avatar
ultimanet committed
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
## initialize the 'FOUND-packages'-dictionary
#FOUND = {}
#try:
#    from mpi4py import MPI
#    FOUND['MPI'] = True
#except(ImportError):
#    import mpi_dummy as MPI
#    FOUND['MPI'] = False
#
#try:
#    import pyfftw
#    FOUND['pyfftw'] = True
#except(ImportError):
#    FOUND['pyfftw'] = False
#
#try:
#    import h5py
#    FOUND['h5py'] = True
#    FOUND['h5py_parallel'] = h5py.get_config().mpi
#except(ImportError):
#    FOUND['h5py'] = False
#    FOUND['h5py_parallel'] = False

_maybe_fftw = ['fftw'] if ('pyfftw' in gdi) else []
STRATEGIES = {
                'all': ['not', 'equal', 'freeform'] + _maybe_fftw,
                'global': ['not', 'equal'] + _maybe_fftw,
                'local': ['freeform'],
                'slicing': ['equal', 'freeform'] + _maybe_fftw,
                'not': ['not'],
                'hdf5': ['equal'] + _maybe_fftw,
             }
68

69

ultimanet's avatar
ultimanet committed
70
71
72
73
74
75
76
77
class distributed_data_object(object):
    """

        NIFTY class for distributed data

        Parameters
        ----------
        global_data : {tuple, list, numpy.ndarray} *at least 1-dimensional*
78
            Initial data which will be casted to a numpy.ndarray and then
ultimanet's avatar
ultimanet committed
79
80
81
82
83
84
            stored according to the distribution strategy. The global_data's
            shape overwrites global_shape.
        global_shape : tuple of ints, *optional*
            If no global_data is supplied, global_shape can be used to
            initialize an empty distributed_data_object
        dtype : type, *optional*
85
86
            If an explicit dtype is supplied, the given global_data will be
            casted to it.
ultimanet's avatar
ultimanet committed
87
        distribution_strategy : {'fftw' (default), 'not'}, *optional*
88
89
            Specifies the way, how global_data will be distributed to the
            individual nodes.
ultimanet's avatar
ultimanet committed
90
            'fftw' follows the distribution strategy of pyfftw.
91
92
            'not' does not distribute the data at all.

ultimanet's avatar
ultimanet committed
93
94
95
96
97
98
99
100
101
102

        Attributes
        ----------
        data : numpy.ndarray
            The numpy.ndarray in which the individual node's data is stored.
        dtype : type
            Data type of the data object.
        distribution_strategy : string
            Name of the used distribution_strategy
        distributor : distributor
103
104
            The distributor object which takes care of all distribution and
            consolidation of the data.
ultimanet's avatar
ultimanet committed
105
106
        shape : tuple of int
            The global shape of the data
107

ultimanet's avatar
ultimanet committed
108
109
        Raises
        ------
110
111
112
        TypeError :
            If the supplied distribution strategy is not known.

ultimanet's avatar
ultimanet committed
113
    """
114

115
    def __init__(self, global_data=None, global_shape=None, dtype=None,
Ultima's avatar
Ultima committed
116
                 local_data=None, local_shape=None,
117
                 distribution_strategy='fftw', hermitian=False,
118
119
120
121
122
                 alias=None, path=None, comm=MPI.COMM_WORLD,
                 copy=True, *args, **kwargs):

        # TODO: allow init with empty shape

123
124
125
126
        if isinstance(global_data, tuple) or isinstance(global_data, list):
            global_data = np.array(global_data, copy=False)
        if isinstance(local_data, tuple) or isinstance(local_data, list):
            local_data = np.array(local_data, copy=False)
127

128
        self.distributor = distributor_factory.get_distributor(
129
130
131
132
133
134
135
136
137
                                distribution_strategy=distribution_strategy,
                                comm=comm,
                                global_data=global_data,
                                global_shape=global_shape,
                                local_data=local_data,
                                local_shape=local_shape,
                                alias=alias,
                                path=path,
                                dtype=dtype,
138
                                **kwargs)
139

ultimanet's avatar
ultimanet committed
140
141
142
        self.distribution_strategy = distribution_strategy
        self.dtype = self.distributor.dtype
        self.shape = self.distributor.global_shape
Ultima's avatar
Ultima committed
143
144
        self.local_shape = self.distributor.local_shape
        self.comm = self.distributor.comm
145
146

        self.init_args = args
147
        self.init_kwargs = kwargs
148

Ultima's avatar
Ultima committed
149
        (self.data, self.hermitian) = self.distributor.initialize_data(
150
151
152
153
154
155
            global_data=global_data,
            local_data=local_data,
            alias=alias,
            path=path,
            hermitian=hermitian,
            copy=copy)
156
        self.index = d2o_librarian.register(self)
157

158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
    @property
    def real(self):
        new_data = self.get_local_data().real
        new_dtype = new_data.dtype
        new_d2o = self.copy_empty(dtype=new_dtype)
        new_d2o.set_local_data(data=new_data,
                               hermitian=self.hermitian)
        return new_d2o

    @property
    def imag(self):
        new_data = self.get_local_data().imag
        new_dtype = new_data.dtype
        new_d2o = self.copy_empty(dtype=new_dtype)
        new_d2o.set_local_data(data=new_data,
                               hermitian=self.hermitian)
        return new_d2o

Ultimanet's avatar
Ultimanet committed
176
    def copy(self, dtype=None, distribution_strategy=None, **kwargs):
177
178
179
        temp_d2o = self.copy_empty(dtype=dtype,
                                   distribution_strategy=distribution_strategy,
                                   **kwargs)
Ultima's avatar
Ultima committed
180
        if distribution_strategy is None or \
181
                distribution_strategy == self.distribution_strategy:
Ultimanet's avatar
Ultimanet committed
182
183
            temp_d2o.set_local_data(self.get_local_data(), copy=True)
        else:
Ultima's avatar
Ultima committed
184
            temp_d2o.inject((slice(None),), self, (slice(None),))
185
        temp_d2o.hermitian = self.hermitian
186
        return temp_d2o
187
188

    def copy_empty(self, global_shape=None, local_shape=None, dtype=None,
189
                   distribution_strategy=None, **kwargs):
Ultima's avatar
Ultima committed
190
        if self.distribution_strategy == 'not' and \
191
                distribution_strategy in STRATEGIES['local'] and \
192
193
194
195
196
                local_shape is None:
            result = self.copy_empty(global_shape=global_shape,
                                     local_shape=local_shape,
                                     dtype=dtype,
                                     distribution_strategy='equal',
Ultima's avatar
Ultima committed
197
                                     **kwargs)
198
            return result.copy_empty(distribution_strategy='freeform')
199

Ultima's avatar
Ultima committed
200
        if global_shape is None:
201
            global_shape = self.shape
Ultima's avatar
Ultima committed
202
203
204
        if local_shape is None:
            local_shape = self.local_shape
        if dtype is None:
205
            dtype = self.dtype
Ultima's avatar
Ultima committed
206
        if distribution_strategy is None:
207
208
209
            distribution_strategy = self.distribution_strategy

        kwargs.update(self.init_kwargs)
210

211
212
213
214
215
216
217
218
        temp_d2o = distributed_data_object(
                                   global_shape=global_shape,
                                   local_shape=local_shape,
                                   dtype=dtype,
                                   distribution_strategy=distribution_strategy,
                                   comm=self.comm,
                                   *self.init_args,
                                   **kwargs)
219
        return temp_d2o
220

221
    def apply_scalar_function(self, function, inplace=False, dtype=None):
222
        remember_hermitianQ = self.hermitian
223

224
        if inplace is True:
Ultimanet's avatar
Ultimanet committed
225
            temp = self
Ultima's avatar
Ultima committed
226
            if dtype is not None and self.dtype != np.dtype(dtype):
227
228
229
                about.warnings.cprint(
                    "WARNING: Inplace dtype conversion is not possible!")

Ultimanet's avatar
Ultimanet committed
230
        else:
231
            temp = self.copy_empty(dtype=dtype)
Ultimanet's avatar
Ultimanet committed
232

Ultima's avatar
Ultima committed
233
        if np.prod(self.local_shape) != 0:
234
            try:
Ultima's avatar
Ultima committed
235
236
237
238
                temp.data[:] = function(self.data)
            except:
                temp.data[:] = np.vectorize(function)(self.data)
        else:
239
240
            # Noting to do here. The value-empty array
            # is also geometrically empty
Ultima's avatar
Ultima committed
241
            pass
242

243
244
245
246
        if function in (np.exp, np.log):
            temp.hermitian = remember_hermitianQ
        else:
            temp.hermitian = False
Ultimanet's avatar
Ultimanet committed
247
        return temp
248

Ultimanet's avatar
Ultimanet committed
249
250
251
    def apply_generator(self, generator):
        self.set_local_data(generator(self.distributor.local_shape))
        self.hermitian = False
252

ultimanet's avatar
ultimanet committed
253
254
    def __str__(self):
        return self.data.__str__()
255

ultimanet's avatar
ultimanet committed
256
    def __repr__(self):
257
        return '<distributed_data_object>\n' + self.data.__repr__()
258

259
    def _compare_helper(self, other, op):
260
        result = self.copy_empty(dtype=np.bool_)
261
262
        # Case 1: 'other' is a scalar
        # -> make point-wise comparison
Ultimanet's avatar
Ultimanet committed
263
        if np.isscalar(other):
264
            result.set_local_data(
265
                getattr(self.get_local_data(copy=False), op)(other))
266
            return result
Ultimanet's avatar
Ultimanet committed
267

268
269
        # Case 2: 'other' is a numpy array or a distributed_data_object
        # -> extract the local data and make point-wise comparison
Ultimanet's avatar
Ultimanet committed
270
        elif isinstance(other, np.ndarray) or\
271
                isinstance(other, distributed_data_object):
Ultimanet's avatar
Ultimanet committed
272
            temp_data = self.distributor.extract_local_data(other)
273
274
            result.set_local_data(
                getattr(self.get_local_data(copy=False), op)(temp_data))
Ultimanet's avatar
Ultimanet committed
275
            return result
276
277

        # Case 3: 'other' is None
Ultima's avatar
Ultima committed
278
        elif other is None:
Ultimanet's avatar
Ultimanet committed
279
            return False
280
281
282

        # Case 4: 'other' is something different
        # -> make a numpy casting and make a recursive call
Ultimanet's avatar
Ultimanet committed
283
284
        else:
            temp_other = np.array(other)
285
            return getattr(self, op)(temp_other)
286

287
288
    def __ne__(self, other):
        return self._compare_helper(other, '__ne__')
289

290
291
    def __lt__(self, other):
        return self._compare_helper(other, '__lt__')
292

293
294
295
296
297
298
    def __le__(self, other):
        return self._compare_helper(other, '__le__')

    def __eq__(self, other):

        return self._compare_helper(other, '__eq__')
299

300
301
302
303
304
305
    def __ge__(self, other):
        return self._compare_helper(other, '__ge__')

    def __gt__(self, other):
        return self._compare_helper(other, '__gt__')

Ultimanet's avatar
Ultimanet committed
306
    def equal(self, other):
Ultimanet's avatar
Ultimanet committed
307
308
309
310
311
312
313
314
315
        if other is None:
            return False
        try:
            assert(self.dtype == other.dtype)
            assert(self.shape == other.shape)
            assert(self.init_args == other.init_args)
            assert(self.init_kwargs == other.init_kwargs)
            assert(self.distribution_strategy == other.distribution_strategy)
            assert(np.all(self.data == other.data))
Ultimanet's avatar
Ultimanet committed
316
        except(AssertionError, AttributeError):
Ultimanet's avatar
Ultimanet committed
317
318
319
320
            return False
        else:
            return True

321
    def __pos__(self):
322
        temp_d2o = self.copy_empty()
323
        temp_d2o.set_local_data(data=self.get_local_data(), copy=True)
324
        return temp_d2o
325

ultimanet's avatar
ultimanet committed
326
    def __neg__(self):
327
        temp_d2o = self.copy_empty()
328
329
        temp_d2o.set_local_data(data=self.get_local_data().__neg__(),
                                copy=True)
ultimanet's avatar
ultimanet committed
330
        return temp_d2o
331

332
    def __abs__(self):
333
        # translate complex dtypes
334
335
336
337
338
339
        if self.dtype == np.dtype('complex64'):
            new_dtype = np.dtype('float32')
        elif self.dtype == np.dtype('complex128'):
            new_dtype = np.dtype('float64')
        elif issubclass(self.dtype.type, np.complexfloating):
            new_dtype = np.dtype('float')
Ultimanet's avatar
Ultimanet committed
340
341
        else:
            new_dtype = self.dtype
342
343
344
        temp_d2o = self.copy_empty(dtype=new_dtype)
        temp_d2o.set_local_data(data=self.get_local_data().__abs__(),
                                copy=True)
345
        return temp_d2o
346

Ultima's avatar
Ultima committed
347
348
349
350
351
    def _builtin_helper(self, operator, other, inplace=False):
        if isinstance(other, distributed_data_object):
            other_is_real = other.isreal()
        else:
            other_is_real = np.isreal(other)
352
353

        # Case 1: other is not a scalar
Ultimanet's avatar
Ultimanet committed
354
        if not (np.isscalar(other) or np.shape(other) == (1,)):
355
            try:
356
                hermitian_Q = (other.hermitian and self.hermitian)
357
358
            except(AttributeError):
                hermitian_Q = False
359
            # extract the local data from the 'other' object
Ultimanet's avatar
Ultimanet committed
360
361
            temp_data = self.distributor.extract_local_data(other)
            temp_data = operator(temp_data)
362
363

        # Case 2: other is a real scalar -> preserve hermitianity
Ultima's avatar
Ultima committed
364
365
        elif other_is_real or (self.dtype not in (np.dtype('complex128'),
                                                  np.dtype('complex256'))):
366
            hermitian_Q = self.hermitian
ultimanet's avatar
ultimanet committed
367
            temp_data = operator(other)
368
        # Case 3: other is complex
369
370
        else:
            hermitian_Q = False
371
372
            temp_data = operator(other)
        # write the new data into a new distributed_data_object
373
        if inplace is True:
374
375
            temp_d2o = self
        else:
376
            # use common datatype for self and other
377
            new_dtype = np.dtype(np.find_common_type((self.dtype,),
378
                                                     (temp_data.dtype,)))
379
            temp_d2o = self.copy_empty(
380
                dtype=new_dtype)
ultimanet's avatar
ultimanet committed
381
        temp_d2o.set_local_data(data=temp_data)
382
        temp_d2o.hermitian = hermitian_Q
ultimanet's avatar
ultimanet committed
383
        return temp_d2o
384

ultimanet's avatar
ultimanet committed
385
    def __add__(self, other):
Ultima's avatar
Ultima committed
386
        return self._builtin_helper(self.get_local_data().__add__, other)
ultimanet's avatar
ultimanet committed
387
388

    def __radd__(self, other):
Ultima's avatar
Ultima committed
389
        return self._builtin_helper(self.get_local_data().__radd__, other)
Ultimanet's avatar
Ultimanet committed
390
391

    def __iadd__(self, other):
392
        return self._builtin_helper(self.get_local_data().__iadd__,
393
394
                                    other,
                                    inplace=True)
Ultimanet's avatar
Ultimanet committed
395

ultimanet's avatar
ultimanet committed
396
    def __sub__(self, other):
Ultima's avatar
Ultima committed
397
        return self._builtin_helper(self.get_local_data().__sub__, other)
398

ultimanet's avatar
ultimanet committed
399
    def __rsub__(self, other):
Ultima's avatar
Ultima committed
400
        return self._builtin_helper(self.get_local_data().__rsub__, other)
401

ultimanet's avatar
ultimanet committed
402
    def __isub__(self, other):
403
        return self._builtin_helper(self.get_local_data().__isub__,
404
405
                                    other,
                                    inplace=True)
406

ultimanet's avatar
ultimanet committed
407
    def __div__(self, other):
Ultima's avatar
Ultima committed
408
        return self._builtin_helper(self.get_local_data().__div__, other)
409

410
411
    def __truediv__(self, other):
        return self.__div__(other)
412

ultimanet's avatar
ultimanet committed
413
    def __rdiv__(self, other):
Ultima's avatar
Ultima committed
414
        return self._builtin_helper(self.get_local_data().__rdiv__, other)
415

416
417
    def __rtruediv__(self, other):
        return self.__rdiv__(other)
ultimanet's avatar
ultimanet committed
418

Ultimanet's avatar
Ultimanet committed
419
    def __idiv__(self, other):
420
        return self._builtin_helper(self.get_local_data().__idiv__,
421
422
423
                                    other,
                                    inplace=True)

424
    def __itruediv__(self, other):
425
        return self.__idiv__(other)
426

ultimanet's avatar
ultimanet committed
427
    def __floordiv__(self, other):
428
        return self._builtin_helper(self.get_local_data().__floordiv__,
429
430
                                    other)

ultimanet's avatar
ultimanet committed
431
    def __rfloordiv__(self, other):
432
        return self._builtin_helper(self.get_local_data().__rfloordiv__,
433
434
                                    other)

Ultimanet's avatar
Ultimanet committed
435
    def __ifloordiv__(self, other):
Ultima's avatar
Ultima committed
436
        return self._builtin_helper(
437
438
            self.get_local_data().__ifloordiv__, other,
            inplace=True)
439

ultimanet's avatar
ultimanet committed
440
    def __mul__(self, other):
Ultima's avatar
Ultima committed
441
        return self._builtin_helper(self.get_local_data().__mul__, other)
442

ultimanet's avatar
ultimanet committed
443
    def __rmul__(self, other):
Ultima's avatar
Ultima committed
444
        return self._builtin_helper(self.get_local_data().__rmul__, other)
ultimanet's avatar
ultimanet committed
445
446

    def __imul__(self, other):
447
        return self._builtin_helper(self.get_local_data().__imul__,
448
449
                                    other,
                                    inplace=True)
Ultimanet's avatar
Ultimanet committed
450

ultimanet's avatar
ultimanet committed
451
    def __pow__(self, other):
Ultima's avatar
Ultima committed
452
        return self._builtin_helper(self.get_local_data().__pow__, other)
453

ultimanet's avatar
ultimanet committed
454
    def __rpow__(self, other):
Ultima's avatar
Ultima committed
455
        return self._builtin_helper(self.get_local_data().__rpow__, other)
ultimanet's avatar
ultimanet committed
456
457

    def __ipow__(self, other):
458
        return self._builtin_helper(self.get_local_data().__ipow__,
459
460
461
                                    other,
                                    inplace=True)

Ultima's avatar
Ultima committed
462
463
    def __mod__(self, other):
        return self._builtin_helper(self.get_local_data().__mod__, other)
464

Ultima's avatar
Ultima committed
465
    def __rmod__(self, other):
466
        return self._builtin_helper(self.get_local_data().__rmod__, other)
467

Ultima's avatar
Ultima committed
468
    def __imod__(self, other):
469
        return self._builtin_helper(self.get_local_data().__imod__,
470
471
472
                                    other,
                                    inplace=True)

473
474
    def __len__(self):
        return self.shape[0]
475

476
    def get_dim(self):
477
        return np.prod(self.shape)
478

479
    def vdot(self, other):
480
        other = self.distributor.extract_local_data(other)
481
482
483
484
        local_vdot = np.vdot(self.get_local_data(), other)
        local_vdot_list = self.distributor._allgather(local_vdot)
        global_vdot = np.sum(local_vdot_list)
        return global_vdot
Ultimanet's avatar
Ultimanet committed
485

ultimanet's avatar
ultimanet committed
486
    def __getitem__(self, key):
Ultima's avatar
Ultima committed
487
        return self.get_data(key)
488

ultimanet's avatar
ultimanet committed
489
490
    def __setitem__(self, key, data):
        self.set_data(data, key)
491

492
    def _contraction_helper(self, function, **kwargs):
Ultima's avatar
Ultima committed
493
494
495
496
497
498
499
        if np.prod(self.data.shape) == 0:
            local = 0
            include = False
        else:
            local = function(self.data, **kwargs)
            include = True

500
        local_list = self.distributor._allgather(local)
501
        local_list = np.array(local_list, dtype=np.dtype(local_list[0]))
Ultima's avatar
Ultima committed
502
503
504
        include_list = np.array(self.distributor._allgather(include))
        work_list = local_list[include_list]
        if work_list.shape[0] == 0:
505
            raise ValueError("ERROR: Zero-size array to reduction operation " +
Ultima's avatar
Ultima committed
506
                             "which has no identity")
507
        else:
Ultima's avatar
Ultima committed
508
509
            result = function(work_list, axis=0)
            return result
510

511
512
513
    def min(self, **kwargs):
        return self.amin(**kwargs)

514
    def amin(self, **kwargs):
515
        return self._contraction_helper(np.amin, **kwargs)
516
517

    def nanmin(self, **kwargs):
518
        return self._contraction_helper(np.nanmin, **kwargs)
519

520
521
522
    def max(self, **kwargs):
        return self.amax(**kwargs)

523
    def amax(self, **kwargs):
524
        return self._contraction_helper(np.amax, **kwargs)
525

526
    def nanmax(self, **kwargs):
527
        return self._contraction_helper(np.nanmax, **kwargs)
528

529
530
531
532
    def sum(self, **kwargs):
        return self._contraction_helper(np.sum, **kwargs)

    def prod(self, **kwargs):
533
534
        return self._contraction_helper(np.prod, **kwargs)

535
    def mean(self, power=1):
536
        # compute the local means and the weights for the mean-mean.
Ultima's avatar
Ultima committed
537
538
539
540
541
542
        if np.prod(self.data.shape) == 0:
            local_mean = 0
            include = False
        else:
            local_mean = np.mean(self.data**power)
            include = True
543

544
        local_weight = np.prod(self.data.shape)
545
        # collect the local means and cast the result to a ndarray
Ultima's avatar
Ultima committed
546
547
        local_mean_list = self.distributor._allgather(local_mean)
        local_weight_list = self.distributor._allgather(local_weight)
548

549
550
        local_mean_list = np.array(local_mean_list,
                                   dtype=np.dtype(local_mean_list[0]))
551
552
        local_weight_list = np.array(local_weight_list)
        # extract the parts from the non-empty nodes
Ultima's avatar
Ultima committed
553
554
555
556
557
        include_list = np.array(self.distributor._allgather(include))
        work_mean_list = local_mean_list[include_list]
        work_weight_list = local_weight_list[include_list]
        if work_mean_list.shape[0] == 0:
            raise ValueError("ERROR:  Mean of empty slice.")
558
559
        else:
            # compute the denominator for the weighted mean-mean
Ultima's avatar
Ultima committed
560
            global_weight = np.sum(work_weight_list)
561
            # compute the numerator
Ultima's avatar
Ultima committed
562
            numerator = np.sum(work_mean_list * work_weight_list)
563
            global_mean = numerator / global_weight
Ultima's avatar
Ultima committed
564
            return global_mean
565
566
567
568
569

    def var(self):
        mean_of_the_square = self.mean(power=2)
        square_of_the_mean = self.mean()**2
        return mean_of_the_square - square_of_the_mean
570

571
572
    def std(self):
        return np.sqrt(self.var())
573

Ultima's avatar
Ultima committed
574
575
576
577
578
579
580
    def argmin(self):
        if np.prod(self.data.shape) == 0:
            local_argmin = np.nan
            local_argmin_value = np.nan
            globalized_local_argmin = np.nan
        else:
            local_argmin = np.argmin(self.data)
581
            local_argmin_value = self.data[np.unravel_index(local_argmin,
Ultima's avatar
Ultima committed
582
                                                            self.data.shape)]
583

Ultima's avatar
Ultima committed
584
            globalized_local_argmin = self.distributor.globalize_flat_index(
585
586
587
588
                local_argmin)
        local_argmin_list = self.distributor._allgather(
                                                    (local_argmin_value,
                                                     globalized_local_argmin))
589
        local_argmin_list = np.array(local_argmin_list, dtype=[
590
591
            ('value', np.dtype('complex128')),
            ('index', np.dtype('float'))])
592
593
        local_argmin_list = np.sort(local_argmin_list,
                                    order=['value', 'index'])
Ultima's avatar
Ultima committed
594
        return np.int(local_argmin_list[0][1])
595

Ultima's avatar
Ultima committed
596
597
598
599
600
601
602
    def argmax(self):
        if np.prod(self.data.shape) == 0:
            local_argmax = np.nan
            local_argmax_value = np.nan
            globalized_local_argmax = np.nan
        else:
            local_argmax = np.argmax(self.data)
603
            local_argmax_value = -self.data[np.unravel_index(local_argmax,
604
                                                             self.data.shape)]
Ultima's avatar
Ultima committed
605
            globalized_local_argmax = self.distributor.globalize_flat_index(
606
607
608
609
                local_argmax)
        local_argmax_list = self.distributor._allgather(
                                                  (local_argmax_value,
                                                   globalized_local_argmax))
610
        local_argmax_list = np.array(local_argmax_list, dtype=[
611
612
            ('value', np.dtype('complex128')),
            ('index', np.dtype('float'))])
613
614
        local_argmax_list = np.sort(local_argmax_list,
                                    order=['value', 'index'])
Ultima's avatar
Ultima committed
615
        return np.int(local_argmax_list[0][1])
616

617
    def argmin_nonflat(self):
Ultima's avatar
Ultima committed
618
        return np.unravel_index(self.argmin(), self.shape)
619

Ultima's avatar
Ultima committed
620
621
    def argmax_nonflat(self):
        return np.unravel_index(self.argmax(), self.shape)
622

623
624
625
626
627
628
629
    def conjugate(self):
        temp_d2o = self.copy_empty()
        temp_data = np.conj(self.get_local_data())
        temp_d2o.set_local_data(temp_data)
        return temp_d2o

    def conj(self):
630
631
        return self.conjugate()

632
    def median(self):
633
        about.warnings.cprint(
634
635
636
            "WARNING: The current implementation of median is very expensive!")
        median = np.median(self.get_full_data())
        return median
637

638
    def _is_helper(self, function):
Ultima's avatar
Ultima committed
639
        temp_d2o = self.copy_empty(dtype=np.dtype('bool'))
640
        temp_d2o.set_local_data(function(self.data))
641
        return temp_d2o
642

643
644
645
    def iscomplex(self):
        return self._is_helper(np.iscomplex)

646
    def isreal(self):
647
648
649
650
651
652
653
654
655
656
657
658
659
660
        return self._is_helper(np.isreal)

    def isnan(self):
        return self._is_helper(np.isnan)

    def isinf(self):
        return self._is_helper(np.isinf)

    def isfinite(self):
        return self._is_helper(np.isfinite)

    def nan_to_num(self):
        temp_d2o = self.copy_empty()
        temp_d2o.set_local_data(np.nan_to_num(self.data))
661
        return temp_d2o
662

663
664
665
666
667
668
669
670
    def all(self):
        local_all = np.all(self.get_local_data())
        global_all = self.distributor._allgather(local_all)
        return np.all(global_all)

    def any(self):
        local_any = np.any(self.get_local_data())
        global_any = self.distributor._allgather(local_any)
671
        return np.any(global_any)
672

673
674
675
676
677
    def unique(self):
        local_unique = np.unique(self.get_local_data())
        global_unique = self.distributor._allgather(local_unique)
        global_unique = np.concatenate(global_unique)
        return np.unique(global_unique)
678

679
    def bincount(self, weights=None, minlength=None):
680
        if self.dtype not in [np.dtype('int16'), np.dtype('int32'),
681
682
                              np.dtype('int64'),  np.dtype('uint16'),
                              np.dtype('uint32'), np.dtype('uint64')]:
683
            raise TypeError(about._errors.cstring(
684
685
                "ERROR: Distributed-data-object must be of integer datatype!"))

686
        minlength = max(self.amax() + 1, minlength)
687

688
689
        if weights is not None:
            local_weights = self.distributor.extract_local_data(weights).\
690
                flatten()
691
692
        else:
            local_weights = None
693

694
        local_counts = np.bincount(self.get_local_data().flatten(),
695
696
                                   weights=local_weights,
                                   minlength=minlength)
Ultima's avatar
Ultima committed
697
698
699
700
        if self.distribution_strategy == 'not':
            return local_counts
        else:
            list_of_counts = self.distributor._allgather(local_counts)
701
            counts = np.sum(list_of_counts, axis=0)
Ultima's avatar
Ultima committed
702
            return counts
703

Ultima's avatar
Ultima committed
704
705
    def where(self):
        return self.distributor.where(self.data)
706

707
    def set_local_data(self, data, hermitian=False, copy=True):
ultimanet's avatar
ultimanet committed
708
        """
709
            Stores data directly in the local data attribute. No distribution
ultimanet's avatar
ultimanet committed
710
711
712
713
714
            is done. The shape of the data must fit the local data attributes
            shape.

            Parameters
            ----------
715
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
716
                The data which should be stored in the local data attribute.
717

ultimanet's avatar
ultimanet committed
718
719
720
            Returns
            -------
            None
721

ultimanet's avatar
ultimanet committed
722
        """
Ultimanet's avatar
Ultimanet committed
723
        self.hermitian = hermitian
724
        if copy is True:
Ultima's avatar
Ultima committed
725
726
            self.data[:] = data
        else:
727
728
729
730
            self.data = np.array(data,
                                 dtype=self.dtype,
                                 copy=False,
                                 order='C').reshape(self.local_shape)
731

Ultima's avatar
Ultima committed
732
    def set_data(self, data, to_key, from_key=None, local_keys=False,
733
                 hermitian=False, copy=True, **kwargs):
ultimanet's avatar
ultimanet committed
734
        """
735
            Stores the supplied data in the region which is specified by key.
ultimanet's avatar
ultimanet committed
736
            The data is distributed according to the distribution strategy. If
737
            the individual nodes get different key-arguments. Their data is
ultimanet's avatar
ultimanet committed
738
            processed one-by-one.
739

ultimanet's avatar
ultimanet committed
740
741
            Parameters
            ----------
742
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
743
744
                The data which should be distributed.
            key : int, slice, tuple of int or slice
745
746
747
                The key is the object which specifies the region, where data
                will be stored in.

ultimanet's avatar
ultimanet committed
748
749
750
            Returns
            -------
            None
751

ultimanet's avatar
ultimanet committed
752
        """
Ultimanet's avatar
Ultimanet committed
753
        self.hermitian = hermitian
754
755
756
757
758
759
        self.distributor.disperse_data(data=self.data,
                                       to_key=to_key,
                                       data_update=data,
                                       from_key=from_key,
                                       local_keys=local_keys,
                                       copy=copy,
Ultima's avatar
Ultima committed
760
                                       **kwargs)
761

762
    def set_full_data(self, data, hermitian=False, copy=True, **kwargs):
ultimanet's avatar
ultimanet committed
763
        """
764
            Distributes the supplied data to the nodes. The shape of data must
ultimanet's avatar
ultimanet committed
765
            match the shape of the distributed_data_object.
766

ultimanet's avatar
ultimanet committed
767
768
            Parameters
            ----------
769
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
770
                The data which should be distributed.
771

ultimanet's avatar
ultimanet committed
772
773
            Notes
            -----
774
            set_full_data(foo) is equivalent to set_data(foo,slice(None)) but
ultimanet's avatar
ultimanet committed
775
            faster.
776

ultimanet's avatar
ultimanet committed
777
778
779
            Returns
            -------
            None
780

ultimanet's avatar
ultimanet committed
781
        """
Ultimanet's avatar
Ultimanet committed
782
        self.hermitian = hermitian
783
        self.data = self.distributor.distribute_data(data=data, copy=copy,
784
                                                     **kwargs)
ultimanet's avatar
ultimanet committed
785

Ultimanet's avatar
Ultimanet committed
786
    def get_local_data(self, key=(slice(None),), copy=True):
ultimanet's avatar
ultimanet committed
787
        """
788
789
            Loads data directly from the local data attribute. No consolidation
            is done.
ultimanet's avatar
ultimanet committed
790
791
792
793

            Parameters
            ----------
            key : int, slice, tuple of int or slice
794
795
                The key which will be used to access the data.

ultimanet's avatar
ultimanet committed
796
797
798
            Returns
            -------
            self.data[key] : numpy.ndarray
799

Ultimanet's avatar
Ultimanet committed
800
        """
801
        if copy is True:
802
            return self.data[key]
803
        if copy is False:
Ultimanet's avatar
Ultimanet committed
804
            return self.data
805

806
    def get_data(self, key, local_keys=False, **kwargs):
ultimanet's avatar
ultimanet committed
807
        """
808
809
            Loads data from the region which is specified by key. The data is
            consolidated according to the distribution strategy. If the
ultimanet's avatar
ultimanet committed
810
            individual nodes get different key-arguments, they get individual
811
812
            data.