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

2import numpy as np 

3 

4 

5__all__ = ["get_ylm_inds", "get_ul_inds", "get_ylmw_inds", "integers"] 

6 

7 

8integers = (int, np.int, np.int16, np.int32, np.int64) 

9 

10 

11def get_ylm_inds(ydeg, ls, ms): 

12 """ 

13 

14 """ 

15 

16 # Turn `ls` and `ms` into slices 

17 if isinstance(ls, integers): 

18 ls = slice(ls, ls + 1) 

19 if isinstance(ms, integers): 

20 ms = slice(ms, ms + 1) 

21 

22 if isinstance(ls, slice) and isinstance(ms, slice): 

23 

24 # List of indices user is accessing 

25 inds = [] 

26 

27 # Fill in the `None`s 

28 if ls.start is None: 

29 ls = slice(0, ls.stop, ls.step) 

30 if ls.stop is None: 

31 ls = slice(ls.start, ydeg + 1, ls.step) 

32 if ls.step is None: 

33 ls = slice(ls.start, ls.stop, 1) 

34 if ms.step is None: 

35 ms = slice(ms.start, ms.stop, 1) 

36 

37 if (ls.start < 0) or (ls.start > ydeg): 

38 raise ValueError("Invalid value for `l`.") 

39 

40 # Loop through all the Ylms 

41 for l in range(ls.start, ls.stop, ls.step): 

42 ms_ = slice(ms.start, ms.stop, ms.step) 

43 if (ms_.start is None) or (ms_.start < -l): 

44 ms_ = slice(-l, ms_.stop, ms_.step) 

45 if (ms_.stop is None) or (ms_.stop > l): 

46 ms_ = slice(ms_.start, l + 1, ms_.step) 

47 for m in range(ms_.start, ms_.stop, ms_.step): 

48 n = l * l + l + m 

49 if ( 

50 (n < 0) or (n >= (ydeg + 1) ** 2) or (m > l) or (m < -l) 

51 ): # pragma: no cover 

52 raise ValueError("Invalid value for `l` and/or `m`.") 

53 inds.append(n) 

54 

55 return inds 

56 

57 else: 

58 

59 # Not a slice, not an int... What is it? 

60 raise ValueError("Invalid value for `l` and/or `m`.") 

61 

62 

63def get_ylmw_inds(ydeg, nw, ls, ms, ws): 

64 """ 

65 

66 """ 

67 

68 # Turn the `ws` into slices 

69 if isinstance(ws, integers): 

70 ws = slice(ws, ws + 1) 

71 

72 if isinstance(ws, slice): 

73 

74 # Get the `l`, `m` indices 

75 lminds = get_ylm_inds(ydeg, ls, ms) 

76 

77 # Process the `w` indices 

78 winds = [] 

79 

80 # Fill in the `None`s 

81 if ws.start is None: 

82 ws = slice(0, ws.stop, ws.step) 

83 if ws.stop is None: 

84 ws = slice(ws.start, nw, ws.step) 

85 if ws.step is None: 

86 ws = slice(ws.start, ws.stop, 1) 

87 

88 if (ws.start < 0) or (ws.start >= nw): 

89 raise ValueError("Invalid value for `w`.") 

90 

91 return tuple((lminds, ws)) 

92 

93 else: 

94 

95 # Not a slice, not an int... What is it? 

96 raise ValueError("Invalid value for `w`.") 

97 

98 

99def get_ul_inds(udeg, ls): 

100 """ 

101 

102 """ 

103 

104 # Turn `ls` into a slice 

105 if isinstance(ls, integers): 

106 ls = slice(ls, ls + 1) 

107 

108 if isinstance(ls, slice): 

109 

110 # List of indices user is accessing 

111 inds = [] 

112 

113 # Fill in the `None`s 

114 if ls.start is None: 

115 ls = slice(0, ls.stop, ls.step) 

116 if ls.stop is None: 

117 ls = slice(ls.start, udeg + 1, ls.step) 

118 if ls.step is None: 

119 ls = slice(ls.start, ls.stop, 1) 

120 

121 if (ls.start < 0) or (ls.start > udeg): 

122 raise ValueError("Invalid value for `l`.") 

123 

124 # Loop through all the `ls` 

125 for l in range(ls.start, ls.stop, ls.step): 

126 inds.append(l) 

127 

128 return inds 

129 

130 else: # pragma: no cover 

131 

132 # Not a slice, not an int... What is it? 

133 raise ValueError("Invalid value for `l`.")