expanded std with BigInteger and Number

Kotlin · Nov 7, 2024 · 2ef56b4 · 2ef56b4
1 parent 4d4ebda
commit 2ef56b4
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.api.ddof_default
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
 import org.jetbrains.kotlinx.dataframe.impl.renderType
 import java.math.BigDecimal
+import java.math.BigInteger
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 
@@ -26,6 +27,8 @@ internal fun <T : Number> Iterable<T?>.std(
         Int::class, Short::class, Byte::class -> (this as Iterable<Int>).std(ddof)
         Long::class -> (this as Iterable<Long>).std(ddof)
         BigDecimal::class -> (this as Iterable<BigDecimal>).std(ddof)
+        BigInteger::class -> (this as Iterable<BigInteger>).std(ddof)
+        Number::class -> (this as Iterable<Number>).map { it.toDouble() }.std(skipNA, ddof)
         Nothing::class -> Double.NaN
         else -> throw IllegalArgumentException("Unable to compute the std for type ${renderType(type)}")
     }
@@ -47,3 +50,6 @@ public fun Iterable<Long>.std(ddof: Int = ddof_default): Double = varianceAndMea
 
 @JvmName("bigDecimalStd")
 public fun Iterable<BigDecimal>.std(ddof: Int = ddof_default): Double = varianceAndMean().std(ddof)
+
+@JvmName("bigIntegerStd")
+public fun Iterable<BigInteger>.std(ddof: Int = ddof_default): Double = varianceAndMean().std(ddof)
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/stdMean.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/stdMean.kt
@@ -4,6 +4,7 @@ package org.jetbrains.kotlinx.dataframe.math
 
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
 import java.math.BigDecimal
+import java.math.BigInteger
 import kotlin.math.sqrt
 
 public data class BasicStats(val count: Int, val mean: Double, val variance: Double) {
@@ -114,3 +115,20 @@ public fun Iterable<BigDecimal>.varianceAndMean(): BasicStats {
     }
     return BasicStats(count, mean.toDouble(), variance.toDouble())
 }
+
+@JvmName("bigIntegerVarianceAndMean")
+public fun Iterable<BigInteger>.varianceAndMean(): BasicStats {
+    var count = 0
+    var sum = BigInteger.ZERO
+    for (element in this) {
+        sum += element
+        count++
+    }
+    val mean = sum.toDouble() / count
+    var variance = .0
+    for (element in this) {
+        val diff = element.toDouble() - mean
+        variance += diff * diff
+    }
+    return BasicStats(count, mean, variance)
+}