1try: 2 from collections.abc import Mapping # noqa 3except ImportError: 4 from collections import Mapping # noqa 5 6import itertools 7try: from future_builtins import zip 8except ImportError: pass 9import numpy as np 10 11from .utils import castarray 12 13# in order to support [:end] syntax, we must make sure 14# start has a non-None value. lineno.indices() would set it 15# to 0, but we don't know if that's a reasonable value or 16# not. If start is None we set it to the first line 17def sanitize_slice(s, source): 18 if all((s.start, s.stop, s.step)): 19 return s 20 21 start, stop, step = s.start, s.stop, s.step 22 increasing = step is None or step > 0 23 24 if start is None: 25 start = min(source) if increasing else max(source) 26 27 if stop is None: 28 stop = max(source) + 1 if increasing else min(source) - 1 29 30 return slice(start, stop, step) 31 32class Line(Mapping): 33 """ 34 The Line implements the dict interface, with a fixed set of int_like keys, 35 the line numbers/labels. Data is read lazily from disk, so iteration does 36 not consume much memory, and are returned as numpy.ndarrays. 37 38 It provides a convenient interface for reading data in a cartesian grid 39 system, provided one exists and is detectable by segyio. 40 41 Lines can be accessed individually or with slices, and writing is done via 42 assignment. Note that accessing lines uses the line numbers, not their 43 position, so if a files has lines [2400..2500], accessing line [0..100] 44 will be an error. Note that since each line is returned as a numpy.ndarray, 45 meaning accessing the intersections of the inline and crossline is 46 0-indexed - orthogonal labels are not preserved. 47 48 Additionally, the line has a concept of offsets, which is useful when 49 dealing with prestack files. Offsets are accessed via sub indexing, meaning 50 iline[10, 4] will give you line 10 at offset 4. Please note that offset, 51 like lines, are accessed via their labels, not their indices. If your file 52 has the offsets [150, 250, 350, 450] and the lines [2400..2500], you can 53 access the third offset with [2403, 350]. Please refer to the examples for 54 more details. If no offset is specified, segyio will give you the first. 55 56 Notes 57 ----- 58 .. versionadded:: 1.1 59 60 .. versionchanged:: 1.6 61 common dict operations (Mapping) 62 """ 63 64 def __init__(self, filehandle, labels, length, stride, offsets, name): 65 self.filehandle = filehandle.xfd 66 self.lines = labels 67 self.length = length 68 self.stride = stride 69 self.shape = (length, len(filehandle.samples)) 70 self.dtype = filehandle.dtype 71 72 # pre-compute all line beginnings 73 from ._segyio import fread_trace0 74 self.heads = { 75 label: fread_trace0(label, 76 length, 77 stride, 78 len(offsets), 79 labels, 80 name) 81 for label in labels 82 } 83 84 self.offsets = { x: i for i, x in enumerate(offsets) } 85 self.default_offset = offsets[0] 86 87 def ranges(self, index, offset): 88 if not isinstance(index, slice): 89 index = slice(index, index + 1) 90 91 if not isinstance(offset, slice): 92 offset = slice(offset, offset + 1) 93 94 index = sanitize_slice(index, self.heads.keys()) 95 offset = sanitize_slice(offset, self.offsets.keys()) 96 irange = range(*index.indices(max(self.heads.keys()) + 1)) 97 orange = range(*offset.indices(max(self.offsets.keys()) + 1)) 98 irange = filter(self.heads.__contains__, irange) 99 orange = filter(self.offsets.__contains__, orange) 100 # offset-range is used in inner loops, so make it a list for 101 # reusability. offsets are usually few, so no real punishment by using 102 # non-generators here 103 return irange, list(orange) 104 105 def __getitem__(self, index): 106 """line[i] or line[i, o] 107 108 The line `i`, or the line `i` at a specific offset `o`. ``line[i]`` 109 returns a numpy array, and changes to this array will *not* be 110 reflected on disk. 111 112 The `i` and `o` are *keys*, and should correspond to the line- and 113 offset labels in your file, and in the `ilines`, `xlines`, and 114 `offsets` attributes. 115 116 Slices can contain lines and offsets not in the file, and like with 117 list slicing, these are handled gracefully and ignored. 118 119 When `i` or `o` is a slice, a generator of numpy arrays is returned. If 120 the slice is defaulted (:), segyio knows enough about the structure to 121 give you all of the respective labels. 122 123 When both `i` and `o` are slices, only one generator is returned, and 124 the lines are yielded offsets-first, roughly equivalent to the double 125 for loop:: 126 127 >>> for line in lines: 128 ... for off in offsets: 129 ... yield line[line, off] 130 ... 131 132 Parameters 133 ---------- 134 i : int or slice 135 o : int or slice 136 137 Returns 138 ------- 139 line : numpy.ndarray of dtype or generator of numpy.ndarray of dtype 140 141 Raises 142 ------ 143 KeyError 144 If `i` or `o` don't exist 145 146 Notes 147 ----- 148 .. versionadded:: 1.1 149 150 Examples 151 -------- 152 153 Read an inline: 154 155 >>> x = line[2400] 156 157 Copy every inline into a list: 158 159 >>> l = [numpy.copy(x) for x in iline[:]] 160 161 Numpy operations on every other inline: 162 163 >>> for line in line[::2]: 164 ... line = line * 2 165 ... avg = np.average(line) 166 167 Read lines up to 2430: 168 169 >>> for line in line[:2430]: 170 ... line.mean() 171 172 Copy all lines at all offsets: 173 174 >>> l = [numpy.copy(x) for x in line[:,:]] 175 176 Copy all offsets of a line: 177 178 >>> x = numpy.copy(iline[10,:]) 179 180 Copy all lines at a fixed offset: 181 182 >>> x = numpy.copy(iline[:, 120]) 183 184 Copy every other line and offset: 185 186 >>> map(numpy.copy, line[::2, ::2]) 187 188 Copy all offsets [200, 250, 300, 350, ...] in the range [200, 800) for 189 all lines [2420,2460): 190 191 >>> l = [numpy.copy(x) for x in line[2420:2460, 200:800:50]] 192 """ 193 194 offset = self.default_offset 195 try: index, offset = index 196 except TypeError: pass 197 198 # prioritise the code path that's potentially in loops externally 199 try: 200 head = self.heads[index] + self.offsets[offset] 201 except TypeError: 202 # index is either unhashable (because it's a slice), or offset is a 203 # slice. 204 pass 205 else: 206 return self.filehandle.getline(head, 207 self.length, 208 self.stride, 209 len(self.offsets), 210 np.empty(self.shape, dtype=self.dtype), 211 ) 212 213 # at this point, either offset or index is a slice (or proper 214 # type-error), so we're definitely making a generator. make them both 215 # slices to unify all code paths 216 irange, orange = self.ranges(index, offset) 217 218 def gen(): 219 x = np.empty(self.shape, dtype=self.dtype) 220 y = np.copy(x) 221 222 # only fetch lines that exist. the slice can generate both offsets 223 # and line numbers that don't exist, so filter out misses before 224 # they happen 225 for line in irange: 226 for off in orange: 227 head = self.heads[line] + self.offsets[off] 228 self.filehandle.getline(head, 229 self.length, 230 self.stride, 231 len(self.offsets), 232 y, 233 ) 234 y, x = x, y 235 yield x 236 237 return gen() 238 239 def __setitem__(self, index, val): 240 """line[i] = val or line[i, o] = val 241 242 Follows the same rules for indexing and slicing as ``line[i]``. 243 244 In either case, if the `val` iterable is exhausted before the line(s), 245 assignment stops with whatever is written so far. If `val` is longer 246 than an individual line, it's essentially truncated. 247 248 Parameters 249 ---------- 250 i : int or slice 251 offset : int or slice 252 val : array_like 253 254 Raises 255 ------ 256 KeyError 257 If `i` or `o` don't exist 258 259 Notes 260 ----- 261 .. versionadded:: 1.1 262 263 Examples 264 -------- 265 Copy a full line: 266 267 >>> line[2400] = other[2834] 268 269 Copy first half of the inlines from g to f: 270 271 >>> line[:] = other[:labels[len(labels) / 2]] 272 273 Copy every other line consecutively: 274 275 >>> line[:] = other[::2] 276 277 Copy every third offset: 278 279 >>> line[:,:] = other[:,::3] 280 281 Copy a line into a set line and offset: 282 283 >>> line[12, 200] = other[21] 284 """ 285 286 offset = self.default_offset 287 try: index, offset = index 288 except TypeError: pass 289 290 try: head = self.heads[index] + self.offsets[offset] 291 except TypeError: pass 292 else: 293 return self.filehandle.putline(head, 294 self.length, 295 self.stride, 296 len(self.offsets), 297 index, 298 offset, 299 castarray(val, dtype = self.dtype), 300 ) 301 302 irange, orange = self.ranges(index, offset) 303 304 val = iter(val) 305 for line in irange: 306 for off in orange: 307 head = self.heads[line] + self.offsets[off] 308 try: self.filehandle.putline(head, 309 self.length, 310 self.stride, 311 len(self.offsets), 312 line, 313 off, 314 next(val), 315 ) 316 except StopIteration: return 317 318 # can't rely on most Mapping default implementations of 319 # dict-like, because iter() does not yield keys for this class, it gives 320 # the lines themselves. that violates some assumptions (but segyio's always 321 # worked that way), and it's the more natural behaviour for segyio, so it's 322 # acceptible. additionally, the default implementations would be very slow 323 # and ineffective because they assume __getitem__ is sufficiently cheap, 324 # but it isn't here since it involves a disk operation 325 def __len__(self): 326 """x.__len__() <==> len(x)""" 327 return len(self.heads) 328 329 def __iter__(self): 330 """x.__iter__() <==> iter(x)""" 331 return self[:] 332 333 def __contains__(self, key): 334 """x.__contains__(y) <==> y in x""" 335 return key in self.heads 336 337 def keys(self): 338 """D.keys() -> a set-like object providing a view on D's keys""" 339 return sorted(self.heads.keys()) 340 341 def values(self): 342 """D.values() -> generator of D's values""" 343 return self[:] 344 345 def items(self): 346 """D.values() -> generator of D's (key,values), as 2-tuples""" 347 return zip(self.keys(), self[:]) 348 349class HeaderLine(Line): 350 """ 351 The Line implements the dict interface, with a fixed set of int_like keys, 352 the line numbers/labels. The values are iterables of Field objects. 353 354 Notes 355 ----- 356 .. versionadded:: 1.1 357 358 .. versionchanged:: 1.6 359 common dict operations (Mapping) 360 """ 361 # a lot of implementation details are shared between reading data traces 362 # line-by-line and trace headers line-by-line, so (ab)use inheritance for 363 # __len__, keys() etc., however, the __getitem__ is way different and is re-implemented 364 365 def __init__(self, header, base, direction): 366 super(HeaderLine, self).__init__(header.segy, 367 base.lines, 368 base.length, 369 base.stride, 370 sorted(base.offsets.keys()), 371 'header.' + direction, 372 ) 373 self.header = header 374 375 def __getitem__(self, index): 376 """line[i] or line[i, o] 377 378 The line `i`, or the line `i` at a specific offset `o`. ``line[i]`` 379 returns an iterable of `Field` objects, and changes to these *will* be 380 reflected on disk. 381 382 The `i` and `o` are *keys*, and should correspond to the line- and 383 offset labels in your file, and in the `ilines`, `xlines`, and 384 `offsets` attributes. 385 386 Slices can contain lines and offsets not in the file, and like with 387 list slicing, these are handled gracefully and ignored. 388 389 When `i` or `o` is a slice, a generator of iterables of headers are 390 returned. 391 392 When both `i` and `o` are slices, one generator is returned for the 393 product `i` and `o`, and the lines are yielded offsets-first, roughly 394 equivalent to the double for loop:: 395 396 >>> for line in lines: 397 ... for off in offsets: 398 ... yield line[line, off] 399 ... 400 401 Parameters 402 ---------- 403 404 i : int or slice 405 o : int or slice 406 407 Returns 408 ------- 409 line : iterable of Field or generator of iterator of Field 410 411 Raises 412 ------ 413 KeyError 414 If `i` or `o` don't exist 415 416 Notes 417 ----- 418 .. versionadded:: 1.1 419 420 """ 421 offset = self.default_offset 422 try: index, offset = index 423 except TypeError: pass 424 425 try: 426 start = self.heads[index] + self.offsets[offset] 427 except TypeError: 428 # index is either unhashable (because it's a slice), or offset is a 429 # slice. 430 pass 431 432 else: 433 step = self.stride * len(self.offsets) 434 stop = start + step * self.length 435 return self.header[start:stop:step] 436 437 def gen(): 438 irange, orange = self.ranges(index, offset) 439 for line in irange: 440 for off in orange: 441 yield self[line, off] 442 443 return gen() 444 445 def __setitem__(self, index, val): 446 """line[i] = val or line[i, o] = val 447 448 Follows the same rules for indexing and slicing as ``line[i]``. If `i` 449 is an int, and `val` is a dict or Field, that value is replicated and 450 assigned to every trace header in the line, otherwise it's treated as 451 an iterable, and each trace in the line is assigned the ``next()`` 452 yielded value. 453 454 If `i` or `o` is a slice, `val` must be an iterable. 455 456 In either case, if the `val` iterable is exhausted before the line(s), 457 assignment stops with whatever is written so far. 458 459 Parameters 460 ---------- 461 i : int or slice 462 offset : int or slice 463 val : dict_like or iterable of dict_like 464 465 Raises 466 ------ 467 KeyError 468 If `i` or `o` don't exist 469 470 Notes 471 ----- 472 .. versionadded:: 1.1 473 474 Examples 475 -------- 476 Rename the iline 3 to 4: 477 478 >>> line[3] = { TraceField.INLINE_3D: 4 } 479 >>> # please note that rewriting the header won't update the 480 >>> # file's interpretation of the file until you reload it, so 481 >>> # the new iline 4 will be considered iline 3 until the file 482 >>> # is reloaded 483 484 Set offset line 3 offset 3 to 5: 485 486 >>> line[3, 3] = { TraceField.offset: 5 } 487 """ 488 offset = self.default_offset 489 try: index, offset = index 490 except TypeError: pass 491 492 try: start = self.heads[index] + self.offsets[offset] 493 except TypeError: pass 494 else: 495 step = self.stride * len(self.offsets) 496 stop = start + step * self.length 497 self.header[start:stop:step] = val 498 return 499 500 # if this is a dict-like, just repeat it 501 if hasattr(val, 'keys'): 502 val = itertools.repeat(val) 503 504 irange, orange = self.ranges(index, offset) 505 val = iter(val) 506 for line in irange: 507 for off in orange: 508 try: 509 self[line, off] = next(val) 510 except StopIteration: 511 return 512