Put the windowing (taper) in tools.py in get_chunk_with_margin.
@@ -6,7 +6,8 @@ import numpy as np
def get_chunk_with_margin(rec_segment, start_frame, end_frame,
- channel_indices, margin, add_zeros=False):
+ channel_indices, margin, add_zeros=False,
+ window_on_margin=False, dtype=None):
"""
Helper to get chunk with margin
"""
@@ -16,6 +17,7 @@ def get_chunk_with_margin(rec_segment, start_frame, end_frame,
channel_indices = slice(None)
if not add_zeros:
+ assert not window_on_margin, 'window_mon_margin can be used only for add_zeros=True'
if start_frame is None:
left_margin = 0
start_frame = 0
@@ -56,8 +58,16 @@ def get_chunk_with_margin(rec_segment, start_frame, end_frame,
traces_chunk = rec_segment.get_traces(start_frame2, end_frame2, channel_indices)
- if left_pad > 0 or right_pad > 0:
- traces_chunk2 = np.zeros((full_size, traces_chunk.shape[1]), dtype=traces_chunk.dtype)
+
+ if dtype is not None or window_on_margin or left_pad > 0 or right_pad > 0:
+ need_copy = True
+ else:
+ need_copy = False
+
+ if need_copy:
+ if dtype is None:
+ dtype = traces_chunk.dtype
+ traces_chunk2 = np.zeros((full_size, traces_chunk.shape[1]), dtype=dtype)
i0 = left_pad
i1 = left_pad + traces_chunk.shape[0]
traces_chunk2[i0: i1, :] = traces_chunk
@@ -66,9 +76,16 @@ def get_chunk_with_margin(rec_segment, start_frame, end_frame,
right_margin = margin
else:
right_margin = end_frame + margin - length
+ if window_on_margin:
+ # apply inplace taper on border
+ taper = (1 - np.cos(np.arange(margin) / margin * np.pi)) / 2
+ taper = taper[:, np.newaxis]
+ traces_chunk2[:margin] *= taper
+ traces_chunk2[-margin:] *= taper[::-1]
traces_chunk = traces_chunk2
else:
left_margin = margin
right_margin = margin
return traces_chunk, left_margin, right_margin
+
