# -*- coding: utf-8 -*-
"""
The obspy.segy Seismic Unix test suite.
"""
from obspy.core.util import NamedTemporaryFile
from obspy.segy.segy import readSU, SEGYTraceReadingError
from StringIO import StringIO
import numpy as np
import os
import unittest
class SUTestCase(unittest.TestCase):
"""
Test cases for SU reading and writing.
Since the Seismic Unix format is a subset of the SEG Y file format a lot of
the SEG Y tests cover certain aspects of the SU format and ensure that the
SU implementation is working correctly.
"""
def setUp(self):
# directory where the test files are located
self.dir = os.path.dirname(__file__)
self.path = os.path.join(self.dir, 'data')
def test_readAndWriteSU(self):
"""
Reading and writing a SU file should not change it.
"""
file = os.path.join(self.path, '1.su_first_trace')
# Read the original file once.
with open(file, 'rb') as f:
org_data = f.read()
outfile = NamedTemporaryFile().name
# Read the SU file.
su = readSU(file)
# Write it.
su.write(outfile)
with open(outfile, 'rb') as f:
new_data = f.read()
os.remove(outfile)
# Should be identical!
self.assertEqual(org_data, new_data)
def test_enforcingByteordersWhileReading(self):
"""
Tests whether or not enforcing the byteorder while reading and writing
does something and works at all. Using the wrong byteorder will most
likely raise an Exception.
"""
# This file is little endian.
file = os.path.join(self.path, '1.su_first_trace')
# The following should both work.
su = readSU(file)
self.assertEqual(su.endian, '<')
su = readSU(file, endian='<')
self.assertEqual(su.endian, '<')
# The following not because it will unpack the header and try to unpack
# the number of data samples specified there which will of course not
# correct.
self.assertRaises(SEGYTraceReadingError, readSU, file, endian='>')
def test_readingAndWritingDifferentByteorders(self):
"""
Writing different byteorders should not change
"""
# This file is little endian.
file = os.path.join(self.path, '1.su_first_trace')
outfile = NamedTemporaryFile().name
# The following should both work.
su = readSU(file)
data = su.traces[0].data
# Also read the original file.
with open(file, 'rb') as f:
org_data = f.read()
self.assertEqual(su.endian, '<')
# Write it little endian.
su.write(outfile, endian='<')
with open(outfile, 'rb') as f:
new_data = f.read()
self.assertEqual(org_data, new_data)
su2 = readSU(outfile)
self.assertEqual(su2.endian, '<')
np.testing.assert_array_equal(data, su2.traces[0].data)
os.remove(outfile)
# Write it big endian.
su.write(outfile, endian='>')
with open(outfile, 'rb') as f:
new_data = f.read()
self.assertFalse(org_data == new_data)
su3 = readSU(outfile)
os.remove(outfile)
self.assertEqual(su3.endian, '>')
np.testing.assert_array_equal(data, su3.traces[0].data)
def test_unpackingSUData(self):
"""
Unpacks data and compares them to data unpacked by Madagascar.
"""
# This file has the same data as 1.sgy_first_trace.
file = os.path.join(self.path, '1.su_first_trace')
data_file = os.path.join(self.path, '1.sgy_first_trace.npy')
su = readSU(file)
data = su.traces[0].data
# The data is written as integer so it is also converted to float32.
correct_data = np.require(np.load(data_file).ravel(), 'float32')
# Compare both.
np.testing.assert_array_equal(correct_data, data)
def test_readStringIO(self):
"""
Tests reading from StringIO instances.
"""
# 1
file = os.path.join(self.path, '1.su_first_trace')
data = open(file, 'rb').read()
st = readSU(StringIO(data))
self.assertEqual(len(st.traces[0].data), 8000)
def suite():
return unittest.makeSuite(SUTestCase, 'test')
if __name__ == '__main__':
unittest.main(defaultTest='suite')
|