Skip to content

Commit fb60ef5

Browse files
committed
Ensure getting true raw bytes for DR array
1 parent 1271bb7 commit fb60ef5

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

0 commit comments

Comments
 (0)