power_space.py 6.31 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# NIFTy
# Copyright (C) 2017  Theo Steininger
#
# Author: Theo Steininger
#
# 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/>.
theos's avatar
theos committed
18

theos's avatar
theos committed
19
20
import numpy as np

21
22
import d2o

23
24
from power_index_factory import PowerIndexFactory

25
from nifty.spaces.space import Space
26
from nifty.spaces.rg_space import RGSpace
27
from nifty.nifty_utilities import cast_axis_to_tuple
theos's avatar
theos committed
28
29


Theo Steininger's avatar
Theo Steininger committed
30
class PowerSpace(Space):
31

32
33
    # ---Overwritten properties and methods---

34
35
    def __init__(self, harmonic_domain=RGSpace((1,)),
                 distribution_strategy='not',
Jait Dixit's avatar
Jait Dixit committed
36
                 log=False, nbin=None, binbounds=None,
37
                 dtype=None):
38
39

        super(PowerSpace, self).__init__(dtype)
40
41
        self._ignore_for_hash += ['_pindex', '_kindex', '_rho', '_pundex',
                                  '_k_array']
42
43

        if not isinstance(harmonic_domain, Space):
44
45
            raise ValueError(
                "harmonic_domain must be a Space.")
46
        if not harmonic_domain.harmonic:
47
48
            raise ValueError(
                "harmonic_domain must be a harmonic space.")
49
50
        self._harmonic_domain = harmonic_domain

Jait Dixit's avatar
Jait Dixit committed
51
52
53
54
55
56
        power_index = PowerIndexFactory.get_power_index(
                        domain=self.harmonic_domain,
                        distribution_strategy=distribution_strategy,
                        log=log,
                        nbin=nbin,
                        binbounds=binbounds)
57
58
59
60
61
62
63
64
65

        config = power_index['config']
        self._log = config['log']
        self._nbin = config['nbin']
        self._binbounds = config['binbounds']

        self._pindex = power_index['pindex']
        self._kindex = power_index['kindex']
        self._rho = power_index['rho']
66
67
        self._pundex = power_index['pundex']
        self._k_array = power_index['k_array']
68

69
70
71
72
73
74
    def pre_cast(self, x, axes=None):
        if callable(x):
            return x(self.kindex)
        else:
            return x

75
76
77
78
79
    # ---Mandatory properties and methods---

    @property
    def harmonic(self):
        return True
80

81
82
    @property
    def shape(self):
83
        return self.kindex.shape
84

85
86
87
88
89
90
91
    @property
    def dim(self):
        return self.shape[0]

    @property
    def total_volume(self):
        # every power-pixel has a volume of 1
Jait Dixit's avatar
Jait Dixit committed
92
        return float(reduce(lambda x, y: x*y, self.pindex.shape))
93
94

    def copy(self):
95
        distribution_strategy = self.pindex.distribution_strategy
96
        return self.__class__(harmonic_domain=self.harmonic_domain,
97
                              distribution_strategy=distribution_strategy,
98
99
100
101
                              log=self.log,
                              nbin=self.nbin,
                              binbounds=self.binbounds,
                              dtype=self.dtype)
102

103
    def weight(self, x, power=1, axes=None, inplace=False):
Jait Dixit's avatar
Jait Dixit committed
104
105
        reshaper = [1, ] * len(x.shape)
        # we know len(axes) is always 1
106
107
        reshaper[axes[0]] = self.shape[0]

108
        weight = self.rho.reshape(reshaper)
109
110
        if power != 1:
            weight = weight ** power
111
112
113
114
115
116

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
117
118
119

        return result_x

120
    def get_distance_array(self, distribution_strategy):
121
122
123
124
        result = d2o.distributed_data_object(
                                self.kindex,
                                distribution_strategy=distribution_strategy)
        return result
theos's avatar
theos committed
125

126
    def get_fft_smoothing_kernel_function(self, sigma):
127
        raise NotImplementedError(
128
            "There is no fft smoothing function for PowerSpace.")
theos's avatar
theos committed
129

130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    # ---Added properties and methods---

    @property
    def harmonic_domain(self):
        return self._harmonic_domain

    @property
    def log(self):
        return self._log

    @property
    def nbin(self):
        return self._nbin

    @property
    def binbounds(self):
        return self._binbounds

    @property
    def pindex(self):
        return self._pindex

    @property
    def kindex(self):
        return self._kindex

    @property
    def rho(self):
        return self._rho
159

160
161
162
163
164
165
166
    @property
    def pundex(self):
        return self._pundex

    @property
    def k_array(self):
        return self._k_array
167
168
169
170

    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Jait Dixit's avatar
Jait Dixit committed
171
172
173
        hdf5_group['kindex'] = self.kindex
        hdf5_group['rho'] = self.rho
        hdf5_group['pundex'] = self.pundex
Theo Steininger's avatar
Theo Steininger committed
174
        hdf5_group.attrs['dtype'] = self.dtype.name
175
        hdf5_group['log'] = self.log
Theo Steininger's avatar
Theo Steininger committed
176
177
178
        # Store nbin as string, since it can be None
        hdf5_group.attrs['nbin'] = str(self.nbin)
        hdf5_group.attrs['binbounds'] = str(self.binbounds)
179
180
181
182
183
184
185
186

        return {
            'harmonic_domain': self.harmonic_domain,
            'pindex': self.pindex,
            'k_array': self.k_array
        }

    @classmethod
Theo Steininger's avatar
Theo Steininger committed
187
    def _from_hdf5(cls, hdf5_group, repository):
Jait Dixit's avatar
Jait Dixit committed
188
189
190
191
        # make an empty PowerSpace object
        new_ps = EmptyPowerSpace()
        # reset class
        new_ps.__class__ = cls
Jait Dixit's avatar
Jait Dixit committed
192
193
        # call instructor so that classes are properly setup
        super(PowerSpace, new_ps).__init__(np.dtype(hdf5_group.attrs['dtype']))
Jait Dixit's avatar
Jait Dixit committed
194
        # set all values
Theo Steininger's avatar
Theo Steininger committed
195
        new_ps._harmonic_domain = repository.get('harmonic_domain', hdf5_group)
196
        new_ps._log = hdf5_group['log'][()]
Theo Steininger's avatar
Theo Steininger committed
197
198
        exec('new_ps._nbin = ' + hdf5_group.attrs['nbin'])
        exec('new_ps._binbounds = ' + hdf5_group.attrs['binbounds'])
Jait Dixit's avatar
Jait Dixit committed
199

Theo Steininger's avatar
Theo Steininger committed
200
        new_ps._pindex = repository.get('pindex', hdf5_group)
Jait Dixit's avatar
Jait Dixit committed
201
202
203
        new_ps._kindex = hdf5_group['kindex'][:]
        new_ps._rho = hdf5_group['rho'][:]
        new_ps._pundex = hdf5_group['pundex'][:]
Theo Steininger's avatar
Theo Steininger committed
204
        new_ps._k_array = repository.get('k_array', hdf5_group)
Jait Dixit's avatar
Jait Dixit committed
205
206
        new_ps._ignore_for_hash += ['_pindex', '_kindex', '_rho', '_pundex',
                                   '_k_array']
Jait Dixit's avatar
Jait Dixit committed
207
208
209
210
211
212
213

        return new_ps


class EmptyPowerSpace(PowerSpace):
    def __init__(self):
        pass