Skip to content

Commit cdf9c6f

Browse files
committed
Ensure getting true raw bytes for DR array
1 parent 1e91fab commit cdf9c6f

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

0 commit comments

Comments
 (0)