From 8ac253b9fe112bfe64d467c07db9dc9b9ed74e52 Mon Sep 17 00:00:00 2001
From: Andrei Kislitsyn <kisandy@yandex.ru>
Date: Thu, 6 May 2021 14:09:47 +0300
Subject: [PATCH] cov + docs

---
 .../tensors/api/AnalyticTensorAlgebra.kt      | 10 +++++++
 .../algebras/DoubleAnalyticTensorAlgebra.kt   | 24 ++++++++++++++++-
 .../core/algebras/DoubleTensorAlgebra.kt      | 27 +++++++++++++------
 3 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/api/AnalyticTensorAlgebra.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/api/AnalyticTensorAlgebra.kt
index aa5678b31..69e88c28f 100644
--- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/api/AnalyticTensorAlgebra.kt
+++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/api/AnalyticTensorAlgebra.kt
@@ -68,6 +68,16 @@ public interface AnalyticTensorAlgebra<T> :
      */
     public fun Tensor<T>.variance(dim: Int, keepDim: Boolean): Tensor<T>
 
+    /**
+     * Returns the covariance matrix M of given vectors.
+     *
+     * M[i, j] contains covariance of i-th and j-th given vectors
+     *
+     * @param tensors the [List] of 1-dimensional tensors with same shape
+     * @return the covariance matrix
+     */
+    public fun cov(tensors: List<Tensor<T>>): Tensor<T>
+
     //For information: https://pytorch.org/docs/stable/generated/torch.exp.html
     public fun Tensor<T>.exp(): Tensor<T>
 
diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleAnalyticTensorAlgebra.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleAnalyticTensorAlgebra.kt
index 4a1f360e3..23a2fa282 100644
--- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleAnalyticTensorAlgebra.kt
+++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleAnalyticTensorAlgebra.kt
@@ -57,6 +57,28 @@ public object DoubleAnalyticTensorAlgebra :
         keepDim
     )
 
+    private fun cov(x: DoubleTensor, y:DoubleTensor): Double{
+        val n = x.shape[0]
+        return ((x - x.mean()) * (y - y.mean())).mean() * n / (n - 1)
+    }
+
+    override fun cov(tensors: List<Tensor<Double>>): DoubleTensor {
+        check(tensors.isNotEmpty()) { "List must have at least 1 element" }
+        val n = tensors.size
+        val m = tensors[0].shape[0]
+        check(tensors.all { it.shape contentEquals intArrayOf(m) }) { "Tensors must have same shapes" }
+        val resTensor = DoubleTensor(
+            intArrayOf(n, n),
+            DoubleArray(n * n) {0.0}
+        )
+        for (i in 0 until n){
+            for (j in 0 until n){
+                resTensor[intArrayOf(i, j)] = cov(tensors[i].tensor, tensors[j].tensor)
+            }
+        }
+        return resTensor
+    }
+
     override fun Tensor<Double>.exp(): DoubleTensor = tensor.map(::exp)
 
     override fun Tensor<Double>.ln(): DoubleTensor = tensor.map(::ln)
@@ -91,4 +113,4 @@ public object DoubleAnalyticTensorAlgebra :
 
     override fun Tensor<Double>.floor(): DoubleTensor = tensor.map(::floor)
 
-}
\ No newline at end of file
+}
diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt
index d220bdd9a..74ef63ab7 100644
--- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt
+++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt
@@ -401,19 +401,31 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra<Double> {
     public fun Tensor<Double>.randomNormalLike(seed: Long = 0): DoubleTensor =
         DoubleTensor(tensor.shape, getRandomNormals(tensor.shape.reduce(Int::times), seed))
 
-    // stack tensors by axis 0
-    public fun stack(tensors: List<DoubleTensor>): DoubleTensor {
-        val shape = tensors.firstOrNull()?.shape
-        check(shape != null) { "Collection must have at least 1 element" }
-        check(tensors.all { it.shape contentEquals shape }) { "Stacking tensors must have same shapes" }
+    /**
+     * Concatenates a sequence of tensors along a new dimension.
+     *
+     * @param tensors the [List] of tensors with same shapes to concatenate
+     * @param dim the dimension to insert
+     * @return tensor with concatenation result
+     */
+    public fun stack(tensors: List<Tensor<Double>>, dim: Int = 0): DoubleTensor {
+        check(dim == 0) { "Stack by non-zero dimension not implemented yet" }
+        check(tensors.isNotEmpty()) { "List must have at least 1 element" }
+        val shape = tensors[0].shape
+        check(tensors.all { it.shape contentEquals shape }) { "Tensors must have same shapes" }
         val resShape = intArrayOf(tensors.size) + shape
         val resBuffer = tensors.flatMap {
-            it.tensor.mutableBuffer.array().drop(it.bufferStart).take(it.numElements)
+            it.tensor.mutableBuffer.array().drop(it.tensor.bufferStart).take(it.tensor.numElements)
         }.toDoubleArray()
         return DoubleTensor(resShape, resBuffer, 0)
     }
 
-    // build tensor from this rows by given indices
+    /**
+     * Build tensor from rows of input tensor
+     *
+     * @param indices the [IntArray] of 1-dimensional indices
+     * @return tensor with rows corresponding to rows by [indices]
+     */
     public fun Tensor<Double>.rowsByIndices(indices: IntArray): DoubleTensor {
         return stack(indices.map { this[it] })
     }
@@ -450,7 +462,6 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra<Double> {
     override fun Tensor<Double>.sum(dim: Int, keepDim: Boolean): DoubleTensor =
         foldDim({ x -> x.sum() }, dim, keepDim)
 
-
     override fun Tensor<Double>.min(): Double = this.fold { it.minOrNull()!! }
 
     override fun Tensor<Double>.min(dim: Int, keepDim: Boolean): DoubleTensor =