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

Go to download
package space.kscience.kmath.nd

import space.kscience.kmath.misc.UnstableKMathAPI
import space.kscience.kmath.operations.ExtendedField
import space.kscience.kmath.operations.RealField
import space.kscience.kmath.operations.RingWithNumbers
import space.kscience.kmath.structures.Buffer
import space.kscience.kmath.structures.RealBuffer
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract

@OptIn(UnstableKMathAPI::class)
public class RealNDField(
    shape: IntArray,
) : BufferedNDField(shape, RealField, Buffer.Companion::real),
    RingWithNumbers>,
    ExtendedField> {

    override val zero: NDBuffer by lazy { produce { zero } }
    override val one: NDBuffer by lazy { produce { one } }

    override fun number(value: Number): NDBuffer {
        val d = value.toDouble() // minimize conversions
        return produce { d }
    }

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

    @Suppress("OVERRIDE_BY_INLINE")
    override inline fun NDStructure.map(
        transform: RealField.(Double) -> Double,
    ): NDBuffer {
        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 {
        val array = DoubleArray(strides.linearSize) { offset ->
            val index = strides.index(offset)
            RealField.initializer(index)
        }
        return NDBuffer(strides, RealBuffer(array))
    }

    @Suppress("OVERRIDE_BY_INLINE")
    override inline fun NDStructure.mapIndexed(
        transform: RealField.(index: IntArray, Double) -> Double,
    ): NDBuffer = NDBuffer(
        strides,
        buffer = RealBuffer(strides.linearSize) { offset ->
            RealField.transform(
                strides.index(offset),
                buffer.array[offset]
            )
        })

    @Suppress("OVERRIDE_BY_INLINE")
    override inline fun combine(
        a: NDStructure,
        b: NDStructure,
        transform: RealField.(Double, Double) -> Double,
    ): NDBuffer {
        val buffer = RealBuffer(strides.linearSize) { offset ->
            RealField.transform(a.buffer.array[offset], b.buffer.array[offset])
        }
        return NDBuffer(strides, buffer)
    }

    override fun power(arg: NDStructure, pow: Number): NDBuffer = arg.map { power(it, pow) }

    override fun exp(arg: NDStructure): NDBuffer = arg.map { exp(it) }

    override fun ln(arg: NDStructure): NDBuffer = arg.map { ln(it) }

    override fun sin(arg: NDStructure): NDBuffer = arg.map { sin(it) }
    override fun cos(arg: NDStructure): NDBuffer = arg.map { cos(it) }
    override fun tan(arg: NDStructure): NDBuffer = arg.map { tan(it) }
    override fun asin(arg: NDStructure): NDBuffer = arg.map { asin(it) }
    override fun acos(arg: NDStructure): NDBuffer = arg.map { acos(it) }
    override fun atan(arg: NDStructure): NDBuffer = arg.map { atan(it) }

    override fun sinh(arg: NDStructure): NDBuffer = arg.map { sinh(it) }
    override fun cosh(arg: NDStructure): NDBuffer = arg.map { cosh(it) }
    override fun tanh(arg: NDStructure): NDBuffer = arg.map { tanh(it) }
    override fun asinh(arg: NDStructure): NDBuffer = arg.map { asinh(it) }
    override fun acosh(arg: NDStructure): NDBuffer = arg.map { acosh(it) }
    override fun atanh(arg: NDStructure): NDBuffer = arg.map { atanh(it) }
}

public fun NDAlgebra.Companion.real(vararg shape: Int): RealNDField = RealNDField(shape)

/**
 * Produce a context for n-dimensional operations inside this real field
 */
public inline fun  RealField.nd(vararg shape: Int, action: RealNDField.() -> R): R {
    contract { callsInPlace(action, InvocationKind.EXACTLY_ONCE) }
    return RealNDField(shape).run(action)
}