diff --git a/docs/changes/153.feature.rst b/docs/changes/153.feature.rst new file mode 100644 index 00000000..a1d5f524 --- /dev/null +++ b/docs/changes/153.feature.rst @@ -0,0 +1 @@ +The UVH5 writer now also writes timestamps to the file under the ``times`` key. diff --git a/src/pyvisgen/io/datawriters.py b/src/pyvisgen/io/datawriters.py index 548676b1..ca5b3d15 100644 --- a/src/pyvisgen/io/datawriters.py +++ b/src/pyvisgen/io/datawriters.py @@ -423,6 +423,8 @@ class UVH5Writer(DataWriter): │ └── n ├── frequency_bands ├── channel_widths + ├── normalize + ├── times └── sky/ └── SI @@ -528,6 +530,9 @@ def write( f.create_dataset("channel_widths", data=self.__to_numpy(obs.bandwidths)) f.create_dataset("normalize", data=np.bool_(normalize)) + times = self.__to_numpy(vis_data.date) + f.create_dataset("times", data=times) + if sky is not None: sky_grp = f.create_group("sky") sky_grp.create_dataset("SI", data=self.__to_numpy(sky)) diff --git a/tests/io/conftest.py b/tests/io/conftest.py index 6801f7d8..c1df2dba 100644 --- a/tests/io/conftest.py +++ b/tests/io/conftest.py @@ -181,6 +181,7 @@ def uvh5_vis_data() -> SimpleNamespace: v=torch.randn(n_baselines, dtype=torch.float64), w=torch.randn(n_baselines, dtype=torch.float64), st_id_pairs=torch.arange(n_baselines * 2).reshape(n_baselines, 2), + date=torch.randn(n_baselines, dtype=torch.float64), ) diff --git a/tests/io/test_datawriters.py b/tests/io/test_datawriters.py index f9024156..0b912902 100644 --- a/tests/io/test_datawriters.py +++ b/tests/io/test_datawriters.py @@ -293,6 +293,15 @@ def test_frequency_bands(self, output_dir, uvh5_vis_data, uvh5_obs) -> None: with h5py.File(output_dir / "train_0.uvh5", "r") as f: np.testing.assert_array_almost_equal(f["frequency_bands"][...], expected) + def test_times(self, output_dir, uvh5_vis_data, uvh5_obs) -> None: + with UVH5Writer(output_path=output_dir, dataset_type="train") as writer: + writer.write(uvh5_vis_data, uvh5_obs, index=0) + + expected = uvh5_vis_data.date + + with h5py.File(output_dir / "train_0.uvh5", "r") as f: + np.testing.assert_array_almost_equal(f["times"], expected) + def test_sky_written(self, output_dir, uvh5_vis_data, uvh5_obs, uvh5_sky) -> None: with UVH5Writer(output_path=output_dir, dataset_type="train") as writer: writer.write(uvh5_vis_data, uvh5_obs, index=0, sky=uvh5_sky)