diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayIterators.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayIterators.kt index 115c78cb9..bba5089a1 100644 --- a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayIterators.kt +++ b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayIterators.kt @@ -16,7 +16,7 @@ internal sealed class INDArrayIteratorBase(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(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(iterateOver) { override fun getSingle(indices: LongArray) = iterateOver.getFloat(*indices) } + +internal fun INDArray.floatIterator() = INDArrayFloatIterator(this) diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayStructures.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayStructures.kt index 351110485..ef7436285 100644 --- a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayStructures.kt +++ b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayStructures.kt @@ -8,7 +8,7 @@ interface INDArrayStructure : NDStructure { val ndArray: INDArray override val shape: IntArray - get() = narrowToIntArray(ndArray.shape()) + get() = ndArray.shape().toIntArray() fun elementsIterator(): Iterator> override fun elements(): Sequence> = Sequence(::elementsIterator) @@ -24,7 +24,7 @@ fun INDArray.asIntStructure(): INDArrayIntStructure = INDArrayIntStructure(this) data class INDArrayLongStructure(override val ndArray: INDArray) : INDArrayStructure { override fun elementsIterator(): Iterator> = 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) diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/arrays.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/arrays.kt index 3d5062a4f..269fc89c2 100644 --- a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/arrays.kt +++ b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/arrays.kt @@ -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() } diff --git a/kmath-nd4j/src/test/kotlin/scientifik/kmath/nd4j/INDArrayAlgebraTest.kt b/kmath-nd4j/src/test/kotlin/scientifik/kmath/nd4j/INDArrayAlgebraTest.kt index f971e7871..4aa40c233 100644 --- a/kmath-nd4j/src/test/kotlin/scientifik/kmath/nd4j/INDArrayAlgebraTest.kt +++ b/kmath-nd4j/src/test/kotlin/scientifik/kmath/nd4j/INDArrayAlgebraTest.kt @@ -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) + } }