Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/scipy/fft/_pocketfft/realtransforms.py: 34%

1import numpy as np 

2from . import pypocketfft as pfft 

3from .helper import (_asfarray, _init_nd_shape_and_axes, _datacopied, 

4 _fix_shape, _fix_shape_1d, _normalization, _workers) 

5import functools 

6 

7 

8def _r2r(forward, transform, x, type=2, n=None, axis=-1, norm=None, 

9 overwrite_x=False, workers=None, orthogonalize=None): 

10 """Forward or backward 1-D DCT/DST 

11 

12 Parameters 

13 ---------- 

14 forward : bool 

15 Transform direction (determines type and normalisation) 

16 transform : {pypocketfft.dct, pypocketfft.dst} 

17 The transform to perform 

18 """ 

19 tmp = _asfarray(x) 

20 overwrite_x = overwrite_x or _datacopied(tmp, x) 

21 norm = _normalization(norm, forward) 

22 workers = _workers(workers) 

23 

24 if not forward: 

25 if type == 2: 

26 type = 3 

27 elif type == 3: 

28 type = 2 

29 

30 if n is not None: 

31 tmp, copied = _fix_shape_1d(tmp, n, axis) 

32 overwrite_x = overwrite_x or copied 

33 elif tmp.shape[axis] < 1: 

34 raise ValueError("invalid number of data points ({0}) specified" 

35 .format(tmp.shape[axis])) 

36 

37 out = (tmp if overwrite_x else None) 

38 

39 # For complex input, transform real and imaginary components separably 

40 if np.iscomplexobj(x): 

41 out = np.empty_like(tmp) if out is None else out 

42 transform(tmp.real, type, (axis,), norm, out.real, workers) 

43 transform(tmp.imag, type, (axis,), norm, out.imag, workers) 

44 return out 

45 

46 return transform(tmp, type, (axis,), norm, out, workers, orthogonalize) 

47 

48 

49dct = functools.partial(_r2r, True, pfft.dct) 

50dct.__name__ = 'dct' 

51idct = functools.partial(_r2r, False, pfft.dct) 

52idct.__name__ = 'idct' 

53 

54dst = functools.partial(_r2r, True, pfft.dst) 

55dst.__name__ = 'dst' 

56idst = functools.partial(_r2r, False, pfft.dst) 

57idst.__name__ = 'idst' 

58 

59 

60def _r2rn(forward, transform, x, type=2, s=None, axes=None, norm=None, 

61 overwrite_x=False, workers=None, orthogonalize=None): 

62 """Forward or backward nd DCT/DST 

63 

64 Parameters 

65 ---------- 

66 forward : bool 

67 Transform direction (determines type and normalisation) 

68 transform : {pypocketfft.dct, pypocketfft.dst} 

69 The transform to perform 

70 """ 

71 tmp = _asfarray(x) 

72 

73 shape, axes = _init_nd_shape_and_axes(tmp, s, axes) 

74 overwrite_x = overwrite_x or _datacopied(tmp, x) 

75 

76 if len(axes) == 0: 

77 return x 

78 

79 tmp, copied = _fix_shape(tmp, shape, axes) 

80 overwrite_x = overwrite_x or copied 

81 

82 if not forward: 

83 if type == 2: 

84 type = 3 

85 elif type == 3: 

86 type = 2 

87 

88 norm = _normalization(norm, forward) 

89 workers = _workers(workers) 

90 out = (tmp if overwrite_x else None) 

91 

92 # For complex input, transform real and imaginary components separably 

93 if np.iscomplexobj(x): 

94 out = np.empty_like(tmp) if out is None else out 

95 transform(tmp.real, type, axes, norm, out.real, workers) 

96 transform(tmp.imag, type, axes, norm, out.imag, workers) 

97 return out 

98 

99 return transform(tmp, type, axes, norm, out, workers, orthogonalize) 

100 

101 

102dctn = functools.partial(_r2rn, True, pfft.dct) 

103dctn.__name__ = 'dctn' 

104idctn = functools.partial(_r2rn, False, pfft.dct) 

105idctn.__name__ = 'idctn' 

106 

107dstn = functools.partial(_r2rn, True, pfft.dst) 

108dstn.__name__ = 'dstn' 

109idstn = functools.partial(_r2rn, False, pfft.dst) 

110idstn.__name__ = 'idstn'