Coverage for src / tracekit / search / context.py: 100%

1"""Context extraction around points of interest. 

2 

3 

4This module provides efficient extraction of signal context around 

5events, maintaining original time references for debugging workflows. 

6""" 

7 

8from typing import Any 

9 

10import numpy as np 

11from numpy.typing import NDArray 

12 

13 

14def extract_context( 

15 trace: NDArray[np.float64], 

16 index: int | list[int] | NDArray[np.int_], 

17 *, 

18 before: int = 100, 

19 after: int = 100, 

20 sample_rate: float | None = None, 

21 include_metadata: bool = True, 

22) -> dict[str, Any] | list[dict[str, Any]]: 

23 """Extract signal context around a point of interest. 

24 

25 : Context extraction with time reference preservation. 

26 Supports batch extraction for multiple indices and optional protocol data. 

27 

28 Args: 

29 trace: Input signal trace 

30 index: Sample index or list of indices to extract context around. 

31 Can be int, list of ints, or numpy array. 

32 before: Number of samples to include before index (default: 100) 

33 after: Number of samples to include after index (default: 100) 

34 sample_rate: Optional sample rate in Hz for time calculations 

35 include_metadata: Include metadata dict with context info (default: True) 

36 

37 Returns: 

38 If index is scalar: Single context dictionary 

39 If index is list/array: List of context dictionaries 

40 

41 Each context dictionary contains: 

42 - data: Extracted sub-trace array 

43 - start_index: Starting index in original trace 

44 - end_index: Ending index in original trace 

45 - center_index: Center index (original query index) 

46 - time_reference: Time offset if sample_rate provided 

47 - length: Number of samples in context 

48 

49 Raises: 

50 ValueError: If index is out of bounds 

51 ValueError: If before or after are negative 

52 

53 Examples: 

54 >>> # Extract context around a glitch 

55 >>> trace = np.random.randn(1000) 

56 >>> glitch_index = 500 

57 >>> context = extract_context( 

58 ... trace, 

59 ... glitch_index, 

60 ... before=50, 

61 ... after=50, 

62 ... sample_rate=1e6 

63 ... ) 

64 >>> print(f"Context length: {len(context['data'])}") 

65 >>> print(f"Time reference: {context['time_reference']*1e6:.2f} µs") 

66 

67 >>> # Batch extraction for multiple events 

68 >>> event_indices = [100, 200, 300] 

69 >>> contexts = extract_context( 

70 ... trace, 

71 ... event_indices, 

72 ... before=25, 

73 ... after=25 

74 ... ) 

75 >>> print(f"Extracted {len(contexts)} contexts") 

76 

77 Notes: 

78 - Handles edge cases at trace boundaries automatically 

79 - Context may be shorter than before+after at boundaries 

80 - Time reference is relative to start of extracted context 

81 - Original trace is not modified 

82 

83 References: 

84 SRCH-003: Context Extraction 

85 """ 

86 if before < 0 or after < 0: 

87 raise ValueError("before and after must be non-negative") 

88 

89 if trace.size == 0: 

90 raise ValueError("Trace cannot be empty") 

91 

92 # Handle single index vs multiple indices 

93 if isinstance(index, int | np.integer): 

94 indices = [int(index)] 

95 return_single = True 

96 else: 

97 indices = [int(i) for i in index] 

98 return_single = False 

99 

100 # Validate indices 

101 for idx in indices: 

102 if idx < 0 or idx >= len(trace): 

103 raise ValueError(f"Index {idx} out of bounds for trace of length {len(trace)}") 

104 

105 # Extract contexts 

106 contexts = [] 

107 

108 for idx in indices: 

109 # Calculate window bounds with boundary handling 

110 start_idx = max(0, idx - before) 

111 end_idx = min(len(trace), idx + after + 1) 

112 

113 # Extract data 

114 data = trace[start_idx:end_idx].copy() 

115 

116 # Build context dictionary 

117 context: dict[str, Any] = { 

118 "data": data, 

119 "start_index": start_idx, 

120 "end_index": end_idx, 

121 "center_index": idx, 

122 "length": len(data), 

123 } 

124 

125 # Add time reference if sample rate provided 

126 if sample_rate is not None: 

127 time_offset = start_idx / sample_rate 

128 context["time_reference"] = time_offset 

129 context["sample_rate"] = sample_rate 

130 

131 # Time array for the context 

132 dt = 1.0 / sample_rate 

133 context["time_array"] = np.arange(len(data)) * dt + time_offset 

134 

135 if include_metadata: 

136 context["metadata"] = { 

137 "samples_before": idx - start_idx, 

138 "samples_after": end_idx - idx - 1, 

139 "at_start_boundary": start_idx == 0, 

140 "at_end_boundary": end_idx == len(trace), 

141 } 

142 

143 contexts.append(context) 

144 

145 # Return single context or list 

146 if return_single: 

147 return contexts[0] 

148 else: 

149 return contexts