Reset to main and add profiling

Author: BrianMichell

src/mdio/core/grid.py

@@ -134,24 +134,19 @@ def build_map(self, index_headers: HeaderArray) -> None:
         # Process headers in batches
         for start in range(0, total_live_traces, batch_size):
             end = min(start + batch_size, total_live_traces)
-
-            # 1) build your per-dimension index arrays
-            live_dim_indices = [
-                np.searchsorted(dim, index_headers[dim.name][start:end])
-                .astype(np.uint32)
-                for dim in self.dims[:-1]
-            ]
-
-            # 2) flatten to 1D indices
-            flat_idx = np.ravel_multi_index(tuple(live_dim_indices), dims=self.map.shape)
-
-            # 3) write into flattened views
-            flat_map  = self.map.reshape(-1)
-            flat_mask = self.live_mask.reshape(-1)
-            trace_indices = np.arange(start, end, dtype=flat_map.dtype)
-
-            flat_map[flat_idx]  = trace_indices
-            flat_mask[flat_idx] = True
+            live_dim_indices = []
+
+            # Compute indices for the batch
+            for dim in self.dims[:-1]:
+                dim_hdr = index_headers[dim.name][start:end]
+                indices = np.searchsorted(dim, dim_hdr).astype(np.uint32)
+                live_dim_indices.append(indices)
+            live_dim_indices = tuple(live_dim_indices)
+
+            # Assign trace indices
+            trace_indices = np.arange(start, end, dtype=np.uint64)
+            self.map.vindex[live_dim_indices] = trace_indices
+            self.live_mask.vindex[live_dim_indices] = True
 
 
 class GridSerializer(Serializer):

src/mdio/segy/_workers.py

@@ -16,6 +16,10 @@
 
     from mdio.core import Grid
 
+import os
+import cProfile
+import pstats
+
 
 def header_scan_worker(segy_file: SegyFile, trace_range: tuple[int, int]) -> HeaderArray:
     """Header scan worker.
@@ -77,45 +81,57 @@ def trace_worker(
     Returns:
         Partial statistics for chunk, or None
     """
-    # Special case where there are no traces inside chunk.
-    live_subset = grid.live_mask[chunk_indices[:-1]]
 
-    if np.count_nonzero(live_subset) == 0:
-        return None
+    profiler = cProfile.Profile()
+    profiler.enable()
+    try:
+
+        # Special case where there are no traces inside chunk.
+        live_subset = grid.live_mask[chunk_indices[:-1]]
+
+        if np.count_nonzero(live_subset) == 0:
+            return None
 
-    # Let's get trace numbers from grid map using the chunk indices.
-    seq_trace_indices = grid.map[chunk_indices[:-1]]
+        # Let's get trace numbers from grid map using the chunk indices.
+        seq_trace_indices = grid.map[chunk_indices[:-1]]
 
-    tmp_data = np.zeros(seq_trace_indices.shape + (grid.shape[-1],), dtype=data_array.dtype)
-    tmp_metadata = np.zeros(seq_trace_indices.shape, dtype=metadata_array.dtype)
+        tmp_data = np.zeros(seq_trace_indices.shape + (grid.shape[-1],), dtype=data_array.dtype)
+        tmp_metadata = np.zeros(seq_trace_indices.shape, dtype=metadata_array.dtype)
 
-    del grid  # To save some memory
+        del grid  # To save some memory
 
-    # Read headers and traces for block
-    valid_indices = seq_trace_indices[live_subset]
+        # Read headers and traces for block
+        valid_indices = seq_trace_indices[live_subset]
 
-    traces = segy_file.trace[valid_indices.tolist()]
-    headers, samples = traces["header"], traces["data"]
+        traces = segy_file.trace[valid_indices.tolist()]
+        headers, samples = traces["header"], traces["data"]
 
-    tmp_metadata[live_subset] = headers.view(tmp_metadata.dtype)
-    tmp_data[live_subset] = samples
+        tmp_metadata[live_subset] = headers.view(tmp_metadata.dtype)
+        tmp_data[live_subset] = samples
 
-    # Flush metadata to zarr
-    metadata_array.set_basic_selection(selection=chunk_indices[:-1], value=tmp_metadata)
+        # Flush metadata to zarr
+        metadata_array.set_basic_selection(selection=chunk_indices[:-1], value=tmp_metadata)
 
-    nonzero_mask = samples != 0
-    nonzero_count = nonzero_mask.sum(dtype="uint32")
+        nonzero_mask = samples != 0
+        nonzero_count = nonzero_mask.sum(dtype="uint32")
 
-    if nonzero_count == 0:
-        return None
+        if nonzero_count == 0:
+            return None
 
-    data_array.set_basic_selection(selection=chunk_indices, value=tmp_data)
+        data_array.set_basic_selection(selection=chunk_indices, value=tmp_data)
 
-    # Calculate statistics
-    tmp_data = samples[nonzero_mask]
-    chunk_sum = tmp_data.sum(dtype="float64")
-    chunk_sum_squares = np.square(tmp_data, dtype="float64").sum()
-    min_val = tmp_data.min()
-    max_val = tmp_data.max()
+        # Calculate statistics
+        tmp_data = samples[nonzero_mask]
+        chunk_sum = tmp_data.sum(dtype="float64")
+        chunk_sum_squares = np.square(tmp_data, dtype="float64").sum()
+        min_val = tmp_data.min()
+        max_val = tmp_data.max()
 
-    return nonzero_count, chunk_sum, chunk_sum_squares, min_val, max_val
+        return nonzero_count, chunk_sum, chunk_sum_squares, min_val, max_val
+    finally:
+        profiler.disable()
+        pid = os.getpid()
+        profile_path = f"/tmp/trace_worker_profile_{pid}.prof"
+        with open(profile_path, "w") as f:
+            ps = pstats.Stats(profiler, stream=f)
+            ps.strip_dirs().sort_stats("cumulative").print_stats()