contraction_operator.py 2.72 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
27
class ContractionOperator(LinearOperator):
    """A linear operator which sums up fields into the direction of subspaces.
Philipp Arras's avatar
Philipp Arras committed
28

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

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

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

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