diff --git a/segpy-ext/segpy-numpy/segpy_numpy/extract.py b/segpy-ext/segpy-numpy/segpy_numpy/extract.py index 052f6d8..e0504de 100644 --- a/segpy-ext/segpy-numpy/segpy_numpy/extract.py +++ b/segpy-ext/segpy-numpy/segpy_numpy/extract.py @@ -1,8 +1,7 @@ """Tools for interoperability between Segpy and Numpy arrays.""" -from functools import singledispatch import numpy as np -from segpy.util import make_sorted_distinct_sequence +from segpy.util import ensure_superset from segpy_numpy.dtypes import make_dtype @@ -18,9 +17,10 @@ def extract_inline_3d(reader, inline_number, xline_numbers=None, sample_numbers= inline_number: The number of the inline to be extracted. - xline_numbers: An optional sorted sequence of crossline numbers at which to - extract samples. If not provided, samples will be extracted at all - crosslines. + xline_numbers: Either a sequence of crossline numbers from which traces are + to be extracted or a slice object indicating that some slice of all + crossline numbers is to be used. If None, traces will be extracted at + all crosslines. sample_numbers: An optional sorted sequence of samples numbers at which to extract samples. If not provided, samples will be extracted at all @@ -38,60 +38,58 @@ def extract_inline_3d(reader, inline_number, xline_numbers=None, sample_numbers= if inline_number not in reader.inline_numbers(): raise ValueError("Inline number {} not present in {}".format(inline_number, reader)) - if xline_numbers is None: - xline_numbers = reader.xline_numbers() - elif isinstance(xline_numbers, slice): - xline_numbers = reader.xline_numbers()[xline_numbers] - else: - xline_numbers = make_sorted_distinct_sequence(xline_numbers) - - if sample_numbers is None: - sample_numbers = range(0, reader.max_num_trace_samples()) - elif isinstance(sample_numbers, slice): - sample_numbers = range(0, reader.max_num_trace_samples())[sample_numbers] - else: - sample_numbers = make_sorted_distinct_sequence(sample_numbers) - + xline_numbers = ensure_superset(reader.xline_numbers(), xline_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) - if null is None: - array = np.ma.masked_all(shape, dtype) - else: - array = np.empty(shape, dtype) - array.fill(null) + src_start = sample_numbers.start - sample_start + + try: + src_step = sample_numbers.step + except AttributeError: + src_step = None for xline_index, xline_number in enumerate(xline_numbers): inline_xline_number = (inline_number, xline_number) if reader.has_trace_index(inline_xline_number): - # Read the trace trace_index = reader.trace_index(inline_xline_number) - - sample_start = sample_numbers[0] - try: - sample_stop = sample_numbers.stop - except AttributeError: - sample_stop = sample_numbers[-1] + 1 - num_trace_samples = reader.num_trace_samples(trace_index) - sample_stop = max(sample_stop, num_trace_samples) + 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_samples = reader.trace_samples(trace_index, sample_start, sample_stop) - - try: - sample_step = sample_numbers.step - except AttributeError: - sample_step = None - - if sample_step is not None: + if src_step is not None: # Assign to a slice of the target array source_slice = slice( - sample_numbers.start - sample_start, - sample_numbers.stop - sample_start, - sample_step) + 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_number, sample_index] = trace_samples[sample_number - sample_start] - return array \ No newline at end of file + array[xline_index, sample_index] = trace_samples[sample_number - sample_start] + return array + + +def make_array(shape, dtype, null=None): + """Make an array""" + if null is None: + return np.ma.masked_all(shape, dtype) + + array = np.empty(shape, dtype) + array.fill(null) + return array + + +def start_and_stop(sequence): + """Obtain start and stop values from a sequence.""" + sample_start = sequence[0] + try: + sample_stop = sequence.stop + except AttributeError: + sample_stop = sequence[-1] + 1 + return sample_start, sample_stop diff --git a/segpy/sorted_set.py b/segpy/sorted_set.py index 507ae93..2b406f0 100644 --- a/segpy/sorted_set.py +++ b/segpy/sorted_set.py @@ -6,8 +6,11 @@ from itertools import chain class SortedFrozenSet(Sequence, Set): - def __init__(self, items=None): - self._items = sorted(set(items)) if items is not None else [] + def __new__(cls, items=None): + if type(items) == cls: + return items + obj = object.__new__(cls) + obj._items = sorted(set(items)) if items is not None else [] def __contains__(self, item): try: diff --git a/segpy/util.py b/segpy/util.py index f30975a..b8ca65d 100644 --- a/segpy/util.py +++ b/segpy/util.py @@ -391,3 +391,53 @@ def hash_for_file(fh, *args): sha1.update(encoded_arg) digest = sha1.hexdigest() return digest + + +def is_range_superset_of_range(superset_range, subset_range): + """Are all the elements of + + """ + if subset_range.start not in superset_range: + return False + if subset_range.step % superset_range.step != 0: + return False + if subset_range[-1] > superset_range[-1]: + return False + assert set(subset_range).issubset(set(superset_range)) + return True + + +def is_subset(superset, subset): + """A more general version of set.issubset that is smart enough to work with ranges.""" + if isinstance(subset, range) and isinstance(superset, range): + return is_range_superset_of_range(superset, subset) + if isinstance(superset, range): + return all(item in superset for item in subset) + return set(superset).issuperset(subset) + + +def ensure_superset(superset, subset): + """Obtains a subset of all items. + + Args: + all_items: A sequence containing all items. + + subset: Subset must either be a collection the elements of which are a subset of + all_items, or a slice object, in which case the subset items will be sliced + from all_items. + Returns: + A sorted, distinct collection which is a subset of all_items. + + Raises: + ValueError: If the items in subset are not a subset of the items in all_items. + """ + if subset is None: + return superset + elif isinstance(subset, slice): + return superset[subset] + else: + subset = make_sorted_distinct_sequence(subset) + if not is_subset(superset, subset): + raise ValueError("subset_or_slice {!r} is not a subset of all_items {!r}" + .format(subset, superset)) + return subset \ No newline at end of file