diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt index 418cf16b9..7fa6e885d 100644 --- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt +++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt @@ -10,6 +10,7 @@ package space.kscience.kmath.tensors.core import space.kscience.kmath.misc.PerformancePitfall import space.kscience.kmath.nd.MutableStructure2D +import space.kscience.kmath.nd.Structure2D import space.kscience.kmath.nd.StructureND import space.kscience.kmath.nd.as1D import space.kscience.kmath.nd.as2D @@ -885,7 +886,7 @@ public open class DoubleTensorAlgebra : return Triple(uTensor.transpose(), sTensor, vTensor.transpose()) } - override fun StructureND.symEig(): Pair = symEig(epsilon = 1e-15) + override fun StructureND.symEig(): Pair = symEigJacobi(epsilon = 1e-10) /** * Returns eigenvalues and eigenvectors of a real symmetric matrix input or a batch of real symmetric matrices, @@ -922,96 +923,99 @@ public open class DoubleTensorAlgebra : return eig to v } - // TODO - // 1. Cyclic method - // 2. Sort eigenvalues - public fun StructureND.symEig(epsilon: Double): Pair { + public fun StructureND.symEigJacobi(epsilon: Double): Pair { checkSymmetric(tensor, epsilon) - val ii = tensor.minusIndex(-2) - val jj = tensor.minusIndex(-1) - val n = tensor.numElements - val size = this.dimension - val commonShape = this.shape.sliceArray(0 until size - 2) + intArrayOf(1, 1) + val s = zeros(this.shape) + val eigenvalues = zeros(this.shape.sliceArray(0 until size - 1)) - var d = this.copy() - var s = diagonalEmbedding(ones(this.shape.sliceArray(0 until size - 1))) + var eigenvalueStart = 0 + var eigenvectorStart = 0 + for (matrix in tensor.matrixSequence()) { + matrix.as2D().jacobiHelper(eigenvalues, s, eigenvalueStart, eigenvectorStart, epsilon) + eigenvalueStart += this.shape.last() + eigenvectorStart += this.shape.last() * this.shape.last() + } + // TODO sort eigenvalues + return eigenvalues to s + } + + private fun MutableStructure2D.jacobiHelper( + eigenvalues: DoubleTensor, + eigenvectors: DoubleTensor, + eigenvalueStart: Int, + eigenvectorStart: Int, + epsilon: Double + ) { + var d = this + var s = eye(this.shape[0]) + + // TODO implement cyclic method do { // 1. Find max element by abs value that is not on diagonal - val buffer = MutableBuffer.boxing(commonShape.reduce(Int::times)) { Triple(0.0, 0, 0) } - val maxOffDiagonalElements = BufferedTensor(commonShape, buffer, 0) - for (offset in 0 until n) { - val multiIndex = d.linearStructure.index(offset) - if (multiIndex[ii] != multiIndex[jj]) { - val value = d.mutableBuffer.array()[offset] - - val commonIndex = multiIndex.sliceArray(0 until size - 2) + intArrayOf(0, 0) - if (abs(value) > maxOffDiagonalElements[commonIndex].first) { - maxOffDiagonalElements[commonIndex] = Triple(abs(value), multiIndex[ii], multiIndex[jj]) + var maxOffDiagonalElement = 0.0 + var maxElementIndex = Pair(0, 0) + for (i in 0 until this.rowNum) { + for (j in 0 until this.colNum) { + if (i == j) continue + if (abs(d[i, j]) > maxOffDiagonalElement) { + maxOffDiagonalElement = abs(d[i, j]) + maxElementIndex = i to j } } } // 2. Evaluate "rotation" angle - val angles = zeros(commonShape) - for (offset in 0 until maxOffDiagonalElements.numElements) { - val (_, i, j) = maxOffDiagonalElements.mutableBuffer[offset] - val multiIndex = maxOffDiagonalElements.linearStructure.index(offset) + val dIJ = d[maxElementIndex.first, maxElementIndex.second] + val dII = d[maxElementIndex.first, maxElementIndex.first] + val dJJ = d[maxElementIndex.second, maxElementIndex.second] - val dIJ = d.mutableBuffer[d.linearStructure.offset(multiIndex.also { it[ii] = i; it[jj] = j })] - val dII = d.mutableBuffer[d.linearStructure.offset(multiIndex.also { it[ii] = i; it[jj] = i })] - val dJJ = d.mutableBuffer[d.linearStructure.offset(multiIndex.also { it[ii] = j; it[jj] = j })] - - angles.mutableBuffer.array()[offset] = if (dII == dJJ) { - if (dIJ > 0) PI / 4 else -PI / 4 - } else { - 0.5 * atan(2 * dIJ / (dJJ - dII)) - } + val angle = if (dII == dJJ) { + if (dIJ > 0) PI / 4 else -PI / 4 + } else { + 0.5 * atan(2 * dIJ / (dJJ - dII)) } // 3. Build rotation tensor `s1` - val s1 = diagonalEmbedding(ones(this.shape.sliceArray(0 until size - 1))) - for (offset in 0 until n) { - val multiIndex = d.linearStructure.index(offset) - - val commonIndex = multiIndex.sliceArray(0 until size - 2) + intArrayOf(0, 0) - val (_, i, j) = maxOffDiagonalElements[commonIndex] - val angleValue = angles[commonIndex] - s1.mutableBuffer.array()[offset] = when { - multiIndex[ii] == i && multiIndex[jj] == i || multiIndex[ii] == j && multiIndex[jj] == j -> cos(angleValue) - multiIndex[ii] == i && multiIndex[jj] == j -> sin(angleValue) - multiIndex[ii] == j && multiIndex[jj] == i -> -sin(angleValue) - else -> s1.mutableBuffer.array()[offset] + val s1 = eye(this.rowNum) + for (i in 0 until this.rowNum) { + for (j in 0 until this.colNum) { + s1.mutableBuffer.array()[i * this.rowNum + j] = when { + maxElementIndex.first == i && maxElementIndex.first == j -> cos(angle) + maxElementIndex.second == i && maxElementIndex.second == j -> cos(angle) + maxElementIndex.first == i && maxElementIndex.second == j -> sin(angle) + maxElementIndex.first == j && maxElementIndex.second == i -> -sin(angle) + else -> s1.mutableBuffer.array()[i * this.rowNum + j] + } } } // 4. Evaluate new tensor - d = (s1.transpose() dot d) dot s1 + d = ((s1.transpose() dot d) dot s1).as2D() s = s dot s1 if (d.isDiagonal(epsilon)) break } while(true) - val eigenvalues = zeros(d.shape.sliceArray(0 until size - 1)) - for (offset in 0 until n) { - val multiIndex = d.linearStructure.index(offset) - if (multiIndex[ii] == multiIndex[jj]) { - eigenvalues[multiIndex.sliceArray(0 until size - 1)] = d.mutableBuffer.array()[offset] + // 5. Copy result + for (i in 0 until this.rowNum) { + for (j in 0 until this.colNum) { + eigenvectors.mutableBuffer.array()[eigenvectorStart + i * this.rowNum + j] = s.mutableBuffer.array()[i * this.rowNum + j] } } - return eigenvalues to s + for (i in 0 until this.rowNum) { + eigenvalues.mutableBuffer.array()[eigenvalueStart + i] = d[i, i] + } } - public fun StructureND.isDiagonal(epsilon: Double = 1e-9): Boolean { - val ii = tensor.minusIndex(-2) - val jj = tensor.minusIndex(-1) - - for (offset in 0 until tensor.numElements) { - val multiIndex = tensor.linearStructure.index(offset) - if (multiIndex[ii] != multiIndex[jj] && abs(tensor.mutableBuffer.array()[offset]) > epsilon) { - return false + public fun Structure2D.isDiagonal(epsilon: Double = 1e-9): Boolean { + for (i in 0 until this.rowNum) { + for (j in 0 until this.colNum) { + if (i != j && abs(this[i, j]) > epsilon) { + return false + } } }