Adds an extract_inline_3d() function to segpy_numpy

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
+                    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

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:

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