Cleanup after nd optimization phase II

benchmarks/src/jmh/kotlin/scientifik/kmath/structures/NDFieldBenchmark.kt

@@ -17,7 +17,7 @@ open class NDFieldBenchmark {
 
     @Benchmark
     fun autoElementAdd() {
-        var res = bufferedField.run { one.toElement() }
+        var res = genericField.one
         repeat(n) {
             res += 1.0
         }

benchmarks/src/main/kotlin/scientifik/kmath/structures/NDFieldBenchmark.kt

@@ -16,54 +16,54 @@ fun main(args: Array<String>) {
     //A generic boxing field. It should be used for objects, not primitives.
     val genericField = NDField.buffered(intArrayOf(dim, dim), RealField)
 
-
-    val autoTime = measureTimeMillis {
-        autoField.run {
-            var res = one
-            repeat(n) {
-                res += 1.0
-            }
-        }
-    }
-
-    println("Buffered addition completed in $autoTime millis")
-
-    val elementTime = measureTimeMillis {
-        var res = genericField.one
-        repeat(n) {
-            res += 1.0
-        }
-    }
-
-    println("Element addition completed in $elementTime millis")
-
-    val specializedTime = measureTimeMillis {
-        specializedField.run {
-            var res:NDBuffer<Double> = one
-            repeat(n) {
-                res += 1.0
-            }
-        }
-    }
-
-    println("Specialized addition completed in $specializedTime millis")
-
-
-    val lazyTime = measureTimeMillis {
-        lazyField.run {
-            val res = one.map {
-                var c = 0.0
-                repeat(n) {
-                    c += 1.0
-                }
-                c
-            }
-
-            res.elements().forEach { it.second }
-        }
-    }
-
-    println("Lazy addition completed in $lazyTime millis")
+//
+//    val autoTime = measureTimeMillis {
+//        autoField.run {
+//            var res = one
+//            repeat(n) {
+//                res += 1.0
+//            }
+//        }
+//    }
+//
+//    println("Buffered addition completed in $autoTime millis")
+//
+//    val elementTime = measureTimeMillis {
+//        var res = genericField.one
+//        repeat(n) {
+//            res += 1.0
+//        }
+//    }
+//
+//    println("Element addition completed in $elementTime millis")
+//
+//    val specializedTime = measureTimeMillis {
+//        specializedField.run {
+//            var res:NDBuffer<Double> = one
+//            repeat(n) {
+//                res += 1.0
+//            }
+//        }
+//    }
+//
+//    println("Specialized addition completed in $specializedTime millis")
+//
+//
+//    val lazyTime = measureTimeMillis {
+//        lazyField.run {
+//            val res = one.map {
+//                var c = 0.0
+//                repeat(n) {
+//                    c += 1.0
+//                }
+//                c
+//            }
+//
+//            res.elements().forEach { it.second }
+//        }
+//    }
+//
+//    println("Lazy addition completed in $lazyTime millis")
 
     val genericTime = measureTimeMillis {
         //genericField.run(action)

kmath-core/src/commonMain/kotlin/scientifik/kmath/operations/Algebra.kt

@@ -33,9 +33,9 @@ interface Space<T> {
     operator fun T.times(k: Number) = multiply(this, k.toDouble())
     operator fun T.div(k: Number) = multiply(this, 1.0 / k.toDouble())
     operator fun Number.times(b: T) = b * this
-    fun Iterable<T>.sum(): T = fold(zero) { left, right -> left + right }
+    fun Iterable<T>.sum(): T = fold(zero) { left, right -> add(left,right) }
 
-    fun Sequence<T>.sum(): T = fold(zero) { left, right -> left + right }
+    fun Sequence<T>.sum(): T = fold(zero) { left, right -> add(left, right) }
 }
 
 /**

kmath-core/src/commonMain/kotlin/scientifik/kmath/operations/AlgebraElements.kt

@@ -2,13 +2,20 @@ package scientifik.kmath.operations
 
 /**
  * The generic mathematics elements which is able to store its context
- * @param S the type of mathematical context for this element
+ * @param T the type of space operation results
+ * @param I self type of the element. Needed for static type checking
+ * @param C the type of mathematical context for this element
  */
-interface MathElement<S> {
+interface MathElement<C> {
     /**
      * The context this element belongs to
      */
-    val context: S
+    val context: C
+}
+
+interface MathWrapper<T, I> {
+    fun unwrap(): T
+    fun T.wrap(): I
 }
 
 /**
@@ -17,13 +24,7 @@ interface MathElement<S> {
  * @param I self type of the element. Needed for static type checking
  * @param S the type of space
  */
-interface SpaceElement<T, I : SpaceElement<T, I, S>, S : Space<T>> : MathElement<S> {
-    /**
-     * Self value. Needed for static type checking.
-     */
-    fun unwrap(): T
-
-    fun T.wrap(): I
+interface SpaceElement<T, I : SpaceElement<T, I, S>, S : Space<T>> : MathElement<S>, MathWrapper<T, I> {
 
     operator fun plus(b: T) = context.add(unwrap(), b).wrap()
     operator fun minus(b: T) = context.add(unwrap(), context.multiply(b, -1.0)).wrap()
@@ -35,7 +36,6 @@ interface SpaceElement<T, I : SpaceElement<T, I, S>, S : Space<T>> : MathElement
  * Ring element
  */
 interface RingElement<T, I : RingElement<T, I, R>, R : Ring<T>> : SpaceElement<T, I, R> {
-    override val context: R
     operator fun times(b: T) = context.multiply(unwrap(), b).wrap()
 }
 
@@ -44,6 +44,5 @@ interface RingElement<T, I : RingElement<T, I, R>, R : Ring<T>> : SpaceElement<T
  */
 interface FieldElement<T, I : FieldElement<T, I, F>, F : Field<T>> : RingElement<T, I, F> {
     override val context: F
-
     operator fun div(b: T) = context.divide(unwrap(), b).wrap()
 }

kmath-core/src/commonMain/kotlin/scientifik/kmath/operations/NumberAlgebra.kt

@@ -32,6 +32,7 @@ inline class Real(val value: Double) : FieldElement<Double, Real, RealField> {
 /**
  * A field for double without boxing. Does not produce appropriate field element
  */
+@Suppress("EXTENSION_SHADOWED_BY_MEMBER")
 object RealField : ExtendedField<Double>, Norm<Double, Double> {
     override val zero: Double = 0.0
     override fun add(a: Double, b: Double): Double = a + b
@@ -45,10 +46,13 @@ object RealField : ExtendedField<Double>, Norm<Double, Double> {
     override fun power(arg: Double, pow: Double): Double = arg.pow(pow)
 
     override fun exp(arg: Double): Double = kotlin.math.exp(arg)
-
     override fun ln(arg: Double): Double = kotlin.math.ln(arg)
 
     override fun norm(arg: Double): Double = kotlin.math.abs(arg)
+
+    override fun Double.unaryMinus(): Double  = -this
+
+    override fun Double.minus(b: Double): Double = this - b
 }
 
 /**

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/BoxingNDField.kt

@@ -1,15 +1,19 @@
 package scientifik.kmath.structures
 
 import scientifik.kmath.operations.Field
+import scientifik.kmath.operations.FieldElement
 
 
-class BufferNDField<T, F : Field<T>>(
-    shape: IntArray,
+class BoxingNDField<T, F : Field<T>>(
+    override val shape: IntArray,
     override val elementContext: F,
     val bufferFactory: BufferFactory<T>
-) : StridedNDField<T, F>(shape), NDField<T, F, NDBuffer<T>> {
+) : BufferedNDField<T, F> {
 
-    override fun buildBuffer(size: Int, initializer: (Int) -> T): Buffer<T> = bufferFactory(size, initializer)
+    override val strides: Strides = DefaultStrides(shape)
+
+    override fun buildBuffer(size: Int, initializer: (Int) -> T): Buffer<T> =
+        bufferFactory(size, initializer)
 
     override fun check(vararg elements: NDBuffer<T>) {
         if (!elements.all { it.strides == this.strides }) error("Element strides are not the same as context strides")
@@ -18,14 +22,14 @@ class BufferNDField<T, F : Field<T>>(
     override val zero by lazy { produce { zero } }
     override val one by lazy { produce { one } }
 
-    override fun produce(initializer: F.(IntArray) -> T): StridedNDFieldElement<T, F> =
-        StridedNDFieldElement(
+    override fun produce(initializer: F.(IntArray) -> T) =
+        BufferedNDFieldElement(
             this,
             buildBuffer(strides.linearSize) { offset -> elementContext.initializer(strides.index(offset)) })
 
-    override fun map(arg: NDBuffer<T>, transform: F.(T) -> T): StridedNDFieldElement<T, F> {
+    override fun map(arg: NDBuffer<T>, transform: F.(T) -> T): BufferedNDFieldElement<T, F> {
         check(arg)
-        return StridedNDFieldElement(
+        return BufferedNDFieldElement(
             this,
             buildBuffer(arg.strides.linearSize) { offset -> elementContext.transform(arg.buffer[offset]) })
     }
@@ -33,9 +37,9 @@ class BufferNDField<T, F : Field<T>>(
     override fun mapIndexed(
         arg: NDBuffer<T>,
         transform: F.(index: IntArray, T) -> T
-    ): StridedNDFieldElement<T, F> {
+    ): BufferedNDFieldElement<T, F> {
         check(arg)
-        return StridedNDFieldElement(
+        return BufferedNDFieldElement(
             this,
             buildBuffer(arg.strides.linearSize) { offset ->
                 elementContext.transform(
@@ -49,10 +53,13 @@ class BufferNDField<T, F : Field<T>>(
         a: NDBuffer<T>,
         b: NDBuffer<T>,
         transform: F.(T, T) -> T
-    ): StridedNDFieldElement<T, F> {
+    ): BufferedNDFieldElement<T, F> {
         check(a, b)
-        return StridedNDFieldElement(
+        return BufferedNDFieldElement(
             this,
             buildBuffer(strides.linearSize) { offset -> elementContext.transform(a.buffer[offset], b.buffer[offset]) })
     }
+
+    override fun NDBuffer<T>.toElement(): FieldElement<NDBuffer<T>, *, out BufferedNDField<T, F>> =
+        BufferedNDFieldElement(this@BoxingNDField, buffer)
 }

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/BufferedNDAlgebra.kt

@@ -0,0 +1,45 @@
+package scientifik.kmath.structures
+
+import scientifik.kmath.operations.*
+
+interface BufferedNDAlgebra<T, C>: NDAlgebra<T, C, NDBuffer<T>>{
+    val strides: Strides
+
+    fun buildBuffer(size: Int, initializer: (Int) -> T): Buffer<T>
+
+    override fun check(vararg elements: NDBuffer<T>) {
+        if (!elements.all { it.strides == this.strides }) error("Strides mismatch")
+    }
+
+    /**
+     * Convert any [NDStructure] to buffered structure using strides from this context.
+     * If the structure is already [NDBuffer], conversion is free. If not, it could be expensive because iteration over indexes
+     *
+     * If the argument is [NDBuffer] with different strides structure, the new element will be produced.
+     */
+    fun NDStructure<T>.toBuffer(): NDBuffer<T> {
+        return if (this is NDBuffer<T> && this.strides == this@BufferedNDAlgebra.strides) {
+            this
+        } else {
+            produce { index -> get(index) }
+        }
+    }
+
+    /**
+     * Convert a buffer to element of this algebra
+     */
+    fun NDBuffer<T>.toElement(): MathElement<out BufferedNDAlgebra<T, C>>
+}
+
+
+interface BufferedNDSpace<T, S : Space<T>> : NDSpace<T, S, NDBuffer<T>>, BufferedNDAlgebra<T,S>  {
+    override fun NDBuffer<T>.toElement(): SpaceElement<NDBuffer<T>, *, out BufferedNDSpace<T, S>>
+}
+
+interface BufferedNDRing<T, R : Ring<T>> : NDRing<T, R, NDBuffer<T>>, BufferedNDSpace<T, R> {
+    override fun NDBuffer<T>.toElement(): RingElement<NDBuffer<T>, *, out BufferedNDRing<T, R>>
+}
+
+interface BufferedNDField<T, F : Field<T>> : NDField<T, F, NDBuffer<T>>, BufferedNDRing<T, F> {
+    override fun NDBuffer<T>.toElement(): FieldElement<NDBuffer<T>, *, out BufferedNDField<T, F>>
+}

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/BufferedNDElement.kt

@@ -0,0 +1,88 @@
+package scientifik.kmath.structures
+
+import scientifik.kmath.operations.*
+
+/**
+ * Base interface for an element with context, containing strides
+ */
+interface BufferedNDElement<T, C> : NDBuffer<T>, NDElement<T, C, NDBuffer<T>> {
+    override val context: BufferedNDAlgebra<T, C>
+
+    override val strides get() = context.strides
+
+    override val shape: IntArray get() = context.shape
+}
+
+class BufferedNDSpaceElement<T, S : Space<T>>(
+    override val context: BufferedNDSpace<T, S>,
+    override val buffer: Buffer<T>
+) : BufferedNDElement<T, S>, SpaceElement<NDBuffer<T>, BufferedNDSpaceElement<T, S>, BufferedNDSpace<T, S>> {
+
+    override fun unwrap(): NDBuffer<T> = this
+
+    override fun NDBuffer<T>.wrap(): BufferedNDSpaceElement<T, S> {
+        context.check(this)
+        return BufferedNDSpaceElement(context, buffer)
+    }
+}
+
+class BufferedNDRingElement<T, R : Ring<T>>(
+    override val context: BufferedNDRing<T, R>,
+    override val buffer: Buffer<T>
+) : BufferedNDElement<T, R>, RingElement<NDBuffer<T>, BufferedNDRingElement<T, R>, BufferedNDRing<T, R>> {
+
+    override fun unwrap(): NDBuffer<T> = this
+
+    override fun NDBuffer<T>.wrap(): BufferedNDRingElement<T, R> {
+        context.check(this)
+        return BufferedNDRingElement(context, buffer)
+    }
+}
+
+class BufferedNDFieldElement<T, F : Field<T>>(
+    override val context: BufferedNDField<T, F>,
+    override val buffer: Buffer<T>
+) : BufferedNDElement<T, F>, FieldElement<NDBuffer<T>, BufferedNDFieldElement<T, F>, BufferedNDField<T, F>> {
+
+    override fun unwrap(): NDBuffer<T> = this
+
+    override fun NDBuffer<T>.wrap(): BufferedNDFieldElement<T, F> {
+        context.check(this)
+        return BufferedNDFieldElement(context, buffer)
+    }
+}
+
+
+/**
+ * Element by element application of any operation on elements to the whole array. Just like in numpy
+ */
+operator fun <T : Any, F : Field<T>> Function1<T, T>.invoke(ndElement: BufferedNDElement<T, F>) =
+    ndElement.context.run { map(ndElement) { invoke(it) }.toElement() }
+
+/* plus and minus */
+
+/**
+ * Summation operation for [BufferedNDElement] and single element
+ */
+operator fun <T : Any, F : Space<T>> BufferedNDElement<T, F>.plus(arg: T) =
+    context.map(this) { it + arg }.wrap()
+
+/**
+ * Subtraction operation between [BufferedNDElement] and single element
+ */
+operator fun <T : Any, F : Space<T>> BufferedNDElement<T, F>.minus(arg: T) =
+    context.map(this) { it - arg }.wrap()
+
+/* prod and div */
+
+/**
+ * Product operation for [BufferedNDElement] and single element
+ */
+operator fun <T : Any, F : Ring<T>> BufferedNDElement<T, F>.times(arg: T) =
+    context.map(this) { it * arg }.wrap()
+
+/**
+ * Division operation between [BufferedNDElement] and single element
+ */
+operator fun <T : Any, F : Field<T>> BufferedNDElement<T, F>.div(arg: T) =
+    context.map(this) { it / arg }.wrap()

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/Buffers.kt

@@ -30,7 +30,7 @@ interface Buffer<T> {
          * Create most appropriate immutable buffer for given type avoiding boxing wherever possible
          */
         @Suppress("UNCHECKED_CAST")
-        inline fun <reified T : Any> auto(size: Int, initializer: (Int) -> T): Buffer<T> {
+        inline fun <reified T : Any> auto(size: Int, crossinline initializer: (Int) -> T): Buffer<T> {
             return when (T::class) {
                 Double::class -> DoubleBuffer(DoubleArray(size) { initializer(it) as Double }) as Buffer<T>
                 Short::class -> ShortBuffer(ShortArray(size) { initializer(it) as Short }) as Buffer<T>
@@ -40,8 +40,10 @@ interface Buffer<T> {
             }
         }
 
-        val DoubleBufferFactory: BufferFactory<Double> = { size, initializer -> DoubleBuffer(DoubleArray(size, initializer)) }
-        val ShortBufferFactory: BufferFactory<Short> = { size, initializer -> ShortBuffer(ShortArray(size, initializer)) }
+        val DoubleBufferFactory: BufferFactory<Double> =
+            { size, initializer -> DoubleBuffer(DoubleArray(size, initializer)) }
+        val ShortBufferFactory: BufferFactory<Short> =
+            { size, initializer -> ShortBuffer(ShortArray(size, initializer)) }
         val IntBufferFactory: BufferFactory<Int> = { size, initializer -> IntBuffer(IntArray(size, initializer)) }
         val LongBufferFactory: BufferFactory<Long> = { size, initializer -> LongBuffer(LongArray(size, initializer)) }
     }

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/NDAlgebra.kt

@@ -11,40 +11,79 @@ import scientifik.kmath.operations.Space
 class ShapeMismatchException(val expected: IntArray, val actual: IntArray) : RuntimeException()
 
 
-interface NDSpace<T, S : Space<T>, N : NDStructure<T>> : Space<N> {
+/**
+ * The base interface for all nd-algebra implementations
+ * @param T the type of nd-structure element
+ * @param C the type of the context
+ * @param N the type of the structure
+ */
+interface NDAlgebra<T, C, N : NDStructure<T>> {
     val shape: IntArray
-    val elementContext: S
+    val elementContext: C
 
     /**
      * Produce a new [N] structure using given initializer function
      */
-    fun produce(initializer: S.(IntArray) -> T): N
+    fun produce(initializer: C.(IntArray) -> T): N
 
-    fun map(arg: N, transform: S.(T) -> T): N
-    fun mapIndexed(arg: N, transform: S.(index: IntArray, T) -> T): N
-    fun combine(a: N, b: N, transform: S.(T, T) -> T): N
+    /**
+     * Map elements from one structure to another one
+     */
+    fun map(arg: N, transform: C.(T) -> T): N
 
+    /**
+     * Map indexed elements
+     */
+    fun mapIndexed(arg: N, transform: C.(index: IntArray, T) -> T): N
+
+    /**
+     * Combine two structures into one
+     */
+    fun combine(a: N, b: N, transform: C.(T, T) -> T): N
+
+    /**
+     * Check if given elements are consistent with this context
+     */
+    fun check(vararg elements: N) {
+        elements.forEach {
+            if (!shape.contentEquals(it.shape)) {
+                throw ShapeMismatchException(shape, it.shape)
+            }
+        }
+    }
+
+    /**
+     * element-by-element invoke a function working on [T] on a [NDStructure]
+     */
+    operator fun Function1<T, T>.invoke(structure: N) = map(structure) { value -> this@invoke(value) }
+}
+
+/**
+ * An nd-space over element space
+ */
+interface NDSpace<T, S : Space<T>, N : NDStructure<T>> : Space<N>, NDAlgebra<T, S, N> {
     /**
      * Element-by-element addition
      */
     override fun add(a: N, b: N): N =
-        combine(a, b) { aValue, bValue -> aValue + bValue }
+        combine(a, b) { aValue, bValue -> add(aValue, bValue) }
 
     /**
      * Multiply all elements by constant
      */
     override fun multiply(a: N, k: Double): N =
-        map(a) { it * k }
+        map(a) { multiply(it, k) }
 
-    operator fun Function1<T, T>.invoke(structure: N) = map(structure) { value -> this@invoke(value) }
-    operator fun N.plus(arg: T) = map(this) { value -> elementContext.run { arg + value } }
-    operator fun N.minus(arg: T) = map(this) { value -> elementContext.run { arg - value } }
-
-    operator fun T.plus(arg: N) = arg + this
-    operator fun T.minus(arg: N) = arg - this
+    operator fun N.plus(arg: T) = map(this) { value -> add(arg, value) }
+    operator fun N.minus(arg: T) = map(this) { value -> add(arg, -value) }
 
+    operator fun T.plus(arg: N) = map(arg) { value -> add(this@plus, value) }
+    operator fun T.minus(arg: N) = map(arg) { value -> add(-this@minus, value) }
 }
 
+/**
+ * An nd-ring over element ring
+ */
 interface NDRing<T, R : Ring<T>, N : NDStructure<T>> : Ring<N>, NDSpace<T, R, N> {
 
     /**
@@ -53,16 +92,16 @@ interface NDRing<T, R : Ring<T>, N : NDStructure<T>> : Ring<N>, NDSpace<T, R, N>
     override fun multiply(a: N, b: N): N =
         combine(a, b) { aValue, bValue -> aValue * bValue }
 
-    operator fun N.times(arg: T) = map(this) { value -> elementContext.run { arg * value } }
-    operator fun T.times(arg: N) = arg * this
+    operator fun N.times(arg: T) = map(this) { value -> multiply(arg, value) }
+    operator fun T.times(arg: N) = map(arg) { value -> multiply(this@times, value) }
 }
 
 /**
- * Field for n-dimensional arrays.
+ * Field for n-dimensional structures.
  * @param shape - the list of dimensions of the array
  * @param elementField - operations field defined on individual array element
  * @param T - the type of the element contained in ND structure
- * @param F - field over structure elements
+ * @param F - field of structure elements
  * @param R - actual nd-element type of this field
  */
 interface NDField<T, F : Field<T>, N : NDStructure<T>> : Field<N>, NDRing<T, F, N> {
@@ -73,17 +112,9 @@ interface NDField<T, F : Field<T>, N : NDStructure<T>> : Field<N>, NDRing<T, F,
     override fun divide(a: N, b: N): N =
         combine(a, b) { aValue, bValue -> aValue / bValue }
 
-    operator fun N.div(arg: T) = map(this) { value -> elementContext.run { arg / value } }
+    operator fun N.div(arg: T) = map(this) { value -> arg / value }
     operator fun T.div(arg: N) = arg / this
 
-    fun check(vararg elements: N) {
-        elements.forEach {
-            if (!shape.contentEquals(it.shape)) {
-                throw ShapeMismatchException(shape, it.shape)
-            }
-        }
-    }
-
     companion object {
         /**
          * Create a nd-field for [Double] values
@@ -94,16 +125,16 @@ interface NDField<T, F : Field<T>, N : NDStructure<T>> : Field<N>, NDRing<T, F,
          * Create a nd-field with boxing generic buffer
          */
         fun <T : Any, F : Field<T>> buffered(shape: IntArray, field: F) =
-            BufferNDField(shape, field, Buffer.Companion::boxing)
+            BoxingNDField(shape, field, Buffer.Companion::boxing)
 
         /**
          * Create a most suitable implementation for nd-field using reified class.
          */
         @Suppress("UNCHECKED_CAST")
-        inline fun <reified T : Any, F : Field<T>> auto(shape: IntArray, field: F): StridedNDField<T, F> =
+        inline fun <reified T : Any, F : Field<T>> auto(shape: IntArray, field: F): BufferedNDField<T, F> =
             when {
-                T::class == Double::class -> real(shape) as StridedNDField<T, F>
-                else -> BufferNDField(shape, field, Buffer.Companion::auto)
+                T::class == Double::class -> real(shape) as BufferedNDField<T, F>
+                else -> BoxingNDField(shape, field, Buffer.Companion::auto)
             }
     }
 }

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/NDElement.kt

@@ -1,16 +1,21 @@
 package scientifik.kmath.structures
 
 import scientifik.kmath.operations.Field
-import scientifik.kmath.operations.FieldElement
 import scientifik.kmath.operations.RealField
+import scientifik.kmath.operations.Ring
 import scientifik.kmath.operations.Space
 
 
-interface NDElement<T, S : Space<T>> : NDStructure<T> {
-    val elementField: S
+interface NDElement<T, C, N : NDStructure<T>> : NDStructure<T> {
 
-    fun mapIndexed(transform: S.(index: IntArray, T) -> T): NDElement<T, S>
-    fun map(action: S.(T) -> T) = mapIndexed { _, value -> action(value) }
+    val context: NDAlgebra<T, C, N>
+
+    fun unwrap(): N
+
+    fun N.wrap(): NDElement<T, C, N>
+
+    fun mapIndexed(transform: C.(index: IntArray, T) -> T) = context.mapIndexed(unwrap(), transform).wrap()
+    fun map(transform: C.(T) -> T) = context.map(unwrap(), transform).wrap()
 
     companion object {
         /**
@@ -38,8 +43,8 @@ interface NDElement<T, S : Space<T>> : NDStructure<T> {
             shape: IntArray,
             field: F,
             initializer: F.(IntArray) -> T
-        ): StridedNDFieldElement<T, F> {
-            val ndField = BufferNDField(shape, field, Buffer.Companion::boxing)
+        ): BufferedNDElement<T, F> {
+            val ndField = BoxingNDField(shape, field, Buffer.Companion::boxing)
             return ndField.produce(initializer)
         }
 
@@ -47,47 +52,46 @@ interface NDElement<T, S : Space<T>> : NDStructure<T> {
             shape: IntArray,
             field: F,
             noinline initializer: F.(IntArray) -> T
-        ): StridedNDFieldElement<T, F> {
+        ): BufferedNDFieldElement<T, F> {
             val ndField = NDField.auto(shape, field)
-            return StridedNDFieldElement(ndField, ndField.produce(initializer).buffer)
+            return BufferedNDFieldElement(ndField, ndField.produce(initializer).buffer)
         }
     }
 }
 
-
 /**
- * Element by element application of any operation on elements to the whole array. Just like in numpy
+ * Element by element application of any operation on elements to the whole [NDElement]
  */
-operator fun <T, F : Field<T>> Function1<T, T>.invoke(ndElement: NDElement<T, F>) =
+operator fun <T, C> Function1<T, T>.invoke(ndElement: NDElement<T, C, *>) =
     ndElement.map { value -> this@invoke(value) }
 
 /* plus and minus */
 
 /**
- * Summation operation for [NDElements] and single element
+ * Summation operation for [NDElement] and single element
  */
-operator fun <T, F : Field<T>> NDElement<T, F>.plus(arg: T): NDElement<T, F> =
-    this.map { value -> elementField.run { arg + value } }
+operator fun <T, S : Space<T>> NDElement<T, S, *>.plus(arg: T) =
+    map { value -> arg + value }
 
 /**
- * Subtraction operation between [NDElements] and single element
+ * Subtraction operation between [NDElement] and single element
  */
-operator fun <T, F : Field<T>> NDElement<T, F>.minus(arg: T): NDElement<T, F> =
-    this.map { value -> elementField.run { arg - value } }
+operator fun <T, S : Space<T>> NDElement<T, S, *>.minus(arg: T) =
+    map { value -> arg - value }
 
 /* prod and div */
 
 /**
- * Product operation for [NDElements] and single element
+ * Product operation for [NDElement] and single element
  */
-operator fun <T, F : Field<T>> NDElement<T, F>.times(arg: T): NDElement<T, F> =
-    this.map { value -> elementField.run { arg * value } }
+operator fun <T, R : Ring<T>> NDElement<T, R, *>.times(arg: T) =
+    map { value -> arg * value }
 
 /**
- * Division operation between [NDElements] and single element
+ * Division operation between [NDElement] and single element
  */
-operator fun <T, F : Field<T>> NDElement<T, F>.div(arg: T): NDElement<T, F> =
-    this.map { value -> elementField.run { arg / value } }
+operator fun <T, F : Field<T>> NDElement<T, F, *>.div(arg: T) =
+    map { value -> arg / value }
 
 
 //    /**
@@ -117,24 +121,3 @@ operator fun <T, F : Field<T>> NDElement<T, F>.div(arg: T): NDElement<T, F> =
 //    operator fun T.div(arg: NDStructure<T>): NDElement<T, F> = produce { index ->
 //        field.run { this@div / arg[index] }
 //    }
-
-
-/**
- *  Read-only [NDStructure] coupled to the context.
- */
-class GenericNDFieldElement<T, F : Field<T>>(
-    override val context: NDField<T, F, NDStructure<T>>,
-    private val structure: NDStructure<T>
-) :
-    NDStructure<T> by structure,
-    NDElement<T, F>,
-    FieldElement<NDStructure<T>, GenericNDFieldElement<T, F>, NDField<T, F, NDStructure<T>>> {
-    override val elementField: F get() = context.elementContext
-
-    override fun unwrap(): NDStructure<T> = structure
-
-    override fun NDStructure<T>.wrap() = GenericNDFieldElement(context, this)
-
-    override fun mapIndexed(transform: F.(index: IntArray, T) -> T) =
-        ndStructure(context.shape) { index: IntArray -> context.elementContext.transform(index, get(index)) }.wrap()
-}

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/RealNDField.kt

@@ -1,13 +1,16 @@
 package scientifik.kmath.structures
 
+import scientifik.kmath.operations.FieldElement
 import scientifik.kmath.operations.RealField
 
-typealias RealNDElement = StridedNDFieldElement<Double, RealField>
+typealias RealNDElement = BufferedNDFieldElement<Double, RealField>
 
-class RealNDField(shape: IntArray) :
-    StridedNDField<Double, RealField>(shape),
+class RealNDField(override val shape: IntArray) :
+    BufferedNDField<Double, RealField>,
     ExtendedNDField<Double, RealField, NDBuffer<Double>> {
 
+    override val strides: Strides = DefaultStrides(shape)
+
     override val elementContext: RealField get() = RealField
     override val zero by lazy { produce { zero } }
     override val one by lazy { produce { one } }
@@ -25,20 +28,20 @@ class RealNDField(shape: IntArray) :
     ): RealNDElement {
         check(arg)
         val array = buildBuffer(arg.strides.linearSize) { offset -> RealField.transform(arg.buffer[offset]) }
-        return StridedNDFieldElement(this, array)
+        return BufferedNDFieldElement(this, array)
     }
 
     override fun produce(initializer: RealField.(IntArray) -> Double): RealNDElement {
         val array = buildBuffer(strides.linearSize) { offset -> elementContext.initializer(strides.index(offset)) }
-        return StridedNDFieldElement(this, array)
+        return BufferedNDFieldElement(this, array)
     }
 
     override fun mapIndexed(
         arg: NDBuffer<Double>,
         transform: RealField.(index: IntArray, Double) -> Double
-    ): StridedNDFieldElement<Double, RealField> {
+    ): RealNDElement {
         check(arg)
-        return StridedNDFieldElement(
+        return BufferedNDFieldElement(
             this,
             buildBuffer(arg.strides.linearSize) { offset ->
                 elementContext.transform(
@@ -52,13 +55,16 @@ class RealNDField(shape: IntArray) :
         a: NDBuffer<Double>,
         b: NDBuffer<Double>,
         transform: RealField.(Double, Double) -> Double
-    ): StridedNDFieldElement<Double, RealField> {
+    ): RealNDElement {
         check(a, b)
-        return StridedNDFieldElement(
+        return BufferedNDFieldElement(
             this,
             buildBuffer(strides.linearSize) { offset -> elementContext.transform(a.buffer[offset], b.buffer[offset]) })
     }
 
+    override fun NDBuffer<Double>.toElement(): FieldElement<NDBuffer<Double>, *, out BufferedNDField<Double, RealField>> =
+        BufferedNDFieldElement(this@RealNDField, buffer)
+
     override fun power(arg: NDBuffer<Double>, pow: Double) = map(arg) { power(it, pow) }
 
     override fun exp(arg: NDBuffer<Double>) = map(arg) { exp(it) }
@@ -75,9 +81,9 @@ class RealNDField(shape: IntArray) :
 /**
  * Fast element production using function inlining
  */
-inline fun StridedNDField<Double, RealField>.produceInline(crossinline initializer: RealField.(Int) -> Double): RealNDElement {
+inline fun BufferedNDField<Double, RealField>.produceInline(crossinline initializer: RealField.(Int) -> Double): RealNDElement {
     val array = DoubleArray(strides.linearSize) { offset -> RealField.initializer(offset) }
-    return StridedNDFieldElement(this, DoubleBuffer(array))
+    return BufferedNDFieldElement(this, DoubleBuffer(array))
 }
 
 /**
@@ -90,13 +96,13 @@ operator fun Function1<Double, Double>.invoke(ndElement: RealNDElement) =
 /* plus and minus */
 
 /**
- * Summation operation for [StridedNDFieldElement] and single element
+ * Summation operation for [BufferedNDElement] and single element
  */
 operator fun RealNDElement.plus(arg: Double) =
     context.produceInline { i -> buffer[i] + arg }
 
 /**
- * Subtraction operation between [StridedNDFieldElement] and single element
+ * Subtraction operation between [BufferedNDElement] and single element
  */
 operator fun RealNDElement.minus(arg: Double) =
     context.produceInline { i -> buffer[i] - arg }

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/ShortNDRing.kt

@@ -1,86 +1,96 @@
-//package scientifik.kmath.structures
-//
-//import scientifik.kmath.operations.ShortRing
-//
-//
-////typealias ShortNDElement = StridedNDFieldElement<Short, ShortRing>
-//
-//class ShortNDRing(shape: IntArray) :
-//    NDRing<Short, ShortRing, NDBuffer<Short>> {
-//
-//    inline fun buildBuffer(size: Int, crossinline initializer: (Int) -> Short): Buffer<Short> =
-//        ShortBuffer(ShortArray(size) { initializer(it) })
-//
-//    /**
-//     * Inline transform an NDStructure to
-//     */
-//    override fun map(
-//        arg: NDBuffer<Short>,
-//        transform: ShortRing.(Short) -> Short
-//    ): ShortNDElement {
-//        check(arg)
-//        val array = buildBuffer(arg.strides.linearSize) { offset -> ShortRing.transform(arg.buffer[offset]) }
-//        return StridedNDFieldElement(this, array)
-//    }
-//
-//    override fun produce(initializer: ShortRing.(IntArray) -> Short): ShortNDElement {
-//        val array = buildBuffer(strides.linearSize) { offset -> elementContext.initializer(strides.index(offset)) }
-//        return StridedNDFieldElement(this, array)
-//    }
-//
-//    override fun mapIndexed(
-//        arg: NDBuffer<Short>,
-//        transform: ShortRing.(index: IntArray, Short) -> Short
-//    ): StridedNDFieldElement<Short, ShortRing> {
-//        check(arg)
-//        return StridedNDFieldElement(
-//            this,
-//            buildBuffer(arg.strides.linearSize) { offset ->
-//                elementContext.transform(
-//                    arg.strides.index(offset),
-//                    arg.buffer[offset]
-//                )
-//            })
-//    }
-//
-//    override fun combine(
-//        a: NDBuffer<Short>,
-//        b: NDBuffer<Short>,
-//        transform: ShortRing.(Short, Short) -> Short
-//    ): StridedNDFieldElement<Short, ShortRing> {
-//        check(a, b)
-//        return StridedNDFieldElement(
-//            this,
-//            buildBuffer(strides.linearSize) { offset -> elementContext.transform(a.buffer[offset], b.buffer[offset]) })
-//    }
-//}
-//
-//
-///**
-// * Fast element production using function inlining
-// */
-//inline fun StridedNDField<Short, ShortRing>.produceInline(crossinline initializer: ShortRing.(Int) -> Short): ShortNDElement {
-//    val array = ShortArray(strides.linearSize) { offset -> ShortRing.initializer(offset) }
-//    return StridedNDFieldElement(this, ShortBuffer(array))
-//}
-//
-///**
-// * Element by element application of any operation on elements to the whole array. Just like in numpy
-// */
-//operator fun Function1<Short, Short>.invoke(ndElement: ShortNDElement) =
-//    ndElement.context.produceInline { i -> invoke(ndElement.buffer[i]) }
-//
-//
-///* plus and minus */
-//
-///**
-// * Summation operation for [StridedNDFieldElement] and single element
-// */
-//operator fun ShortNDElement.plus(arg: Short) =
-//    context.produceInline { i -> buffer[i] + arg }
-//
-///**
-// * Subtraction operation between [StridedNDFieldElement] and single element
-// */
-//operator fun ShortNDElement.minus(arg: Short) =
-//    context.produceInline { i -> buffer[i] - arg }
+package scientifik.kmath.structures
+
+import scientifik.kmath.operations.RingElement
+import scientifik.kmath.operations.ShortRing
+
+
+typealias ShortNDElement = BufferedNDRingElement<Short, ShortRing>
+
+class ShortNDRing(override val shape: IntArray) :
+    BufferedNDRing<Short, ShortRing> {
+
+    override val strides: Strides = DefaultStrides(shape)
+
+    override val elementContext: ShortRing get() = ShortRing
+    override val zero by lazy { produce { ShortRing.zero } }
+    override val one by lazy { produce { ShortRing.one } }
+
+    override inline fun buildBuffer(size: Int, crossinline initializer: (Int) -> Short): Buffer<Short> =
+        ShortBuffer(ShortArray(size) { initializer(it) })
+
+    /**
+     * Inline transform an NDStructure to
+     */
+    override fun map(
+        arg: NDBuffer<Short>,
+        transform: ShortRing.(Short) -> Short
+    ): ShortNDElement {
+        check(arg)
+        val array = buildBuffer(arg.strides.linearSize) { offset -> ShortRing.transform(arg.buffer[offset]) }
+        return BufferedNDRingElement(this, array)
+    }
+
+    override fun produce(initializer: ShortRing.(IntArray) -> Short): ShortNDElement {
+        val array = buildBuffer(strides.linearSize) { offset -> elementContext.initializer(strides.index(offset)) }
+        return BufferedNDRingElement(this, array)
+    }
+
+    override fun mapIndexed(
+        arg: NDBuffer<Short>,
+        transform: ShortRing.(index: IntArray, Short) -> Short
+    ): ShortNDElement {
+        check(arg)
+        return BufferedNDRingElement(
+            this,
+            buildBuffer(arg.strides.linearSize) { offset ->
+                elementContext.transform(
+                    arg.strides.index(offset),
+                    arg.buffer[offset]
+                )
+            })
+    }
+
+    override fun combine(
+        a: NDBuffer<Short>,
+        b: NDBuffer<Short>,
+        transform: ShortRing.(Short, Short) -> Short
+    ): ShortNDElement {
+        check(a, b)
+        return BufferedNDRingElement(
+            this,
+            buildBuffer(strides.linearSize) { offset -> elementContext.transform(a.buffer[offset], b.buffer[offset]) })
+    }
+
+    override fun NDBuffer<Short>.toElement(): RingElement<NDBuffer<Short>, *, out BufferedNDRing<Short, ShortRing>> =
+        BufferedNDRingElement(this@ShortNDRing, buffer)
+}
+
+
+/**
+ * Fast element production using function inlining
+ */
+inline fun BufferedNDRing<Short, ShortRing>.produceInline(crossinline initializer: ShortRing.(Int) -> Short): ShortNDElement {
+    val array = ShortArray(strides.linearSize) { offset -> ShortRing.initializer(offset) }
+    return BufferedNDRingElement(this, ShortBuffer(array))
+}
+
+/**
+ * Element by element application of any operation on elements to the whole array. Just like in numpy
+ */
+operator fun Function1<Short, Short>.invoke(ndElement: ShortNDElement) =
+    ndElement.context.produceInline { i -> invoke(ndElement.buffer[i]) }
+
+
+/* plus and minus */
+
+/**
+ * Summation operation for [StridedNDFieldElement] and single element
+ */
+operator fun ShortNDElement.plus(arg: Short) =
+    context.produceInline { i -> (buffer[i] + arg).toShort() }
+
+/**
+ * Subtraction operation between [StridedNDFieldElement] and single element
+ */
+operator fun ShortNDElement.minus(arg: Short) =
+    context.produceInline { i -> (buffer[i] - arg).toShort() }

kmath-core/src/commonMain/kotlin/scientifik/kmath/structures/StridedContext.kt

@@ -1,97 +0,0 @@
-package scientifik.kmath.structures
-
-import scientifik.kmath.operations.Field
-import scientifik.kmath.operations.FieldElement
-
-abstract class StridedNDField<T, F : Field<T>>(final override val shape: IntArray) : NDField<T, F, NDBuffer<T>> {
-    val strides = DefaultStrides(shape)
-
-    abstract fun buildBuffer(size: Int, initializer: (Int) -> T): Buffer<T>
-
-    /**
-     * Convert any [NDStructure] to buffered structure using strides from this context.
-     * If the structure is already [NDBuffer], conversion is free. If not, it could be expensive because iteration over indexes
-     *
-     * If the argument is [NDBuffer] with different strides structure, the new element will be produced.
-     */
-    fun NDStructure<T>.toBuffer(): NDBuffer<T> {
-        return if (this is NDBuffer<T> && this.strides == this@StridedNDField.strides) {
-            this
-        } else {
-            produce { index -> get(index) }
-        }
-    }
-
-    fun NDBuffer<T>.toElement(): StridedNDFieldElement<T, F> =
-        StridedNDFieldElement(this@StridedNDField, buffer)
-}
-
-class StridedNDFieldElement<T, F : Field<T>>(
-    override val context: StridedNDField<T, F>,
-    override val buffer: Buffer<T>
-) :
-    NDBuffer<T>,
-    FieldElement<NDBuffer<T>, StridedNDFieldElement<T, F>, StridedNDField<T, F>>,
-    NDElement<T, F> {
-
-    override val elementField: F
-        get() = context.elementContext
-
-    override fun unwrap(): NDBuffer<T> =
-        this
-
-    override fun NDBuffer<T>.wrap(): StridedNDFieldElement<T, F> =
-        StridedNDFieldElement(context, this.buffer)
-
-    override val strides
-        get() = context.strides
-
-    override val shape: IntArray
-        get() = context.shape
-
-    override fun get(index: IntArray): T =
-        buffer[strides.offset(index)]
-
-    override fun elements(): Sequence<Pair<IntArray, T>> =
-        strides.indices().map { it to get(it) }
-
-    override fun map(action: F.(T) -> T) =
-        context.run { map(this@StridedNDFieldElement, action) }.wrap()
-
-    override fun mapIndexed(transform: F.(index: IntArray, T) -> T) =
-        context.run { mapIndexed(this@StridedNDFieldElement, transform) }.wrap()
-}
-
-/**
- * Element by element application of any operation on elements to the whole array. Just like in numpy
- */
-operator fun <T : Any, F : Field<T>> Function1<T, T>.invoke(ndElement: StridedNDFieldElement<T, F>) =
-    ndElement.context.run { ndElement.map { invoke(it) } }
-
-/* plus and minus */
-
-/**
- * Summation operation for [StridedNDFieldElement] and single element
- */
-operator fun <T : Any, F : Field<T>> StridedNDFieldElement<T, F>.plus(arg: T) =
-    context.run { map { it + arg } }
-
-/**
- * Subtraction operation between [StridedNDFieldElement] and single element
- */
-operator fun <T : Any, F : Field<T>> StridedNDFieldElement<T, F>.minus(arg: T) =
-    context.run { map { it - arg } }
-
-/* prod and div */
-
-/**
- * Product operation for [StridedNDFieldElement] and single element
- */
-operator fun <T : Any, F : Field<T>> StridedNDFieldElement<T, F>.times(arg: T) =
-    context.run { map { it * arg } }
-
-/**
- * Division operation between [StridedNDFieldElement] and single element
- */
-operator fun <T : Any, F : Field<T>> StridedNDFieldElement<T, F>.div(arg: T) =
-    context.run { map { it / arg } }

kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt

@@ -29,7 +29,7 @@ fun MutableBuffer.Companion.complex(size: Int) =
     ObjectBuffer.create(ComplexBufferSpec, size)
 
 fun NDField.Companion.complex(shape: IntArray) =
-    BufferNDField(shape, ComplexField) { size, init -> ObjectBuffer.create(ComplexBufferSpec, size, init) }
+    BoxingNDField(shape, ComplexField) { size, init -> ObjectBuffer.create(ComplexBufferSpec, size, init) }
 
 fun NDElement.Companion.complex(shape: IntArray, initializer: ComplexField.(IntArray) -> Complex) =
     NDField.complex(shape).produce(initializer)

kmath-coroutines/src/commonMain/kotlin/scientifik/kmath/structures/LazyNDField.kt

@@ -8,8 +8,7 @@ class LazyNDField<T, F : Field<T>>(
     override val shape: IntArray,
     override val elementContext: F,
     val scope: CoroutineScope = GlobalScope
-) :
-    NDField<T, F, NDStructure<T>> {
+) : NDField<T, F, NDStructure<T>> {
 
     override val zero by lazy { produce { zero } }
 
@@ -65,7 +64,8 @@ class LazyNDField<T, F : Field<T>>(
 class LazyNDStructure<T, F : Field<T>>(
     override val context: LazyNDField<T, F>,
     val function: suspend (IntArray) -> T
-) : FieldElement<NDStructure<T>, LazyNDStructure<T, F>, LazyNDField<T, F>>, NDElement<T, F> {
+) : FieldElement<NDStructure<T>, LazyNDStructure<T, F>, LazyNDField<T, F>>,
+    NDElement<T, F, NDStructure<T>> {
 
 
     override fun unwrap(): NDStructure<T> = this
@@ -73,10 +73,6 @@ class LazyNDStructure<T, F : Field<T>>(
     override fun NDStructure<T>.wrap(): LazyNDStructure<T, F> = LazyNDStructure(context) { await(it) }
 
     override val shape: IntArray get() = context.shape
-    override val elementField: F get() = context.elementContext
-
-    override fun mapIndexed(transform: F.(index: IntArray, T) -> T): NDElement<T, F> =
-        context.run { mapIndexed(this@LazyNDStructure, transform) }
 
     private val cache = HashMap<IntArray, Deferred<T>>()