nifty_mpi_data.py 145 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
_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,
             }
Ultima's avatar
Ultima committed
44
45
46
47
if _maybe_fftw != []:
    _default_strategy = 'fftw'
else:
    _default_strategy = 'equal'
48

49

ultimanet's avatar
ultimanet committed
50
class distributed_data_object(object):
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
    """A multidimensional array with modular MPI-based distribution schemes.

    The purpose of a distributed_data_object (d2o) is to provide the user
    with a numpy.ndarray like interface while storing the data on an arbitrary
    number of MPI nodes. The logic of a certain distribution strategy is
    implemented by an associated distributor.

    Parameters
    ----------
    global_data : array-like, at least 1-dimensional
        Used with global-type distribution strategies in order to fill the
        d2o with data during initialization.
    global_shape : tuple of ints
        Used with global-type distribution strategies. If no global_data is
        supplied, it will be used.
    dtype : {np.dtype, type}
        Used as the d2o's datatype. Overwrites the data-type of any init data.
    local_data : array-like, at least 1-dimensional
        Used with local-type distribution strategies in order to fill the
        d2o with data during initialization.
    local_shape : tuple of ints
        Used with local-type distribution strategies. If no local_data is
        supplied, local_shape will be used.
ultimanet's avatar
ultimanet committed
74
    distribution_strategy : optional[{'fftw', 'equal', 'not', 'freeform'}]
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
        Specifies which distributor will be created and used.
        'fftw'      uses the distribution strategy of pyfftw,
        'equal'     tries to  distribute the data as uniform as possible
        'not'       does not distribute the data at all
        'freeform'  distribute the data according to the given local data/shape
    hermitian : boolean
        Specifies if the given init-data is hermitian or not. The
        self.hermitian attribute will be set accordingly.
    alias : String
        Used in order to initialize the d2o from a hdf5 file.
    path : String
        Used in order to initialize the d2o from a hdf5 file. If no path is
        given, '$working_directory/alias' is used.
    comm : mpi4py.MPI.Intracomm
        The MPI communicator on which the d2o lives.
    copy : boolean
        If true it is guaranteed that the input data will be copied. If false
        copying is tried to be avoided.
    *args
        Although not directly used during the init process, further parameters
        are stored in the self.init_args attribute.
    **kwargs
        Additional keyword arguments are passed to the distributor_factory and
        furthermore get stored in the self.init_kwargs attribute.
    skip_parsing : boolean (optional keyword argument)
        If true, the distribution_factory will skip all sanity checks and
        completions of the given (keyword-)arguments. It just uses what it
        gets. Hence the user is fully responsible for supplying complete and
        consistent parameters. This can be used in order to speed up the init
        process. Also see notes section.

    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
        The distributor object which takes care of all distribution and
        consolidation of the data.
    shape : tuple of int
        The global shape of the data.
    local_shape : tuple of int
        The nodes individual local shape of the stored data.
    comm : mpi4py.MPI.Intracomm
        The MPI communicator on which the d2o lives.
    hermitian : boolean
        Specfies whether the d2o's data definitely possesses hermitian
        symmetry.
    index : int
        The d2o's registration index it got from the d2o_librarian.
    init_args : list
        Any additional initialization arguments are stored here.
    init_kwargs : dict
        Any additional initialization keyword arguments are stored here.

    Raises
    ------
    ValueError
        Raised if
ultimanet's avatar
ultimanet committed
137
138
            * the supplied distribution strategy is not known
            * comm is None
139
            * different distribution strategies where given on the
ultimanet's avatar
ultimanet committed
140
141
              individual nodes
            * different dtypes where given on the individual nodes
142
            * neither a non-0-dimensional global_data nor global_shape nor
ultimanet's avatar
ultimanet committed
143
144
145
              hdf5 file supplied
            * global_shape == ()
            * different global_shapes where given on the individual nodes
146
            * neither non-0-dimensional local_data nor local_shape nor
ultimanet's avatar
ultimanet committed
147
              global d2o supplied
148
            * local_shape == ()
ultimanet's avatar
ultimanet committed
149
            * the first entry of local_shape is not the same on all nodes
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

    Notes
    -----
    The index is the d2o's global unique indentifier. One may use it in order
    to assemble the corresponding local d2o objects on different nodes if
    only one local object on a specific node is given.

    In order to speed up the init process the distributor_factory checks
    if the global_configuration object gc yields gc['d2o_init_checks'] == True.
    If yes, all checks expensive checks are skipped; namely those which  need
    mpi communication. Use this in order to get a fast init speed without
    loosing d2o's init parsing logic.

    Examples
    --------
    >>> a = np.arange(16, dtype=np.float).reshape((4,4))
    >>> obj = distributed_data_object(a, dtype=np.complex)
    >>> obj
    <distributed_data_object>
    array([[  0.+0.j,   1.+0.j,   2.+0.j,   3.+0.j],
           [  4.+0.j,   5.+0.j,   6.+0.j,   7.+0.j],
           [  8.+0.j,   9.+0.j,  10.+0.j,  11.+0.j],
           [ 12.+0.j,  13.+0.j,  14.+0.j,  15.+0.j]])
ultimanet's avatar
ultimanet committed
173

174
175
176
    See Also
    --------
    distributor 
ultimanet's avatar
ultimanet committed
177
    """
178

179
    def __init__(self, global_data=None, global_shape=None, dtype=None,
Ultima's avatar
Ultima committed
180
                 local_data=None, local_shape=None,
Ultima's avatar
Ultima committed
181
                 distribution_strategy=_default_strategy, hermitian=False,
182
183
184
185
186
                 alias=None, path=None, comm=MPI.COMM_WORLD,
                 copy=True, *args, **kwargs):

        # TODO: allow init with empty shape

187
188
189
190
        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)
191

192
        self.distributor = distributor_factory.get_distributor(
193
194
195
196
197
198
199
200
201
                                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,
202
                                **kwargs)
203

ultimanet's avatar
ultimanet committed
204
205
206
        self.distribution_strategy = distribution_strategy
        self.dtype = self.distributor.dtype
        self.shape = self.distributor.global_shape
Ultima's avatar
Ultima committed
207
208
        self.local_shape = self.distributor.local_shape
        self.comm = self.distributor.comm
209
210

        self.init_args = args
211
        self.init_kwargs = kwargs
212

Ultima's avatar
Ultima committed
213
        (self.data, self.hermitian) = self.distributor.initialize_data(
214
215
216
217
218
219
            global_data=global_data,
            local_data=local_data,
            alias=alias,
            path=path,
            hermitian=hermitian,
            copy=copy)
220
        self.index = d2o_librarian.register(self)
221

222
223
    @property
    def real(self):
224
225
226
227
228
229
230
231
232
        """ Returns a d2o containing the real part of the d2o's elements.

        Returns
        -------
        out : distributed_data_object
            The output object. The new datatype is the one numpy yields when
            taking the real part on the local data.
        """

233
234
235
236
237
238
239
240
241
        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):
242
243
244
245
246
247
248
249
250
        """ Returns a d2o containing the imaginary part of the d2o's elements.

        Returns
        -------
        out : distributed_data_object
            The output object. The new datatype is the one numpy yields when
            taking the imaginary part on the local data.
        """

251
252
253
254
255
256
257
        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

258
259
260
261
262
263
264
265
266
    @property 
    def hermitian(self):
        return self._hermitian

    @hermitian.setter
    def hermitian(self, value):
        self._hermitian = bool(value)


267
    def _fast_copy_empty(self):
268
269
270
271
272
273
        """ Make a very fast low level copy of the d2o without its data.

        This function is fast, because it uses EmptyD2o - a derived class from
        distributed_data_object and then copies the __dict__ directly. Unlike
        copy_empty, _fast_copy_empty will copy all attributes unchanged.
        """
274
275
276
277
278
279
280
281
282
283
        # make an empty d2o
        new_copy = EmptyD2o()
        # repair its class
        new_copy.__class__ = self.__class__
        # now copy everthing in the __dict__ except for the data array
        for key, value in self.__dict__.items():
            if key != 'data':
                new_copy.__dict__[key] = value
            else:
                new_copy.__dict__[key] = np.empty_like(value)
284
        # Register the new d2o at the librarian in order to get a unique index
285
        new_copy.index = d2o_librarian.register(new_copy)
286
287
        return new_copy

Ultimanet's avatar
Ultimanet committed
288
    def copy(self, dtype=None, distribution_strategy=None, **kwargs):
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
        """ Returns a full copy of the distributed data object.

        If no keyword arguments are given, the returned object will be an
        identical copy of the original d2o. By explicit specification one is
        able to define the dtype and the distribution_strategy of the returned
        d2o.

        Parameters
        ----------
        dtype : type
            The dtype that the new d2o will have. The data of the primary
            d2o will be casted.
        distribution_strategy : all supported distribution strategies
            The distribution strategy the new d2o should have. If not None and
            different from the original one, there will certainly be inter-node
            communication.
        **kwargs
            Additional keyword arguments get passed to the used copy_empty
            routine.

        Returns
        -------
        out : distributed_data_object
            The output object. It containes the old data, possibly casted to a
            new datatype and distributed according to a new distribution
            strategy

        See Also
        --------
        copy_empty

        """
321
322
323
        temp_d2o = self.copy_empty(dtype=dtype,
                                   distribution_strategy=distribution_strategy,
                                   **kwargs)
Ultima's avatar
Ultima committed
324
        if distribution_strategy is None or \
325
                distribution_strategy == self.distribution_strategy:
Ultimanet's avatar
Ultimanet committed
326
327
            temp_d2o.set_local_data(self.get_local_data(), copy=True)
        else:
Ultima's avatar
Ultima committed
328
            temp_d2o.inject((slice(None),), self, (slice(None),))
329
        temp_d2o.hermitian = self.hermitian
330
        return temp_d2o
331
332

    def copy_empty(self, global_shape=None, local_shape=None, dtype=None,
333
                   distribution_strategy=None, **kwargs):
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
        """ Returns an empty copy of the distributed data object.

        If no keyword arguments are given, the returned object will be an
        identical copy of the original d2o containing random data. By explicit
        specification one is able to define the new dtype and
        distribution_strategy of the returned d2o and to modify the new shape.

        Parameters
        ----------
        global_shape : tuple of ints
            The global shape that the new d2o shall have. Relevant for
            global-type distribution strategies like 'equal' or 'fftw'.
        local_shape : tuple of ints
            The local shape that the new d2o shall have. Relevant for
            local-type distribution strategies like 'freeform'.
        dtype : type
            The dtype that the new d2o will have.
        distribution_strategy : all supported distribution strategies
            The distribution strategy the new d2o should have.
        **kwargs
            Additional keyword arguments get passed to the init-call if the
            full initialization of a new distributed_data_object is necessary

        Returns
        -------
        out : distributed_data_object
            The output object. It contains random data.

        See Also
        --------
        copy

        """
Ultima's avatar
Ultima committed
367
        if self.distribution_strategy == 'not' and \
368
                distribution_strategy in STRATEGIES['local'] and \
369
370
371
372
373
                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
374
                                     **kwargs)
375
376
            return result.copy_empty(
                distribution_strategy=distribution_strategy)
377

Ultima's avatar
Ultima committed
378
        if global_shape is None:
379
            global_shape = self.shape
Ultima's avatar
Ultima committed
380
381
382
        if local_shape is None:
            local_shape = self.local_shape
        if dtype is None:
383
            dtype = self.dtype
384
385
        else:
            dtype = np.dtype(dtype)
Ultima's avatar
Ultima committed
386
        if distribution_strategy is None:
387
388
            distribution_strategy = self.distribution_strategy

389
390
391
392
393
394
395
396
        # check if all parameters remain the same -> use the _fast_copy_empty
        if (global_shape == self.shape and
                local_shape == self.local_shape and
                dtype == self.dtype and
                distribution_strategy == self.distribution_strategy and
                kwargs == self.init_kwargs):
            return self._fast_copy_empty()

397
        kwargs.update(self.init_kwargs)
398

399
400
401
402
403
404
405
406
        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)
407
        return temp_d2o
408

409
    def apply_scalar_function(self, function, inplace=False, dtype=None):
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
        """ Maps a scalar function on each entry of an array.

        The result of the function evaluation may be stored in the original
        array or in a new array (default). Furthermore the dtype of the
        returned array can be specified explicitly if inplace is set to False.

        Parameters
        ----------
        function : callable
            Will be applied to the array's entries. It will be the node's local
            data array into function as a whole. If this fails, the numpy
            vectorize function will be used.
        inplace : boolean
            Specifies if the result of the function evaluation should be stored
            in the original array or not.
        dtype : type
            If inplace is set to False, it is possible to specify the return
            d2o's dtype explicitly.

        Returns
        -------
        out : distributed_data_object
            Resulting d2o. This is either a newly created array or the primary
            d2o itself.
        """
435
        remember_hermitianQ = self.hermitian
436

437
        if inplace is True:
Ultimanet's avatar
Ultimanet committed
438
            temp = self
Ultima's avatar
Ultima committed
439
            if dtype is not None and self.dtype != np.dtype(dtype):
440
441
442
                about.warnings.cprint(
                    "WARNING: Inplace dtype conversion is not possible!")

Ultimanet's avatar
Ultimanet committed
443
        else:
444
            temp = self.copy_empty(dtype=dtype)
Ultimanet's avatar
Ultimanet committed
445

Ultima's avatar
Ultima committed
446
        if np.prod(self.local_shape) != 0:
447
            try:
Ultima's avatar
Ultima committed
448
449
                temp.data[:] = function(self.data)
            except:
450
451
                about.warnings.cprint(
                    "WARNING: Trying to use np.vectorize!")
Ultima's avatar
Ultima committed
452
453
                temp.data[:] = np.vectorize(function)(self.data)
        else:
454
455
            # Noting to do here. The value-empty array
            # is also geometrically empty
Ultima's avatar
Ultima committed
456
            pass
457

458
459
460
461
        if function in (np.exp, np.log):
            temp.hermitian = remember_hermitianQ
        else:
            temp.hermitian = False
Ultimanet's avatar
Ultimanet committed
462
        return temp
463

464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
    def apply_generator(self, generator, copy=False):
        """ Evaluates generator(local_shape) and stores the result locally.

        Parameters
        ----------
        generator : callable
            This function must be able to process the node's local data shape
            and return a numpy.ndarray of this very shape. This array is then 
            stored as the local data array on each node.
        copy : boolean
            Specifies whether the self.set_local_data method is instructed to 
            copy the result from generator or not.

        Notes
        -----
        The generator function yields node-local results. Therefore it is 
        assumed that the resulting overall d2o does not possess hermitian 
        symmetry anymore. Therefore self.hermitian is set to False.
                        
        """
        self.set_local_data(generator(self.distributor.local_shape), copy=copy)
Ultimanet's avatar
Ultimanet committed
485
        self.hermitian = False
486

ultimanet's avatar
ultimanet committed
487
    def __str__(self):
488
        """ x.__str__() <==> str(x)"""
ultimanet's avatar
ultimanet committed
489
        return self.data.__str__()
490

ultimanet's avatar
ultimanet committed
491
    def __repr__(self):
492
        """ x.__repr__() <==> repr(x)"""
493
        return '<distributed_data_object>\n' + self.data.__repr__()
494

495
    def _compare_helper(self, other, op):
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
        """ _compare_helper is used for <, <=, ==, !=, >= and >.
        
        It checks the class of `other` and then utilizes the appropriate 
        methods of self. If `other` is not a scalar, numpy.ndarray or 
        distributed_data_object this method will use numpy casting.
        
        Parameters
        ----------
        other : scalar, numpy.ndarray, distributed_data_object, array_like
            This is the object that will be compared to self.
        op : string
            The name of the comparison function, e.g. '__ne__'.
        
        Returns
        -------
        result : boolean, distributed_data_object
            If `other` was None, False will be returned. This follows the 
            behaviour of numpy but will changed as soon as numpy changed their 
            convention. In every other case a distributed_data_object with 
            element-wise comparison results will be returned.

        """
        
        if other is not None:
            result = self.copy_empty(dtype=np.bool_)

522
        # Case 1: 'other' is a scalar
523
        # -> make element-wise comparison
Ultimanet's avatar
Ultimanet committed
524
        if np.isscalar(other):
525
            result.set_local_data(
526
                getattr(self.get_local_data(copy=False), op)(other))
527
            return result
Ultimanet's avatar
Ultimanet committed
528

529
        # Case 2: 'other' is a numpy array or a distributed_data_object
530
        # -> extract the local data and make element-wise comparison
Ultimanet's avatar
Ultimanet committed
531
        elif isinstance(other, np.ndarray) or\
532
                isinstance(other, distributed_data_object):
Ultimanet's avatar
Ultimanet committed
533
            temp_data = self.distributor.extract_local_data(other)
534
535
            result.set_local_data(
                getattr(self.get_local_data(copy=False), op)(temp_data))
Ultimanet's avatar
Ultimanet committed
536
            return result
537
538

        # Case 3: 'other' is None
Ultima's avatar
Ultima committed
539
        elif other is None:
Ultimanet's avatar
Ultimanet committed
540
            return False
541
542
543

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

548
    def __ne__(self, other):
549
550
551
552
553
554
555
        """ x.__ne__(y) <==> x != y

        See Also
        --------
        _compare_helper

        """
556
        return self._compare_helper(other, '__ne__')
557

558
    def __lt__(self, other):
559
560
561
562
563
564
565
566
        """ x.__lt__(y) <==> x < y

        See Also
        --------
        _compare_helper

        """

567
        return self._compare_helper(other, '__lt__')
568

569
    def __le__(self, other):
570
571
572
573
574
575
576
577
        """ x.__le__(y) <==> x <= y

        See Also
        --------
        _compare_helper

        """

578
579
580
        return self._compare_helper(other, '__le__')

    def __eq__(self, other):
581
582
583
584
585
586
587
        """ x.__eq__(y) <==> x == y

        See Also
        --------
        _compare_helper

        """
588
589

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

591
    def __ge__(self, other):
592
593
594
595
596
597
598
599
        """ x.__ge__(y) <==> x >= y

        See Also
        --------
        _compare_helper

        """

600
601
602
        return self._compare_helper(other, '__ge__')

    def __gt__(self, other):
603
604
605
606
607
608
609
610
        """ x.__gt__(y) <==> x > y

        See Also
        --------
        _compare_helper

        """

611
612
        return self._compare_helper(other, '__gt__')

Ultima's avatar
Ultima committed
613
    def __iter__(self):
614
615
616
617
618
619
620
621
622
        """ x.__iter__() <==> iter(x)

        The __iter__ call returns an iterator it got from self.distributor.

        See Also
        --------
        distributor.get_iter

        """
Ultima's avatar
Ultima committed
623
624
        return self.distributor.get_iter(self)

Ultimanet's avatar
Ultimanet committed
625
    def equal(self, other):
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
        """  Checks if `other` and `self` are structurally the same. 
        
        In contrast to the element-wise comparison with `__eq__`, `equal` 
        checks more than only the equality of the array data. 
        It checks the equality of
            * shape
            * dtype
            * init_args
            * init_kwargs
            * distribution_strategy
            * node's local data

        Parameters
        ----------
        other : object
            The object that will be compared to `self`.

        Returns
        -------
        result : boolean
            True if above conditions are met, False otherwise. 

        """

Ultimanet's avatar
Ultimanet committed
650
651
652
653
654
655
656
657
658
        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
659
        except(AssertionError, AttributeError):
Ultimanet's avatar
Ultimanet committed
660
661
662
663
            return False
        else:
            return True

664
    def __pos__(self):
665
666
667
668
669
        """ x.__pos__() <==> +x 
        
        Returns a (positive) copy of `self`.
        """
        
670
        temp_d2o = self.copy_empty()
671
672
        temp_d2o.set_local_data(data=self.get_local_data().__pos__(), 
                                copy=False)
673
        return temp_d2o
674

ultimanet's avatar
ultimanet committed
675
    def __neg__(self):
676
677
678
679
680
        """ x.__neg__() <==> -x 
        
        Returns a negative copy of `self`.
        """
        
681
        temp_d2o = self.copy_empty()
682
        temp_d2o.set_local_data(data=self.get_local_data().__neg__(),
683
                                copy=False)
ultimanet's avatar
ultimanet committed
684
        return temp_d2o
685

686
    def __abs__(self):
687
688
689
690
691
        """ x.__abs__() <==> abs(x)

        Returns an absolute valued copy of `self`.
        """

692
        # translate complex dtypes
693
694
695
696
697
698
        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
699
700
        else:
            new_dtype = self.dtype
701
702
        temp_d2o = self.copy_empty(dtype=new_dtype)
        temp_d2o.set_local_data(data=self.get_local_data().__abs__(),
703
                                copy=False)
704
        return temp_d2o
705

Ultima's avatar
Ultima committed
706
    def _builtin_helper(self, operator, other, inplace=False):
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
        """ Used for various binary operations like +, -, *, /, **, *=, +=,...

        _builtin_helper checks whether `other` is a scalar or an array and 
        based on that extracts the locally relevant data from it. If `self`
        is hermitian, _builtin_helper tries to conserve this flag; but without 
        checking hermitianity explicitly. 

        Parameters
        ----------
        operator : callable

        other : scalar, array-like

        inplace : boolean
            If the result shall be saved in the data array of `self`. Used for 
            +=, -=, etc...
        Returns
        -------
        out : distributed_data_object
            The distributed_data_object containing the computation's result.
            Equals `self` if `inplace is True`.

        """
730
        # Case 1: other is not a scalar
Ultimanet's avatar
Ultimanet committed
731
        if not (np.isscalar(other) or np.shape(other) == (1,)):
732
            try:
733
                hermitian_Q = (other.hermitian and self.hermitian)
734
735
            except(AttributeError):
                hermitian_Q = False
736
            # extract the local data from the 'other' object
Ultimanet's avatar
Ultimanet committed
737
738
            temp_data = self.distributor.extract_local_data(other)
            temp_data = operator(temp_data)
739

740
        # Case 2: other is a scalar
741
        else:
742
            # if other is a scalar packed in a d2o, extract its value.
743
            if isinstance(other, distributed_data_object):
744
                other = other[0]
745

746
            if np.isrealobj(other): 
747
748
749
750
                hermitian_Q = self.hermitian
            else:
                hermitian_Q = False

751
            temp_data = operator(other)
752
753
        
        # select the return-distributed_data_object
754
        if inplace is True:
755
756
            temp_d2o = self
        else:
757
            # use common datatype for self and other
758
            new_dtype = np.dtype(np.find_common_type((self.dtype,),
759
                                                     (temp_data.dtype,)))
760
            temp_d2o = self.copy_empty(dtype=new_dtype)
761
762
        
        # write the new data into the return-distributed_data_object
763
        temp_d2o.set_local_data(data=temp_data, copy=False)
764
        temp_d2o.hermitian = hermitian_Q
ultimanet's avatar
ultimanet committed
765
        return temp_d2o
766

ultimanet's avatar
ultimanet committed
767
    def __add__(self, other):
768
769
770
771
772
773
774
        """ x.__add__(y) <==> x+y

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
775
        return self._builtin_helper(self.get_local_data().__add__, other)
ultimanet's avatar
ultimanet committed
776
777

    def __radd__(self, other):
778
779
780
781
782
783
784
        """ x.__radd__(y) <==> y+x

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
785
        return self._builtin_helper(self.get_local_data().__radd__, other)
Ultimanet's avatar
Ultimanet committed
786
787

    def __iadd__(self, other):
788
789
790
791
792
793
794
        """ x.__iadd__(y) <==> x+=y

        See Also
        --------
        _builtin_helper
        """

795
        return self._builtin_helper(self.get_local_data().__iadd__,
796
797
                                    other,
                                    inplace=True)
Ultimanet's avatar
Ultimanet committed
798

ultimanet's avatar
ultimanet committed
799
    def __sub__(self, other):
800
801
802
803
804
805
806
        """ x.__sub__(y) <==> x-y

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
807
        return self._builtin_helper(self.get_local_data().__sub__, other)
808

ultimanet's avatar
ultimanet committed
809
    def __rsub__(self, other):
810
811
812
813
814
815
816
        """ x.__rsub__(y) <==> y-x

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
817
        return self._builtin_helper(self.get_local_data().__rsub__, other)
818

ultimanet's avatar
ultimanet committed
819
    def __isub__(self, other):
820
821
822
823
824
825
826
        """ x.__isub__(y) <==> x-=y

        See Also
        --------
        _builtin_helper
        """

827
        return self._builtin_helper(self.get_local_data().__isub__,
828
829
                                    other,
                                    inplace=True)
830

ultimanet's avatar
ultimanet committed
831
    def __div__(self, other):
832
833
834
835
836
837
838
        """ x.__div__(y) <==> x/y

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
839
        return self._builtin_helper(self.get_local_data().__div__, other)
840

841
    def __truediv__(self, other):
842
843
844
845
846
847
848
        """ x.__truediv__(y) <==> x/y

        See Also
        --------
        _builtin_helper
        """

849
        return self.__div__(other)
850

ultimanet's avatar
ultimanet committed
851
    def __rdiv__(self, other):
852
853
854
855
856
857
858
        """ x.__rdiv__(y) <==> y/x

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
859
        return self._builtin_helper(self.get_local_data().__rdiv__, other)
860

861
    def __rtruediv__(self, other):
862
863
864
865
866
867
868
        """ x.__rtruediv__(y) <==> y/x

        See Also
        --------
        _builtin_helper
        """

869
        return self.__rdiv__(other)
ultimanet's avatar
ultimanet committed
870

Ultimanet's avatar
Ultimanet committed
871
    def __idiv__(self, other):
872
873
874
875
876
877
878
        """ x.__idiv__(y) <==> x/=y

        See Also
        --------
        _builtin_helper
        """

879
        return self._builtin_helper(self.get_local_data().__idiv__,
880
881
882
                                    other,
                                    inplace=True)

883
    def __itruediv__(self, other):
884
885
886
887
888
889
890
        """ x.__itruediv__(y) <==> x/=y

        See Also
        --------
        _builtin_helper
        """

891
        return self.__idiv__(other)
892

ultimanet's avatar
ultimanet committed
893
    def __floordiv__(self, other):
894
895
896
897
898
899
900
        """ x.__floordiv__(y) <==> x//y

        See Also
        --------
        _builtin_helper
        """

901
        return self._builtin_helper(self.get_local_data().__floordiv__,
902
903
                                    other)

ultimanet's avatar
ultimanet committed
904
    def __rfloordiv__(self, other):
905
906
907
908
909
910
911
        """ x.__rfloordiv__(y) <==> y//x

        See Also
        --------
        _builtin_helper
        """

912
        return self._builtin_helper(self.get_local_data().__rfloordiv__,
913
914
                                    other)

Ultimanet's avatar
Ultimanet committed
915
    def __ifloordiv__(self, other):
916
917
918
919
920
921
922
        """ x.__ifloordiv__(y) <==> x//=y

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
923
        return self._builtin_helper(
924
925
            self.get_local_data().__ifloordiv__, other,
            inplace=True)
926

ultimanet's avatar
ultimanet committed
927
    def __mul__(self, other):
928
929
930
931
932
933
934
        """ x.__mul__(y) <==> x*y

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
935
        return self._builtin_helper(self.get_local_data().__mul__, other)
936

ultimanet's avatar
ultimanet committed
937
    def __rmul__(self, other):
938
939
940
941
942
943
944
        """ x.__rmul__(y) <==> y*x

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
945
        return self._builtin_helper(self.get_local_data().__rmul__, other)
ultimanet's avatar
ultimanet committed
946
947

    def __imul__(self, other):
948
949
950
951
952
953
954
        """ x.__imul__(y) <==> x*=y

        See Also
        --------
        _builtin_helper
        """

955
        return self._builtin_helper(self.get_local_data().__imul__,
956
957
                                    other,
                                    inplace=True)
Ultimanet's avatar
Ultimanet committed
958

ultimanet's avatar
ultimanet committed
959
    def __pow__(self, other):
960
961
962
963
964
965
966
        """ x.__pow__(y) <==> x**y

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
967
        return self._builtin_helper(self.get_local_data().__pow__, other)
968

ultimanet's avatar
ultimanet committed
969
    def __rpow__(self, other):
970
971
972
973
974
975
976
        """ x.__rpow__(y) <==> y**x

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
977
        return self._builtin_helper(self.get_local_data().__rpow__, other)
ultimanet's avatar
ultimanet committed
978
979

    def __ipow__(self, other):
980
981
982
983
984
985
986
        """ x.__ipow__(y) <==> x**=y

        See Also
        --------
        _builtin_helper
        """

987
        return self._builtin_helper(self.get_local_data().__ipow__,
988
989
990
                                    other,
                                    inplace=True)

Ultima's avatar
Ultima committed
991
    def __mod__(self, other):
992
993
994
995
996
997
998
        """ x.__mod__(y) <==> x%y

        See Also
        --------
        _builtin_helper
        """

Ultima's avatar
Ultima committed
999
        return self._builtin_helper(self.get_local_data().__mod__, other)
1000

Ultima's avatar
Ultima committed
1001
    def __rmod__(self, other):
1002
1003
1004
1005
1006
1007
1008
        """ x.__rmod__(y) <==> y%x

        See Also
        --------
        _builtin_helper
        """

1009
        return self._builtin_helper(self.get_local_data().__rmod__, other)
1010

Ultima's avatar
Ultima committed
1011
    def __imod__(self, other):
1012
1013
1014
1015
1016
1017
1018
        """ x.__imod__(y) <==> x%=y

        See Also
        --------
        _builtin_helper
        """

1019
        return self._builtin_helper(self.get_local_data().__imod__,
1020
1021
1022
                                    other,
                                    inplace=True)

1023
    def __len__(self):
1024
1025
        """ Returns the length of the first axis."""

1026
        return self.shape[0]
1027

1028
    def get_dim(self):
1029
1030
1031
1032
1033
        """" Returns the total number of entries in the array.

        This is equivalent to the product of the shape.
        """

1034
        return np.prod(self.shape)
1035

1036
    def vdot(self, other):
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
        """ Returns the numpy.vdot analogous product of two arrays.

        If `self` is a complex array, the complex conjugate of it will be used.
        Internally the numpy.vdot function is used for the d2o's local data, 
        and the individual results get MPI-reduced. 

        See Also
        --------
        numpy.vdot
        """

1048
        other = self.distributor.extract_local_data(other)
1049
1050
1051
1052
1053
1054
1055
1056
1057
        local_vdot = np.array([np.vdot(self.get_local_data(), other)])
        global_vdot = np.empty_like(local_vdot)
        self.distributor._Allreduce_sum(sendbuf=local_vdot,
                                        recvbuf=global_vdot)

#        local_vdot = np.vdot(self.get_local_data(), other)
#        local_vdot_list = self.distributor._allgather(local_vdot)
#        global_vdot = np.result_type(self.dtype,
#                                     other.dtype).type(np.sum(local_vdot_list))
1058
        return global_vdot[0]
Ultimanet's avatar
Ultimanet committed
1059

ultimanet's avatar
ultimanet committed
1060
    def __getitem__(self, key):
1061
        """ x.__getitem__(y) <==> x[y] <==> x.get_data(y) """
Ultima's avatar
Ultima committed
1062
        return self.get_data(key)
1063

ultimanet's avatar
ultimanet committed
1064
    def __setitem__(self, key, data):
1065
        """ x.__setitem__(i, y) <==> x[i]=y <==> x.set_data(y, i) """
ultimanet's avatar
ultimanet committed
1066
        self.set_data(data, key)
1067

1068
    def _contraction_helper(self, function, **kwargs):
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
        """ Used for various operations like min, max, sum, prod, mean,...

        _builtin_helper checks whether the local node's data array is empty,
        then applies the given function on the local data, collects the 
        results, and then applies the function to this with the keyword axis=0.

        Parameters
        ----------
        function : callable
            This object will be applied to the local data array.

        **kwargs
            Additional keyword arguments will be passed to the first `function`
            call. 

        Returns
        -------
        out : 
            The return object of `function`.

        Raises
        ------
        ValueError 
            Raised if the d2o's shape equals (0,).
        """

1095
1096
1097
        if self.shape == (0,):
            raise ValueError("ERROR: Zero-size array to reduction operation " +
                             "which has no identity")
Ultima's avatar
Ultima committed
1098
1099
1100
1101
1102
1103
1104
        if np.prod(self.data.shape) == 0:
            local = 0
            include = False
        else:
            local = function(self.data, **kwargs)
            include = True

1105
        local_list = self.distributor._allgather(local)
1106
        local_list = np.array(local_list, dtype=np.dtype(local_list[0]))
Ultima's avatar
Ultima committed
1107
1108
1109
        include_list = np.array(self.distributor._allgather(include))
        work_list = local_list[include_list]
        if work_list.shape[0] == 0:
1110
            raise ValueError("ERROR: Zero-size array to reduction operation " +
Ultima's avatar
Ultima committed
1111
                             "which has no identity")
1112
        else:
Ultima's avatar
Ultima committed
1113
1114
            result = function(work_list, axis=0)
            return result
1115

1116
    def min(self, **kwargs):
1117
        """ x.min() <==> x.amin() """
1118
1119
        return self.amin(**kwargs)

1120
    def amin(self, **kwargs):
1121
1122
1123
1124
        """ Returns the minimum of an array. 
        
        See Also
        --------
1125
        numpy.amin
1126
1127
        """

1128
        return self._contraction_helper(np.amin, **kwargs)
1129
1130

    def nanmin(self, **kwargs):
1131
1132
1133
1134
        """ Returns the minimum of an array ignoring all NaNs.
        
        See Also
        --------
1135
        numpy.nanmin
1136
1137
        """

1138
        return self._contraction_helper(np.nanmin, **kwargs)
1139

1140
    def max(self, **kwargs):
1141
        """ x.max() <==> x.amax() """
1142
1143
        return self.amax(**kwargs)

1144
    def amax(self, **kwargs):
1145
1146
1147
1148
        """ Returns the maximum of an array. 
        
        See Also
        --------
1149
        numpy.amax
1150
1151
        """
 
1152
        return self._contraction_helper(np.amax, **kwargs)
1153

1154
    def nanmax(self, **kwargs):
1155
1156
1157
1158
        """ Returns the maximum of an array ignoring all NaNs.
        
        See Also
        --------
1159
        numpy.nanmax
1160
        """
1161

1162
        return self._contraction_helper(np.nanmax, **kwargs)
1163

1164
    def sum(self, **kwargs):
1165
1166
1167
1168
1169
1170
1171
        """ Sums the array elements.
        
        See Also
        --------
        numpy.sum 
        """

1172
1173
        if self.shape == (0,):
            return self.dtype.type(0)
1174
1175
1176
        return self._contraction_helper(np.sum, **kwargs)

    def prod(self, **kwargs):
1177
1178
1179
1180
1181
1182
1183
        """ Multiplies the array elements.

        See Also
        --------
        numpy.prod
        """

1184
1185
        if self.shape == (0,):
            return self.dtype.type(1)
1186
1187
        return self._contraction_helper(np.prod, **kwargs)

1188
    def mean(self, power=1):
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
        """ Returns the mean of the d2o's elements. 

        Parameters
        ----------
        power : scalar
            Used for point-wise exponentiation of the array's elements before
            computing the mean: mean(data**power)
        
        Returns
        -------
        mean : scalar
            The (pre-powered) mean-value of the array.
        """
1202
1203
        if self.shape == (0,):
            return np.mean(np.array([], dtype=self.dtype))
1204
        # compute the local means and the weights for the mean-mean.
Ultima's avatar
Ultima committed
1205
1206
1207
1208
        if np.prod(self.data.shape) == 0:
            local_mean = 0
            include = False
        else:
1209
1210
1211
1212
            if power != 1:
                local_mean = np.mean(self.data**power)
            else:
                local_mean = np.mean(self.data)
Ultima's avatar
Ultima committed
1213
            include = True
1214

1215
        local_weight = np.prod(self.data.shape)
1216
        # collect the local means and cast the result to a ndarray
Ultima's avatar
Ultima committed
1217
1218
        local_mean_list = self.distributor._allgather(local_mean)
        local_weight_list = self.distributor._allgather(local_weight)
1219

1220
1221
        local_mean_list = np.array(local_mean_list,
                                   dtype=np.dtype(local_mean_list[0]))
1222
1223
        local_weight_list = np.array(local_weight_list)
        # extract the parts from the non-empty nodes
Ultima's avatar
Ultima committed
1224
1225
1226
1227
1228
        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.")
1229
1230
        else:
            # compute the denominator for the weighted mean-mean
Ultima's avatar
Ultima committed
1231
            global_weight = np.sum(work_weight_list)