fix 2d lazy back end error and add tests

trhallam · May 15, 2024 · 8f62d7b · 8f62d7b
1 parent d24d190
commit 8f62d7b
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 25 deletions.
diff --git a/segysak/segy/_xarray.py b/segysak/segy/_xarray.py
@@ -136,13 +136,25 @@ def _raw_indexing_method(self, key: tuple) -> np.typing.ArrayLike:
         for dim in self.dnames[:-1]:
             ds[dim] = Variable(("tracen",), tracen_slice_df[dim].values)
         ds["sgy"] = Variable(("tracen", VerticalKeyDim.samples), volume)
-        ds = ds.set_index(tracen=list(self.dnames[:-1]))
-        data = (
-            ds["sgy"]
-            .unstack("tracen")
-            .transpose(*self.dnames)
-            .values.squeeze(axis=squeeze_me)
-        )
+
+        unstack_dims = list(self.dnames[:-1])
+        if len(unstack_dims) <= 1:
+            # this is usually 2d data
+            data = (
+                ds["sgy"]
+                .swap_dims({"tracen": unstack_dims[0]})
+                .drop("tracen")
+                .transpose(*self.dnames)
+                .values
+            )
+        else:
+            ds = ds.set_index(tracen=unstack_dims)
+            data = (
+                ds["sgy"]
+                .unstack("tracen")
+                .transpose(*self.dnames)
+                .values.squeeze(axis=squeeze_me)
+            )
         return data
 
 

diff --git a/tests/test_segy_xarray.py b/tests/test_segy_xarray.py
@@ -9,22 +9,19 @@
 TEST_DATA_SEGYSAK = TESTS_PATH / "test-data-segysak"
 
 
-def split_segyio_kwargs(segyio_kwargs):
-    valid_dims = ("iline", "xline", "offset", "cdp")
-    valid_kwargs = ("endian",)
-    dims = dict()
-    other = dict()
-    for kwarg, value in segyio_kwargs.items():
-        if kwarg in valid_dims:
-            dims[kwarg] = value
-        elif kwarg in valid_kwargs:
-            other[kwarg] = value
-    return dims, other
+def test_SegyBackendEntrypoint_3d(segyio3d_test_files):
+    path, dims, extra, segyio_kwargs = segyio3d_test_files
+    ds = xr.open_dataset(path, dim_byte_fields=dims, segyio_kwargs=segyio_kwargs)
+    assert isinstance(ds, xr.Dataset)
+    for dim in dims:
+        assert dim in ds.sizes
 
+    data = ds.data.compute()
+    pass
 
-def test_SegyBackendEntrypoint_3d(segyio3d_test_files):
-    path, segyio_kwargs = segyio3d_test_files
-    dims, segyio_kwargs = split_segyio_kwargs(segyio_kwargs)
+
+def test_SegyBackendEntrypoint_2d(volve_segy2d):
+    path, dims, extra, segyio_kwargs = volve_segy2d
     ds = xr.open_dataset(path, dim_byte_fields=dims, segyio_kwargs=segyio_kwargs)
     assert isinstance(ds, xr.Dataset)
     for dim in dims:
@@ -35,17 +32,15 @@ def test_SegyBackendEntrypoint_3d(segyio3d_test_files):
 
 
 def test_SegyBackendEntrypoint_3dps(segyio3dps_test_files):
-    path, segyio_kwargs = segyio3dps_test_files
-    dims, segyio_kwargs = split_segyio_kwargs(segyio_kwargs)
+    path, dims, extra, segyio_kwargs = segyio3dps_test_files
     ds = xr.open_dataset(path, dim_byte_fields=dims, segyio_kwargs=segyio_kwargs)
     assert isinstance(ds, xr.Dataset)
     for dim in dims:
         assert dim in ds.sizes
 
 
 def test_SegyBackendEntrypoint_3dps_subsel(segyio3dps_test_files):
-    path, segyio_kwargs = segyio3dps_test_files
-    dims, segyio_kwargs = split_segyio_kwargs(segyio_kwargs)
+    path, dims, extra, segyio_kwargs = segyio3dps_test_files
     ds = xr.open_dataset(path, dim_byte_fields=dims, segyio_kwargs=segyio_kwargs)
     assert isinstance(ds, xr.Dataset)
     for dim in dims: