SUR.hpp

#pragma once
#include "creativity/data/Equation.hpp"
#include <Eigen/Core>
#include <vector>
#include <algorithm>
#include <ostream>
#include <limits>
#include <map>

namespace creativity { namespace data {

class SUR {
    public:
        template <class... Eqns>
        explicit SUR(Eqns... eqns);

        template <class... MoreEqns>
        void add(const Equation &eq1, MoreEqns... eqns);

        template <class... MoreEqns>
        void add(Equation &&eq1, MoreEqns... eqns);

        const std::vector<Equation>& equations() const;

        unsigned int n() const { return eqs_.empty() ? 0 : eqs_.front().depVar()->size(); }

        void gather();

        void solve();

        void clear();

        Eigen::VectorBlock<const Eigen::VectorXd> beta(unsigned i) const;

        std::vector<std::string> varNames(unsigned i) const;

        const unsigned& k(unsigned i) const;

        int df(unsigned i) const;

        Eigen::Block<const Eigen::MatrixXd> covariance(unsigned i) const;

        Eigen::VectorBlock<const Eigen::VectorXd> se(unsigned i) const;

        Eigen::VectorBlock<const Eigen::VectorXd> tRatios(unsigned i) const;

        Eigen::VectorBlock<const Eigen::VectorXd> pValues(unsigned i) const;

        static const std::map<double, std::string> default_pstar_threshold;

        std::vector<std::string> pStars(unsigned i, const std::map<double, std::string> &threshold = default_pstar_threshold) const;

        const double& s2(unsigned i) const;

        Eigen::VectorBlock<const Eigen::VectorXd> residuals(unsigned i) const;

        const double& ssr(unsigned i) const;

        const double& Rsq(unsigned i) const;

        const Eigen::VectorXd& y() const;

        Eigen::VectorBlock<const Eigen::VectorXd> y(unsigned i) const;

        const std::vector<Eigen::MatrixXd>& X() const;

        Eigen::MatrixX4d summary(unsigned i) const;

        friend std::ostream& operator<<(std::ostream &os, const SUR &ols);

    protected:
        void add() {}

        std::vector<Equation> eqs_;

        bool gathered_ = false;

        Eigen::VectorXd y_;

        std::vector<Eigen::MatrixXd> X_;

        std::vector<unsigned> k_;

        std::vector<unsigned> offset_;

        bool solved_ = false;

        Eigen::VectorXd beta_;

        Eigen::MatrixXd var_beta_;

        Eigen::VectorXd se_;

        Eigen::VectorXd t_ratios_;

        Eigen::VectorXd p_values_;

        Eigen::VectorXd residuals_;

        std::vector<double>
            ssr_, 
            s2_, 
            R2_; 

        void requireGathered() const { if (!gathered_) throw std::logic_error("Cannot access model data before calling gather()"); }

        void requireSolved() const { if (!solved_) throw std::logic_error("Cannot obtain model estimates before calling solve()"); }

        void requireEquation(unsigned i) const { if (i >= eqs_.size()) throw std::out_of_range("Equation " + std::to_string(i) + " does not exist"); }

        void requireGathered(unsigned i) const { requireGathered(); requireEquation(i); }

        void requireSolved(unsigned i) const { requireSolved(); requireEquation(i); }
};


template <class... Eqns>
SUR::SUR(Eqns... eqns) {
    add(std::forward<Eqns>(eqns)...);
}

template <class... MoreEqns>
void SUR::add(const Equation &eq1, MoreEqns... eqns) {
    if (offset_.empty()) offset_.push_back(0);
    else offset_.push_back(offset_.back() + k_.back());
    k_.push_back(eq1.numVars());
    eqs_.push_back(eq1);
    add(std::forward<MoreEqns>(eqns)...);
}

template <class... MoreEqns>
void SUR::add(Equation &&eq1, MoreEqns... eqns) {
    if (offset_.empty()) offset_.push_back(0);
    else offset_.push_back(offset_.back() + k_.back());
    k_.push_back(eq1.numVars());
    eqs_.push_back(std::move(eq1));
    add(std::forward<MoreEqns>(eqns)...);
}

}}