Refactor array functions

This commit is contained in:
Iaroslav 2020-06-29 21:31:08 +07:00
parent 23b2ba9950
commit d87dd3e717
No known key found for this signature in database
GPG Key ID: 46E15E4A31B3BCD7
4 changed files with 22 additions and 8 deletions

View File

@ -16,7 +16,7 @@ internal sealed class INDArrayIteratorBase<T>(protected val iterateOver: INDArra
else
Shape.ind2sub(iterateOver, i++.toLong())!!
return narrowToIntArray(la) to getSingle(la)
return la.toIntArray() to getSingle(la)
}
}
@ -30,13 +30,16 @@ internal class INDArrayLongIterator(iterateOver: INDArray) : INDArrayIteratorBas
override fun getSingle(indices: LongArray) = iterateOver.getLong(*indices)
}
// TODO
//internal fun INDArray.longI
internal fun INDArray.longIterator(): INDArrayLongIterator = INDArrayLongIterator(this)
internal class INDArrayIntIterator(iterateOver: INDArray) : INDArrayIteratorBase<Int>(iterateOver) {
override fun getSingle(indices: LongArray) = iterateOver.getInt(*narrowToIntArray(indices))
override fun getSingle(indices: LongArray) = iterateOver.getInt(*indices.toIntArray())
}
internal fun INDArray.intIterator(): INDArrayIntIterator = INDArrayIntIterator(this)
internal class INDArrayFloatIterator(iterateOver: INDArray) : INDArrayIteratorBase<Float>(iterateOver) {
override fun getSingle(indices: LongArray) = iterateOver.getFloat(*indices)
}
internal fun INDArray.floatIterator() = INDArrayFloatIterator(this)

View File

@ -8,7 +8,7 @@ interface INDArrayStructure<T> : NDStructure<T> {
val ndArray: INDArray
override val shape: IntArray
get() = narrowToIntArray(ndArray.shape())
get() = ndArray.shape().toIntArray()
fun elementsIterator(): Iterator<Pair<IntArray, T>>
override fun elements(): Sequence<Pair<IntArray, T>> = Sequence(::elementsIterator)
@ -24,7 +24,7 @@ fun INDArray.asIntStructure(): INDArrayIntStructure = INDArrayIntStructure(this)
data class INDArrayLongStructure(override val ndArray: INDArray) : INDArrayStructure<Long> {
override fun elementsIterator(): Iterator<Pair<IntArray, Long>> = INDArrayLongIterator(ndArray)
override fun get(index: IntArray): Long = ndArray.getLong(*widenToLongArray(index))
override fun get(index: IntArray): Long = ndArray.getLong(*index.toLongArray())
}
fun INDArray.asLongStructure(): INDArrayLongStructure = INDArrayLongStructure(this)

View File

@ -1,4 +1,4 @@
package scientifik.kmath.nd4j
internal fun widenToLongArray(ia: IntArray): LongArray = LongArray(ia.size) { ia[it].toLong() }
internal fun narrowToIntArray(la: LongArray): IntArray = IntArray(la.size) { la[it].toInt() }
internal fun IntArray.toLongArray(): LongArray = LongArray(size) { this[it].toLong() }
internal fun LongArray.toIntArray(): IntArray = IntArray(size) { this[it].toInt() }

View File

@ -27,4 +27,15 @@ internal class INDArrayAlgebraTest {
expected[intArrayOf(1, 1)] = 3
assertEquals(expected, res)
}
@Test
fun testAdd() {
val res = (IntINDArrayRing(intArrayOf(2, 2))) { one + 25 }
val expected = Nd4j.create(2, 2)!!.asIntStructure()
expected[intArrayOf(0, 0)] = 26
expected[intArrayOf(0, 1)] = 26
expected[intArrayOf(1, 0)] = 26
expected[intArrayOf(1, 1)] = 26
assertEquals(expected, res)
}
}