glmollweide.py 1.68 KB
Newer Older
1 2 3 4 5 6
# -*- coding: utf-8 -*-

from nifty import dependency_injector as gdi
from heatmap import Heatmap
import numpy as np

7 8
from .mollweide_helper import mollweide_helper

9 10 11 12
pyHealpix = gdi.get('pyHealpix')


class GLMollweide(Heatmap):
Theo Steininger's avatar
Theo Steininger committed
13 14 15
    def __init__(self, data, xsize=800, color_map=None,
                 webgl=False, smoothing=False):
        # smoothing 'best', 'fast', False
16
        if pyHealpix is None:
17 18
            raise ImportError(
                "The module pyHealpix is needed but not available.")
Theo Steininger's avatar
Theo Steininger committed
19 20 21 22 23
        self.xsize = xsize

        super(GLMollweide, self).__init__(data, color_map, webgl, smoothing)

    def at(self, data):
24 25 26
        if isinstance(data, list):
            data = [self._mollview(d) for d in data]
        else:
Theo Steininger's avatar
Theo Steininger committed
27 28 29 30 31 32
            data = self._mollview(data)
        return GLMollweide(data=data,
                           xsize=self.xsize,
                           color_map=self.color_map,
                           webgl=self.webgl,
                           smoothing=self.smoothing)
33

Theo Steininger's avatar
Theo Steininger committed
34
    @staticmethod
35 36 37 38 39 40 41 42 43
    def _find_closest(A, target):
        # A must be sorted
        idx = A.searchsorted(target)
        idx = np.clip(idx, 1, len(A)-1)
        left = A[idx-1]
        right = A[idx]
        idx -= target - left < right - target
        return idx

Theo Steininger's avatar
Theo Steininger committed
44 45 46 47 48
    def _mollview(self, x):
        xsize = self.xsize
        nlat = x.shape[0]
        nlon = x.shape[1]

49
        res, mask, theta, phi = mollweide_helper(xsize)
Martin Reinecke's avatar
Martin Reinecke committed
50 51

        ra = np.linspace(0, 2*np.pi, nlon+1)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
52 53
        dec = pyHealpix.GL_thetas(nlat)
        ilat = self._find_closest(dec, theta)
Martin Reinecke's avatar
Martin Reinecke committed
54
        ilon = self._find_closest(ra, phi)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
55 56
        ilon = np.where(ilon == nlon, 0, ilon)
        res[mask] = x[ilat, ilon]
Martin Reinecke's avatar
Martin Reinecke committed
57
        return res