Write flags back to SD HDF5 data sets

imagesets/sdhdfimageset.cpp

@@ -86,18 +86,40 @@ std::string SdhdfImageSet::Description(const ImageSetIndex& index) const {
          _beams[beam].bandParams[band].label;
 }
 
+bool SdhdfImageSet::tryOpen(H5::DataSet& dataSet, H5::H5File& file,
+                            const std::string& name) {
+#if H5_VERSION_GE(1, 10, 7)
+  if (file.exists(name)) {
+    dataSet = file.openDataSet(name);
+    return true;
+  } else
+    return false;
+#else
+  // file.exists(name) doesn't work in older HDF5 libs, so use trial
+  // and error. Unfortunately this generates output on the cmdline
+  // when the file does not exist, hence only use it as fall back.
+  try {
+    dataSet = file.openDataSet(name);
+    return true;
+  } catch (H5::FileIException&) {
+    // The data set did not exist (probably)
+    return false;
+  }
+#endif
+}
+
 BaselineData SdhdfImageSet::loadData(ProgressListener& progress,
                                      const ImageSetIndex& index) {
-  size_t beam = _indexTable[index.Value()].first;
-  size_t band = _indexTable[index.Value()].second;
+  const size_t beam = _indexTable[index.Value()].first;
+  const size_t band = _indexTable[index.Value()].second;
   const BandParams& bandParams = _beams[beam].bandParams[band];
   const std::string label = bandParams.label;
   Logger::Info << "Reading beam " << beam << ", " << label << "...\n";
   H5::H5File file(_path, H5F_ACC_RDONLY, H5P_DEFAULT);
 
-  const std::string prefix =
-      "beam_" + std::to_string(beam) + "/" + label + "/astronomy_data/";
-  H5::DataSet dataSet = file.openDataSet(prefix + "data");
+  const std::string beamPrefix = "beam_" + std::to_string(beam) + "/" + label;
+  const std::string dataPrefix = beamPrefix + "/astronomy_data/";
+  H5::DataSet dataSet = file.openDataSet(dataPrefix + "data");
   H5::DataSpace space = dataSet.getSpace();
   std::vector<hsize_t> ddims(space.getSimpleExtentNdims());
   space.getSimpleExtentDims(ddims.data());
@@ -135,17 +157,8 @@ BaselineData SdhdfImageSet::loadData(ProgressListener& progress,
   }
 
   Mask2DPtr mask;
-  // file.exists(prefix + "flag") doesn't work in older HDF5 libs, so use trial
-  // and error:
   H5::DataSet flagsSet;
-  bool hasFlags = false;
-  try {
-    flagsSet = file.openDataSet(prefix + "flag");
-    hasFlags = true;
-  } catch (H5::FileIException&) {
-    // Ignore
-  }
-  if (hasFlags) {
+  if (tryOpen(flagsSet, file, dataPrefix + "flag")) {
     progress.OnProgress(5, 10);
     Logger::Info << "Reading flags...\n";
     H5::DataSpace space = flagsSet.getSpace();
@@ -199,8 +212,8 @@ BaselineData SdhdfImageSet::loadData(ProgressListener& progress,
   H5::CompType obsParamsType(sizeof(ObsParams));
   obsParamsType.insertMember("MJD", HOFFSET(ObsParams, mjd),
                              H5::PredType::NATIVE_DOUBLE);
-  H5::DataSet obsParamsSet = file.openDataSet(
-      "beam_" + std::to_string(beam) + "/" + label + "/metadata/obs_params");
+  H5::DataSet obsParamsSet =
+      file.openDataSet(beamPrefix + "/metadata/obs_params");
   H5::DataSpace obsParamsSpace = obsParamsSet.getSpace();
   if (obsParamsSpace.getSimpleExtentNdims() != 1)
     throw std::runtime_error(
@@ -221,8 +234,7 @@ BaselineData SdhdfImageSet::loadData(ProgressListener& progress,
   metaData->SetObservationTimes(std::move(times));
 
   H5::DataSet frequencySet =
-      file.openDataSet("beam_" + std::to_string(beam) + "/" + label +
-                       "/astronomy_data/frequency");
+      file.openDataSet(beamPrefix + "/astronomy_data/frequency");
   H5::DataSpace frequencySpace = frequencySet.getSpace();
   if (frequencySpace.getSimpleExtentNdims() != 1)
     throw std::runtime_error(
@@ -259,4 +271,44 @@ void SdhdfImageSet::PerformReadRequests(ProgressListener& progress) {
   progress.OnFinish();
 }
 
+void SdhdfImageSet::AddWriteFlagsTask(const ImageSetIndex& index,
+                                      std::vector<Mask2DCPtr>& flags) {
+  assert(!flags.empty());
+  const size_t beam = _indexTable[index.Value()].first;
+  const size_t band = _indexTable[index.Value()].second;
+  const BandParams& bandParams = _beams[beam].bandParams[band];
+  const std::string label = bandParams.label;
+  const size_t nTimes = flags.front()->Width();
+  Logger::Info << "Writing flags for beam " << beam << ", " << label << "...\n";
+  H5::H5File file(_path, H5F_ACC_RDWR, H5P_DEFAULT);
+
+  const std::string beamPrefix =
+      "beam_" + std::to_string(beam) + '/' + label + '/';
+  const std::string dataPrefix = beamPrefix + "astronomy_data/";
+
+  H5::DataSet flagsSet;
+  if (!tryOpen(flagsSet, file, dataPrefix + "flag")) {
+    // Create the data set
+    H5::DataType dataType(H5::PredType::NATIVE_UINT8);
+    const hsize_t dimsf[2]{nTimes, hsize_t(bandParams.nChannels)};
+    H5::DataSpace dataSpace(2, dimsf);
+    flagsSet = file.createDataSet(dataPrefix + "flag", dataType, dataSpace);
+  }
+
+  Logger::Info << "Writing flag data...\n";
+  aocommon::UVector<unsigned char> flagData(nTimes * bandParams.nChannels);
+  unsigned char* valuePtr = flagData.data();
+  for (size_t timeIndex = 0; timeIndex != nTimes; ++timeIndex) {
+    for (int channel = 0; channel != bandParams.nChannels; ++channel) {
+      bool flag = false;
+      for (const Mask2DCPtr& mask : flags)
+        flag = flag || mask->Value(timeIndex, channel);
+      *valuePtr = flag ? 1 : 0;
+      ++valuePtr;
+    }
+  }
+  flagsSet.write(flagData.data(), H5::PredType::NATIVE_UINT8, H5S_ALL, H5S_ALL,
+                 H5P_DEFAULT);
+}
+
 }  // namespace imagesets

imagesets/sdhdfimageset.h

@@ -7,6 +7,11 @@
 
 #include "imageset.h"
 
+namespace H5 {
+class DataSet;
+class H5File;
+}  // namespace H5
+
 namespace imagesets {
 
 class SdhdfImageSet final : public ImageSet {
@@ -41,7 +46,18 @@ class SdhdfImageSet final : public ImageSet {
     return std::unique_ptr<BaselineData>(new BaselineData(baseline));
   }
 
+  void AddWriteFlagsTask(const ImageSetIndex &index,
+                         std::vector<Mask2DCPtr> &flags) override;
+  void PerformWriteFlagsTask() override {}
+
  private:
+  /**
+   * Open the data set @c name inside @c file, and store the result in @c
+   * dataSet. If the data set does not exist, @c false is returned.
+   */
+  static bool tryOpen(H5::DataSet &dataSet, H5::H5File &file,
+                      const std::string &name);
+
   struct Header {
     int nBeams;
     char telescopeName[64];