selection_operator.py 1.03 KB
Newer Older
1
2
3
4
5
6
from ..operators import LinearOperator
from ..sugar import full


class SelectionOperator(LinearOperator):
    def __init__(self, domain, key):
Philipp Arras's avatar
Fixups    
Philipp Arras committed
7
        from ..multi import MultiDomain
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
        if not isinstance(domain, MultiDomain):
            raise TypeError("Domain must be a MultiDomain")
        self._target = domain[key]
        self._domain = domain
        self._key = key

    @property
    def domain(self):
        return self._domain

    @property
    def target(self):
        return self._target

    @property
    def capability(self):
        return self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
        if mode == self.TIMES:
            return x[self._key].copy()
        else:
            result = {}
            for key, val in self.domain.items():
                if key != self._key:
                    result[key] = full(val, 0.)
                else:
                    result[key] = x.copy()
Philipp Arras's avatar
Fixups    
Philipp Arras committed
37
            from ..multi import MultiField
38
            return MultiField(result)