commonMain.jetbrains.datalore.plot.builder.data.GroupingContext.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) 2019. 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.builder.data

import jetbrains.datalore.plot.base.DataFrame
import jetbrains.datalore.plot.base.DataFrame.Variable
import jetbrains.datalore.plot.base.stat.Stats
import jetbrains.datalore.plot.builder.VarBinding
import jetbrains.datalore.plot.builder.data.DataProcessing.findOptionalVariable

class GroupingContext(
    private val myData: DataFrame,
    bindings: List,
    groupingVarName: String?,
    pathIdVarName: String?,
    private val myExpectMultiple: Boolean
) {

    private val myBindings: List = ArrayList(bindings)
    internal val optionalGroupingVar: Variable? = findOptionalVariable(myData, groupingVarName)
    private val pathIdVar: Variable? = findOptionalVariable(myData, pathIdVarName)

    private var myGroupSizeList: List? = null
    private var myGroupMapper: ((Int) -> Int)? = null

    //myGroupMapper = DataProcessing.computeGroups(myData, myBindings, myOptionalGroupingVar, myExpectMultiple);
    val groupMapper: (Int) -> Int
        get() = { index ->
            if (myGroupMapper == null) {
                myGroupMapper = computeGroups()
            }
            myGroupMapper!!(index)
        }

    private fun computeGroups(): (Int) -> Int {
        if (myData.has(Stats.GROUP)) {
            val list = myData.getNumeric(Stats.GROUP)
            return GroupUtil.wrap(list)
        } else if (myGroupSizeList != null) {
            if (myGroupSizeList!!.size == myData.rowCount()) {
                return GroupUtil.SINGLE_GROUP
            } else {
                val groupByPointIndex =
                    toIndexMap(myGroupSizeList!!)
                return GroupUtil.wrap(groupByPointIndex)
            }
        } else if (myExpectMultiple) {
            return DataProcessing.computeGroups(
                myData,
                myBindings,
                optionalGroupingVar,
                pathIdVar
            )
        }
        return GroupUtil.SINGLE_GROUP
    }

    companion object {
        internal fun withOrderedGroups(data: DataFrame, groupSizeList: List): GroupingContext {
            val groupingContext = GroupingContext(data, emptyList(), null, null, false)
            groupingContext.myGroupSizeList = ArrayList(groupSizeList)
            return groupingContext
        }

        private fun toIndexMap(groupSizeList: List): Map {
            val result = HashMap()
            var currentGroup = 0
            var currentGroupIndexOffset = 0
            for (groupSize in groupSizeList) {
                for (i in 0 until groupSize) {
                    result[currentGroupIndexOffset + i] = currentGroup
                }
                currentGroup++
                currentGroupIndexOffset += groupSize
            }
            return result
        }
    }
}