|
9 | 9 |
|
10 | 10 | import numpy as np |
11 | 11 | from segy import SegyFile |
| 12 | +from segy.indexing import merge_cat_file |
12 | 13 |
|
13 | 14 | from mdio.api.io import to_mdio |
14 | 15 | from mdio.builder.schemas.dtype import ScalarType |
@@ -151,13 +152,39 @@ def trace_worker( # noqa: PLR0913 |
151 | 152 | ) |
152 | 153 | if raw_header_key in worker_variables: |
153 | 154 | 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 | + |
155 | 181 | ds_to_write[raw_header_key] = Variable( |
156 | 182 | ds_to_write[raw_header_key].dims, |
157 | 183 | tmp_raw_headers, |
158 | 184 | 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, |
160 | 186 | ) |
| 187 | + |
161 | 188 | data_variable = ds_to_write[data_variable_name] |
162 | 189 | fill_value = _get_fill_value(ScalarType(data_variable.dtype.name)) |
163 | 190 | tmp_samples = np.full_like(data_variable, fill_value=fill_value) |
|
0 commit comments