Fixes to tests and lazy structures

benchmarks/build.gradle

@@ -5,6 +5,7 @@ plugins {
 }
 
 dependencies {
-    compile project(':kmath-core')
+    compile project(":kmath-core")
+    compile project(":kmath-coroutines")
     //jmh project(':kmath-core')
 }

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

@@ -1,15 +1,16 @@
 package scientifik.kmath.structures
 
-import scientifik.kmath.operations.DoubleField
+import scientifik.kmath.operations.RealField
 import kotlin.system.measureTimeMillis
 
 fun main(args: Array<String>) {
     val dim = 1000
-    val n = 10000
+    val n = 100
 
-    val bufferedField = NDField.buffered(intArrayOf(dim, dim), DoubleField)
+    val bufferedField = NDField.buffered(intArrayOf(dim, dim), RealField)
     val specializedField = NDField.real(intArrayOf(dim, dim))
-    val genericField = NDField.generic(intArrayOf(dim, dim), DoubleField)
+    val genericField = NDField.generic(intArrayOf(dim, dim), RealField)
+    val lazyNDField = NDField.lazy(intArrayOf(dim, dim), RealField)
 
 //    val action: NDField<Double, DoubleField, NDStructure<Double>>.() -> Unit = {
 //        var res = one
@@ -55,6 +56,23 @@ fun main(args: Array<String>) {
     println("Specialized addition completed in $specializedTime millis")
 
 
+    val lazyTime = measureTimeMillis {
+        val tr : RealField.(Double)->Double = {arg->
+            var r = arg
+            repeat(n) {
+                r += 1.0
+            }
+            r
+        }
+        lazyNDField.run {
+            val res = one.map(tr)
+
+            res.elements().sumByDouble { it.second }
+        }
+    }
+
+    println("Lazy addition completed in $lazyTime millis")
+
     val genericTime = measureTimeMillis {
         //genericField.run(action)
         genericField.run {

kmath-core/src/commonMain/kotlin/scientifik/kmath/histogram/PhantomHistogram.kt

@@ -42,7 +42,7 @@ class PhantomBin<T : Comparable<T>>(val template: BinTemplate<T>, override val v
 /**
  * Immutable histogram with explicit structure for content and additional external bin description.
  * Bin search is slow, but full histogram algebra is supported.
- * @param bins map a template into structure index
+ * @param bins transform a template into structure index
  */
 class PhantomHistogram<T : Comparable<T>>(
     val bins: Map<BinTemplate<T>, IntArray>,

kmath-core/src/commonMain/kotlin/scientifik/kmath/linear/LUDecomposition.kt

@@ -1,7 +1,7 @@
 package scientifik.kmath.linear
 
-import scientifik.kmath.operations.DoubleField
 import scientifik.kmath.operations.Field
+import scientifik.kmath.operations.RealField
 import scientifik.kmath.structures.*
 import kotlin.math.absoluteValue
 
@@ -184,7 +184,7 @@ abstract class LUDecomposition<T : Comparable<T>, F : Field<T>>(val matrix: Matr
 }
 
 class RealLUDecomposition(matrix: RealMatrix, private val singularityThreshold: Double = DEFAULT_TOO_SMALL) :
-    LUDecomposition<Double, DoubleField>(matrix) {
+    LUDecomposition<Double, RealField>(matrix) {
     override fun isSingular(value: Double): Boolean {
         return value.absoluteValue < singularityThreshold
     }
@@ -197,9 +197,9 @@ class RealLUDecomposition(matrix: RealMatrix, private val singularityThreshold:
 
 
 /** Specialized solver.  */
-object RealLUSolver : LinearSolver<Double, DoubleField> {
+object RealLUSolver : LinearSolver<Double, RealField> {
 
-    fun decompose(mat: Matrix<Double, DoubleField>, threshold: Double = 1e-11): RealLUDecomposition =
+    fun decompose(mat: Matrix<Double, RealField>, threshold: Double = 1e-11): RealLUDecomposition =
         RealLUDecomposition(mat, threshold)
 
     override fun solve(a: RealMatrix, b: RealMatrix): RealMatrix {

kmath-core/src/commonMain/kotlin/scientifik/kmath/linear/LinearAlgrebra.kt

@@ -1,8 +1,8 @@
 package scientifik.kmath.linear
 
-import scientifik.kmath.operations.DoubleField
 import scientifik.kmath.operations.Field
 import scientifik.kmath.operations.Norm
+import scientifik.kmath.operations.RealField
 import scientifik.kmath.operations.Ring
 import scientifik.kmath.structures.asSequence
 import scientifik.kmath.structures.boxingBuffer
@@ -61,5 +61,5 @@ object VectorL2Norm : Norm<Vector<out Number, *>, Double> {
     }
 }
 
-typealias RealVector = Vector<Double, DoubleField>
+typealias RealVector = Vector<Double, RealField>
-typealias RealMatrix = Matrix<Double, DoubleField>
+typealias RealMatrix = Matrix<Double, RealField>

kmath-core/src/commonMain/kotlin/scientifik/kmath/linear/Matrix.kt

@@ -1,6 +1,6 @@
 package scientifik.kmath.linear
 
-import scientifik.kmath.operations.DoubleField
+import scientifik.kmath.operations.RealField
 import scientifik.kmath.operations.Ring
 import scientifik.kmath.operations.Space
 import scientifik.kmath.operations.SpaceElement
@@ -42,8 +42,8 @@ interface MatrixSpace<T : Any, R : Ring<T>> : Space<Matrix<T, R>> {
         /**
          * Non-boxing double matrix
          */
-        fun real(rows: Int, columns: Int): MatrixSpace<Double, DoubleField> =
+        fun real(rows: Int, columns: Int): MatrixSpace<Double, RealField> =
-            StructureMatrixSpace(rows, columns, DoubleField, DoubleBufferFactory)
+            StructureMatrixSpace(rows, columns, RealField, DoubleBufferFactory)
 
         /**
          * A structured matrix with custom buffer

kmath-core/src/commonMain/kotlin/scientifik/kmath/linear/Vector.kt

@@ -1,6 +1,6 @@
 package scientifik.kmath.linear
 
-import scientifik.kmath.operations.DoubleField
+import scientifik.kmath.operations.RealField
 import scientifik.kmath.operations.Space
 import scientifik.kmath.operations.SpaceElement
 import scientifik.kmath.structures.*
@@ -34,13 +34,13 @@ interface VectorSpace<T : Any, S : Space<T>> : Space<Point<T>> {
 
     companion object {
 
-        private val realSpaceCache = HashMap<Int, BufferVectorSpace<Double, DoubleField>>()
+        private val realSpaceCache = HashMap<Int, BufferVectorSpace<Double, RealField>>()
 
         /**
          * Non-boxing double vector space
          */
-        fun real(size: Int): BufferVectorSpace<Double, DoubleField> {
+        fun real(size: Int): BufferVectorSpace<Double, RealField> {
-            return realSpaceCache.getOrPut(size) { BufferVectorSpace(size, DoubleField, DoubleBufferFactory) }
+            return realSpaceCache.getOrPut(size) { BufferVectorSpace(size, RealField, DoubleBufferFactory) }
         }
 
         /**
@@ -79,10 +79,10 @@ interface Vector<T : Any, S : Space<T>> : SpaceElement<Point<T>, Vector<T, S>, V
         fun <T : Any, S : Space<T>> generic(size: Int, field: S, initializer: (Int) -> T): Vector<T, S> =
             VectorSpace.buffered(size, field).produceElement(initializer)
 
-        fun real(size: Int, initializer: (Int) -> Double): Vector<Double, DoubleField> =
+        fun real(size: Int, initializer: (Int) -> Double): Vector<Double, RealField> =
             VectorSpace.real(size).produceElement(initializer)
 
-        fun ofReal(vararg elements: Double): Vector<Double, DoubleField> =
+        fun ofReal(vararg elements: Double): Vector<Double, RealField> =
             VectorSpace.real(elements.size).produceElement { elements[it] }
 
     }

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

@@ -17,12 +17,12 @@ interface ExtendedField<T : Any> :
  *
  * TODO inline does not work due to compiler bug. Waiting for fix for KT-27586
  */
-inline class Real(val value: Double) : FieldElement<Double, Real, DoubleField> {
+inline class Real(val value: Double) : FieldElement<Double, Real, RealField> {
     override fun unwrap(): Double = value
 
     override fun Double.wrap(): Real = Real(value)
 
-    override val context get() = DoubleField
+    override val context get() = RealField
 
     companion object {
 
@@ -32,7 +32,7 @@ inline class Real(val value: Double) : FieldElement<Double, Real, DoubleField> {
 /**
  * A field for double without boxing. Does not produce appropriate field element
  */
-object DoubleField : AbstractField<Double>(),ExtendedField<Double>, Norm<Double, Double> {
+object RealField : AbstractField<Double>(),ExtendedField<Double>, Norm<Double, Double> {
     override val zero: Double = 0.0
     override fun add(a: Double, b: Double): Double = a + b
     override fun multiply(a: Double, @Suppress("PARAMETER_NAME_CHANGED_ON_OVERRIDE") b: Double): Double = a * b

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

@@ -5,18 +5,19 @@ import scientifik.kmath.operations.FieldElement
 
 abstract class StridedNDField<T, F : Field<T>>(shape: IntArray, elementField: F) :
     AbstractNDField<T, F, NDBuffer<T>>(shape, elementField) {
-
-    abstract val bufferFactory: BufferFactory<T>
     val strides = DefaultStrides(shape)
+
+    abstract fun buildBuffer(size: Int, initializer: (Int) -> T): Buffer<T>
 }
 
 
 class BufferNDField<T, F : Field<T>>(
     shape: IntArray,
     elementField: F,
-    override val bufferFactory: BufferFactory<T>
+    val bufferFactory: BufferFactory<T>
-) :
+) : StridedNDField<T, F>(shape, elementField) {
-    StridedNDField<T, F>(shape, elementField) {
+
+    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")
@@ -32,22 +33,25 @@ class BufferNDField<T, F : Field<T>>(
             bufferFactory(strides.linearSize) { offset -> elementField.initializer(strides.index(offset)) })
 
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun NDBuffer<T>.map(crossinline transform: F.(T) -> T): BufferNDElement<T, F> {
+    override inline fun map(arg: NDBuffer<T>, crossinline transform: F.(T) -> T): BufferNDElement<T, F> {
-        check(this)
+        check(arg)
         return BufferNDElement(
-            this@BufferNDField,
+            this,
-            bufferFactory(strides.linearSize) { offset -> elementField.transform(buffer[offset]) })
+            bufferFactory(arg.strides.linearSize) { offset -> elementField.transform(arg.buffer[offset]) })
     }
 
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun NDBuffer<T>.mapIndexed(crossinline transform: F.(index: IntArray, T) -> T): BufferNDElement<T, F> {
+    override inline fun mapIndexed(
-        check(this)
+        arg: NDBuffer<T>,
+        crossinline transform: F.(index: IntArray, T) -> T
+    ): BufferNDElement<T, F> {
+        check(arg)
         return BufferNDElement(
-            this@BufferNDField,
+            this,
-            bufferFactory(strides.linearSize) { offset ->
+            bufferFactory(arg.strides.linearSize) { offset ->
                 elementField.transform(
-                    strides.index(offset),
+                    arg.strides.index(offset),
-                    buffer[offset]
+                    arg.buffer[offset]
                 )
             })
     }
@@ -105,11 +109,11 @@ class BufferNDElement<T, F : Field<T>>(override val context: StridedNDField<T, F
     override fun elements(): Sequence<Pair<IntArray, T>> =
         strides.indices().map { it to get(it) }
 
-    override fun map(action: F.(T) -> T): BufferNDElement<T, F> =
+    override fun map(action: F.(T) -> T) =
-        context.run { map(action) }
+        context.run { map(this@BufferNDElement, action) }.wrap()
 
-    override fun mapIndexed(transform: F.(index: IntArray, T) -> T): BufferNDElement<T, F> =
+    override fun mapIndexed(transform: F.(index: IntArray, T) -> T) =
-        context.run { mapIndexed(transform) }
+        context.run { mapIndexed(this@BufferNDElement, transform) }.wrap()
 }
 
 /**

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

@@ -6,7 +6,7 @@ import scientifik.kmath.operations.PowerOperations
 import scientifik.kmath.operations.TrigonometricOperations
 
 
-interface ExtendedNDField<T : Any, F : ExtendedField<T>, N : NDStructure<out T>> :
+interface ExtendedNDField<T : Any, F : ExtendedField<T>, N : NDStructure<T>> :
     NDField<T, F, N>,
     TrigonometricOperations<N>,
     PowerOperations<N>,
@@ -16,7 +16,7 @@ interface ExtendedNDField<T : Any, F : ExtendedField<T>, N : NDStructure<out T>>
 /**
  * NDField that supports [ExtendedField] operations on its elements
  */
-class ExtendedNDFieldWrapper<T : Any, F : ExtendedField<T>, N : NDStructure<out T>>(private val ndField: NDField<T, F, N>) :
+class ExtendedNDFieldWrapper<T : Any, F : ExtendedField<T>, N : NDStructure<T>>(private val ndField: NDField<T, F, N>) :
     ExtendedNDField<T, F, N>, NDField<T, F, N> by ndField {
 
     override val shape: IntArray get() = ndField.shape

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

@@ -1,8 +1,8 @@
 package scientifik.kmath.structures
 
-import scientifik.kmath.operations.DoubleField
 import scientifik.kmath.operations.Field
 import scientifik.kmath.operations.FieldElement
+import scientifik.kmath.operations.RealField
 
 
 interface NDElement<T, F : Field<T>> : NDStructure<T> {
@@ -17,7 +17,7 @@ object NDElements {
     /**
      * Create a optimized NDArray of doubles
      */
-    fun real(shape: IntArray, initializer: DoubleField.(IntArray) -> Double = { 0.0 }) =
+    fun real(shape: IntArray, initializer: RealField.(IntArray) -> Double = { 0.0 }) =
         NDField.real(shape).produce(initializer)
 
 

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

@@ -23,9 +23,9 @@ interface NDField<T, F : Field<T>, N : NDStructure<T>> : Field<N> {
 
     fun produce(initializer: F.(IntArray) -> T): N
 
-    fun N.map(transform: F.(T) -> T): N
+    fun map(arg: N, transform: F.(T) -> T): N
 
-    fun N.mapIndexed(transform: F.(index: IntArray, T) -> T): N
+    fun mapIndexed(arg: N, transform: F.(index: IntArray, T) -> T): N
 
     fun combine(a: N, b: N, transform: F.(T, T) -> T): N
 
@@ -87,11 +87,11 @@ abstract class AbstractNDField<T, F : Field<T>, N : NDStructure<T>>(
 
     override val one: N by lazy { produce { one } }
 
-    final override operator fun Function1<T, T>.invoke(structure: N) = structure.map { value -> this@invoke(value) }
+    final override operator fun Function1<T, T>.invoke(structure: N) = map(structure) { value -> this@invoke(value) }
-    final override operator fun N.plus(arg: T) = this.map { value -> elementField.run { arg + value } }
+    final override operator fun N.plus(arg: T) = map(this) { value -> elementField.run { arg + value } }
-    final override operator fun N.minus(arg: T) = this.map { value -> elementField.run { arg - value } }
+    final override operator fun N.minus(arg: T) = map(this) { value -> elementField.run { arg - value } }
-    final override operator fun N.times(arg: T) = this.map { value -> elementField.run { arg * value } }
+    final override operator fun N.times(arg: T) = map(this) { value -> elementField.run { arg * value } }
-    final override operator fun N.div(arg: T) = this.map { value -> elementField.run { arg / value } }
+    final override operator fun N.div(arg: T) = map(this) { value -> elementField.run { arg / value } }
 
     final override operator fun T.plus(arg: N) = arg + this
     final override operator fun T.minus(arg: N) = arg - this
@@ -109,7 +109,7 @@ abstract class AbstractNDField<T, F : Field<T>, N : NDStructure<T>>(
      * Multiply all elements by cinstant
      */
     override fun multiply(a: N, k: Double): N =
-        a.map { it * k }
+        map(a) { it * k }
 
 
     /**
@@ -140,17 +140,16 @@ class GenericNDField<T : Any, F : Field<T>>(
     shape: IntArray,
     elementField: F,
     val bufferFactory: BufferFactory<T> = ::boxingBuffer
-) :
+) : AbstractNDField<T, F, NDStructure<T>>(shape, elementField) {
-    AbstractNDField<T, F, NDStructure<T>>(shape, elementField) {
 
     override fun produce(initializer: F.(IntArray) -> T): NDStructure<T> =
         ndStructure(shape, bufferFactory) { elementField.initializer(it) }
 
-    override fun NDStructure<T>.map(transform: F.(T) -> T): NDStructure<T> =
+    override fun map(arg: NDStructure<T>, transform: F.(T) -> T): NDStructure<T> =
-        produce { index -> transform(get(index)) }
+        produce { index -> transform(arg.get(index)) }
 
-    override fun NDStructure<T>.mapIndexed(transform: F.(index: IntArray, T) -> T): NDStructure<T> =
+    override fun mapIndexed(arg: NDStructure<T>, transform: F.(index: IntArray, T) -> T): NDStructure<T> =
-        produce { index -> transform(index, get(index)) }
+        produce { index -> transform(index, arg.get(index)) }
 
     override fun combine(a: NDStructure<T>, b: NDStructure<T>, transform: F.(T, T) -> T): NDStructure<T> =
         produce { index -> transform(a[index], b[index]) }

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

@@ -1,41 +1,47 @@
 package scientifik.kmath.structures
 
-import scientifik.kmath.operations.DoubleField
+import scientifik.kmath.operations.RealField
 
-typealias RealNDElement = BufferNDElement<Double, DoubleField>
+typealias RealNDElement = BufferNDElement<Double, RealField>
 
 class RealNDField(shape: IntArray) :
-    StridedNDField<Double, DoubleField>(shape, DoubleField),
+    StridedNDField<Double, RealField>(shape, RealField),
-    ExtendedNDField<Double, DoubleField, NDBuffer<Double>> {
+    ExtendedNDField<Double, RealField, NDBuffer<Double>> {
 
-    override val bufferFactory: BufferFactory<Double>
+    override fun buildBuffer(size: Int, initializer: (Int) -> Double): Buffer<Double> =
-        get() = DoubleBufferFactory
+        DoubleBuffer(DoubleArray(size, initializer))
 
     /**
-     * Inline map an NDStructure to
+     * Inline transform an NDStructure to
      */
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun NDBuffer<Double>.map(crossinline transform: DoubleField.(Double) -> Double): RealNDElement {
+    override inline fun map(
-        check(this)
+        arg: NDBuffer<Double>,
-        val array = DoubleArray(strides.linearSize) { offset -> DoubleField.transform(buffer[offset]) }
+        crossinline transform: RealField.(Double) -> Double
-        return BufferNDElement(this@RealNDField, DoubleBuffer(array))
+    ): RealNDElement {
+        check(arg)
+        val array = DoubleArray(arg.strides.linearSize) { offset -> RealField.transform(arg.buffer[offset]) }
+        return BufferNDElement(this, DoubleBuffer(array))
     }
 
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun produce(crossinline initializer: DoubleField.(IntArray) -> Double): RealNDElement {
+    override inline fun produce(crossinline initializer: RealField.(IntArray) -> Double): RealNDElement {
         val array = DoubleArray(strides.linearSize) { offset -> elementField.initializer(strides.index(offset)) }
         return BufferNDElement(this, DoubleBuffer(array))
     }
 
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun NDBuffer<Double>.mapIndexed(crossinline transform: DoubleField.(index: IntArray, Double) -> Double): BufferNDElement<Double, DoubleField> {
+    override inline fun mapIndexed(
-        check(this)
+        arg: NDBuffer<Double>,
+        crossinline transform: RealField.(index: IntArray, Double) -> Double
+    ): BufferNDElement<Double, RealField> {
+        check(arg)
         return BufferNDElement(
-            this@RealNDField,
+            this,
-            bufferFactory(strides.linearSize) { offset ->
+            buildBuffer(arg.strides.linearSize) { offset ->
                 elementField.transform(
-                    strides.index(offset),
+                    arg.strides.index(offset),
-                    buffer[offset]
+                    arg.buffer[offset]
                 )
             })
     }
@@ -44,23 +50,23 @@ class RealNDField(shape: IntArray) :
     override inline fun combine(
         a: NDBuffer<Double>,
         b: NDBuffer<Double>,
-        crossinline transform: DoubleField.(Double, Double) -> Double
+        crossinline transform: RealField.(Double, Double) -> Double
-    ): BufferNDElement<Double, DoubleField> {
+    ): BufferNDElement<Double, RealField> {
         check(a, b)
         return BufferNDElement(
             this,
-            bufferFactory(strides.linearSize) { offset -> elementField.transform(a.buffer[offset], b.buffer[offset]) })
+            buildBuffer(strides.linearSize) { offset -> elementField.transform(a.buffer[offset], b.buffer[offset]) })
     }
 
-    override fun power(arg: NDBuffer<Double>, pow: Double) = arg.map { power(it, pow) }
+    override fun power(arg: NDBuffer<Double>, pow: Double) = map(arg) { power(it, pow) }
 
-    override fun exp(arg: NDBuffer<Double>) = arg.map { exp(it) }
+    override fun exp(arg: NDBuffer<Double>) = map(arg) { exp(it) }
 
-    override fun ln(arg: NDBuffer<Double>) = arg.map { ln(it) }
+    override fun ln(arg: NDBuffer<Double>) = map(arg) { ln(it) }
 
-    override fun sin(arg: NDBuffer<Double>) = arg.map { sin(it) }
+    override fun sin(arg: NDBuffer<Double>) = map(arg) { sin(it) }
 
-    override fun cos(arg: NDBuffer<Double>) = arg.map { cos(it) }
+    override fun cos(arg: NDBuffer<Double>) = map(arg) { cos(it) }
 //
 //    override fun NDBuffer<Double>.times(k: Number) = mapInline { value -> value * k.toDouble() }
 //
@@ -75,7 +81,7 @@ class RealNDField(shape: IntArray) :
 /**
  * Fast element production using function inlining
  */
-inline fun StridedNDField<Double, DoubleField>.produceInline(crossinline initializer: DoubleField.(Int) -> Double): RealNDElement {
+inline fun StridedNDField<Double, RealField>.produceInline(crossinline initializer: RealField.(Int) -> Double): RealNDElement {
     val array = DoubleArray(strides.linearSize) { offset -> elementField.initializer(offset) }
     return BufferNDElement(this, DoubleBuffer(array))
 }

kmath-core/src/commonTest/kotlin/scientifik/kmath/expressions/ExpressionFieldTest.kt

@@ -2,14 +2,14 @@ package scientifik.kmath.expressions
 
 import scientifik.kmath.operations.Complex
 import scientifik.kmath.operations.ComplexField
-import scientifik.kmath.operations.DoubleField
+import scientifik.kmath.operations.RealField
 import kotlin.test.Test
 import kotlin.test.assertEquals
 
 class ExpressionFieldTest {
     @Test
     fun testExpression() {
-        val context = ExpressionField(DoubleField)
+        val context = ExpressionField(RealField)
         val expression = with(context) {
             val x = variable("x", 2.0)
             x * x + 2 * x + 1.0
@@ -36,7 +36,7 @@ class ExpressionFieldTest {
             return x * x + 2 * x + one
         }
 
-        val expression = ExpressionField(DoubleField).expression()
+        val expression = ExpressionField(RealField).expression()
         assertEquals(expression("x" to 1.0), 4.0)
     }
 
@@ -47,7 +47,7 @@ class ExpressionFieldTest {
             x * x + 2 * x + 1.0
         }
 
-        val expression = ExpressionField(DoubleField).expressionBuilder()
+        val expression = ExpressionField(RealField).expressionBuilder()
         assertEquals(expression("x" to 1.0), 4.0)
     }
 }

kmath-core/src/commonTest/kotlin/scientifik/kmath/operations/RealFieldTest.kt

@@ -6,7 +6,7 @@ import kotlin.test.assertEquals
 class RealFieldTest {
     @Test
     fun testSqrt() {
-        val sqrt = with(DoubleField) {
+        val sqrt = with(RealField) {
             sqrt(25 * one)
         }
         assertEquals(5.0, sqrt)

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

@@ -2,54 +2,80 @@ package scientifik.kmath.structures
 
 import kotlinx.coroutines.*
 import scientifik.kmath.operations.Field
+import scientifik.kmath.operations.FieldElement
 
 class LazyNDField<T, F : Field<T>>(shape: IntArray, field: F, val scope: CoroutineScope = GlobalScope) :
-    BufferNDField<T, F>(shape, field, ::boxingBuffer) {
+    AbstractNDField<T, F, NDStructure<T>>(shape, field) {
 
-    override fun add(a: NDStructure<T>, b: NDStructure<T>): NDElements<T, F> {
+    override val zero by lazy { produce { zero } }
-        return LazyNDStructure(this) { index ->
+
-            val aDeferred = a.deferred(index)
+    override val one by lazy { produce { one } }
-            val bDeferred = b.deferred(index)
+
-            aDeferred.await() + bDeferred.await()
+    override fun produce(initializer: F.(IntArray) -> T) =
+        LazyNDStructure(this) { elementField.initializer(it) }
+
+    override fun mapIndexed(
+        arg: NDStructure<T>,
+        transform: F.(index: IntArray, T) -> T
+    ): LazyNDStructure<T, F> {
+        check(arg)
+        return if (arg is LazyNDStructure<T, *>) {
+            LazyNDStructure(this) { index ->
+                this.elementField.transform(index, arg.function(index))
+            }
+        } else {
+            LazyNDStructure(this) { elementField.transform(it, arg.await(it)) }
         }
     }
 
-    override fun multiply(a: NDStructure<T>, k: Double): NDElements<T, F> {
+    override fun map(arg: NDStructure<T>, transform: F.(T) -> T) =
-        return LazyNDStructure(this) { index -> a.await(index) * k }
+        mapIndexed(arg) { _, t -> transform(t) }
-    }
 
-    override fun multiply(a: NDStructure<T>, b: NDStructure<T>): NDElements<T, F> {
+    override fun combine(a: NDStructure<T>, b: NDStructure<T>, transform: F.(T, T) -> T): LazyNDStructure<T, F> {
-        return LazyNDStructure(this) { index ->
+        check(a, b)
-            val aDeferred = a.deferred(index)
+        return if (a is LazyNDStructure<T, *> && b is LazyNDStructure<T, *>) {
-            val bDeferred = b.deferred(index)
+            LazyNDStructure(this@LazyNDField) { index ->
-            aDeferred.await() * bDeferred.await()
+                elementField.transform(
+                    a.function(index),
+                    b.function(index)
+                )
+            }
+        } else {
+            LazyNDStructure(this@LazyNDField) { elementField.transform(a.await(it), b.await(it)) }
         }
     }
 
-    override fun divide(a: NDStructure<T>, b: NDStructure<T>): NDElements<T, F> {
+    fun NDStructure<T>.lazy(): LazyNDStructure<T, F> {
-        return LazyNDStructure(this) { index ->
+        check(this)
-            val aDeferred = a.deferred(index)
+        return if (this is LazyNDStructure<T, *>) {
-            val bDeferred = b.deferred(index)
+            LazyNDStructure(this@LazyNDField, function)
-            aDeferred.await() / bDeferred.await()
+        } else {
+            LazyNDStructure(this@LazyNDField) { get(it) }
         }
     }
 }
 
 class LazyNDStructure<T, F : Field<T>>(
     override val context: LazyNDField<T, F>,
-    val function: suspend F.(IntArray) -> T
+    val function: suspend (IntArray) -> T
-) : NDElements<T, F>, NDStructure<T> {
+) : FieldElement<NDStructure<T>, LazyNDStructure<T, F>, LazyNDField<T, F>>, NDElement<T, F> {
-    override val self: NDElements<T, F> get() = this
+
+
+    override fun unwrap(): NDStructure<T> = this
+
+    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.elementField
+
+    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>>()
 
     fun deferred(index: IntArray) = cache.getOrPut(index) {
         context.scope.async(context = Dispatchers.Math) {
-            function.invoke(
+            function(index)
-                context.elementField,
-                index
-            )
         }
     }
 
@@ -61,7 +87,10 @@ class LazyNDStructure<T, F : Field<T>>(
 
     override fun elements(): Sequence<Pair<IntArray, T>> {
         val strides = DefaultStrides(shape)
-        return strides.indices().map { index -> index to runBlocking { await(index) } }
+        val res = runBlocking {
+            strides.indices().toList().map { index -> index to await(index) }
+        }
+        return res.asSequence()
     }
 }
 
@@ -71,21 +100,11 @@ fun <T> NDStructure<T>.deferred(index: IntArray) =
 suspend fun <T> NDStructure<T>.await(index: IntArray) =
     if (this is LazyNDStructure<T, *>) this.await(index) else get(index)
 
-fun <T, F : Field<T>> NDElements<T, F>.lazy(scope: CoroutineScope = GlobalScope): LazyNDStructure<T, F> {
-    return if (this is LazyNDStructure<T, F>) {
-        this
-    } else {
-        val context = LazyNDField(context.shape, context.field, scope)
-        LazyNDStructure(context) { get(it) }
-    }
-}
 
-inline fun <T, F : Field<T>> LazyNDStructure<T, F>.mapIndexed(crossinline action: suspend F.(IntArray, T) -> T) =
+fun <T : Any, F : Field<T>> NDField.Companion.lazy(shape: IntArray, field: F, scope: CoroutineScope = GlobalScope) =
-    LazyNDStructure(context) { index ->
+    LazyNDField(shape, field, scope)
-        action.invoke(this, index, await(index))
-    }
 
-inline fun <T, F : Field<T>> LazyNDStructure<T, F>.map(crossinline action: suspend F.(T) -> T) =
+fun <T, F : Field<T>> NDStructure<T>.lazy(field: F, scope: CoroutineScope = GlobalScope): LazyNDStructure<T, F> {
-    LazyNDStructure(context) { index ->
+    val context: LazyNDField<T, F> = LazyNDField(shape, field, scope)
-        action.invoke(this, await(index))
+    return LazyNDStructure(context) { get(it) }
 }

kmath-coroutines/src/commonTest/kotlin/scientifik/kmath/structures/LazyNDFieldTest.kt

@@ -10,7 +10,7 @@ class LazyNDFieldTest {
     fun testLazyStructure() {
         var counter = 0
         val regularStructure = NDField.generic(intArrayOf(2, 2, 2), IntField).produce { it[0] + it[1] - it[2] }
-        val result = (regularStructure.lazy() + 2).transform {
+        val result = (regularStructure.lazy(IntField) + 2).map {
             counter++
             it * it
         }