diff --git a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt index 35ee5e6b7..516653817 100644 --- a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt +++ b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt @@ -1,6 +1,7 @@ package scientifik.kmath.structures import scientifik.kmath.operations.Complex +import scientifik.kmath.operations.ComplexField import java.nio.ByteBuffer object ComplexBufferSpec : FixedSizeBufferSpec { @@ -19,8 +20,18 @@ object ComplexBufferSpec : FixedSizeBufferSpec { } /** - * Create a mutable buffer which ignores boxing + * Create a read-only/mutable buffer which ignores boxing */ -fun Complex.Companion.createBuffer(size: Int) = ObjectBuffer.create(ComplexBufferSpec, size) +fun Buffer.Companion.complex(size: Int): Buffer = + ObjectBuffer.create(ComplexBufferSpec, size) + +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) } + +fun NDElement.Companion.complex(shape: IntArray, initializer: ComplexField.(IntArray) -> Complex) = + NDField.complex(shape).produce(initializer) diff --git a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt index fceda1b25..8d7bece0c 100644 --- a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt +++ b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt @@ -23,7 +23,13 @@ class ObjectBuffer(private val buffer: ByteBuffer, private val spec: Fi } companion object { - fun create(spec: FixedSizeBufferSpec, size: Int) = - ObjectBuffer(ByteBuffer.allocate(size * spec.unitSize), spec) + fun create(spec: FixedSizeBufferSpec, size: Int, initializer: ((Int) -> T)? = null) = + ObjectBuffer(ByteBuffer.allocate(size * spec.unitSize), spec).also { buffer -> + if (initializer != null) { + (0 until size).forEach { + buffer[it] = initializer(it) + } + } + } } } \ No newline at end of file diff --git a/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt b/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt index 350f06848..1a3aea10f 100644 --- a/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt +++ b/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt @@ -7,7 +7,7 @@ import kotlin.test.assertEquals class ComplexBufferSpecTest { @Test fun testComplexBuffer() { - val buffer = Complex.createBuffer(20) + val buffer = MutableBuffer.complex(20) (0 until 20).forEach { buffer[it] = Complex(it.toDouble(), -it.toDouble()) }