Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

# -*- coding: utf-8 -*- 

""" 

Functions to generate benchmark plots from given SU files. 

 

.. versionadded:: 0.5.1 

 

:copyright: 

    The ObsPy Development Team (devs@obspy.org) 

:license: 

    GNU Lesser General Public License, Version 3 

    (http://www.gnu.org/copyleft/lesser.html) 

""" 

 

from obspy.segy.segy import SUFile, readSU 

import StringIO 

import math 

import matplotlib.pylab as plt 

import matplotlib.cm as cm 

import numpy as np 

import os 

 

 

def _calcOffset(trace): 

    """ 

    Calculating offset for given trace. 

    """ 

    th = trace.header 

    # Some programs out there use the scalco field in an non-standard 

    # way. These programs store the scaling factor or divisor in terms 

    # of an exponent to the base of 10. This conditional handles these 

    # non-standard values. 

    scalco = abs(th.scalar_to_be_applied_to_all_coordinates) 

    if scalco < 10 and scalco != 1: 

        scalco = pow(10, scalco) 

    offset = 1.0 / scalco * \ 

        math.sqrt(pow(th.group_coordinate_x - th.source_coordinate_x, 2) + \ 

                  pow(th.group_coordinate_y - th.source_coordinate_y, 2)) 

    return offset 

 

 

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) 

 

 

if __name__ == '__main__': 

    import doctest 

    doctest.testmod(exclude_empty=True)