Coverage for install/scipp/scipy/signal/__init__.py: 100%

1# SPDX-License-Identifier: BSD-3-Clause 

2# Copyright (c) 2023 Scipp contributors (https://github.com/scipp) 

3# @author Simon Heybrock 

4"""Sub-package for signal processing such as smoothing. 

5 

6This subpackage provides wrappers for a subset of functions from 

7:py:mod:`scipy.signal`. 

8""" 

9 

10from dataclasses import dataclass 

11from typing import Union 

12 

13from numpy import ndarray 

14 

15from ...compat.wrapping import wrap1d 

16from ...core import ( 

17 CoordError, 

18 DataArray, 

19 UnitError, 

20 Variable, 

21 array, 

22 identical, 

23 islinspace, 

24 to_unit, 

25) 

26from ...units import one 

27 

28 

29@dataclass 

30class SOS: 

31 """ 

32 Second-order sections representation returned by :py:func:`scipp.signal.butter`. 

33 

34 This is used as input to :py:func:`scipp.signal.sosfiltfilt`. 

35 """ 

36 

37 coord: Variable 

38 sos: ndarray 

39 

40 def filtfilt( 

41 self, obj: Union[Variable, DataArray], dim: str, **kwargs 

42 ) -> DataArray: 

43 """ 

44 Forwards to :py:func:`scipp.signal.sosfiltfilt` with sos argument set to the SOS 

45 instance. 

46 """ 

47 return sosfiltfilt(obj, dim=dim, sos=self, **kwargs) 

48 

49 

50def _frequency(coord: Variable) -> Variable: 

51 if not islinspace(coord).value: 

52 raise CoordError("Data is not regularly sampled, cannot compute frequency.") 

53 return (len(coord) - 1) / (coord[-1] - coord[0]) 

54 

55 

56def butter(coord: Variable, *, N: int, Wn: Variable, **kwargs) -> SOS: 

57 """ 

58 Butterworth digital and analog filter design. 

59 

60 Design an Nth-order digital or analog Butterworth filter and return the filter 

61 coefficients. 

62 

63 This is intended for use with :py:func:`scipp.scipy.signal.sosfiltfilt`. See there 

64 for an example. 

65 

66 This is a wrapper around :py:func:`scipy.signal.butter`. See there for a 

67 complete description of parameters. The differences are: 

68 

69 - Instead of a sampling frequency fs, this wrapper takes a variable ``coord`` as 

70 input. The sampling frequency is then computed from this coordinate. Only data 

71 sampled at regular intervals are supported. 

72 - The critical frequency or frequencies must be provided as a variable with correct 

73 unit. 

74 - Only 'sos' output is supported. 

75 

76 :seealso: :py:func:`scipp.scipy.signal.sosfiltfilt` 

77 """ 

78 fs = _frequency(coord).value 

79 try: 

80 Wn = to_unit(Wn, one / coord.unit) 

81 except UnitError: 

82 raise UnitError( 

83 f"Critical frequency unit '{Wn.unit}' incompatible with sampling unit " 

84 f"'{one / coord.unit}'" 

85 ) from None 

86 import scipy.signal 

87 

88 return SOS( 

89 coord=coord.copy(), 

90 sos=scipy.signal.butter(N=N, Wn=Wn.values, fs=fs, output='sos', **kwargs), 

91 ) 

92 

93 

94@wrap1d(keep_coords=True) 

95def _sosfiltfilt(da: DataArray, dim: str, *, sos: SOS, **kwargs) -> DataArray: 

96 if not identical(da.coords[dim], sos.coord): 

97 raise CoordError( 

98 f"Coord\n{da.coords[dim]}\nof filter dimension '{dim}' does " 

99 f"not match coord\n{sos.coord}\nused for creating the " 

100 "second-order sections representation by " 

101 "scipp.scipy.signal.butter." 

102 ) 

103 import scipy.signal 

104 

105 data = array( 

106 dims=da.dims, 

107 unit=da.unit, 

108 values=scipy.signal.sosfiltfilt(sos.sos, da.values, **kwargs), 

109 ) 

110 return DataArray(data=data) 

111 

112 

113def sosfiltfilt( 

114 obj: Union[Variable, DataArray], dim: str, *, sos: SOS, **kwargs 

115) -> DataArray: 

116 """ 

117 A forward-backward digital filter using cascaded second-order sections. 

118 

119 This is a wrapper around :py:func:`scipy.signal.sosfiltfilt`. See there for a 

120 complete description of parameters. The differences are: 

121 

122 - Instead of an array ``x`` and an optional axis, the input must be a data array 

123 (or variable) and dimension label. If it is a variable, the coord used for setting 

124 up the second-order sections is used as output coordinate. 

125 - The array of second-order filter coefficients ``sos`` includes the coordinate 

126 used for computing the frequency used for creating the coefficients. This is 

127 compared to the corresponding coordinate of the input data array to ensure that 

128 compatible coefficients are used. 

129 

130 Examples: 

131 

132 .. plot:: :context: close-figs 

133 

134 >>> from scipp.scipy.signal import butter, sosfiltfilt 

135 >>> x = sc.linspace(dim='x', start=1.1, stop=4.0, num=1000, unit='m') 

136 >>> y = sc.sin(x * sc.scalar(1.0, unit='rad/m')) 

137 >>> y += sc.sin(x * sc.scalar(400.0, unit='rad/m')) 

138 >>> da = sc.DataArray(data=y, coords={'x': x}) 

139 >>> sos = butter(da.coords['x'], N=4, Wn=20 / x.unit) 

140 >>> out = sosfiltfilt(da, 'x', sos=sos) 

141 >>> sc.plot({'input':da, 'sosfiltfilt':out}) 

142 

143 Instead of calling sosfiltfilt the more convenient filtfilt method of 

144 :py:class:`scipp.scipy.signal.SOS` can be used: 

145 

146 >>> out = butter(da.coords['x'], N=4, Wn=20 / x.unit).filtfilt(da, 'x') 

147 """ 

148 da = ( 

149 obj 

150 if isinstance(obj, DataArray) 

151 else DataArray(data=obj, coords={dim: sos.coord}) 

152 ) 

153 return _sosfiltfilt(da, dim=dim, sos=sos, **kwargs) 

154 

155 

156__all__ = ['butter', 'sosfiltfilt']