Skip to content

Commit eb6bed2

Browse files
BrianMichelltasansal
authored andcommitted
Ensure getting true raw bytes for DR array
1 parent e4dd2f1 commit eb6bed2

File tree

1 file changed

+29
-2
lines changed

1 file changed

+29
-2
lines changed

src/mdio/segy/_workers.py

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
import numpy as np
1111
from segy import SegyFile
12+
from segy.indexing import merge_cat_file
1213

1314
from mdio.api.io import to_mdio
1415
from mdio.builder.schemas.dtype import ScalarType
@@ -153,13 +154,39 @@ def trace_worker( # noqa: PLR0913
153154
)
154155
if raw_header_key in worker_variables:
155156
tmp_raw_headers = np.zeros_like(dataset[raw_header_key])
156-
tmp_raw_headers[not_null] = traces.header.view("|V240") # TODO: Ensure this is using the RAW view and not an interpreted view.
157+
158+
# Get the indices where we need to place results
159+
live_mask = not_null
160+
live_positions = np.where(live_mask.ravel())[0]
161+
162+
if len(live_positions) > 0:
163+
# Calculate byte ranges for headers
164+
HEADER_SIZE = 240
165+
trace_offset = segy_file.spec.trace.offset
166+
trace_itemsize = segy_file.spec.trace.itemsize
167+
168+
starts = []
169+
ends = []
170+
for global_trace_idx in live_trace_indexes:
171+
header_start = trace_offset + global_trace_idx * trace_itemsize
172+
header_end = header_start + HEADER_SIZE
173+
starts.append(header_start)
174+
ends.append(header_end)
175+
176+
# Capture raw bytes
177+
raw_header_bytes = merge_cat_file(segy_file.fs, segy_file.url, starts, ends)
178+
179+
# Convert and place results
180+
raw_headers_array = np.frombuffer(bytes(raw_header_bytes), dtype="|V240")
181+
tmp_raw_headers.ravel()[live_positions] = raw_headers_array
182+
157183
ds_to_write[raw_header_key] = Variable(
158184
ds_to_write[raw_header_key].dims,
159185
tmp_raw_headers,
160186
attrs=ds_to_write[raw_header_key].attrs,
161-
encoding=ds_to_write[raw_header_key].encoding, # Not strictly necessary, but safer than not doing it.
187+
encoding=ds_to_write[raw_header_key].encoding,
162188
)
189+
163190
data_variable = ds_to_write[data_variable_name]
164191
fill_value = _get_fill_value(ScalarType(data_variable.dtype.name))
165192
tmp_samples = np.full_like(data_variable, fill_value=fill_value)

0 commit comments

Comments
 (0)