Coverage for python/lsst/shapelet/generator.py: 31%

1from .constants import HERMITE, LAGUERRE, computeSize 

2 

3__all__ = ['IndexGenerator', 'HermiteIndexGenerator', 'LaguerreIndexGenerator'] 

4 

5 

6class IndexGenerator: 

7 """Base class for shapelet index generators. 

8 

9 Parameters 

10 ---------- 

11 order: `int` 

12 Order of the function indices to generate. 

13 """ 

14 

15 __slots__ = "order", "size" 

16 

17 def __init__(self, order): 

18 self.order = order 

19 self.size = computeSize(self.order) 

20 

21 @staticmethod 

22 def make(self, order, basisType): 

23 if basisType == HERMITE: 

24 return HermiteIndexGenerator(order) 

25 elif basisType == LAGUERRE: 

26 return LaguerreIndexGenerator(order) 

27 

28 def __len__(self): 

29 return self.size 

30 

31 

32class HermiteIndexGenerator(IndexGenerator): 

33 """Iterable that generates tuples of indices. 

34 

35 Yields 

36 ------ 

37 i : `int` 

38 the overall coefficient index for a 2-d shapelet expansion (just counts from zero) 

39 nx : `int` 

40 the order of the x expansion 

41 ny : `int` 

42 the order of the y expansion 

43 """ 

44 

45 def __iter__(self): 

46 i = 0 

47 for n in range(0, self.order+1): 

48 for nx in range(0, n+1): 

49 yield (i, nx, n - nx) 

50 i += 1 

51 

52 

53class LaguerreIndexGenerator(IndexGenerator): 

54 """Iterable that generates tuples of indices. 

55 

56 Yields 

57 ------ 

58 i : `int` 

59 the overall coefficient index for a 2-d shapelet expansion (just counts from zero) 

60 p, q : `int` 

61 indices of the polar shapelet expansion (see BasisTypeEnum). 

62 re : `bool` 

63 True, if this is the real part of the coefficient. 

64 """ 

65 

66 def __iter__(self): 

67 i = 0 

68 for n in range(0, self.order+1): 

69 p = n 

70 q = 0 

71 while p > q: 

72 yield (i, p, q, True) 

73 i += 1 

74 yield (i, p, q, False) 

75 i += 1 

76 p -= 1 

77 q += 1 

78 if p == q: 

79 yield (i, p, q, True) 

80 i += 1