started std rework

Author: Jolanrensen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/Defaults.kt

@@ -3,5 +3,16 @@ package org.jetbrains.kotlinx.dataframe.api
 @PublishedApi
 internal val skipNaN_default: Boolean = false
 
+/**
+ * Default delta degrees of freedom for the standard deviation (std).
+ *
+ * The default is set to `1`,
+ * meaning DataFrame uses [Bessel’s correction](https://en.wikipedia.org/wiki/Bessel%27s_correction) to calculate the
+ * "unbiased sample standard deviation" by default.
+ * This is also the standard in languages like R.
+ *
+ * This is different from the "population standard deviation" (where `ddof = 0`),
+ * which is used in libraries like Numpy.
+ */
 @PublishedApi
 internal val ddof_default: Int = 1

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt

@@ -66,6 +66,15 @@ internal object Aggregators {
         multipleColumnsHandler = FlatteningMultipleColumnsHandler(),
     )
 
+    private fun <Return : Number?> flattenReducingForNumbers(
+        getReturnType: CalculateReturnType,
+        reducer: Reducer<Number, Return>,
+    ) = Aggregator(
+        aggregationHandler = ReducingAggregationHandler(reducer, getReturnType),
+        inputHandler = NumberInputHandler(),
+        multipleColumnsHandler = FlatteningMultipleColumnsHandler(),
+    )
+
     private fun <Return : Number?> twoStepReducingForNumbers(
         getReturnType: CalculateReturnType,
         reducer: Reducer<Number, Return>,
@@ -111,8 +120,8 @@ internal object Aggregators {
 
     // T: Number? -> Double
     val std by withTwoOptions { skipNA: Boolean, ddof: Int ->
-        flattenReducingForAny<Number, Double>(stdTypeConversion) { type ->
-            asIterable().std(type, skipNA, ddof)
+        flattenReducingForNumbers(stdTypeConversion) { type ->
+           std(type, skipNA, ddof)
         }
     }
 

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/corr.kt

@@ -13,7 +13,7 @@ import org.jetbrains.kotlinx.dataframe.api.isSuitableForCorr
 import org.jetbrains.kotlinx.dataframe.api.toValueColumn
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
-import org.jetbrains.kotlinx.dataframe.math.varianceAndMean
+import org.jetbrains.kotlinx.dataframe.math.calculateBasicStatsOrNull
 import org.jetbrains.kotlinx.dataframe.nrow
 import kotlin.math.sqrt
 
@@ -51,7 +51,7 @@ internal fun <T, C, R> Corr<T, C>.corrImpl(otherColumns: ColumnsSelector<T, R>):
     }
 
     val stdMeans = cols.mapValues {
-        it.value.toList().varianceAndMean()
+        it.value.toList().calculateBasicStatsOrNull()
     }
 
     val cache = mutableMapOf<Pair<ColumnPath, ColumnPath>, Double>()

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/basicStats.kt

@@ -0,0 +1,53 @@
+@file:Suppress("DuplicatedCode")
+
+package org.jetbrains.kotlinx.dataframe.math
+
+import org.jetbrains.kotlinx.dataframe.api.ddof_default
+import org.jetbrains.kotlinx.dataframe.api.skipNaN_default
+import kotlin.math.sqrt
+
+internal data class BasicStats(val count: Int, val mean: Double, val variance: Double)
+
+/**
+ * Calculates the standard deviation from a [BasicStats] with optional delta degrees of freedom.
+ *
+ * @param ddof delta degrees of freedom, the bias-correction of std.
+ *   Default is [ddof_default], so `ddof = 1`, the "unbiased sample standard deviation", but alternatively,
+ *   the "population standard deviation", so `ddof = 0`, can be used.
+ */
+internal fun BasicStats.std(ddof: Int): Double =
+    if (count <= ddof) {
+        Double.NaN
+    } else {
+        sqrt(variance / (count - ddof))
+    }
+
+/**
+ * Creates [BasicStats] instance for [this] sequence.
+ *
+ * This contains the [count][BasicStats.count], [mean][BasicStats.mean], and [variance][BasicStats.variance] and
+ * can be used to efficiently calculate the [standard deviation][std].
+ */
+internal fun Sequence<Double>.calculateBasicStatsOrNull(skipNaN: Boolean = skipNaN_default): BasicStats? {
+    var count = 0
+    var sum = .0
+    for (element in this) {
+        if (element.isNaN()) {
+            if (skipNaN) {
+                continue
+            } else {
+                return null
+            }
+        }
+        sum += element
+        count++
+    }
+    val mean = sum / count
+    var variance = .0
+    for (element in this) {
+        if (element.isNaN()) continue
+        val diff = element - mean
+        variance += diff * diff
+    }
+    return BasicStats(count = count, mean = mean, variance = variance)
+}

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt

@@ -7,25 +7,16 @@ import org.jetbrains.kotlinx.dataframe.impl.renderType
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
-import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
 @Suppress("UNCHECKED_CAST")
 @PublishedApi
-internal fun <T : Number> Iterable<T?>.std(
-    type: KType,
-    skipNaN: Boolean = skipNaN_default,
-    ddof: Int = ddof_default,
-): Double {
+internal fun <T : Number> Sequence<T?>.std(type: KType, skipNaN: Boolean, ddof: Int): Double {
     if (type.isMarkedNullable) {
-        return when {
-            skipNaN -> filterNotNull().std(type = type.withNullability(false), skipNaN = true, ddof = ddof)
-            contains(null) -> Double.NaN
-            else -> std(type = type.withNullability(false), skipNaN = false, ddof = ddof)
-        }
+        error("Encountered nullable type ${renderType(type)} in std function. This should not occur.")
     }
-    return when (type.classifier) {
-        Double::class -> (this as Iterable<Double>).std(skipNaN, ddof)
+    return when (type) {
+        typeOf<Double>() -> (this as Iterable<Double>).std(skipNaN, ddof)
         Float::class -> (this as Iterable<Float>).std(skipNaN, ddof)
         Int::class, Short::class, Byte::class -> (this as Iterable<Int>).std(ddof)
         Long::class -> (this as Iterable<Long>).std(ddof)
@@ -44,20 +35,20 @@ internal val stdTypeConversion: CalculateReturnType = { _, _ ->
 
 @JvmName("doubleStd")
 internal fun Iterable<Double>.std(skipNaN: Boolean = skipNaN_default, ddof: Int = ddof_default): Double =
-    varianceAndMean(skipNaN)?.std(ddof) ?: Double.NaN
+    calculateBasicStatsOrNull(skipNaN)?.std(ddof) ?: Double.NaN
 
 @JvmName("floatStd")
 internal fun Iterable<Float>.std(skipNaN: Boolean = skipNaN_default, ddof: Int = ddof_default): Double =
-    varianceAndMean(skipNaN)?.std(ddof) ?: Double.NaN
+    calculateBasicStatsOrNull(skipNaN)?.std(ddof) ?: Double.NaN
 
 @JvmName("intStd")
-internal fun Iterable<Int>.std(ddof: Int = ddof_default): Double = varianceAndMean().std(ddof)
+internal fun Iterable<Int>.std(ddof: Int = ddof_default): Double = calculateBasicStatsOrNull().std(ddof)
 
 @JvmName("longStd")
-internal fun Iterable<Long>.std(ddof: Int = ddof_default): Double = varianceAndMean().std(ddof)
+internal fun Iterable<Long>.std(ddof: Int = ddof_default): Double = calculateBasicStatsOrNull().std(ddof)
 
 @JvmName("bigDecimalStd")
-internal fun Iterable<BigDecimal>.std(ddof: Int = ddof_default): Double = varianceAndMean().std(ddof)
+internal fun Iterable<BigDecimal>.std(ddof: Int = ddof_default): Double = calculateBasicStatsOrNull().std(ddof)
 
 @JvmName("bigIntegerStd")
-internal fun Iterable<BigInteger>.std(ddof: Int = ddof_default): Double = varianceAndMean().std(ddof)
+internal fun Iterable<BigInteger>.std(ddof: Int = ddof_default): Double = calculateBasicStatsOrNull().std(ddof)