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')
with NamedTemporaryFile() as tf:
outfile = tf.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)
# 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)
self.assertEqual(su3.endian, '>')
np.testing.assert_array_equal(data, su3.traces[0].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_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 test_readBytesIO(self):
"""
Tests reading from BytesIO instances.
"""
# 1
filename = os.path.join(self.path, '1.su_first_trace')
with open(filename, 'rb') as fp:
data = fp.read()
st = readSU(io.BytesIO(data))
self.assertEqual(len(st.traces[0].data), 8000)
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_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()
with NamedTemporaryFile() as tf:
outfile = tf.name
# Read the SU file.
su = readSU(file)
# Write it.
su.write(outfile)
with open(outfile, 'rb') as f:
new_data = f.read()
# Should be identical!
self.assertEqual(org_data, new_data)
def plotBenchmark(sufiles,
normalize='traces',
clip_partial_traces=True,
scale=3.0,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
fig=None,
plot_legend=True,
title="",
size=(800, 600),
dpi=100,
outfile=None,
format=None,
trim_to_smallest_trace=True):
"""
Plot a benchmark plot from given SU files.
:type sufiles: List of SU file names or :class:`~obspy.segy.segy.SUFile`
objects.
:param sufiles: SU files to plot.
:type normalize: ``None``, ``'stream'`` or ``'traces'``, optional
:param normalize: If ``'stream'`` is given it will normalize per stream.
The keyword ``'traces'`` normalizes all traces in all streams. ``None``
will skip normalization. Defaults to ``'traces'``.
:type clip_partial_traces: bool, optional
:param clip_partial_traces: Clips traces which are not completely plotted.
Defaults to ``True``.
:type trim_to_smallest_trace: bool, optional
:param trim_to_smallest_trace: Trims all traces to shortest available
trace. Defaults to ``True``.
:type plot_legend: bool, optional
:param plot_legend: If enabled plots a legend generated by given SU files.
Defaults to ``True``.
:type title: string, optional
:param title: Plots a title if given. Disabled by default.
:type scale: float, optional
:param scale: Scales all amplitudes by this factor. Defaults to ``3.0``.
:type xmin: float, optional
:param xmin: Minimum of time axis.
:type xmax: float, optional
:param xmax: Maximum of time axis.
:type ymin: float, optional
:param ymin: Minimum of offset axis.
:type ymax: float, optional
:param ymax: Maximum of offset axis.
:type fig: :class:`matplotlib.figure.Figure`
:param fig: Use an existing matplotlib figure instance.
Default to ``None``.
:type size: tuple, optional
:param size: Size tuple in pixel for the output file. This corresponds
to the resolution of the graph for vector formats. Defaults to
``(800, 800)`` pixel.
:type dpi: int, optional
:param dpi: Dots per inch of the output file. This also affects the
size of most elements in the graph (text, linewidth, ...).
Defaults to ``100``.
:type outfile: string, optional
:param outfile: Output file name. Also used to automatically
determine the output format. Supported file formats depend on your
matplotlib backend. Most backends support png, pdf, ps, eps and
svg. Defaults to ``None``.
:type format: string, optional
:param format: Format of the graph picture. If no format is given the
outfile parameter will be used to try to automatically determine
the output format. If no format is found it defaults to png output.
If no outfile is specified but a format is, than a binary
imagestring will be returned.
Defaults to ``None``.
.. versionadded:: 0.5.1
.. rubric:: Example
The following example plots four seismic unix files in one benchmark image.
>>> import glob
>>> sufiles = glob.glob('seismic01_*_vz.su')
>>> from obspy.segy.benchmark import plotBenchmark
>>> plotBenchmark(sufiles, title="Homogenous halfspace") # doctest: +SKIP
.. plot::
from obspy.core.util import getExampleFile
files = [getExampleFile('seismic01_fdmpi_vz.su'),
getExampleFile('seismic01_gemini_vz.su'),
getExampleFile('seismic01_sofi2D_transformed_vz.su'),
getExampleFile('seismic01_specfem_vz.su')]
from obspy.segy.segy import readSU
from obspy.segy.benchmark import plotBenchmark
sufiles = [readSU(file) for file in files]
plotBenchmark(sufiles, title="Homogenous halfspace", xmax=0.14)
"""
if not sufiles:
return
# ensure we have SUFile objects
streams = []
for sufile in sufiles:
if isinstance(sufile, SUFile):
streams.append(sufile)
else:
streams.append(readSU(sufile))
# get delta from first trace only and assume it for all traces
delta = streams[0].traces[0].header.sample_interval_in_ms_for_this_trace
# trim to smallest trace
if trim_to_smallest_trace:
# search smallest trace
npts = []
for st in streams:
npts.append(min([len(tr.data) for tr in st.traces]))
npts = min(npts)
# trim all traces to max_npts
for st in streams:
for tr in st.traces:
tr.data = tr.data[0:npts]
# get offsets
offsets = []
for st in streams:
temp = []
for tr in st.traces:
temp.append(_calcOffset(tr))
offsets.append((max(temp) - min(temp)) / len(st.traces))
min_offset = min(offsets)
# normalize
if normalize != 'stream':
maximums = []
minimums = []
for st in streams:
maximums.append(max([_i.data.max() for _i in st.traces]))
minimums.append(min([_i.data.min() for _i in st.traces]))
minimum = min(minimums)
maximum = max(maximums)
data_range = maximum - minimum
for st in streams:
if normalize == 'stream':
data_range = max([_i.data.max() for _i in st.traces]) - \
min([_i.data.min() for _i in st.traces])
for tr in st.traces:
if normalize == 'traces':
data_range = tr.data.max() - tr.data.min()
data_range = abs(data_range)
tr.data /= (data_range / (min_offset * scale))
# Setup the figure if not passed explicitly.
if not fig:
# Setup figure and axes
_fig = plt.figure(num=None,
dpi=dpi,
figsize=(float(size[0]) / dpi, float(size[1]) / dpi))
# XXX: Figure out why this is needed sometimes.
# Set size and dpi.
_fig.set_dpi(dpi)
_fig.set_figwidth(float(size[0]) / dpi)
_fig.set_figheight(float(size[1]) / dpi)
# set title
if title:
if '\n' in title:
title, subtitle = title.split('\n', 1)
_fig.suptitle(title, y=0.97)
_fig.suptitle(subtitle, y=0.935, fontsize='x-small')
else:
_fig.suptitle(title, y=0.95)
else:
_fig = fig
# get current axis
ax = _fig.gca()
# labels - either file names or stream numbers
try:
labels = [os.path.basename(trace.file.name) for trace in streams]
except:
labels = ['Stream #' + str(i) for i in range(len(streams))]
# colors - either auto generated or use a preset
if len(streams) > 5:
colors = cm.get_cmap('hsv', len(streams))
colors = [colors[i] for i in len(streams)]
else:
colors = ['grey', 'red', 'green', 'blue', 'black']
# set first min and max
min_y = np.Inf
max_y = -np.Inf
max_x = -np.Inf
# plot
for _i, st in enumerate(streams):
color = colors[_i]
legend = True
for _j, tr in enumerate(st.traces):
# calculating offset for each trace
offset = _calcOffset(tr)
# create x and y arrays
y = (tr.data) + offset
x = np.arange(len(tr.data)) * delta / 1000000.
# get boundaries
if max(y) > max_y:
max_y = max(y)
if min(y) < min_y:
min_y = min(y)
if max(x) > max_x:
max_x = max(x)
# test if in image
if clip_partial_traces:
if ymin is not None and min(y) < ymin:
continue
if ymax is not None and max(y) > ymax:
continue
# plot, add labels only at new streams
if legend:
ax.plot(x, y, color=color, label=labels[_i], lw=0.5)
else:
ax.plot(x, y, color=color, lw=0.5)
legend = False
# limit offset axis
spacing = (max_y - min_y) / 50.0
if ymax is None or ymax > max_y:
ymax = max_y + spacing
if ymin is None or ymin < min_y:
ymin = min_y - spacing
ax.set_ylim(ymin, ymax)
# limit time axis
if xmin is None or xmin < 0:
xmin = 0
if xmax is None or xmax > max_x:
xmax = max_x
ax.set_xlim(xmin, xmax)
# axis labels
ax.set_xlabel('time [s]')
ax.set_ylabel('offset [m]')
# add legend
if plot_legend:
plt.legend(loc=4,
bbox_to_anchor=(0, 0.1, 0.9, 1),
bbox_transform=plt.gcf().transFigure)
try:
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize='x-small')
except:
pass
# handle output
if outfile is None:
# Return an binary imagestring if not outfile but format.
if format:
imgdata = StringIO.StringIO()
_fig.savefig(imgdata, format=format, dpi=dpi)
imgdata.seek(0)
return imgdata.read()
elif fig is None:
plt.show()
else:
return fig
else:
# If format is set use it.
if format:
_fig.savefig(outfile, format=format, dpi=dpi)
# Otherwise use format from self.outfile or default to PNG.
else:
_fig.savefig(outfile, dpi=dpi)
def plotBenchmark(sufiles, normalize='traces', clip_partial_traces=True,
scale=3.0, xmin=None, xmax=None, ymin=None, ymax=None,
fig=None, plot_legend=True, title="", size=(800, 600),
dpi=100, outfile=None, format=None,
trim_to_smallest_trace=True):
"""
Plot a benchmark plot from given SU files.
:type sufiles: list of str or :class:`~obspy.segy.segy.SUFile`
:param sufiles: SU files to plot.
:type normalize: str or ``None``, optional
:param normalize: Normalization method to perform:
``'stream'``
Normalize per stream.
``'traces'``
Normalizes all traces in all streams.
``None``
Skip normalization.
Defaults to ``'traces'``.
:type clip_partial_traces: bool, optional
:param clip_partial_traces: Clips traces which are not completely plotted.
Defaults to ``True``.
:type trim_to_smallest_trace: bool, optional
:param trim_to_smallest_trace: Trims all traces to shortest available
trace. Defaults to ``True``.
:type plot_legend: bool, optional
:param plot_legend: If enabled plots a legend generated by given SU files.
Defaults to ``True``.
:type title: str, optional
:param title: Plots a title if given. Disabled by default.
:type scale: float, optional
:param scale: Scales all amplitudes by this factor. Defaults to ``3.0``.
:type xmin: float, optional
:param xmin: Minimum of time axis.
:type xmax: float, optional
:param xmax: Maximum of time axis.
:type ymin: float, optional
:param ymin: Minimum of offset axis.
:type ymax: float, optional
:param ymax: Maximum of offset axis.
:type fig: :class:`matplotlib.figure.Figure`, optional
:param fig: Use an existing matplotlib figure instance.
Default to ``None``.
:type size: tuple, optional
:param size: Size tuple in pixel for the output file. This corresponds
to the resolution of the graph for vector formats. Defaults to
``(800, 800)`` pixel.
:type dpi: int, optional
:param dpi: Dots per inch of the output file. This also affects the
size of most elements in the graph (text, linewidth, ...).
Defaults to ``100``.
:type outfile: str, optional
:param outfile: Output file name. Also used to automatically
determine the output format. Supported file formats depend on your
matplotlib backend. Most backends support png, pdf, ps, eps and
svg. Defaults to ``None``.
:type format: str, optional
:param format: Format of the graph picture. If no format is given the
outfile parameter will be used to try to automatically determine
the output format. If no format is found it defaults to png output.
If no outfile is specified but a format is, than a binary
imagestring will be returned.
Defaults to ``None``.
.. versionadded:: 0.5.1
.. rubric:: Example
The following example plots four seismic unix files in one benchmark image.
>>> import glob
>>> sufiles = glob.glob('seismic01_*_vz.su')
>>> from obspy.segy.benchmark import plotBenchmark
>>> plotBenchmark(sufiles, title="Homogenous halfspace") # doctest: +SKIP
.. plot::
from obspy.core.util import getExampleFile
files = [getExampleFile('seismic01_fdmpi_vz.su'),
getExampleFile('seismic01_gemini_vz.su'),
getExampleFile('seismic01_sofi2D_transformed_vz.su'),
getExampleFile('seismic01_specfem_vz.su')]
from obspy.segy.segy import readSU
from obspy.segy.benchmark import plotBenchmark
sufiles = [readSU(file) for file in files]
plotBenchmark(sufiles, title="Homogenous halfspace", xmax=0.14)
"""
if not sufiles:
return
# ensure we have SUFile objects
streams = []
for sufile in sufiles:
if isinstance(sufile, SUFile):
streams.append(sufile)
else:
streams.append(readSU(sufile))
# get delta from first trace only and assume it for all traces
delta = streams[0].traces[0].header.sample_interval_in_ms_for_this_trace
# trim to smallest trace
if trim_to_smallest_trace:
# search smallest trace
npts = []
for st in streams:
npts.append(min([len(tr.data) for tr in st.traces]))
npts = min(npts)
# trim all traces to max_npts
for st in streams:
for tr in st.traces:
tr.data = tr.data[0:npts]
# get offsets
offsets = []
for st in streams:
temp = []
for tr in st.traces:
temp.append(_calcOffset(tr))
offsets.append((max(temp) - min(temp)) / len(st.traces))
min_offset = min(offsets)
# normalize
if normalize != 'stream':
maximums = []
minimums = []
for st in streams:
maximums.append(max([_i.data.max() for _i in st.traces]))
minimums.append(min([_i.data.min() for _i in st.traces]))
minimum = min(minimums)
maximum = max(maximums)
data_range = maximum - minimum
for st in streams:
if normalize == 'stream':
data_range = max([_i.data.max() for _i in st.traces]) - \
min([_i.data.min() for _i in st.traces])
for tr in st.traces:
if normalize == 'traces':
data_range = tr.data.max() - tr.data.min()
data_range = abs(data_range)
tr.data /= (data_range / (min_offset * scale))
# Setup the figure if not passed explicitly.
if not fig:
# Setup figure and axes
_fig = plt.figure(num=None, dpi=dpi, figsize=(float(size[0]) / dpi,
float(size[1]) / dpi))
# XXX: Figure out why this is needed sometimes.
# Set size and dpi.
_fig.set_dpi(dpi)
_fig.set_figwidth(float(size[0]) / dpi)
_fig.set_figheight(float(size[1]) / dpi)
# set title
if title:
if '\n' in title:
title, subtitle = title.split('\n', 1)
_fig.suptitle(title, y=0.97)
_fig.suptitle(subtitle, y=0.935, fontsize='x-small')
else:
_fig.suptitle(title, y=0.95)
else:
_fig = fig
# get current axis
ax = _fig.gca()
# labels - either file names or stream numbers
try:
labels = [os.path.basename(trace.file.name) for trace in streams]
except:
labels = ['Stream #' + str(i) for i in range(len(streams))]
# colors - either auto generated or use a preset
if len(streams) > 5:
colors = cm.get_cmap('hsv', len(streams))
colors = [colors[i] for i in len(streams)]
else:
colors = ['grey', 'red', 'green', 'blue', 'black']
# set first min and max
min_y = np.Inf
max_y = -np.Inf
max_x = -np.Inf
# plot
for _i, st in enumerate(streams):
color = colors[_i]
legend = True
for _j, tr in enumerate(st.traces):
# calculating offset for each trace
offset = _calcOffset(tr)
# create x and y arrays
y = (tr.data) + offset
x = np.arange(len(tr.data)) * delta / 1000000.
# get boundaries
if max(y) > max_y:
max_y = max(y)
if min(y) < min_y:
min_y = min(y)
if max(x) > max_x:
max_x = max(x)
# test if in image
if clip_partial_traces:
if ymin is not None and min(y) < ymin:
continue
if ymax is not None and max(y) > ymax:
continue
# plot, add labels only at new streams
if legend:
ax.plot(x, y, color=color, label=labels[_i], lw=0.5)
else:
ax.plot(x, y, color=color, lw=0.5)
legend = False
# limit offset axis
spacing = (max_y - min_y) / 50.0
if ymax is None or ymax > max_y:
ymax = max_y + spacing
if ymin is None or ymin < min_y:
ymin = min_y - spacing
ax.set_ylim(ymin, ymax)
# limit time axis
if xmin is None or xmin < 0:
xmin = 0
if xmax is None or xmax > max_x:
xmax = max_x
ax.set_xlim(xmin, xmax)
# axis labels
ax.set_xlabel('time [s]')
ax.set_ylabel('offset [m]')
# add legend
if plot_legend:
plt.legend(loc=4, bbox_to_anchor=(0, 0.1, 0.9, 1),
bbox_transform=plt.gcf().transFigure)
try:
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize='x-small')
except:
pass
# handle output
if outfile is None:
# Return an binary imagestring if not outfile but format.
if format:
imgdata = io.StringIO()
_fig.savefig(imgdata, format=format, dpi=dpi)
imgdata.seek(0)
return imgdata.read()
elif fig is None:
plt.show()
else:
return fig
else:
# If format is set use it.
if format:
_fig.savefig(outfile, format=format, dpi=dpi)
# Otherwise use format from self.outfile or default to PNG.
else:
_fig.savefig(outfile, dpi=dpi)
from obspy.core.util import getExampleFile
from obspy.segy.segy import readSU
from obspy.segy.benchmark import plotBenchmark
files = [getExampleFile('seismic01_fdmpi_vz.su'),
getExampleFile('seismic01_gemini_vz.su'),
getExampleFile('seismic01_sofi2D_transformed_vz.su'),
getExampleFile('seismic01_specfem_vz.su')]
sufiles = [readSU(file) for file in files]
plotBenchmark(sufiles, title="Homogenous halfspace", xmax=0.14)
