Add SolTab::GetSubArray() generalization

.clang-tidy

@@ -29,7 +29,7 @@ CheckOptions:
   - key:             modernize-make-shared.MakeSmartPtrFunction
     value:           'std::make_shared'
   - key:             modernize-loop-convert.UseCxx20ReverseRanges
-    value:           'true'
+    value:           'false'
   - key:             modernize-use-override.FinalSpelling
     value:           final
   - key:             modernize-use-using.IgnoreMacros

include/schaapcommon/h5parm/soltab.h

@@ -197,10 +197,17 @@ class SolTab : private H5::Group {
                                              size_t pol, size_t dir,
                                              bool nearest) const;
 
-  /// Get the values or weights of this SolTab for a given antenna given an
-  /// antenna name, a direction index, and a (range of) times and frequencies.
-  /// In the returned vector, the freq will be the fastest changing index,
-  /// irrespective of the axis ordering in the underlying h5 data structure.
+  std::tuple<int, int, int, int, int> GetSubArray(
+      const std::string& val_or_weight, const size_t start_antenna_index,
+      const size_t n_antennas, const size_t antenna_step,
+      const size_t start_time_index, const size_t n_times,
+      const size_t time_step, const size_t start_freq_index,
+      const size_t n_freqs, const size_t freq_step,
+      const size_t start_polarization_index, const size_t n_polarizations,
+      const size_t polarization_step, const size_t start_direction_index,
+      const size_t n_directions, const size_t direction_step,
+      std::vector<double>& buffer) const;
+
   std::vector<double> GetSubArray(const std::string& val_or_weight,
                                   const std::string& antenna_name,
                                   size_t start_time_index, size_t n_times,
@@ -212,6 +219,10 @@ class SolTab : private H5::Group {
       const std::string& value_or_weight, const size_t n_axes,
       const std::array<size_t, 5>& size_axis) const;
 
+  /// Get the values or weights of this SolTab for a given antenna given an
+  /// antenna name, a direction index, and a (range of) times and frequencies.
+  /// In the returned vector, the freq will be the fastest changing index,
+  /// irrespective of the axis ordering in the underlying h5 data structure.
   std::vector<double> GetValuesOrWeights(const std::string& val_or_weight,
                                          const std::string& antenna_name,
                                          const std::vector<double>& times,

src/h5parm/soltab.cc

@@ -411,12 +411,15 @@ TimesAndFrequencies SolTab::GetTimesAndFrequencies(
   return tf;
 }
 
-std::vector<double> SolTab::GetSubArray(
-    const std::string& val_or_weight, const std::string& antenna_name,
+std::tuple<int, int, int, int, int> SolTab::GetSubArray(
+    const std::string& val_or_weight, const size_t start_antenna_index,
+    const size_t n_antennas, const size_t antenna_step,
     const size_t start_time_index, const size_t n_times, const size_t time_step,
     const size_t start_freq_index, const size_t n_freqs, const size_t freq_step,
-    const size_t polarization, const size_t direction) const {
-  std::vector<double> res(n_times * n_freqs);
+    const size_t start_polarization_index, const size_t n_polarizations,
+    const size_t polarization_step, const size_t start_direction_index,
+    const size_t n_directions, const size_t direction_step,
+    std::vector<double>& buffer) const {
   H5::DataSet val = openDataSet(val_or_weight);
 
   // Set offsets and strides
@@ -432,34 +435,64 @@ std::vector<double> SolTab::GetSubArray(
   for (const AxisInfo& axis_info : axes_) {
     hsize_t stride = 1;
     hsize_t count = 1;
-    hsize_t memdim = 1;
     hsize_t offset = 0;
     if (axis_info.name == "time") {
       offset = start_time_index;
       stride = time_step;
       count = n_times;
-      memdim = n_times;
     } else if (axis_info.name == "freq") {
       offset = start_freq_index;
       stride = freq_step;
       count = n_freqs;
-      memdim = n_freqs;
     } else if (axis_info.name == "ant") {
-      offset = GetAntIndex(antenna_name);
+      offset = start_antenna_index;
+      stride = antenna_step;
+      count = n_antennas;
     } else if (axis_info.name == "dir") {
-      offset = direction;
+      offset = start_direction_index;
+      stride = direction_step;
+      count = n_directions;
     } else if (axis_info.name == "pol") {
-      offset = polarization;
+      offset = start_polarization_index;
+      stride = polarization_step;
+      count = n_polarizations;
     } else if (axis_info.size != 1) {
       throw std::runtime_error("Axis \"" + axis_info.name +
                                "\" in H5Parm is not understood");
     }
-    memdims.push_back(memdim);
+    memdims.push_back(count);
     offsets.push_back(offset);
     counts.push_back(count);
     strides.push_back(stride);
   }
 
+  size_t antenna_stride = 0;
+  size_t time_stride = 0;
+  size_t freq_stride = 0;
+  size_t polarization_stride = 0;
+  size_t direction_stride = 0;
+  size_t stride = 1;
+  for (auto axis_info = axes_.rbegin(); axis_info != axes_.rend();
+       ++axis_info) {
+    if (axis_info->name == "time") {
+      time_stride = stride;
+      stride *= n_times;
+    } else if (axis_info->name == "freq") {
+      freq_stride = stride;
+      stride *= n_freqs;
+    } else if (axis_info->name == "ant") {
+      antenna_stride = stride;
+      stride *= n_antennas;
+    } else if (axis_info->name == "dir") {
+      direction_stride = stride;
+      stride *= n_directions;
+    } else if (axis_info->name == "pol") {
+      polarization_stride = stride;
+      stride *= n_polarizations;
+    }
+  }
+  buffer.resize(stride);
+
   H5::DataSpace dataspace = val.getSpace();
 
   dataspace.selectHyperslab(H5S_SELECT_SET, counts.data(), offsets.data(),
@@ -467,12 +500,36 @@ std::vector<double> SolTab::GetSubArray(
 
   // Setup memory dataspace
   H5::DataSpace memspace(axes_.size(), memdims.data());
+
   try {
-    val.read(res.data(), H5::PredType::NATIVE_DOUBLE, memspace, dataspace);
+    val.read(buffer.data(), H5::PredType::NATIVE_DOUBLE, memspace, dataspace);
   } catch (H5::DataSetIException& e) {
     e.printErrorStack();
     throw std::runtime_error("Could not read data");
   }
+  return {antenna_stride, time_stride, freq_stride, polarization_stride,
+          direction_stride};
+}
+
+std::vector<double> SolTab::GetSubArray(
+    const std::string& val_or_weight, const std::string& antenna_name,
+    const size_t start_time_index, const size_t n_times, const size_t time_step,
+    const size_t start_freq_index, const size_t n_freqs, const size_t freq_step,
+    const size_t polarization, const size_t direction) const {
+  std::vector<double> res(n_times * n_freqs);
+
+  size_t antenna_index = 0;
+  try {
+    antenna_index = GetAntIndex(antenna_name);
+  } catch (const std::runtime_error& error) {
+    // A SolTab does not necessarily have an "ant" table.
+    // In that case just leave antenna_index at 0.
+  }
+
+  GetSubArray(val_or_weight, antenna_index, 1, 1, start_time_index, n_times,
+              time_step, start_freq_index, n_freqs, freq_step, polarization, 1,
+              1, direction, 1, 1, res);
+
   return res;
 }