Update Donut to handle arrays of colors

docs/package.json

@@ -30,6 +30,7 @@
     "@carbonplan/prism": "^1.0.3",
     "@carbonplan/theme": "^6.0.1",
     "@mdx-js/loader": "^1.6.22",
+    "@mdx-js/react": "^1.6.22",
     "@next/mdx": "^10.2.3",
     "d3-scale": "^3.3.0",
     "d3-shape": "^2.1.0",

docs/pages/stacked-bar.md

@@ -10,16 +10,43 @@ import {
   StackedBar,
 } from '@carbonplan/charts'
 
-export const data = Array(11)
-  .fill(null)
-  .map((_, i) => [i, 0, (i + 1) * 2, (i + 1) * 5, (i + 1) * 9])
+export const data = [
+  [0, 0, 2, 5, 9],
+  [1, 0, 4, 10, 18],
+  [2, 0, 6, 15, 27],
+  [3, 0, 8, 20, 36],
+  [4, 0, 10, 25, 45],
+  [5, 0, 12, 30, 54],
+  [6, 0, 14, 35, 63],
+  [7, 0, 16, 40, 72],
+  [8, 0, 18, 45, 81],
+  [9, 0, 20, 50, 90],
+  [10, 0, 22, 55, 99],
+]
 export const getRandomColor = () =>
   ['pink', 'red', 'orange', 'yellow', 'green'][Math.floor(Math.random() * 5)]
 
 # StackedBar
 
 This is a simple bar chart.
 
+```js
+const data = [
+  // x,  y0, y1, y2, y3
+  [0, 0, 2, 5, 9],
+  [1, 0, 4, 10, 18],
+  [2, 0, 6, 15, 27],
+  [3, 0, 8, 20, 36],
+  [4, 0, 10, 25, 45],
+  [5, 0, 12, 30, 54],
+  [6, 0, 14, 35, 63],
+  [7, 0, 16, 40, 72],
+  [8, 0, 18, 45, 81],
+  [9, 0, 20, 50, 90],
+  [10, 0, 22, 55, 99],
+]
+```
+
 <Box sx={{ width: '100%', height: '400px' }}>
   <Chart x={[-1, 11]} y={[0, 100]} padding={{ left: 60, top: 50 }}>
     <Ticks left bottom />
@@ -166,6 +193,66 @@ This is a simple bar chart.
 </Box>
 ```
 
+### Custom opacity
+
+<Box sx={{ width: '100%', height: '400px' }}>
+  <Chart x={[-1, 11]} y={[0, 100]} padding={{ left: 60, top: 50 }}>
+    <Ticks left bottom />
+    <TickLabels left bottom />
+    <Axis left bottom />
+    <Plot>
+      <StackedBar data={data} color='purple' opacity={[0.1, 0.2, 0.3]} />
+    </Plot>
+  </Chart>
+</Box>
+
+```jsx
+<Box sx={{ width: '100%', height: '400px' }}>
+  <Chart x={[-1, 11]} y={[0, 100]} padding={{ left: 60, top: 50 }}>
+    <Ticks left bottom />
+    <TickLabels left bottom />
+    <Axis left bottom />
+    <Plot>
+      <StackedBar data={data} color='purple' opacity={[0.1, 0.2, 0.3]} />
+    </Plot>
+  </Chart>
+</Box>
+```
+
+#### Completely customized opacity
+
+<Box sx={{ width: '100%', height: '400px' }}>
+  <Chart x={[-1, 11]} y={[0, 100]} padding={{ left: 60, top: 50 }}>
+    <Ticks left bottom />
+    <TickLabels left bottom />
+    <Axis left bottom />
+    <Plot>
+      <StackedBar
+        data={data}
+        color='purple'
+        opacity={data.map((d) => d.slice(2).map((_, i) => Math.random()))}
+      />
+    </Plot>
+  </Chart>
+</Box>
+
+```jsx
+<Box sx={{ width: '100%', height: '400px' }}>
+  <Chart x={[-1, 11]} y={[0, 100]} padding={{ left: 60, top: 50 }}>
+    <Ticks left bottom />
+    <TickLabels left bottom />
+    <Axis left bottom />
+    <Plot>
+      <StackedBar
+        data={data}
+        color='purple'
+        opacity={data.map((d) => d.slice(2).map(() => Math.random()))}
+      />
+    </Plot>
+  </Chart>
+</Box>
+```
+
 export default ({ children }) => (
   <Section name='stacked-bar'>{children}</Section>
 )

src/bar.js

@@ -1,12 +1,14 @@
 import React, { memo, useMemo } from 'react'
 import { useThemeUI } from 'theme-ui'
 import { useChart } from './chart'
+import { getColorAtIndex, getPropAtIndex } from './utils'
 
 const Bar = ({
   data,
   width = 0.8,
   direction = 'vertical',
   color = 'primary',
+  opacity = 1,
   ...props
 }) => {
   const { x: _x, y: _y } = useChart()
@@ -33,11 +35,6 @@ const Bar = ({
     [xValues.join(',')]
   )
   const fixedWidth = minDelta * width
-  if (Array.isArray(color) && color.length !== data.length) {
-    throw new Error(
-      `Unexpected color array provided. Expected length ${data.length}, received length ${color.length}`
-    )
-  }
 
   return (
     <>
@@ -59,14 +56,16 @@ const Bar = ({
         const [x, y] = position
         const [width, height] = dimensions
 
-        const colorString = typeof color === 'string' ? color : color[i]
-        const fill = theme.rawColors[colorString] || colorString
-
         return (
           <path
             key={i}
             d={`M ${x} ${y} h ${width} v ${height} h -${width} Z`}
-            fill={fill}
+            fill={getColorAtIndex(color, data, i, {
+              colors: theme.rawColors,
+            })}
+            fillOpacity={getPropAtIndex(opacity, data, i, {
+              propName: 'opacity',
+            })}
             stroke='none'
             {...props}
           />

src/donut.js

@@ -2,23 +2,35 @@ import React, { memo } from 'react'
 import { arc, pie } from 'd3-shape'
 import { scaleLinear } from 'd3-scale'
 import { Box } from 'theme-ui'
+import { getPropAtIndex, getColorAtIndex } from './utils'
 
 const Donut = ({
   data,
-  domain,
-  range,
   innerRadius = 0.3,
   outerRadius = 50,
+  opacity,
   color = 'primary',
+  preserveOrder = false,
   sx,
 }) => {
-  domain = domain || [0, data.length - 1]
-  range = range || [0.3, 0.9]
-  const arcs = pie()(data)
+  const arcs = pie()
+    .sort(preserveOrder ? (a, b) => a.i - b.i : null)
+    .value((d) => d.value)(data.map((d, i) => ({ value: d, i })))
   const generator = arc()
     .innerRadius(innerRadius * 100)
     .outerRadius(outerRadius)
-  const opacity = scaleLinear().domain(domain).range(range)
+
+  let defaultOpacity
+  if (opacity == null) {
+    if (Array.isArray(color)) {
+      defaultOpacity = 1
+    } else {
+      const opacityScale = scaleLinear()
+        .domain([0, data.length - 1])
+        .range([0.3, 0.9])
+      defaultOpacity = arcs.map((d) => opacityScale(d.index))
+    }
+  }
 
   return (
     <g transform='translate(50,50)'>
@@ -30,8 +42,10 @@ const Donut = ({
             d={generator(d)}
             sx={{
               stroke: 'none',
-              fillOpacity: opacity(d.index),
-              fill: color,
+              fillOpacity: getPropAtIndex(opacity ?? defaultOpacity, data, i, {
+                propName: 'opacity',
+              }),
+              fill: getColorAtIndex(color, data, i),
               ...sx,
             }}
           />

src/stacked-bar.js

@@ -3,51 +3,67 @@ import { scaleLinear } from 'd3-scale'
 
 import Bar from './bar'
 
-const getBarColors = (color, barLength, dataLength, opacityRange) => {
-  let colors
-  let opacities = []
-  if (typeof color === 'string') {
-    // single color for all bars has been provided, use same color with different opacity per bar
-    colors = new Array(barLength).fill(color)
-    const opacity = scaleLinear()
-      .domain([barLength - 1, 0])
-      .range(opacityRange)
-    opacities = new Array(barLength).fill(null).map((_, i) => opacity(i))
-  } else if (color.every((c) => typeof c === 'string')) {
-    // color has been specified for each bar
-    if (color.length !== barLength) {
-      throw new Error(
-        `Unexpected 1D color array provided. Expected length equal to number of bars: ${barLength}, received length: ${color.length}`
-      )
-    }
-    colors = color
-  } else {
-    // color has been specified for each datum
-    if (color.length !== dataLength) {
-      throw new Error(
-        `Unexpected 2D color array provided. Expected length equal to data: ${dataLength}, received length: ${color.length}`
-      )
-    }
-    const invalidSubarray = color.find((c) => c.length !== barLength)
-    if (invalidSubarray) {
-      throw new Error(
-        `Unexpected 2D color array provided. Expected all subarrays to have length equal to number of bars: ${barLength}, received length: ${invalidSubarray.length}`
+const normalizeProps = (props, barLength, dataLength) => {
+  return Object.keys(props).reduce((accum, propName) => {
+    const prop = props[propName]
+    if (!Array.isArray(prop)) {
+      // single prop for all bars has been provided, use same prop with different opacity per bar
+      accum[propName] = new Array(barLength).fill(prop)
+    } else if (prop.every((p) => !Array.isArray(p))) {
+      // prop has been specified for each bar
+      if (prop.length !== barLength) {
+        throw new Error(
+          `Unexpected 1D ${propName} array provided. Expected length equal to number of bars: ${barLength}, received length: ${prop.length}`
+        )
+      }
+
+      accum[propName] = prop
+    } else {
+      // prop has been specified for each datum
+      if (prop.length !== dataLength) {
+        throw new Error(
+          `Unexpected 2D ${propName} array provided. Expected length equal to data: ${dataLength}, received length: ${prop.length}`
+        )
+      }
+      const invalidSubarray = prop.find((p) => p.length !== barLength)
+      if (invalidSubarray) {
+        throw new Error(
+          `Unexpected 2D ${propName} array provided. Expected all subarrays to have length equal to number of bars: ${barLength}, received length: ${invalidSubarray.length}`
+        )
+      }
+
+      accum[propName] = prop.reduce(
+        (accum, datum) => {
+          datum.forEach((barProp, i) => accum[i].push(barProp))
+          return accum
+        },
+        new Array(barLength).fill(null).map(() => [])
       )
     }
 
-    colors = color.reduce(
-      (accum, datum) => {
-        datum.forEach((barColor, i) => accum[i].push(barColor))
-        return accum
-      },
-      new Array(barLength).fill(null).map(() => [])
-    )
-  }
+    return accum
+  }, {})
+}
+
+const getBarColors = (color, opacity, barLength, dataLength) => {
+  const normalized = normalizeProps({ color, opacity }, barLength, dataLength)
+
+  // if same opacity and color provided, use same color with different opacity per bar
+  if (!Array.isArray(color) && !Array.isArray(opacity)) {
+    const opacityScale = scaleLinear()
+      .domain([barLength - 1, 0])
+      .range([0.3, 0.9])
+    const opacities = new Array(barLength)
+      .fill(null)
+      .map((_, i) => opacityScale(i))
 
-  return { colors, opacities }
+    return { colors: normalized.color, opacities }
+  } else {
+    return { colors: normalized.color, opacities: normalized.opacity }
+  }
 }
 
-const StackedBar = ({ data, color = 'primary', range, ...props }) => {
+const StackedBar = ({ data, color = 'primary', opacity, ...props }) => {
   const bars = useMemo(() => {
     const stackedData = data[0].slice(2).map(() => [])
     return data.reduce((accum, datum) => {
@@ -71,8 +87,8 @@ const StackedBar = ({ data, color = 'primary', range, ...props }) => {
   }, [data])
 
   const { colors, opacities } = useMemo(
-    () => getBarColors(color, bars.length, data.length, range || [0.3, 0.9]),
-    [color, bars.length, data.length, range]
+    () => getBarColors(color, opacity, bars.length, data.length),
+    [color, opacity, bars.length, data.length]
   )
 
   return (
@@ -83,7 +99,7 @@ const StackedBar = ({ data, color = 'primary', range, ...props }) => {
             key={i}
             data={bar}
             color={colors[i]}
-            fillOpacity={opacities[i]}
+            opacity={opacities[i]}
             {...props}
           />
         )

src/utils/index-getters.js

@@ -0,0 +1,24 @@
+export const getPropAtIndex = (prop, data, i, options = {}) => {
+  if (Array.isArray(prop)) {
+    if (prop.length !== data.length) {
+      throw new Error(
+        `Unexpected ${options.propName} array provided. Expected length ${data.length}, received length ${prop.length}`
+      )
+    }
+    return prop[i]
+  } else {
+    return prop
+  }
+}
+
+export const getColorAtIndex = (color, data, i, options = {}) => {
+  const rawValue = getPropAtIndex(color, data, i, {
+    ...options,
+    propName: options.propName || 'color',
+  })
+  if (options.colors && options.colors[rawValue]) {
+    return options.colors[rawValue]
+  } else {
+    return rawValue
+  }
+}

src/utils/index.js

@@ -1,5 +1,2 @@
-import getTicks from './get-ticks'
-
-export default {
-  getTicks,
-}
+export { default as getTicks } from './get-ticks'
+export { getPropAtIndex, getColorAtIndex } from './index-getters'