commonMain.jetbrains.datalore.plot.builder.assemble.PlotFacets.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.assemble

import jetbrains.datalore.base.interval.DoubleSpan
import jetbrains.datalore.plot.base.DataFrame
import jetbrains.datalore.plot.base.data.DataFrameUtil
import jetbrains.datalore.plot.builder.assemble.facet.FacetGrid
import jetbrains.datalore.plot.common.data.SeriesUtil

abstract class PlotFacets {

    abstract val isDefined: Boolean
    abstract val colCount: Int
    abstract val rowCount: Int
    abstract val numTiles: Int
    abstract val variables: List
    abstract val freeHScale: Boolean
    abstract val freeVScale: Boolean

    /**
     * @return List of Dataframes, one Dataframe per tile.
     *          Tiles are enumerated by rows, i.e.:
     *          the index is computed like: row * nCols + col
     */
    abstract fun dataByTile(data: DataFrame): List


    /**
     * @return List of FacetTileInfo.
     *          Tiles are enumerated by rows, i.e.:
     *          the index is computed like: row * nCols + col
     */
    abstract fun tileInfos(): List

    /**
     * @param domains Transformed X-mapped data ranges by tile.
     */
    open fun adjustHDomains(domains: List): List = domains

    /**
     * @param domains Transformed Y-mapped data ranges by tile.
     */
    open fun adjustVDomains(domains: List): List = domains

    companion object {
        const val DEF_ORDER_DIR = 0 // no ordering
        val DEF_FORMATTER: (Any) -> String = { it.toString() }

        fun undefined(): PlotFacets {
            return FacetGrid(null, null, emptyList(), emptyList(), 1, 1)
        }

        fun dataByLevelTuple(
            data: DataFrame,
            varNames: List,
            varLevels: List>
        ): List, DataFrame>> {
            // This also checks invariants.
            val nameLevelTuples = createNameLevelTuples(varNames, varLevels)

            val indicesByVarByLevel = dataIndicesByVarByLevel(data, varNames, varLevels)

            val dataByLevelKey = ArrayList, DataFrame>>()
            for (nameLevelTuple in nameLevelTuples) {
                val topName = nameLevelTuple.first().first
                val topLevel = nameLevelTuple.first().second
                val indices = ArrayList(indicesByVarByLevel.getValue(topName).getValue(topLevel))
                for (i in 1 until nameLevelTuple.size) {
                    val name = nameLevelTuple[i].first
                    val level = nameLevelTuple[i].second
                    val levelIndices = indicesByVarByLevel.getValue(name).getValue(level)
                    indices.retainAll(HashSet(levelIndices))
                }

                val levelKey = nameLevelTuple.map { it.second }

                // build the data subset
                val levelData = data.slice(indices)
                dataByLevelKey.add(levelKey to levelData)
            }

            return dataByLevelKey
        }

        private fun dataIndicesByVarByLevel(
            data: DataFrame,
            varNames: List,
            varLevels: List>
        ): Map>> {

            val indicesByVarByLevel = HashMap>>()
            for ((i, varName) in varNames.withIndex()) {
                val levels = varLevels[i]

                val indicesByLevel = HashMap>()
                for (level in levels) {
                    val indices = when {
                        // 'empty' data in layers with no aes mapping (only constants)
                        data.isEmpty -> emptyList()
                        DataFrameUtil.hasVariable(data, varName) -> {
                            val variable = DataFrameUtil.findVariableOrFail(data, varName)
                            SeriesUtil.matchingIndices(data[variable], level)
                        }
                        else -> {
                            // 'data' has no column 'varName' -> the entire data should be shown in each facet.
                            (0 until data.rowCount()).toList()
                        }
                    }
                    indicesByLevel[level] = indices
                }

                indicesByVarByLevel[varName] = indicesByLevel
            }

            return indicesByVarByLevel
        }

        fun createNameLevelTuples(
            varNames: List,
            varLevels: List>
        ): List>> {
            require(varNames.isNotEmpty()) { "Empty list of facet variables." }
            require(varNames.size == varNames.distinct().size) { "Facet variables must be distinct, were: $varNames." }
            check(varNames.size == varLevels.size)
            return createNameLevelTuplesIntern(varNames, varLevels)
        }

        private fun createNameLevelTuplesIntern(
            varNames: List,
            varLevels: List>
        ): List>> {
            val name = varNames.first()
            val levels = varLevels.first()

            val levelKeys = ArrayList>>()
            for (level in levels) {
                if (varNames.size > 1) {
                    val subKeys = createNameLevelTuples(
                        varNames.subList(1, varNames.size),
                        varLevels.subList(1, varLevels.size)
                    )
                    for (subKey in subKeys) {
                        levelKeys.add(listOf(name to level) + subKey)
                    }
                } else {
                    // exit
                    levelKeys.add(listOf(name to level))
                }
            }

            return levelKeys
        }

        fun reorderLevels(
            varNames: List,
            varLevels: List>,
            ordering: List
        ): List> {
            val orderingByFacet = varNames.zip(ordering).toMap()

            val result = ArrayList>()
            for ((i, name) in varNames.withIndex()) {
                if (i >= varLevels.size) break
                result.add(reorderVarLevels(name, varLevels[i], orderingByFacet.getValue(name)))
            }

            return result
        }

        fun reorderVarLevels(
            name: String?,
            levels: List,
            order: Int
        ): List {
            if (name == null) return levels

            // We expect either a list of Doubles or a list of Strings.
            @Suppress("UNCHECKED_CAST", "NAME_SHADOWING")
            levels as List>

            return when {
                order <= -1 -> levels.sortedDescending()
                order >= 1 -> levels.sorted()
                else -> levels  // not ordered
            }
        }
    }

    class FacetTileInfo constructor(
        val col: Int,
        val row: Int,
        val colLabs: List,
        val rowLab: String?,
        val hasHAxis: Boolean,
        val hasVAxis: Boolean,
        val isBottom: Boolean,  // true is the tile is the last one in its respective column.
        val trueIndex: Int     // tile index before re-ordering (in facet wrap)
    ) {
        override fun toString(): String {
            return "FacetTileInfo(col=$col, row=$row, colLabs=$colLabs, rowLab=$rowLab)"
        }
    }
}