Program Listing for File reflection.hpp

↰ Return to documentation for file (include/dx2/reflection.hpp)

#pragma once

#include "dx2/h5/h5dispatch.hpp"
#include "dx2/h5/h5read_processed.hpp"
#include "dx2/h5/h5utils.hpp"
#include "dx2/h5/h5write.hpp"
#include "dx2/logging.hpp"
#include <array>
#include <chrono>
#include <experimental/mdspan>
#include <functional>
#include <iomanip>
#include <optional>
#include <random>
#include <sstream>
#include <string>
#include <type_traits>
#include <unordered_set>
#include <variant>
#include <vector>

#pragma region Type Helpers
/*
 * Type aliases for centralised control of storage types.
 */

// Internal storage type for the data
template <typename T> using storage_type = std::vector<T>;
// Type alias for mdspan (currently using backport std::experimental::mdspan)
template <typename T>
using mdspan_type =
    std::experimental::mdspan<T, std::experimental::dextents<size_t, 2>>;
#pragma endregion

#pragma region Column Definition
template <typename T> struct TypedColumn {
  using value_type = T;

  TypedColumn(std::string name, std::vector<size_t> shape, storage_type<T> data)
      : name(std::move(name)), shape(std::move(shape)), data(std::move(data)) {
    if (this->shape.size() != 1 && this->shape.size() != 2) {
      std::string got = "{";
      for (size_t i = 0; i < this->shape.size(); ++i) {
        got += (i ? ", " : "") + std::to_string(this->shape[i]);
      }
      throw std::runtime_error(
          "TypedColumn: column shape must be {N} or {N, M}, got " + got + "}");
    }
  }

  size_t rows() const { return shape[0]; }
  size_t cols() const { return shape.size() > 1 ? shape[1] : 1; }

  mdspan_type<T> span() { return mdspan_type<T>(data.data(), rows(), cols()); }

  mdspan_type<const T> span() const {
    return mdspan_type<const T>(data.data(), rows(), cols());
  }

  TypedColumn<T> filtered(const std::vector<size_t> &selected_rows) const {
    storage_type<T> new_data;
    new_data.reserve(selected_rows.size() * cols());
    for (size_t row : selected_rows) {
      auto first = data.begin() + row * cols();
      new_data.insert(new_data.end(), first, first + cols());
    }
    return TypedColumn<T>(name, {selected_rows.size(), cols()},
                          std::move(new_data));
  }

  std::string name;
  std::vector<size_t> shape;
  storage_type<T> data;
};
#pragma endregion

#pragma region Predicate Helpers
template <typename T> struct ColumnPredicate {
  std::string column_name;
  std::function<bool(const mdspan_type<const T> &, size_t)> predicate;
};

namespace logic {

enum class LogicalOp {
  Or,
  And
};

} // namespace logic
#pragma endregion

#pragma region ReflectionTable
class ReflectionTable {
private:
  using Column = h5dispatch::variant_of_t<TypedColumn, h5dispatch::ColumnTypes>;

  // Views returned by column<T>() point into each column's buffer, which
  // survives the column vector reallocating only if columns are moved
  // rather than copied.
  static_assert(std::is_nothrow_move_constructible_v<Column>);

  std::vector<Column> data;
  std::string h5_filepath;
  const std::string DEFAULT_REFL_GROUP = "/dials/processing/group_0";
  std::vector<uint64_t> experiment_ids;
  std::vector<std::string> identifiers;
  uint64_t max_experiment_id = 0;

  size_t get_row_count() const;

  static const std::string &column_name(const Column &col) {
    return std::visit(
        [](const auto &c) -> const std::string & { return c.name; }, col);
  }

  static size_t column_rows(const Column &col) {
    return std::visit([](const auto &c) { return c.rows(); }, col);
  }

  template <typename T, typename Self>
  static auto find_column(Self &self, const std::string &name)
      -> std::conditional_t<std::is_const_v<Self>, const TypedColumn<T>,
                            TypedColumn<T>> * {
    static_assert(h5dispatch::is_column_type_v<T>,
                  "Unsupported column type; see h5dispatch::ColumnTypes");
    for (auto &col : self.data) {
      auto *typed = std::get_if<TypedColumn<T>>(&col);
      if (typed && typed->name == name) {
        return typed;
      }
    }
    return nullptr;
  }

  void merge_into_set(std::unordered_set<size_t> &set,
                      const std::vector<size_t> &rows) const;

  template <typename T>
  std::vector<size_t>
  evaluate_column_predicate(const ColumnPredicate<T> &condition) const {
    auto opt = column<T>(condition.column_name);
    if (!opt)
      throw std::runtime_error("Column not found or wrong type: " +
                               condition.column_name);

    const mdspan_type<const T> &span = *opt;
    std::vector<size_t> matches;
    for (size_t i = 0; i < span.extent(0); ++i) {
      if (condition.predicate(span, i)) {
        matches.push_back(i);
      }
    }
    return matches;
  }

  template <typename T> struct is_column_predicate : std::false_type {};

  template <typename T>
  struct is_column_predicate<ColumnPredicate<T>> : std::true_type {};

public:
#pragma region Constructors
  using BoolEnum = h5dispatch::BoolEnum;

  ReflectionTable();

  ReflectionTable(const std::vector<uint64_t> &experiment_ids,
                  const std::vector<std::string> &identifiers);

  ReflectionTable(const std::string &h5_filepath);
#pragma endregion

#pragma region Metadata Access
  const std::vector<uint64_t> &get_experiment_ids() const;

  void set_experiment_ids(const std::vector<uint64_t> &ids);

  const std::vector<std::string> &get_identifiers() const;

  void set_identifiers(const std::vector<std::string> &ids);

  std::pair<uint64_t, std::string> generate_new_attributes();

  std::vector<std::string> get_column_names() const;

  size_t size() const;
#pragma endregion

#pragma region Column Access
  template <typename T>
  std::optional<mdspan_type<T>> column(const std::string &name) {
    if (auto *typed = find_column<T>(*this, name)) {
      return typed->span();
    }
    return std::nullopt;
  }

  template <typename T>
  std::optional<mdspan_type<const T>> column(const std::string &name) const {
    if (const auto *typed = find_column<T>(*this, name)) {
      return typed->span();
    }
    return std::nullopt;
  }

  template <typename... Predicates>
  std::vector<size_t> find_rows(logic::LogicalOp op,
                                const Predicates &...conds) const {
    // Ensure all predicates are of the correct type
    static_assert((is_column_predicate<Predicates>::value && ...),
                  "All arguments to find_rows must be ColumnPredicate<T>");

    // Expand the parameter pack into a vector of row-index vectors
    std::vector<std::vector<size_t>> all_matches = {
        evaluate_column_predicate(conds)...};
    std::unordered_set<size_t> result_set;

    if (op == logic::LogicalOp::Or) {
      // Combine all matching row indices (set union)
      for (const auto &matches : all_matches) {
        result_set.insert(matches.begin(), matches.end());
      }
    } else if (op == logic::LogicalOp::And) {
      // Combine only common row indices (set intersection)
      if (all_matches.empty())
        return {};

      // Start with the first match set as the initial intersection base
      std::unordered_set<size_t> intersection(all_matches[0].begin(),
                                              all_matches[0].end());

      // Iteratively refine the intersection by retaining only elements
      // that are also present in the next predicate's results
      for (size_t i = 1; i < all_matches.size(); ++i) {
        std::unordered_set<size_t> next(all_matches[i].begin(),
                                        all_matches[i].end());

        // Iterate over current intersection set
        for (auto it = intersection.begin(); it != intersection.end();) {
          // If the current value is not present in the next result set,
          // erase it from the intersection set
          if (next.find(*it) == next.end()) {
            it = intersection.erase(it); // erase returns the next iterator
          } else {
            ++it; // move to the next element if it's still present
          }
        }
      }

      result_set = std::move(intersection);
    }

    return std::vector<size_t>(result_set.begin(), result_set.end());
  }

  template <typename... Predicates>
  std::vector<size_t> find_rows(const Predicates &...conds) const {
    return find_rows(logic::LogicalOp::And, conds...);
  }

  ReflectionTable select(const std::vector<size_t> &selected_rows) const;

  ReflectionTable select(const std::vector<bool> &mask) const;
#pragma endregion

#pragma region Column Modification
  template <typename T>
  void add_column(const std::string &name, const std::vector<size_t> &shape,
                  const std::vector<T> &column_data) {
    // Check if the type T is a bool. If so, convert to BoolEnum and add column.
    if constexpr (std::is_same_v<T, bool>) {
      std::vector<h5dispatch::BoolEnum> converted(column_data.size());
      for (size_t i = 0; i < column_data.size(); ++i) {
        converted[i] = column_data[i] ? h5dispatch::BoolEnum::TRUE
                                      : h5dispatch::BoolEnum::FALSE;
      }
      add_column<h5dispatch::BoolEnum>(name, shape, converted);
    } else {
      static_assert(h5dispatch::is_column_type_v<T>,
                    "Unsupported column type; see h5dispatch::ColumnTypes");
      TypedColumn<T> col(name, shape, column_data);

      // Check for duplicate column names
      for (const auto &existing_col : data) {
        if (column_name(existing_col) == name) {
          throw std::runtime_error("Column with name already exists: " + name);
        }
      }

      // Ensure row count consistency
      if (!data.empty() && col.rows() != get_row_count()) {
        throw std::runtime_error("Row count mismatch when adding column: " +
                                 name);
      }

      // Add the new column to the table
      data.emplace_back(std::move(col));
    }
  }

  template <typename T>
  void add_column(const std::string &name, const std::vector<T> &column_data) {
    add_column(name, std::vector<size_t>{column_data.size()}, column_data);
  }

  template <typename T>
  void add_column(const std::string &name, const size_t rows, const size_t cols,
                  const std::vector<T> &column_data) {
    add_column(name, std::vector<size_t>{rows, cols}, column_data);
  }

  template <typename T>
  void update_column(const std::string &name, const std::vector<size_t> &shape,
                     const std::vector<T> &column_data) {
    // Check if the type T is a bool. If so, convert to BoolEnum and update.
    if constexpr (std::is_same_v<T, bool>) {
      std::vector<h5dispatch::BoolEnum> converted(column_data.size());
      for (size_t i = 0; i < column_data.size(); ++i) {
        converted[i] = column_data[i] ? h5dispatch::BoolEnum::TRUE
                                      : h5dispatch::BoolEnum::FALSE;
      }
      update_column<h5dispatch::BoolEnum>(name, shape, converted);
    } else {
      static_assert(h5dispatch::is_column_type_v<T>,
                    "Unsupported column type; see h5dispatch::ColumnTypes");

      // Locate the existing column (strict: it must already exist)
      auto it = std::find_if(data.begin(), data.end(), [&](const Column &c) {
        return column_name(c) == name;
      });
      if (it == data.end()) {
        throw std::runtime_error("Column not found for update: " + name);
      }

      TypedColumn<T> col(name, shape, column_data);

      // Ensure row count consistency with the other columns in the table
      for (const auto &other : data) {
        if (column_name(other) != name) {
          if (col.rows() != column_rows(other)) {
            throw std::runtime_error(
                "Row count mismatch when updating column: " + name);
          }
          break;
        }
      }

      // Replace in place, preserving the column's position
      *it = std::move(col);
    }
  }

  template <typename T>
  void update_column(const std::string &name,
                     const std::vector<T> &column_data) {
    update_column(name, std::vector<size_t>{column_data.size()}, column_data);
  }

  template <typename T>
  void update_column(const std::string &name, const size_t rows,
                     const size_t cols, const std::vector<T> &column_data) {
    update_column(name, std::vector<size_t>{rows, cols}, column_data);
  }
#pragma endregion

#pragma region Write
  void write(std::string_view filename,
             std::string_view group = "/dials/processing/group_0") const;
#pragma endregion
};

#pragma endregion