feat(heatmap): restore forceSquare feature

packages/heatmap/src/HeatMap.tsx

@@ -1,4 +1,4 @@
-import { ReactNode, Fragment, createElement } from 'react'
+import { ReactNode, Fragment, createElement, useMemo } from 'react'
 import { SvgWrapper, Container, useDimensions } from '@nivo/core'
 import { Axes, Grid } from '@nivo/axes'
 import { AnchoredContinuousColorsLegendSvg } from '@nivo/legends'
@@ -27,7 +27,7 @@ const InnerHeatMap = <Datum extends HeatMapDatum, ExtraProps extends object>({
     width,
     height,
     margin: partialMargin,
-    // forceSquare = svgDefaultProps.forceSquare,
+    forceSquare = svgDefaultProps.forceSquare,
     xInnerPadding = svgDefaultProps.xInnerPadding,
     xOuterPadding = svgDefaultProps.xOuterPadding,
     yInnerPadding = svgDefaultProps.yInnerPadding,
@@ -67,20 +67,31 @@ const InnerHeatMap = <Datum extends HeatMapDatum, ExtraProps extends object>({
     ariaLabelledBy,
     ariaDescribedBy,
 }: InnerHeatMapProps<Datum, ExtraProps>) => {
-    const { margin, innerWidth, innerHeight, outerWidth, outerHeight } = useDimensions(
-        width,
-        height,
-        partialMargin
-    )
+    const {
+        margin: _margin,
+        innerWidth: _innerWidth,
+        innerHeight: _innerHeight,
+        outerWidth,
+        outerHeight,
+    } = useDimensions(width, height, partialMargin)
 
-    const { xScale, yScale, cells, colorScale, activeCell, setActiveCell } = useHeatMap<
-        Datum,
-        ExtraProps
-    >({
-        data,
-        valueFormat,
+    const {
         width: innerWidth,
         height: innerHeight,
+        offsetX,
+        offsetY,
+        xScale,
+        yScale,
+        cells,
+        colorScale,
+        activeCell,
+        setActiveCell,
+    } = useHeatMap<Datum, ExtraProps>({
+        data,
+        valueFormat,
+        width: _innerWidth,
+        height: _innerHeight,
+        forceSquare,
         xInnerPadding,
         xOuterPadding,
         yInnerPadding,
@@ -97,6 +108,15 @@ const InnerHeatMap = <Datum extends HeatMapDatum, ExtraProps extends object>({
         hoverTarget,
     })
 
+    const margin = useMemo(
+        () => ({
+            ..._margin,
+            top: _margin.top + offsetY,
+            left: _margin.left + offsetX,
+        }),
+        [_margin, offsetX, offsetY]
+    )
+
     const layerById: Record<LayerId, ReactNode> = {
         grid: null,
         axes: null,

packages/heatmap/src/HeatMapCanvas.tsx

@@ -29,11 +29,11 @@ const InnerHeatMapCanvas = <Datum extends HeatMapDatum, ExtraProps extends objec
     width,
     height,
     margin: partialMargin,
-    // forceSquare = canvasDefaultProps.forceSquare,
     xInnerPadding = canvasDefaultProps.xInnerPadding,
     xOuterPadding = canvasDefaultProps.xOuterPadding,
     yInnerPadding = canvasDefaultProps.yInnerPadding,
     yOuterPadding = canvasDefaultProps.yOuterPadding,
+    forceSquare = canvasDefaultProps.forceSquare,
     sizeVariation = canvasDefaultProps.sizeVariation,
     renderCell: _renderCell = canvasDefaultProps.renderCell as CellShape,
     opacity = canvasDefaultProps.opacity,
@@ -66,24 +66,35 @@ const InnerHeatMapCanvas = <Datum extends HeatMapDatum, ExtraProps extends objec
 }: InnerNetworkCanvasProps<Datum, ExtraProps>) => {
     const canvasEl = useRef<HTMLCanvasElement | null>(null)
 
-    const { margin, innerWidth, innerHeight, outerWidth, outerHeight } = useDimensions(
-        width,
-        height,
-        partialMargin
-    )
+    const {
+        margin: _margin,
+        innerWidth: _innerWidth,
+        innerHeight: _innerHeight,
+        outerWidth,
+        outerHeight,
+    } = useDimensions(width, height, partialMargin)
 
-    const { xScale, yScale, cells, activeCell, setActiveCell, colorScale } = useHeatMap<
-        Datum,
-        ExtraProps
-    >({
-        data,
-        valueFormat,
+    const {
         width: innerWidth,
         height: innerHeight,
+        offsetX,
+        offsetY,
+        xScale,
+        yScale,
+        cells,
+        colorScale,
+        activeCell,
+        setActiveCell,
+    } = useHeatMap<Datum, ExtraProps>({
+        data,
+        valueFormat,
+        width: _innerWidth,
+        height: _innerHeight,
         xInnerPadding,
         xOuterPadding,
         yInnerPadding,
         yOuterPadding,
+        forceSquare,
         sizeVariation,
         colors,
         emptyColor,
@@ -96,6 +107,15 @@ const InnerHeatMapCanvas = <Datum extends HeatMapDatum, ExtraProps extends objec
         hoverTarget,
     })
 
+    const margin = useMemo(
+        () => ({
+            ..._margin,
+            top: _margin.top + offsetY,
+            left: _margin.left + offsetX,
+        }),
+        [_margin, offsetX, offsetY]
+    )
+
     const boundAnnotations = useCellAnnotations(cells, annotations)
     const computedAnnotations = useComputedAnnotations({
         annotations: boundAnnotations,

packages/heatmap/src/compute.ts

@@ -1,24 +1,69 @@
-import { scaleBand } from 'd3-scale'
+import { scaleBand, scaleLinear } from 'd3-scale'
 import { castBandScale } from '@nivo/scales'
-import { ComputedCell, HeatMapCommonProps, HeatMapDataProps, HeatMapDatum } from './types'
+import {
+    ComputedCell,
+    HeatMapCommonProps,
+    HeatMapDataProps,
+    HeatMapDatum,
+    SizeVariationConfig,
+} from './types'
+
+export const computeLayout = ({
+    width: _width,
+    height: _height,
+    rows,
+    columns,
+    forceSquare,
+}: {
+    width: number
+    height: number
+    rows: number
+    columns: number
+    forceSquare: boolean
+}) => {
+    let width = _width
+    let height = _height
+
+    let offsetX = 0
+    let offsetY = 0
+
+    if (forceSquare) {
+        const cellWidth = Math.max(_width / columns, 0)
+        const cellHeight = Math.max(_height / rows, 0)
+        const cellSize = Math.min(cellWidth, cellHeight)
+
+        width = cellSize * columns
+        height = cellSize * rows
+
+        offsetX = (_width - width) / 2
+        offsetY = (_height - height) / 2
+    }
+
+    return {
+        offsetX,
+        offsetY,
+        width,
+        height,
+    }
+}
 
 export const computeCells = <Datum extends HeatMapDatum, ExtraProps extends object>({
     data,
-    width,
-    height,
+    width: _width,
+    height: _height,
     xInnerPadding,
     xOuterPadding,
     yInnerPadding,
     yOuterPadding,
+    forceSquare,
 }: {
     data: HeatMapDataProps<Datum, ExtraProps>['data']
     width: number
     height: number
-    xInnerPadding: HeatMapCommonProps<Datum>['xInnerPadding']
-    xOuterPadding: HeatMapCommonProps<Datum>['xOuterPadding']
-    yInnerPadding: HeatMapCommonProps<Datum>['yInnerPadding']
-    yOuterPadding: HeatMapCommonProps<Datum>['yOuterPadding']
-}) => {
+} & Pick<
+    HeatMapCommonProps<Datum>,
+    'xOuterPadding' | 'xInnerPadding' | 'yOuterPadding' | 'yInnerPadding' | 'forceSquare'
+>) => {
     const xValuesSet = new Set<Datum['x']>()
     const serieIds: string[] = []
     const allValues: number[] = []
@@ -47,6 +92,15 @@ export const computeCells = <Datum extends HeatMapDatum, ExtraProps extends obje
     })
 
     const xValues = Array.from(xValuesSet)
+
+    const { width, height, offsetX, offsetY } = computeLayout({
+        width: _width,
+        height: _height,
+        columns: xValues.length,
+        rows: serieIds.length,
+        forceSquare,
+    })
+
     const xScale = castBandScale<Datum['x']>(
         scaleBand<Datum['x']>()
             .domain(xValues)
@@ -78,10 +132,46 @@ export const computeCells = <Datum extends HeatMapDatum, ExtraProps extends obje
     }))
 
     return {
+        width,
+        height,
+        offsetX,
+        offsetY,
         xScale,
         yScale,
         minValue: Math.min(...allValues),
         maxValue: Math.max(...allValues),
         cells: cellsWithPosition,
     }
 }
+
+export const computeSizeScale = (
+    size: false | SizeVariationConfig,
+    min: number,
+    max: number
+): ((value: number | null) => number) => {
+    if (!size) return () => 1
+
+    const scale = scaleLinear()
+        .domain(size.values ? size.values : [min, max])
+        .range(size.sizes)
+
+    return (value: number | null) => {
+        if (value === null) return 1
+        return scale(value)
+    }
+}
+
+export const getCellAnnotationPosition = <Datum extends HeatMapDatum>(
+    cell: ComputedCell<Datum>
+) => ({
+    x: cell.x,
+    y: cell.y,
+})
+
+export const getCellAnnotationDimensions = <Datum extends HeatMapDatum>(
+    cell: ComputedCell<Datum>
+) => ({
+    size: Math.max(cell.width, cell.height),
+    width: cell.width,
+    height: cell.height,
+})