diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/Arrays.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/Arrays.kt index 5d341dd68..3d5062a4f 100644 --- a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/Arrays.kt +++ b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/Arrays.kt @@ -1,3 +1,4 @@ package scientifik.kmath.nd4j -internal fun narrowToIntArray(la: LongArray): IntArray = IntArray(la.size) { la[it].toInt() } \ No newline at end of file +internal fun widenToLongArray(ia: IntArray): LongArray = LongArray(ia.size) { ia[it].toLong() } +internal fun narrowToIntArray(la: LongArray): IntArray = IntArray(la.size) { la[it].toInt() } diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayScalarsIterator.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayScalarsIterator.kt deleted file mode 100644 index 2c2dc970f..000000000 --- a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/INDArrayScalarsIterator.kt +++ /dev/null @@ -1,19 +0,0 @@ -package scientifik.kmath.nd4j - -import org.nd4j.linalg.api.ndarray.INDArray -import org.nd4j.linalg.api.shape.Shape - -internal class INDArrayScalarsIterator(private val iterateOver: INDArray) : Iterator> { - private var i: Int = 0 - - override fun hasNext(): Boolean = i < iterateOver.length() - - override fun next(): Pair { - val idx = if (iterateOver.ordering() == 'c') - Shape.ind2subC(iterateOver, i++.toLong())!! - else - Shape.ind2sub(iterateOver, i++.toLong())!! - - return narrowToIntArray(idx) to iterateOver.getScalar(*idx) - } -} diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/ND4JStructure.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/ND4JStructure.kt deleted file mode 100644 index 1d0301ff9..000000000 --- a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/ND4JStructure.kt +++ /dev/null @@ -1,12 +0,0 @@ -package scientifik.kmath.nd4j - -import org.nd4j.linalg.api.ndarray.INDArray -import scientifik.kmath.structures.NDStructure - -data class ND4JStructure(val ndArray: INDArray) : NDStructure { - override val shape: IntArray - get() = narrowToIntArray(ndArray.shape()) - - override fun get(index: IntArray): INDArray = ndArray.getScalar(*index) - override fun elements(): Sequence> = Sequence { INDArrayScalarsIterator(ndArray) } -} diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/NDArrayIterators.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/NDArrayIterators.kt new file mode 100644 index 000000000..426b1ec2d --- /dev/null +++ b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/NDArrayIterators.kt @@ -0,0 +1,37 @@ +package scientifik.kmath.nd4j + +import org.nd4j.linalg.api.ndarray.INDArray +import org.nd4j.linalg.api.shape.Shape + +internal sealed class INDArrayIteratorBase(protected val iterateOver: INDArray) : Iterator> { + private var i: Int = 0 + + override fun hasNext(): Boolean = i < iterateOver.length() + + abstract fun getSingle(indices: LongArray): T + + final override fun next(): Pair { + val la = if (iterateOver.ordering() == 'c') + Shape.ind2subC(iterateOver, i++.toLong())!! + else + Shape.ind2sub(iterateOver, i++.toLong())!! + + return narrowToIntArray(la) to getSingle(la) + } +} + +internal class INDArrayDoubleIterator(iterateOver: INDArray) : INDArrayIteratorBase(iterateOver) { + override fun getSingle(indices: LongArray): Double = iterateOver.getDouble(*indices) +} + +internal class INDArrayLongIterator(iterateOver: INDArray) : INDArrayIteratorBase(iterateOver) { + override fun getSingle(indices: LongArray) = iterateOver.getLong(*indices) +} + +internal class INDArrayIntIterator(iterateOver: INDArray) : INDArrayIteratorBase(iterateOver) { + override fun getSingle(indices: LongArray) = iterateOver.getInt(*narrowToIntArray(indices)) +} + +internal class INDArrayFloatIterator(iterateOver: INDArray) : INDArrayIteratorBase(iterateOver) { + override fun getSingle(indices: LongArray) = iterateOver.getFloat(*indices) +} diff --git a/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/ScalarsND4JStructure.kt b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/ScalarsND4JStructure.kt new file mode 100644 index 000000000..ef8c3ec2e --- /dev/null +++ b/kmath-nd4j/src/main/kotlin/scientifik.kmath.nd4j/ScalarsND4JStructure.kt @@ -0,0 +1,34 @@ +package scientifik.kmath.nd4j + +import org.nd4j.linalg.api.ndarray.INDArray +import scientifik.kmath.structures.NDStructure + +interface INDArrayStructureBase : NDStructure { + val ndArray: INDArray + + override val shape: IntArray + get() = narrowToIntArray(ndArray.shape()) + + fun elementsIterator(): Iterator> + override fun elements(): Sequence> = Sequence { elementsIterator() } +} + +data class INDArrayIntStructure(override val ndArray: INDArray) : INDArrayStructureBase { + override fun elementsIterator(): Iterator> = INDArrayIntIterator(ndArray) + override fun get(index: IntArray): Int = ndArray.getInt(*index) +} + +data class INDArrayLongStructure(override val ndArray: INDArray) : INDArrayStructureBase { + override fun elementsIterator(): Iterator> = INDArrayLongIterator(ndArray) + override fun get(index: IntArray): Long = ndArray.getLong(*widenToLongArray(index)) +} + +data class INDArrayDoubleStructure(override val ndArray: INDArray) : INDArrayStructureBase { + override fun elementsIterator(): Iterator> = INDArrayDoubleIterator(ndArray) + override fun get(index: IntArray): Double = ndArray.getDouble(*index) +} + +data class INDArrayFloatStructure(override val ndArray: INDArray) : INDArrayStructureBase { + override fun elementsIterator(): Iterator> = INDArrayFloatIterator(ndArray) + override fun get(index: IntArray): Float = ndArray.getFloat(*index) +}