feat(swarmplot): add support for log axis and legends

packages/swarmplot/src/SwarmPlot.tsx

@@ -10,6 +10,7 @@ import { useSwarmPlot, useSwarmPlotLayerContext, useNodeMouseHandlers } from './
 import { Circles } from './Circles'
 import { CircleSvg } from './CircleSvg'
 import { SwarmPlotAnnotations } from './SwarmPlotAnnotations'
+import { SwarmPlotLegends } from './SwarmPlotLegends'
 
 type InnerSwarmPlotProps<RawDatum> = Partial<
     Omit<
@@ -52,6 +53,8 @@ const InnerSwarmPlot = <RawDatum,>({
     axisRight = defaultProps.axisRight,
     axisBottom = defaultProps.axisBottom,
     axisLeft = defaultProps.axisLeft,
+    legendLabel,
+    legends = defaultProps.legends,
     isInteractive,
     onMouseEnter,
     onMouseMove,
@@ -67,7 +70,7 @@ const InnerSwarmPlot = <RawDatum,>({
         partialMargin
     )
 
-    const { nodes, ...props } = useSwarmPlot<RawDatum>({
+    const { nodes, legendsData, ...props } = useSwarmPlot<RawDatum>({
         width: innerWidth,
         height: innerHeight,
         data,
@@ -85,6 +88,8 @@ const InnerSwarmPlot = <RawDatum,>({
         colorBy,
         forceStrength,
         simulationIterations,
+        legendLabel,
+        legends,
     })
 
     const xScale = props.xScale as Exclude<typeof props.xScale, ComputedDatum<RawDatum>[]>
@@ -105,6 +110,7 @@ const InnerSwarmPlot = <RawDatum,>({
         circles: null,
         annotations: null,
         mesh: null,
+        legends: null,
     }
 
     if (layers.includes('grid')) {
@@ -165,6 +171,17 @@ const InnerSwarmPlot = <RawDatum,>({
         )
     }
 
+    if (layers.includes('legends')) {
+        layerById.legends = (
+            <SwarmPlotLegends
+                key="legends"
+                width={innerWidth}
+                height={innerHeight}
+                legends={legendsData}
+            />
+        )
+    }
+
     if (isInteractive && useMesh) {
         layerById.mesh = (
             <Mesh

packages/swarmplot/src/SwarmPlotCanvas.tsx

@@ -1,4 +1,4 @@
-import { createElement, useCallback, useEffect, useRef, useState } from 'react'
+import { createElement, useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import * as React from 'react'
 import isNumber from 'lodash/isNumber'
 import { Container, getRelativeCursor, isCursorInRect, useDimensions, useTheme } from '@nivo/core'
@@ -10,6 +10,7 @@ import { useVoronoiMesh, renderVoronoiToCanvas, renderVoronoiCellToCanvas } from
 import { ComputedDatum, SwarmPlotCanvasProps } from './types'
 import { defaultProps } from './props'
 import { useSwarmPlot } from './hooks'
+import { renderLegendToCanvas } from '@nivo/legends'
 
 export const renderCircleDefault = <RawDatum,>(
     ctx: CanvasRenderingContext2D,
@@ -78,6 +79,8 @@ export const InnerSwarmPlotCanvas = <RawDatum,>({
     axisRight = defaultProps.axisRight,
     axisBottom = defaultProps.axisBottom,
     axisLeft = defaultProps.axisLeft,
+    legendLabel,
+    legends = defaultProps.legends,
     isInteractive,
     onMouseMove,
     onClick,
@@ -95,7 +98,7 @@ export const InnerSwarmPlotCanvas = <RawDatum,>({
         partialMargin
     )
 
-    const { nodes, ...scales } = useSwarmPlot<RawDatum>({
+    const { nodes, legendsData, ...scales } = useSwarmPlot<RawDatum>({
         width: innerWidth,
         height: innerHeight,
         data,
@@ -113,6 +116,8 @@ export const InnerSwarmPlotCanvas = <RawDatum,>({
         colorBy,
         forceStrength,
         simulationIterations,
+        legendLabel,
+        legends,
     })
 
     const { xScale, yScale } = scales as Record<'xScale' | 'yScale', AnyScale>
@@ -125,7 +130,8 @@ export const InnerSwarmPlotCanvas = <RawDatum,>({
     })
 
     const getBorderColor = useInheritedColor(borderColor, theme)
-    const getBorderWidth = () => 1
+    // memoize getBorderWidth to provide a stable object for the rendering
+    const getBorderWidth = useMemo(() => () => 1, [])
 
     useEffect(() => {
         if (!canvasEl.current) return
@@ -203,6 +209,18 @@ export const InnerSwarmPlotCanvas = <RawDatum,>({
                     renderVoronoiCellToCanvas(ctx, voronoi, currentNode.index)
                 }
             }
+
+            if (layer === 'legends') {
+                legendsData.forEach(([legend, data]) => {
+                    renderLegendToCanvas(ctx, {
+                        ...legend,
+                        data,
+                        containerWidth: innerWidth,
+                        containerHeight: innerHeight,
+                        theme,
+                    })
+                })
+            }
         })
     }, [
         canvasEl,
@@ -231,6 +249,7 @@ export const InnerSwarmPlotCanvas = <RawDatum,>({
         renderCircle,
         getBorderWidth,
         getBorderColor,
+        legendsData,
     ])
 
     const getNodeFromMouseEvent = useCallback(

packages/swarmplot/src/SwarmPlotLegends.tsx

@@ -0,0 +1,23 @@
+import { BoxLegendSvg } from '@nivo/legends'
+import { LegendProps } from '@nivo/legends'
+import { SwarmPlotLegendData } from './types'
+
+interface SwarmPlotLegendsProps {
+    width: number
+    height: number
+    legends: [LegendProps, SwarmPlotLegendData[]][]
+}
+
+export const SwarmPlotLegends = ({ width, height, legends }: SwarmPlotLegendsProps) => (
+    <>
+        {legends.map(([legend, data], i) => (
+            <BoxLegendSvg
+                key={i}
+                {...legend}
+                containerWidth={width}
+                containerHeight={height}
+                data={data}
+            />
+        ))}
+    </>
+)

packages/swarmplot/src/compute.ts

@@ -1,21 +1,22 @@
 import isNumber from 'lodash/isNumber'
 import isPlainObject from 'lodash/isPlainObject'
 import isString from 'lodash/isString'
+import uniqBy from 'lodash/uniqBy'
 import get from 'lodash/get'
 import { scaleLinear, ScaleOrdinal, scaleOrdinal } from 'd3-scale'
 import { forceSimulation, forceX, forceY, forceCollide, ForceX, ForceY } from 'd3-force'
+import { computeScale, createDateNormalizer, generateSeriesAxis } from '@nivo/scales'
 import {
-    computeScale,
-    createDateNormalizer,
-    generateSeriesAxis,
-    ScaleLinear,
-    ScaleLinearSpec,
-    ScaleTime,
-    ScaleTimeSpec,
-} from '@nivo/scales'
-import { ComputedDatum, PreSimulationDatum, SizeSpec, SimulationForces } from './types'
-
-const getParsedValue = (scaleSpec: ScaleLinearSpec | ScaleTimeSpec) => {
+    ComputedDatum,
+    PreSimulationDatum,
+    SizeSpec,
+    SimulationForces,
+    SwarmPlotLegendData,
+    SwarmPlotValueScaleSpec,
+    SwarmPlotValueScale,
+} from './types'
+
+const getParsedValue = (scaleSpec: SwarmPlotValueScaleSpec) => {
     if (scaleSpec.type === 'time' && scaleSpec.format !== 'native') {
         return createDateNormalizer(scaleSpec) as <T>(value: T) => T
     }
@@ -66,9 +67,9 @@ export const computeValueScale = <RawDatum>({
     height: number
     axis: 'x' | 'y'
     getValue: (datum: RawDatum) => number | Date
-    scale: ScaleLinearSpec | ScaleTimeSpec
+    scale: SwarmPlotValueScaleSpec
     data: RawDatum[]
-}) => {
+}): SwarmPlotValueScale => {
     const values = data.map(getValue)
 
     if (scale.type === 'time') {
@@ -77,9 +78,7 @@ export const computeValueScale = <RawDatum>({
         ]
         const axes = generateSeriesAxis(series, axis, scale)
 
-        return computeScale(scale, axes, axis === 'x' ? width : height, axis) as ScaleTime<
-            Date | string
-        >
+        return computeScale(scale, axes, axis === 'x' ? width : height, axis) as SwarmPlotValueScale
     }
 
     const min = Math.min(...(values as number[]))
@@ -90,7 +89,7 @@ export const computeValueScale = <RawDatum>({
         { all: values, min, max },
         axis === 'x' ? width : height,
         axis
-    ) as ScaleLinear<number>
+    ) as SwarmPlotValueScale
 }
 
 export const getSizeGenerator = <RawDatum>(size: SizeSpec<RawDatum>) => {
@@ -140,7 +139,7 @@ export const computeForces = <RawDatum>({
     forceStrength,
 }: {
     axis: 'x' | 'y'
-    valueScale: ScaleLinear<number> | ScaleTime<string | Date>
+    valueScale: SwarmPlotValueScale
     ordinalScale: ScaleOrdinal<string, number>
     spacing: number
     forceStrength: number
@@ -183,13 +182,13 @@ export const computeNodes = <RawDatum>({
     getId: (datum: RawDatum) => string
     layout: 'vertical' | 'horizontal'
     getValue: (datum: RawDatum) => number | Date
-    valueScale: ScaleLinear<number> | ScaleTime<string | Date>
+    valueScale: SwarmPlotValueScale
     getGroup: (datum: RawDatum) => string
     ordinalScale: ScaleOrdinal<string, number>
     getSize: (datum: RawDatum) => number
     forces: SimulationForces<RawDatum>
     simulationIterations: number
-    valueScaleConfig: ScaleLinearSpec | ScaleTimeSpec
+    valueScaleConfig: SwarmPlotValueScaleSpec
 }) => {
     const config = {
         horizontal: ['x', 'y'],
@@ -220,3 +219,21 @@ export const computeNodes = <RawDatum>({
         nodes: simulation.nodes() as ComputedDatum<RawDatum>[],
     }
 }
+
+export const getBaseLegendData = <RawDatum>({
+    nodes,
+    getLegendLabel,
+}: {
+    nodes: ComputedDatum<RawDatum>[]
+    getLegendLabel: (datum: RawDatum) => string
+}) => {
+    const nodeData = nodes.map(
+        node =>
+            ({
+                id: node.group,
+                label: getLegendLabel(node.data),
+                color: node?.color,
+            } as SwarmPlotLegendData)
+    )
+    return uniqBy(nodeData, ({ id }) => id)
+}

packages/swarmplot/src/hooks.ts

@@ -4,21 +4,25 @@ import { usePropertyAccessor, useValueFormatter } from '@nivo/core'
 import { useOrdinalColorScale } from '@nivo/colors'
 import { AnnotationMatcher, useAnnotations } from '@nivo/annotations'
 import { useTooltip } from '@nivo/tooltip'
-import { ScaleLinear, ScaleLinearSpec, ScaleTime, ScaleTimeSpec } from '@nivo/scales'
 import {
     computeValueScale,
     computeOrdinalScale,
     getSizeGenerator,
+    getBaseLegendData,
     computeForces,
     computeNodes,
 } from './compute'
 import {
     SwarmPlotCommonProps,
+    SwarmPlotLegendData,
     ComputedDatum,
     SizeSpec,
     SwarmPlotCustomLayerProps,
     MouseHandlers,
+    SwarmPlotValueScale,
+    SwarmPlotValueScaleSpec,
 } from './types'
+import { LegendProps } from '@nivo/legends'
 
 export const useValueScale = <RawDatum>({
     width,
@@ -32,7 +36,7 @@ export const useValueScale = <RawDatum>({
     height: number
     axis: 'x' | 'y'
     getValue: (datum: RawDatum) => number | Date
-    scale: ScaleLinearSpec | ScaleTimeSpec
+    scale: SwarmPlotValueScaleSpec
     data: RawDatum[]
 }) =>
     useMemo(
@@ -77,7 +81,7 @@ export const useForces = <RawDatum>({
     forceStrength,
 }: {
     axis: 'x' | 'y'
-    valueScale: ScaleLinear<number> | ScaleTime<string | Date>
+    valueScale: SwarmPlotValueScale
     ordinalScale: ScaleOrdinal<string, number>
     spacing: number
     forceStrength: number
@@ -112,6 +116,8 @@ export const useSwarmPlot = <RawDatum>({
     simulationIterations,
     colors,
     colorBy,
+    legendLabel,
+    legends,
 }: {
     data: RawDatum[]
     width: number
@@ -130,6 +136,8 @@ export const useSwarmPlot = <RawDatum>({
     simulationIterations: SwarmPlotCommonProps<RawDatum>['simulationIterations']
     colors: SwarmPlotCommonProps<RawDatum>['colors']
     colorBy: SwarmPlotCommonProps<RawDatum>['colorBy']
+    legendLabel: SwarmPlotCommonProps<RawDatum>['legendLabel']
+    legends: SwarmPlotCommonProps<RawDatum>['legends']
 }) => {
     const axis = layout === 'horizontal' ? 'x' : 'y'
 
@@ -143,6 +151,7 @@ export const useSwarmPlot = <RawDatum>({
         colors,
         getColorId
     )
+    const getLegendLabel = usePropertyAccessor<RawDatum, string>(legendLabel ?? groupBy)
 
     const valueScale = useValueScale({
         width,
@@ -209,11 +218,25 @@ export const useSwarmPlot = <RawDatum>({
         [nodes, formatValue, getColor]
     )
 
+    const legendsData: [LegendProps, SwarmPlotLegendData[]][] = useMemo(
+        () =>
+            legends.map(legend => {
+                if (legend.data) return [legend, legend.data as SwarmPlotLegendData[]]
+                const data = getBaseLegendData({
+                    nodes: augmentedNodes,
+                    getLegendLabel,
+                })
+                return [legend, data]
+            }),
+        [legends, augmentedNodes, getLegendLabel]
+    )
+
     return {
         nodes: augmentedNodes,
         xScale,
         yScale,
         getColor,
+        legendsData,
     }
 }
 
@@ -306,10 +329,7 @@ export const useSwarmPlotAnnotations = <RawDatum>(
 
 export const useSwarmPlotLayerContext = <
     RawDatum,
-    Scale extends
-        | ScaleLinear<number>
-        | ScaleTime<string | Date>
-        | ScaleOrdinal<string, number, never>
+    Scale extends SwarmPlotValueScale | ScaleOrdinal<string, number, never>
 >({
     nodes,
     xScale,