diff --git a/packages/heatmap/src/HeatMap.tsx b/packages/heatmap/src/HeatMap.tsx
index 32b1764dc5..42f481b9fa 100644
--- a/packages/heatmap/src/HeatMap.tsx
+++ b/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,
diff --git a/packages/heatmap/src/HeatMapCanvas.tsx b/packages/heatmap/src/HeatMapCanvas.tsx
index 631afa5ade..e0d0580eda 100644
--- a/packages/heatmap/src/HeatMapCanvas.tsx
+++ b/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,
diff --git a/packages/heatmap/src/compute.ts b/packages/heatmap/src/compute.ts
index 4bdf097a21..67af3aa571 100644
--- a/packages/heatmap/src/compute.ts
+++ b/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,6 +132,10 @@ export const computeCells = <Datum extends HeatMapDatum, ExtraProps extends obje
     }))
 
     return {
+        width,
+        height,
+        offsetX,
+        offsetY,
         xScale,
         yScale,
         minValue: Math.min(...allValues),
@@ -85,3 +143,35 @@ export const computeCells = <Datum extends HeatMapDatum, ExtraProps extends obje
         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,
+})
diff --git a/packages/heatmap/src/hooks.ts b/packages/heatmap/src/hooks.ts
index 533f45716d..513e07a244 100644
--- a/packages/heatmap/src/hooks.ts
+++ b/packages/heatmap/src/hooks.ts
@@ -1,5 +1,4 @@
 import { useMemo, useCallback, useState } from 'react'
-import { scaleLinear } from 'd3-scale'
 import { useTheme, usePropertyAccessor, useValueFormatter } from '@nivo/core'
 import { useInheritedColor, getContinuousColorScale } from '@nivo/colors'
 import { AnnotationMatcher, useAnnotations } from '@nivo/annotations'
@@ -12,7 +11,12 @@ import {
     SizeVariationConfig,
 } from './types'
 import { commonDefaultProps } from './defaults'
-import { computeCells } from './compute'
+import {
+    computeCells,
+    computeSizeScale,
+    getCellAnnotationPosition,
+    getCellAnnotationDimensions,
+} from './compute'
 
 export const useComputeCells = <Datum extends HeatMapDatum, ExtraProps extends object>({
     data,
@@ -22,15 +26,15 @@ export const useComputeCells = <Datum extends HeatMapDatum, ExtraProps extends o
     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'
+>) =>
     useMemo(
         () =>
             computeCells<Datum, ExtraProps>({
@@ -41,8 +45,18 @@ export const useComputeCells = <Datum extends HeatMapDatum, ExtraProps extends o
                 xOuterPadding,
                 yInnerPadding,
                 yOuterPadding,
+                forceSquare,
             }),
-        [data, width, height, xInnerPadding, xOuterPadding, yInnerPadding, yOuterPadding]
+        [
+            data,
+            width,
+            height,
+            xInnerPadding,
+            xOuterPadding,
+            yInnerPadding,
+            yOuterPadding,
+            forceSquare,
+        ]
     )
 
 const isHoverTargetByType = {
@@ -81,74 +95,45 @@ const useSizeScale = (
     min: number,
     max: number
 ): ((value: number | null) => number) =>
-    useMemo(() => {
-        if (!size) return () => 1
+    useMemo(() => computeSizeScale(size, min, max), [size, min, max])
 
-        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)
-        }
-    }, [size, min, max])
-
-export const useHeatMap = <
-    Datum extends HeatMapDatum = DefaultHeatMapDatum,
-    ExtraProps extends object = Record<string, never>
->({
-    data,
+const useCellsStyle = <Datum extends HeatMapDatum = DefaultHeatMapDatum>({
+    cells,
+    minValue,
+    maxValue,
+    sizeVariation,
+    colors,
+    emptyColor,
+    opacity,
+    activeOpacity,
+    inactiveOpacity,
+    borderColor,
+    label,
+    labelTextColor,
     valueFormat,
-    width,
-    height,
-    // forceSquare = commonDefaultProps.forceSquare,
-    xOuterPadding = commonDefaultProps.xOuterPadding,
-    xInnerPadding = commonDefaultProps.xInnerPadding,
-    yOuterPadding = commonDefaultProps.yOuterPadding,
-    yInnerPadding = commonDefaultProps.yInnerPadding,
-    sizeVariation = commonDefaultProps.sizeVariation,
-    colors = commonDefaultProps.colors as HeatMapCommonProps<Datum>['colors'],
-    emptyColor = commonDefaultProps.emptyColor,
-    opacity = commonDefaultProps.opacity,
-    activeOpacity = commonDefaultProps.activeOpacity,
-    inactiveOpacity = commonDefaultProps.inactiveOpacity,
-    borderColor = commonDefaultProps.borderColor as HeatMapCommonProps<Datum>['borderColor'],
-    label = commonDefaultProps.label as HeatMapCommonProps<Datum>['label'],
-    labelTextColor = commonDefaultProps.labelTextColor as HeatMapCommonProps<Datum>['labelTextColor'],
-    hoverTarget = commonDefaultProps.hoverTarget,
+    activeIds,
 }: {
-    data: HeatMapDataProps<Datum, ExtraProps>['data']
+    cells: Omit<
+        ComputedCell<Datum>,
+        'formattedValue' | 'color' | 'opacity' | 'borderColor' | 'label' | 'labelTextColor'
+    >[]
+    minValue: number
+    maxValue: number
     valueFormat?: HeatMapCommonProps<Datum>['valueFormat']
-    width: number
-    height: number
-    forceSquare?: HeatMapCommonProps<Datum>['forceSquare']
-    xOuterPadding?: HeatMapCommonProps<Datum>['xOuterPadding']
-    xInnerPadding?: HeatMapCommonProps<Datum>['xInnerPadding']
-    yOuterPadding?: HeatMapCommonProps<Datum>['yOuterPadding']
-    yInnerPadding?: HeatMapCommonProps<Datum>['yInnerPadding']
-    sizeVariation?: HeatMapCommonProps<Datum>['sizeVariation']
-    colors?: HeatMapCommonProps<Datum>['colors']
-    emptyColor?: HeatMapCommonProps<Datum>['emptyColor']
-    opacity?: HeatMapCommonProps<Datum>['opacity']
-    activeOpacity?: HeatMapCommonProps<Datum>['activeOpacity']
-    inactiveOpacity?: HeatMapCommonProps<Datum>['inactiveOpacity']
-    borderColor?: HeatMapCommonProps<Datum>['borderColor']
-    label?: HeatMapCommonProps<Datum>['label']
-    labelTextColor?: HeatMapCommonProps<Datum>['labelTextColor']
-    hoverTarget?: HeatMapCommonProps<Datum>['hoverTarget']
-}) => {
-    const [activeCell, setActiveCell] = useState<ComputedCell<Datum> | null>(null)
-
-    const { cells, xScale, yScale, minValue, maxValue } = useComputeCells<Datum, ExtraProps>({
-        data,
-        width,
-        height,
-        xOuterPadding,
-        xInnerPadding,
-        yOuterPadding,
-        yInnerPadding,
-    })
+    activeIds: string[]
+} & Pick<
+    HeatMapCommonProps<Datum>,
+    | 'sizeVariation'
+    | 'colors'
+    | 'emptyColor'
+    | 'opacity'
+    | 'activeOpacity'
+    | 'inactiveOpacity'
+    | 'borderColor'
+    | 'label'
+    | 'labelTextColor'
+>) => {
+    const getSize = useSizeScale(sizeVariation, minValue, maxValue)
 
     const colorScale = useMemo(() => {
         if (typeof colors === 'function') return null
@@ -159,14 +144,6 @@ export const useHeatMap = <
         })
     }, [colors, minValue, maxValue])
 
-    const activeIds = useMemo(() => {
-        if (!activeCell) return []
-
-        const isHoverTarget = isHoverTargetByType[hoverTarget]
-
-        return cells.filter(cell => isHoverTarget(cell, activeCell)).map(cell => cell.id)
-    }, [cells, activeCell, hoverTarget])
-
     const getColor = useCallback(
         (cell: Omit<ComputedCell<Datum>, 'color' | 'opacity' | 'borderColor'>) => {
             if (cell.value !== null) {
@@ -178,14 +155,14 @@ export const useHeatMap = <
         },
         [colors, colorScale, emptyColor]
     )
-    const getSize = useSizeScale(sizeVariation, minValue, maxValue)
     const theme = useTheme()
     const getBorderColor = useInheritedColor(borderColor, theme)
     const getLabelTextColor = useInheritedColor(labelTextColor, theme)
+
     const formatValue = useValueFormatter(valueFormat)
     const getLabel = usePropertyAccessor(label)
 
-    const computedCells = useMemo(
+    const styledCells = useMemo(
         () =>
             cells.map(cell => {
                 let computedOpacity = opacity
@@ -226,53 +203,111 @@ export const useHeatMap = <
     )
 
     return {
-        cells: computedCells,
-        xScale,
-        yScale,
+        cells: styledCells,
         colorScale,
-        activeCell,
-        setActiveCell,
     }
+}
 
-    /*
-    const layoutConfig = useMemo(() => {
-        const columns = keys.length
-        const rows = data.length
+export const useHeatMap = <
+    Datum extends HeatMapDatum = DefaultHeatMapDatum,
+    ExtraProps extends object = Record<string, never>
+>({
+    data,
+    valueFormat,
+    width: _width,
+    height: _height,
+    xOuterPadding = commonDefaultProps.xOuterPadding,
+    xInnerPadding = commonDefaultProps.xInnerPadding,
+    yOuterPadding = commonDefaultProps.yOuterPadding,
+    yInnerPadding = commonDefaultProps.yInnerPadding,
+    forceSquare = commonDefaultProps.forceSquare,
+    sizeVariation = commonDefaultProps.sizeVariation,
+    colors = commonDefaultProps.colors as HeatMapCommonProps<Datum>['colors'],
+    emptyColor = commonDefaultProps.emptyColor,
+    opacity = commonDefaultProps.opacity,
+    activeOpacity = commonDefaultProps.activeOpacity,
+    inactiveOpacity = commonDefaultProps.inactiveOpacity,
+    borderColor = commonDefaultProps.borderColor as HeatMapCommonProps<Datum>['borderColor'],
+    label = commonDefaultProps.label as HeatMapCommonProps<Datum>['label'],
+    labelTextColor = commonDefaultProps.labelTextColor as HeatMapCommonProps<Datum>['labelTextColor'],
+    hoverTarget = commonDefaultProps.hoverTarget,
+}: {
+    data: HeatMapDataProps<Datum, ExtraProps>['data']
+    width: number
+    height: number
+} & Partial<
+    Pick<
+        HeatMapCommonProps<Datum>,
+        | 'valueFormat'
+        | 'xOuterPadding'
+        | 'xInnerPadding'
+        | 'yOuterPadding'
+        | 'yInnerPadding'
+        | 'forceSquare'
+        | 'sizeVariation'
+        | 'colors'
+        | 'emptyColor'
+        | 'opacity'
+        | 'activeOpacity'
+        | 'inactiveOpacity'
+        | 'borderColor'
+        | 'label'
+        | 'labelTextColor'
+        | 'hoverTarget'
+    >
+>) => {
+    const [activeCell, setActiveCell] = useState<ComputedCell<Datum> | null>(null)
 
-        let cellWidth = Math.max((width - padding * (columns + 1)) / columns, 0)
-        let cellHeight = Math.max((height - padding * (rows + 1)) / rows, 0)
+    const { width, height, offsetX, offsetY, cells, xScale, yScale, minValue, maxValue } =
+        useComputeCells<Datum, ExtraProps>({
+            data,
+            width: _width,
+            height: _height,
+            xOuterPadding,
+            xInnerPadding,
+            yOuterPadding,
+            yInnerPadding,
+            forceSquare,
+        })
 
-        let offsetX = 0
-        let offsetY = 0
-        if (forceSquare === true) {
-            const cellSize = Math.min(cellWidth, cellHeight)
-            cellWidth = cellSize
-            cellHeight = cellSize
+    const activeIds = useMemo(() => {
+        if (!activeCell) return []
 
-            offsetX = (width - ((cellWidth + padding) * columns + padding)) / 2
-            offsetY = (height - ((cellHeight + padding) * rows + padding)) / 2
-        }
+        const isHoverTarget = isHoverTargetByType[hoverTarget]
 
-        return {
-            cellWidth,
-            cellHeight,
-            offsetX,
-            offsetY,
-        }
-    }, [data, keys, width, height, padding, forceSquare])
-    */
-}
+        return cells.filter(cell => isHoverTarget(cell, activeCell)).map(cell => cell.id)
+    }, [cells, activeCell, hoverTarget])
 
-const getCellAnnotationPosition = <Datum extends HeatMapDatum>(cell: ComputedCell<Datum>) => ({
-    x: cell.x,
-    y: cell.y,
-})
+    const { cells: computedCells, colorScale } = useCellsStyle<Datum>({
+        cells,
+        minValue,
+        maxValue,
+        sizeVariation,
+        colors,
+        emptyColor,
+        opacity,
+        activeOpacity,
+        inactiveOpacity,
+        borderColor,
+        label,
+        labelTextColor,
+        valueFormat,
+        activeIds,
+    })
 
-const getCellAnnotationDimensions = <Datum extends HeatMapDatum>(cell: ComputedCell<Datum>) => ({
-    size: Math.max(cell.width, cell.height),
-    width: cell.width,
-    height: cell.height,
-})
+    return {
+        width,
+        height,
+        offsetX,
+        offsetY,
+        cells: computedCells,
+        xScale,
+        yScale,
+        colorScale,
+        activeCell,
+        setActiveCell,
+    }
+}
 
 export const useCellAnnotations = <Datum extends HeatMapDatum>(
     cells: ComputedCell<Datum>[],
diff --git a/packages/heatmap/src/index.ts b/packages/heatmap/src/index.ts
index 7ad7f8d3ee..0cae12227e 100644
--- a/packages/heatmap/src/index.ts
+++ b/packages/heatmap/src/index.ts
@@ -5,3 +5,4 @@ export * from './ResponsiveHeatMapCanvas'
 export * from './hooks'
 export * from './defaults'
 export * from './types'
+export * from './compute'