multi_field.py 7.82 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# 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/>.
#
# Copyright(C) 2013-2018 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.

Martin Reinecke's avatar
Martin Reinecke committed
19
20
21
from ..field import Field
import numpy as np
from .multi_domain import MultiDomain
22
from ..utilities import frozendict
Martin Reinecke's avatar
Martin Reinecke committed
23
24
25
26
27
28
29
30
31


class MultiField(object):
    def __init__(self, val):
        """
        Parameters
        ----------
        val : dict
        """
32
        self._val = frozendict(val)
33
34
        self._domain = MultiDomain.make(
            {key: val.domain for key, val in self._val.items()})
Martin Reinecke's avatar
Martin Reinecke committed
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

    def __getitem__(self, key):
        return self._val[key]

    def keys(self):
        return self._val.keys()

    def items(self):
        return self._val.items()

    def values(self):
        return self._val.values()

    @property
    def domain(self):
50
        return self._domain
Martin Reinecke's avatar
Martin Reinecke committed
51

Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
52
53
54
55
    @property
    def dtype(self):
        return {key: val.dtype for key, val in self._val.items()}

56
57
58
59
60
61
62
63
64
65
    @property
    def real(self):
        """MultiField : The real part of the multi field"""
        return MultiField({key: field.real for key, field in self.items()})

    @property
    def imag(self):
        """MultiField : The imaginary part of the multi field"""
        return MultiField({key: field.imag for key, field in self.items()})

66
67
    @staticmethod
    def from_random(random_type, domain, dtype=np.float64, **kwargs):
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
68
        dtype = MultiField.build_dtype(dtype, domain)
69
70
        return MultiField({key: Field.from_random(random_type, domain[key],
                                                  dtype[key], **kwargs)
71
                           for key in sorted(domain.keys())})
72

Martin Reinecke's avatar
Martin Reinecke committed
73
    def _check_domain(self, other):
74
        if other._domain != self._domain:
Martin Reinecke's avatar
Martin Reinecke committed
75
76
77
78
79
80
81
82
83
            raise ValueError("domains are incompatible.")

    def vdot(self, x):
        result = 0.
        self._check_domain(x)
        for key, sub_field in self.items():
            result += sub_field.vdot(x[key])
        return result

Martin Reinecke's avatar
Martin Reinecke committed
84
85
86
87
88
89
90
91
    @staticmethod
    def build_dtype(dtype, domain):
        if isinstance(dtype, dict):
            return dtype
        if dtype is None:
            dtype = np.float64
        return {key: dtype for key in domain.keys()}

Martin Reinecke's avatar
Martin Reinecke committed
92
    @staticmethod
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
93
94
    def full(domain, val):
        return MultiField({key: Field.full(dom, val)
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
95
                           for key, dom in domain.items()})
Martin Reinecke's avatar
Martin Reinecke committed
96

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
97
98
99
100
101
102
103
104
105
    def to_global_data(self):
        return {key: val.to_global_data() for key, val in self._val.items()}

    @staticmethod
    def from_global_data(domain, arr, sum_up=False):
        return MultiField({key: Field.from_global_data(domain[key],
                                                       val, sum_up)
                           for key, val in arr.items()})

Martin Reinecke's avatar
Martin Reinecke committed
106
107
108
109
110
111
112
113
114
115
    def norm(self):
        """ Computes the L2-norm of the field values.

        Returns
        -------
        norm : float
            The L2-norm of the field values.
        """
        return np.sqrt(np.abs(self.vdot(x=self)))

116
117
118
119
120
121
122
123
124
125
    def squared_norm(self):
        """ Computes the square of the L2-norm of the field values.

        Returns
        -------
        float
            The square of the L2-norm of the field values.
        """
        return abs(self.vdot(x=self))

Martin Reinecke's avatar
Martin Reinecke committed
126
127
128
    def __neg__(self):
        return MultiField({key: -val for key, val in self.items()})

129
130
131
    def __abs__(self):
        return MultiField({key: abs(val) for key, val in self.items()})

Martin Reinecke's avatar
Martin Reinecke committed
132
    def conjugate(self):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
133
134
        return MultiField({key: sub_field.conjugate()
                           for key, sub_field in self.items()})
Martin Reinecke's avatar
Martin Reinecke committed
135

136
137
138
139
140
141
142
143
144
145
146
147
    def all(self):
        for v in self.values():
            if not v.all():
                return False
        return True

    def any(self):
        for v in self.values():
            if v.any():
                return True
        return False

148
149
150
    def isEquivalentTo(self, other):
        """Determines (as quickly as possible) whether `self`'s content is
        identical to `other`'s content."""
151
152
153
154
155
156
157
        if self is other:
            return True
        if not isinstance(other, MultiField):
            return False
        if self._domain != other._domain:
            return False
        for key, val in self._val.items():
158
            if not val.isEquivalentTo(other[key]):
159
160
                return False
        return True
Martin Reinecke's avatar
Martin Reinecke committed
161

162
163
164
165
166
167
168
    def isSubsetOf(self, other):
        """Determines (as quickly as possible) whether `self`'s content is
        a subset of `other`'s content."""
        if self is other:
            return True
        if not isinstance(other, MultiField):
            return False
Martin Reinecke's avatar
Martin Reinecke committed
169
170
171
        if len(set(self._domain.keys()) - set(other._domain.keys())) > 0:
            return False
        for key in self._domain.keys():
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
172
            if other._domain[key] is not self._domain[key]:
173
                return False
Martin Reinecke's avatar
Martin Reinecke committed
174
            if not other[key].isSubsetOf(self[key]):
175
176
                return False
        return True
177

178

Martin Reinecke's avatar
Martin Reinecke committed
179
180
181
182
183
184
185
for op in ["__add__", "__radd__",
           "__sub__", "__rsub__",
           "__mul__", "__rmul__",
           "__div__", "__rdiv__",
           "__truediv__", "__rtruediv__",
           "__floordiv__", "__rfloordiv__",
           "__pow__", "__rpow__",
Martin Reinecke's avatar
Martin Reinecke committed
186
187
188
           "__lt__", "__le__", "__gt__", "__ge__", "__eq__", "__ne__"]:
    def func(op):
        def func2(self, other):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
189
            if isinstance(other, MultiField):
Martin Reinecke's avatar
Martin Reinecke committed
190
191
192
193
194
195
                if self._domain == other._domain:
                    result_val = {key: getattr(sub_field, op)(other[key])
                                  for key, sub_field in self.items()}
                else:
                    if not self._domain.compatibleTo(other.domain):
                        raise ValueError("domain mismatch")
Martin Reinecke's avatar
Martin Reinecke committed
196
197
198
199
200
                    s1 = set(self._domain.keys())
                    s2 = set(other._domain.keys())
                    common_keys = s1 & s2
                    only_self_keys = s1 - s2
                    only_other_keys = s2 - s1
Martin Reinecke's avatar
Martin Reinecke committed
201
                    result_val = {}
Martin Reinecke's avatar
Martin Reinecke committed
202
203
204
205
                    for key in common_keys:
                        result_val[key] = getattr(self[key], op)(other[key])
                    if op in ("__add__", "__radd__"):
                        for key in only_self_keys:
206
                            result_val[key] = self[key]
Martin Reinecke's avatar
Martin Reinecke committed
207
                        for key in only_other_keys:
208
                            result_val[key] = other[key]
Martin Reinecke's avatar
Martin Reinecke committed
209
210
                    elif op in ("__mul__", "__rmul__"):
                        pass
Philipp Arras's avatar
Philipp Arras committed
211
                    else:
Martin Reinecke's avatar
Martin Reinecke committed
212
                        for key in only_self_keys:
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
213
214
                            result_val[key] = getattr(
                                self[key], op)(self[key]*0.)
Martin Reinecke's avatar
Martin Reinecke committed
215
                        for key in only_other_keys:
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
216
217
                            result_val[key] = getattr(
                                other[key]*0., op)(other[key])
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
218
219
220
221
            else:
                result_val = {key: getattr(val, op)(other)
                              for key, val in self.items()}
            return MultiField(result_val)
Martin Reinecke's avatar
Martin Reinecke committed
222
223
        return func2
    setattr(MultiField, op, func(op))
224
225
226
227
228
229
230
231
232

for op in ["__iadd__", "__isub__", "__imul__", "__idiv__",
           "__itruediv__", "__ifloordiv__", "__ipow__"]:
    def func(op):
        def func2(self, other):
            raise TypeError(
                "In-place operations are deliberately not supported")
        return func2
    setattr(MultiField, op, func(op))