nifty_lm.py 79.8 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
# NIFTY (Numerical Information Field Theory) has been developed at the
# Max-Planck-Institute for Astrophysics.
#
# Copyright (C) 2015 Max-Planck-Society
#
# Author: Marco Selig
# Project homepage: <http://www.mpa-garching.mpg.de/ift/nifty/>
#
# 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/>.
Marco Selig's avatar
Marco Selig committed
21 22 23 24 25 26 27 28 29 30 31 32 33 34

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

    NIFTY submodule for grids on the two-sphere.

"""
from __future__ import division
35

Marco Selig's avatar
Marco Selig committed
36 37 38 39 40
import os
import numpy as np
import pylab as pl
from matplotlib.colors import LogNorm as ln
from matplotlib.ticker import LogFormatter as lf
41

42 43
from d2o import STRATEGIES as DISTRIBUTION_STRATEGIES

44
from nifty.nifty_core import Space
csongor's avatar
csongor committed
45

46
from nifty.field import Field
csongor's avatar
csongor committed
47

48 49 50
from nifty.config import about,\
                         nifty_configuration as gc,\
                         dependency_injector as gdi
Ultimanet's avatar
Ultimanet committed
51
from nifty.nifty_paradict import lm_space_paradict,\
52 53 54
                                 gl_space_paradict,\
                                 hp_space_paradict
from nifty.nifty_power_indices import lm_power_indices
Ultimanet's avatar
Ultimanet committed
55
from nifty.nifty_random import random
56

Ultima's avatar
Ultima committed
57 58
gl = gdi.get('libsharp_wrapper_gl')
hp = gdi.get('healpy')
59

60 61 62
LM_DISTRIBUTION_STRATEGIES = DISTRIBUTION_STRATEGIES['global']
GL_DISTRIBUTION_STRATEGIES = DISTRIBUTION_STRATEGIES['global']
HP_DISTRIBUTION_STRATEGIES = DISTRIBUTION_STRATEGIES['global']
Marco Selig's avatar
Marco Selig committed
63 64


65
class lm_space(Space):
Marco Selig's avatar
Marco Selig committed
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84
    """
        ..       __
        ..     /  /
        ..    /  /    __ ____ ___
        ..   /  /   /   _    _   |
        ..  /  /_  /  / /  / /  /
        ..  \___/ /__/ /__/ /__/  space class

        NIFTY subclass for spherical harmonics components, for representations
        of fields on the two-sphere.

        Parameters
        ----------
        lmax : int
            Maximum :math:`\ell`-value up to which the spherical harmonics
            coefficients are to be used.
        mmax : int, *optional*
            Maximum :math:`m`-value up to which the spherical harmonics
            coefficients are to be used (default: `lmax`).
85
        dtype : numpy.dtype, *optional*
Marco Selig's avatar
Marco Selig committed
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
            Data type of the field values (default: numpy.complex128).

        See Also
        --------
        hp_space : A class for the HEALPix discretization of the sphere [#]_.
        gl_space : A class for the Gauss-Legendre discretization of the
            sphere [#]_.

        Notes
        -----
        Hermitian symmetry, i.e. :math:`a_{\ell -m} = \overline{a}_{\ell m}` is
        always assumed for the spherical harmonics components, i.e. only fields
        on the two-sphere with real-valued representations in position space
        can be handled.

        References
        ----------
        .. [#] 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.
        .. [#] M. Reinecke and D. Sverre Seljebotn, 2013, "Libsharp - spherical
               harmonic transforms revisited";
               `arXiv:1303.4945 <http://www.arxiv.org/abs/1303.4945>`_

        Attributes
        ----------
        para : numpy.ndarray
            One-dimensional array containing the two numbers `lmax` and
            `mmax`.
115
        dtype : numpy.dtype
Marco Selig's avatar
Marco Selig committed
116 117 118 119 120 121 122
            Data type of the field values.
        discrete : bool
            Parameter captioning the fact that an :py:class:`lm_space` is
            always discrete.
        vol : numpy.ndarray
            Pixel volume of the :py:class:`lm_space`, which is always 1.
    """
123

csongor's avatar
csongor committed
124
    def __init__(self, lmax, mmax=None, dtype=np.dtype('complex128')):
Marco Selig's avatar
Marco Selig committed
125 126 127 128 129 130 131 132 133 134 135
        """
            Sets the attributes for an lm_space class instance.

            Parameters
            ----------
            lmax : int
                Maximum :math:`\ell`-value up to which the spherical harmonics
                coefficients are to be used.
            mmax : int, *optional*
                Maximum :math:`m`-value up to which the spherical harmonics
                coefficients are to be used (default: `lmax`).
136
            dtype : numpy.dtype, *optional*
Marco Selig's avatar
Marco Selig committed
137 138 139 140 141 142 143 144 145 146 147 148 149 150 151
                Data type of the field values (default: numpy.complex128).

            Returns
            -------
            None.

            Raises
            ------
            ImportError
                If neither the libsharp_wrapper_gl nor the healpy module are
                available.
            ValueError
                If input `nside` is invaild.

        """
152

153
        # check imports
Ultima's avatar
Ultima committed
154
        if not gc['use_libsharp'] and not gc['use_healpy']:
155
            raise ImportError(about._errors.cstring(
Ultima's avatar
Ultima committed
156
                "ERROR: neither libsharp_wrapper_gl nor healpy activated."))
157

Ultima's avatar
Ultima committed
158 159
        self._cache_dict = {'check_codomain': {}}

160
        self.paradict = lm_space_paradict(lmax=lmax, mmax=mmax)
Marco Selig's avatar
Marco Selig committed
161

162 163 164
        # check data type
        dtype = np.dtype(dtype)
        if dtype not in [np.dtype('complex64'), np.dtype('complex128')]:
theos's avatar
theos committed
165
            about.warnings.cprint("WARNING: data type set to complex128.")
166 167
            dtype = np.dtype('complex128')
        self.dtype = dtype
168

Marco Selig's avatar
Marco Selig committed
169
        self.discrete = True
170
        self.harmonic = True
171 172 173 174
        self.distances = (np.float(1),)

        self.power_indices = lm_power_indices(
                    lmax=self.paradict['lmax'],
175
                    dim=self.dim,
176
                    allowed_distribution_strategies=LM_DISTRIBUTION_STRATEGIES)
Marco Selig's avatar
Marco Selig committed
177

178 179
    @property
    def para(self):
180
        temp = np.array([self.paradict['lmax'],
181 182
                         self.paradict['mmax']], dtype=int)
        return temp
183

184 185 186 187 188
    @para.setter
    def para(self, x):
        self.paradict['lmax'] = x[0]
        self.paradict['mmax'] = x[1]

Ultima's avatar
Ultima committed
189 190 191
    def __hash__(self):
        result_hash = 0
        for (key, item) in vars(self).items():
Ultima's avatar
Ultima committed
192
            if key in ['_cache_dict', 'power_indices']:
Ultima's avatar
Ultima committed
193 194 195 196
                continue
            result_hash ^= item.__hash__() * hash(key)
        return result_hash

Ultima's avatar
Ultima committed
197 198 199 200 201 202
    def _identifier(self):
        # Extract the identifying parts from the vars(self) dict.
        temp = [(ii[0],
                 ((lambda x: tuple(x) if
                  isinstance(x, np.ndarray) else x)(ii[1])))
                for ii in vars(self).iteritems()
csongor's avatar
csongor committed
203
                if ii[0] not in ['_cache_dict', 'power_indices']]
Ultima's avatar
Ultima committed
204 205 206
        # Return the sorted identifiers as a tuple.
        return tuple(sorted(temp))

207
    def copy(self):
208 209 210 211
        return lm_space(lmax=self.paradict['lmax'],
                        mmax=self.paradict['mmax'],
                        dtype=self.dtype)

212 213
    @property
    def shape(self):
Ultima's avatar
Ultima committed
214
        lmax = self.paradict['lmax']
Ultima's avatar
Ultima committed
215 216
        mmax = self.paradict['mmax']
        return (np.int((mmax + 1) * (lmax + 1) - ((mmax + 1) * mmax) // 2),)
217

218 219
    @property
    def dof(self):
Marco Selig's avatar
Marco Selig committed
220 221 222 223 224 225 226 227 228 229 230 231 232 233
        """
            Computes the number of degrees of freedom of the space, taking into
            account symmetry constraints and complex-valuedness.

            Returns
            -------
            dof : int
                Number of degrees of freedom of the space.

            Notes
            -----
            The number of degrees of freedom is reduced due to the hermitian
            symmetry, which is assumed for the spherical harmonics components.
        """
234 235
        # dof = 2*dim-(lmax+1) = (lmax+1)*(2*mmax+1)*(mmax+1)*mmax
        lmax = self.paradict['lmax']
Ultima's avatar
Ultima committed
236
        mmax = self.paradict['mmax']
237
        dof = np.int((lmax + 1) * (2 * mmax + 1) - (mmax + 1) * mmax)
238 239 240 241 242
        return dof

    @property
    def dof_split(self):
        return (self.dof,)
Marco Selig's avatar
Marco Selig committed
243

244 245
    @property
    def meta_volume(self):
Marco Selig's avatar
Marco Selig committed
246
        """
247
            Calculates the meta volumes.
Marco Selig's avatar
Marco Selig committed
248

249 250 251 252
            The meta volumes are the volumes associated with each component of
            a field, taking into account field components that are not
            explicitly included in the array of field values but are determined
            by symmetry conditions.
Marco Selig's avatar
Marco Selig committed
253 254 255

            Parameters
            ----------
256 257 258
            total : bool, *optional*
                Whether to return the total meta volume of the space or the
                individual ones of each field component (default: False).
Marco Selig's avatar
Marco Selig committed
259 260 261

            Returns
            -------
262 263
            mol : {numpy.ndarray, float}
                Meta volume of the field components or the complete space.
Marco Selig's avatar
Marco Selig committed
264

265 266 267 268 269
            Notes
            -----
            The spherical harmonics components with :math:`m=0` have meta
            volume 1, the ones with :math:`m>0` have meta volume 2, sinnce they
            each determine another component with negative :math:`m`.
Marco Selig's avatar
Marco Selig committed
270
        """
271 272 273 274 275 276 277
        return np.float(self.dof())

    @property
    def meta_volume_split(self):
        mol = self.cast(1, dtype=np.float)
        mol[self.paradict['lmax'] + 1:] = 2  # redundant: (l,m) and (l,-m)
        return mol
Marco Selig's avatar
Marco Selig committed
278

279
    def complement_cast(self, x, axis=None, **kwargs):
csongor's avatar
csongor committed
280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295
        if axis is None:
            lmax = self.paradict['lmax']
            complexity_mask = x[:lmax+1].iscomplex()
            if complexity_mask.any():
                about.warnings.cprint("WARNING: Taking the absolute values for " +
                                      "all complex entries where lmax==0")
                x[:lmax+1] = abs(x[:lmax+1])
        else:
            # TODO hermitianize only on specific axis
            lmax = self.paradict['lmax']
            complexity_mask = x[:lmax+1].iscomplex()
            if complexity_mask.any():
                about.warnings.cprint("WARNING: Taking the absolute values for " +
                                      "all complex entries where lmax==0")
                x[:lmax+1] = abs(x[:lmax+1])
        return x
296

297
    # TODO: Extend to binning/log
298 299 300 301 302
    def enforce_power(self, spec, size=None, kindex=None):
        if kindex is None:
            kindex_size = self.paradict['lmax'] + 1
            kindex = np.arange(kindex_size,
                               dtype=np.array(self.distances).dtype)
303 304 305 306
        return self._enforce_power_helper(spec=spec,
                                          size=size,
                                          kindex=kindex)

Ultima's avatar
Ultima committed
307
    def _check_codomain(self, codomain):
Marco Selig's avatar
Marco Selig committed
308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326
        """
            Checks whether a given codomain is compatible to the
            :py:class:`lm_space` or not.

            Parameters
            ----------
            codomain : nifty.space
                Space to be checked for compatibility.

            Returns
            -------
            check : bool
                Whether or not the given codomain is compatible to the space.

            Notes
            -----
            Compatible codomains are instances of :py:class:`lm_space`,
            :py:class:`gl_space`, and :py:class:`hp_space`.
        """
327 328
        if codomain is None:
            return False
329

330 331 332
        if not isinstance(codomain, space):
            raise TypeError(about._errors.cstring(
                "ERROR: The given codomain must be a nifty lm_space."))
Marco Selig's avatar
Marco Selig committed
333

334 335 336
        elif isinstance(codomain, gl_space):
            # lmax==mmax
            # nlat==lmax+1
337
            # nlon==2*lmax+1
338 339 340
            if ((self.paradict['lmax'] == self.paradict['mmax']) and
                    (codomain.paradict['nlat'] == self.paradict['lmax']+1) and
                    (codomain.paradict['nlon'] == 2*self.paradict['lmax']+1)):
Marco Selig's avatar
Marco Selig committed
341 342
                return True

343 344 345 346 347
        elif isinstance(codomain, hp_space):
            # lmax==mmax
            # 3*nside-1==lmax
            if ((self.paradict['lmax'] == self.paradict['mmax']) and
                    (3*codomain.paradict['nside']-1 == self.paradict['lmax'])):
Marco Selig's avatar
Marco Selig committed
348 349 350 351
                return True

        return False

352
    def get_codomain(self, coname=None, **kwargs):
Marco Selig's avatar
Marco Selig committed
353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
        """
            Generates a compatible codomain to which transformations are
            reasonable, i.e.\  a pixelization of the two-sphere.

            Parameters
            ----------
            coname : string, *optional*
                String specifying a desired codomain (default: None).

            Returns
            -------
            codomain : nifty.space
                A compatible codomain.

            Notes
            -----
            Possible arguments for `coname` are ``'gl'`` in which case a Gauss-
            Legendre pixelization [#]_ of the sphere is generated, and ``'hp'``
            in which case a HEALPix pixelization [#]_ is generated.

            References
            ----------
            .. [#] 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.
378 379
            .. [#] M. Reinecke and D. Sverre Seljebotn, 2013,
                   "Libsharp - spherical
Marco Selig's avatar
Marco Selig committed
380 381 382 383
                   harmonic transforms revisited";
                   `arXiv:1303.4945 <http://www.arxiv.org/abs/1303.4945>`_

        """
384 385 386 387 388
        if coname == 'gl' or (coname is None and gc['lm2gl']):
            if self.dtype == np.dtype('complex64'):
                new_dtype = np.float32
            elif self.dtype == np.dtype('complex128'):
                new_dtype = np.float64
Marco Selig's avatar
Marco Selig committed
389
            else:
390 391 392
                raise NotImplementedError
            nlat = self.paradict['lmax'] + 1
            nlon = self.paradict['lmax'] * 2 + 1
csongor's avatar
csongor committed
393
            return gl_space(nlat=nlat, nlon=nlon, dtype=new_dtype)
394

395 396
        elif coname == 'hp' or (coname is None and not gc['lm2gl']):
            nside = (self.paradict['lmax']+1) // 3
csongor's avatar
csongor committed
397
            return hp_space(nside=nside)
398

Marco Selig's avatar
Marco Selig committed
399
        else:
400
            raise ValueError(about._errors.cstring(
401 402 403 404 405 406 407 408 409 410 411 412
                "ERROR: unsupported or incompatible codomain '"+coname+"'."))

    def get_random_values(self, **kwargs):
        """
            Generates random field values according to the specifications given
            by the parameters, taking into account complex-valuedness and
            hermitian symmetry.

            Returns
            -------
            x : numpy.ndarray
                Valid field values.
Marco Selig's avatar
Marco Selig committed
413

414 415 416 417 418 419
            Other parameters
            ----------------
            random : string, *optional*
                Specifies the probability distribution from which the random
                numbers are to be drawn.
                Supported distributions are:
Marco Selig's avatar
Marco Selig committed
420

421 422 423 424 425 426
                - "pm1" (uniform distribution over {+1,-1} or {+1,+i,-1,-i}
                - "gau" (normal distribution with zero-mean and a given
                    standard
                    deviation or variance)
                - "syn" (synthesizes from a given power spectrum)
                - "uni" (uniform distribution over [vmin,vmax[)
Marco Selig's avatar
Marco Selig committed
427

428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443
                (default: None).
            dev : float, *optional*
                Standard deviation (default: 1).
            var : float, *optional*
                Variance, overriding `dev` if both are specified
                (default: 1).
            spec : {scalar, list, numpy.array, nifty.field, function},
                *optional*
                Power spectrum (default: 1).
            vmin : float, *optional*
                Lower limit for a uniform distribution (default: 0).
            vmax : float, *optional*
                Upper limit for a uniform distribution (default: 1).
        """
        arg = random.parse_arguments(self, **kwargs)

444 445 446 447
#        if arg is None:
#            x = 0
#
#        elif arg['random'] == "pm1":
448
#            x = random.pm1(dtype=self.dtype, shape=self.shape)
449 450 451
#
#        elif arg['random'] == "gau":
#            x = random.gau(dtype=self.dtype,
452
#                           shape=self.shape,
453 454
#                           mean=arg['mean'],
#                           std=arg['std'])
455

456
        if arg['random'] == "syn":
457 458 459
            lmax = self.paradict['lmax']
            mmax = self.paradict['mmax']
            if self.dtype == np.dtype('complex64'):
Ultima's avatar
Ultima committed
460
                if gc['use_libsharp']:
461
                    sample = gl.synalm_f(arg['spec'], lmax=lmax, mmax=mmax)
462
                else:
463 464 465 466
                    sample = hp.synalm(
                                arg['spec'].astype(np.complex128),
                                lmax=lmax, mmax=mmax).astype(np.complex64,
                                                             copy=False)
467
            else:
Ultima's avatar
Ultima committed
468
                if gc['use_healpy']:
469
                    sample = hp.synalm(arg['spec'], lmax=lmax, mmax=mmax)
470
                else:
471
                    sample = gl.synalm(arg['spec'], lmax=lmax, mmax=mmax)
472 473

        else:
474
            sample = super(lm_space, self).get_random_values(**arg)
Marco Selig's avatar
Marco Selig committed
475

476 477
#        elif arg['random'] == "uni":
#            x = random.uni(dtype=self.dtype,
478
#                           shape=self.shape,
479 480 481 482 483 484 485 486
#                           vmin=arg['vmin'],
#                           vmax=arg['vmax'])
#
#        else:
#            raise KeyError(about._errors.cstring(
#                "ERROR: unsupported random key '" + str(arg['random']) + "'."))
        sample = self.cast(sample)
        return sample
487

488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522
#    def calc_dot(self, x, y):
#        """
#            Computes the discrete inner product of two given arrays of field
#            values.
#
#            Parameters
#            ----------
#            x : numpy.ndarray
#                First array
#            y : numpy.ndarray
#                Second array
#
#            Returns
#            -------
#            dot : scalar
#                Inner product of the two arrays.
#        """
#        x = self.cast(x)
#        y = self.cast(y)
#
#        lmax = self.paradict['lmax']
#
#        x_low = x[:lmax + 1]
#        x_high = x[lmax + 1:]
#        y_low = y[:lmax + 1]
#        y_high = y[lmax + 1:]
#
#        dot = (x_low.real * y_low.real).sum()
#        dot += 2 * (x_high.real * y_high.real).sum()
#        dot += 2 * (x_high.imag * y_high.imag).sum()
#        return dot

    def dot_contraction(self, x, axes):
        assert len(axes) == 1
        axis = axes[0]
523
        lmax = self.paradict['lmax']
524

525 526 527 528 529
        # extract the low and high parts of x
        extractor = ()
        extractor += (slice(None),)*axis
        low_extractor = extractor + (slice(None, lmax+1), )
        high_extractor = extractor + (slice(lmax+1), )
530

531 532
        result = x[low_extractor].sum(axes) + 2 * x[high_extractor].sum(axes)
        return result
533

534
    def calc_transform(self, x, codomain=None, **kwargs):
Marco Selig's avatar
Marco Selig committed
535 536 537 538 539 540 541 542
        """
            Computes the transform of a given array of field values.

            Parameters
            ----------
            x : numpy.ndarray
                Array to be transformed.
            codomain : nifty.space, *optional*
543
                codomain space to which the transformation shall map
Marco Selig's avatar
Marco Selig committed
544 545 546 547 548 549 550
                (default: self).

            Returns
            -------
            Tx : numpy.ndarray
                Transformed array
        """
551
        x = self.cast(x)
Marco Selig's avatar
Marco Selig committed
552

553 554
        if codomain is None:
            codomain = self.get_codomain()
Marco Selig's avatar
Marco Selig committed
555

556 557 558 559
        # Check if the given codomain is suitable for the transformation
        if not self.check_codomain(codomain):
            raise ValueError(about._errors.cstring(
                "ERROR: unsupported codomain."))
Marco Selig's avatar
Marco Selig committed
560

csongor's avatar
csongor committed
561 562 563 564
        # if self.datamodel != 'not':
        #     about.warnings.cprint(
        #         "WARNING: Field data is consolidated to all nodes for "
        #         "external alm2map method!")
565 566 567

        np_x = x.get_full_data()

568 569 570 571 572
        if isinstance(codomain, gl_space):
            nlat = codomain.paradict['nlat']
            nlon = codomain.paradict['nlon']
            lmax = self.paradict['lmax']
            mmax = self.paradict['mmax']
Marco Selig's avatar
Marco Selig committed
573

574
            # transform
575
            if self.dtype == np.dtype('complex64'):
576 577
                np_Tx = gl.alm2map_f(np_x, nlat=nlat, nlon=nlon,
                                     lmax=lmax, mmax=mmax, cl=False)
Marco Selig's avatar
Marco Selig committed
578
            else:
579 580 581
                np_Tx = gl.alm2map(np_x, nlat=nlat, nlon=nlon,
                                   lmax=lmax, mmax=mmax, cl=False)
            Tx = codomain.cast(np_Tx)
Marco Selig's avatar
Marco Selig committed
582

583 584 585 586 587
        elif isinstance(codomain, hp_space):
            nside = codomain.paradict['nside']
            lmax = self.paradict['lmax']
            mmax = self.paradict['mmax']

588
            # transform
589 590 591 592 593
            np_x = np_x.astype(np.complex128, copy=False)
            np_Tx = hp.alm2map(np_x, nside, lmax=lmax,
                               mmax=mmax, pixwin=False, fwhm=0.0, sigma=None,
                               pol=True, inplace=False)
            Tx = codomain.cast(np_Tx)
Marco Selig's avatar
Marco Selig committed
594 595

        else:
596 597
            raise ValueError(about._errors.cstring(
                "ERROR: unsupported transformation."))
Marco Selig's avatar
Marco Selig committed
598

599 600 601 602 603
        # re-weight if discrete
        if codomain.discrete:
            Tx = codomain.calc_weight(Tx, power=0.5)

        return codomain.cast(Tx)
Marco Selig's avatar
Marco Selig committed
604

605
    def calc_smooth(self, x, sigma=0, **kwargs):
Marco Selig's avatar
Marco Selig committed
606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623
        """
            Smoothes an array of field values by convolution with a Gaussian
            kernel in position space.

            Parameters
            ----------
            x : numpy.ndarray
                Array of field values to be smoothed.
            sigma : float, *optional*
                Standard deviation of the Gaussian kernel, specified in units
                of length in position space; for testing: a sigma of -1 will be
                reset to a reasonable value (default: 0).

            Returns
            -------
            Gx : numpy.ndarray
                Smoothed array.
        """
624
        x = self.cast(x)
625
        # check sigma
626
        if sigma == 0:
Ultima's avatar
Ultima committed
627
            return self.unary_operation(x, op='copy')
628
        elif sigma == -1:
Marco Selig's avatar
Marco Selig committed
629
            about.infos.cprint("INFO: invalid sigma reset.")
630 631
            sigma = np.sqrt(2) * np.pi / (self.paradict['lmax'] + 1)
        elif sigma < 0:
Marco Selig's avatar
Marco Selig committed
632
            raise ValueError(about._errors.cstring("ERROR: invalid sigma."))
633

csongor's avatar
csongor committed
634 635 636 637
        # if self.datamodel != 'not':
        #     about.warnings.cprint(
        #         "WARNING: Field data is consolidated to all nodes for "
        #         "external smoothalm method!")
theos's avatar
theos committed
638 639 640

        np_x = x.get_full_data()

Ultima's avatar
Ultima committed
641
        if gc['use_healpy']:
theos's avatar
theos committed
642 643 644 645 646 647 648
            np_smoothed_x = hp.smoothalm(np_x,
                                         fwhm=0.0,
                                         sigma=sigma,
                                         pol=True,
                                         mmax=self.paradict['mmax'],
                                         verbose=False,
                                         inplace=False)
Marco Selig's avatar
Marco Selig committed
649
        else:
theos's avatar
theos committed
650 651 652 653 654 655 656
            np_smoothed_x = gl.smoothalm(np_x,
                                         lmax=self.paradict['lmax'],
                                         mmax=self.paradict['mmax'],
                                         fwhm=0.0,
                                         sigma=sigma,
                                         overwrite=False)
        return self.cast(np_smoothed_x)
Marco Selig's avatar
Marco Selig committed
657

658
    def calc_power(self, x, **kwargs):
Marco Selig's avatar
Marco Selig committed
659 660 661 662 663 664 665 666 667 668 669 670 671 672
        """
            Computes the power of an array of field values.

            Parameters
            ----------
            x : numpy.ndarray
                Array containing the field values of which the power is to be
                calculated.

            Returns
            -------
            spec : numpy.ndarray
                Power contained in the input array.
        """
673 674 675 676
        x = self.cast(x)
        lmax = self.paradict['lmax']
        mmax = self.paradict['mmax']

csongor's avatar
csongor committed
677 678 679 680
        # if self.datamodel != 'not':
        #     about.warnings.cprint(
        #         "WARNING: Field data is consolidated to all nodes for "
        #         "external anaalm/alm2cl method!")
theos's avatar
theos committed
681 682 683

        np_x = x.get_full_data()

684
        # power spectrum
685
        if self.dtype == np.dtype('complex64'):
Ultima's avatar
Ultima committed
686
            if gc['use_libsharp']:
theos's avatar
theos committed
687
                result = gl.anaalm_f(np_x, lmax=lmax, mmax=mmax)
Marco Selig's avatar
Marco Selig committed
688
            else:
theos's avatar
theos committed
689 690 691 692 693 694 695
                np_x = np_x.astype(np.complex128, copy=False)
                result = hp.alm2cl(np_x,
                                   alms2=None,
                                   lmax=lmax,
                                   mmax=mmax,
                                   lmax_out=lmax,
                                   nspec=None)
Marco Selig's avatar
Marco Selig committed
696
        else:
Ultima's avatar
Ultima committed
697
            if gc['use_healpy']:
theos's avatar
theos committed
698 699 700 701 702 703
                result = hp.alm2cl(np_x,
                                   alms2=None,
                                   lmax=lmax,
                                   mmax=mmax,
                                   lmax_out=lmax,
                                   nspec=None)
Marco Selig's avatar
Marco Selig committed
704
            else:
theos's avatar
theos committed
705 706 707
                result = gl.anaalm(np_x,
                                   lmax=lmax,
                                   mmax=mmax)
theos's avatar
theos committed
708 709 710 711 712 713 714 715 716

        if self.dtype == np.dtype('complex64'):
            result = result.astype(np.float32, copy=False)
        elif self.dtype == np.dtype('complex128'):
            result = result.astype(np.float64, copy=False)
        else:
            raise NotImplementedError(about._errors.cstring(
                "ERROR: dtype %s not known to calc_power method." %
                str(self.dtype)))
Marco Selig's avatar
Marco Selig committed
717

718 719 720
    def get_plot(self, x, title="", vmin=None, vmax=None, power=True,
                 norm=None, cmap=None, cbar=True, other=None, legend=False,
                 mono=True, **kwargs):
Marco Selig's avatar
Marco Selig committed
721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764
        """
            Creates a plot of field values according to the specifications
            given by the parameters.

            Parameters
            ----------
            x : numpy.ndarray
                Array containing the field values.

            Returns
            -------
            None

            Other parameters
            ----------------
            title : string, *optional*
                Title of the plot (default: "").
            vmin : float, *optional*
                Minimum value to be displayed (default: ``min(x)``).
            vmax : float, *optional*
                Maximum value to be displayed (default: ``max(x)``).
            power : bool, *optional*
                Whether to plot the power contained in the field or the field
                values themselves (default: True).
            unit : string, *optional*
                Unit of the field values (default: "").
            norm : string, *optional*
                Scaling of the field values before plotting (default: None).
            cmap : matplotlib.colors.LinearSegmentedColormap, *optional*
                Color map to be used for two-dimensional plots (default: None).
            cbar : bool, *optional*
                Whether to show the color bar or not (default: True).
            other : {single object, tuple of objects}, *optional*
                Object or tuple of objects to be added, where objects can be
                scalars, arrays, or fields (default: None).
            legend : bool, *optional*
                Whether to show the legend or not (default: False).
            mono : bool, *optional*
                Whether to plot the monopole or not (default: True).
            save : string, *optional*
                Valid file name where the figure is to be stored, by default
                the figure is not saved (default: False).

        """
theos's avatar
theos committed
765 766 767 768 769
        try:
            x = x.get_full_data()
        except AttributeError:
            pass

770
        if(not pl.isinteractive())and(not bool(kwargs.get("save", False))):
Marco Selig's avatar
Marco Selig committed
771 772 773 774 775
            about.warnings.cprint("WARNING: interactive mode off.")

        if(power):
            x = self.calc_power(x)

776 777 778
            fig = pl.figure(num=None, figsize=(6.4, 4.8), dpi=None, facecolor="none",
                            edgecolor="none", frameon=False, FigureClass=pl.Figure)
            ax0 = fig.add_axes([0.12, 0.12, 0.82, 0.76])
Marco Selig's avatar
Marco Selig committed
779

780
            xaxes = np.arange(self.para[0] + 1, dtype=np.int)
Marco Selig's avatar
Marco Selig committed
781
            if(vmin is None):
782 783
                vmin = np.min(x[:mono].tolist(
                ) + (xaxes * (2 * xaxes + 1) * x)[1:].tolist(), axis=None, out=None)
Marco Selig's avatar
Marco Selig committed
784
            if(vmax is None):
785 786 787 788
                vmax = np.max(x[:mono].tolist(
                ) + (xaxes * (2 * xaxes + 1) * x)[1:].tolist(), axis=None, out=None)
            ax0.loglog(xaxes[1:], (xaxes * (2 * xaxes + 1) * x)[1:], color=[0.0,
                                                                            0.5, 0.0], label="graph 0", linestyle='-', linewidth=2.0, zorder=1)
Marco Selig's avatar
Marco Selig committed
789
            if(mono):
790 791
                ax0.scatter(0.5 * (xaxes[1] + xaxes[2]), x[0], s=20, color=[0.0, 0.5, 0.0], marker='o',
                            cmap=None, norm=None, vmin=None, vmax=None, alpha=None, linewidths=None, verts=None, zorder=1)
Marco Selig's avatar
Marco Selig committed
792 793

            if(other is not None):
794
                if(isinstance(other, tuple)):
Marco Selig's avatar
Marco Selig committed
795 796
                    other = list(other)
                    for ii in xrange(len(other)):
797
                        if(isinstance(other[ii], field)):
Marco Selig's avatar
Marco Selig committed
798 799 800
                            other[ii] = other[ii].power(**kwargs)
                        else:
                            other[ii] = self.enforce_power(other[ii])
801
                elif(isinstance(other, field)):
Marco Selig's avatar
Marco Selig committed
802 803 804
                    other = [other.power(**kwargs)]
                else:
                    other = [self.enforce_power(other)]
805
                imax = max(1, len(other) - 1)
Marco Selig's avatar
Marco Selig committed
806
                for ii in xrange(len(other)):
807 808
                    ax0.loglog(xaxes[1:], (xaxes * (2 * xaxes + 1) * other[ii])[1:], color=[max(0.0, 1.0 - (2 * ii / imax)**2), 0.5 * ((2 * ii - imax) / imax)
                                                                                            ** 2, max(0.0, 1.0 - (2 * (ii - imax) / imax)**2)], label="graph " + str(ii + 1), linestyle='-', linewidth=1.0, zorder=-ii)
Marco Selig's avatar
Marco Selig committed
809
                    if(mono):
810 811
                        ax0.scatter(0.5 * (xaxes[1] + xaxes[2]), other[ii][0], s=20, color=[max(0.0, 1.0 - (2 * ii / imax)**2), 0.5 * ((2 * ii - imax) / imax)**2, max(
                            0.0, 1.0 - (2 * (ii - imax) / imax)**2)], marker='o', cmap=None, norm=None, vmin=None, vmax=None, alpha=None, linewidths=None, verts=None, zorder=-ii)
Marco Selig's avatar
Marco Selig committed
812 813 814
                if(legend):
                    ax0.legend()

815
            ax0.set_xlim(xaxes[1], xaxes[-1])
Marco Selig's avatar
Marco Selig committed
816
            ax0.set_xlabel(r"$\ell$")
817
            ax0.set_ylim(vmin, vmax)
Marco Selig's avatar
Marco Selig committed
818 819 820 821 822 823 824
            ax0.set_ylabel(r"$\ell(2\ell+1) C_\ell$")
            ax0.set_title(title)

        else:
            if(np.iscomplexobj(x)):
                if(title):
                    title += " "
825 826 827 828 829 830
                if(bool(kwargs.get("save", False))):
                    save_ = os.path.splitext(
                        os.path.basename(str(kwargs.get("save"))))
                    kwargs.update(save=save_[0] + "_absolute" + save_[1])
                self.get_plot(np.absolute(x), title=title + "(absolute)", vmin=vmin, vmax=vmax,
                              power=False, norm=norm, cmap=cmap, cbar=cbar, other=None, legend=False, **kwargs)
Marco Selig's avatar
Marco Selig committed
831 832 833 834
#                self.get_plot(np.real(x),title=title+"(real part)",vmin=vmin,vmax=vmax,power=False,norm=norm,cmap=cmap,cbar=cbar,other=None,legend=False,**kwargs)
#                self.get_plot(np.imag(x),title=title+"(imaginary part)",vmin=vmin,vmax=vmax,power=False,norm=norm,cmap=cmap,cbar=cbar,other=None,legend=False,**kwargs)
                if(cmap is None):
                    cmap = pl.cm.hsv_r
835 836 837 838 839
                if(bool(kwargs.get("save", False))):
                    kwargs.update(save=save_[0] + "_phase" + save_[1])
                self.get_plot(np.angle(x, deg=False), title=title + "(phase)", vmin=-3.1416, vmax=3.1416, power=False,
                              norm=None, cmap=cmap, cbar=cbar, other=None, legend=False, **kwargs)  # values in [-pi,pi]
                return None  # leave method
Marco Selig's avatar
Marco Selig committed
840 841
            else:
                if(vmin is None):
842
                    vmin = np.min(x, axis=None, out=None)
Marco Selig's avatar
Marco Selig committed
843
                if(vmax is None):
844 845 846 847 848 849 850 851 852 853
                    vmax = np.max(x, axis=None, out=None)
                if(norm == "log")and(vmin <= 0):
                    raise ValueError(about._errors.cstring(
                        "ERROR: nonpositive value(s)."))

                # not a number
                xmesh = np.nan * \
                    np.empty(self.para[::-1] + 1, dtype=np.float16, order='C')
                xmesh[4, 1] = None
                xmesh[1, 4] = None
Marco Selig's avatar
Marco Selig committed
854
                lm = 0
855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870
                for mm in xrange(self.para[1] + 1):
                    xmesh[mm][mm:] = x[lm:lm + self.para[0] + 1 - mm]
                    lm += self.para[0] + 1 - mm

                s_ = np.array([1, self.para[1] / self.para[0]
                               * (1.0 + 0.159 * bool(cbar))])
                fig = pl.figure(num=None, figsize=(
                    6.4 * s_[0], 6.4 * s_[1]), dpi=None, facecolor="none", edgecolor="none", frameon=False, FigureClass=pl.Figure)
                ax0 = fig.add_axes(
                    [0.06 / s_[0], 0.06 / s_[1], 1.0 - 0.12 / s_[0], 1.0 - 0.12 / s_[1]])
                ax0.set_axis_bgcolor([0.0, 0.0, 0.0, 0.0])

                xaxes = np.arange(self.para[0] + 2, dtype=np.int) - 0.5
                yaxes = np.arange(self.para[1] + 2, dtype=np.int) - 0.5
                if(norm == "log"):
                    n_ = ln(vmin=vmin, vmax=vmax)
Marco Selig's avatar
Marco Selig committed
871 872
                else:
                    n_ = None
873 874 875 876
                sub = ax0.pcolormesh(xaxes, yaxes, np.ma.masked_where(np.isnan(
                    xmesh), xmesh), cmap=cmap, norm=n_, vmin=vmin, vmax=vmax, clim=(vmin, vmax))
                ax0.set_xlim(xaxes[0], xaxes[-1])
                ax0.set_xticks([0], minor=False)
Marco Selig's avatar
Marco Selig committed
877
                ax0.set_xlabel(r"$\ell$")
878 879
                ax0.set_ylim(yaxes[0], yaxes[-1])
                ax0.set_yticks([0], minor=False)
Marco Selig's avatar
Marco Selig committed
880 881 882
                ax0.set_ylabel(r"$m$")
                ax0.set_aspect("equal")
                if(cbar):
883 884 885 886 887 888
                    if(norm == "log"):
                        f_ = lf(10, labelOnlyBase=False)
                        b_ = sub.norm.inverse(
                            np.linspace(0, 1, sub.cmap.N + 1))
                        v_ = np.linspace(
                            sub.norm.vmin, sub.norm.vmax, sub.cmap.N)
Marco Selig's avatar
Marco Selig committed
889 890 891 892
                    else:
                        f_ = None
                        b_ = None
                        v_ = None
893 894
                    fig.colorbar(sub, ax=ax0, orientation="horizontal", fraction=0.1, pad=0.05, shrink=0.75, aspect=20, ticks=[
                                 vmin, vmax], format=f_, drawedges=False, boundaries=b_, values=v_)
Marco Selig's avatar
Marco Selig committed
895 896
                ax0.set_title(title)

897 898 899
        if(bool(kwargs.get("save", False))):
            fig.savefig(str(kwargs.get("save")), dpi=None, facecolor="none", edgecolor="none", orientation="portrait",
                        papertype=None, format=None, transparent=False, bbox_inches=None, pad_inches=0.1)
Marco Selig's avatar
Marco Selig committed
900 901 902 903
            pl.close(fig)
        else:
            fig.canvas.draw()

904
    def getlm(self):  # > compute all (l,m)
905
        index = np.arange(self.dim)
906 907 908 909 910 911
        n = 2 * self.paradict['lmax'] + 1
        m = np.ceil(
            (n - np.sqrt(n**2 - 8 * (index - self.paradict['lmax']))) / 2
                    ).astype(np.int)
        l = index - self.paradict['lmax'] * m + m * (m - 1) // 2
        return l, m
Marco Selig's avatar
Marco Selig committed
912 913


914
class gl_space(Space):
Marco Selig's avatar
Marco Selig committed
915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931
    """
        ..                 __
        ..               /  /
        ..     ____ __  /  /
        ..   /   _   / /  /
        ..  /  /_/  / /  /_
        ..  \___   /  \___/  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``).
932
        dtype : numpy.dtype, *optional*
Marco Selig's avatar
Marco Selig committed
933 934 935 936 937 938 939 940 941 942
            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.
943
        `dtype` has to be either numpy.float64 or numpy.float32.
Marco Selig's avatar
Marco Selig committed
944 945 946 947 948 949 950 951 952 953 954 955 956 957

        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`.
958
        dtype : numpy.dtype
Marco Selig's avatar
Marco Selig committed
959 960 961 962 963 964 965
            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.
    """
966

csongor's avatar
csongor committed
967
    def __init__(self, nlat, nlon=None, dtype=np.dtype('float64')):
Marco Selig's avatar
Marco Selig committed
968 969 970 971 972 973 974 975 976
        """
            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``).
977
            dtype : numpy.dtype, *optional*
Marco Selig's avatar
Marco Selig committed
978 979 980 981 982 983 984 985 986 987 988 989 990 991
                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.

        """
992
        # check imports
Ultima's avatar
Ultima committed
993
        if not gc['use_libsharp']:
994
            raise ImportError(about._errors.cstring(
Ultima's avatar
Ultima committed
995
                "ERROR: libsharp_wrapper_gl not loaded."))
996

Ultima's avatar
Ultima committed
997
        self._cache_dict = {'check_codomain': {}}
998
        self.paradict = gl_space_paradict(nlat=nlat, nlon=nlon)
Marco Selig's avatar
Marco Selig committed
999

1000 1001 1002
        # check data type
        dtype = np.dtype(dtype)
        if dtype not in [np.dtype('float32'), np.dtype('float64')]:
Marco Selig's avatar
Marco Selig committed
1003
            about.warnings.cprint("WARNING: data type set to default.")
1004 1005
            dtype = np.dtype('float')
        self.dtype = dtype
1006

Marco Selig's avatar
Marco Selig committed
1007
        self.discrete = False
1008
        self.harmonic = False
csongor's avatar
csongor committed
1009
        self.distances = tuple(gl.vol(self.paradict['nlat'],
1010
                                      nlon=self.paradict['nlon']
csongor's avatar
csongor committed
1011
                                      ).astype(np.float))
1012 1013 1014

    @property
    def para(self):
1015
        temp = np.array([self.paradict['nlat'],
1016 1017
                         self.paradict['nlon']], dtype=int)
        return temp
1018

1019 1020 1021 1022
    @para.setter
    def para(self, x):
        self.paradict['nlat'] = x[0]
        self.paradict['nlon'] = x[1]
1023

1024
    def copy(self):
1025 1026 1027 1028
        return gl_space(nlat=self.paradict['nlat'],
                        nlon=self.paradict['nlon'],
                        dtype=self.dtype)

1029 1030
    @property
    def shape(self):
1031 1032
        return (np.int((self.paradict['nlat'] * self.paradict['nlon'])),)

1033 1034
    @property
    def meta_volume(self):
1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058
        """
            Calculates the meta volumes.

            The meta volumes are the volumes associated with each component of
            a field, taking into account field components that are not
            explicitly included in the array of field values but are determined
            by symmetry conditions.

            Parameters
            ----------
            total : bool, *optional*
                Whether to return the total meta volume of the space or the
                individual ones of each field component (default: False).

            Returns
            -------
            mol : {numpy.ndarray, float}
                Meta volume of the field components or the complete space.

            Notes
            -----
            For Gauss-Legendre pixelizations, the meta volumes are the pixel
            sizes.
        """