power_space.py 6.19 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',
Martin Reinecke's avatar
Martin Reinecke committed
36
                 log=False, nbin=None, binbounds=None):
37

Martin Reinecke's avatar
Martin Reinecke committed
38
        super(PowerSpace, self).__init__()
39
40
        self._ignore_for_hash += ['_pindex', '_kindex', '_rho', '_pundex',
                                  '_k_array']
41
42

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

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

        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']
65
66
        self._pundex = power_index['pundex']
        self._k_array = power_index['k_array']
67

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

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

    @property
    def harmonic(self):
        return True
79

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

84
85
86
87
88
89
90
    @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
91
        return float(reduce(lambda x, y: x*y, self.pindex.shape))
92
93

    def copy(self):
94
        distribution_strategy = self.pindex.distribution_strategy
95
        return self.__class__(harmonic_domain=self.harmonic_domain,
96
                              distribution_strategy=distribution_strategy,
97
98
                              log=self.log,
                              nbin=self.nbin,
Martin Reinecke's avatar
Martin Reinecke committed
99
                              binbounds=self.binbounds)
100

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

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

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
115
116
117

        return result_x

118
    def get_distance_array(self, distribution_strategy):
119
        result = d2o.distributed_data_object(
Martin Reinecke's avatar
Martin Reinecke committed
120
                                self.kindex, dtype=np.float64,
121
122
                                distribution_strategy=distribution_strategy)
        return result
theos's avatar
theos committed
123

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

128
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
    # ---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
157

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

    @property
    def k_array(self):
        return self._k_array
165
166
167
168

    # ---Serialization---

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

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

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

Theo Steininger's avatar
Theo Steininger committed
197
        new_ps._pindex = repository.get('pindex', hdf5_group)
Jait Dixit's avatar
Jait Dixit committed
198
199
200
        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
201
        new_ps._k_array = repository.get('k_array', hdf5_group)
Jait Dixit's avatar
Jait Dixit committed
202
203
        new_ps._ignore_for_hash += ['_pindex', '_kindex', '_rho', '_pundex',
                                   '_k_array']
Jait Dixit's avatar
Jait Dixit committed
204
205
206
207

        return new_ps


208
209
210
    def plot(self):
        pass

Jait Dixit's avatar
Jait Dixit committed
211
212
213
class EmptyPowerSpace(PowerSpace):
    def __init__(self):
        pass