continuos_cache.hpp

Go to the documentation of this file.

#pragma once

 
#include "../../utils/ClusterInfo.hpp"

 
class ContinuosCache : public ClusterInfo {
public:
    // Important to declare the struct useful.
    // HINT: If the name changed modified the file which uses it accordingly.
    struct ClusterStats {
        int n = 0;
        double sum = 0.0;
        double sumsq = 0.0;
    };
 
private:
    std::vector<ClusterStats> cluster_stats;
 
public:
    const Eigen::VectorXd continuos_covariates;
 
    ContinuosCache(Eigen::VectorXi &initial_allocations, const Eigen::VectorXd &continuos_covariates)
        : continuos_covariates(continuos_covariates) {
        cluster_stats.clear();
 
        const int K = initial_allocations.maxCoeff() + 1;
        cluster_stats.resize(K);
 
        // Initialize cluster statistics based on initial allocations
        for (int i = 0; i < initial_allocations.size(); ++i) {
            int cluster = initial_allocations(i);
            if (cluster < 0)
                continue;
 
            double value = continuos_covariates(i);
            ClusterStats &stats = cluster_stats[cluster];
            stats.n++;
            stats.sum += value;
            stats.sumsq += value * value;
        }
    }

    void set_allocation(int index, int cluster, int old_cluster) override;

    inline ClusterStats get_cluster_stats(int cluster) const { return cluster_stats[cluster]; }

    inline const ClusterStats &get_cluster_stats_ref(int cluster) const { return cluster_stats[cluster]; }

    void recompute(const int K, const Eigen::VectorXi &allocations) override;

    inline void move_cluster_info(int from_cluster, int to_cluster) override {
        cluster_stats[to_cluster] = std::move(cluster_stats[from_cluster]);
    };

    void remove_info(int cluster) override;
};