diff --git a/cpp/purify/h5reader.h b/cpp/purify/h5reader.h index 965a636b9..be56357f3 100644 --- a/cpp/purify/h5reader.h +++ b/cpp/purify/h5reader.h @@ -160,7 +160,7 @@ class H5Handler { size_t _datalen, _slicepos, _slicelen, _batchpos; }; -/// @brief Reads an HDF5 file with u,v visibilities, constructs a vis_params objects and returns it. +/// @brief Reads an HDF5 file with u,v visibilities, constructs a vis_params object and returns it. /// /// @note vis_name: name of input HDF5 file containing [u, v, real(V), imag(V)]. utilities::vis_params read_visibility(const std::string& vis_name, const bool w_term) { @@ -240,6 +240,50 @@ utilities::vis_params stochread_visibility(H5Handler& file, const size_t N, cons return uv_vis; } +/// @brief Write an HDF5 file with u,v visibilities from a vis_params object. +void write_visibility(const utilities::vis_params& uv_vis, const std::string& h5name, + const bool w_term, const size_t chunksize = 0) { + // Set up HDF5 file + HighFive::File h5file(h5name, HighFive::File::OpenOrCreate | HighFive::File::Truncate); + // Set up file properties, such as chunking and compression + // Note: the I/O is minimised if compression is disabled (obvs) + // If using decompressed data is not an option, then the + // I/O performance can be optimised by chunking the dataset + // in such a way that each MPI rank only has to decompress + // its allocated segment (or a subset thereof) + HighFive::DataSetCreateProps props; + if (uv_vis.u.size()) { + if (chunksize > 0) { + props.add(HighFive::Chunking(std::vector{chunksize})); + } else { + props.add(HighFive::Chunking(std::vector{static_cast(uv_vis.u.size())})); + } + props.add(HighFive::Deflate(9)); // maximal compression + } + // Create the H5 datasets + h5file.createDataSet("u", std::vector(uv_vis.u.data(), uv_vis.u.data() + uv_vis.u.size()), + props); + h5file.createDataSet("v", std::vector(uv_vis.v.data(), uv_vis.v.data() + uv_vis.v.size()), + props); + if (w_term) { + h5file.createDataSet( + "w", std::vector(uv_vis.w.data(), uv_vis.w.data() + uv_vis.w.size()), props); + } + + std::vector redata, imdata, sigma; + redata.reserve(uv_vis.vis.size()); + imdata.reserve(uv_vis.vis.size()); + sigma.reserve(uv_vis.weights.size()); + for (size_t i = 0; i < uv_vis.vis.size(); ++i) { + redata.push_back(uv_vis.vis(i).real()); + imdata.push_back(uv_vis.vis(i).imag()); + sigma.push_back(1.0 / uv_vis.weights(i).real()); + } + h5file.createDataSet("re", std::move(redata), props); + h5file.createDataSet("im", std::move(imdata), props); + h5file.createDataSet("sigma", std::move(imdata), props); +} + } // namespace purify::H5 #endif diff --git a/cpp/tests/purify_h5.cc b/cpp/tests/purify_h5.cc index 50bd42e01..fd0bed453 100644 --- a/cpp/tests/purify_h5.cc +++ b/cpp/tests/purify_h5.cc @@ -3,6 +3,7 @@ #include "purify/config.h" #include "purify/types.h" #include "purify/logging.h" +#include "purify/read_measurements.h" #include "purify/directories.h" #include "purify/h5reader.h" @@ -11,7 +12,7 @@ using namespace purify; -TEST_CASE("Purify H5", "[HDF5]") { +TEST_CASE("Purify H5 reader", "[HDF5]") { H5::H5Handler f(atca_filename("0332-391.h5")); const std::vector u = f.read("u"); @@ -32,3 +33,31 @@ TEST_CASE("Purify H5", "[HDF5]") { u.size() == re.size() && u.size() == im.size() && u.size() == sigma.size(); CHECK(pass); } + +TEST_CASE("Purify H5 writer", "[HDF5]") { + const auto uvfits = read_measurements::read_measurements(atca_filename("0332-391.uvfits")); + + H5::write_visibility(uvfits, "test-h5.h5", false); + + H5::H5Handler f("test-h5.h5"); + + const std::vector u = f.read("u"); + const std::vector v = f.read("v"); + // const std::vector w = f.read("w"); + const std::vector re = f.read("re"); + const std::vector im = f.read("im"); + const std::vector sigma = f.read("sigma"); + + CAPTURE(u.size()); + CAPTURE(v.size()); + // CAPTURE(w.size()); + CAPTURE(re.size()); + CAPTURE(im.size()); + CAPTURE(sigma.size()); + + const bool pass = u.size() == uvfits.u.size() && + u.size() == v.size() && /*u.size() == w.size() &&*/ + u.size() == re.size() && u.size() == im.size() && u.size() == sigma.size(); + + CHECK(pass); +}