[segyio] 215/376: tools.collect and tools.cube

[Jørgen Kvalsvik]

This is an automated email from the git hooks/post-receive script.

jokva-guest pushed a commit to branch debian
in repository segyio.

commit af7393258787de95bd43e0765c408b73250ef1d6
Author: Jørgen Kvalsvik <jokva at statoil.com>
Date:   Tue Feb 28 10:25:50 2017 +0100

    tools.collect and tools.cube
    
    Collect handles the common case of copying a series of traces, lines,
    slices into one stacked numpy array.
    
    Cube attempts to read a full cube from a file or file handle and returns
    the data as a 3D numpy array.
---
 python/segyio/__init__.py |  1 +
 python/segyio/tools.py    | 44 ++++++++++++++++++++++++++++++++++++++++++++
 python/test/tools.py      | 19 +++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/python/segyio/__init__.py b/python/segyio/__init__.py
index e3599ef..c52b93f 100644
--- a/python/segyio/__init__.py
+++ b/python/segyio/__init__.py
@@ -87,6 +87,7 @@ from .open import open
 from .create import create
 from .segy import SegyFile, spec
 from .tools import dt, sample_indexes, create_text_header, native
+from .tools import collect, cube
 
 __version__    = '1.0.4'
 __copyright__  = 'Copyright 2016, Statoil ASA'
diff --git a/python/segyio/tools.py b/python/segyio/tools.py
index 348c246..2b51679 100644
--- a/python/segyio/tools.py
+++ b/python/segyio/tools.py
@@ -76,3 +76,47 @@ def native(data,
 
     format = int(segyio.SegySampleFormat(format))
     return segyio._segyio.native(data, format)
+
+def collect(itr):
+    """ Collect traces or lines into one ndarray
+
+    Eagerly copy a series of traces, lines or depths into one numpy ndarray. If
+    collecting traces or fast-direction over a post-stacked file, reshaping the
+    resulting array is equivalent to calling `tools.cube`.
+
+    Examples:
+
+    collect-cube identity::
+        >>> f = segyio.open('post-stack.sgy')
+        >>> x = segyio.tools.collect(f.traces[:])
+        >>> x = x.reshape((len(f.ilines), len(f.xlines), f.samples))
+        >>> numpy.all(x == segyio.tools.cube(f))
+
+    :type itr: iterable[numpy.ndarray]
+    :rtype: numpy.ndarray
+    """
+    return np.stack([np.copy(x) for x in itr])
+
+def cube(f):
+    """ Read a full cube from a file
+
+    Takes an open segy file (created with segyio.open) or a file name.
+
+    If the file is a prestack file, the cube returned has the dimensions
+    (fast,slow,offset,sample). If it is post-stack (i.e. only the one offset),
+    the dimensions are normalised to (fast,slow,sample)
+
+    :type f: SegyFile|str
+    :rtype numpy.ndarray
+    """
+
+    if not isinstance(f, segyio.SegyFile):
+        with segyio.open(f) as fl:
+            return cube(fl)
+
+    ilsort = f.sorting == segyio.TraceSortingFormat.INLINE_SORTING
+    fast = f.ilines if ilsort else f.xlines
+    slow = f.xlines if ilsort else f.ilines
+    fast, slow, offs, samples = len(fast), len(slow), len(f.offsets), f.samples
+    dims = (fast, slow, samples) if offs == 1 else (fast, slow, offs, samples)
+    return f.trace.raw[:].reshape(dims)
diff --git a/python/test/tools.py b/python/test/tools.py
index faba4f3..5c2a7b7 100644
--- a/python/test/tools.py
+++ b/python/test/tools.py
@@ -10,6 +10,7 @@ import segyio
 class ToolsTest(TestCase):
     def setUp(self):
         self.filename = "test-data/small.sgy"
+        self.prestack = "test-data/small-ps.sgy"
 
     def test_dt_fallback(self):
         with TestContext("dt_fallback") as context:
@@ -74,3 +75,21 @@ class ToolsTest(TestCase):
             filetr = np.frombuffer(filetr, dtype = np.single)
             self.assertFalse(np.array_equal(segytr, filetr))
             self.assertTrue(np.array_equal(segytr, segyio.tools.native(filetr)))
+
+    def test_cube_filename(self):
+        with segyio.open(self.filename) as f:
+            c1 = segyio.tools.cube(f)
+            c2 = segyio.tools.cube(self.filename)
+            self.assertTrue(np.all(c1 == c2))
+
+    def test_cube_identity(self):
+        with segyio.open(self.filename) as f:
+            x = segyio.tools.collect(f.trace[:])
+            x = x.reshape((len(f.ilines), len(f.xlines), f.samples))
+            self.assertTrue(np.all(x == segyio.tools.cube(f)))
+
+    def test_cube_identity_prestack(self):
+        with segyio.open(self.prestack) as f:
+            dims = (len(f.ilines), len(f.xlines), len(f.offsets), f.samples)
+            x = segyio.tools.collect(f.trace[:]).reshape(dims)
+            self.assertTrue(np.all(x == segyio.tools.cube(f)))

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-science/packages/segyio.git