diff --git a/segpy-ext/segpy-numpy/segpy_numpy/extract.py b/segpy-ext/segpy-numpy/segpy_numpy/extract.py index e0504de..137acdb 100644 --- a/segpy-ext/segpy-numpy/segpy_numpy/extract.py +++ b/segpy-ext/segpy-numpy/segpy_numpy/extract.py @@ -9,7 +9,7 @@ class DimensionalityError: pass -def extract_inline_3d(reader, inline_number, xline_numbers=None, sample_numbers=None, null=None, packed=False): +def extract_inline_3d(reader, inline_number, xline_numbers=None, sample_numbers=None, null=None): """Extract an inline as a two-dimensional array. Args: @@ -42,37 +42,39 @@ def extract_inline_3d(reader, inline_number, xline_numbers=None, sample_numbers= sample_numbers = ensure_superset(range(0, reader.max_num_trace_samples()), sample_numbers) shape = (len(xline_numbers), len(sample_numbers)) dtype = make_dtype(reader.data_sample_format) - sample_start, sample_stop = start_and_stop(sample_numbers) array = make_array(shape, dtype, null) - src_start = sample_numbers.start - sample_start + if isinstance(sample_numbers, range): + _populate_inline_array_over_sample_range(reader, inline_number, xline_numbers, sample_numbers, array) + else: + _populate_inline_array_numbered_samples(reader, inline_number, xline_numbers, sample_numbers, array) - try: - src_step = sample_numbers.step - except AttributeError: - src_step = None + return array + +def _populate_inline_array_numbered_samples(reader, inline_number, xline_numbers, sample_numbers, array): for xline_index, xline_number in enumerate(xline_numbers): inline_xline_number = (inline_number, xline_number) if reader.has_trace_index(inline_xline_number): trace_index = reader.trace_index(inline_xline_number) num_trace_samples = reader.num_trace_samples(trace_index) - trace_sample_stop = max(sample_stop, num_trace_samples) - trace_samples = reader.trace_samples(trace_index, sample_start, trace_sample_stop) - src_stop = trace_sample_stop - sample_start + trace_sample_start = sample_numbers[0] + trace_sample_stop = min(sample_numbers[-1] + 1, num_trace_samples) + trace_samples = reader.trace_samples(trace_index, trace_sample_start, trace_sample_stop) + for sample_index, sample_number in enumerate(sample_numbers): + array[xline_index, sample_index] = trace_samples[sample_number - trace_sample_start] - if src_step is not None: - # Assign to a slice of the target array - source_slice = slice( - src_start, - src_stop, - src_step) - array[xline_index, :] = trace_samples[source_slice] - else: - # Assign element by element - for sample_index, sample_number in enumerate(sample_numbers): - array[xline_index, sample_index] = trace_samples[sample_number - sample_start] - return array + +def _populate_inline_array_over_sample_range(reader, inline_number, xline_numbers, sample_numbers, array): + for xline_index, xline_number in enumerate(xline_numbers): + inline_xline_number = (inline_number, xline_number) + if reader.has_trace_index(inline_xline_number): + trace_index = reader.trace_index(inline_xline_number) + num_trace_samples = reader.num_trace_samples(trace_index) + trace_sample_stop = min(sample_numbers.stop, num_trace_samples) + trace_samples = reader.trace_samples(trace_index, sample_numbers.start, trace_sample_stop) + source_slice = slice(sample_numbers.start, trace_sample_stop, sample_numbers.step) + array[xline_index, :] = trace_samples[source_slice] def make_array(shape, dtype, null=None):