From 8d2cf082d6c9213cfc25d188aea31158e05a2453 Mon Sep 17 00:00:00 2001 From: Johan Larsson Date: Thu, 30 Nov 2023 13:43:29 +0100 Subject: [PATCH] feat: support sparse X input --- sortedl1/estimators.py | 8 +++++++- src/main.cpp | 22 +++++++++++++++++----- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/sortedl1/estimators.py b/sortedl1/estimators.py index 4494a9b..b420cf1 100644 --- a/sortedl1/estimators.py +++ b/sortedl1/estimators.py @@ -1,5 +1,6 @@ import _sortedl1 as sl1 import numpy as np +from scipy import sparse from sklearn.base import BaseEstimator, RegressorMixin from sklearn.utils.validation import check_array, check_is_fitted, check_X_y @@ -53,7 +54,12 @@ def fit(self, X, y): alpha = np.atleast_1d(self.alpha).astype(np.float64) - result = sl1.fit_slope(X, y, lam, alpha) + if sparse.issparse(X): + fit_slope = sl1.fit_slope_sparse + else: + fit_slope = sl1.fit_slope_dense + + result = fit_slope(X, y, lam, alpha) self.intercept_ = result[0] self.sparse_coef_ = result[1] diff --git a/src/main.cpp b/src/main.cpp index 509763b..4f2d40b 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -9,10 +9,21 @@ using namespace pybind11::literals; namespace py = pybind11; pybind11::tuple -fit_slope(const Eigen::MatrixXd x, - const Eigen::MatrixXd y, - Eigen::ArrayXd lambda, - Eigen::ArrayXd alpha) +fit_slope_dense(const Eigen::MatrixXd x, + const Eigen::MatrixXd y, + Eigen::ArrayXd lambda, + Eigen::ArrayXd alpha) +{ + auto result = slope::slope(x, y, alpha, lambda); + + return py::make_tuple(result.beta0s, result.betas); +} + +pybind11::tuple +fit_slope_sparse(const Eigen::SparseMatrix x, + const Eigen::MatrixXd y, + Eigen::ArrayXd lambda, + Eigen::ArrayXd alpha) { auto result = slope::slope(x, y, alpha, lambda); @@ -21,5 +32,6 @@ fit_slope(const Eigen::MatrixXd x, PYBIND11_MODULE(_sortedl1, m) { - m.def("fit_slope", &fit_slope); + m.def("fit_slope_dense", &fit_slope_dense); + m.def("fit_slope_sparse", &fit_slope_sparse); }