rg_space.py 12.3 KB
Newer Older
1 2 3 4
# 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.
5
#
6 7
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
8 9 10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
11
# You should have received a copy of the GNU General Public License
12
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
13 14 15 16 17
#
# Copyright(C) 2013-2017 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.
Marco Selig's avatar
Marco Selig committed
18 19 20 21 22 23 24 25 26 27

"""
    ..                  __   ____   __
    ..                /__/ /   _/ /  /_
    ..      __ ___    __  /  /_  /   _/  __   __
    ..    /   _   | /  / /   _/ /  /   /  / /  /
    ..   /  / /  / /  / /  /   /  /_  /  /_/  /
    ..  /__/ /__/ /__/ /__/    \___/  \___   /  rg
    ..                               /______/

Marco Selig's avatar
Marco Selig committed
28
    NIFTY submodule for regular Cartesian grids.
Marco Selig's avatar
Marco Selig committed
29 30 31

"""
from __future__ import division
32

Marco Selig's avatar
Marco Selig committed
33
import numpy as np
Theo Steininger's avatar
Theo Steininger committed
34

35 36
from d2o import distributed_data_object,\
                STRATEGIES as DISTRIBUTION_STRATEGIES
37

38
from nifty.spaces.space import Space
39

Marco Selig's avatar
Marco Selig committed
40

41
class RGSpace(Space):
Marco Selig's avatar
Marco Selig committed
42 43 44 45 46 47 48 49 50
    """
        ..      _____   _______
        ..    /   __/ /   _   /
        ..   /  /    /  /_/  /
        ..  /__/     \____  /  space class
        ..          /______/

        NIFTY subclass for spaces of regular Cartesian grids.

Theo Steininger's avatar
Theo Steininger committed
51 52 53 54
        Parameters
        ----------
        shape : {int, numpy.ndarray}
            Number of grid points or numbers of gridpoints along each axis.
55 56 57 58
        zerocenter : {bool, numpy.ndarray} *optional*
            Whether x==0 (or k==0, respectively) is located in the center of
            the grid (or the center of each axis speparately) or not.
            (default: False).
Theo Steininger's avatar
Theo Steininger committed
59 60 61 62 63 64 65 66 67 68
        distances : {float, numpy.ndarray}, *optional*
            Distance between two grid points along each axis
            (default: None).
            If distances==None:
                if harmonic==True, all distances will be set to 1
                if harmonic==False, the distance along each axis will be
                  set to the inverse of the number of points along that
                  axis.
        harmonic : bool, *optional*
        Whether the space represents a grid in position or harmonic space.
Theo Steininger's avatar
Theo Steininger committed
69
            (default: False).
Marco Selig's avatar
Marco Selig committed
70 71 72

        Attributes
        ----------
Martin Reinecke's avatar
Martin Reinecke committed
73
        harmonic : bool
Theo Steininger's avatar
Theo Steininger committed
74 75 76 77 78
            Whether or not the grid represents a position or harmonic space.
        zerocenter : tuple of bool
            Whether x==0 (or k==0, respectively) is located in the center of
            the grid (or the center of each axis speparately) or not.
        distances : tuple of floats
79 80 81 82 83 84 85 86 87
            Distance between two grid points along the correponding axis.
        dim : np.int
            Total number of dimensionality, i.e. the number of pixels.
        harmonic : bool
            Specifies whether the space is a signal or harmonic space.
        total_volume : np.float
            The total volume of the space.
        shape : tuple of np.ints
            The shape of the space's data array.
Theo Steininger's avatar
Theo Steininger committed
88

Marco Selig's avatar
Marco Selig committed
89 90
    """

91 92
    # ---Overwritten properties and methods---

93
    def __init__(self, shape, zerocenter=False, distances=None,
94
                 harmonic=False):
95 96
        self._harmonic = bool(harmonic)

97
        super(RGSpace, self).__init__()
98

99 100 101
        self._shape = self._parse_shape(shape)
        self._distances = self._parse_distances(distances)
        self._zerocenter = self._parse_zerocenter(zerocenter)
Marco Selig's avatar
Marco Selig committed
102

103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
# This code is unused but may be useful to keep around if it is ever needed
# again in the future ...

#    def hermitian_fixed_points(self):
#        dimensions = len(self.shape)
#        mid_index = np.array(self.shape)//2
#        ndlist = [1]*dimensions
#        for k in range(dimensions):
#            if self.shape[k] % 2 == 0:
#                ndlist[k] = 2
#        ndlist = tuple(ndlist)
#        fixed_points = []
#        for index in np.ndindex(ndlist):
#            for k in range(dimensions):
#                if self.shape[k] % 2 != 0 and self.zerocenter[k]:
#                    index = list(index)
#                    index[k] = 1
#                    index = tuple(index)
#            fixed_points += [tuple(index * mid_index)]
#        return fixed_points

124
    def hermitianize_inverter(self, x, axes):
125
        # calculate the number of dimensions the input array has
Martin Reinecke's avatar
Martin Reinecke committed
126
        dimensions = len(x.shape)
127 128 129 130 131 132
        # prepare the slicing object which will be used for mirroring
        slice_primitive = [slice(None), ] * dimensions
        # copy the input data
        y = x.copy()

        # flip in the desired directions
133 134
        for k in range(len(axes)):
            i = axes[k]
135 136
            slice_picker = slice_primitive[:]
            slice_inverter = slice_primitive[:]
137
            if (not self.zerocenter[k]) or self.shape[k] % 2 == 0:
Martin Reinecke's avatar
Martin Reinecke committed
138
                slice_picker[i] = slice(1, None, None)
139 140
                slice_inverter[i] = slice(None, 0, -1)
            else:
Martin Reinecke's avatar
Martin Reinecke committed
141
                slice_picker[i] = slice(None)
142
                slice_inverter[i] = slice(None, None, -1)
Martin Reinecke's avatar
Martin Reinecke committed
143
            slice_picker = tuple(slice_picker)
144 145 146
            slice_inverter = tuple(slice_inverter)

            try:
Martin Reinecke's avatar
Martin Reinecke committed
147 148
                y.set_data(to_key=slice_picker, data=y,
                           from_key=slice_inverter)
149 150 151 152
            except(AttributeError):
                y[slice_picker] = y[slice_inverter]
        return y

153 154
    # ---Mandatory properties and methods---

155 156 157 158
    def __repr__(self):
        return ("RGSpace(shape=%r, zerocenter=%r, distances=%r, harmonic=%r)"
                % (self.shape, self.zerocenter, self.distances, self.harmonic))

159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178
    @property
    def harmonic(self):
        return self._harmonic

    @property
    def shape(self):
        return self._shape

    @property
    def dim(self):
        return reduce(lambda x, y: x*y, self.shape)

    @property
    def total_volume(self):
        return self.dim * reduce(lambda x, y: x*y, self.distances)

    def copy(self):
        return self.__class__(shape=self.shape,
                              zerocenter=self.zerocenter,
                              distances=self.distances,
179
                              harmonic=self.harmonic)
180 181

    def weight(self, x, power=1, axes=None, inplace=False):
182
        weight = reduce(lambda x, y: x*y, self.distances) ** np.float(power)
183 184 185 186 187 188 189
        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
        return result_x

190
    def get_distance_array(self, distribution_strategy):
Theo Steininger's avatar
Theo Steininger committed
191 192
        """ Calculates an n-dimensional array with its entries being the
        lengths of the vectors from the zero point of the grid.
193

Theo Steininger's avatar
Theo Steininger committed
194 195
        Parameters
        ----------
Theo Steininger's avatar
Theo Steininger committed
196 197 198
        distribution_strategy : str
            The distribution_strategy which shall be used the returned
            distributed_data_object.
199

Theo Steininger's avatar
Theo Steininger committed
200 201
        Returns
        -------
Theo Steininger's avatar
Theo Steininger committed
202
        distributed_data_object
Theo Steininger's avatar
Theo Steininger committed
203 204
            A d2o containing the distances.

205
        """
Theo Steininger's avatar
Theo Steininger committed
206

207 208 209
        shape = self.shape
        # prepare the distributed_data_object
        nkdict = distributed_data_object(
Martin Reinecke's avatar
Martin Reinecke committed
210
                        global_shape=shape, dtype=np.float64,
211 212 213 214 215 216 217 218 219
                        distribution_strategy=distribution_strategy)

        if distribution_strategy in DISTRIBUTION_STRATEGIES['slicing']:
            # get the node's individual slice of the first dimension
            slice_of_first_dimension = slice(
                                    *nkdict.distributor.local_slice[0:2])
        elif distribution_strategy in DISTRIBUTION_STRATEGIES['not']:
            slice_of_first_dimension = slice(0, shape[0])
        else:
220 221
            raise ValueError(
                "Unsupported distribution strategy")
222 223 224 225 226 227 228 229 230 231 232 233 234 235 236
        dists = self._distance_array_helper(slice_of_first_dimension)
        nkdict.set_local_data(dists)

        return nkdict

    def _distance_array_helper(self, slice_of_first_dimension):
        dk = self.distances
        shape = self.shape

        inds = []
        for a in shape:
            inds += [slice(0, a)]

        cords = np.ogrid[inds]

237 238
        dists = (cords[0] - shape[0]//2)*dk[0]
        dists *= dists
239
        # apply zerocenterQ shift
240 241
        if not self.zerocenter[0]:
            dists = np.fft.ifftshift(dists)
242 243 244
        # only save the individual slice
        dists = dists[slice_of_first_dimension]
        for ii in range(1, len(shape)):
245 246
            temp = (cords[ii] - shape[ii] // 2) * dk[ii]
            temp *= temp
247
            if not self.zerocenter[ii]:
Martin Reinecke's avatar
Martin Reinecke committed
248
                temp = np.fft.ifftshift(temp)
249 250 251 252
            dists = dists + temp
        dists = np.sqrt(dists)
        return dists

Martin Reinecke's avatar
Martin Reinecke committed
253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282
    def get_unique_distances(self):
        dimensions = len(self.shape)
        if dimensions == 1:  # extra easy
            maxdist = self.shape[0]//2
            return np.arange(maxdist+1, dtype=np.float64) * self.distances[0]
        if np.all(self.distances == self.distances[0]):  # shortcut
            maxdist = np.asarray(self.shape)//2
            tmp = np.sum(maxdist*maxdist)
            tmp = np.zeros(tmp+1, dtype=np.bool)
            t2 = np.arange(maxdist[0]+1, dtype=np.int64)
            t2 *= t2
            for i in range(1, dimensions):
                t3 = np.arange(maxdist[i]+1, dtype=np.int64)
                t3 *= t3
                t2 = np.add.outer(t2, t3)
            tmp[t2] = True
            return np.sqrt(np.nonzero(tmp)[0])*self.distances[0]
        else:  # do it the hard way
            tmp = self.get_distance_array('not').unique()  # expensive!
            tol = 1e-12*tmp[-1]
            # remove all points that are closer than tol to their right
            # neighbors.
            # I'm appending the last value*2 to the array to treat the
            # rightmost point correctly.
            return tmp[np.diff(np.r_[tmp, 2*tmp[-1]]) > tol]

    def get_natural_binbounds(self):
        tmp = self.get_unique_distances()
        return 0.5*(tmp[:-1]+tmp[1:])

283
    def get_fft_smoothing_kernel_function(self, sigma):
284
        return lambda x: np.exp(-2. * np.pi*np.pi * x*x * sigma*sigma)
285

286 287 288 289
    # ---Added properties and methods---

    @property
    def distances(self):
Theo Steininger's avatar
Theo Steininger committed
290 291 292
        """Distance between two grid points along each axis. It is a tuple
        of positive floating point numbers with the n-th entry giving the
        distances of grid points along the n-th dimension.
293
        """
Theo Steininger's avatar
Theo Steininger committed
294

295 296 297 298
        return self._distances

    @property
    def zerocenter(self):
299
        """Returns True if grid points lie symmetrically around zero.
Theo Steininger's avatar
Theo Steininger committed
300

301 302
        Returns
        -------
Theo Steininger's avatar
Theo Steininger committed
303 304 305 306 307
        bool
            True if the grid points are centered around the 0 grid point. This
            option is most common for harmonic spaces (where both conventions
            are used) but may be used for position spaces, too.

308
        """
Theo Steininger's avatar
Theo Steininger committed
309

310 311 312 313 314 315 316 317 318 319 320 321
        return self._zerocenter

    def _parse_shape(self, shape):
        if np.isscalar(shape):
            shape = (shape,)
        temp = np.empty(len(shape), dtype=np.int)
        temp[:] = shape
        return tuple(temp)

    def _parse_distances(self, distances):
        if distances is None:
            if self.harmonic:
322
                temp = np.ones_like(self.shape, dtype=np.float64)
323
            else:
324
                temp = 1 / np.array(self.shape, dtype=np.float64)
325
        else:
326
            temp = np.empty(len(self.shape), dtype=np.float64)
327 328 329 330 331 332
            temp[:] = distances
        return tuple(temp)

    def _parse_zerocenter(self, zerocenter):
        temp = np.empty(len(self.shape), dtype=bool)
        temp[:] = zerocenter
333 334
        if np.any(np.logical_and(temp, np.array(self.shape) % 2)):
            raise ValueError("All zerocentered axis must have even length!")
335
        return tuple(temp)
336 337 338 339

    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
340 341 342
        hdf5_group['shape'] = self.shape
        hdf5_group['zerocenter'] = self.zerocenter
        hdf5_group['distances'] = self.distances
343
        hdf5_group['harmonic'] = self.harmonic
344

345 346 347
        return None

    @classmethod
348
    def _from_hdf5(cls, hdf5_group, repository):
349
        result = cls(
350 351 352
            shape=hdf5_group['shape'][:],
            zerocenter=hdf5_group['zerocenter'][:],
            distances=hdf5_group['distances'][:],
353
            harmonic=hdf5_group['harmonic'][()],
354
            )
355
        return result