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 Any 

12 

13import numpy.typing as npt 

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: npt.NDArray[Any] 

39 

40 def filtfilt(self, obj: Variable | DataArray, dim: str, **kwargs: Any) -> DataArray: 

41 """ 

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

43 instance. 

44 """ 

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

46 

47 

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

49 if not islinspace(coord).value: 

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

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

52 

53 

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

55 """ 

56 Butterworth digital and analog filter design. 

57 

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

59 coefficients. 

60 

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

62 for an example. 

63 

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

65 complete description of parameters. The differences are: 

66 

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

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

69 sampled at regular intervals are supported. 

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

71 unit. 

72 - Only 'sos' output is supported. 

73 

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

75 """ 

76 fs = _frequency(coord).value 

77 try: 

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

79 except UnitError: 

80 raise UnitError( 

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

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

83 ) from None 

84 import scipy.signal 

85 

86 return SOS( 

87 coord=coord.copy(), 

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

89 ) 

90 

91 

92@wrap1d(keep_coords=True) 

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

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

95 raise CoordError( 

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

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

98 "second-order sections representation by " 

99 "scipp.scipy.signal.butter." 

100 ) 

101 import scipy.signal 

102 

103 data = array( 

104 dims=da.dims, 

105 unit=da.unit, 

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

107 ) 

108 return DataArray(data=data) 

109 

110 

111def sosfiltfilt( 

112 obj: Variable | DataArray, dim: str, *, sos: SOS, **kwargs: Any 

113) -> DataArray: 

114 """ 

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

116 

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

118 complete description of parameters. The differences are: 

119 

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

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

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

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

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

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

126 compatible coefficients are used. 

127 

128 Examples: 

129 

130 .. plot:: :context: close-figs 

131 

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

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

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

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

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

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

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

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

140 

141 Instead of calling sosfiltfilt the more convenient filtfilt method of 

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

143 

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

145 """ 

146 da = ( 

147 obj 

148 if isinstance(obj, DataArray) 

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

150 ) 

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

152 

153 

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