diff --git a/src/equisolve/numpy/models/linear_model.py b/src/equisolve/numpy/models/linear_model.py index cb21433..6b58d1d 100644 --- a/src/equisolve/numpy/models/linear_model.py +++ b/src/equisolve/numpy/models/linear_model.py @@ -9,6 +9,7 @@ from typing import List, Optional, Union import numpy as np +import scipy.linalg from equistore import Labels, TensorBlock, TensorMap from equistore.operations import dot, multiply, ones_like, slice from equistore.operations._utils import _check_blocks, _check_maps @@ -118,7 +119,7 @@ def _validate_params( "properties." ) - def _solver( + def _lstsq_solver( self, X: TensorBlock, y: TensorBlock, @@ -168,6 +169,46 @@ def _solver( return weights_block + def _solve_solver( + self, + X: TensorBlock, + y: TensorBlock, + alpha: TensorBlock, + sample_weight: TensorBlock, + ) -> TensorBlock: + """A regularized solver using ``scipy.linalg.lstsq``.""" + + # Convert TensorMaps into arrays for processing them with NumPy. + + # X_arr has shape of (n_targets, n_properties) + X_arr = block_to_array(X, self.parameter_keys) + + # y_arr has shape lentgth of n_targets + y_arr = block_to_array(y, self.parameter_keys) + + # sw_arr has shape of (n_samples, 1) + sw_arr = block_to_array(sample_weight, self.parameter_keys) + + # alpha_arr has shape of (1, n_properties) + alpha_arr = block_to_array(alpha, ["values"]) + + X_eff = X_arr / sw_arr + y_eff = y_arr / sw_arr + + X_eff = X_arr.T @ X_arr + np.diag(alpha_arr[0, :]) + y_eff = X_arr.T @ y_arr[:, 0] + + w = scipy.linalg.solve(X_eff, y_eff, assume_a="pos", overwrite_a=True).ravel() + + weights_block = TensorBlock( + values=w.reshape(1, -1), + samples=y.properties, + components=[], + properties=X.properties, + ) + + return weights_block + def fit( self, X: TensorMap, @@ -175,6 +216,7 @@ def fit( alpha: Union[float, TensorMap] = 1.0, sample_weight: Union[float, TensorMap] = 1.0, rcond: float = 1e-13, + solver="solve", ) -> None: """Fit a regression model to each block in `X`. @@ -229,9 +271,16 @@ def fit( alpha_block = alpha.block(key) sw_block = sample_weight.block(key) - weight_block = self._solver(X_block, y_block, alpha_block, sw_block, rcond) + if solver == "lstsq": + weights = self._lstsq_solver( + X_block, y_block, alpha_block, sw_block, rcond + ) + elif solver == "solve": + weights = self._solve_solver(X_block, y_block, alpha_block, sw_block) + else: + raise ValueError("Unknown solver") - weights_blocks.append(weight_block) + weights_blocks.append(weights) # convert weights to a dictionary allowing pickle dump of an instance self._weights = tensor_map_to_dict(TensorMap(X.keys, weights_blocks)) diff --git a/tests/equisolve_tests/numpy/models/test_linear_model.py b/tests/equisolve_tests/numpy/models/test_linear_model.py index 615e13e..1667ea5 100644 --- a/tests/equisolve_tests/numpy/models/test_linear_model.py +++ b/tests/equisolve_tests/numpy/models/test_linear_model.py @@ -59,9 +59,7 @@ def to_equistore(self, X_arr=None, y_arr=None, alpha_arr=None, sw_arr=None): def equisolve_solver_from_numpy_arrays(self, X_arr, y_arr, alpha_arr, sw_arr=None): X, y, alpha, sw = self.to_equistore(X_arr, y_arr, alpha_arr, sw_arr) - clf = Ridge( - parameter_keys="values", - ) + clf = Ridge(parameter_keys="values") clf.fit(X=X, y=y, alpha=alpha, sample_weight=sw) return clf