commonMain.jetbrains.datalore.plot.base.stat.QQLineStat.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of lets-plot-common Show documentation
Lets-Plot JVM package without rendering part
There is a newer version: 4.5.3-alpha1
/*
 * Copyright (c) 2022. JetBrains s.r.o.
 * Use of this source code is governed by the MIT license that can be found in the LICENSE file.
 */

package jetbrains.datalore.plot.base.stat

import jetbrains.datalore.plot.base.Aes
import jetbrains.datalore.plot.base.DataFrame
import jetbrains.datalore.plot.base.StatContext
import jetbrains.datalore.plot.base.data.TransformVar
import kotlin.math.min
import kotlin.math.max

class QQLineStat(
    private val distribution: QQStat.Distribution,
    private val distributionParameters: List,
    private val lineQuantiles: Pair
) : BaseStat(DEF_MAPPING) {

    override fun consumes(): List> {
        return listOf(Aes.SAMPLE)
    }

    override fun apply(data: DataFrame, statCtx: StatContext, messageConsumer: (s: String) -> Unit): DataFrame {
        if (!hasRequiredValues(data, Aes.SAMPLE)) {
            return withEmptyStatValues()
        }

        val statData = buildStat(data.getNumeric(TransformVar.SAMPLE))

        return DataFrame.Builder()
            .putNumeric(Stats.THEORETICAL, statData.getValue(Stats.THEORETICAL))
            .putNumeric(Stats.SAMPLE, statData.getValue(Stats.SAMPLE))
            .build()
    }

    private fun buildStat(
        sampleSeries: List
    ): MutableMap> {
        val sortedSample = sampleSeries.filter { it?.isFinite() ?: false }.map { it!! }.sorted()
        if (!sortedSample.any()) {
            return mutableMapOf(
                Stats.THEORETICAL to emptyList(),
                Stats.SAMPLE to emptyList()
            )
        }

        val quantilesSample = Pair(
            QQStatUtil.getEstimatedQuantile(sortedSample, lineQuantiles.first),
            QQStatUtil.getEstimatedQuantile(sortedSample, lineQuantiles.second)
        )
        val dist = QQStatUtil.getDistribution(distribution, distributionParameters)
        // Use min/max to avoid an infinity
        val quantilesTheoretical = Pair(
            dist.inverseCumulativeProbability(max(0.5 / sortedSample.size, lineQuantiles.first)),
            dist.inverseCumulativeProbability(min(1.0 - 0.5 / sortedSample.size, lineQuantiles.second))
        )
        val endpointsTheoretical = listOf(
            dist.inverseCumulativeProbability(0.5 / sortedSample.size),
            dist.inverseCumulativeProbability(1.0 - 0.5 / sortedSample.size)
        )
        if (quantilesTheoretical.first == quantilesTheoretical.second) {
            return mutableMapOf(
                Stats.THEORETICAL to listOf(quantilesTheoretical.first, quantilesTheoretical.second),
                Stats.SAMPLE to listOf(sortedSample.first(), sortedSample.last())
            )
        }

        val line = QQStatUtil.lineByPoints(quantilesTheoretical, quantilesSample)
        return mutableMapOf(
            Stats.THEORETICAL to endpointsTheoretical,
            Stats.SAMPLE to endpointsTheoretical.map { line(it) }
        )
    }

    companion object {
        val DEF_LINE_QUANTILES = Pair(0.25, 0.75)

        private val DEF_MAPPING: Map, DataFrame.Variable> = mapOf(
            Aes.X to Stats.THEORETICAL,
            Aes.Y to Stats.SAMPLE
        )
    }
}