Repository: secretflow/spu (language: python)

Issue:
add sml/kmeans
add sml/kmeans


Discussion / hints:
****CLA Assistant Lite bot**** All contributors have signed the CLA  ✍️ ✅
I have read the CLA Document and I hereby sign the CLA
> I have read the CLA Document and I hereby sign the CLA

> I have read the CLA Document and I hereby sign the CLA

Thanks for contributing KMeans algorithm. The correctness of this algo seems ok to me, here are some optimization tips and     recommendations:

1. you can use the feature of broadcast in `jax.numpy` to avoid for-loop in `fit` and `predict` method 
2. you can avoid some redundant `.T` op in `fit` method.
3. I find you use something like float compare trick(< eps) rather equal op directly. Maybe it's a good idea for `ABY3` protocol, because `mul` is cheap,  but if you use `CHEETAH`, maybe the total cost of your implementation will be larger?(I think you can do some more experiments under more config, e.g. `CHEETAH`, `FM128`, etc.)
4. maybe you can add some doc and comment about the details of your implementation?

you can refer to my #todo note for my recommendations.
```python
import jax
import jax.numpy as jnp

class KMEANS:
    def __init__(self, k, data_len, max_iter=300, tol=1e-4):
        # todo: some doc about your hyper-param
        self.k = k
        self.max_iter = max_iter
        self.tol = tol
        self.init_params = jax.random.randint(jax.random.PRNGKey(1),shape=[self.k],minval=0,maxval=data_len)
        self._centers = jnp.zeros(())

    def fit(self, x):
        """Fit KMEANS.
        # todo: some doc about you implementation

        Parameters
        ----------
        x : {array-like}, shape (n_samples, n_features)
            Input data.

        Returns
        -------
        self : object
            Returns an instance of self.
        """

        centers = jnp.array([x[i] for i in self.init_params]) # k,p
        for i in range(self.max_iter):
            # print(f"iter: {i+1:3d}/{self.max_iter}\r")
            # C = []
            # for j in range(self.k):
            #     C.append(x - centers[j])
            # C = jnp.array(C)
            # todo: use broadcast
            C = x.reshape((1, x.shape[0], x.shape[1])) - centers.reshape((centers.shape[0], 1, centers.shape[1]))

            C = jnp.argmin(jnp.sum(jnp.square(C), axis=2), axis=0)

            S = jnp.tile(C,(self.k,1))

            # todo: simple simplify
            # ks = jnp.array([i for i in range(self.k)])
            ks = jnp.arange(self.k)
            
            # todo: simple simplify
            # equals_raw = (S.T - ks.T).T 
            equals_raw = (S.T - ks).T  

            # todo: maybe more experiments about two implementation ?
            equals = jnp.less(jnp.square(equals_raw), self.tol) 
            # equals = jnp.equal(equals_raw, 0)


            # centers_raw = []
            # for j in range(self.k):
            #     centers_raw.append(jnp.multiply(x.T,equals[j]).T)  
            #     # centers_raw.append(jnp.multiply(x.T,equals[j].T).T) 
            # centers_raw = jnp.array(centers_raw)
            # todo: use broadcast
            centers_raw = x.reshape((1, x.shape[0], x.shape[1])) * equals.reshape((equals.shape[0],equals.shape[1],1))

            equals_sum = jnp.sum(equals,axis=1) 
            centers_sum = jnp.sum(centers_raw, axis=1)
            # todo: simple simplify
            centers = jnp.divide(centers_sum.T, equals_sum).T 
            # centers = jnp.divide(centers_sum.T, equals_sum.T).T

        self._centers = centers
        return self

    def predict(self, x):
        """result estimates.

        Parameters
        ----------
        x : {array-like}, shape (n_samples, n_features)
            Input data for prediction.

        Returns
        -------
        ndarray of shape (n_samples)
            Returns the result of the sample for each class in the model.
        """
        centers = self._centers
        # todo: can use broadcast, leave this for you
        y = []
        for j in range(self.k):
            y.append(x - centers[j])
        y = jnp.array(y)
        y = jnp.argmin(jnp.sum(jnp.square(y), axis=2), axis=0)
        return y 
```



> I have read the CLA Document and I hereby sign the CLA

Besides, the param `k` can change to `n_clusters` to align the param in sklearn.
I make revisions based on the suggestions. The comparison with eps  is also used to avoid the impact of possible  small errors in some MPC schemes, and a strict equality comparisons may get things wrong in some case, I think.
> I make revisions based on the suggestions. The comparison with eps is also used to avoid the impact of possible small errors in some MPC schemes, and a strict equality comparisons may get things wrong in some case, I think.

Thanks for your revisions, the algorithm part seems good to me. 
However here still exists some minor errosr: for emul and test file, your params do not change to the newest version, please **must** make sure your program can runs correctly before doing commits.

Besides, I **highly recommend** that you compare your output with the plaintext sklearn  output just like https://github.com/secretflow/spu/pull/240 do in both emul and test file.(You can use `sklearn.datasets.make_blobs` to generate very distinct datasets).

Finally, you say "The comparison with eps is also used to avoid the impact of possible small errors in some MPC schemes", but I think it's not a problem at all. You can find the dtype of  the matrix `S` and `ks` are both int, so SPU can do it lossless. The integer equal will call `a2b` which costs a lot, but for fxp less, `msb` is sufficient(of course, one more `mul`).So I think for `aby3` protocol, comparison with eps will do better, but if you runs in `cheetah`, I am not sure which way can do better(for `mul` is much more expensive for `cheetah`, maybe you can do some more experiments about this). 


Thanks for your suggestions and I add `sklearn.datasets.make_blobs`.  
I do experiments on the using of  `equal` and `compare with eps`. The results (in seconds) are shown in the table.
It seems the results of the two methods are similar for smaller sample sets, . When the data volume is large, it appears that the advantage of `compare with eps` in ABY3 is smaller than that of `equal ` in CHEETAH, so maybe using `equal` is better.
<html xmlns:v="urn:schemas-microsoft-com:vml"
xmlns:o="urn:schemas-microsoft-com:office:office"
xmlns:x="urn:schemas-microsoft-com:office:excel"
xmlns="http://www.w3.org/TR/REC-html40">

<head>

<meta name=ProgId content=Excel.Sheet>
<meta name=Generator content="Microsoft Excel 15">

</head>

<body link="#0563C1" vlink="#954F72">



  | aby3 |   | cheetah |  
-- | -- | -- | -- | --
n_samples | equal | compare with eps| equal | compare with eps
100 | 0.12 | 0.12 | 0.12 | 0.12
1000 | 0.44 | 0.42 | 0.44 | 0.41
10000 | 2.52 | 2.45 | 2.54 | 2.47
100000 | 27.81 | 27.93 | 26.01 | 28.05



</body>

</html>
















































































> Thanks for your suggestions and I add `sklearn.datasets.make_blobs`. I do experiments on the using of `equal` and `compare with eps`. The results (in seconds) are shown in the table. It seems the results of the two methods are similar for smaller sample sets, . When the data volume is large, it appears that the advantage of `compare with eps` in ABY3 is smaller than that of `equal ` in CHEETAH, so maybe using `equal` is better.
> 
>  	aby3	 	cheetah	 
> n_samples	equal	compare with eps	equal	compare with eps
> 100	0.12	0.12	0.12	0.12
> 1000	0.44	0.42	0.44	0.41
> 10000	2.52	2.45	2.54	2.47
> 100000	27.81	27.93	26.01	28.05

Thanks for doing elaborated experiments! All other things are good to me! 

But I am confused how you run your experiments?( In my thoughts, cheetah is much much slower than aby3, but the output seems similar. )
> Thanks for your suggestions and I add `sklearn.datasets.make_blobs`. I do experiments on the using of `equal` and `compare with eps`. The results (in seconds) are shown in the table. It seems the results of the two methods are similar for smaller sample sets, . When the data volume is large, it appears that the advantage of `compare with eps` in ABY3 is smaller than that of `equal ` in CHEETAH, so maybe using `equal` is better.
> 
>  	aby3	 	cheetah	 
> n_samples	equal	compare with eps	equal	compare with eps
> 100	0.12	0.12	0.12	0.12
> 1000	0.44	0.42	0.44	0.41
> 10000	2.52	2.45	2.54	2.47
> 100000	27.81	27.93	26.01	28.05

您好，我最近发现之前emul部分的代码有点问题，即在运行emul程序之前，是需要先将数据load到PYU上的，但之前emulation并没有提供相应的api；所以如果使用了自定义数据，实际上emul部分都是在明文下运行的，具体可以参考：（其实主要就是调用一下`emulator.seal`哈）
https://github.com/secretflow/spu/blob/main/sml/linear_model/simple_sgd_emul.py

我理解这个修改以后，cheetah和aby3就会有比较明显的性能区别了～

Thanks
我修改了代码，这是新的结果，看起来两种方法仍然是接近的，并且aby3比cheetah快很多，使用`equal`是没问题的。
<html xmlns:v="urn:schemas-microsoft-com:vml"
xmlns:o="urn:schemas-microsoft-com:office:office"
xmlns:x="urn:schemas-microsoft-com:office:excel"
xmlns="http://www.w3.org/TR/REC-html40">

<head>

<meta name=ProgId content=Excel.Sheet>
<meta name=Generator content="Microsoft Excel 15">

</head>

<body link="#0563C1" vlink="#954F72">



  | aby3 |   | cheetah |  
-- | -- | -- | -- | --
n_samples | equal | compare with eps| equal | compare with eps
100 | 1.10 | 1.10 | 11.26 | 11.30
200 | 1.60 | 1.48 | 15.34 | 15.81
1000 | 2.02 | 2.01 | 56.49 | 56.65
2000 | 2.94 | 2.93 | 87.91 | 88.89
</body>

</html>


























































































Fix patch:
diff --git a/sml/kmeans/BUILD.bazel b/sml/kmeans/BUILD.bazel
new file mode 100644
index 00000000..b6c92646
--- /dev/null
+++ b/sml/kmeans/BUILD.bazel
@@ -0,0 +1,37 @@
+load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "kmeans",
+    srcs = ["kmeans.py"],
+    deps = [
+        "//sml/utils:fxp_approx",
+    ],
+)
+
+py_binary(
+    name = "kmeans_emul",
+    srcs = ["kmeans_emul.py"],
+    deps = [
+        ":kmeans",
+        "//examples/python/utils:dataset_utils",  # FIXME: remove examples dependency
+        "//sml/utils:emulation"
+    ],
+)
+
+
+
+py_test(
+    name = "kmeans_test",
+    srcs = ["kmeans_test.py"],
+    data = [
+        "//examples/python/conf",  # FIXME: remove examples dependency
+    ],
+    deps = [
+        ":kmeans",
+        "//examples/python/utils:dataset_utils",  # FIXME: remove examples dependency
+        "//spu:init",
+        "//spu/utils:simulation",
+    ],
+)
diff --git a/sml/kmeans/kmeans.py b/sml/kmeans/kmeans.py
new file mode 100644
index 00000000..0145e25e
--- /dev/null
+++ b/sml/kmeans/kmeans.py
@@ -0,0 +1,86 @@
+import jax
+import jax.numpy as jnp
+
+class KMEANS:
+    """
+    Parameters
+    ----------
+    n_clusters : int
+        The number of clusters to form as well as the number of
+        centroids to generate.
+
+    n_samples : int
+        The number of samples. 
+
+    max_iter : int, default=300
+        Maximum number of iterations of the k-means algorithm for a
+        single run.   
+    
+    tol : float, default=1e-4
+        Acceptable error to consider the two to be equal. 
+    """
+    def __init__(self, n_clusters, n_samples, max_iter=300, tol=1e-4):
+        self.n_clusters = n_clusters
+        self.max_iter = max_iter
+        self.tol = tol
+        self.init_params = jax.random.randint(jax.random.PRNGKey(1),shape=[self.n_clusters],minval=0,maxval=n_samples)
+        self._centers = jnp.zeros(())
+    
+    def fit(self, x):
+        """Fit KMEANS.
+
+        Firstly, randomly select the initial centers. Then calculate the distance between each sample and each center, 
+        and assign each sample to the nearest center. Use an `aligned_array` to indicate the samples in a cluster, 
+        where unrelated samples will be set to 0. Once all samples are assigned, the center of each cluster will 
+        be updated to the average. The average could be got by `sum(data * aligned_array) / sum(aligned_array)`.
+        Different clusters could use broadcast for better performance.
+
+        Parameters
+        ----------
+        x : {array-like}, shape (n_samples, n_features)
+            Input data.
+
+        Returns
+        -------
+        self : object
+            Returns an instance of self.
+        """
+        
+        centers = jnp.array([x[i] for i in self.init_params])
+        for _ in range(self.max_iter):
+            C = x.reshape((1, x.shape[0], x.shape[1])) - centers.reshape((centers.shape[0], 1, centers.shape[1]))
+            C = jnp.argmin(jnp.sum(jnp.square(C), axis=2), axis=0)
+
+            S = jnp.tile(C,(self.n_clusters,1))
+            ks = jnp.arange(self.n_clusters)
+            aligned_array_raw = (S.T - ks).T
+            aligned_array = jnp.equal(aligned_array_raw, 0)
+
+            centers_raw = x.reshape((1, x.shape[0], x.shape[1])) * aligned_array.reshape((aligned_array.shape[0],aligned_array.shape[1],1))
+            equals_sum = jnp.sum(aligned_array,axis=1)
+            centers_sum = jnp.sum(centers_raw, axis=1)
+            centers = jnp.divide(centers_sum.T, equals_sum).T 
+
+        self._centers = centers
+        return self
+    
+    def predict(self, x):
+        """Result estimates.
+
+        Calculate the distance between each sample and each center, 
+        and assign each sample to the nearest center.
+
+        Parameters
+        ----------
+        x : {array-like}, shape (n_samples, n_features)
+            Input data for prediction.
+
+        Returns
+        -------
+        ndarray of shape (n_samples)
+            Returns the result of the sample for each class in the model.
+        """
+        centers = self._centers
+        y = x.reshape((1, x.shape[0], x.shape[1])) - centers.reshape((centers.shape[0], 1, centers.shape[1]))
+        y = jnp.argmin(jnp.sum(jnp.square(y), axis=2), axis=0)
+        return y
\ No newline at end of file
diff --git a/sml/kmeans/kmeans_emul.py b/sml/kmeans/kmeans_emul.py
new file mode 100644
index 00000000..5b3289f5
--- /dev/null
+++ b/sml/kmeans/kmeans_emul.py
@@ -0,0 +1,54 @@
+import jax.numpy as jnp
+
+# from sklearn.metrics import roc_auc_score, explained_variance_score
+import sml.utils.emulation as emulation
+
+from sml.kmeans.kmeans import KMEANS
+from sklearn.datasets import make_blobs
+
+# TODO: design the enumation framework, just like py.unittest
+# all emulation action should begin with `emul_` (for reflection)
+def emul_KMEANS(mode: emulation.Mode.MULTIPROCESS):
+    def proc(x1, x2):
+        x = jnp.concatenate((x1, x2), axis=1)
+        model = KMEANS(
+            n_clusters=2,
+            n_samples=x.shape[0],
+            max_iter=10
+        )
+
+        return model.fit(x).predict(x)
+    
+    def load_data():
+        n_samples = 1000
+        n_features = 100
+        X, _ = make_blobs(n_samples=n_samples,n_features=n_features,centers=2)
+        split_index = n_features//2
+        return X[:, :split_index], X[:, split_index:]
+    
+    try:
+        # bandwidth and latency only work for docker mode
+        emulator = emulation.Emulator(
+            "examples/python/conf/3pc.json", mode, bandwidth=300, latency=20
+        )
+        emulator.up()
+
+        # load mock data
+        x1, x2  = load_data()
+        X = jnp.concatenate((x1, x2), axis=1)
+
+        # mark these data to be protected in SPU
+        x1, x2 = emulator.seal(x1, x2)
+        result = emulator.run(proc)(x1, x2)
+        print("result\n",result)
+
+        # Compare with sklearn
+        from sklearn.cluster import KMeans
+        model = KMeans(n_clusters=2)
+        print("sklearn:\n",model.fit(X).predict(X))
+    finally:
+        emulator.down()
+
+
+if __name__ == "__main__":
+    emul_KMEANS(emulation.Mode.MULTIPROCESS)


Test patch:
diff --git a/sml/kmeans/kmeans_test.py b/sml/kmeans/kmeans_test.py
new file mode 100644
index 00000000..b455615e
--- /dev/null
+++ b/sml/kmeans/kmeans_test.py
@@ -0,0 +1,50 @@
+import unittest
+import json
+import jax.numpy as jnp
+
+# from sklearn.metrics import roc_auc_score, explained_variance_score
+import spu.utils.simulation as spsim
+import spu.spu_pb2 as spu_pb2  # type: ignore
+
+# TODO: unify this.
+import examples.python.utils.dataset_utils as dsutil
+
+from sml.kmeans.kmeans import KMEANS
+from sklearn.datasets import make_blobs
+
+class UnitTests(unittest.TestCase):
+    def test_kmeans(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
+        )
+
+        def proc(x1, x2):
+            x = jnp.concatenate((x1, x2), axis=1)
+            model = KMEANS(
+                n_clusters=2,
+                n_samples=x.shape[0],
+                max_iter=10
+            )
+
+            return model.fit(x).predict(x)
+        
+        def load_data():
+            n_samples = 1000
+            n_features = 100
+            X, _ = make_blobs(n_samples=n_samples,n_features=n_features,centers=2)
+            split_index = n_features//2
+            return X[:, :split_index], X[:, split_index:]
+
+        x1, x2 = load_data()
+        X = jnp.concatenate((x1, x2), axis=1)
+        result = spsim.sim_jax(sim, proc)(x1, x2)
+        print("result\n",result)
+
+        # Compare with sklearn
+        from sklearn.cluster import KMeans
+        model = KMeans(n_clusters=2)
+        print("sklearn:\n",model.fit(X).predict(X))
+
+
+if __name__ == "__main__":
+    unittest.main()