rg_space.py 7.84 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/>.
Theo Steininger's avatar
Theo Steininger committed
13
#
14
# Copyright(C) 2013-2019 Max-Planck-Society
Theo Steininger's avatar
Theo Steininger committed
15
#
16
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
Philipp Arras's avatar
Philipp Arras committed
17

Martin Reinecke's avatar
Martin Reinecke committed
18
from functools import reduce
Marco Selig's avatar
Marco Selig committed
19
import numpy as np
Philipp Arras's avatar
Philipp Arras committed
20 21 22

from ..field import Field
from .structured_domain import StructuredDomain
csongor's avatar
csongor committed
23

Marco Selig's avatar
Marco Selig committed
24

Martin Reinecke's avatar
Martin Reinecke committed
25
class RGSpace(StructuredDomain):
26
    """Represents a regular Cartesian grid.
Martin Reinecke's avatar
Martin Reinecke committed
27 28 29

    Parameters
    ----------
Martin Reinecke's avatar
Martin Reinecke committed
30
    shape : int or tuple of int
Martin Reinecke's avatar
Martin Reinecke committed
31
        Number of grid points or numbers of gridpoints along each axis.
Martin Reinecke's avatar
Martin Reinecke committed
32
    distances : None or float or tuple of float, optional
Philipp Arras's avatar
Philipp Arras committed
33
        Distance between two grid points along each axis.
Martin Reinecke's avatar
Martin Reinecke committed
34

Philipp Arras's avatar
Philipp Arras committed
35 36 37
        By default (distances=None):
          - If harmonic==True, all distances will be set to 1
          - If harmonic==False, the distance along each axis will be
Martin Reinecke's avatar
Martin Reinecke committed
38 39
            set to the inverse of the number of points along that axis.

Martin Reinecke's avatar
Martin Reinecke committed
40
    harmonic : bool, optional
41
        Whether the space represents a grid in position or harmonic space.
Philipp Arras's avatar
Philipp Arras committed
42
        Default: False.
Philipp Arras's avatar
Philipp Arras committed
43 44 45 46 47

    Notes
    -----
    Topologically, a n-dimensional RGSpace is a n-Torus, i.e. it has periodic
    boundary conditions.
Marco Selig's avatar
Marco Selig committed
48
    """
Martin Reinecke's avatar
Martin Reinecke committed
49
    _needed_for_hash = ["_distances", "_shape", "_harmonic"]
50

Martin Reinecke's avatar
Martin Reinecke committed
51
    def __init__(self, shape, distances=None, harmonic=False):
Martin Reinecke's avatar
Martin Reinecke committed
52
        self._harmonic = bool(harmonic)
Martin Reinecke's avatar
Martin Reinecke committed
53 54 55
        if np.isscalar(shape):
            shape = (shape,)
        self._shape = tuple(int(i) for i in shape)
56 57
        if min(self._shape) < 0:
            raise ValueError('Negative number of pixels encountered')
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
58 59 60 61 62 63 64 65 66 67 68 69

        if distances is None:
            if self.harmonic:
                self._distances = (1.,) * len(self._shape)
            else:
                self._distances = tuple(1./s for s in self._shape)
        elif np.isscalar(distances):
            self._distances = (float(distances),) * len(self._shape)
        else:
            temp = np.empty(len(self.shape), dtype=np.float64)
            temp[:] = distances
            self._distances = tuple(temp)
70 71
        if min(self._distances) <= 0:
            raise ValueError('Non-positive distances encountered')
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
72

73
        self._dvol = float(reduce(lambda x, y: x*y, self._distances))
Martin Reinecke's avatar
Martin Reinecke committed
74
        self._size = int(reduce(lambda x, y: x*y, self._shape))
Marco Selig's avatar
Marco Selig committed
75

76
    def __repr__(self):
Martin Reinecke's avatar
Martin Reinecke committed
77 78
        return ("RGSpace(shape={}, distances={}, harmonic={})"
                .format(self.shape, self.distances, self.harmonic))
79

80 81 82 83 84 85 86 87 88
    @property
    def harmonic(self):
        return self._harmonic

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

    @property
Martin Reinecke's avatar
Martin Reinecke committed
89 90
    def size(self):
        return self._size
91

Martin Reinecke's avatar
Martin Reinecke committed
92
    @property
93 94
    def scalar_dvol(self):
        return self._dvol
95

96
    def _get_dist_array(self):
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
97
        res = np.arange(self.shape[0], dtype=np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
98 99
        res = np.minimum(res, self.shape[0]-res)*self.distances[0]
        if len(self.shape) == 1:
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
100
            return Field.from_arr(self, res)
Martin Reinecke's avatar
Martin Reinecke committed
101 102
        res *= res
        for i in range(1, len(self.shape)):
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
103
            tmp = np.arange(self.shape[i], dtype=np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
104 105 106
            tmp = np.minimum(tmp, self.shape[i]-tmp)*self.distances[i]
            tmp *= tmp
            res = np.add.outer(res, tmp)
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
107
        return Field.from_arr(self, np.sqrt(res))
108

109 110 111 112 113
    def get_k_length_array(self):
        if (not self.harmonic):
            raise NotImplementedError
        return self._get_dist_array()

114
    def get_unique_k_lengths(self):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
115 116
        if (not self.harmonic):
            raise NotImplementedError
Martin Reinecke's avatar
Martin Reinecke committed
117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133
        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
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
134
            tmp = self.get_k_length_array().val
Martin Reinecke's avatar
Martin Reinecke committed
135
            tmp = np.unique(tmp)
Martin Reinecke's avatar
Martin Reinecke committed
136 137 138 139 140 141 142
            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]

Martin Reinecke's avatar
Martin Reinecke committed
143 144
    @staticmethod
    def _kernel(x, sigma):
145
        from ..sugar import exp
146
        return exp(x*x * (-2.*np.pi*np.pi*sigma*sigma))
Martin Reinecke's avatar
Martin Reinecke committed
147

148
    def get_fft_smoothing_kernel_function(self, sigma):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
149 150
        if (not self.harmonic):
            raise NotImplementedError
Martin Reinecke's avatar
Martin Reinecke committed
151
        return lambda x: self._kernel(x, sigma)
152

153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169
    def get_conv_kernel_from_func(self, func):
        """Creates a convolution kernel defined by a function.

        Parameters
        ----------
        func: function
            This function needs to take exactly one argument, which is
            distance from center (in the same units as the RGSpace distances),
            and return the kernel amplitude at that distance.

        Assumes the function to be radially symmetric,
        e.g. only dependant on distance"""
        from ..operators.harmonic_operators import HarmonicTransformOperator
        if (not self.harmonic):
            raise NotImplementedError
        op = HarmonicTransformOperator(self, self.get_default_codomain())
        dist = op.target[0]._get_dist_array()
Martin Reinecke's avatar
stage2  
Martin Reinecke committed
170
        kernel = Field(op.target, func(dist.val))
171 172 173
        kernel = kernel / kernel.integrate()
        return op.adjoint_times(kernel.weight(1))

Martin Reinecke's avatar
Martin Reinecke committed
174
    def get_default_codomain(self):
Martin Reinecke's avatar
Martin Reinecke committed
175 176 177 178 179 180
        """Returns a :class:`RGSpace` object representing the (position or
        harmonic) partner domain of `self`, depending on `self.harmonic`.

        Returns
        -------
        RGSpace
Martin Reinecke's avatar
Martin Reinecke committed
181
            The partner domain
Martin Reinecke's avatar
Martin Reinecke committed
182
        """
Martin Reinecke's avatar
Martin Reinecke committed
183 184 185 186
        distances = 1. / (np.array(self.shape)*np.array(self.distances))
        return RGSpace(self.shape, distances, not self.harmonic)

    def check_codomain(self, codomain):
Martin Reinecke's avatar
Martin Reinecke committed
187 188 189
        """Raises `TypeError` if `codomain` is not a matching partner domain
        for `self`.
        """
Martin Reinecke's avatar
Martin Reinecke committed
190 191 192 193 194 195 196 197 198 199 200 201
        if not isinstance(codomain, RGSpace):
            raise TypeError("domain is not a RGSpace")

        if self.shape != codomain.shape:
            raise AttributeError("The shapes of domain and codomain must be "
                                 "identical.")

        if self.harmonic == codomain.harmonic:
            raise AttributeError("domain.harmonic and codomain.harmonic must "
                                 "not be the same.")

        # Check if the distances match, i.e. dist' = 1 / (num * dist)
202 203 204
        if not np.all(abs(np.array(self.shape) *
                          np.array(self.distances) *
                          np.array(codomain.distances)-1) < 1e-7):
Martin Reinecke's avatar
Martin Reinecke committed
205 206 207
            raise AttributeError("The grid-distances of domain and codomain "
                                 "do not match.")

208 209
    @property
    def distances(self):
Martin Reinecke's avatar
Martin Reinecke committed
210 211 212
        """tuple of float : Distance between grid points along each axis.
        The n-th entry of the tuple is the distance between neighboring
        grid points along the n-th dimension.
213
        """
214
        return self._distances