gl_space.py 6.68 KB
Newer Older
csongor's avatar
csongor committed
1 2
from __future__ import division

3 4
import pickle

Jait Dixit's avatar
Jait Dixit committed
5
import itertools
csongor's avatar
csongor committed
6 7
import numpy as np

8 9
import d2o
from d2o import STRATEGIES as DISTRIBUTION_STRATEGIES
10
from keepers import Versionable
csongor's avatar
csongor committed
11

12
from nifty.spaces.space import Space
13
from nifty.config import nifty_configuration as gc,\
14
                         dependency_injector as gdi
15
import nifty.nifty_utilities as utilities
csongor's avatar
csongor committed
16 17 18 19 20

gl = gdi.get('libsharp_wrapper_gl')

GL_DISTRIBUTION_STRATEGIES = DISTRIBUTION_STRATEGIES['global']

21

22
class GLSpace(Versionable, Space):
csongor's avatar
csongor committed
23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
    """
        ..                 __
        ..               /  /
        ..     ____ __  /  /
        ..   /   _   / /  /
        ..  /  /_/  / /  /_
        ..  \___   /  \___/  space class
        .. /______/

        NIFTY subclass for Gauss-Legendre pixelizations [#]_ of the two-sphere.

        Parameters
        ----------
        nlat : int
            Number of latitudinal bins, or rings.
        nlon : int, *optional*
            Number of longitudinal bins (default: ``2*nlat - 1``).
        dtype : numpy.dtype, *optional*
            Data type of the field values (default: numpy.float64).

        See Also
        --------
        hp_space : A class for the HEALPix discretization of the sphere [#]_.
        lm_space : A class for spherical harmonic components.

        Notes
        -----
        Only real-valued fields on the two-sphere are supported, i.e.
        `dtype` has to be either numpy.float64 or numpy.float32.

        References
        ----------
        .. [#] M. Reinecke and D. Sverre Seljebotn, 2013, "Libsharp - spherical
               harmonic transforms revisited";
               `arXiv:1303.4945 <http://www.arxiv.org/abs/1303.4945>`_
        .. [#] K.M. Gorski et al., 2005, "HEALPix: A Framework for
               High-Resolution Discretization and Fast Analysis of Data
               Distributed on the Sphere", *ApJ* 622..759G.

        Attributes
        ----------
        para : numpy.ndarray
            One-dimensional array containing the two numbers `nlat` and `nlon`.
        dtype : numpy.dtype
            Data type of the field values.
        discrete : bool
            Whether or not the underlying space is discrete, always ``False``
            for spherical spaces.
        vol : numpy.ndarray
            An array containing the pixel sizes.
    """

75 76
    # ---Overwritten properties and methods---

77
    def __init__(self, nlat=2, nlon=None, dtype=np.dtype('float')):
csongor's avatar
csongor committed
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
        """
            Sets the attributes for a gl_space class instance.

            Parameters
            ----------
            nlat : int
                Number of latitudinal bins, or rings.
            nlon : int, *optional*
                Number of longitudinal bins (default: ``2*nlat - 1``).
            dtype : numpy.dtype, *optional*
                Data type of the field values (default: numpy.float64).

            Returns
            -------
            None

            Raises
            ------
            ImportError
                If the libsharp_wrapper_gl module is not available.
            ValueError
                If input `nlat` is invaild.

        """
        # check imports
        if not gc['use_libsharp']:
104 105
            raise ImportError(
                "libsharp_wrapper_gl not available or not loaded.")
106 107

        super(GLSpace, self).__init__(dtype)
csongor's avatar
csongor committed
108

109 110
        self._nlat = self._parse_nlat(nlat)
        self._nlon = self._parse_nlon(nlon)
csongor's avatar
csongor committed
111

112
    # ---Mandatory properties and methods---
csongor's avatar
csongor committed
113

114 115 116
    @property
    def harmonic(self):
        return False
csongor's avatar
csongor committed
117 118 119

    @property
    def shape(self):
120
        return (np.int((self.nlat * self.nlon)),)
csongor's avatar
csongor committed
121

122
    @property
123
    def dim(self):
124
        return np.int((self.nlat * self.nlon))
125 126 127 128

    @property
    def total_volume(self):
        return 4 * np.pi
129

130 131 132 133 134
    def copy(self):
        return self.__class__(nlat=self.nlat,
                              nlon=self.nlon,
                              dtype=self.dtype)

Jait Dixit's avatar
Jait Dixit committed
135
    def weight(self, x, power=1, axes=None, inplace=False):
136
        axes = utilities.cast_axis_to_tuple(axes, length=1)
137

138 139
        nlon = self.nlon
        nlat = self.nlat
140 141

        weight = np.array(list(itertools.chain.from_iterable(
142 143
            itertools.repeat(x ** power, nlon)
            for x in gl.vol(nlat))))
Jait Dixit's avatar
Jait Dixit committed
144 145 146

        if axes is not None:
            # reshape the weight array to match the input shape
147
            new_shape = np.ones(len(x.shape), dtype=np.int)
Jait Dixit's avatar
Jait Dixit committed
148 149 150 151 152 153 154
            for index in range(len(axes)):
                new_shape[index] = len(weight)
            weight = weight.reshape(new_shape)

        if inplace:
            x *= weight
            result_x = x
csongor's avatar
csongor committed
155
        else:
Jait Dixit's avatar
Jait Dixit committed
156
            result_x = x * weight
csongor's avatar
csongor committed
157

Jait Dixit's avatar
Jait Dixit committed
158
        return result_x
159

160
    def get_distance_array(self, distribution_strategy):
161 162
        dists = d2o.arange(start=0, stop=self.shape[0],
                           distribution_strategy=distribution_strategy)
163

164
        dists = dists.apply_scalar_function(
165
            lambda x: self._distance_array_helper(divmod(x, self.nlon)),
166
            dtype=np.float)
167 168 169

        return dists

theos's avatar
theos committed
170
    def _distance_array_helper(self, qr_tuple):
171 172
        lat = qr_tuple[0]*(np.pi/(self.nlat-1))
        lon = qr_tuple[1]*(2*np.pi/(self.nlon-1))
173 174 175
        numerator = np.sqrt(np.sin(lat)**2 +
                            (np.sin(lon) * np.cos(lat))**2)
        denominator = np.cos(lon) * np.cos(lat)
176

theos's avatar
theos committed
177
        return np.arctan(numerator / denominator)
178

179
    def get_fft_smoothing_kernel_function(self, sigma):
Jait Dixit's avatar
Jait Dixit committed
180
        if sigma is None:
181
            sigma = np.sqrt(2) * np.pi
Jait Dixit's avatar
Jait Dixit committed
182 183

        return lambda x: np.exp((-0.5 * x**2) / sigma**2)
184

185 186 187 188 189 190 191 192 193 194 195 196 197
    # ---Added properties and methods---

    @property
    def nlat(self):
        return self._nlat

    @property
    def nlon(self):
        return self._nlon

    def _parse_nlat(self, nlat):
        nlat = int(nlat)
        if nlat < 2:
198 199
            raise ValueError(
                "nlat must be a positive number.")
200
        elif nlat % 2 != 0:
201 202
            raise ValueError(
                "nlat must be a multiple of 2.")
203 204 205 206 207 208 209 210
        return nlat

    def _parse_nlon(self, nlon):
        if nlon is None:
            nlon = 2 * self.nlat - 1
        else:
            nlon = int(nlon)
            if nlon != 2 * self.nlat - 1:
211 212
                self.logger.warn("nlon was set to an unrecommended value: "
                                 "nlon <> 2*nlat-1.")
213
        return nlon
214 215 216 217

    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Jait Dixit's avatar
Jait Dixit committed
218 219 220 221
        hdf5_group['nlat'] = self.nlat
        hdf5_group['nlon'] = self.nlon
        hdf5_group['dtype'] = pickle.dumps(self.dtype)

222 223 224 225 226
        return None

    @classmethod
    def _from_hdf5(cls, hdf5_group, loopback_get):
        result = cls(
Jait Dixit's avatar
Jait Dixit committed
227 228 229 230 231
            nlat=hdf5_group['nlat'][()],
            nlon=hdf5_group['nlon'][()],
            dtype=pickle.loads(hdf5_group['dtype'][()])
            )

232
        return result