Slight adjustment to tensor internals

This commit is contained in:
Alexander Nozik 2021-06-02 21:10:05 +03:00
parent c8621ee5b7
commit 7f32348e7a
2 changed files with 84 additions and 63 deletions

View File

@ -5,8 +5,10 @@
package space.kscience.kmath.tensors.core package space.kscience.kmath.tensors.core
import space.kscience.kmath.nd.MutableStructure2D
import space.kscience.kmath.nd.as1D import space.kscience.kmath.nd.as1D
import space.kscience.kmath.nd.as2D import space.kscience.kmath.nd.as2D
import space.kscience.kmath.structures.indices
import space.kscience.kmath.tensors.api.AnalyticTensorAlgebra import space.kscience.kmath.tensors.api.AnalyticTensorAlgebra
import space.kscience.kmath.tensors.api.LinearOpsTensorAlgebra import space.kscience.kmath.tensors.api.LinearOpsTensorAlgebra
import space.kscience.kmath.tensors.api.Tensor import space.kscience.kmath.tensors.api.Tensor
@ -813,28 +815,32 @@ public open class DoubleTensorAlgebra :
val sTensor = zeros(commonShape + intArrayOf(min(n, m))) val sTensor = zeros(commonShape + intArrayOf(min(n, m)))
val vTensor = zeros(commonShape + intArrayOf(min(n, m), m)) val vTensor = zeros(commonShape + intArrayOf(min(n, m), m))
tensor.matrixSequence() val matrices = tensor.matrices
.zip( val uTensors = uTensor.matrices
uTensor.matrixSequence() val sTensorVectors = sTensor.vectors
.zip( val vTensors = vTensor.matrices
sTensor.vectorSequence()
.zip(vTensor.matrixSequence()) for (index in matrices.indices) {
) val matrix = matrices[index]
).forEach { (matrix, USV) -> val usv = Triple(
val matrixSize = matrix.shape.reduce { acc, i -> acc * i } uTensors[index],
val curMatrix = DoubleTensor( sTensorVectors[index],
matrix.shape, vTensors[index]
matrix.mutableBuffer.array().slice(matrix.bufferStart until matrix.bufferStart + matrixSize) )
.toDoubleArray() val matrixSize = matrix.shape.reduce { acc, i -> acc * i }
) val curMatrix = DoubleTensor(
svdHelper(curMatrix, USV, m, n, epsilon) matrix.shape,
} matrix.mutableBuffer.array()
.slice(matrix.bufferStart until matrix.bufferStart + matrixSize)
.toDoubleArray()
)
svdHelper(curMatrix, usv, m, n, epsilon)
}
return Triple(uTensor.transpose(), sTensor, vTensor.transpose()) return Triple(uTensor.transpose(), sTensor, vTensor.transpose())
} }
override fun Tensor<Double>.symEig(): Pair<DoubleTensor, DoubleTensor> = override fun Tensor<Double>.symEig(): Pair<DoubleTensor, DoubleTensor> = symEig(epsilon = 1e-15)
symEig(epsilon = 1e-15)
/** /**
* Returns eigenvalues and eigenvectors of a real symmetric matrix input or a batch of real symmetric matrices, * Returns eigenvalues and eigenvectors of a real symmetric matrix input or a batch of real symmetric matrices,
@ -846,12 +852,26 @@ public open class DoubleTensorAlgebra :
*/ */
public fun Tensor<Double>.symEig(epsilon: Double): Pair<DoubleTensor, DoubleTensor> { public fun Tensor<Double>.symEig(epsilon: Double): Pair<DoubleTensor, DoubleTensor> {
checkSymmetric(tensor, epsilon) checkSymmetric(tensor, epsilon)
fun MutableStructure2D<Double>.cleanSym(n: Int) {
for (i in 0 until n) {
for (j in 0 until n) {
if (i == j) {
this[i, j] = sign(this[i, j])
} else {
this[i, j] = 0.0
}
}
}
}
val (u, s, v) = tensor.svd(epsilon) val (u, s, v) = tensor.svd(epsilon)
val shp = s.shape + intArrayOf(1) val shp = s.shape + intArrayOf(1)
val utv = u.transpose() dot v val utv = u.transpose() dot v
val n = s.shape.last() val n = s.shape.last()
for (matrix in utv.matrixSequence()) for (matrix in utv.matrixSequence()) {
cleanSymHelper(matrix.as2D(), n) matrix.as2D().cleanSym(n)
}
val eig = (utv dot s.view(shp)).view(s.shape) val eig = (utv dot s.view(shp)).view(s.shape)
return eig to v return eig to v

View File

@ -10,41 +10,54 @@ import space.kscience.kmath.nd.MutableStructure2D
import space.kscience.kmath.nd.as1D import space.kscience.kmath.nd.as1D
import space.kscience.kmath.nd.as2D import space.kscience.kmath.nd.as2D
import space.kscience.kmath.operations.invoke import space.kscience.kmath.operations.invoke
import space.kscience.kmath.tensors.core.* import space.kscience.kmath.structures.VirtualBuffer
import space.kscience.kmath.structures.asSequence
import space.kscience.kmath.tensors.core.BufferedTensor
import space.kscience.kmath.tensors.core.DoubleTensor
import space.kscience.kmath.tensors.core.DoubleTensorAlgebra import space.kscience.kmath.tensors.core.DoubleTensorAlgebra
import space.kscience.kmath.tensors.core.DoubleTensorAlgebra.Companion.valueOrNull import space.kscience.kmath.tensors.core.IntTensor
import kotlin.math.abs import kotlin.math.abs
import kotlin.math.min import kotlin.math.min
import kotlin.math.sign
import kotlin.math.sqrt import kotlin.math.sqrt
internal val <T> BufferedTensor<T>.vectors: VirtualBuffer<BufferedTensor<T>>
get() {
val n = shape.size
val vectorOffset = shape[n - 1]
val vectorShape = intArrayOf(shape.last())
internal fun <T> BufferedTensor<T>.vectorSequence(): Sequence<BufferedTensor<T>> = sequence { return VirtualBuffer(numElements / vectorOffset) { index ->
val n = shape.size val offset = index * vectorOffset
val vectorOffset = shape[n - 1] BufferedTensor(vectorShape, mutableBuffer, bufferStart + offset)
val vectorShape = intArrayOf(shape.last()) }
for (offset in 0 until numElements step vectorOffset) {
val vector = BufferedTensor(vectorShape, mutableBuffer, bufferStart + offset)
yield(vector)
} }
}
internal fun <T> BufferedTensor<T>.matrixSequence(): Sequence<BufferedTensor<T>> = sequence {
val n = shape.size internal fun <T> BufferedTensor<T>.vectorSequence(): Sequence<BufferedTensor<T>> = vectors.asSequence()
check(n >= 2) { "Expected tensor with 2 or more dimensions, got size $n" }
val matrixOffset = shape[n - 1] * shape[n - 2] /**
val matrixShape = intArrayOf(shape[n - 2], shape[n - 1]) * A random access alternative to [matrixSequence]
for (offset in 0 until numElements step matrixOffset) { */
val matrix = BufferedTensor(matrixShape, mutableBuffer, bufferStart + offset) internal val <T> BufferedTensor<T>.matrices: VirtualBuffer<BufferedTensor<T>>
yield(matrix) get() {
val n = shape.size
check(n >= 2) { "Expected tensor with 2 or more dimensions, got size $n" }
val matrixOffset = shape[n - 1] * shape[n - 2]
val matrixShape = intArrayOf(shape[n - 2], shape[n - 1])
return VirtualBuffer(numElements / matrixOffset) { index ->
val offset = index * matrixOffset
BufferedTensor(matrixShape, mutableBuffer, bufferStart + offset)
}
} }
}
internal fun <T> BufferedTensor<T>.matrixSequence(): Sequence<BufferedTensor<T>> = matrices.asSequence()
internal fun dotHelper( internal fun dotHelper(
a: MutableStructure2D<Double>, a: MutableStructure2D<Double>,
b: MutableStructure2D<Double>, b: MutableStructure2D<Double>,
res: MutableStructure2D<Double>, res: MutableStructure2D<Double>,
l: Int, m: Int, n: Int l: Int, m: Int, n: Int,
) { ) {
for (i in 0 until l) { for (i in 0 until l) {
for (j in 0 until n) { for (j in 0 until n) {
@ -60,7 +73,7 @@ internal fun dotHelper(
internal fun luHelper( internal fun luHelper(
lu: MutableStructure2D<Double>, lu: MutableStructure2D<Double>,
pivots: MutableStructure1D<Int>, pivots: MutableStructure1D<Int>,
epsilon: Double epsilon: Double,
): Boolean { ): Boolean {
val m = lu.rowNum val m = lu.rowNum
@ -122,7 +135,7 @@ internal fun <T> BufferedTensor<T>.setUpPivots(): IntTensor {
internal fun DoubleTensorAlgebra.computeLU( internal fun DoubleTensorAlgebra.computeLU(
tensor: DoubleTensor, tensor: DoubleTensor,
epsilon: Double epsilon: Double,
): Pair<DoubleTensor, IntTensor>? { ): Pair<DoubleTensor, IntTensor>? {
checkSquareMatrix(tensor.shape) checkSquareMatrix(tensor.shape)
@ -139,7 +152,7 @@ internal fun DoubleTensorAlgebra.computeLU(
internal fun pivInit( internal fun pivInit(
p: MutableStructure2D<Double>, p: MutableStructure2D<Double>,
pivot: MutableStructure1D<Int>, pivot: MutableStructure1D<Int>,
n: Int n: Int,
) { ) {
for (i in 0 until n) { for (i in 0 until n) {
p[i, pivot[i]] = 1.0 p[i, pivot[i]] = 1.0
@ -150,7 +163,7 @@ internal fun luPivotHelper(
l: MutableStructure2D<Double>, l: MutableStructure2D<Double>,
u: MutableStructure2D<Double>, u: MutableStructure2D<Double>,
lu: MutableStructure2D<Double>, lu: MutableStructure2D<Double>,
n: Int n: Int,
) { ) {
for (i in 0 until n) { for (i in 0 until n) {
for (j in 0 until n) { for (j in 0 until n) {
@ -170,7 +183,7 @@ internal fun luPivotHelper(
internal fun choleskyHelper( internal fun choleskyHelper(
a: MutableStructure2D<Double>, a: MutableStructure2D<Double>,
l: MutableStructure2D<Double>, l: MutableStructure2D<Double>,
n: Int n: Int,
) { ) {
for (i in 0 until n) { for (i in 0 until n) {
for (j in 0 until i) { for (j in 0 until i) {
@ -200,7 +213,7 @@ internal fun luMatrixDet(lu: MutableStructure2D<Double>, pivots: MutableStructur
internal fun luMatrixInv( internal fun luMatrixInv(
lu: MutableStructure2D<Double>, lu: MutableStructure2D<Double>,
pivots: MutableStructure1D<Int>, pivots: MutableStructure1D<Int>,
invMatrix: MutableStructure2D<Double> invMatrix: MutableStructure2D<Double>,
) { ) {
val m = lu.shape[0] val m = lu.shape[0]
@ -227,7 +240,7 @@ internal fun luMatrixInv(
internal fun DoubleTensorAlgebra.qrHelper( internal fun DoubleTensorAlgebra.qrHelper(
matrix: DoubleTensor, matrix: DoubleTensor,
q: DoubleTensor, q: DoubleTensor,
r: MutableStructure2D<Double> r: MutableStructure2D<Double>,
) { ) {
checkSquareMatrix(matrix.shape) checkSquareMatrix(matrix.shape)
val n = matrix.shape[0] val n = matrix.shape[0]
@ -280,12 +293,11 @@ internal fun DoubleTensorAlgebra.svd1d(a: DoubleTensor, epsilon: Double = 1e-10)
internal fun DoubleTensorAlgebra.svdHelper( internal fun DoubleTensorAlgebra.svdHelper(
matrix: DoubleTensor, matrix: DoubleTensor,
USV: Pair<BufferedTensor<Double>, Pair<BufferedTensor<Double>, BufferedTensor<Double>>>, USV: Triple<BufferedTensor<Double>, BufferedTensor<Double>, BufferedTensor<Double>>,
m: Int, n: Int, epsilon: Double m: Int, n: Int, epsilon: Double,
) { ) {
val res = ArrayList<Triple<Double, DoubleTensor, DoubleTensor>>(0) val res = ArrayList<Triple<Double, DoubleTensor, DoubleTensor>>(0)
val (matrixU, SV) = USV val (matrixU, matrixS, matrixV) = USV
val (matrixS, matrixV) = SV
for (k in 0 until min(n, m)) { for (k in 0 until min(n, m)) {
var a = matrix.copy() var a = matrix.copy()
@ -329,14 +341,3 @@ internal fun DoubleTensorAlgebra.svdHelper(
matrixV.mutableBuffer.array()[matrixV.bufferStart + i] = vBuffer[i] matrixV.mutableBuffer.array()[matrixV.bufferStart + i] = vBuffer[i]
} }
} }
internal fun cleanSymHelper(matrix: MutableStructure2D<Double>, n: Int) {
for (i in 0 until n)
for (j in 0 until n) {
if (i == j) {
matrix[i, j] = sign(matrix[i, j])
} else {
matrix[i, j] = 0.0
}
}
}