Repository: secretflow/spu (language: python)

Issue:
[linear-ridge-svd]使用 SPU 实现线性回归模型
> 此 ISSUE 为 [隐语开源共建计划（SecretFlow Open Source Contribution Plan，简称 SF OSCP）第二期](https://github.com/orgs/secretflow/discussions/741)  https://github.com/secretflow/spu/issues/249 子任务，若对该任务感兴趣，欢迎在 https://github.com/secretflow/spu/issues/249 提交你的不同设计思路。
> 也欢迎其他开发者点击[这里](https://github.com/orgs/secretflow/discussions/741)，选择其他你感兴趣的任务，热烈欢迎社区开发者参与共建～

## 设计思路
by：[magic-hya](https://github.com/magic-hya)
1. 参考sklearn的ridge算法，在jax上实现
2. 实现后与明文sklearn ridge算法结果做对比


## 任务介绍

- 任务名称：使用 SPU 实现线性回归模型
- 技术方向：SPU/SML
- 任务难度：热身🌟

## 详细要求
- 安全性（尽量少 reveal）
- 功能性：实现算法的基本功能，包括：
  - 支持 fit（至少一种 eg：sgd, lbfgs, svd...），predict 
  - 提供配置是否需要拟合 bias 项，是否需要正则等基本参数
  - 支持起码一种正则方式（L2, L1, elastic...）
- 收敛性：包含 simulator 跑出的实验数据并且证明收敛性/准确性（最好有与明文 sklearn 结果的对比，可详见[参考 PR](https://github.com/secretflow/spu/pull/240)）
- 代码规范：Python 代码需要使用 black+isort 进行格式化（流水线包含代码规范检查卡点）
- 提交说明：关联该 isuue 并提交代码至 https://github.com/secretflow/spu/tree/main/sml
- 特殊说明：若某个特性有特殊的限制，如需要 FM128，需要更多 fxp 等需要在注释文档中明确说明

## 能力要求
- 熟悉经典的机器学习算法
- 熟悉 JAX 或 NumPy，可以使用 NumPy 实现算法

## 操作说明
- 操作指引：https://www.secretflow.org.cn/docs/spu/latest/en-US/getting_started/tutorials/develop_your_first_mpc_application
- 可参考明文实现：
  - [linear_model/_base](https://github.com/scikit-learn/scikit-learn/blob/7f9bad99d/sklearn/linear_model/_base.py#L534)
  - [linear_model/_ridge](https://github.com/scikit-learn/scikit-learn/blob/7f9bad99d/sklearn/linear_model/_ridge.py#L915)
-  [范例](https://github.com/secretflow/spu/blob/main/sml/linear_model/simple_sgd.py)
-  [参考 PR](https://github.com/secretflow/spu/pull/240)


## 认领说明
本任务可有多种实现方式，故支持一个任务有多位开发者进行认领，请在认领任务后，在该 issue 下 comment 你的具体设计思路。

**设计思路说明**：简单说明计划使用什么算法 or 什么优化器实现任务需求即可

> 当同一任务有多位开发者提交设计思路时：
> -  若多位开发者的设计思路均类似，则将按照 comment 的时间 assign 给第一位 comment 的开发者；
> -  若多位开发者的设计思路均不同，则该 issue 将会拆分为以不同设计思路进行实现的子 issue，并 assign 给对应开发者。

## 开发须知
以下部分代码请必须增加代码注释，对对应代码模块进行说明，包括：
-  __init__函数的超参数含义
-  fit 的具体算法实现说明


Discussion / hints:
magic-hya Give it to me
hello，，请问您当前实现了哪种solver去实现ridge regression？
> hello，，请问您当前实现了哪种solver去实现ridge regression？

hello, 我是参考sklearn里面的做法实现，里面实现的方法很多，我也打算逐步实现里面的内容，丰富算法，目前第一个版本先实现简单的svd(奇异值分解)
实现一个solver以后就可以直接先发pr哈，不然这个pr的持续时间就太长了～

如果您还有实现其他solver的打算，可以再发起其他pr～

Thanks
大概明后天吧，我会把我第一版的代码提交pr
> > hello，，请问您当前实现了哪种solver去实现ridge regression？
> 
> hello, 我是参考sklearn里面的做法实现，里面实现的方法很多，我也打算逐步实现里面的内容，丰富算法，目前第一个版本先实现简单的svg

sorry，，svg？是svd还是sag呀？
> > > hello，，请问您当前实现了哪种solver去实现ridge regression？
> > 
> > 
> > hello, 我是参考sklearn里面的做法实现，里面实现的方法很多，我也打算逐步实现里面的内容，丰富算法，目前第一个版本先实现简单的svg
> 
> sorry，，svg？是svd还是sag呀？

sorry，svd 奇异值分解
目前本地测试没有问题，但是执行emul和test会报错。
执行`python sml/linear_model/ridge_emul.py`
错误：
`Traceback (most recent call last):
  File "sml/linear_model/ridge_emul.py", line 31, in <module>
    import sml.utils.emulation as emulation
  File "/data/linear-ridge/spuxx/sml/utils/emulation.py", line 85, in <module>
    class Emulator:
  File "/data/linear-ridge/spuxx/sml/utils/emulation.py", line 163, in Emulator
    copts=spu_pb2.CompilerOptions(),
AttributeError: module 'spu.spu_pb2' has no attribute 'CompilerOptions'`

执行`python sml/linear_model/ridge_emul.py`
错误：
`Traceback (most recent call last):
  File "sml/linear_model/ridge_test.py", line 49, in test_simple
    result = spsim.sim_jax(sim, proc)(x1, x2, y)
  File "/usr/local/lib/python3.8/site-packages/spu/utils/simulation.py", line 152, in wrapper
    executable, output = spu_fe.compile(
  File "/usr/local/lib/python3.8/site-packages/spu/utils/frontend.py", line 112, in compile
    ir_text, output = _jax_compilation(fn, static_argnums, args, kwargs)
  File "/usr/local/lib/python3.8/site-packages/cachetools/__init__.py", line 737, in wrapper
    v = func(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/spu/utils/frontend.py", line 41, in _jax_compilation
    cfn, output = jax.xla_computation(
  File "sml/linear_model/ridge_test.py", line 42, in proc
    model.fit(x, y).predict_proba(x)
  File "/data/linear-ridge/spuxx/sml/linear_model/ridge.py", line 78, in fit
    self.coef = _solve_svd(x, y, alpha)
  File "/data/linear-ridge/spuxx/sml/linear_model/ridge.py", line 105, in _solve_svd
    s_nnz = s[idx][:, jnp.newaxis]
  File "/usr/local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3907, in _rewriting_take
    treedef, static_idx, dynamic_idx = _split_index_for_jit(idx, arr.shape)
  File "/usr/local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3985, in _split_index_for_jit
    idx = _expand_bool_indices(idx, shape)
  File "/usr/local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 4299, in _expand_bool_indices
    raise errors.NonConcreteBooleanIndexError(abstract_i)
jax._src.errors.NonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[10])`

详细代码见附件ridge.zip
[ridge.zip](https://github.com/secretflow/spu/files/12290089/ridge.zip)

请用 bazel test/run 的方式运行 test 和 emul
目前遇到问题，无法通过测试，执行测试报错，主要代码为以下函数：
`

def _solve_svd(x, y, alpha): 

     U, s, Vt = jsci.linalg.svd(x, full_matrices=False)
     s_nnz = s[:, jnp.newaxis]
     UTy = jnp.dot(U.T, y)
     d = s_nnz / (s_nnz**2 + alpha)
     d_UT_y = d * UTy
     return jnp.dot(Vt.T, d_UT_y).T`

错误如下：`
jax._src.traceback_util.UnfilteredStackTrace: NotImplementedError: MLIR translation rule for primitive 'eigh' not found for platform interpreter
`
spu暂时不支持svd，eigh等于特征值相关的api，近期可能会commit一个utils文件，里面会有一个简单的svd实现。您可以：

1. 使用其他的解法
2. 自己用jax实现一个svd
@magic-hya 可以参考下这个 issue #213 
> spu暂时不支持svd，eigh等于特征值相关的api，近期可能会commit一个utils文件，里面会有一个简单的svd实现。您可以：
> 
> 1. 使用其他的解法
> 2. 自己用jax实现一个svd

使用cholesky实现了，也跑通了emul和test测试，跟sklearn结果误差在0.01-0.02之间，我先提交个pr，然后看下还有什么要改的

Fix patch:
diff --git a/sml/linear_model/BUILD.bazel b/sml/linear_model/BUILD.bazel
index e3942c7c..9c3e564e 100644
--- a/sml/linear_model/BUILD.bazel
+++ b/sml/linear_model/BUILD.bazel
@@ -31,3 +31,8 @@ py_library(
         "//sml/utils:fxp_approx",
     ],
 )
+
+py_library(
+    name = "ridge",
+    srcs = ["ridge.py"],
+)
diff --git a/sml/linear_model/emulations/BUILD.bazel b/sml/linear_model/emulations/BUILD.bazel
index 6c8f6329..8a6c5245 100644
--- a/sml/linear_model/emulations/BUILD.bazel
+++ b/sml/linear_model/emulations/BUILD.bazel
@@ -34,3 +34,13 @@ py_binary(
         "//sml/utils:emulation",
     ],
 )
+
+py_binary(
+    name = "ridge_emul",
+    srcs = ["ridge_emul.py"],
+    deps = [
+        "//examples/python/utils:dataset_utils",  # FIXME: remove examples dependency
+        "//sml/linear_model:ridge",
+        "//sml/utils:emulation",
+    ],
+)
diff --git a/sml/linear_model/emulations/ridge_emul.py b/sml/linear_model/emulations/ridge_emul.py
new file mode 100644
index 00000000..f8980853
--- /dev/null
+++ b/sml/linear_model/emulations/ridge_emul.py
@@ -0,0 +1,79 @@
+# Copyright 2023 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import jax.numpy as jnp
+from sklearn.linear_model import Ridge as skRidge
+
+import examples.python.utils.dataset_utils as dsutil
+import sml.utils.emulation as emulation
+from sml.linear_model.ridge import Ridge
+
+
+# TODO: design the enumation framework, just like py.unittest
+# all emulation action should begin with `emul_` (for reflection)
+def emul_Ridge(mode: emulation.Mode.MULTIPROCESS):
+    def proc(x1, x2, y):
+        model = Ridge(alpha=1.0, solver="cholesky")
+
+        x = jnp.concatenate((x1, x2), axis=1)
+        y = y.reshape((y.shape[0], 1))
+
+        return model.fit(x, y).predict(x)
+
+    def load_data():
+        dataset_config = {
+            "use_mock_data": False,
+            "problem_type": "regression",
+            "builtin_dataset_name": "diabetes",
+            "left_slice_feature_ratio": 0.5,
+        }
+
+        x1, x2, y = dsutil.load_dataset_by_config(dataset_config)
+
+        return x1, x2, y
+
+    try:
+        # bandwidth and latency only work for docker mode
+        emulator = emulation.Emulator(
+            emulation.CLUSTER_ABY3_3PC, mode, bandwidth=300, latency=20
+        )
+        emulator.up()
+
+        # load mock data
+        x1, x2, y = load_data()
+
+        # sklearn test
+        x = jnp.concatenate((x1, x2), axis=1)
+        sklearn_result = (
+            skRidge(alpha=1, solver='cholesky', fit_intercept=True).fit(x, y).predict(x)
+        )
+        print("[sklearn_result]---------------------------------------------")
+        print(sklearn_result[:10])
+
+        # mark these data to be protected in SPU
+        x1, x2, y = emulator.seal(x1, x2, y)
+
+        # run
+        result = emulator.run(proc)(x1, x2, y)
+        print("[emul_result]------------------------------------------------")
+        print(result[:10])
+
+        # absolute_error
+        print("[absolute_error]---------------------------------------------")
+        print(jnp.round(jnp.abs(result - sklearn_result)[:20], 5))
+    finally:
+        emulator.down()
+
+
+if __name__ == "__main__":
+    emul_Ridge(emulation.Mode.MULTIPROCESS)
diff --git a/sml/linear_model/ridge.py b/sml/linear_model/ridge.py
new file mode 100644
index 00000000..4e4cb560
--- /dev/null
+++ b/sml/linear_model/ridge.py
@@ -0,0 +1,140 @@
+# Copyright 2023 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from enum import Enum
+
+import jax.numpy as jnp
+import jax.scipy as jsci
+
+
+class Solver(Enum):
+    SVD = 'svd'  # not supported
+    CHOLESKY = 'cholesky'
+
+
+class Ridge:
+    """Linear least squares with l2 regularization.
+
+    Minimizes the objective function::
+
+    ||y - Xw||^2_2 + alpha * ||w||^2_2
+
+    This model solves a regression model where the loss function is
+    the linear least squares function and regularization is given by
+    the l2-norm. Also known as Ridge Regression or Tikhonov regularization.
+
+    Parameters
+    ----------
+    alpha : {float}, default=1.0
+        Constant that multiplies the L2 term, controlling regularization
+        strength. `alpha` must be a non-negative float i.e. in `[0, inf)`.
+
+    fit_bias : bool, default=True
+        Whether to fit the bias for this model. If set
+        to false, no bias will be used in calculations
+
+    solver : {'svd', 'cholesky'}, default='cholesky'
+        Solver to use in the computational routines:
+
+        - 'svd' uses a Singular Value Decomposition of X to compute the Ridge
+          coefficients.
+
+        - 'cholesky' uses the standard jax.scipy.linalg.solve function to
+          obtain a closed-form solution via a Cholesky decomposition of
+          dot(X.T, X)
+    """
+
+    def __init__(self, alpha=1.0, fit_bias=True, solver="cholesky") -> None:
+        self.alpha = alpha
+        self.solver = solver
+        self.fit_bias = fit_bias
+
+    def fit(self, x, y):
+        """Fit Ridge regression model.
+
+        Parameters
+        ----------
+        X : {array-like}, shape (n_samples, n_features)
+            Training data.
+
+        y : ndarray of shape (n_samples,)
+            Target values.
+
+        Returns
+        -------
+        self : object
+            Returns an instance of self.
+        """
+        if y.ndim == 1:
+            y = y.reshape(-1, 1)
+        alpha = float(self.alpha)
+
+        x, y, x_offset, y_offset = self.preprocess_data(x, y)
+
+        if self.solver == Solver.CHOLESKY.value:
+            self.coef = _solve_cholesky(x, y, alpha)
+        self.coef = self.coef.ravel()
+
+        self.set_bias(x_offset, y_offset)
+
+        return self
+
+    def predict(self, x):
+        """
+        Predict using the linear model.
+
+        Parameters
+        ----------
+        X : {array-like}, shape (n_samples, n_features)
+            Input data for prediction.
+
+        Returns
+        -------
+        C : array, shape (n_samples,)
+            Returns predicted values.
+        """
+        a = x
+        b = self.coef.T
+        ret = jnp.dot(a, b) + self.bias
+        return ret
+
+    def preprocess_data(self, x, y):
+        # Center and scale data.
+        if self.fit_bias:
+            x_offset = jnp.average(x, axis=0)
+            x -= x_offset
+            y_offset = jnp.average(y, axis=0)
+            y -= y_offset
+        else:
+            x_offset = None
+            y_offset = None
+        return x, y, x_offset, y_offset
+
+    def set_bias(self, x_offset, y_offset):
+        if self.fit_bias:
+            self.bias = y_offset - jnp.dot(x_offset, self.coef.T)
+        else:
+            self.bias = 0.0
+
+
+def _solve_cholesky(x, y, alpha):
+    # w = inv(X^t X + alpha*Id) * X.T y
+    n_features = x.shape[1]
+
+    A = jnp.dot(x.T, x)
+    Xy = jnp.dot(x.T, y)
+
+    A += jnp.diag(jnp.ones(n_features) * alpha)
+
+    coefs = jsci.linalg.solve(A, Xy, assume_a="pos", overwrite_a=True).T
+    return coefs


Test patch:
diff --git a/sml/linear_model/tests/BUILD.bazel b/sml/linear_model/tests/BUILD.bazel
index 4a4ad38e..2b9e7c10 100644
--- a/sml/linear_model/tests/BUILD.bazel
+++ b/sml/linear_model/tests/BUILD.bazel
@@ -39,3 +39,14 @@ py_test(
         "//spu/utils:simulation",
     ],
 )
+
+py_test(
+    name = "ridge_test",
+    srcs = ["ridge_test.py"],
+    deps = [
+        "//examples/python/utils:dataset_utils",  # FIXME: remove examples dependency
+        "//sml/linear_model:ridge",
+        "//spu:init",
+        "//spu/utils:simulation",
+    ],
+)
diff --git a/sml/linear_model/tests/ridge_test.py b/sml/linear_model/tests/ridge_test.py
new file mode 100644
index 00000000..99c468a7
--- /dev/null
+++ b/sml/linear_model/tests/ridge_test.py
@@ -0,0 +1,67 @@
+# Copyright 2023 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import unittest
+
+import jax.numpy as jnp
+from sklearn.linear_model import Ridge as skRidge
+
+import examples.python.utils.dataset_utils as dsutil
+import spu.spu_pb2 as spu_pb2  # type: ignore
+import spu.utils.simulation as spsim
+from sml.linear_model.ridge import Ridge
+
+
+class UnitTests(unittest.TestCase):
+    def test_ridge(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
+        )
+
+        def proc(x1, x2, y):
+            model = Ridge(alpha=1.0, solver="cholesky")
+
+            x = jnp.concatenate((x1, x2), axis=1)
+            y = y.reshape((y.shape[0], 1))
+
+            result = model.fit(x, y).predict(x)
+            return result
+
+        dataset_config = {
+            "use_mock_data": False,
+            "problem_type": "regression",
+            "builtin_dataset_name": "diabetes",
+            "left_slice_feature_ratio": 0.5,
+        }
+
+        x1, x2, y = dsutil.load_dataset_by_config(dataset_config)
+
+        # sklearn test
+        x = jnp.concatenate((x1, x2), axis=1)
+        sklearn_result = (
+            skRidge(alpha=1, solver='cholesky', fit_intercept=True).fit(x, y).predict(x)
+        )
+        print("[sklearn_result]---------------------------------------------")
+        print(sklearn_result[:10])
+
+        result = spsim.sim_jax(sim, proc)(x1, x2, y)
+        print("[spsim_result]-----------------------------------------------")
+        print(result[:10])
+
+        # absolute_error
+        print("[absolute_error]---------------------------------------------")
+        print(jnp.round(jnp.abs(result - sklearn_result)[:20], 5))
+
+
+if __name__ == "__main__":
+    unittest.main()