Coverage for mlprodict/onnxrt/ops_cpu/op_window.py: 100%

1# -*- encoding: utf-8 -*- 

2# pylint: disable=E0203,E1101,C0111 

3""" 

4@file 

5@brief Runtime operator. 

6""" 

7import numpy 

8from onnx.onnx_pb import TensorProto # pylint: disable=E0611 

9from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE 

10from ._op import OpRun 

11 

12 

13class _CommonWindow: 

14 

15 def _begin(self, size): 

16 if self.periodic == 1: 

17 N_1 = size 

18 else: 

19 N_1 = size - 1 

20 ni = numpy.arange(size, dtype=self.dtype) 

21 return ni, N_1 

22 

23 def _end(self, size, res): 

24 return (res.astype(self.dtype), ) 

25 

26 

27class BlackmanWindow(OpRun, _CommonWindow): 

28 """ 

29 Returns 

30 :math:`\\omega_n = 0.42 - 0.5 \\cos \\left( \\frac{2\\pi n}{N-1} \\right) + 

31 0.08 \\cos \\left( \\frac{4\\pi n}{N-1} \\right)` 

32 where *N* is the window length. 

33 See `blackman_window 

34 <https://pytorch.org/docs/stable/generated/torch.blackman_window.html>`_ 

35 """ 

36 

37 atts = {'output_datatype': TensorProto.FLOAT, 'periodic': 1} 

38 

39 def __init__(self, onnx_node, desc=None, **options): 

40 OpRun.__init__(self, onnx_node, desc=desc, 

41 expected_attributes=BlackmanWindow.atts, 

42 **options) 

43 self.dtype = TENSOR_TYPE_TO_NP_TYPE[self.output_datatype] 

44 

45 def _run(self, size, attributes=None, verbose=0, fLOG=None): # pylint: disable=W0221 

46 # ni, N_1 = self._begin(size) 

47 ni, N_1 = numpy.arange(size, dtype=self.dtype), size 

48 alpha = 0.42 

49 beta = 0.08 

50 pi = 3.1415 

51 y = alpha 

52 y -= numpy.cos((ni * (pi * 2)) / N_1) / 2 

53 y += numpy.cos((ni * (pi * 4)) / N_1) * beta 

54 return (self._end(size, y), ) 

55 

56 

57class HannWindow(OpRun, _CommonWindow): 

58 """ 

59 Returns 

60 :math:`\\omega_n = \\sin^2\\left( \\frac{\\pi n}{N-1} \\right)` 

61 where *N* is the window length. 

62 See `hann_window 

63 <https://pytorch.org/docs/stable/generated/torch.hann_window.html>`_ 

64 """ 

65 

66 atts = {'output_datatype': TensorProto.FLOAT, 'periodic': 1} 

67 

68 def __init__(self, onnx_node, desc=None, **options): 

69 OpRun.__init__(self, onnx_node, desc=desc, 

70 expected_attributes=HannWindow.atts, 

71 **options) 

72 self.dtype = TENSOR_TYPE_TO_NP_TYPE[self.output_datatype] 

73 

74 def _run(self, size, attributes=None, verbose=0, fLOG=None): # pylint: disable=W0221 

75 ni, N_1 = self._begin(size) 

76 res = numpy.sin(ni * 3.1415 / N_1) ** 2 

77 return self._end(size, res) 

78 

79 

80class HammingWindow(OpRun, _CommonWindow): 

81 """ 

82 Returns 

83 :math:`\\omega_n = \\alpha - \\beta \\cos \\left( \\frac{\\pi n}{N-1} \\right)` 

84 where *N* is the window length. 

85 See `hamming_window 

86 <https://pytorch.org/docs/stable/generated/torch.hamming_window.html>`_. 

87 `alpha=0.54, beta=0.46` 

88 """ 

89 

90 atts = {'output_datatype': TensorProto.FLOAT, 'periodic': 1} 

91 

92 def __init__(self, onnx_node, desc=None, **options): 

93 OpRun.__init__(self, onnx_node, desc=desc, 

94 expected_attributes=HammingWindow.atts, 

95 **options) 

96 self.dtype = TENSOR_TYPE_TO_NP_TYPE[self.output_datatype] 

97 

98 def _run(self, size, attributes=None, verbose=0, fLOG=None): # pylint: disable=W0221 

99 ni, N_1 = self._begin(size) 

100 alpha = 25. / 46. 

101 beta = 1 - alpha 

102 res = alpha - numpy.cos(ni * 3.1415 * 2 / N_1) * beta 

103 return self._end(size, res)