field.py 17.8 KB
Newer Older
csongor's avatar
csongor committed
1
2
3
4
from __future__ import division
import numpy as np
import pylab as pl

5
6
from d2o import distributed_data_object, \
    STRATEGIES as DISTRIBUTION_STRATEGIES
csongor's avatar
csongor committed
7

8
9
10
from nifty.config import about, \
    nifty_configuration as gc, \
    dependency_injector as gdi
csongor's avatar
csongor committed
11

12
13
from nifty.field_types import FieldType,\
                              FieldArray
14

15
from nifty.spaces.space import Space
csongor's avatar
csongor committed
16

csongor's avatar
csongor committed
17
18
19
import nifty.nifty_utilities as utilities

POINT_DISTRIBUTION_STRATEGIES = DISTRIBUTION_STRATEGIES['global']
theos's avatar
theos committed
20
COMM = getattr(gdi[gc['mpi_module']], gc['default_comm'])
csongor's avatar
csongor committed
21
22


23
class Field(object):
theos's avatar
theos committed
24
25
26
    # ---Initialization methods---
    def __init__(self, domain=None, val=None, dtype=None, field_type=None,
                 datamodel=None, copy=False):
27
        if isinstance(val, Field):
theos's avatar
theos committed
28
29
30
31
32
33
34
35
            if domain is None:
                domain = val.domain
            if dtype is None:
                dtype = val.dtype
            if field_type is None:
                field_type = val.field_type
            if datamodel is None:
                datamodel = val.datamodel
csongor's avatar
csongor committed
36

37
        self.domain = self._parse_domain(domain=domain)
38
        self.domain_axes = self._get_axes_tuple(self.domain)
csongor's avatar
csongor committed
39

40
        self.field_type = self._parse_field_type(field_type)
41

theos's avatar
theos committed
42
43
44
45
46
47
        try:
            start = len(reduce(lambda x, y: x+y, self.domain_axes))
        except TypeError:
            start = 0
        self.field_type_axes = self._get_axes_tuple(self.field_type,
                                                    start=start)
48

theos's avatar
theos committed
49
50
51
        self.dtype = self._infer_dtype(dtype=dtype,
                                       domain=self.domain,
                                       field_type=self.field_type)
52

theos's avatar
theos committed
53
54
        self.datamodel = self._parse_datamodel(datamodel=datamodel,
                                               val=val)
csongor's avatar
csongor committed
55
56
57

        self.set_val(new_val=val, copy=copy)

58
    def _parse_domain(self, domain):
59
60
61
        if domain is None:
            domain = ()
        elif not isinstance(domain, tuple):
62
            domain = (domain,)
csongor's avatar
csongor committed
63
        for d in domain:
64
            if not isinstance(d, Space):
csongor's avatar
csongor committed
65
                raise TypeError(about._errors.cstring(
66
67
                    "ERROR: Given domain contains something that is not a "
                    "nifty.space."))
csongor's avatar
csongor committed
68
69
        return domain

70
71
72
73
74
75
    def _parse_field_type(self, field_type):
        if field_type is None:
            field_type = ()
        elif not isinstance(field_type, tuple):
            field_type = (field_type,)
        for ft in field_type:
76
            if not isinstance(ft, FieldType):
77
                raise TypeError(about._errors.cstring(
78
                    "ERROR: Given object is not a nifty.FieldType."))
79
80
        return field_type

theos's avatar
theos committed
81
82
83
84
85
86
87
88
89
90
    def _get_axes_tuple(self, things_with_shape, start=0):
        i = start
        axes_list = []
        for thing in things_with_shape:
            l = []
            for j in range(len(thing.shape)):
                l += [i]
                i += 1
            axes_list += [tuple(l)]
        return tuple(axes_list)
91

theos's avatar
theos committed
92
    def _infer_dtype(self, dtype=None, domain=None, field_type=None):
csongor's avatar
csongor committed
93
        if dtype is None:
theos's avatar
theos committed
94
95
96
97
98
99
100
            dtype_tuple = (np.dtype(gc['default_field_dtype']),)
        else:
            dtype_tuple = (np.dtype(dtype),)
        if domain is not None:
            dtype_tuple += tuple(np.dtype(sp.dtype) for sp in domain)
        if field_type is not None:
            dtype_tuple += tuple(np.dtype(ft.dtype) for ft in field_type)
csongor's avatar
csongor committed
101

theos's avatar
theos committed
102
103
        dtype = reduce(lambda x, y: np.result_type(x, y), dtype_tuple)
        return dtype
104

theos's avatar
theos committed
105
106
107
108
109
    def _parse_datamodel(self, datamodel, val):
        if datamodel in DISTRIBUTION_STRATEGIES['all']:
            pass
        elif isinstance(val, distributed_data_object):
            datamodel = val.distribution_strategy
csongor's avatar
csongor committed
110
        else:
theos's avatar
theos committed
111
            datamodel = gc['default_datamodel']
csongor's avatar
csongor committed
112

theos's avatar
theos committed
113
        return datamodel
csongor's avatar
csongor committed
114

theos's avatar
theos committed
115
116
    # ---Properties---
    def set_val(self, new_val=None, copy=False):
117
118
        new_val = self.cast(new_val)
        if copy:
theos's avatar
theos committed
119
120
121
            new_val = new_val.copy()
        self._val = new_val
        return self._val
csongor's avatar
csongor committed
122

123
124
    def get_val(self, copy=False):
        if copy:
theos's avatar
theos committed
125
            return self._val.copy()
126
        else:
theos's avatar
theos committed
127
            return self._val
csongor's avatar
csongor committed
128

theos's avatar
theos committed
129
130
131
    @property
    def val(self):
        return self._val
csongor's avatar
csongor committed
132

theos's avatar
theos committed
133
134
135
    @val.setter
    def val(self, new_val):
        self._val = self.cast(new_val)
csongor's avatar
csongor committed
136

137
138
    @property
    def shape(self):
139
140
141
142
143
144
145
        shape_tuple = ()
        shape_tuple += tuple(sp.shape for sp in self.domain)
        shape_tuple += tuple(ft.shape for ft in self.field_type)
        try:
            global_shape = reduce(lambda x, y: x + y, shape_tuple)
        except TypeError:
            global_shape = ()
csongor's avatar
csongor committed
146

147
        return global_shape
csongor's avatar
csongor committed
148

149
150
    @property
    def dim(self):
theos's avatar
theos committed
151
152
153
154
155
156
157
        dim_tuple = ()
        dim_tuple += tuple(sp.dim for sp in self.domain)
        dim_tuple += tuple(ft.dim for ft in self.field_type)
        try:
            return reduce(lambda x, y: x * y, dim_tuple)
        except TypeError:
            return 0
csongor's avatar
csongor committed
158

159
160
    @property
    def dof(self):
theos's avatar
theos committed
161
162
163
164
165
166
167
168
        dof = self.dim
        if issubclass(self.dtype.type, np.complexfloating):
            dof *= 2
        return dof

    @property
    def total_volume(self):
        volume_tuple = tuple(sp.total_volume for sp in self.domain)
169
        try:
theos's avatar
theos committed
170
            return reduce(lambda x, y: x * y, volume_tuple)
171
        except TypeError:
theos's avatar
theos committed
172
            return 0
173

theos's avatar
theos committed
174
    # ---Special unary/binary operations---
csongor's avatar
csongor committed
175
176
177
    def cast(self, x=None, dtype=None):
        if dtype is None:
            dtype = self.dtype
178
179
        else:
            dtype = np.dtype(dtype)
180

theos's avatar
theos committed
181
        casted_x = self._actual_cast(x, dtype=dtype)
182
183

        for ind, sp in enumerate(self.domain):
184
            casted_x = sp.complement_cast(casted_x,
theos's avatar
theos committed
185
                                          axes=self.domain_axes[ind])
186
187
188

        for ind, ft in enumerate(self.field_type):
            casted_x = ft.complement_cast(casted_x,
theos's avatar
theos committed
189
                                          axes=self.field_type_axes[ind])
190
191

        return casted_x
csongor's avatar
csongor committed
192

theos's avatar
theos committed
193
    def _actual_cast(self, x, dtype=None):
194
        if isinstance(x, Field):
csongor's avatar
csongor committed
195
196
197
198
199
            x = x.get_val()

        if dtype is None:
            dtype = self.dtype

theos's avatar
theos committed
200
201
202
        x = distributed_data_object(x,
                                    global_shape=self.shape,
                                    dtype=dtype,
csongor's avatar
csongor committed
203
204
                                    distribution_strategy=self.datamodel)

theos's avatar
theos committed
205
206
207
208
209
210
211
212
213
214
215
        return x

    def copy(self, domain=None, dtype=None, field_type=None,
             datamodel=None):
        copied_val = self.get_val(copy=True)
        new_field = self.copy_empty(domain=domain,
                                    dtype=dtype,
                                    field_type=field_type,
                                    datamodel=datamodel)
        new_field.set_val(new_val=copied_val, copy=False)
        return new_field
csongor's avatar
csongor committed
216

theos's avatar
theos committed
217
218
219
220
    def copy_empty(self, domain=None, dtype=None, field_type=None,
                   datamodel=None):
        if domain is None:
            domain = self.domain
csongor's avatar
csongor committed
221
        else:
theos's avatar
theos committed
222
            domain = self._parse_domain(domain)
csongor's avatar
csongor committed
223

theos's avatar
theos committed
224
225
226
227
        if dtype is None:
            dtype = self.dtype
        else:
            dtype = np.dtype(dtype)
csongor's avatar
csongor committed
228

theos's avatar
theos committed
229
230
231
232
        if field_type is None:
            field_type = self.field_type
        else:
            field_type = self._parse_field_type(field_type)
csongor's avatar
csongor committed
233

theos's avatar
theos committed
234
235
        if datamodel is None:
            datamodel = self.datamodel
csongor's avatar
csongor committed
236

theos's avatar
theos committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
        fast_copyable = True
        try:
            for i in xrange(len(self.domain)):
                if self.domain[i] is not domain[i]:
                    fast_copyable = False
                    break
            for i in xrange(len(self.field_type)):
                if self.field_type[i] is not field_type[i]:
                    fast_copyable = False
                    break
        except IndexError:
            fast_copyable = False

        if (fast_copyable and dtype == self.dtype and
                datamodel == self.datamodel):
            new_field = self._fast_copy_empty()
        else:
            new_field = Field(domain=domain,
                              dtype=dtype,
                              field_type=field_type,
                              datamodel=datamodel)
        return new_field
csongor's avatar
csongor committed
259

theos's avatar
theos committed
260
261
262
263
264
265
266
267
268
269
270
271
272
273
    def _fast_copy_empty(self):
        # make an empty field
        new_field = EmptyField()
        # repair its class
        new_field.__class__ = self.__class__
        # copy domain, codomain and val
        for key, value in self.__dict__.items():
            if key != 'val':
                new_field.__dict__[key] = value
            else:
                new_field.__dict__[key] = self.val.copy_empty()
        return new_field

    def weight(self, power=1, inplace=False, spaces=None):
274
        if inplace:
csongor's avatar
csongor committed
275
276
277
278
            new_field = self
        else:
            new_field = self.copy_empty()

279
        new_val = self.get_val(copy=False)
csongor's avatar
csongor committed
280

csongor's avatar
csongor committed
281
        if spaces is None:
theos's avatar
theos committed
282
283
284
            spaces = range(len(self.domain))
        else:
            spaces = utilities.cast_axis_to_tuple(spaces, len(self.domain))
csongor's avatar
csongor committed
285

286
        for ind, sp in enumerate(self.domain):
theos's avatar
theos committed
287
288
289
290
291
            if ind in spaces:
                new_val = sp.weight(new_val,
                                    power=power,
                                    axes=self.domain_axes[ind],
                                    inplace=inplace)
292
293

        new_field.set_val(new_val=new_val, copy=False)
csongor's avatar
csongor committed
294
295
        return new_field

theos's avatar
theos committed
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
321
322
323
324
    def dot(self, x=None, bare=False):
        if isinstance(x, Field):
            try:
                assert len(x.domain) == len(self.domain)
                for index in xrange(len(self.domain)):
                    assert x.domain[index] == self.domain[index]
                for index in xrange(len(self.field_type)):
                    assert x.field_type[index] == self.field_type[index]
            except AssertionError:
                raise ValueError(about._errors.cstring(
                    "ERROR: domains are incompatible."))
            # extract the data from x and try to dot with this
            x = x.get_val(copy=False)

        # Compute the dot respecting the fact of discrete/continous spaces
        if bare:
            y = self
        else:
            y = self.weight(power=1)

        y = y.get_val(copy=False)

        # Cast the input in order to cure dtype and shape differences
        x = self.cast(x)

        dotted = x.conjugate() * y

        return dotted.sum()

325
    def norm(self, q=2):
csongor's avatar
csongor committed
326
327
328
329
330
331
332
333
334
335
336
337
338
339
        """
            Computes the Lq-norm of the field values.

            Parameters
            ----------
            q : scalar
                Parameter q of the Lq-norm (default: 2).

            Returns
            -------
            norm : scalar
                The Lq-norm of the field values.

        """
340
        if q == 2:
341
            return (self.dot(x=self)) ** (1 / 2)
csongor's avatar
csongor committed
342
        else:
343
            return self.dot(x=self ** (q - 1)) ** (1 / q)
csongor's avatar
csongor committed
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359

    def conjugate(self, inplace=False):
        """
            Computes the complex conjugate of the field.

            Returns
            -------
            cc : field
                The complex conjugated field.

        """
        if inplace:
            work_field = self
        else:
            work_field = self.copy_empty()

360
        new_val = self.get_val(copy=False)
theos's avatar
theos committed
361
        new_val = new_val.conjugate()
362
        work_field.set_val(new_val=new_val, copy=False)
csongor's avatar
csongor committed
363
364
365

        return work_field

theos's avatar
theos committed
366
367
368
    # ---General unary/contraction methods---
    def __pos__(self):
        return self.copy()
369

theos's avatar
theos committed
370
371
372
373
    def __neg__(self):
        return_field = self.copy_empty()
        new_val = -self.get_val(copy=False)
        return_field.set_val(new_val, copy=False)
csongor's avatar
csongor committed
374
375
        return return_field

theos's avatar
theos committed
376
377
378
379
380
    def __abs__(self):
        return_field = self.copy_empty()
        new_val = abs(self.get_val(copy=False))
        return_field.set_val(new_val, copy=False)
        return return_field
csongor's avatar
csongor committed
381

theos's avatar
theos committed
382
383
384
385
386
387
    def _contraction_helper(self, op, spaces, types):
        # build a list of all axes
        if spaces is None:
            spaces = xrange(len(self.domain))
        else:
            spaces = utilities.cast_axis_to_tuple(spaces, len(self.domain))
csongor's avatar
csongor committed
388

theos's avatar
theos committed
389
390
391
392
        if types is None:
            types = xrange(len(self.field_type))
        else:
            types = utilities.cast_axis_to_tuple(types, len(self.field_type))
393

theos's avatar
theos committed
394
395
396
397
        axes_list = ()
        axes_list += tuple(self.domain_axes[sp_index] for sp_index in spaces)
        axes_list += tuple(self.field_type_axes[ft_index] for
                           ft_index in types)
398
        try:
theos's avatar
theos committed
399
            axes_list = reduce(lambda x, y: x+y, axes_list)
400
        except TypeError:
theos's avatar
theos committed
401
            axes_list = ()
csongor's avatar
csongor committed
402

theos's avatar
theos committed
403
404
405
        # perform the contraction on the d2o
        data = self.get_val(copy=False)
        data = getattr(data, op)(axis=axes_list)
csongor's avatar
csongor committed
406

theos's avatar
theos committed
407
408
409
        # check if the result is scalar or if a result_field must be constr.
        if np.isscalar(data):
            return data
csongor's avatar
csongor committed
410
        else:
theos's avatar
theos committed
411
412
413
414
415
416
417
418
419
420
421
            return_domain = tuple(self.domain[i]
                                  for i in xrange(len(self.domain))
                                  if i not in spaces)
            return_field_type = tuple(self.field_type[i]
                                      for i in xrange(len(self.field_type))
                                      if i not in types)
            return_field = Field(domain=return_domain,
                                 val=data,
                                 field_type=return_field_type,
                                 copy=False)
            return return_field
csongor's avatar
csongor committed
422

theos's avatar
theos committed
423
424
    def sum(self, spaces=None, types=None):
        return self._contraction_helper('sum', spaces, types)
csongor's avatar
csongor committed
425

theos's avatar
theos committed
426
427
    def prod(self, spaces=None, types=None):
        return self._contraction_helper('prod', spaces, types)
csongor's avatar
csongor committed
428

theos's avatar
theos committed
429
430
    def all(self, spaces=None, types=None):
        return self._contraction_helper('all', spaces, types)
csongor's avatar
csongor committed
431

theos's avatar
theos committed
432
433
    def any(self, spaces=None, types=None):
        return self._contraction_helper('any', spaces, types)
csongor's avatar
csongor committed
434

theos's avatar
theos committed
435
436
    def min(self, spaces=None, types=None):
        return self._contraction_helper('min', spaces, types)
csongor's avatar
csongor committed
437

theos's avatar
theos committed
438
439
    def nanmin(self, spaces=None, types=None):
        return self._contraction_helper('nanmin', spaces, types)
csongor's avatar
csongor committed
440

theos's avatar
theos committed
441
442
    def max(self, spaces=None, types=None):
        return self._contraction_helper('max', spaces, types)
csongor's avatar
csongor committed
443

theos's avatar
theos committed
444
445
    def nanmax(self, spaces=None, types=None):
        return self._contraction_helper('nanmax', spaces, types)
csongor's avatar
csongor committed
446

theos's avatar
theos committed
447
448
    def mean(self, spaces=None, types=None):
        return self._contraction_helper('mean', spaces, types)
csongor's avatar
csongor committed
449

theos's avatar
theos committed
450
451
    def var(self, spaces=None, types=None):
        return self._contraction_helper('var', spaces, types)
csongor's avatar
csongor committed
452

theos's avatar
theos committed
453
454
    def std(self, spaces=None, types=None):
        return self._contraction_helper('std', spaces, types)
csongor's avatar
csongor committed
455

theos's avatar
theos committed
456
    # ---General binary methods---
csongor's avatar
csongor committed
457

theos's avatar
theos committed
458
    def _binary_helper(self, other, op, inplace=False):
csongor's avatar
csongor committed
459
        # if other is a field, make sure that the domains match
460
        if isinstance(other, Field):
theos's avatar
theos committed
461
462
463
464
465
466
467
468
469
470
            try:
                assert len(other.domain) == len(self.domain)
                for index in xrange(len(self.domain)):
                    assert other.domain[index] == self.domain[index]
                for index in xrange(len(self.field_type)):
                    assert other.field_type[index] == self.field_type[index]
            except AssertionError:
                raise ValueError(about._errors.cstring(
                    "ERROR: domains are incompatible."))
            other = other.get_val(copy=False)
csongor's avatar
csongor committed
471

theos's avatar
theos committed
472
473
        self_val = self.get_val(copy=False)
        return_val = getattr(self_val, op)(other)
csongor's avatar
csongor committed
474
475
476
477
478
479

        if inplace:
            working_field = self
        else:
            working_field = self.copy_empty()

theos's avatar
theos committed
480
        working_field.set_val(return_val, copy=False)
csongor's avatar
csongor committed
481
482
483
        return working_field

    def __add__(self, other):
theos's avatar
theos committed
484
        return self._binary_helper(other, op='__add__')
485

486
    def __radd__(self, other):
theos's avatar
theos committed
487
        return self._binary_helper(other, op='__radd__')
csongor's avatar
csongor committed
488
489

    def __iadd__(self, other):
theos's avatar
theos committed
490
        return self._binary_helper(other, op='__iadd__', inplace=True)
csongor's avatar
csongor committed
491
492

    def __sub__(self, other):
theos's avatar
theos committed
493
        return self._binary_helper(other, op='__sub__')
csongor's avatar
csongor committed
494
495

    def __rsub__(self, other):
theos's avatar
theos committed
496
        return self._binary_helper(other, op='__rsub__')
csongor's avatar
csongor committed
497
498

    def __isub__(self, other):
theos's avatar
theos committed
499
        return self._binary_helper(other, op='__isub__', inplace=True)
csongor's avatar
csongor committed
500
501

    def __mul__(self, other):
theos's avatar
theos committed
502
        return self._binary_helper(other, op='__mul__')
503

504
    def __rmul__(self, other):
theos's avatar
theos committed
505
        return self._binary_helper(other, op='__rmul__')
csongor's avatar
csongor committed
506
507

    def __imul__(self, other):
theos's avatar
theos committed
508
        return self._binary_helper(other, op='__imul__', inplace=True)
csongor's avatar
csongor committed
509
510

    def __div__(self, other):
theos's avatar
theos committed
511
        return self._binary_helper(other, op='__div__')
csongor's avatar
csongor committed
512
513

    def __rdiv__(self, other):
theos's avatar
theos committed
514
        return self._binary_helper(other, op='__rdiv__')
csongor's avatar
csongor committed
515
516

    def __idiv__(self, other):
theos's avatar
theos committed
517
        return self._binary_helper(other, op='__idiv__', inplace=True)
518

csongor's avatar
csongor committed
519
    def __pow__(self, other):
theos's avatar
theos committed
520
        return self._binary_helper(other, op='__pow__')
csongor's avatar
csongor committed
521
522

    def __rpow__(self, other):
theos's avatar
theos committed
523
        return self._binary_helper(other, op='__rpow__')
csongor's avatar
csongor committed
524
525

    def __ipow__(self, other):
theos's avatar
theos committed
526
        return self._binary_helper(other, op='__ipow__', inplace=True)
csongor's avatar
csongor committed
527
528

    def __lt__(self, other):
theos's avatar
theos committed
529
        return self._binary_helper(other, op='__lt__')
csongor's avatar
csongor committed
530
531

    def __le__(self, other):
theos's avatar
theos committed
532
        return self._binary_helper(other, op='__le__')
csongor's avatar
csongor committed
533
534
535
536
537

    def __ne__(self, other):
        if other is None:
            return True
        else:
theos's avatar
theos committed
538
            return self._binary_helper(other, op='__ne__')
csongor's avatar
csongor committed
539
540
541
542
543

    def __eq__(self, other):
        if other is None:
            return False
        else:
theos's avatar
theos committed
544
            return self._binary_helper(other, op='__eq__')
csongor's avatar
csongor committed
545
546

    def __ge__(self, other):
theos's avatar
theos committed
547
        return self._binary_helper(other, op='__ge__')
csongor's avatar
csongor committed
548
549

    def __gt__(self, other):
theos's avatar
theos committed
550
551
552
553
554
555
556
557
558
559
560
561
562
        return self._binary_helper(other, op='__gt__')

    def __repr__(self):
        return "<nifty_core.field>"

    def __str__(self):
        minmax = [self.min(), self.max()]
        mean = self.mean()
        return "nifty_core.field instance\n- domain      = " + \
               repr(self.domain) + \
               "\n- val         = " + repr(self.get_val()) + \
               "\n  - min.,max. = " + str(minmax) + \
               "\n  - mean = " + str(mean)
csongor's avatar
csongor committed
563

564

565
class EmptyField(Field):
csongor's avatar
csongor committed
566
567
    def __init__(self):
        pass