1# -*- coding: utf-8 -*- 

2"""Healpy-based spherical harmonic transform utilities for starry.""" 

3import numpy as np 

4from PIL import Image 

5from matplotlib.image import pil_to_array 

6import os 

7from scipy import ndimage 

8 

9try: 

10 import healpy as hp 

11except ImportError: 

12 hp = None 

13 

14__all__ = ["image2map", "healpix2map", "array2map", "array2healpix"] 

15 

16 

17def healpix2map(healpix_map, lmax=10, **kwargs): 

18 """Return a map vector corresponding to a healpix array.""" 

19 if hp is None: 

20 raise ImportError( 

21 "Please install the `healpy` Python package to " 

22 "enable this feature. See `https://healpy.readthedocs.io`." 

23 ) 

24 # Get the complex spherical harmonic coefficients 

25 alm = hp.sphtfunc.map2alm(healpix_map, lmax=lmax) 

26 

27 # We first need to do a rotation to get our axes aligned correctly, 

28 # since we use a different convention than `healpy` 

29 alm = hp.rotator.Rotator((-90, 0, -90)).rotate_alm(alm) 

30 

31 # Smooth the map? 

32 sigma = kwargs.pop("sigma", None) 

33 if sigma is not None: 

34 alm = hp.sphtfunc.smoothalm(alm, sigma=sigma, verbose=False) 

35 

36 # Convert them to real coefficients 

37 ylm = np.zeros(lmax ** 2 + 2 * lmax + 1, dtype="float") 

38 i = 0 

39 for l in range(0, lmax + 1): 

40 for m in range(-l, l + 1): 

41 j = hp.sphtfunc.Alm.getidx(lmax, l, np.abs(m)) 

42 if m < 0: 

43 ylm[i] = np.sqrt(2) * (-1) ** m * alm[j].imag 

44 elif m == 0: 

45 ylm[i] = alm[j].real 

46 else: 

47 ylm[i] = np.sqrt(2) * (-1) ** m * alm[j].real 

48 i += 1 

49 

50 return ylm 

51 

52 

53def image2map(image, **kwargs): 

54 """Return a map vector corresponding to a lat-long map image.""" 

55 # If image doesn't exist, check for it in `img` directory 

56 if not os.path.exists(image): 

57 dn = os.path.dirname 

58 image = os.path.join(dn(os.path.abspath(__file__)), "img", image) 

59 if not image.endswith(".jpg"): 

60 image += ".jpg" 

61 if not os.path.exists(image): 

62 raise ValueError("File not found: %s." % image) 

63 

64 # Get the image array 

65 grayscale_pil_image = Image.open(image).convert("L") 

66 image_array = pil_to_array(grayscale_pil_image) 

67 image_array = np.array(image_array, dtype=float) 

68 image_array /= 255.0 

69 

70 # Convert it to a map 

71 return array2map(image_array, **kwargs) 

72 

73 

74def array2healpix(image_array, nside=16, max_iter=3, **kwargs): 

75 """Return a healpix ring-ordered map corresponding to a lat-lon map image array.""" 

76 if hp is None: 

77 raise ImportError( 

78 "Please install the `healpy` Python package to " 

79 "enable this feature. See `https://healpy.readthedocs.io`." 

80 ) 

81 

82 # Starting value for the zoom 

83 zoom = 2 

84 

85 # Keep track of the number of unseen pixels 

86 unseen = 1 

87 ntries = 0 

88 while unseen > 0: 

89 

90 # Make the image bigger so we have good angular coverage 

91 image_array = ndimage.zoom(image_array, zoom) 

92 

93 # Convert to a healpix map 

94 theta = np.linspace(0, np.pi, image_array.shape[0])[:, None] 

95 phi = np.linspace(-np.pi, np.pi, image_array.shape[1])[::-1] 

96 pix = hp.ang2pix(nside, theta, phi, nest=False) 

97 healpix_map = ( 

98 np.ones(hp.nside2npix(nside), dtype=np.float64) * hp.UNSEEN 

99 ) 

100 healpix_map[pix] = image_array 

101 

102 # Count the unseen pixels 

103 unseen = np.count_nonzero(healpix_map == hp.UNSEEN) 

104 

105 # Did we do this too many times? 

106 ntries += 1 

107 if ntries > max_iter: 

108 raise ValueError( 

109 "Maximum number of iterations exceeded. Either decreaser `nside` or increase `max_iter`." 

110 ) 

111 

112 return healpix_map 

113 

114 

115def array2map(image_array, **kwargs): 

116 """Return a map vector corresponding to a lat-lon map image array.""" 

117 # Get the healpix map 

118 healpix_map = array2healpix(image_array, **kwargs) 

119 

120 # Now convert it to a spherical harmonic map 

121 return healpix2map(healpix_map, **kwargs)