diff --git a/aidrin/file_handling/readers/hdf5_reader.py b/aidrin/file_handling/readers/hdf5_reader.py index 382f444e..b9e06530 100644 --- a/aidrin/file_handling/readers/hdf5_reader.py +++ b/aidrin/file_handling/readers/hdf5_reader.py @@ -242,14 +242,24 @@ def _get_selected_dataset_keys(self): return [] def _column_name_from_path(self, path, used_names): - """Pick a unique DataFrame column name for an HDF5 dataset path.""" + """Prefer the full dataset path so group/station context is preserved. + + Examples: ``S_01_01/X``, ``D1.fill_starts``. Falls back only if the full + path is already used as a column name (duplicate selection). + """ + full = path.strip("/") or path + if full not in used_names: + return full short = path.split("/")[-1] or path if short not in used_names: return short - dotted = path.strip("/").replace("/", ".") + dotted = full.replace("/", ".") if dotted not in used_names: return dotted - return path + suffix = 2 + while f"{full}_{suffix}" in used_names: + suffix += 1 + return f"{full}_{suffix}" def _read_compatible_dataset_paths(self, paths): """Merge multiple same-length 1D datasets into one DataFrame.""" @@ -353,7 +363,7 @@ def _read_dataset_path(self, path): data = obj[()] data = self._apply_fill_values(data, obj, path) - col_name = path.split("/")[-1] or path + col_name = self._column_name_from_path(path, set()) if getattr(data, "ndim", 0) == 0: df = pd.DataFrame({col_name: [data]}) elif data.ndim == 1: diff --git a/tests/unit/test_hdf5_reader.py b/tests/unit/test_hdf5_reader.py index 795297f1..67cbfd94 100644 --- a/tests/unit/test_hdf5_reader.py +++ b/tests/unit/test_hdf5_reader.py @@ -480,9 +480,21 @@ def test_read_selected_waveforms_preserves_zeros(self, tmp_path, logger): assert df is not None assert df.shape == (100, 3) - assert list(df.columns) == ["X", "Y", "Z"] - assert (df["X"] == 0.0).any() - assert df["X"].isna().sum() == 0 + assert list(df.columns) == ["S_01_01/X", "S_01_01/Y", "S_01_01/Z"] + assert (df["S_01_01/X"] == 0.0).any() + assert df["S_01_01/X"].isna().sum() == 0 + + def test_read_selected_waveforms_keeps_station_path_context(self, tmp_path, logger): + """Selecting the same short names from two stations stays distinguishable.""" + fpath = str(tmp_path / "rechdf5.h5") + _make_mock_rechdf5(fpath, nx=2, ny=2, npts=50) + + keys = ["S_01_01/X", "S_01_02/X"] + df = hdf5Reader(fpath, logger, selected_keys=keys).read() + + assert df is not None + assert df.shape == (50, 2) + assert list(df.columns) == ["S_01_01/X", "S_01_02/X"] def test_pandas_hdf5_stays_legacy_not_multi_dataset(self, logger): from pathlib import Path