TorchScript compatibility

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,13 +169,54 @@ 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,
         y: TensorMap,
         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))

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