contraction_operator.py 2.83 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
# 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/>.
#
14
# Copyright(C) 2013-2019 Max-Planck-Society
Martin Reinecke's avatar
Martin Reinecke committed
15
#
16
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
17

Philipp Arras's avatar
Fixups    
Philipp Arras committed
18
import numpy as np
19

Philipp Arras's avatar
Philipp Arras committed
20
from .. import utilities
Philipp Arras's avatar
Fixups    
Philipp Arras committed
21
from ..domain_tuple import DomainTuple
22
from ..field import Field
Philipp Arras's avatar
Fixups    
Philipp Arras committed
23
24
25
from .linear_operator import LinearOperator


26
class ContractionOperator(LinearOperator):
Martin Reinecke's avatar
fix    
Martin Reinecke committed
27
    """A :class:`LinearOperator` which sums up fields into the direction of
Martin Reinecke's avatar
Martin Reinecke committed
28
    subspaces.
Philipp Arras's avatar
Philipp Arras committed
29

30
    This Operator sums up a field with is defined on a :class:`DomainTuple`
Martin Reinecke's avatar
fix    
Martin Reinecke committed
31
    to a :class:`DomainTuple` which is a subset of the former.
Philipp Arras's avatar
Philipp Arras committed
32
33
34

    Parameters
    ----------
35
    domain : Domain, tuple of Domain or DomainTuple
Martin Reinecke's avatar
fix    
Martin Reinecke committed
36
    spaces : None, int or tuple of int
37
        The elements of "domain" which are contracted.
Martin Reinecke's avatar
fix    
Martin Reinecke committed
38
        If `None`, everything is contracted
39
    weight : int, default=0
40
        If nonzero, the fields defined on self.domain are weighted with the
Martin Reinecke's avatar
fix    
Martin Reinecke committed
41
        specified power along the submdomains which are contracted.
Philipp Arras's avatar
Philipp Arras committed
42
43
    """

44
    def __init__(self, domain, spaces, weight=0):
45
46
47
        self._domain = DomainTuple.make(domain)
        self._spaces = utilities.parse_spaces(spaces, len(self._domain))
        self._target = [
48
            dom for i, dom in enumerate(self._domain) if i not in self._spaces
Philipp Arras's avatar
Philipp Arras committed
49
        ]
50
        self._target = DomainTuple.make(self._target)
51
        self._weight = weight
Martin Reinecke's avatar
Martin Reinecke committed
52
        self._capability = self.TIMES | self.ADJOINT_TIMES
Philipp Arras's avatar
Fixups    
Philipp Arras committed
53
54
55

    def apply(self, x, mode):
        self._check_input(x, mode)
56
        if mode == self.ADJOINT_TIMES:
57
            ldat = x.to_global_data() if 0 in self._spaces else x.local_data
58
            shp = []
59
60
            for i, dom in enumerate(self._domain):
                tmp = dom.shape if i > 0 else dom.local_shape
61
                shp += tmp if i not in self._spaces else (1,)*len(dom.shape)
62
            ldat = np.broadcast_to(ldat.reshape(shp), self._domain.local_shape)
63
64
65
66
            res = Field.from_local_data(self._domain, ldat)
            if self._weight != 0:
                res = res.weight(self._weight, spaces=self._spaces)
            return res
Philipp Arras's avatar
Fixups    
Philipp Arras committed
67
        else:
68
69
70
71
            if self._weight != 0:
                x = x.weight(self._weight, spaces=self._spaces)
            res = x.sum(self._spaces)
            return res if isinstance(res, Field) else Field.scalar(res)