From bf071bcdc1f1040bf644fe9638827c3ffe36491e Mon Sep 17 00:00:00 2001
From: Iaroslav <postovalovya@gmail.com>
Date: Mon, 29 Jun 2020 22:30:08 +0700
Subject: [PATCH] Minor refactor

---
 kmath-nd4j/README.md                             |  2 +-
 .../scientifik.kmath.nd4j/INDArrayAlgebra.kt     | 16 +++++++++-------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/kmath-nd4j/README.md b/kmath-nd4j/README.md
index 63687a880..ad76799c3 100644
--- a/kmath-nd4j/README.md
+++ b/kmath-nd4j/README.md
@@ -50,7 +50,7 @@ array[intArrayOf(0, 0)] = 24.0
 println(array[0, 0]) // 24.0
 ```
 
-Fast element-wise arithmetics for INDArray:
+Fast element-wise and in-place arithmetics for INDArray:
 
 ```kotlin
 import org.nd4j.linalg.factory.*
diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayAlgebra.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayAlgebra.kt
index f476af0d5..14fe202c3 100644
--- a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayAlgebra.kt
+++ b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayAlgebra.kt
@@ -9,24 +9,22 @@ import scientifik.kmath.structures.NDRing
 
 interface INDArrayRing<T, R, N> :
     NDRing<T, R, N> where R : Ring<T>, N : INDArrayStructure<T>, N : MutableNDStructure<T> {
-    fun INDArray.wrap(): N
-
     override val zero: N
         get() = Nd4j.zeros(*shape).wrap()
 
     override val one: N
         get() = Nd4j.ones(*shape).wrap()
 
+    fun INDArray.wrap(): N
+
     override fun produce(initializer: R.(IntArray) -> T): N {
-        val struct = Nd4j.create(*shape).wrap()
+        val struct = Nd4j.create(*shape)!!.wrap()
         struct.elements().map(Pair<IntArray, T>::first).forEach { struct[it] = elementContext.initializer(it) }
         return struct
     }
 
     override fun map(arg: N, transform: R.(T) -> T): N {
-        val new = Nd4j.create(*shape)
-        Nd4j.copy(arg.ndArray, new)
-        val newStruct = new.wrap()
+        val newStruct = arg.ndArray.dup().wrap()
         newStruct.elements().forEach { (idx, value) -> newStruct[idx] = elementContext.transform(value) }
         return newStruct
     }
@@ -52,6 +50,7 @@ interface INDArrayRing<T, R, N> :
     override fun N.minus(b: Number): N = ndArray.subi(b).wrap()
     override fun N.plus(b: Number): N = ndArray.addi(b).wrap()
     override fun N.times(k: Number): N = ndArray.muli(k).wrap()
+    override fun Number.minus(b: N): N = b.ndArray.rsubi(this).wrap()
 }
 
 interface INDArrayField<T, F, N> : NDField<T, F, N>,
@@ -61,13 +60,14 @@ interface INDArrayField<T, F, N> : NDField<T, F, N>,
 }
 
 class RealINDArrayField(override val shape: IntArray, override val elementContext: Field<Double> = RealField) :
-     INDArrayField<Double, Field<Double>, INDArrayRealStructure> {
+    INDArrayField<Double, Field<Double>, INDArrayRealStructure> {
     override fun INDArray.wrap(): INDArrayRealStructure = asRealStructure()
     override fun INDArrayRealStructure.div(arg: Double): INDArrayRealStructure = ndArray.divi(arg).wrap()
     override fun INDArrayRealStructure.plus(arg: Double): INDArrayRealStructure = ndArray.addi(arg).wrap()
     override fun INDArrayRealStructure.minus(arg: Double): INDArrayRealStructure = ndArray.subi(arg).wrap()
     override fun INDArrayRealStructure.times(arg: Double): INDArrayRealStructure = ndArray.muli(arg).wrap()
     override fun Double.div(arg: INDArrayRealStructure): INDArrayRealStructure = arg.ndArray.rdivi(this).wrap()
+    override fun Double.minus(arg: INDArrayRealStructure): INDArrayRealStructure = arg.ndArray.rsubi(this).wrap()
 }
 
 class FloatINDArrayField(override val shape: IntArray, override val elementContext: Field<Float> = FloatField) :
@@ -78,6 +78,7 @@ class FloatINDArrayField(override val shape: IntArray, override val elementConte
     override fun INDArrayFloatStructure.minus(arg: Float): INDArrayFloatStructure = ndArray.subi(arg).wrap()
     override fun INDArrayFloatStructure.times(arg: Float): INDArrayFloatStructure = ndArray.muli(arg).wrap()
     override fun Float.div(arg: INDArrayFloatStructure): INDArrayFloatStructure = arg.ndArray.rdivi(this).wrap()
+    override fun Float.minus(arg: INDArrayFloatStructure): INDArrayFloatStructure = arg.ndArray.rsubi(this).wrap()
 }
 
 class IntINDArrayRing(override val shape: IntArray, override val elementContext: Ring<Int> = IntRing) :
@@ -86,4 +87,5 @@ class IntINDArrayRing(override val shape: IntArray, override val elementContext:
     override fun INDArrayIntStructure.plus(arg: Int): INDArrayIntStructure = ndArray.addi(arg).wrap()
     override fun INDArrayIntStructure.minus(arg: Int): INDArrayIntStructure = ndArray.subi(arg).wrap()
     override fun INDArrayIntStructure.times(arg: Int): INDArrayIntStructure = ndArray.muli(arg).wrap()
+    override fun Int.minus(arg: INDArrayIntStructure): INDArrayIntStructure = arg.ndArray.rsubi(this).wrap()
 }