contraction_operator.py 2.74 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
Martin Reinecke committed
27
28
    """A  :class:`LinearOperator` which sums up fields into the direction of
    subspaces.
Philipp Arras's avatar
Philipp Arras committed
29

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

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

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

    def apply(self, x, mode):
        self._check_input(x, mode)
55
        if mode == self.ADJOINT_TIMES:
56
            ldat = x.to_global_data() if 0 in self._spaces else x.local_data
57
            shp = []
58
59
            for i, dom in enumerate(self._domain):
                tmp = dom.shape if i > 0 else dom.local_shape
60
                shp += tmp if i not in self._spaces else (1,)*len(dom.shape)
61
            ldat = np.broadcast_to(ldat.reshape(shp), self._domain.local_shape)
62
63
64
65
            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
66
        else:
67
68
69
70
            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)