C++: Linear Regression and Polynomial Regression

Question

I wrote a simple linear/polynomial regressor based on my previous matrix project (https://github.com/frozenca/Ndim-Matrix).

#include <Matrix.h>
#include <algorithm>
#include <concepts>
#include <cctype>
#include <functional>
#include <stdexcept>
#include <numeric>
#include <ranges>
#include <string>
#include <vector>

namespace frozenca {

class LinearRegression {
private:
    std::vector<float> theta_;
    std::vector<std::vector<std::size_t>> monomials_;

    enum class Penalty {
        None,
        L1,
        L2
    };

public:
    void fit(const std::vector<std::vector<float>>& X, const std::vector<float>& y,
             const std::pair<std::string, float>& penalty = {}, std::size_t degree = 1,
             std::size_t batch_size = 0) {
        if (X.empty() || y.empty()) {
            throw std::invalid_argument("Data set is empty");
        }
        if (X.size() != y.size()) {
            throw std::invalid_argument("Data set sizes do not match");
        }
        const std::size_t n = X[0].size();
        if (std::ranges::any_of(X, [&n](const auto& row) { return row.size() != n;})) {
            throw std::invalid_argument("Feature set sizes are not equal");
        }
        auto [p_type, alpha] = penalty;
        Penalty penalty_type = Penalty::None;
        std::ranges::transform(p_type, p_type.begin(), [](auto c){ return std::tolower(c);});
        if (p_type == "") {
            penalty_type = Penalty::None;
            alpha = 0.0f;
        } else if (p_type == "l1" || p_type == "lasso") {
            penalty_type = Penalty::L1;
            if (!batch_size) {
                batch_size = 1; // L1 regression only can be fit via gradient descent
            }
        } else if (p_type == "l2" || p_type == "ridge") {
            penalty_type = Penalty::L2;
        } else {
            throw std::invalid_argument("Unknown penalty type");
        }
        if (alpha < 0.0f) {
            throw std::invalid_argument("Alpha value must be non-negative");
        }
        if (degree == 0) {
            throw std::invalid_argument("Degree must be nonzero");
        }
        const std::size_t num_samples = X.size();
        if (batch_size && num_samples % batch_size) {
            throw std::invalid_argument("Batch size must divide number of samples");
        }

        std::size_t num_features = 1;
        for (std::size_t k = 1; k <= degree; ++k) {
            num_features *= (n + k);
            num_features /= k;
        }

        monomials_ = get_monomials(n, degree);

        Mat<float> M(num_samples, num_features);
        for (std::size_t i = 0; i < num_samples; ++i) {
            auto m_i = M.row(i);
            construct_features(m_i, X[i], monomials_);
        }

        Vec<float> Y(num_samples);
        for (std::size_t i = 0; i < num_samples; ++i) {
            Y[i] = y[i];
        }

        fit_theta(M, Y, penalty_type, alpha, batch_size);
    }

    [[nodiscard]] std::vector<float> getTheta() {
        return theta_;
    }

    float predict(const std::vector<float>& x) {
        if (theta_.empty()) {
            throw std::runtime_error("The model has not been fit yet!");
        }
        auto res = 0.0f;
        for (std::size_t i = 0; i < theta_.size(); ++i) {
            res += theta_[i] * eval_monomial(x, monomials_[i]);
        }
        return res;
    }

private:

    void fit_theta(const Mat<float>& X, const Vec<float>& Y,
                   const Penalty& penalty_type, float alpha, std::size_t batch_size) {
        const std::size_t num_samples = X.dims(0);
        const std::size_t num_features = X.dims(1);
        Vec<float> Theta (num_features);
        if (!batch_size) {
            Mat<float> l2_pen = identity<float>(num_features);
            l2_pen *= alpha;
            Theta = dot(dot(inv(dot(transpose(X), X) + l2_pen), transpose(X)), Y);
        } else {
            // gradient descent

            // random initialization
            std::mt19937 gen(std::random_device{}());
            std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
            for (std::size_t i = 0; i < num_features; ++i) {
                Theta[i] = dist(gen);
            }

            constexpr std::size_t n_epochs = 100;
            constexpr float tolerance = 1e-7;
            std::size_t num_batches = num_samples / batch_size;
            assert(!(num_samples % batch_size));
            std::vector<std::size_t> batch_indices (num_batches);
            std::iota(batch_indices.begin(), batch_indices.end(), 0lu);
            std::ranges::shuffle(batch_indices, gen);

            bool converged = false;
            for (std::size_t i = 0; i < n_epochs; ++i) {
                if (converged) {
                    break;
                }
                for (std::size_t b = 0; b < num_batches; ++b) {
                    auto xi = X.submatrix({batch_indices[b] * batch_size, 0}, {(batch_indices[b] + 1) * batch_size, num_features});
                    auto yi = Y.submatrix(batch_indices[b] * batch_size, (batch_indices[b] + 1) * batch_size);
                    auto gradient = dot(transpose(xi), dot(xi, Theta) - yi);
                    float lr = 0.1f * n_epochs * num_batches / (n_epochs + i) / (num_batches + b);
                    gradient *= lr;
                    if (penalty_type == Penalty::None) {
                        Theta = Theta - gradient;
                    } else if (penalty_type == Penalty::L2) {
                        auto pen = Theta;
                        pen *= lr * alpha;
                        Theta = Theta - gradient - pen;
                    } else if (penalty_type == Penalty::L1) {
                        Vec<float> pen (num_features);
                        for (std::size_t f = 0; f < num_features; ++f) {
                            if (Theta[f] > 0.0f) {
                                pen[f] = 1.0f;
                            } else if (Theta[f] < 0.0f) {
                                pen[f] = -1.0f;
                            }
                        }
                        pen *= lr * alpha;
                        Theta = Theta - gradient - pen;
                    }
                    if (norm(gradient) < tolerance) {
                        converged = true;
                        break;
                    }
                }
            }
        }
        theta_.clear();
        for (std::size_t i = 0; i < num_features; ++i) {
            theta_.push_back(Theta[i]);
        }
    }

    // for n = 2 and degree 2,
    // 1, x, y, x^2, xy, y^2
    // {(0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (0, 2)}
    static std::vector<std::vector<std::size_t>> get_monomials(std::size_t n, std::size_t deg) {
        std::vector<std::vector<std::size_t>> d_monomials;
        d_monomials.push_back({0});
        construct_monomials(d_monomials, n, deg);
        std::ranges::sort(d_monomials);
        return d_monomials;
    }

    // recursive construct.
    // 0 -> 0
    //   -> 1
    //   -> 2
    // 1 -> 0
    //   -> 1
    // 2 -> 0
    static void construct_monomials(std::vector<std::vector<std::size_t>>& d_monomials, std::size_t n,
                               std::size_t deg) {
        if (!n) {
            return;
        }
        std::vector<std::vector<std::size_t>> new_monomials;
        while (!d_monomials.empty()) {
            auto back = d_monomials.back();
            std::size_t s = back[0];
            d_monomials.pop_back();
            for (std::size_t d = 0; d <= deg - s && d <= deg; ++d) {
                new_monomials.push_back(back);
                new_monomials.back().push_back(d);
                new_monomials.back()[0] += d;
            }
        }
        d_monomials = new_monomials;
        construct_monomials(d_monomials, n - 1, deg);
    }

    static float eval_monomial(const std::vector<float>& sample, const std::vector<std::size_t>& monomial) {
        return std::inner_product(sample.begin(), sample.end(), monomial.begin() + 1, 0.0f,
                                  std::plus<>(), [](auto s, auto exp) {
                    return std::pow(s, 1.0f * exp);
        });
    }

    static void construct_features(VecView<float>& row, const std::vector<float>& sample,
                            const std::vector<std::vector<std::size_t>>& monomials) {
        const std::size_t num_features = row.dims(0);
        const std::size_t n = sample.size();
        for (std::size_t i = 0; i < num_features; ++i) {
            row[i] = eval_monomial(sample, monomials[i]);
        }
    }

};

Test code:

    std::mt19937 gen(std::random_device{}());

    // linear regression example
    constexpr std::size_t num_samples = 100;
    std::uniform_real_distribution<float> noise(-1.0f, 1.0f);

    {
        std::vector<std::vector<float>> X;
        std::vector<float> y;

        // y = 3x + 4
        for (std::size_t i = 0; i < num_samples; ++i) {
            X.push_back({2.0f + noise(gen)});
            y.push_back(4.0f + X.back()[0] * 3.0f + noise(gen));
        }

        fc::LinearRegression linreg;
        linreg.fit(X, y);
        auto theta = linreg.getTheta();
        for (auto t: theta) {
            std::cout << t << ' ';
        }
        std::cout << '\n';

        std::cout << linreg.predict({1.0}) << '\n';
    }

    {
        std::vector<std::vector<float>> X;
        std::vector<float> y;

        // y = a^2 + 2ab + 3b^2 + 4a + 5b + 6
        for (std::size_t i = 0; i < num_samples; ++i) {
            auto a = 0.0f + noise(gen);
            auto b = 0.0f + noise(gen);
            X.push_back({a, b});
            y.push_back(a * a + 2.0f * a * b + 3.0f * b * b + 4.0f * a + 5.0f * b + 6.0f + noise(gen));
        }

        fc::LinearRegression linreg;
        linreg.fit(X, y, {}, 2);
        auto theta = linreg.getTheta();
        for (auto t: theta) {
            std::cout << t << ' ';
        }
        std::cout << '\n';
    }

    {
        std::vector<std::vector<float>> X;
        std::vector<float> y;

        // y = a^2 + 2ab + 3b^2 + 4a + 5b + 6
        // L1 regression
        for (std::size_t i = 0; i < num_samples; ++i) {
            auto a = 0.0f + noise(gen);
            auto b = 0.0f + noise(gen);
            X.push_back({a, b});
            y.push_back(a * a + 2.0f * a * b + 3.0f * b * b + 4.0f * a + 5.0f * b + 6.0f + noise(gen));
        }

        fc::LinearRegression linreg;
        linreg.fit(X, y, {"L1", 0.1f}, 2);
        auto theta = linreg.getTheta();
        for (auto t: theta) {
            std::cout << t << ' ';
        }
        std::cout << '\n';
    }

Result

3.78924 3.12022
6.90945
-1.13539 4.26815 2.09389 2.83393 0.149067 0.781616
0.327844 1.17 0.244501 2.76427 3.45505 0.812552

Feel free to comment anything!

Answer 1

My first thought is to change the LinearRegression::fit() method into a constructor called with the same arguments--that is, rename void fit() to LinearRegression(). The entire point of this class is to create a fit for data, so have the class create the fit when it is first created. As it is now, you need to lines to set up the class:

fc::LinearRegression linreg;
linreg.fit(X, y);

Before the call to fit(), linreg is in a useless zombie state where none of its methods return anything useful. By making the constructor create the fit, these two lines become

fc::LinearRegression linreg(X, y);

Now, there's no chance of trying to use the linear regression before feeding it the data. Plus, linreg can be declared const if needed. Speaking of which, ...

---

All of the other non-static classes methods should be marked with const because they do not change the member variable data. This will allow these methods to be called when a user creates a const instance of LinearRegression or passes an instance to a function as a const LinearRegression& or other parameter with const.

---

bool converged = false;
for (std::size_t i = 0; i < n_epochs; ++i) {
    if (converged) {
        break;
    }
    // ...
}

This can also be written as

bool converged = false;
for (std::size_t i = 0; i < n_epochs && !converged; ++i) {
    // ...
}

---

Put an empty line before a new if() block. Otherwise, a reader might suspect that the if() should have been an else if() since it is grouped with the previous if() block.

Answer 2

Separate fitting from storing the result of the fit

Mark H already pointed out in his answer that it weird to have a class that you first have to construct, then have to call the member function fit() for it to store the result, and the suggestion was to make the constructor do the fitting. I would go further than that, and split up class LinearRegression into a free function fit(), and a class that holds the result of the fit. However, the result of a polynomial fit is just a polynomial. So I would create a class Polynomial that stores the result of the fit, so that you could write:

std::vector<std::vector<float>> X = {...};
std::vector<float> y = {...};

Polynomial polynomial = fc::fit(X, y);

for (auto coeff: polynomial.getCoeffs())
    std::cout << coeff << '\n';

std::cout << polynomial({1.0}) << '\n';

It would be really nice if Polynomial behaved like an arithmetic type, so you could add two polynomials together for example using +, and so on.