From d850b3ed10536c0fd5345082b5647daaaca116fa Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Wed, 27 Aug 2025 15:48:07 +0000 Subject: [PATCH 1/9] Coca template take 2 --- src/mdio/converters/segy.py | 97 ++++++++++++++++--- .../v1/templates/seismic_3d_prestack_coca.py | 8 +- .../test_segy_import_export_masked.py | 4 +- tests/unit/v1/converters/test_segy_to_mdio.py | 78 +++++++++++++++ .../test_seismic_3d_prestack_coca.py | 2 +- 5 files changed, 171 insertions(+), 18 deletions(-) create mode 100644 tests/unit/v1/converters/test_segy_to_mdio.py diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 6a1d7df7e..da8128e3c 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -214,20 +214,89 @@ def populate_dim_coordinates( return dataset, drop_vars_delayed -def populate_non_dim_coordinates( +def _check_dimensions_values_identical(arr: np.ndarray, axes_to_check: tuple[int, ...]) -> np.ndarray: + """ + Check if all values along specified dimensions are identical for each + sub-array defined by the other dimensions. + + Args: + arr (np.ndarray): an N-dimensional array. For example, an array of all 'cdp_x' segy + header values for coordinates "inline", "crossline", "offset", "azimuth". + axes_to_check (tuple): A tuple of integers representing the axes to check for + identical values. For example, (2, 3) would check the "offset", "azimuth" + dimensions. + + Returns: + bool: True indicates the all values in the dimensions selected by axes_to_check + are identical, and False otherwise. + """ + + # Create a slicing tuple to get the first element along the axes to check + full_slice = [slice(None)] * arr.ndim + for axis in axes_to_check: + full_slice[axis] = 0 + + # Broadcast the first element along the specified axes for comparison + first_element_slice = arr[tuple(full_slice)] + + # Add new axes to the slice to enable broadcasting against the original array + for axis in axes_to_check: + first_element_slice = np.expand_dims(first_element_slice, axis) + + # Compare the array with the broadcasted slice and use np.all() + # to collapse the dimensions being checked + identical = np.all(np.isclose(arr, first_element_slice), axis=axes_to_check) + return np.all(identical).item() + + +def _populate_non_dim_coordinates( dataset: xr_Dataset, grid: Grid, - coordinates: dict[str, SegyHeaderArray], + coordinate_headers: dict[str, SegyHeaderArray], drop_vars_delayed: list[str], ) -> tuple[xr_Dataset, list[str]]: """Populate the xarray dataset with coordinate variables.""" - not_null = grid.map[:] != UINT32_MAX - for c_name, c_values in coordinates.items(): - c_tmp_array = dataset[c_name].values - c_tmp_array[not_null] = c_values - dataset[c_name][:] = c_tmp_array + coord_is_good = {} + # Load the grid map values into memory. + # Q: Should we be using the trace_mask boolean array instead of grid map? + grid_map_values = grid.map[:] + for c_name, coord_headers_values in coordinate_headers.items(): + + # In the case of Coca and some other templates, the coordinate header values, + # coord_headers_values, have a full set of dimensions (e.g. a 4-tuple of "inline", + # "crossline", "offset", "azimuth"), while the non-dimensional coordinates, (e.g., + # dataset["cdp_x"]) are defined over only a subset of the dimensions (e.g. 2-tuple of + # "inline", "crossline"). + # Thus, every coordinate 2-tuple has multiple duplicate values of the "cdp_x" coordinates + # stored in coord_headers_values. Those needs to be reduced to a unique value. + # We assume (and check) that all the duplicate values are (near) identical. + headers_dims = grid.dim_names[:-1] # e.g.: "inline", "crossline", "offset", "azimuth" + coord_dims = dataset[c_name].dims # e.g.: "inline", "crossline" + # This will create a temporary array in memory with the same shape as the coordinate + # defined in the dataset. Since the coordinate variable has not yet been populated, + # the temporary array will be populated with NaN from the current coordinate values. + tmp_nd_coord_values = dataset[c_name].values + # Create slices for the all grid dimensions that are also the coordinate dimensions. + # For other dimension, select the first element (with index 0) + slices = tuple(slice(None) if name in coord_dims else 0 for name in headers_dims) + # Create a boolean mask for the live trace values with the dimensions of the coordinate + # Q: Should we be using the trace_mask boolean array instead of grid map? + not_null = grid_map_values[slices] != UINT32_MAX + + ch_reshaped = coord_headers_values.reshape(grid_map_values.shape) + # Select a subset of the coordinate_headers that have unique values + cr_reduced = ch_reshaped[slices] + + # Validate the all reduced dimensions had identical values + axes_to_check = tuple(i for i, dim in enumerate(headers_dims) if dim not in coord_dims) + coord_is_good[c_name] = _check_dimensions_values_identical(ch_reshaped, axes_to_check) + + # Save the unique coordinate values for live traces only + tmp_nd_coord_values[not_null] = cr_reduced.ravel() + dataset[c_name][:] = tmp_nd_coord_values + drop_vars_delayed.append(c_name) - return dataset, drop_vars_delayed + return coord_is_good, dataset, drop_vars_delayed def _get_horizontal_coordinate_unit(segy_headers: list[Dimension]) -> AllUnits | None: @@ -268,13 +337,19 @@ def _populate_coordinates( """ drop_vars_delayed = [] # Populate the dimension coordinate variables (1-D arrays) - dataset, vars_to_drop_later = populate_dim_coordinates(dataset, grid, drop_vars_delayed=drop_vars_delayed) + dataset, drop_vars_delayed = populate_dim_coordinates(dataset, grid, drop_vars_delayed=drop_vars_delayed) # Populate the non-dimension coordinate variables (N-dim arrays) - dataset, vars_to_drop_later = populate_non_dim_coordinates( - dataset, grid, coordinates=coords, drop_vars_delayed=drop_vars_delayed + is_good, dataset, drop_vars_delayed = _populate_non_dim_coordinates( + dataset, grid, coordinate_headers=coords, drop_vars_delayed=drop_vars_delayed ) + if not all(is_good.values()): + failed_dims = [key for key, value in is_good.items() if not value] + err = f"Non-dim coordinate(s) {failed_dims} have non-identical values " + \ + f"along reduced dimensions." + raise ValueError(err) + dataset = dataset.drop_vars(drop_vars_delayed) return dataset, drop_vars_delayed diff --git a/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py b/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py index 0903fc770..3f159b4e6 100644 --- a/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py +++ b/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py @@ -12,8 +12,8 @@ class Seismic3DPreStackCocaTemplate(AbstractDatasetTemplate): def __init__(self, domain: str): super().__init__(domain=domain) - self._coord_dim_names = ["inline", "crossline", "offset", "azimuth"] - self._dim_names = [*self._coord_dim_names, self._trace_domain] + self._coord_dim_names = ["inline", "crossline"] # This could be a dictionary specific to every coord + self._dim_names = [*self._coord_dim_names, "offset", "azimuth", self._trace_domain] self._coord_names = ["cdp_x", "cdp_y"] self._var_chunk_shape = [8, 8, 32, 1, 1024] @@ -64,13 +64,13 @@ def _add_coordinates(self) -> None: # Add non-dimension coordinates self._builder.add_coordinate( "cdp_x", - dimensions=["inline", "crossline"], + dimensions=self._coord_dim_names, data_type=ScalarType.FLOAT64, metadata_info=[self._horizontal_coord_unit], ) self._builder.add_coordinate( "cdp_y", - dimensions=["inline", "crossline"], + dimensions=self._coord_dim_names, data_type=ScalarType.FLOAT64, metadata_info=[self._horizontal_coord_unit], ) diff --git a/tests/integration/test_segy_import_export_masked.py b/tests/integration/test_segy_import_export_masked.py index 5c310865f..4bb3f26a3 100644 --- a/tests/integration/test_segy_import_export_masked.py +++ b/tests/integration/test_segy_import_export_masked.py @@ -278,8 +278,8 @@ def test_import(self, test_conf: MaskedExportConfig, export_masked_path: Path) - template_name = "PreStackShotGathers2D" + domain case "3d_streamer": template_name = "PreStackShotGathers3D" + domain - # case "3d_coca": - # templateName = "PostStack3D" + domain + case "3d_coca": + template_name = "PreStackCocaGathers3D" + domain case _: err = f"Unsupported test configuration: {grid_conf.name}" raise ValueError(err) diff --git a/tests/unit/v1/converters/test_segy_to_mdio.py b/tests/unit/v1/converters/test_segy_to_mdio.py new file mode 100644 index 000000000..e0d983899 --- /dev/null +++ b/tests/unit/v1/converters/test_segy_to_mdio.py @@ -0,0 +1,78 @@ +import math +import random +import time +import numpy as np +import xarray as xr + +from segy.arrays import HeaderArray as SegyHeaderArray + +from mdio.converters.segy import _populate_non_dim_coordinates +from mdio.core.dimension import Dimension +from mdio.core.grid import Grid + +def test__populate_non_dim_coordinates() ->None: + + dim_in, dim_xl, dim_of, dim_az, dim_tm = 2, 3, 4, 5, 1 + il = Dimension(name="inline", coords=np.arange(dim_in)) + xl = Dimension(name="crossline", coords=np.arange(dim_xl)) + of = Dimension(name="offset", coords=np.arange(dim_of)) + az = Dimension(name="azimuth", coords=np.arange(dim_az)) + tm = Dimension(name="time", coords=np.arange(dim_tm)) + + r = np.random.rand(dim_in, dim_xl, dim_of, dim_az) + + il_, xl_, of_, az_ = np.meshgrid(il.coords, xl.coords, of.coords, az.coords, indexing='ij') + diff = 1000*il_ + 100*xl_ + 10*of_ + 1*az_ # Different values for the same (il, xl) + same = 1000*il_ + 100*xl_ # Same values for the same (il, xl) + near = 1000*il_ + 100*xl_ + 1e-09 * r # Near same values for the same (il, xl) + # NOTE: near[1][1][1][1]: np.float64(1100.0000000003308) + + data_type=[('inline', ' None: t = Seismic3DPreStackCocaTemplate(domain="time") # Template attributes - assert t._coord_dim_names == ["inline", "crossline", "offset", "azimuth"] + assert t._coord_dim_names == ["inline", "crossline"] assert t._dim_names == ["inline", "crossline", "offset", "azimuth", "time"] assert t._coord_names == ["cdp_x", "cdp_y"] assert t._var_chunk_shape == [8, 8, 32, 1, 1024] From 7ecf44d07c480a92789b9bae1efd019528ded597 Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Wed, 27 Aug 2025 16:16:17 +0000 Subject: [PATCH 2/9] Fix the normal templates --- src/mdio/converters/segy.py | 70 +++++++++++++++++++++---------------- 1 file changed, 39 insertions(+), 31 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index da8128e3c..4580fbec0 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -261,40 +261,48 @@ def _populate_non_dim_coordinates( # Q: Should we be using the trace_mask boolean array instead of grid map? grid_map_values = grid.map[:] for c_name, coord_headers_values in coordinate_headers.items(): - - # In the case of Coca and some other templates, the coordinate header values, - # coord_headers_values, have a full set of dimensions (e.g. a 4-tuple of "inline", - # "crossline", "offset", "azimuth"), while the non-dimensional coordinates, (e.g., - # dataset["cdp_x"]) are defined over only a subset of the dimensions (e.g. 2-tuple of - # "inline", "crossline"). - # Thus, every coordinate 2-tuple has multiple duplicate values of the "cdp_x" coordinates - # stored in coord_headers_values. Those needs to be reduced to a unique value. - # We assume (and check) that all the duplicate values are (near) identical. headers_dims = grid.dim_names[:-1] # e.g.: "inline", "crossline", "offset", "azimuth" coord_dims = dataset[c_name].dims # e.g.: "inline", "crossline" - # This will create a temporary array in memory with the same shape as the coordinate - # defined in the dataset. Since the coordinate variable has not yet been populated, - # the temporary array will be populated with NaN from the current coordinate values. - tmp_nd_coord_values = dataset[c_name].values - # Create slices for the all grid dimensions that are also the coordinate dimensions. - # For other dimension, select the first element (with index 0) - slices = tuple(slice(None) if name in coord_dims else 0 for name in headers_dims) - # Create a boolean mask for the live trace values with the dimensions of the coordinate - # Q: Should we be using the trace_mask boolean array instead of grid map? - not_null = grid_map_values[slices] != UINT32_MAX - - ch_reshaped = coord_headers_values.reshape(grid_map_values.shape) - # Select a subset of the coordinate_headers that have unique values - cr_reduced = ch_reshaped[slices] - - # Validate the all reduced dimensions had identical values axes_to_check = tuple(i for i, dim in enumerate(headers_dims) if dim not in coord_dims) - coord_is_good[c_name] = _check_dimensions_values_identical(ch_reshaped, axes_to_check) - - # Save the unique coordinate values for live traces only - tmp_nd_coord_values[not_null] = cr_reduced.ravel() - dataset[c_name][:] = tmp_nd_coord_values - + if axes_to_check == (): + # In case the coordinate has the same dimensions as grid map + coord_is_good[c_name] = True + not_null = grid_map_values != UINT32_MAX + tmp_coord_values = dataset[c_name].values + tmp_coord_values[not_null] = coord_headers_values + else: + # In the case of Coca and some other templates, the coordinate header values, + # coord_headers_values, have a full set of dimensions (e.g. a 4-tuple of "inline", + # "crossline", "offset", "azimuth"), while the non-dimensional coordinates, (e.g., + # dataset["cdp_x"]) are defined over only a subset of the dimensions (e.g. 2-tuple of + # "inline", "crossline"). + # Thus, every coordinate 2-tuple has multiple duplicate values of the "cdp_x" coordinates + # stored in coord_headers_values. Those needs to be reduced to a unique value. + # We assume (and check) that all the duplicate values are (near) identical. + # + # The following will create a temporary array in memory with the same shape as the + # coordinate defined in the dataset. Since the coordinate variable has not yet been + # populated, the temporary array will be populated with NaN from the current coordinate + # values. + tmp_coord_values = dataset[c_name].values + # Create slices for the all grid dimensions that are also the coordinate dimensions. + # For other dimension, select the first element (with index 0) + slices = tuple(slice(None) if name in coord_dims else 0 for name in headers_dims) + # Create a boolean mask for the live trace values with the dimensions of the coordinate + # Q: Should we be using the trace_mask boolean array instead of grid map? + not_null = grid_map_values[slices] != UINT32_MAX + + ch_reshaped = coord_headers_values.reshape(grid_map_values.shape) + # Select a subset of the coordinate_headers that have unique values + cr_reduced = ch_reshaped[slices] + + # Validate the all reduced dimensions had identical values + coord_is_good[c_name] = _check_dimensions_values_identical(ch_reshaped, axes_to_check) + + # Save the unique coordinate values for live traces only + tmp_coord_values[not_null] = cr_reduced.ravel() + + dataset[c_name][:] = tmp_coord_values drop_vars_delayed.append(c_name) return coord_is_good, dataset, drop_vars_delayed From d7f33ee2693175c2b1e17b31ed972e3446f19dc5 Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Wed, 27 Aug 2025 19:36:20 +0000 Subject: [PATCH 3/9] Update tests --- src/mdio/converters/segy.py | 1 - .../test_segy_import_export_masked.py | 62 ++++++++++++++----- 2 files changed, 47 insertions(+), 16 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 4580fbec0..8b0f187aa 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -357,7 +357,6 @@ def _populate_coordinates( f"along reduced dimensions." raise ValueError(err) - dataset = dataset.drop_vars(drop_vars_delayed) return dataset, drop_vars_delayed diff --git a/tests/integration/test_segy_import_export_masked.py b/tests/integration/test_segy_import_export_masked.py index 4bb3f26a3..0e072dca2 100644 --- a/tests/integration/test_segy_import_export_masked.py +++ b/tests/integration/test_segy_import_export_masked.py @@ -198,17 +198,32 @@ def mock_nd_segy(path: str, grid_conf: GridConfig, segy_factory_conf: SegyFactor headers[dim.name] = dim_grid[dim_idx].ravel() # Fill coordinates (e.g. {SRC-REC-CDP}-X/Y - headers["coord_scalar"] = -100 - for field in ["cdp_x", "source_coord_x", "group_coord_x"]: - start = 700000 - step = 100 - stop = start + step * (trace_numbers.size - 0) - headers[field] = np.arange(start=start, stop=stop, step=step) - for field in ["cdp_y", "source_coord_y", "group_coord_y"]: - start = 4000000 - step = 100 - stop = start + step * (trace_numbers.size - 0) - headers[field] = np.arange(start=start, stop=stop, step=step) + if grid_conf.name == "3d_coca": + coords = [] + for name in ["inline", "crossline", "offset", "azimuth"]: + dim = next(dim for dim in grid_conf.dims if dim.name == name) + coords.append(np.arange(start=dim.start, stop=dim.start + dim.size * dim.step, step=dim.step)) + inline, crossline, _, _ = np.meshgrid(*coords, indexing='ij') + # Make sure that multiple cdp_x and cdp_y values are the same for each (il, xl) + cdp_x = 700000 + 100*inline + crossline + cdp_y = 4000000 + inline + 1000*crossline + headers["coord_scalar"] = -100 + for field in ["cdp_x", "source_coord_x", "group_coord_x"]: + headers[field] = cdp_x.ravel() + for field in ["cdp_y", "source_coord_y", "group_coord_y"]: + headers[field] = cdp_y.ravel() + else: + headers["coord_scalar"] = -100 + for field in ["cdp_x", "source_coord_x", "group_coord_x"]: + start = 700000 + step = 100 + stop = start + step * (trace_numbers.size - 0) + headers[field] = np.arange(start=start, stop=stop, step=step) + for field in ["cdp_y", "source_coord_y", "group_coord_y"]: + start = 4000000 + step = 100 + stop = start + step * (trace_numbers.size - 0) + headers[field] = np.arange(start=start, stop=stop, step=step) # Array filled with repeating sequence sequence = np.array([1, 2, 3]) @@ -328,11 +343,28 @@ def test_ingested_mdio(self, test_conf: MaskedExportConfig, export_masked_path: assert headers.shape == live_mask.shape assert set(headers.dims) == {dim.name for dim in grid_conf.dims} # Validate header values - trace_index = 0 - trace_header = headers.values.flatten()[trace_index] - expected_x = 700000 + trace_index * 100 - expected_y = 4000000 + trace_index * 100 + if grid_conf.name == "3d_coca": + # Select an index of the trace, which header to check: + il, xl, of, az = 0, 0, 0, 0 + trace_index = np.ravel_multi_index((il, xl, of, az), dims=live_mask.shape) + # Get the trace header + trace_header = headers.values[il][xl][of][az] + # calculate the expected values + coords = {} + for name in ["inline", "crossline", "offset", "azimuth"]: + dim = next(dim for dim in grid_conf.dims if dim.name == name) + coords[name] = np.arange(start=dim.start, stop=dim.start + dim.size * dim.step, step=dim.step) + inline = coords["inline"][il] + crossline = coords["crossline"][xl] + expected_x = 700000 + 100*inline + crossline + expected_y = 4000000 + inline + 1000*crossline + else: + trace_index = 0 + trace_header = headers.values.flatten()[trace_index] + expected_x = 700000 + trace_index * 100 + expected_y = 4000000 + trace_index * 100 expected_gun = 1 + (trace_index % 3) + # validate assert trace_header["coord_scalar"] == -100 assert trace_header["source_coord_x"] == expected_x assert trace_header["source_coord_y"] == expected_y From 7161d9276344f1dd6362354457036ae08c9b976d Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Wed, 27 Aug 2025 20:06:24 +0000 Subject: [PATCH 4/9] Fix pre-commit warnings --- src/mdio/converters/segy.py | 45 +++++----- .../v1/templates/seismic_3d_prestack_coca.py | 2 +- .../test_segy_import_export_masked.py | 6 +- tests/unit/v1/converters/test_segy_to_mdio.py | 89 ++++++++++--------- 4 files changed, 76 insertions(+), 66 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 8b0f187aa..74c6a5711 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -214,23 +214,23 @@ def populate_dim_coordinates( return dataset, drop_vars_delayed -def _check_dimensions_values_identical(arr: np.ndarray, axes_to_check: tuple[int, ...]) -> np.ndarray: - """ - Check if all values along specified dimensions are identical for each - sub-array defined by the other dimensions. +def _check_dimensions_values_identical(arr: np.ndarray, axes_to_check: tuple[int, ...]) -> bool: + """Check if all values along specified dimensions are identical. + + Check if all values along specified dimensions are identical for each sub-array + defined by the other dimensions. Args: - arr (np.ndarray): an N-dimensional array. For example, an array of all 'cdp_x' segy + arr: an N-dimensional array. For example, an array of all 'cdp_x' segy header values for coordinates "inline", "crossline", "offset", "azimuth". - axes_to_check (tuple): A tuple of integers representing the axes to check for + axes_to_check: A tuple of integers representing the axes to check for identical values. For example, (2, 3) would check the "offset", "azimuth" dimensions. Returns: - bool: True indicates the all values in the dimensions selected by axes_to_check + True indicates the all values in the dimensions selected by axes_to_check are identical, and False otherwise. """ - # Create a slicing tuple to get the first element along the axes to check full_slice = [slice(None)] * arr.ndim for axis in axes_to_check: @@ -262,7 +262,7 @@ def _populate_non_dim_coordinates( grid_map_values = grid.map[:] for c_name, coord_headers_values in coordinate_headers.items(): headers_dims = grid.dim_names[:-1] # e.g.: "inline", "crossline", "offset", "azimuth" - coord_dims = dataset[c_name].dims # e.g.: "inline", "crossline" + coord_dims = dataset[c_name].dims # e.g.: "inline", "crossline" axes_to_check = tuple(i for i, dim in enumerate(headers_dims) if dim not in coord_dims) if axes_to_check == (): # In case the coordinate has the same dimensions as grid map @@ -271,18 +271,18 @@ def _populate_non_dim_coordinates( tmp_coord_values = dataset[c_name].values tmp_coord_values[not_null] = coord_headers_values else: - # In the case of Coca and some other templates, the coordinate header values, - # coord_headers_values, have a full set of dimensions (e.g. a 4-tuple of "inline", - # "crossline", "offset", "azimuth"), while the non-dimensional coordinates, (e.g., - # dataset["cdp_x"]) are defined over only a subset of the dimensions (e.g. 2-tuple of - # "inline", "crossline"). + # In the case of Coca and some other templates, the coordinate header values, + # coord_headers_values, have a full set of dimensions (e.g. a 4-tuple of "inline", + # "crossline", "offset", "azimuth"), while the non-dimensional coordinates, (e.g., + # dataset["cdp_x"]) are defined over only a subset of the dimensions (e.g. 2-tuple of + # "inline", "crossline"). # Thus, every coordinate 2-tuple has multiple duplicate values of the "cdp_x" coordinates - # stored in coord_headers_values. Those needs to be reduced to a unique value. + # stored in coord_headers_values. Those needs to be reduced to a unique value. # We assume (and check) that all the duplicate values are (near) identical. # - # The following will create a temporary array in memory with the same shape as the - # coordinate defined in the dataset. Since the coordinate variable has not yet been - # populated, the temporary array will be populated with NaN from the current coordinate + # The following will create a temporary array in memory with the same shape as the + # coordinate defined in the dataset. Since the coordinate variable has not yet been + # populated, the temporary array will be populated with NaN from the current coordinate # values. tmp_coord_values = dataset[c_name].values # Create slices for the all grid dimensions that are also the coordinate dimensions. @@ -340,8 +340,12 @@ def _populate_coordinates( grid: The grid object containing the grid map. coords: The non-dim coordinates to populate. + Raises: + ValueError: If coordinate headers have non-identical values for the reduced dimensions. + Returns: - Xarray dataset with filled coordinates and updated variables to drop after writing + A tuple of the Xarray dataset with filled coordinates and updated variables to drop + after writing """ drop_vars_delayed = [] # Populate the dimension coordinate variables (1-D arrays) @@ -353,8 +357,7 @@ def _populate_coordinates( ) if not all(is_good.values()): failed_dims = [key for key, value in is_good.items() if not value] - err = f"Non-dim coordinate(s) {failed_dims} have non-identical values " + \ - f"along reduced dimensions." + err = f"Non-dim coordinate(s) {failed_dims} have non-identical values " + "along reduced dimensions." raise ValueError(err) return dataset, drop_vars_delayed diff --git a/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py b/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py index 3f159b4e6..b6c4b9a28 100644 --- a/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py +++ b/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py @@ -12,7 +12,7 @@ class Seismic3DPreStackCocaTemplate(AbstractDatasetTemplate): def __init__(self, domain: str): super().__init__(domain=domain) - self._coord_dim_names = ["inline", "crossline"] # This could be a dictionary specific to every coord + self._coord_dim_names = ["inline", "crossline"] # This could be a dictionary specific to every coord self._dim_names = [*self._coord_dim_names, "offset", "azimuth", self._trace_domain] self._coord_names = ["cdp_x", "cdp_y"] self._var_chunk_shape = [8, 8, 32, 1, 1024] diff --git a/tests/integration/test_segy_import_export_masked.py b/tests/integration/test_segy_import_export_masked.py index 0e072dca2..1129a1264 100644 --- a/tests/integration/test_segy_import_export_masked.py +++ b/tests/integration/test_segy_import_export_masked.py @@ -203,10 +203,10 @@ def mock_nd_segy(path: str, grid_conf: GridConfig, segy_factory_conf: SegyFactor for name in ["inline", "crossline", "offset", "azimuth"]: dim = next(dim for dim in grid_conf.dims if dim.name == name) coords.append(np.arange(start=dim.start, stop=dim.start + dim.size * dim.step, step=dim.step)) - inline, crossline, _, _ = np.meshgrid(*coords, indexing='ij') + inline, crossline, _, _ = np.meshgrid(*coords, indexing="ij") # Make sure that multiple cdp_x and cdp_y values are the same for each (il, xl) - cdp_x = 700000 + 100*inline + crossline - cdp_y = 4000000 + inline + 1000*crossline + cdp_x = 700000 + 100 * inline + crossline + cdp_y = 4000000 + inline + 1000 * crossline headers["coord_scalar"] = -100 for field in ["cdp_x", "source_coord_x", "group_coord_x"]: headers[field] = cdp_x.ravel() diff --git a/tests/unit/v1/converters/test_segy_to_mdio.py b/tests/unit/v1/converters/test_segy_to_mdio.py index e0d983899..1d1c4fe0c 100644 --- a/tests/unit/v1/converters/test_segy_to_mdio.py +++ b/tests/unit/v1/converters/test_segy_to_mdio.py @@ -1,40 +1,47 @@ -import math -import random -import time +"""Tests for segY-to_mdio methods.""" + import numpy as np import xarray as xr - from segy.arrays import HeaderArray as SegyHeaderArray from mdio.converters.segy import _populate_non_dim_coordinates from mdio.core.dimension import Dimension from mdio.core.grid import Grid -def test__populate_non_dim_coordinates() ->None: +def test__populate_non_dim_coordinates() -> None: + """Test population of non-dimensional coordinates.""" dim_in, dim_xl, dim_of, dim_az, dim_tm = 2, 3, 4, 5, 1 il = Dimension(name="inline", coords=np.arange(dim_in)) xl = Dimension(name="crossline", coords=np.arange(dim_xl)) of = Dimension(name="offset", coords=np.arange(dim_of)) az = Dimension(name="azimuth", coords=np.arange(dim_az)) tm = Dimension(name="time", coords=np.arange(dim_tm)) - + r = np.random.rand(dim_in, dim_xl, dim_of, dim_az) - il_, xl_, of_, az_ = np.meshgrid(il.coords, xl.coords, of.coords, az.coords, indexing='ij') - diff = 1000*il_ + 100*xl_ + 10*of_ + 1*az_ # Different values for the same (il, xl) - same = 1000*il_ + 100*xl_ # Same values for the same (il, xl) - near = 1000*il_ + 100*xl_ + 1e-09 * r # Near same values for the same (il, xl) + il_, xl_, of_, az_ = np.meshgrid(il.coords, xl.coords, of.coords, az.coords, indexing="ij") + diff = 1000 * il_ + 100 * xl_ + 10 * of_ + 1 * az_ # Different values for the same (il, xl) + same = 1000 * il_ + 100 * xl_ # Same values for the same (il, xl) + near = 1000 * il_ + 100 * xl_ + 1e-09 * r # Near same values for the same (il, xl) # NOTE: near[1][1][1][1]: np.float64(1100.0000000003308) - data_type=[('inline', 'None: ds = xr.Dataset( data_vars={ - 'amplitude': (["inline", "crossline", "offset", "azimuth", "time"], - np.zeros((dim_in, dim_xl, dim_of, dim_az, dim_tm), dtype=np.float32)), + "amplitude": ( + ["inline", "crossline", "offset", "azimuth", "time"], + np.zeros((dim_in, dim_xl, dim_of, dim_az, dim_tm), dtype=np.float32), + ), }, coords={ # Define coordinates with their dimensions and values - 'inline': il.coords, - 'crossline': xl.coords, - 'offset': of.coords, - 'azimuth': az.coords, - 'time': tm.coords, - 'cdp_x': (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), - 'cdp_y': (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), - 'cdp_z': (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)) - } + "inline": il.coords, + "crossline": xl.coords, + "offset": of.coords, + "azimuth": az.coords, + "time": tm.coords, + "cdp_x": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), + "cdp_y": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), + "cdp_z": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), + }, ) coordinate_headers: dict[str, SegyHeaderArray] = { - 'cdp_x': segy_headers['cdp_x'], - 'cdp_y': segy_headers['cdp_y'], - 'cdp_z': segy_headers['cdp_z'] + "cdp_x": segy_headers["cdp_x"], + "cdp_y": segy_headers["cdp_y"], + "cdp_z": segy_headers["cdp_z"], } is_good, ds_populated, _ = _populate_non_dim_coordinates(ds, grid, coordinate_headers, []) - assert is_good['cdp_x'] is False - assert is_good['cdp_y'] is True - assert is_good['cdp_z'] is True - expected_values = np.array([[0., 100., 200.], [1000., 1100., 1200.]], dtype=np.float32) - assert np.allclose(ds_populated['cdp_x'].values, expected_values) - assert np.allclose(ds_populated['cdp_y'].values, expected_values) - assert np.allclose(ds_populated['cdp_z'].values, expected_values) + assert is_good["cdp_x"] is False + assert is_good["cdp_y"] is True + assert is_good["cdp_z"] is True + expected_values = np.array([[0.0, 100.0, 200.0], [1000.0, 1100.0, 1200.0]], dtype=np.float32) + assert np.allclose(ds_populated["cdp_x"].values, expected_values) + assert np.allclose(ds_populated["cdp_y"].values, expected_values) + assert np.allclose(ds_populated["cdp_z"].values, expected_values) # NOTE: ds_populated['cdp_z'].values[1][1]: np.float64(1100.0000000008847) - - pass From bf4fafd9395eaa892135402e1cbf5cac37f9b68c Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Thu, 28 Aug 2025 20:51:11 +0000 Subject: [PATCH 5/9] Address PR review issues --- src/mdio/converters/segy.py | 39 ++++++----------- .../v1/templates/seismic_3d_prestack_coca.py | 2 +- tests/unit/v1/converters/test_segy_to_mdio.py | 43 +++++++++++-------- 3 files changed, 38 insertions(+), 46 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 74c6a5711..c2f9f9545 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -245,7 +245,7 @@ def _check_dimensions_values_identical(arr: np.ndarray, axes_to_check: tuple[int # Compare the array with the broadcasted slice and use np.all() # to collapse the dimensions being checked - identical = np.all(np.isclose(arr, first_element_slice), axis=axes_to_check) + identical = np.all(np.isclose(arr, first_element_slice, rtol=1e-05, atol=1e-08), axis=axes_to_check) return np.all(identical).item() @@ -256,18 +256,14 @@ def _populate_non_dim_coordinates( drop_vars_delayed: list[str], ) -> tuple[xr_Dataset, list[str]]: """Populate the xarray dataset with coordinate variables.""" - coord_is_good = {} # Load the grid map values into memory. - # Q: Should we be using the trace_mask boolean array instead of grid map? - grid_map_values = grid.map[:] for c_name, coord_headers_values in coordinate_headers.items(): headers_dims = grid.dim_names[:-1] # e.g.: "inline", "crossline", "offset", "azimuth" coord_dims = dataset[c_name].dims # e.g.: "inline", "crossline" axes_to_check = tuple(i for i, dim in enumerate(headers_dims) if dim not in coord_dims) if axes_to_check == (): # In case the coordinate has the same dimensions as grid map - coord_is_good[c_name] = True - not_null = grid_map_values != UINT32_MAX + not_null = grid.map[:] != UINT32_MAX tmp_coord_values = dataset[c_name].values tmp_coord_values[not_null] = coord_headers_values else: @@ -282,29 +278,27 @@ def _populate_non_dim_coordinates( # # The following will create a temporary array in memory with the same shape as the # coordinate defined in the dataset. Since the coordinate variable has not yet been - # populated, the temporary array will be populated with NaN from the current coordinate - # values. + # populated, the temporary array will be populated with _FillValue from the current + # coordinate values. tmp_coord_values = dataset[c_name].values # Create slices for the all grid dimensions that are also the coordinate dimensions. # For other dimension, select the first element (with index 0) slices = tuple(slice(None) if name in coord_dims else 0 for name in headers_dims) # Create a boolean mask for the live trace values with the dimensions of the coordinate - # Q: Should we be using the trace_mask boolean array instead of grid map? - not_null = grid_map_values[slices] != UINT32_MAX - - ch_reshaped = coord_headers_values.reshape(grid_map_values.shape) + not_null = grid.map[slices] != UINT32_MAX + ch_reshaped = coord_headers_values.reshape(grid.map.shape) # Select a subset of the coordinate_headers that have unique values - cr_reduced = ch_reshaped[slices] + # and save the unique coordinate values for live traces only + tmp_coord_values[not_null] = ch_reshaped[slices].ravel() # Validate the all reduced dimensions had identical values - coord_is_good[c_name] = _check_dimensions_values_identical(ch_reshaped, axes_to_check) - - # Save the unique coordinate values for live traces only - tmp_coord_values[not_null] = cr_reduced.ravel() + if not _check_dimensions_values_identical(ch_reshaped, axes_to_check): + err = f"Coordinate '{c_name}' has non-identical values along reduced dimensions." + raise ValueError(err) dataset[c_name][:] = tmp_coord_values drop_vars_delayed.append(c_name) - return coord_is_good, dataset, drop_vars_delayed + return dataset, drop_vars_delayed def _get_horizontal_coordinate_unit(segy_headers: list[Dimension]) -> AllUnits | None: @@ -340,9 +334,6 @@ def _populate_coordinates( grid: The grid object containing the grid map. coords: The non-dim coordinates to populate. - Raises: - ValueError: If coordinate headers have non-identical values for the reduced dimensions. - Returns: A tuple of the Xarray dataset with filled coordinates and updated variables to drop after writing @@ -352,13 +343,9 @@ def _populate_coordinates( dataset, drop_vars_delayed = populate_dim_coordinates(dataset, grid, drop_vars_delayed=drop_vars_delayed) # Populate the non-dimension coordinate variables (N-dim arrays) - is_good, dataset, drop_vars_delayed = _populate_non_dim_coordinates( + dataset, drop_vars_delayed = _populate_non_dim_coordinates( dataset, grid, coordinate_headers=coords, drop_vars_delayed=drop_vars_delayed ) - if not all(is_good.values()): - failed_dims = [key for key, value in is_good.items() if not value] - err = f"Non-dim coordinate(s) {failed_dims} have non-identical values " + "along reduced dimensions." - raise ValueError(err) return dataset, drop_vars_delayed diff --git a/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py b/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py index b6c4b9a28..94328a13c 100644 --- a/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py +++ b/src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py @@ -12,7 +12,7 @@ class Seismic3DPreStackCocaTemplate(AbstractDatasetTemplate): def __init__(self, domain: str): super().__init__(domain=domain) - self._coord_dim_names = ["inline", "crossline"] # This could be a dictionary specific to every coord + self._coord_dim_names = ["inline", "crossline"] self._dim_names = [*self._coord_dim_names, "offset", "azimuth", self._trace_domain] self._coord_names = ["cdp_x", "cdp_y"] self._var_chunk_shape = [8, 8, 32, 1, 1024] diff --git a/tests/unit/v1/converters/test_segy_to_mdio.py b/tests/unit/v1/converters/test_segy_to_mdio.py index 1d1c4fe0c..61f87288b 100644 --- a/tests/unit/v1/converters/test_segy_to_mdio.py +++ b/tests/unit/v1/converters/test_segy_to_mdio.py @@ -1,6 +1,7 @@ -"""Tests for segY-to_mdio methods.""" +"""Tests for segy-to_mdio convenience functions.""" import numpy as np +import pytest import xarray as xr from segy.arrays import HeaderArray as SegyHeaderArray @@ -31,9 +32,9 @@ def test__populate_non_dim_coordinates() -> None: ("crossline", " None: "offset": of.coords, "azimuth": az.coords, "time": tm.coords, - "cdp_x": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), - "cdp_y": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), - "cdp_z": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), + "cdp_diff": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), + "cdp_same": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), + "cdp_near": (["inline", "crossline"], np.zeros((dim_in, dim_xl), dtype=np.float64)), }, ) + # "cdp_diff" has different values for the same (il, xl) coordinate_headers: dict[str, SegyHeaderArray] = { - "cdp_x": segy_headers["cdp_x"], - "cdp_y": segy_headers["cdp_y"], - "cdp_z": segy_headers["cdp_z"], + "cdp_diff": segy_headers["cdp_diff"], } + with pytest.raises(ValueError) as exc_info: + ds_populated, _ = _populate_non_dim_coordinates(ds, grid, coordinate_headers, []) + err = str(exc_info.value) + assert "Coordinate 'cdp_diff' have non-identical values along reduced dimensions." in err - is_good, ds_populated, _ = _populate_non_dim_coordinates(ds, grid, coordinate_headers, []) - - assert is_good["cdp_x"] is False - assert is_good["cdp_y"] is True - assert is_good["cdp_z"] is True + # "cdp_same" has identical values for the same (il, xl) + # "cdp_near" has near identical values for the same (il, xl) + coordinate_headers: dict[str, SegyHeaderArray] = { + "cdp_same": segy_headers["cdp_same"], + "cdp_near": segy_headers["cdp_near"], + } + ds_populated, _ = _populate_non_dim_coordinates(ds, grid, coordinate_headers, []) expected_values = np.array([[0.0, 100.0, 200.0], [1000.0, 1100.0, 1200.0]], dtype=np.float32) - assert np.allclose(ds_populated["cdp_x"].values, expected_values) - assert np.allclose(ds_populated["cdp_y"].values, expected_values) - assert np.allclose(ds_populated["cdp_z"].values, expected_values) - # NOTE: ds_populated['cdp_z'].values[1][1]: np.float64(1100.0000000008847) + assert np.allclose(ds_populated["cdp_same"].values, expected_values) + assert np.allclose(ds_populated["cdp_near"].values, expected_values) + # NOTE: ds_populated['cdp_near'].values[1][1]: np.float64(1100.0000000008847) From 1986e77532f69f4fb8b7b71b2ab97d12e78e128a Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Tue, 2 Sep 2025 15:36:15 +0000 Subject: [PATCH 6/9] Fix PR review comment --- src/mdio/converters/segy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index c2f9f9545..7f3e16242 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -263,7 +263,7 @@ def _populate_non_dim_coordinates( axes_to_check = tuple(i for i, dim in enumerate(headers_dims) if dim not in coord_dims) if axes_to_check == (): # In case the coordinate has the same dimensions as grid map - not_null = grid.map[:] != UINT32_MAX + not_null = grid.map[:] != grid.map.fill_value tmp_coord_values = dataset[c_name].values tmp_coord_values[not_null] = coord_headers_values else: @@ -285,7 +285,7 @@ def _populate_non_dim_coordinates( # For other dimension, select the first element (with index 0) slices = tuple(slice(None) if name in coord_dims else 0 for name in headers_dims) # Create a boolean mask for the live trace values with the dimensions of the coordinate - not_null = grid.map[slices] != UINT32_MAX + not_null = grid.map[slices] != grid.map.fill_value ch_reshaped = coord_headers_values.reshape(grid.map.shape) # Select a subset of the coordinate_headers that have unique values # and save the unique coordinate values for live traces only From 96040b414631258506c12b5795effec327c7fab4 Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Tue, 2 Sep 2025 15:42:31 +0000 Subject: [PATCH 7/9] Fix pre-commit error --- src/mdio/converters/segy.py | 1 - tests/unit/v1/converters/test_segy_to_mdio.py | 5 ++--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 7f3e16242..28a6be07b 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -12,7 +12,6 @@ from segy.standards.codes import MeasurementSystem as segy_MeasurementSystem from segy.standards.fields.trace import Rev0 as TraceHeaderFieldsRev0 -from mdio.constants import UINT32_MAX from mdio.converters.exceptions import EnvironmentFormatError from mdio.converters.exceptions import GridTraceCountError from mdio.converters.exceptions import GridTraceSparsityError diff --git a/tests/unit/v1/converters/test_segy_to_mdio.py b/tests/unit/v1/converters/test_segy_to_mdio.py index 61f87288b..6c06d6ac0 100644 --- a/tests/unit/v1/converters/test_segy_to_mdio.py +++ b/tests/unit/v1/converters/test_segy_to_mdio.py @@ -72,10 +72,9 @@ def test__populate_non_dim_coordinates() -> None: coordinate_headers: dict[str, SegyHeaderArray] = { "cdp_diff": segy_headers["cdp_diff"], } - with pytest.raises(ValueError) as exc_info: + expected_err = "Coordinate 'cdp_diff' has non-identical values along reduced dimensions." + with pytest.raises(ValueError, match=expected_err): ds_populated, _ = _populate_non_dim_coordinates(ds, grid, coordinate_headers, []) - err = str(exc_info.value) - assert "Coordinate 'cdp_diff' have non-identical values along reduced dimensions." in err # "cdp_same" has identical values for the same (il, xl) # "cdp_near" has near identical values for the same (il, xl) From ffa6abff797632ce91e4efd10bcd390516b3916b Mon Sep 17 00:00:00 2001 From: Dmitriy Repin Date: Tue, 2 Sep 2025 20:55:53 +0000 Subject: [PATCH 8/9] Merge 'segy' --- src/mdio/converters/mdio.py | 86 +++++++---------- src/mdio/core/utils_read.py | 16 ++++ .../v1/templates/abstract_dataset_template.py | 4 +- src/mdio/segy/creation.py | 82 +++++++--------- tests/conftest.py | 25 +++-- tests/integration/conftest.py | 38 ++++---- tests/integration/test_segy_import_export.py | 95 +++++++++--------- .../test_segy_import_export_masked.py | 96 ++++++++++++------- tests/integration/testing_data.py | 33 +++++++ tests/integration/testing_helpers.py | 36 ++++++- 10 files changed, 294 insertions(+), 217 deletions(-) create mode 100644 src/mdio/core/utils_read.py diff --git a/src/mdio/converters/mdio.py b/src/mdio/converters/mdio.py index 9726cd6b8..41fc4dd33 100644 --- a/src/mdio/converters/mdio.py +++ b/src/mdio/converters/mdio.py @@ -5,34 +5,37 @@ import os from pathlib import Path from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING +import dask.array as da import numpy as np +import xarray as xr from psutil import cpu_count from tqdm.dask import TqdmCallback -from mdio import MDIOReader +from mdio.core.utils_read import open_zarr_dataset from mdio.segy.blocked_io import to_segy from mdio.segy.creation import concat_files from mdio.segy.creation import mdio_spec_to_segy -from mdio.segy.utilities import segy_export_rechunker +from mdio.segy.utilities import segy_export_rechunker try: import distributed except ImportError: distributed = None +if TYPE_CHECKING: + from segy.schema import SegySpec + + from mdio.core.storage_location import StorageLocation default_cpus = cpu_count(logical=True) NUM_CPUS = int(os.getenv("MDIO__EXPORT__CPU_COUNT", default_cpus)) - -def mdio_to_segy( # noqa: PLR0912, PLR0913 - mdio_path_or_buffer: str, - output_segy_path: str, - endian: str = "big", - access_pattern: str = "012", - storage_options: dict = None, - new_chunks: tuple[int, ...] = None, +def mdio_to_segy( # noqa: PLR0912, PLR0913, PLR0915 + segy_spec: SegySpec, + input_location: StorageLocation, + output_location: StorageLocation, selection_mask: np.ndarray = None, client: distributed.Client = None, ) -> None: @@ -47,13 +50,9 @@ def mdio_to_segy( # noqa: PLR0912, PLR0913 A `selection_mask` can be provided (same shape as spatial grid) to export a subset. Args: - mdio_path_or_buffer: Input path where the MDIO is located. - output_segy_path: Path to the output SEG-Y file. - endian: Endianness of the input SEG-Y. Rev.2 allows little endian. Default is 'big'. - access_pattern: This specificies the chunk access pattern. Underlying zarr.Array must - exist. Examples: '012', '01' - storage_options: Storage options for the cloud storage backend. Default: None (anonymous) - new_chunks: Set manual chunksize. For development purposes only. + segy_spec: The SEG-Y specification to use for the conversion. + input_location: Store or URL (and cloud options) for MDIO file. + output_location: Path to the output SEG-Y file. selection_mask: Array that lists the subset of traces client: Dask client. If `None` we will use local threaded scheduler. If `auto` is used we will create multiple processes (with 8 threads each). @@ -64,8 +63,7 @@ def mdio_to_segy( # noqa: PLR0912, PLR0913 Examples: To export an existing local MDIO file to SEG-Y we use the code snippet below. This will - export the full MDIO (without padding) to SEG-Y format using IBM floats and big-endian - byte order. + export the full MDIO (without padding) to SEG-Y format. >>> from mdio import mdio_to_segy >>> @@ -74,51 +72,32 @@ def mdio_to_segy( # noqa: PLR0912, PLR0913 ... mdio_path_or_buffer="prefix2/file.mdio", ... output_segy_path="prefix/file.segy", ... ) - - If we want to export this as an IEEE big-endian, using a selection mask, we would run: - - >>> mdio_to_segy( - ... mdio_path_or_buffer="prefix2/file.mdio", - ... output_segy_path="prefix/file.segy", - ... selection_mask=boolean_mask, - ... ) - """ - backend = "dask" - - output_segy_path = Path(output_segy_path) + output_segy_path = Path(output_location.uri) - mdio = MDIOReader( - mdio_path_or_buffer=mdio_path_or_buffer, - access_pattern=access_pattern, - storage_options=storage_options, - ) + mdio_xr = open_zarr_dataset(input_location) - if new_chunks is None: - new_chunks = segy_export_rechunker(mdio.chunks, mdio.shape, mdio._traces.dtype) + trace_variable_name = mdio_xr.attrs["attributes"]["traceVariableName"] + amplitude = mdio_xr[trace_variable_name] + chunks: tuple[int, ...] = amplitude.data.chunksize + shape: tuple[int, ...] = amplitude.data.shape + dtype = amplitude.dtype + new_chunks = segy_export_rechunker(chunks, shape, dtype) - creation_args = [ - mdio_path_or_buffer, - output_segy_path, - access_pattern, - endian, - storage_options, - new_chunks, - backend, - ] + creation_args = [segy_spec, input_location, output_location] if client is not None: if distributed is not None: # This is in case we work with big data feature = client.submit(mdio_spec_to_segy, *creation_args) - mdio, segy_factory = feature.result() + mdio_xr, segy_factory = feature.result() else: msg = "Distributed client was provided, but `distributed` is not installed" raise ImportError(msg) else: - mdio, segy_factory = mdio_spec_to_segy(*creation_args) + mdio_xr, segy_factory = mdio_spec_to_segy(*creation_args) - live_mask = mdio.live_mask.compute() + live_mask = mdio_xr["trace_mask"].data.compute() if selection_mask is not None: live_mask = live_mask & selection_mask @@ -138,8 +117,9 @@ def mdio_to_segy( # noqa: PLR0912, PLR0913 dim_slices += (slice(start, stop),) # Lazily pull the data with limits now, and limit mask so its the same shape. - live_mask, headers, samples = mdio[dim_slices] - live_mask = live_mask.rechunk(headers.chunks) + live_mask = mdio_xr["trace_mask"].data.rechunk(new_chunks[:-1])[dim_slices] + headers = mdio_xr["headers"].data.rechunk(new_chunks[:-1])[dim_slices] + samples = mdio_xr["amplitude"].data.rechunk(new_chunks)[dim_slices] if selection_mask is not None: selection_mask = selection_mask[dim_slices] @@ -170,4 +150,4 @@ def mdio_to_segy( # noqa: PLR0912, PLR0913 if client is not None: _ = client.submit(concat_files, paths=ordered_files).result() else: - concat_files(paths=ordered_files, progress=True) + concat_files(paths=ordered_files, progress=True) \ No newline at end of file diff --git a/src/mdio/core/utils_read.py b/src/mdio/core/utils_read.py new file mode 100644 index 000000000..6c76d2474 --- /dev/null +++ b/src/mdio/core/utils_read.py @@ -0,0 +1,16 @@ +"""Utils for reading MDIO dataset""" + +from typing import TYPE_CHECKING + +from xarray import open_dataset as xr_open_dataset +from mdio.core.storage_location import StorageLocation +from xarray import Dataset as xr_Dataset + +def open_zarr_dataset(storage_location: StorageLocation) -> xr_Dataset: + """ + Open a Zarr dataset from the specified storage location. + """ + # NOTE: If mask_and_scale is not set, + # Xarray will convert int to float and replace _FillValue with NaN + # We are using chunks={} to force Xarray to create Dask arrays + return xr_open_dataset(storage_location.uri, engine="zarr", chunks={}, mask_and_scale=False) \ No newline at end of file diff --git a/src/mdio/schemas/v1/templates/abstract_dataset_template.py b/src/mdio/schemas/v1/templates/abstract_dataset_template.py index 7bd39ff2d..ef10ecbe7 100644 --- a/src/mdio/schemas/v1/templates/abstract_dataset_template.py +++ b/src/mdio/schemas/v1/templates/abstract_dataset_template.py @@ -84,7 +84,9 @@ def build_dataset( self._dim_sizes = sizes self._horizontal_coord_unit = horizontal_coord_unit - self._builder = MDIODatasetBuilder(name=name, attributes=self._load_dataset_attributes()) + attr = self._load_dataset_attributes() or UserAttributes(attributes={}) + attr.attributes["traceVariableName"] = self._trace_variable_name + self._builder = MDIODatasetBuilder(name=name, attributes=attr) self._add_dimensions() self._add_coordinates() self._add_variables() diff --git a/src/mdio/segy/creation.py b/src/mdio/segy/creation.py index b1daf7d5c..903a1c07a 100644 --- a/src/mdio/segy/creation.py +++ b/src/mdio/segy/creation.py @@ -7,98 +7,80 @@ from pathlib import Path from shutil import copyfileobj from typing import TYPE_CHECKING -from typing import Any import numpy as np +import xarray as xr from segy.factory import SegyFactory -from segy.schema import Endianness from segy.schema import SegySpec from tqdm.auto import tqdm +from xarray import Dataset as xr_Dataset -from mdio.api.accessor import MDIOReader -from mdio.segy.compat import mdio_segy_spec +from mdio.core.utils_read import open_zarr_dataset from mdio.segy.compat import revision_encode if TYPE_CHECKING: from numpy.typing import NDArray + from mdio.core.storage_location import StorageLocation + logger = logging.getLogger(__name__) -def make_segy_factory(mdio: MDIOReader, spec: SegySpec) -> SegyFactory: +def make_segy_factory(mdio_xr: xr_Dataset, spec: SegySpec) -> SegyFactory: """Generate SEG-Y factory from MDIO metadata.""" - grid = mdio.grid - sample_dim = grid.select_dim("sample") - sample_interval = sample_dim[1] - sample_dim[0] - samples_per_trace = len(sample_dim) - + binary_header = mdio_xr.attrs["attributes"]["binaryHeader"] + sample_interval = binary_header["sample_interval"] + samples_per_trace = binary_header["samples_per_trace"] return SegyFactory( spec=spec, - sample_interval=sample_interval * 1000, + sample_interval=sample_interval, # Sample Interval is already in milliseconds samples_per_trace=samples_per_trace, ) -def mdio_spec_to_segy( # noqa: PLR0913 - mdio_path_or_buffer: str, - output_segy_path: Path, - access_pattern: str, - output_endian: str, - storage_options: dict[str, Any], - new_chunks: tuple[int, ...], - backend: str, -) -> tuple[MDIOReader, SegyFactory]: +def mdio_spec_to_segy( + segy_spec: SegySpec, + input_location: StorageLocation, + output_location: StorageLocation +) -> tuple[xr_Dataset, SegyFactory]: """Create SEG-Y file without any traces given MDIO specification. This function opens an MDIO file, gets some relevant information for SEG-Y files, then creates a SEG-Y file with the specification it read from the MDIO file. - It then returns the `MDIOReader` instance, and the parsed floating point format `sample_format` - for further use. + It then returns the Xarray Dataset instance and SegyFactory for further use. - Function will attempt to read text, and binary headers, and some grid information from the MDIO - file. If these don't exist, the process will fail. + Function will attempt to read text, and binary headers information from the MDIO file. + If these don't exist, the process will fail. Args: - mdio_path_or_buffer: Store or URL for MDIO file. - output_segy_path: Path to the output SEG-Y file. - access_pattern: Chunk access pattern, optional. Default is "012". Examples: '012', '01'. - output_endian: Endianness of the output file. - storage_options: Options for the storage backend. By default, system-wide credentials - will be used. - new_chunks: Set manual chunksize. For development purposes only. - backend: Backend selection, optional. Default is "zarr". Must be in {'zarr', 'dask'}. + segy_spec: The SEG-Y specification to use for the conversion. + input_location: Store or URL (and cloud options) for MDIO file. + output_location: Path to the output SEG-Y file. Returns: - Initialized MDIOReader for MDIO file and return SegyFactory + Opened Xarray Dataset for MDIO file and SegyFactory """ - mdio = MDIOReader( - mdio_path_or_buffer=mdio_path_or_buffer, - access_pattern=access_pattern, - storage_options=storage_options, - return_metadata=True, - new_chunks=new_chunks, - backend=backend, - disk_cache=False, # Making sure disk caching is disabled - ) + mdio_xr = open_zarr_dataset(input_location) + factory = make_segy_factory(mdio_xr, spec=segy_spec) - mdio_file_version = mdio.root.attrs["api_version"] - spec = mdio_segy_spec(mdio_file_version) - spec.endianness = Endianness(output_endian) - factory = make_segy_factory(mdio, spec=spec) + attr = mdio_xr.attrs["attributes"] - text_str = "\n".join(mdio.text_header) + txt_header = attr["textHeader"] + text_str = "\n".join(txt_header) text_bytes = factory.create_textual_header(text_str) - binary_header = revision_encode(mdio.binary_header, mdio_file_version) + bin_header = attr["binaryHeader"] + mdio_file_version = mdio_xr.attrs["apiVersion"] + binary_header = revision_encode(bin_header, mdio_file_version) bin_hdr_bytes = factory.create_binary_header(binary_header) - with output_segy_path.open(mode="wb") as fp: + with Path(output_location.uri).open(mode="wb") as fp: fp.write(text_bytes) fp.write(bin_hdr_bytes) - return mdio, factory + return mdio_xr, factory @dataclass(slots=True) diff --git a/tests/conftest.py b/tests/conftest.py index 6f4f2d1f3..24bc06ef9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,14 +3,12 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING +from pathlib import Path from urllib.request import urlretrieve import pytest -if TYPE_CHECKING: - from pathlib import Path - +DEBUG_MODE = False # Suppress Dask's chunk balancing warning warnings.filterwarnings( @@ -24,6 +22,8 @@ @pytest.fixture(scope="session") def fake_segy_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path: """Make a temp file for the fake SEG-Y files we are going to create.""" + if DEBUG_MODE: + return Path("TMP/fake_segy") return tmp_path_factory.mktemp(r"fake_segy") @@ -36,27 +36,38 @@ def segy_input_uri() -> str: @pytest.fixture(scope="session") def segy_input(segy_input_uri: str, tmp_path_factory: pytest.TempPathFactory) -> Path: """Download teapot dome dataset for testing.""" - tmp_dir = tmp_path_factory.mktemp("segy") + if DEBUG_MODE: + tmp_dir = Path("TMP/segy") + tmp_dir.mkdir(parents=True, exist_ok=True) + else: + tmp_dir = tmp_path_factory.mktemp("segy") tmp_file = tmp_dir / "teapot.segy" urlretrieve(segy_input_uri, tmp_file) # noqa: S310 - return tmp_file @pytest.fixture(scope="module") def zarr_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path: """Make a temp file for the output MDIO.""" + if DEBUG_MODE: + return Path("TMP/mdio") return tmp_path_factory.mktemp(r"mdio") @pytest.fixture(scope="module") def zarr_tmp2(tmp_path_factory: pytest.TempPathFactory) -> Path: """Make a temp file for the output MDIO.""" + if DEBUG_MODE: + return Path("TMP/mdio2") return tmp_path_factory.mktemp(r"mdio2") @pytest.fixture(scope="session") def segy_export_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path: """Make a temp file for the round-trip IBM SEG-Y.""" - tmp_dir = tmp_path_factory.mktemp("segy") + if DEBUG_MODE: + tmp_dir = Path("TMP/segy") + tmp_dir.mkdir(parents=True, exist_ok=True) + else: + tmp_dir = tmp_path_factory.mktemp("segy") return tmp_dir / "teapot_roundtrip.segy" diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index f361e2e60..6140b2958 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -16,6 +16,27 @@ if TYPE_CHECKING: from pathlib import Path + from segy.schema import SegySpec + + +def _segy_spec_mock_4d() -> SegySpec: + """Create a mock SEG-Y spec for 4D data.""" + trace_header_fields = [ + HeaderField(name="field_rec_no", byte=9, format="int32"), + HeaderField(name="channel", byte=13, format="int32"), + HeaderField(name="shot_point", byte=17, format="int32"), + HeaderField(name="offset", byte=37, format="int32"), + HeaderField(name="samples_per_trace", byte=115, format="int32"), + HeaderField(name="sample_interval", byte=117, format="int32"), + HeaderField(name="shot_line", byte=133, format="int16"), + HeaderField(name="cable", byte=137, format="int16"), + HeaderField(name="gun", byte=171, format="int16"), + ] + rev1_spec = get_segy_standard(1.0) + spec = rev1_spec.customize(trace_header_fields=trace_header_fields) + spec.segy_standard = SegyStandard.REV1 + return spec + def create_segy_mock_4d( # noqa: PLR0913 fake_segy_tmp: Path, @@ -61,23 +82,8 @@ def create_segy_mock_4d( # noqa: PLR0913 cable_headers = np.tile(cable_headers, shot_count) channel_headers = np.tile(channel_headers, shot_count) - trace_header_fields = [ - HeaderField(name="field_rec_no", byte=9, format="int32"), - HeaderField(name="channel", byte=13, format="int32"), - HeaderField(name="shot_point", byte=17, format="int32"), - HeaderField(name="offset", byte=37, format="int32"), - HeaderField(name="samples_per_trace", byte=115, format="int32"), - HeaderField(name="sample_interval", byte=117, format="int32"), - HeaderField(name="shot_line", byte=133, format="int16"), - HeaderField(name="cable", byte=137, format="int16"), - HeaderField(name="gun", byte=171, format="int16"), - ] - - rev1_spec = get_segy_standard(1.0) - spec = rev1_spec.customize(trace_header_fields=trace_header_fields) - spec.segy_standard = SegyStandard.REV1 factory = SegyFactory( - spec=spec, + spec=_segy_spec_mock_4d(), sample_interval=1000, samples_per_trace=num_samples, ) diff --git a/tests/integration/test_segy_import_export.py b/tests/integration/test_segy_import_export.py index e80028df1..54842c6fe 100644 --- a/tests/integration/test_segy_import_export.py +++ b/tests/integration/test_segy_import_export.py @@ -12,10 +12,9 @@ import pytest import xarray as xr from segy import SegyFile -from segy.standards import get_segy_standard from tests.integration.testing_data import binary_header_teapot_dome +from tests.integration.testing_data import custom_teapot_dome_segy_spec from tests.integration.testing_data import text_header_teapot_dome -from tests.integration.testing_helpers import customize_segy_specs from tests.integration.testing_helpers import get_inline_header_values from tests.integration.testing_helpers import get_values from tests.integration.testing_helpers import validate_variable @@ -23,6 +22,7 @@ from mdio import MDIOReader from mdio import mdio_to_segy from mdio.converters.exceptions import GridTraceSparsityError +from mdio.core.utils_read import open_zarr_dataset from mdio.converters.segy import segy_to_mdio from mdio.core import Dimension from mdio.core.storage_location import StorageLocation @@ -246,27 +246,16 @@ def test_import_6d_segy( # noqa: PLR0913 @pytest.mark.dependency -@pytest.mark.parametrize("index_bytes", [(17, 13, 81, 85)]) -@pytest.mark.parametrize("index_names", [("inline", "crossline", "cdp_x", "cdp_y")]) -@pytest.mark.parametrize("index_types", [("int32", "int32", "int32", "int32")]) def test_3d_import( - segy_input: Path, - zarr_tmp: Path, - index_bytes: tuple[int, ...], - index_names: tuple[str, ...], - index_types: tuple[str, ...], + segy_input: Path, + zarr_tmp: Path, ) -> None: - """Test importing a SEG-Y file to MDIO.""" - segy_spec = get_segy_standard(1.0) - segy_spec = customize_segy_specs( - segy_spec=segy_spec, - index_bytes=index_bytes, - index_names=index_names, - index_types=index_types, - ) - + """Test importing a SEG-Y file to MDIO. + + NOTE: This test must be executed before the 'TestReader' and 'TestExport' tests. + """ segy_to_mdio( - segy_spec=segy_spec, + segy_spec=custom_teapot_dome_segy_spec(keep_unaltered=True), mdio_template=TemplateRegistry().get("PostStack3DTime"), input_location=StorageLocation(str(segy_input)), output_location=StorageLocation(str(zarr_tmp)), @@ -276,13 +265,17 @@ def test_3d_import( @pytest.mark.dependency("test_3d_import") class TestReader: - """Test reader functionality.""" + """Test reader functionality. - def test_meta_dataset_read(self, zarr_tmp: Path) -> None: + NOTE: This tests must be executed after the 'test_3d_import' successfully completes and + before running 'TestExport' tests. + """ + + def test_meta_dataset(self, zarr_tmp: Path) -> None: """Metadata reading tests.""" # NOTE: If mask_and_scale is not set, # Xarray will convert int to float and replace _FillValue with NaN - ds = xr.open_dataset(zarr_tmp, engine="zarr", mask_and_scale=False) + ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) expected_attrs = { "apiVersion": "1.0.0a1", "createdOn": "2025-08-06 16:21:54.747880+00:00", @@ -299,23 +292,23 @@ def test_meta_dataset_read(self, zarr_tmp: Path) -> None: attributes = ds.attrs["attributes"] assert attributes is not None - - # Validate attributes provided by the template + assert len(attributes) == 7 + # Validate all attribute provided by the abstract template + assert attributes["traceVariableName"] == "amplitude" + # Validate attributes provided by the PostStack3DTime template assert attributes["surveyDimensionality"] == "3D" assert attributes["ensembleType"] == "line" assert attributes["processingStage"] == "post-stack" - # Validate text header assert attributes["textHeader"] == text_header_teapot_dome() - # Validate binary header assert attributes["binaryHeader"] == binary_header_teapot_dome() - def test_meta_variable_read(self, zarr_tmp: Path) -> None: + def test_meta_variable(self, zarr_tmp: Path) -> None: """Metadata reading tests.""" # NOTE: If mask_and_scale is not set, # Xarray will convert int to float and replace _FillValue with NaN - ds = xr.open_dataset(zarr_tmp, engine="zarr", mask_and_scale=False) + ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) expected_attrs = { "count": 97354860, "sum": -8594.551666259766, @@ -332,7 +325,7 @@ def test_grid(self, zarr_tmp: Path) -> None: # Load Xarray dataset from the MDIO file # NOTE: If mask_and_scale is not set, # Xarray will convert int to float and replace _FillValue with NaN - ds = xr.open_dataset(zarr_tmp, engine="zarr", mask_and_scale=False) + ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) # Note: in order to create the dataset we used the Time template, so the # sample dimension is called "time" @@ -347,9 +340,10 @@ def test_grid(self, zarr_tmp: Path) -> None: validate_variable(ds, "cdp_y", (345, 188), ["inline", "crossline"], np.float64, None, None) # Validate the headers - # We have a subset of headers since we used customize_segy_specs() providing the values only - # for "inline", "crossline", "cdp_x", "cdp_y" - data_type = np.dtype([("inline", " None: """Read and compare every 75 inlines' mean and std. dev.""" # NOTE: If mask_and_scale is not set, # Xarray will convert int to float and replace _FillValue with NaN - ds = xr.open_dataset(zarr_tmp, engine="zarr", mask_and_scale=False) + ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) inlines = ds["amplitude"][::75, :, :] mean, std = inlines.mean(), inlines.std() npt.assert_allclose([mean, std], [1.0555277e-04, 6.0027051e-01]) @@ -387,7 +381,7 @@ def test_crossline(self, zarr_tmp: Path) -> None: """Read and compare every 75 crosslines' mean and std. dev.""" # NOTE: If mask_and_scale is not set, # Xarray will convert int to float and replace _FillValue with NaN - ds = xr.open_dataset(zarr_tmp, engine="zarr", mask_and_scale=False) + ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) xlines = ds["amplitude"][:, ::75, :] mean, std = xlines.mean(), xlines.std() @@ -397,7 +391,7 @@ def test_zslice(self, zarr_tmp: Path) -> None: """Read and compare every 225 z-slices' mean and std. dev.""" # NOTE: If mask_and_scale is not set, # Xarray will convert int to float and replace _FillValue with NaN - ds = xr.open_dataset(zarr_tmp, engine="zarr", mask_and_scale=False) + ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) slices = ds["amplitude"][:, :, ::225] mean, std = slices.mean(), slices.std() npt.assert_allclose([mean, std], [0.005236923, 0.61279935]) @@ -405,23 +399,26 @@ def test_zslice(self, zarr_tmp: Path) -> None: @pytest.mark.dependency("test_3d_import") class TestExport: - """Test SEG-Y exporting functionaliy.""" + """Test SEG-Y exporting functionality. + + NOTE: This test(s) must be executed after the 'test_3d_import' and 'TestReader' tests + successfully complete. + """ - def test_3d_export(self, zarr_tmp: Path, segy_export_tmp: Path) -> None: + def test_3d_export(self, segy_input: Path, zarr_tmp: Path, segy_export_tmp: Path) -> None: """Test 3D export to IBM and IEEE.""" + spec = custom_teapot_dome_segy_spec(keep_unaltered=True) + mdio_to_segy( - mdio_path_or_buffer=zarr_tmp.__str__(), - output_segy_path=segy_export_tmp.__str__(), + segy_spec=spec, + input_location=StorageLocation(str(zarr_tmp)), + output_location=StorageLocation(str(segy_export_tmp)), ) - def test_size_equal(self, segy_input: Path, segy_export_tmp: Path) -> None: - """Check if file sizes match on IBM file.""" + # Check if file sizes match on IBM file. assert segy_input.stat().st_size == segy_export_tmp.stat().st_size - def test_rand_equal(self, segy_input: Path, segy_export_tmp: Path) -> None: - """IBM. Is random original traces and headers match round-trip file?""" - spec = mdio_segy_spec() - + # IBM. Is random original traces and headers match round-trip file? in_segy = SegyFile(segy_input, spec=spec) out_segy = SegyFile(segy_export_tmp, spec=spec) @@ -433,5 +430,7 @@ def test_rand_equal(self, segy_input: Path, segy_export_tmp: Path) -> None: assert in_segy.num_traces == out_segy.num_traces assert in_segy.text_header == out_segy.text_header assert in_segy.binary_header == out_segy.binary_header - npt.assert_array_equal(in_traces.header, out_traces.header) - npt.assert_array_equal(in_traces.sample, out_traces.sample) + # TODO (Dmitriy Repin): Reconcile custom SegySpecs used in the roundtrip SEGY -> MDIO -> SEGY tests + # https://github.com/TGSAI/mdio-python/issues/610 + npt.assert_array_equal(desired=in_traces.header, actual=out_traces.header) + npt.assert_array_equal(desired=in_traces.sample, actual=out_traces.sample) diff --git a/tests/integration/test_segy_import_export_masked.py b/tests/integration/test_segy_import_export_masked.py index 1129a1264..1edd2f140 100644 --- a/tests/integration/test_segy_import_export_masked.py +++ b/tests/integration/test_segy_import_export_masked.py @@ -21,17 +21,16 @@ from segy.schema import HeaderField from segy.schema import SegySpec from segy.standards import get_segy_standard +from tests.conftest import DEBUG_MODE -from mdio import MDIOReader from mdio import mdio_to_segy +from mdio.core.utils_read import open_zarr_dataset from mdio.converters.segy import segy_to_mdio from mdio.core.storage_location import StorageLocation from mdio.schemas.v1.templates.template_registry import TemplateRegistry -from mdio.segy.utilities import segy_export_rechunker if TYPE_CHECKING: from collections.abc import Iterable - from pathlib import Path from numpy.typing import NDArray @@ -150,10 +149,8 @@ def __iter__(self) -> Iterable[MaskedExportConfigTypes]: # fmt: on -def mock_nd_segy(path: str, grid_conf: GridConfig, segy_factory_conf: SegyFactoryConfig) -> SegySpec: - """Create a fake SEG-Y file with a multidimensional grid.""" +def _segy_spec_mock_nd_segy(grid_conf: GridConfig, segy_factory_conf: SegyFactoryConfig) -> SegySpec: spec = get_segy_standard(segy_factory_conf.revision) - header_flds = [] for dim in grid_conf.dims: byte_loc = segy_factory_conf.header_byte_map[dim.name] @@ -179,6 +176,12 @@ def mock_nd_segy(path: str, grid_conf: GridConfig, segy_factory_conf: SegyFactor spec = spec.customize(trace_header_fields=header_flds) spec.segy_standard = segy_factory_conf.revision + return spec + + +def mock_nd_segy(path: str, grid_conf: GridConfig, segy_factory_conf: SegyFactoryConfig) -> SegySpec: + """Create a fake SEG-Y file with a multidimensional grid.""" + spec = _segy_spec_mock_nd_segy(grid_conf, segy_factory_conf) factory = SegyFactory(spec=spec, samples_per_trace=segy_factory_conf.num_samples) dim_coords = () @@ -262,6 +265,8 @@ def generate_selection_mask(selection_conf: SelectionMaskConfig, grid_conf: Grid @pytest.fixture def export_masked_path(tmp_path_factory: pytest.TempPathFactory) -> Path: """Fixture that generates temp directory for export tests.""" + if DEBUG_MODE: + return Path("TMP/export_masked") return tmp_path_factory.getbasetemp() / "export_masked" @@ -276,7 +281,11 @@ class TestNdImportExport: """Test import/export of n-D SEG-Ys to MDIO, with and without selection mask.""" def test_import(self, test_conf: MaskedExportConfig, export_masked_path: Path) -> None: - """Test import of an n-D SEG-Y file to MDIO.""" + """Test import of an n-D SEG-Y file to MDIO. + + NOTE: This test must be executed before running 'test_ingested_mdio', 'test_export', and + 'test_export_masked' tests. + """ grid_conf, segy_factory_conf, segy_to_mdio_conf, _ = test_conf domain = "Time" @@ -313,14 +322,16 @@ def test_import(self, test_conf: MaskedExportConfig, export_masked_path: Path) - ) def test_ingested_mdio(self, test_conf: MaskedExportConfig, export_masked_path: Path) -> None: - """Verify if ingested data is correct.""" + """Verify if ingested data is correct. + + NOTE: This test must be executed after the 'test_import' successfully completes + and before running 'test_export' and 'test_export_masked' tests. + """ grid_conf, segy_factory_conf, segy_to_mdio_conf, _ = test_conf mdio_path = export_masked_path / f"{grid_conf.name}.mdio" # Open the MDIO file - # NOTE: If mask_and_scale is not set, - # Xarray will convert int to float and replace _FillValue with NaN - ds = xr.open_dataset(mdio_path, engine="zarr", mask_and_scale=False) + ds = open_zarr_dataset(StorageLocation(str(mdio_path))) # Test dimensions and ingested dimension headers expected_dims = grid_conf.dims @@ -385,56 +396,67 @@ def test_ingested_mdio(self, test_conf: MaskedExportConfig, export_masked_path: def test_export(self, test_conf: MaskedExportConfig, export_masked_path: Path) -> None: - """Test export of an n-D MDIO file back to SEG-Y.""" + """Test export of an n-D MDIO file back to SEG-Y. + + NOTE: This test must be executed after the 'test_import' and 'test_ingested_mdio' + successfully complete. + """ grid_conf, segy_factory_conf, segy_to_mdio_conf, _ = test_conf segy_path = export_masked_path / f"{grid_conf.name}.sgy" mdio_path = export_masked_path / f"{grid_conf.name}.mdio" segy_rt_path = export_masked_path / f"{grid_conf.name}_rt.sgy" - index_names = segy_factory_conf.header_byte_map.keys() - access_pattern = "".join(map(str, range(len(index_names) + 1))) - mdio = MDIOReader(mdio_path.__str__(), access_pattern=access_pattern) - - chunks, shape = mdio.chunks, mdio.shape - new_chunks = segy_export_rechunker(chunks, shape, dtype="float32", limit="0.3M") - mdio_to_segy( - mdio_path.__str__(), - segy_rt_path.__str__(), - access_pattern=access_pattern, - new_chunks=new_chunks, + segy_spec=_segy_spec_mock_nd_segy(grid_conf, segy_factory_conf), + input_location=StorageLocation(str(mdio_path)), + output_location=StorageLocation(str(segy_rt_path)) ) expected_sgy = SegyFile(segy_path) actual_sgy = SegyFile(segy_rt_path) - assert_array_equal(actual_sgy.trace[:], expected_sgy.trace[:]) + + num_traces = expected_sgy.num_traces + random_indices = np.random.choice(num_traces, 10, replace=False) + expected_traces = expected_sgy.trace[random_indices] + actual_traces = actual_sgy.trace[random_indices] + + assert expected_sgy.num_traces == actual_sgy.num_traces + assert expected_sgy.text_header == actual_sgy.text_header + assert expected_sgy.binary_header == actual_sgy.binary_header + + # TODO (Dmitriy Repin): Reconcile custom SegySpecs used in the roundtrip SEGY -> MDIO -> SEGY tests + # https://github.com/TGSAI/mdio-python/issues/610 + assert_array_equal(desired=expected_traces.header, actual=actual_traces.header) + assert_array_equal(desired=expected_traces.sample, actual=actual_traces.sample) + def test_export_masked(self, test_conf: MaskedExportConfig, export_masked_path: Path) -> None: - """Test export of an n-D MDIO file back to SEG-Y with masked export.""" + """Test export of an n-D MDIO file back to SEG-Y with masked export. + + NOTE: This test must be executed after the 'test_import' and 'test_ingested_mdio' + successfully complete. + """ grid_conf, segy_factory_conf, segy_to_mdio_conf, selection_conf = test_conf segy_path = export_masked_path / f"{grid_conf.name}.sgy" mdio_path = export_masked_path / f"{grid_conf.name}.mdio" segy_rt_path = export_masked_path / f"{grid_conf.name}_rt.sgy" - index_names = segy_factory_conf.header_byte_map.keys() - access_pattern = "".join(map(str, range(len(index_names) + 1))) - mdio = MDIOReader(mdio_path.__str__(), access_pattern=access_pattern) - export_chunks = segy_export_rechunker( - mdio.chunks, mdio.shape, dtype="float32", limit="0.3M" - ) selection_mask = generate_selection_mask(selection_conf, grid_conf) mdio_to_segy( - mdio_path.__str__(), - segy_rt_path.__str__(), - access_pattern=access_pattern, - new_chunks=export_chunks, - selection_mask=selection_mask, + segy_spec=_segy_spec_mock_nd_segy(grid_conf, segy_factory_conf), + input_location=StorageLocation(str(mdio_path)), + output_location=StorageLocation(str(segy_rt_path)), + selection_mask=selection_mask ) expected_trc_idx = selection_mask.ravel().nonzero()[0] expected_sgy = SegyFile(segy_path) actual_sgy = SegyFile(segy_rt_path) - assert_array_equal(actual_sgy.trace[:], expected_sgy.trace[expected_trc_idx]) + + # TODO (Dmitriy Repin): Reconcile custom SegySpecs used in the roundtrip SEGY -> MDIO -> SEGY tests + # https://github.com/TGSAI/mdio-python/issues/610 + assert_array_equal(actual_sgy.trace[:].header, expected_sgy.trace[expected_trc_idx].header) + assert_array_equal(actual_sgy.trace[:].sample, expected_sgy.trace[expected_trc_idx].sample) diff --git a/tests/integration/testing_data.py b/tests/integration/testing_data.py index 3696a7cf1..b61263e47 100644 --- a/tests/integration/testing_data.py +++ b/tests/integration/testing_data.py @@ -1,5 +1,38 @@ """Integration tests data for teapot dome SEG-Y.""" +from segy.schema import SegySpec +from segy.standards import get_segy_standard +from tests.integration.testing_helpers import customize_segy_specs + + +def custom_teapot_dome_segy_spec(keep_unaltered: bool) -> SegySpec: + """Return the minimum customized SEG-Y specification for the teapot dome dataset. + + In SEG-Y spec rev 1.0: + inline = (189, "int32") + crossline = (193, "int32") + cdp_x = (181, "int32") + cdp_y = (185, "int32") + and + trace_num_orig_record = (13, "int32") + energy_source_point_num = (17, "int32") + group_coord_x = (81, "int32") + group_coord_y = (85, "int32") + + In SEGY 1.0 - 2.1, the trace header contains unassigned bytes 181-240. + """ + index_bytes: tuple[int, ...] = (17, 13, 81, 85) + index_names: tuple[str, ...] = ("inline", "crossline", "cdp_x", "cdp_y") + index_types: tuple[str, ...] = ("int32", "int32", "int32", "int32") + segy_spec = get_segy_standard(1.0) + return customize_segy_specs( + segy_spec=segy_spec, + index_bytes=index_bytes, + index_names=index_names, + index_types=index_types, + keep_unaltered = keep_unaltered + ) + def text_header_teapot_dome() -> list[str]: """Return the teapot dome expected text header.""" diff --git a/tests/integration/testing_helpers.py b/tests/integration/testing_helpers.py index 0a33fbd69..79fdf7c86 100644 --- a/tests/integration/testing_helpers.py +++ b/tests/integration/testing_helpers.py @@ -13,6 +13,7 @@ def customize_segy_specs( index_bytes: tuple[int, ...] | None = None, index_names: tuple[int, ...] | None = None, index_types: tuple[int, ...] | None = None, + keep_unaltered: bool = False ) -> SegySpec: """Customize SEG-Y specifications with user-defined index fields.""" if not index_bytes: @@ -26,13 +27,19 @@ def customize_segy_specs( err = "All index fields must have the same length." raise ValueError(err) + fields = {} + if keep_unaltered: + # Keep unaltered fields, but remove the fields that are being customized + # to avoid field name duplication + for f in segy_spec.trace.header.fields: + if f.name not in index_names: + fields[f.byte] = f + # Index the dataset using a spec that interprets the user provided index headers. - index_fields = [] for name, byte, format_ in zip(index_names, index_bytes, index_types, strict=True): - index_fields.append(HeaderField(name=name, byte=byte, format=format_)) - - return segy_spec.customize(trace_header_fields=index_fields) + fields[byte] = HeaderField(name=name, byte=byte, format=format_) + return segy_spec.customize(trace_header_fields=fields.values()) def get_values(arr: xr.DataArray) -> np.ndarray: """Extract actual values from an Xarray DataArray.""" @@ -57,7 +64,26 @@ def validate_variable( # noqa PLR0913 arr = dataset[name] assert shape == arr.shape assert set(dims) == set(arr.dims) - assert data_type == arr.dtype + if hasattr(data_type, "fields") and data_type.fields is not None: + # The following assertion will fail because of differences in offsets + # assert data_type == arr.dtype + + # Compare field names + expected_names = [name for name in data_type.names] + actual_names = [name for name in arr.dtype.names] + assert expected_names == actual_names + + # Compare field types + expected_types = [data_type[name] for name in data_type.names] + actual_types = [arr.dtype[name] for name in arr.dtype.names] + assert expected_types == actual_types + + # Compare field offsets fails. + # However, we believe this is acceptable and do not compare offsets + # name: 'shot_point' dt_exp: (dtype('>i4'), 196) dt_act: (dtype(' Date: Tue, 2 Sep 2025 20:59:33 +0000 Subject: [PATCH 9/9] Adjust TestReader tests --- tests/integration/test_segy_import_export.py | 8 ++++---- tests/integration/testing_helpers.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_segy_import_export.py b/tests/integration/test_segy_import_export.py index 54842c6fe..5fa68e890 100644 --- a/tests/integration/test_segy_import_export.py +++ b/tests/integration/test_segy_import_export.py @@ -292,7 +292,7 @@ def test_meta_dataset(self, zarr_tmp: Path) -> None: attributes = ds.attrs["attributes"] assert attributes is not None - assert len(attributes) == 7 + assert len(attributes) == 6 # Validate all attribute provided by the abstract template assert attributes["traceVariableName"] == "amplitude" # Validate attributes provided by the PostStack3DTime template @@ -375,7 +375,7 @@ def test_inline(self, zarr_tmp: Path) -> None: ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) inlines = ds["amplitude"][::75, :, :] mean, std = inlines.mean(), inlines.std() - npt.assert_allclose([mean, std], [1.0555277e-04, 6.0027051e-01]) + npt.assert_allclose([mean, std], [1.0555277e-04, 6.0027051e-01], rtol=1e-05) def test_crossline(self, zarr_tmp: Path) -> None: """Read and compare every 75 crosslines' mean and std. dev.""" @@ -385,7 +385,7 @@ def test_crossline(self, zarr_tmp: Path) -> None: xlines = ds["amplitude"][:, ::75, :] mean, std = xlines.mean(), xlines.std() - npt.assert_allclose([mean, std], [-5.0329847e-05, 5.9406823e-01]) + npt.assert_allclose([mean, std], [-5.0329847e-05, 5.9406823e-01], rtol=1e-06) def test_zslice(self, zarr_tmp: Path) -> None: """Read and compare every 225 z-slices' mean and std. dev.""" @@ -394,7 +394,7 @@ def test_zslice(self, zarr_tmp: Path) -> None: ds = open_zarr_dataset(StorageLocation(str(zarr_tmp))) slices = ds["amplitude"][:, :, ::225] mean, std = slices.mean(), slices.std() - npt.assert_allclose([mean, std], [0.005236923, 0.61279935]) + npt.assert_allclose([mean, std], [0.005236923, 0.61279935], rtol=1e-06) @pytest.mark.dependency("test_3d_import") diff --git a/tests/integration/testing_helpers.py b/tests/integration/testing_helpers.py index 79fdf7c86..a9d40475b 100644 --- a/tests/integration/testing_helpers.py +++ b/tests/integration/testing_helpers.py @@ -74,8 +74,8 @@ def validate_variable( # noqa PLR0913 assert expected_names == actual_names # Compare field types - expected_types = [data_type[name] for name in data_type.names] - actual_types = [arr.dtype[name] for name in arr.dtype.names] + expected_types = [data_type[name].newbyteorder('=') for name in data_type.names] + actual_types = [arr.dtype[name].newbyteorder('=') for name in arr.dtype.names] assert expected_types == actual_types # Compare field offsets fails.