feat: Save timestamps in UVH5 writer

docs/changes/153.feature.rst

@@ -0,0 +1 @@
+The UVH5 writer now also writes timestamps to the file under the ``times`` key.

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))

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),
     )
 
 

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)