multi_field.py 3.78 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from ..field import Field
import numpy as np
from .multi_domain import MultiDomain


class MultiField(object):
    def __init__(self, val):
        """
        Parameters
        ----------
        val : dict
        """
        self._val = val

    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):
        return MultiDomain({key: val.domain for key, val in self._val.items()})

Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
31
32
33
34
    @property
    def dtype(self):
        return {key: val.dtype for key, val in self._val.items()}

35
36
37
38
39
40
41
    @staticmethod
    def from_random(random_type, domain, dtype=np.float64, **kwargs):
        dtype = self.build_dtype(dtype)
        return MultiField({key: Field.from_random(random_type, domain[key],
                                                  dtype[key], **kwargs)
                           for key in domain.keys})

Martin Reinecke's avatar
Martin Reinecke committed
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
    def _check_domain(self, other):
        if other.domain != self.domain:
            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

    def lock(self):
        for v in self.values():
            v.lock()
        return self

    def copy(self):
        return MultiField({key: val.copy() for key, val in self.items()})

Martin Reinecke's avatar
Martin Reinecke committed
61
62
63
64
65
66
67
68
    @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
69
70
    @staticmethod
    def zeros(domain, dtype=None):
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
71
        dtype = MultiField.build_dtype(dtype, domain)
Martin Reinecke's avatar
Martin Reinecke committed
72
        return MultiField({key: Field.zeros(dom, dtype=dtype[key])
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
73
                           for key, dom in domain.items()})
Martin Reinecke's avatar
Martin Reinecke committed
74
75
76

    @staticmethod
    def ones(domain, dtype=None):
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
77
        dtype = MultiField.build_dtype(dtype, domain)
Martin Reinecke's avatar
Martin Reinecke committed
78
        return MultiField({key: Field.ones(dom, dtype=dtype[key])
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
79
                           for key, dom in domain.items()})
Martin Reinecke's avatar
Martin Reinecke committed
80
81
82

    @staticmethod
    def empty(domain, dtype=None):
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
83
        dtype = MultiField.build_dtype(dtype, domain)
Martin Reinecke's avatar
Martin Reinecke committed
84
        return MultiField({key: Field.empty(dom, dtype=dtype[key])
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
85
                           for key, dom in domain.items()})
Martin Reinecke's avatar
Martin Reinecke committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

    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)))

    def __neg__(self):
        return MultiField({key: -val for key, val in self.items()})

    def conjugate(self):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
101
102
        return MultiField({key: sub_field.conjugate()
                           for key, sub_field in self.items()})
Martin Reinecke's avatar
Martin Reinecke committed
103
104
105
106
107
108
109
110
111
112
113
114


for op in ["__add__", "__radd__", "__iadd__",
           "__sub__", "__rsub__", "__isub__",
           "__mul__", "__rmul__", "__imul__",
           "__div__", "__rdiv__", "__idiv__",
           "__truediv__", "__rtruediv__", "__itruediv__",
           "__floordiv__", "__rfloordiv__", "__ifloordiv__",
           "__pow__", "__rpow__", "__ipow__",
           "__lt__", "__le__", "__gt__", "__ge__", "__eq__", "__ne__"]:
    def func(op):
        def func2(self, other):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
115
116
117
118
119
120
121
122
            if isinstance(other, MultiField):
                self._check_domain(other)
                result_val = {key: getattr(sub_field, op)(other[key])
                              for key, sub_field in self.items()}
            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
123
124
        return func2
    setattr(MultiField, op, func(op))