From d87dd3e717cc18851b7922d1e31b7ac81698c082 Mon Sep 17 00:00:00 2001
From: Iaroslav <postovalovya@gmail.com>
Date: Mon, 29 Jun 2020 21:31:08 +0700
Subject: [PATCH] Refactor array functions

---
 .../kotlin/scientifik.kmath.nd4j/INDArrayIterators.kt | 11 +++++++----
 .../scientifik.kmath.nd4j/INDArrayStructures.kt       |  4 ++--
 .../src/main/kotlin/scientifik.kmath.nd4j/arrays.kt   |  4 ++--
 .../scientifik/kmath/nd4j/INDArrayAlgebraTest.kt      | 11 +++++++++++
 4 files changed, 22 insertions(+), 8 deletions(-)

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<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)
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<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)
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)
+    }
 }