fun geomImshow()

in plot-api/src/commonMain/kotlin/org/jetbrains/letsPlot/geom/geom_imshow.kt [63:213]
119 lines of code
22 McCabe index (conditional complexity)

fun geomImshow(
    rasterData: RasterData,
    norm: Boolean = true,
    vmin: Number? = null,
    vmax: Number? = null,
    extent: List<Number>? = null,
    showLegend: Boolean = true,
    colorBy: String = "paint_c",
): Feature {
    require(extent == null || extent.size == 4) { "Invalid `extent`: list of 4 numbers expected: ${extent!!.size}" }
    val colorAesthetics = listOf("fill", "color", "paint_a", "paint_b", "paint_c")
    require(colorBy in colorAesthetics) { "Invalid colorBy value \"$colorBy\". Use: \"color\", \"fill\", \"paint_a\", \"paint_b\" or \"paint_c\"." }

    var raster = rasterData.createRaster()

    require(raster.nChannels in 1..4) {
        "Invalid rasterData: num of channels expected to be 1 (G) or 2 (GA) for greyscale image, 3 (RGB) or 4 (RGBA) for color image, but was ${raster.nChannels}"
    }

    val cmap: String? = null // TODO: add palettes support

    val greyscale = raster.nChannels == 1
    var greyScaleDataMin: Double = Double.NaN
    var greyScaleDataMax: Double = Double.NaN

    if (greyscale) {
        var hasNan = raster.hasNan()
        val maxLum = if (!(hasNan && cmap != null)) 255 else 254  // index 255 reserved for NaN-s
        normalize2d(raster, norm, vmin?.toFloat(), vmax?.toFloat(), maxLum).let {
            greyScaleDataMin = it.first
            greyScaleDataMax = it.second
        }
        hasNan = raster.hasNan()

        if (hasNan && cmap.isNullOrEmpty()) {
            // add alpha
            raster = raster.addChannel()

            require(raster.nChannels == 2)
            raster.updatePixels { pix ->
                pix[1] = when (pix[0].isNaN()) {
                    true -> vmin?.toFloat() ?: Float.NaN
                    false -> 255f
                }
            }
        } else if (hasNan && !cmap.isNullOrEmpty()) {
            raster.updateChannels { it.takeUnless { it.isNaN() } ?: 255f }
        }
    } else {
        if (raster.isDTypeF) {
            raster.updateChannels { it * 255f + 0.5f }
        }
    }

    raster.updateChannels { it.coerceIn(0.0f, 255.0f) }

    var (extX0, extX1, extY0, extY1) =
        extent?.map(Number::toDouble)
            ?: listOf(-.5, raster.width - .5, -.5, raster.height - .5)

    val flipColumns: Boolean = extX0 > extX1
    if (flipColumns) {
        extX0 = extX1.also { extX1 = extX0 }
    }

    val flipRows = extY0 > extY1
    if (flipRows) {
        extY0 = extY1.also { extY1 = extY0 }
    }

    val outputStream = OutputPngStream()
    val png = PngWriter(
        outputStream, ImageInfo(
            raster.width,
            raster.height,
            bitdepth = 8,
            alpha = (raster.nChannels == 4 || raster.nChannels == 2),
            greyscale = raster.nChannels < 3
        )
    )

    val iLine = ImageLineByte(png.imgInfo)
    val px = raster.pixel()
    val rows = (0 until raster.height).let { it.takeIf { !flipRows } ?: it.reversed() }
    val columns = (0 until raster.width).let { it.takeIf { !flipColumns } ?: it.reversed() }
    for (row in rows) {
        var p = 0
        for (col in columns) {
            px.atXY(col, row).channels().forEach {
                iLine.scanline[p++] = it.toInt().toByte()
            }
        }
        png.writeRow(iLine)
    }
    png.end()


    // Show Legend (color-bar) if applicable.
    val layerMapping: Options = if (greyscale && showLegend) {
        // Provide two imaginable data-points to build a legend.
        Options.of(
            colorBy to listOf(greyScaleDataMin, greyScaleDataMax)
        )
    } else {
        Options.empty()
    }

    val legendTitle = ""
    val colorScale: Scale? = if (greyscale && showLegend) {
        if (cmap != null) when (norm) {
            true -> null  // ToDo
            else -> null  // ToDo
        } else {
            val start = if (norm) 0.0 else greyScaleDataMin / 255
            val end = if (norm) 1.0 else greyScaleDataMax / 255
            scaleGrey(aesthetic = colorBy, start = start, end = end, name = legendTitle)
        }
    } else {
        null
    }


    val geomLayer = object : Layer(
        geom = GeomOptions(GeomKind.IMAGE),
        data = null,
        mapping = layerMapping,
        stat = Stat.identity,
        position = null,
        showLegend = showLegend,
        inheritAes = false,
        sampling = null,
        orientation = null,
        tooltips = null,
        labels = null
    ) {
        override fun seal(): Options {
            return Options.of(
                Option.Geom.Image.HREF to "data:image/png;base64," + Base64.encode(outputStream.byteArray),
                Option.Geom.Image.XMIN to extX0,
                Option.Geom.Image.YMIN to extY0,
                Option.Geom.Image.XMAX to extX1,
                Option.Geom.Image.YMAX to extY1,
                Option.Layer.COLOR_BY to colorBy  // for the legend
            )
        }
    }

    return colorScale?.let {
        geomLayer + it
    } ?: geomLayer
}