Matrix performance optimization

examples/src/benchmarks/kotlin/kscience/kmath/benchmarks/DotBenchmark.kt

@@ -31,38 +31,35 @@ class DotBenchmark {
     }
 
     @Benchmark
-    fun commonsMathMultiplication() {
+    fun cmDot() {
         CMMatrixContext {
             cmMatrix1 dot cmMatrix2
         }
     }
 
     @Benchmark
-    fun ejmlMultiplication() {
+    fun ejmlDot() {
         EjmlMatrixContext {
             ejmlMatrix1 dot ejmlMatrix2
         }
     }
 
     @Benchmark
-    fun ejmlMultiplicationwithConversion() {
+    fun ejmlDotWithConversion() {
         EjmlMatrixContext {
-            val ejmlMatrix1 = matrix1.toEjml()
+            matrix1 dot matrix2
-            val ejmlMatrix2 = matrix2.toEjml()
-
-            ejmlMatrix1 dot ejmlMatrix2
         }
     }
 
     @Benchmark
-    fun bufferedMultiplication() {
+    fun bufferedDot() {
         BufferMatrixContext(RealField, Buffer.Companion::real).invoke {
             matrix1 dot matrix2
         }
     }
 
     @Benchmark
-    fun realMultiplication() {
+    fun realDot() {
         RealMatrixContext {
             matrix1 dot matrix2
         }

examples/src/benchmarks/kotlin/kscience/kmath/benchmarks/LinearAlgebraBenchmark.kt

@@ -26,7 +26,9 @@ class LinearAlgebraBenchmark {
 
     @Benchmark
     fun kmathLUPInversion() {
-        MatrixContext.real.inverseWithLUP(matrix)
+        MatrixContext.real{
+            inverseWithLUP(matrix)
+        }
     }
 
     @Benchmark

examples/src/benchmarks/kotlin/kscience/kmath/benchmarks/ViktorBenchmark.kt

@@ -23,7 +23,7 @@ internal class ViktorBenchmark {
     fun automaticFieldAddition() {
         autoField {
             var res: NDStructure<Double> = one
-            repeat(n) { res += one }
+            repeat(n) { res += 1.0 }
         }
     }
 
@@ -31,7 +31,7 @@ internal class ViktorBenchmark {
     fun realFieldAddition() {
         realField {
             var res: NDStructure<Double> = one
-            repeat(n) { res += one }
+            repeat(n) { res += 1.0 }
         }
     }
 
@@ -39,7 +39,7 @@ internal class ViktorBenchmark {
     fun viktorFieldAddition() {
         viktorField {
             var res = one
-            repeat(n) { res += one }
+            repeat(n) { res += 1.0 }
         }
     }
 

examples/src/benchmarks/kotlin/kscience/kmath/benchmarks/ViktorLogBenchmark.kt

@@ -15,7 +15,7 @@ internal class ViktorLogBenchmark {
     final val n: Int = 100
 
     // automatically build context most suited for given type.
-    final val autoField: BufferedNDField<Double, RealField> = NDAlgebra.auto(RealField, dim, dim)
+    final val autoField: NDField<Double, RealField> = NDAlgebra.auto(RealField, dim, dim)
     final val realField: RealNDField = NDAlgebra.real(dim, dim)
     final val viktorField: ViktorNDField = ViktorNDField(intArrayOf(dim, dim))
 
@@ -29,6 +29,15 @@ internal class ViktorLogBenchmark {
         }
     }
 
+    @Benchmark
+    fun viktorFieldLog() {
+        viktorField {
+            val fortyTwo = produce { 42.0 }
+            var res = one
+            repeat(n) { res = ln(fortyTwo) }
+        }
+    }
+
     @Benchmark
     fun rawViktorLog() {
         val fortyTwo = F64Array.full(dim, dim, init = 42.0)

examples/src/main/kotlin/kscience/kmath/structures/ParallelRealNDField.kt

@@ -48,22 +48,20 @@ class StreamRealNDField(
         return NDBuffer(strides, array.asBuffer())
     }
 
-    override fun map(
+    override fun NDStructure<Double>.map(
-        arg: NDStructure<Double>,
         transform: RealField.(Double) -> Double,
     ): NDBuffer<Double> {
-        val array = Arrays.stream(arg.buffer.array).parallel().map { RealField.transform(it) }.toArray()
+        val array = Arrays.stream(buffer.array).parallel().map { RealField.transform(it) }.toArray()
         return NDBuffer(strides, array.asBuffer())
     }
 
-    override fun mapIndexed(
+    override fun NDStructure<Double>.mapIndexed(
-        arg: NDStructure<Double>,
         transform: RealField.(index: IntArray, Double) -> Double,
     ): NDBuffer<Double> {
         val array = IntStream.range(0, strides.linearSize).parallel().mapToDouble { offset ->
             RealField.transform(
                 strides.index(offset),
-                arg.buffer.array[offset]
+                buffer.array[offset]
             )
         }.toArray()
 
@@ -81,25 +79,25 @@ class StreamRealNDField(
         return NDBuffer(strides, array.asBuffer())
     }
 
-    override fun power(arg: NDStructure<Double>, pow: Number): NDBuffer<Double> = map(arg) { power(it, pow) }
+    override fun power(arg: NDStructure<Double>, pow: Number): NDBuffer<Double> = arg.map() { power(it, pow) }
 
-    override fun exp(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { exp(it) }
+    override fun exp(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { exp(it) }
 
-    override fun ln(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { ln(it) }
+    override fun ln(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { ln(it) }
 
-    override fun sin(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { sin(it) }
+    override fun sin(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { sin(it) }
-    override fun cos(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { cos(it) }
+    override fun cos(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { cos(it) }
-    override fun tan(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { tan(it) }
+    override fun tan(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { tan(it) }
-    override fun asin(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { asin(it) }
+    override fun asin(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { asin(it) }
-    override fun acos(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { acos(it) }
+    override fun acos(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { acos(it) }
-    override fun atan(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { atan(it) }
+    override fun atan(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { atan(it) }
 
-    override fun sinh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { sinh(it) }
+    override fun sinh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { sinh(it) }
-    override fun cosh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { cosh(it) }
+    override fun cosh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { cosh(it) }
-    override fun tanh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { tanh(it) }
+    override fun tanh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { tanh(it) }
-    override fun asinh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { asinh(it) }
+    override fun asinh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { asinh(it) }
-    override fun acosh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { acosh(it) }
+    override fun acosh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { acosh(it) }
-    override fun atanh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { atanh(it) }
+    override fun atanh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map() { atanh(it) }
 }
 
 fun NDAlgebra.Companion.realWithStream(vararg shape: Int): StreamRealNDField = StreamRealNDField(shape)

kmath-core/src/commonMain/kotlin/kscience/kmath/linear/BufferMatrix.kt

@@ -3,6 +3,7 @@ package kscience.kmath.linear
 import kscience.kmath.nd.NDStructure
 import kscience.kmath.nd.Structure2D
 import kscience.kmath.operations.Ring
+import kscience.kmath.operations.invoke
 import kscience.kmath.structures.Buffer
 import kscience.kmath.structures.BufferFactory
 import kscience.kmath.structures.asSequence
@@ -28,6 +29,62 @@ public class BufferMatrixContext<T : Any, R : Ring<T>>(
 
     public override fun point(size: Int, initializer: (Int) -> T): Point<T> = bufferFactory(size, initializer)
 
+    private fun Matrix<T>.toBufferMatrix(): BufferMatrix<T> = if (this is BufferMatrix) this else {
+        produce(rowNum, colNum) { i, j -> get(i, j) }
+    }
+
+    public fun one(rows: Int, columns: Int): Matrix<Double> = VirtualMatrix(rows, columns) { i, j ->
+        if (i == j) 1.0 else 0.0
+    } + DiagonalFeature
+
+    public override infix fun Matrix<T>.dot(other: Matrix<T>): BufferMatrix<T> {
+        require(colNum == other.rowNum) { "Matrix dot operation dimension mismatch: ($rowNum, $colNum) x (${other.rowNum}, ${other.colNum})" }
+        val bufferMatrix = toBufferMatrix()
+        val otherBufferMatrix = other.toBufferMatrix()
+        return elementContext {
+            produce(rowNum, other.colNum) { i, j ->
+                var res = one
+                for (l in 0 until colNum) {
+                    res += bufferMatrix[i, l] * otherBufferMatrix[l, j]
+                }
+                res
+            }
+        }
+    }
+
+    public override infix fun Matrix<T>.dot(vector: Point<T>): Point<T> {
+        require(colNum == vector.size) { "Matrix dot vector operation dimension mismatch: ($rowNum, $colNum) x (${vector.size})" }
+        val bufferMatrix = toBufferMatrix()
+        return elementContext {
+            bufferFactory(rowNum) { i ->
+                var res = one
+                for (j in 0 until colNum) {
+                    res += bufferMatrix[i, j] * vector[j]
+                }
+                res
+            }
+        }
+    }
+
+    override fun add(a: Matrix<T>, b: Matrix<T>): BufferMatrix<T> {
+        require(a.rowNum == b.rowNum) { "Row number mismatch in matrix addition. Left side: ${a.rowNum}, right side: ${b.rowNum}" }
+        require(a.colNum == b.colNum) { "Column number mismatch in matrix addition. Left side: ${a.colNum}, right side: ${b.colNum}" }
+        val aBufferMatrix = a.toBufferMatrix()
+        val bBufferMatrix = b.toBufferMatrix()
+        return elementContext {
+            produce(a.rowNum, a.colNum) { i, j ->
+                aBufferMatrix[i, j] + bBufferMatrix[i, j]
+            }
+        }
+    }
+
+    override fun multiply(a: Matrix<T>, k: Number): BufferMatrix<T> {
+        val aBufferMatrix = a.toBufferMatrix()
+        return elementContext {
+            produce(a.rowNum, a.colNum) { i, j -> aBufferMatrix[i, j] * k.toDouble() }
+        }
+    }
+
     public companion object
 }
 

kmath-core/src/commonMain/kotlin/kscience/kmath/linear/MatrixWrapper.kt

@@ -60,12 +60,6 @@ public operator fun <T : Any> Matrix<T>.plus(newFeatures: Collection<MatrixFeatu
         MatrixWrapper(this, newFeatures.toSet())
     }
 
-public inline fun Structure2D.Companion.real(
-    rows: Int,
-    columns: Int,
-    initializer: (Int, Int) -> Double,
-): BufferMatrix<Double> = MatrixContext.real.produce(rows, columns, initializer)
-
 /**
  * Build a square matrix from given elements.
  */

kmath-core/src/commonMain/kotlin/kscience/kmath/linear/RealMatrixContext.kt

@@ -2,10 +2,9 @@ package kscience.kmath.linear
 
 import kscience.kmath.structures.RealBuffer
 
-@Suppress("OVERRIDE_BY_INLINE")
 public object RealMatrixContext : MatrixContext<Double, BufferMatrix<Double>> {
 
-    public override inline fun produce(
+    public override fun produce(
         rows: Int,
         columns: Int,
         initializer: (i: Int, j: Int) -> Double,
@@ -14,7 +13,7 @@ public object RealMatrixContext : MatrixContext<Double, BufferMatrix<Double>> {
         return BufferMatrix(rows, columns, buffer)
     }
 
-    private fun Matrix<Double>.wrap(): BufferMatrix<Double> = if (this is BufferMatrix) this else {
+    public fun Matrix<Double>.toBufferMatrix(): BufferMatrix<Double> = if (this is BufferMatrix) this else {
         produce(rowNum, colNum) { i, j -> get(i, j) }
     }
 
@@ -24,10 +23,12 @@ public object RealMatrixContext : MatrixContext<Double, BufferMatrix<Double>> {
 
     public override infix fun Matrix<Double>.dot(other: Matrix<Double>): BufferMatrix<Double> {
         require(colNum == other.rowNum) { "Matrix dot operation dimension mismatch: ($rowNum, $colNum) x (${other.rowNum}, ${other.colNum})" }
+        val bufferMatrix = toBufferMatrix()
+        val otherBufferMatrix = other.toBufferMatrix()
         return produce(rowNum, other.colNum) { i, j ->
             var res = 0.0
             for (l in 0 until colNum) {
-                res += get(i, l) * other.get(l, j)
+                res += bufferMatrix[i, l] * otherBufferMatrix[l, j]
             }
             res
         }
@@ -35,10 +36,11 @@ public object RealMatrixContext : MatrixContext<Double, BufferMatrix<Double>> {
 
     public override infix fun Matrix<Double>.dot(vector: Point<Double>): Point<Double> {
         require(colNum == vector.size) { "Matrix dot vector operation dimension mismatch: ($rowNum, $colNum) x (${vector.size})" }
+        val bufferMatrix = toBufferMatrix()
         return RealBuffer(rowNum) { i ->
             var res = 0.0
             for (j in 0 until colNum) {
-                res += get(i, j) * vector[j]
+                res += bufferMatrix[i, j] * vector[j]
             }
             res
         }
@@ -47,17 +49,23 @@ public object RealMatrixContext : MatrixContext<Double, BufferMatrix<Double>> {
     override fun add(a: Matrix<Double>, b: Matrix<Double>): BufferMatrix<Double> {
         require(a.rowNum == b.rowNum) { "Row number mismatch in matrix addition. Left side: ${a.rowNum}, right side: ${b.rowNum}" }
         require(a.colNum == b.colNum) { "Column number mismatch in matrix addition. Left side: ${a.colNum}, right side: ${b.colNum}" }
+        val aBufferMatrix = a.toBufferMatrix()
+        val bBufferMatrix = b.toBufferMatrix()
         return produce(a.rowNum, a.colNum) { i, j ->
-            a[i, j] + b[i, j]
+            aBufferMatrix[i, j] + bBufferMatrix[i, j]
         }
     }
 
-    override fun Matrix<Double>.times(value: Double): BufferMatrix<Double> =
+    override fun Matrix<Double>.times(value: Double): BufferMatrix<Double> {
-        produce(rowNum, colNum) { i, j -> get(i, j) * value }
+        val bufferMatrix = toBufferMatrix()
+        return produce(rowNum, colNum) { i, j -> bufferMatrix[i, j] * value }
+    }
 
 
-    override fun multiply(a: Matrix<Double>, k: Number): BufferMatrix<Double> =
+    override fun multiply(a: Matrix<Double>, k: Number): BufferMatrix<Double> {
-        produce(a.rowNum, a.colNum) { i, j -> a[i, j] * k.toDouble() }
+        val aBufferMatrix = a.toBufferMatrix()
+        return produce(a.rowNum, a.colNum) { i, j -> aBufferMatrix[i, j] * k.toDouble() }
+    }
 }
 
 

kmath-core/src/commonMain/kotlin/kscience/kmath/nd/BufferNDAlgebra.kt

@@ -28,24 +28,24 @@ public interface BufferNDAlgebra<T, C> : NDAlgebra<T, C> {
             else -> bufferFactory(strides.linearSize) { offset -> get(strides.index(offset)) }
         }
 
-    override fun map(arg: NDStructure<T>, transform: C.(T) -> T): NDBuffer<T> {
+    override fun NDStructure<T>.map(transform: C.(T) -> T): NDBuffer<T> {
         val buffer = bufferFactory(strides.linearSize) { offset ->
-            elementContext.transform(arg.buffer[offset])
+            elementContext.transform(buffer[offset])
         }
         return NDBuffer(strides, buffer)
     }
 
-    override fun mapIndexed(arg: NDStructure<T>, transform: C.(index: IntArray, T) -> T): NDStructure<T> {
+    override fun NDStructure<T>.mapIndexed(transform: C.(index: IntArray, T) -> T): NDBuffer<T> {
         val buffer = bufferFactory(strides.linearSize) { offset ->
             elementContext.transform(
                 strides.index(offset),
-                arg.buffer[offset]
+                buffer[offset]
             )
         }
         return NDBuffer(strides, buffer)
     }
 
-    override fun combine(a: NDStructure<T>, b: NDStructure<T>, transform: C.(T, T) -> T): NDStructure<T> {
+    override fun combine(a: NDStructure<T>, b: NDStructure<T>, transform: C.(T, T) -> T): NDBuffer<T> {
         val buffer = bufferFactory(strides.linearSize) { offset ->
             elementContext.transform(a.buffer[offset], b.buffer[offset])
         }

kmath-core/src/commonMain/kotlin/kscience/kmath/nd/ComplexNDField.kt

@@ -70,25 +70,25 @@ public class ComplexNDField(
 //        return BufferedNDFieldElement(this, buffer)
 //    }
 
-    override fun power(arg: NDStructure<Complex>, pow: Number): NDBuffer<Complex> = map(arg) { power(it, pow) }
+    override fun power(arg: NDStructure<Complex>, pow: Number): NDBuffer<Complex> = arg.map() { power(it, pow) }
 
-    override fun exp(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { exp(it) }
+    override fun exp(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { exp(it) }
 
-    override fun ln(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { ln(it) }
+    override fun ln(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { ln(it) }
 
-    override fun sin(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { sin(it) }
+    override fun sin(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { sin(it) }
-    override fun cos(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { cos(it) }
+    override fun cos(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { cos(it) }
-    override fun tan(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { tan(it) }
+    override fun tan(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { tan(it) }
-    override fun asin(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { asin(it) }
+    override fun asin(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { asin(it) }
-    override fun acos(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { acos(it) }
+    override fun acos(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { acos(it) }
-    override fun atan(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { atan(it) }
+    override fun atan(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { atan(it) }
 
-    override fun sinh(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { sinh(it) }
+    override fun sinh(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { sinh(it) }
-    override fun cosh(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { cosh(it) }
+    override fun cosh(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { cosh(it) }
-    override fun tanh(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { tanh(it) }
+    override fun tanh(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { tanh(it) }
-    override fun asinh(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { asinh(it) }
+    override fun asinh(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { asinh(it) }
-    override fun acosh(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { acosh(it) }
+    override fun acosh(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { acosh(it) }
-    override fun atanh(arg: NDStructure<Complex>): NDBuffer<Complex> = map(arg) { atanh(it) }
+    override fun atanh(arg: NDStructure<Complex>): NDBuffer<Complex> = arg.map() { atanh(it) }
 }
 
 

kmath-core/src/commonMain/kotlin/kscience/kmath/nd/NDAlgebra.kt

@@ -40,12 +40,12 @@ public interface NDAlgebra<T, C> {
     /**
      * Maps elements from one structure to another one by applying [transform] to them.
      */
-    public fun map(arg: NDStructure<T>, transform: C.(T) -> T): NDStructure<T>
+    public fun NDStructure<T>.map(transform: C.(T) -> T): NDStructure<T>
 
     /**
      * Maps elements from one structure to another one by applying [transform] to them alongside with their indices.
      */
-    public fun mapIndexed(arg: NDStructure<T>, transform: C.(index: IntArray, T) -> T): NDStructure<T>
+    public fun NDStructure<T>.mapIndexed(transform: C.(index: IntArray, T) -> T): NDStructure<T>
 
     /**
      * Combines two structures into one.
@@ -56,7 +56,7 @@ public interface NDAlgebra<T, C> {
      * Element-wise invocation of function working on [T] on a [NDStructure].
      */
     public operator fun Function1<T, T>.invoke(structure: NDStructure<T>): NDStructure<T> =
-        map(structure) { value -> this@invoke(value) }
+        structure.map() { value -> this@invoke(value) }
 
     public companion object
 }
@@ -109,7 +109,7 @@ public interface NDSpace<T, S : Space<T>> : Space<NDStructure<T>>, NDAlgebra<T,
      * @param k the multiplier.
      * @return the product.
      */
-    public override fun multiply(a: NDStructure<T>, k: Number): NDStructure<T> = map(a) { multiply(it, k) }
+    public override fun multiply(a: NDStructure<T>, k: Number): NDStructure<T> = a.map() { multiply(it, k) }
 
     // TODO move to extensions after KEEP-176
 
@@ -120,7 +120,7 @@ public interface NDSpace<T, S : Space<T>> : Space<NDStructure<T>>, NDAlgebra<T,
      * @param arg the augend.
      * @return the sum.
      */
-    public operator fun NDStructure<T>.plus(arg: T): NDStructure<T> = map(this) { value -> add(arg, value) }
+    public operator fun NDStructure<T>.plus(arg: T): NDStructure<T> = this.map() { value -> add(arg, value) }
 
     /**
      * Subtracts an element from ND structure of it.
@@ -129,7 +129,7 @@ public interface NDSpace<T, S : Space<T>> : Space<NDStructure<T>>, NDAlgebra<T,
      * @param arg the divisor.
      * @return the quotient.
      */
-    public operator fun NDStructure<T>.minus(arg: T): NDStructure<T> = map(this) { value -> add(arg, -value) }
+    public operator fun NDStructure<T>.minus(arg: T): NDStructure<T> = this.map() { value -> add(arg, -value) }
 
     /**
      * Adds an element to ND structure of it.
@@ -138,7 +138,7 @@ public interface NDSpace<T, S : Space<T>> : Space<NDStructure<T>>, NDAlgebra<T,
      * @param arg the augend.
      * @return the sum.
      */
-    public operator fun T.plus(arg: NDStructure<T>): NDStructure<T> = map(arg) { value -> add(this@plus, value) }
+    public operator fun T.plus(arg: NDStructure<T>): NDStructure<T> = arg.map() { value -> add(this@plus, value) }
 
     /**
      * Subtracts an ND structure from an element of it.
@@ -147,7 +147,7 @@ public interface NDSpace<T, S : Space<T>> : Space<NDStructure<T>>, NDAlgebra<T,
      * @param arg the divisor.
      * @return the quotient.
      */
-    public operator fun T.minus(arg: NDStructure<T>): NDStructure<T> = map(arg) { value -> add(-this@minus, value) }
+    public operator fun T.minus(arg: NDStructure<T>): NDStructure<T> = arg.map() { value -> add(-this@minus, value) }
 
     public companion object
 }
@@ -179,7 +179,7 @@ public interface NDRing<T, R : Ring<T>> : Ring<NDStructure<T>>, NDSpace<T, R> {
      * @param arg the multiplier.
      * @return the product.
      */
-    public operator fun NDStructure<T>.times(arg: T): NDStructure<T> = map(this) { value -> multiply(arg, value) }
+    public operator fun NDStructure<T>.times(arg: T): NDStructure<T> = this.map() { value -> multiply(arg, value) }
 
     /**
      * Multiplies an element by a ND structure of it.
@@ -188,7 +188,7 @@ public interface NDRing<T, R : Ring<T>> : Ring<NDStructure<T>>, NDSpace<T, R> {
      * @param arg the multiplier.
      * @return the product.
      */
-    public operator fun T.times(arg: NDStructure<T>): NDStructure<T> = map(arg) { value -> multiply(this@times, value) }
+    public operator fun T.times(arg: NDStructure<T>): NDStructure<T> = arg.map() { value -> multiply(this@times, value) }
 
     public companion object
 }
@@ -219,7 +219,7 @@ public interface NDField<T, F : Field<T>> : Field<NDStructure<T>>, NDRing<T, F>
      * @param arg the divisor.
      * @return the quotient.
      */
-    public operator fun NDStructure<T>.div(arg: T): NDStructure<T> = map(this) { value -> divide(arg, value) }
+    public operator fun NDStructure<T>.div(arg: T): NDStructure<T> = this.map() { value -> divide(arg, value) }
 
     /**
      * Divides an element by an ND structure of it.
@@ -228,7 +228,7 @@ public interface NDField<T, F : Field<T>> : Field<NDStructure<T>>, NDRing<T, F>
      * @param arg the divisor.
      * @return the quotient.
      */
-    public operator fun T.div(arg: NDStructure<T>): NDStructure<T> = map(arg) { divide(it, this@div) }
+    public operator fun T.div(arg: NDStructure<T>): NDStructure<T> = arg.map() { divide(it, this@div) }
 
 //    @ThreadLocal
 //    public companion object {

kmath-core/src/commonMain/kotlin/kscience/kmath/nd/RealNDField.kt

@@ -35,16 +35,15 @@ public class RealNDField(
         }
 
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun map(
+    override inline fun NDStructure<Double>.map(
-        arg: NDStructure<Double>,
         transform: RealField.(Double) -> Double,
     ): NDBuffer<Double> {
-        val buffer = RealBuffer(strides.linearSize) { offset -> RealField.transform(arg.buffer.array[offset]) }
+        val buffer = RealBuffer(strides.linearSize) { offset -> RealField.transform(buffer.array[offset]) }
         return NDBuffer(strides, buffer)
     }
 
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun produce(initializer: RealField.(IntArray) -> Double): NDBuffer<Double>  {
+    override inline fun produce(initializer: RealField.(IntArray) -> Double): NDBuffer<Double> {
         val array = DoubleArray(strides.linearSize) { offset ->
             val index = strides.index(offset)
             RealField.initializer(index)
@@ -53,15 +52,14 @@ public class RealNDField(
     }
 
     @Suppress("OVERRIDE_BY_INLINE")
-    override inline fun mapIndexed(
+    override inline fun NDStructure<Double>.mapIndexed(
-        arg: NDStructure<Double>,
         transform: RealField.(index: IntArray, Double) -> Double,
     ): NDBuffer<Double> = NDBuffer(
         strides,
         buffer = RealBuffer(strides.linearSize) { offset ->
             RealField.transform(
                 strides.index(offset),
-                arg.buffer.array[offset]
+                buffer.array[offset]
             )
         })
 
@@ -70,32 +68,32 @@ public class RealNDField(
         a: NDStructure<Double>,
         b: NDStructure<Double>,
         transform: RealField.(Double, Double) -> Double,
-    ): NDBuffer<Double>  {
+    ): NDBuffer<Double> {
         val buffer = RealBuffer(strides.linearSize) { offset ->
             RealField.transform(a.buffer.array[offset], b.buffer.array[offset])
         }
-        return  NDBuffer(strides, buffer)
+        return NDBuffer(strides, buffer)
     }
 
-    override fun power(arg: NDStructure<Double>, pow: Number): NDBuffer<Double> = map(arg) { power(it, pow) }
+    override fun power(arg: NDStructure<Double>, pow: Number): NDBuffer<Double> = arg.map { power(it, pow) }
 
-    override fun exp(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { exp(it) }
+    override fun exp(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { exp(it) }
 
-    override fun ln(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { ln(it) }
+    override fun ln(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { ln(it) }
 
-    override fun sin(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { sin(it) }
+    override fun sin(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { sin(it) }
-    override fun cos(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { cos(it) }
+    override fun cos(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { cos(it) }
-    override fun tan(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { tan(it) }
+    override fun tan(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { tan(it) }
-    override fun asin(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { asin(it) }
+    override fun asin(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { asin(it) }
-    override fun acos(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { acos(it) }
+    override fun acos(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { acos(it) }
-    override fun atan(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { atan(it) }
+    override fun atan(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { atan(it) }
 
-    override fun sinh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { sinh(it) }
+    override fun sinh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { sinh(it) }
-    override fun cosh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { cosh(it) }
+    override fun cosh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { cosh(it) }
-    override fun tanh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { tanh(it) }
+    override fun tanh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { tanh(it) }
-    override fun asinh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { asinh(it) }
+    override fun asinh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { asinh(it) }
-    override fun acosh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { acosh(it) }
+    override fun acosh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { acosh(it) }
-    override fun atanh(arg: NDStructure<Double>): NDBuffer<Double> = map(arg) { atanh(it) }
+    override fun atanh(arg: NDStructure<Double>): NDBuffer<Double> = arg.map { atanh(it) }
 }
 
 public fun NDAlgebra.Companion.real(vararg shape: Int): RealNDField = RealNDField(shape)

kmath-core/src/commonMain/kotlin/kscience/kmath/nd/Structure2D.kt

@@ -1,6 +1,7 @@
 package kscience.kmath.nd
 
-import kscience.kmath.linear.Matrix
+import kscience.kmath.linear.BufferMatrix
+import kscience.kmath.linear.RealMatrixContext
 import kscience.kmath.structures.Buffer
 import kscience.kmath.structures.VirtualBuffer
 
@@ -58,9 +59,9 @@ public interface Structure2D<T> : NDStructure<T> {
             rows: Int,
             columns: Int,
             crossinline init: (i: Int, j: Int) -> Double,
-        ): Matrix<Double> = NDAlgebra.real(rows, columns).produce { (i, j) ->
+        ): BufferMatrix<Double> = RealMatrixContext.produce(rows,columns) { i, j ->
             init(i, j)
-        }.as2D()
+        }
     }
 }
 

kmath-for-real/src/commonMain/kotlin/kscience/kmath/real/RealMatrix.kt

@@ -140,7 +140,7 @@ public fun RealMatrix.min(): Double? = elements().map { (_, value) -> value }.mi
 public fun RealMatrix.max(): Double? = elements().map { (_, value) -> value }.maxOrNull()
 public fun RealMatrix.average(): Double = elements().map { (_, value) -> value }.average()
 
-public inline fun RealMatrix.map(transform: (Double) -> Double): RealMatrix =
+public inline fun RealMatrix.map(crossinline transform: (Double) -> Double): RealMatrix =
     MatrixContext.real.produce(rowNum, colNum) { i, j ->
         transform(get(i, j))
     }

kmath-nd4j/src/main/kotlin/kscience.kmath.nd4j/Nd4jArrayAlgebra.kt

@@ -44,18 +44,17 @@ public interface Nd4jArrayAlgebra<T, C> : NDAlgebra<T, C> {
         return struct
     }
 
-    public override fun map(arg: NDStructure<T>, transform: C.(T) -> T): Nd4jArrayStructure<T> {
+    public override fun NDStructure<T>.map(transform: C.(T) -> T): Nd4jArrayStructure<T> {
-        val newStruct = arg.ndArray.dup().wrap()
+        val newStruct = ndArray.dup().wrap()
         newStruct.elements().forEach { (idx, value) -> newStruct[idx] = elementContext.transform(value) }
         return newStruct
     }
 
-    public override fun mapIndexed(
+    public override fun NDStructure<T>.mapIndexed(
-        arg: NDStructure<T>,
         transform: C.(index: IntArray, T) -> T,
     ): Nd4jArrayStructure<T> {
-        val new = Nd4j.create(*shape).wrap()
+        val new = Nd4j.create(*this@Nd4jArrayAlgebra.shape).wrap()
-        new.indicesIterator().forEach { idx -> new[idx] = elementContext.transform(idx, arg[idx]) }
+        new.indicesIterator().forEach { idx -> new[idx] = elementContext.transform(idx, this[idx]) }
         return new
     }
 

kmath-nd4j/src/test/kotlin/kscience/kmath/nd4j/Nd4jArrayAlgebraTest.kt

@@ -1,7 +1,7 @@
 package kscience.kmath.nd4j
 
-import org.nd4j.linalg.factory.Nd4j
 import kscience.kmath.operations.invoke
+import org.nd4j.linalg.factory.Nd4j
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.fail
@@ -20,7 +20,7 @@ internal class Nd4jArrayAlgebraTest {
 
     @Test
     fun testMap() {
-        val res = (IntNd4jArrayRing(intArrayOf(2, 2))) { map(one) { it + it * 2 } }
+        val res = (IntNd4jArrayRing(intArrayOf(2, 2))) { one.map() { it + it * 2 } }
         val expected = (Nd4j.create(2, 2) ?: fail()).asIntStructure()
         expected[intArrayOf(0, 0)] = 3
         expected[intArrayOf(0, 1)] = 3

kmath-viktor/src/main/kotlin/kscience/kmath/viktor/ViktorNDStructure.kt

@@ -1,7 +1,10 @@
 package kscience.kmath.viktor
 
+import kscience.kmath.misc.UnstableKMathAPI
 import kscience.kmath.nd.*
+import kscience.kmath.operations.ExtendedField
 import kscience.kmath.operations.RealField
+import kscience.kmath.operations.RingWithNumbers
 import org.jetbrains.bio.viktor.F64Array
 
 @Suppress("OVERRIDE_BY_INLINE", "NOTHING_TO_INLINE")
@@ -20,8 +23,10 @@ public inline class ViktorNDStructure(public val f64Buffer: F64Array) : MutableN
 
 public fun F64Array.asStructure(): ViktorNDStructure = ViktorNDStructure(this)
 
+@OptIn(UnstableKMathAPI::class)
 @Suppress("OVERRIDE_BY_INLINE", "NOTHING_TO_INLINE")
-public class ViktorNDField(public override val shape: IntArray) : NDField<Double, RealField> {
+public class ViktorNDField(public override val shape: IntArray) : NDField<Double, RealField>,
+    RingWithNumbers<NDStructure<Double>>, ExtendedField<NDStructure<Double>> {
 
     public val NDStructure<Double>.f64Buffer: F64Array
         get() = when {
@@ -50,26 +55,25 @@ public class ViktorNDField(public override val shape: IntArray) : NDField<Double
             }
         }.asStructure()
 
-    public override fun map(arg: NDStructure<Double>, transform: RealField.(Double) -> Double): ViktorNDStructure =
+    public override fun NDStructure<Double>.map(transform: RealField.(Double) -> Double): ViktorNDStructure =
-        F64Array(*shape).apply {
+        F64Array(*this@ViktorNDField.shape).apply {
             this@ViktorNDField.strides.indices().forEach { index ->
-                set(value = RealField.transform(arg[index]), indices = index)
+                set(value = RealField.transform(this@map[index]), indices = index)
             }
         }.asStructure()
 
-    public override fun mapIndexed(
+    public override fun NDStructure<Double>.mapIndexed(
-        arg: NDStructure<Double>,
+        transform: RealField.(index: IntArray, Double) -> Double,
-        transform: RealField.(index: IntArray, Double) -> Double
+    ): ViktorNDStructure = F64Array(*this@ViktorNDField.shape).apply {
-    ): ViktorNDStructure = F64Array(*shape).apply {
         this@ViktorNDField.strides.indices().forEach { index ->
-            set(value = RealField.transform(index, arg[index]), indices = index)
+            set(value = RealField.transform(index, this@mapIndexed[index]), indices = index)
         }
     }.asStructure()
 
     public override fun combine(
         a: NDStructure<Double>,
         b: NDStructure<Double>,
-        transform: RealField.(Double, Double) -> Double
+        transform: RealField.(Double, Double) -> Double,
     ): ViktorNDStructure = F64Array(*shape).apply {
         this@ViktorNDField.strides.indices().forEach { index ->
             set(value = RealField.transform(a[index], b[index]), indices = index)
@@ -93,6 +97,25 @@ public class ViktorNDField(public override val shape: IntArray) : NDField<Double
 
     public override inline fun NDStructure<Double>.plus(arg: Double): ViktorNDStructure =
         (f64Buffer.plus(arg)).asStructure()
+
+    override fun number(value: Number): ViktorNDStructure =
+        F64Array.full(init = value.toDouble(), shape = shape).asStructure()
+
+    override fun sin(arg: NDStructure<Double>): ViktorNDStructure = arg.map { sin(it) }
+
+    override fun cos(arg: NDStructure<Double>): ViktorNDStructure = arg.map { cos(it) }
+
+    override fun asin(arg: NDStructure<Double>): ViktorNDStructure = arg.map { asin(it) }
+
+    override fun acos(arg: NDStructure<Double>): ViktorNDStructure = arg.map { acos(it) }
+
+    override fun atan(arg: NDStructure<Double>): ViktorNDStructure = arg.map { atan(it) }
+
+    override fun power(arg: NDStructure<Double>, pow: Number): ViktorNDStructure = arg.map { it.pow(pow) }
+
+    override fun exp(arg: NDStructure<Double>): ViktorNDStructure = arg.f64Buffer.exp().asStructure()
+
+    override fun ln(arg: NDStructure<Double>): ViktorNDStructure = arg.f64Buffer.log().asStructure()
 }
 
 public fun ViktorNDField(vararg shape: Int): ViktorNDField = ViktorNDField(shape)