commonMain.space.kscience.kmath.nd.BufferNDAlgebra.kt Maven / Gradle / Ivy

Go to download
package space.kscience.kmath.nd

import space.kscience.kmath.operations.Field
import space.kscience.kmath.operations.RealField
import space.kscience.kmath.operations.Ring
import space.kscience.kmath.operations.Space
import space.kscience.kmath.structures.Buffer
import space.kscience.kmath.structures.BufferFactory
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract

public interface BufferNDAlgebra : NDAlgebra {
    public val strides: Strides
    public val bufferFactory: BufferFactory

    override fun produce(initializer: C.(IntArray) -> T): NDBuffer = NDBuffer(
        strides,
        bufferFactory(strides.linearSize) { offset ->
            elementContext.initializer(strides.index(offset))
        }
    )

    public val NDStructure.buffer: Buffer
        get() = when {
            !shape.contentEquals([email protected]) -> throw ShapeMismatchException(
                [email protected],
                shape
            )
            this is NDBuffer && this.strides == [email protected] -> this.buffer
            else -> bufferFactory(strides.linearSize) { offset -> get(strides.index(offset)) }
        }

    override fun NDStructure.map(transform: C.(T) -> T): NDBuffer {
        val buffer = bufferFactory(strides.linearSize) { offset ->
            elementContext.transform(buffer[offset])
        }
        return NDBuffer(strides, buffer)
    }

    override fun NDStructure.mapIndexed(transform: C.(index: IntArray, T) -> T): NDBuffer {
        val buffer = bufferFactory(strides.linearSize) { offset ->
            elementContext.transform(
                strides.index(offset),
                buffer[offset]
            )
        }
        return NDBuffer(strides, buffer)
    }

    override fun combine(a: NDStructure, b: NDStructure, transform: C.(T, T) -> T): NDBuffer {
        val buffer = bufferFactory(strides.linearSize) { offset ->
            elementContext.transform(a.buffer[offset], b.buffer[offset])
        }
        return NDBuffer(strides, buffer)
    }
}

public open class BufferedNDSpace>(
    final override val shape: IntArray,
    final override val elementContext: R,
    final override val bufferFactory: BufferFactory,
) : NDSpace, BufferNDAlgebra {
    override val strides: Strides = DefaultStrides(shape)
    override val zero: NDBuffer by lazy { produce { zero } }
}

public open class BufferedNDRing>(
    shape: IntArray,
    elementContext: R,
    bufferFactory: BufferFactory,
) : BufferedNDSpace(shape, elementContext, bufferFactory), NDRing {
    override val one: NDBuffer by lazy { produce { one } }
}

public open class BufferedNDField>(
    shape: IntArray,
    elementContext: R,
    bufferFactory: BufferFactory,
) : BufferedNDRing(shape, elementContext, bufferFactory), NDField

// space factories
public fun > NDAlgebra.Companion.space(
    space: A,
    bufferFactory: BufferFactory,
    vararg shape: Int,
): BufferedNDSpace = BufferedNDSpace(shape, space, bufferFactory)

public inline fun , R> A.ndSpace(
    noinline bufferFactory: BufferFactory,
    vararg shape: Int,
    action: BufferedNDSpace.() -> R,
): R {
    contract { callsInPlace(action, InvocationKind.EXACTLY_ONCE) }
    return NDAlgebra.space(this, bufferFactory, *shape).run(action)
}

//ring factories
public fun > NDAlgebra.Companion.ring(
    ring: A,
    bufferFactory: BufferFactory,
    vararg shape: Int,
): BufferedNDRing = BufferedNDRing(shape, ring, bufferFactory)

public inline fun , R> A.ndRing(
    noinline bufferFactory: BufferFactory,
    vararg shape: Int,
    action: BufferedNDRing.() -> R,
): R {
    contract { callsInPlace(action, InvocationKind.EXACTLY_ONCE) }
    return NDAlgebra.ring(this, bufferFactory, *shape).run(action)
}

//field factories
public fun > NDAlgebra.Companion.field(
    field: A,
    bufferFactory: BufferFactory,
    vararg shape: Int,
): BufferedNDField = BufferedNDField(shape, field, bufferFactory)

@Suppress("UNCHECKED_CAST")
public inline fun > NDAlgebra.Companion.auto(
    field: A,
    vararg shape: Int,
): NDField = when (field) {
    RealField -> RealNDField(shape) as NDField
    else -> BufferedNDField(shape, field, Buffer.Companion::auto)
}

public inline fun , R> A.ndField(
    noinline bufferFactory: BufferFactory,
    vararg shape: Int,
    action: BufferedNDField.() -> R,
): R {
    contract { callsInPlace(action, InvocationKind.EXACTLY_ONCE) }
    return NDAlgebra.field(this, bufferFactory, *shape).run(action)
}