GroupPartition.hpp

Go to the documentation of this file.

#pragma once
 
 
#include "utils/Types.hpp"
#include "utils/Invariant.hpp"
 
#include <map>
#include <optional>
#include <set>
#include <vector>
#include <Eigen/Dense>
 
namespace ppforest2::stats {
  class GroupPartition {
    using Group       = types::GroupId;
    using GroupSet    = std::set<types::GroupId>;
    using GroupMap    = std::map<types::GroupId, types::GroupId>;
    using GroupInvMap = std::map<types::GroupId, GroupSet>;
    using GroupVector = types::GroupIdVector;
 
  public:
    explicit GroupPartition(types::GroupIdVector const& y);

    explicit GroupPartition(types::OutcomeVector const& y);

    GroupPartition(int start, int end);

    GroupPartition bisect(int mid) const;

    int group_start(Group const& group) const;
    int group_end(Group const& group) const;
    int group_size(Group const& group) const;

    Group first_group() const {
      invariant(!groups.empty(), "GroupPartition::first_group: partition is empty");
      return *groups.begin();
    }

    int total_size() const;

    template<typename Derived> auto group(Eigen::MatrixBase<Derived> const& x, Group const& group) const {
      std::vector<int> indices;
 
      auto const& subs = this->subgroups.at(group);
 
      for (auto const& g : subs) {
        for (int i = group_start(g); i <= group_end(g); ++i) {
          invariant(i >= 0 && i < x.rows(), "GroupPartition::group: index out of bounds");
          indices.push_back(i);
        }
      }
 
      return x(indices, Eigen::all);
    }

    Eigen::VectorXi group_indices(Group const& group) const {
      std::vector<int> indices;
 
      auto const& subs = this->subgroups.at(group);
 
      for (auto const& g : subs) {
        for (int i = group_start(g); i <= group_end(g); ++i) {
          indices.push_back(i);
        }
      }
 
      return Eigen::VectorXi::Map(indices.data(), static_cast<Eigen::Index>(indices.size()));
    }

    template<typename Derived> auto data(Eigen::MatrixBase<Derived> const& x) const {
      std::vector<int> indices;
 
      for (auto const& kv : blocks) {
        auto const& g = kv.first;
        for (int i = group_start(g); i <= group_end(g); ++i) {
          indices.push_back(i);
        }
      }
 
      return x(indices, Eigen::all);
    }

    types::FeatureVector mean(types::FeatureMatrix const& x) const;
    types::FeatureMatrix bgss(types::FeatureMatrix const& x) const;
    types::FeatureMatrix wgss(types::FeatureMatrix const& x) const;

    GroupPartition subset(GroupSet const& groups) const;
 
    using SplitSizes = std::map<types::GroupId, int>;

    std::pair<GroupPartition, GroupPartition> split(SplitSizes const& left_sizes) const;

    GroupPartition remap(GroupMap const& mapping) const;

    GroupPartition collapse() const;

    GroupSet const groups;
    GroupMap const supergroups;
    GroupInvMap const subgroups;
 
  private:
    struct Block {
      int start                          = 0;
      int end                            = 0;
      int size                           = 0;
      std::optional<types::GroupId> next = std::nullopt;
      std::optional<types::GroupId> prev = std::nullopt;
    };
 
    using BlockMap = std::map<types::GroupId, Block>;
    BlockMap const blocks;
 
    static BlockMap init_blocks(GroupVector const& y);
    static GroupMap init_supergroups(GroupSet const& groups);
 
    explicit GroupPartition(BlockMap const& blocks);
 
    GroupPartition(BlockMap const& blocks, GroupSet const& groups);
 
    GroupPartition(BlockMap const& blocks, GroupMap const& supergroups);
 
    GroupPartition(BlockMap const& blocks, GroupSet const& groups, GroupMap const& supergroups);
  };
}