d3-sankey-diagram/src/diagram.js

Sankey diagram d3 plugin

// The reusable SVG component for the sliced Sankey diagram

import sankeyLink from './linkPath.js'
import sankeyNode from './node.js'
import positionGroup from './positionGroup.js'

import { select, event } from 'd3-selection'
import { transition } from 'd3-transition'
import { dispatch } from 'd3-dispatch'
import { format } from 'd3-format'
import { interpolate } from 'd3-interpolate'
import { map } from 'd3-collection'

export function linkTitleGenerator (nodeTitle, typeTitle, fmt) {
  return function (d) {
    const parts = []
    const sourceTitle = nodeTitle(d.source)
    const targetTitle = nodeTitle(d.target)
    const matTitle = typeTitle(d)

    parts.push(`${sourceTitle} → ${targetTitle}`)
    if (matTitle) parts.push(matTitle)
    parts.push(fmt(d.value))
    return parts.join('\n')
  }
}

export default function sankeyDiagram () {
  let margin = { top: 0, right: 0, bottom: 0, left: 0 }

  // eslint-disable-next-line no-unused-vars -- FIXME should set class for nodes
  let selectedNode = null
  let selectedEdge = null

  let groups = []

  const fmt = format('.3s')
  const node = sankeyNode()
  const link = sankeyLink()

  let linkColor = d => null
  let linkTitle = linkTitleGenerator(node.nodeTitle(), d => d.type, fmt)
  let linkLabel = defaultLinkLabel
  let linkImportance = defaultLinkImportance
  let linkImportanceAgg = defaultLinkImportanceAgg

  const listeners = dispatch('selectNode', 'selectGroup', 'selectLink')

  /* Main chart */

  function exports (context) {
    const selection = context.selection ? context.selection() : context

    selection.each(function (G) {
      // Create the skeleton, if it doesn't already exist
      const svg = select(this)

      let sankey = svg.selectAll('.sankey')
        .data([{ type: 'sankey' }])

      const sankeyEnter = sankey.enter()
        .append('g')
        .classed('sankey', true)

      sankeyEnter.append('g').classed('groups', true)
      sankeyEnter.append('g').classed('links', true) // Links below nodes
      sankeyEnter.append('g').classed('nodes', true)
      sankeyEnter.append('g').classed('slice-titles', true) // Slice titles

      sankey = sankey.merge(sankeyEnter)

      // Update margins
      sankey
        .attr('transform', 'translate(' + margin.left + ',' + margin.top + ')')
        // .select('.slice-titles')
        // .attr('transform', 'translate(' + margin.left + ',0)')

      // Groups of nodes
      const nodeMap = map(G.nodes, n => n.id)
      const groupsPositioned = (groups || []).map(g => positionGroup(nodeMap, g))

      // All links -- including "from elsewhere" and "to elsewhere" ones
      const links = Array.from(G.links)
      G.nodes.forEach(node => {
        Array.prototype.push.apply(links, node.fromElsewhere || [])
        Array.prototype.push.apply(links, node.toElsewhere || [])
      })

      // Render
      updateNodes(sankey, context, G.nodes)
      updateLinks(sankey, context, links)
      updateGroups(svg, groupsPositioned)
      // updateSlices(svg, layout.slices(nodes));

      // Events
      svg.on('click', function () {
        listeners.call('selectNode', this, null)
        listeners.call('selectLink', this, null)
      })
    })
  }

  function updateNodes (sankey, context, nodes) {
    let nodeSel = sankey
      .select('.nodes')
      .selectAll('.node')
      .data(nodes, d => d.id)

    // EXIT
    nodeSel.exit().remove()

    nodeSel = nodeSel.merge(
      nodeSel.enter()
        .append('g')
        .attr('class', 'node')
        .call(node)
        .on('click', selectNode))

    if (context instanceof transition) {
      nodeSel.transition(context)
        .call(node)
    } else {
      nodeSel.call(node)
    }
  }

  function updateLinks (sankey, context, edges) {
    let linkSel = sankey
      .select('.links')
      .selectAll('.link')
      .data(edges, d => d.source.id + '-' + d.target.id + '-' + d.type)

    // EXIT

    linkSel.exit().remove()

    // ENTER

    const linkEnter = linkSel.enter()
      .append('g')
      .attr('class', 'link')
      .on('click', selectLink)

    linkEnter.append('path')
      .attr('d', link)
      .style('fill', 'white')
      .each(function (d) { this._current = d })

    linkEnter.append('title')

    linkEnter.append('text')
      .attr('class', 'label')
      .attr('dy', '0.35em')
      .attr('x', d => d.points[0].x + 4)
      .attr('y', d => d.points[0].y)

    // UPDATE

    linkSel = linkSel.merge(linkEnter)

    // Calculate group importance for sorting
    calculateGroupImportance(edges)

    // Non-transition updates
    linkSel.classed('selected', (d) => d.id === selectedEdge)
    linkSel.sort(linkOrder)

    // Transition updates, if available
    if (context instanceof transition) {
      linkSel = linkSel.transition(context)
      linkSel
        .select('path')
        .style('fill', linkColor)
        .each(function (d) {
          select(this)
            .transition(context)
            .attrTween('d', interpolateLink)
        })
    } else {
      linkSel
        .select('path')
        .style('fill', linkColor)
        .attr('d', link)
    }

    linkSel.select('title')
      .text(linkTitle)

    linkSel.select('.label')
      .text(linkLabel)
      .attr('x', d => d.points[0].x + 4)
      .attr('y', d => d.points[0].y)
  }

  function calculateGroupImportance (edges) {
    // Group links by source-target pair
    const groups = new Map()

    // Calculate individual importance for each link
    edges.forEach(edge => {
      const groupKey = edge.source.id + '-' + edge.target.id
      if (!groups.has(groupKey)) {
        groups.set(groupKey, [])
      }
      const importance = linkImportance(edge)
      edge._importance = importance
      groups.get(groupKey).push({ edge, importance })
    })

    // Calculate aggregated importance for each group and assign to all links in group
    groups.forEach(linkGroup => {
      const importanceValues = linkGroup.map(item => item.importance)
      const groupImportance = linkImportanceAgg(importanceValues)
      linkGroup.forEach(item => {
        item.edge._groupImportance = groupImportance
      })
    })
  }

  // function updateSlices(svg, slices) {
  //   var slice = svg.select('.slice-titles').selectAll('.slice')
  //         .data(slices, function(d) { return d.id; });

  //   var textWidth = (slices.length > 1 ?
  //                    0.9 * (slices[1].x - slices[0].x) :
  //                    null);

  //   slice.enter().append('g')
  //     .attr('class', 'slice')
  //     .append('foreignObject')
  //     .attr('requiredFeatures',
  //           'http://www.w3.org/TR/SVG11/feature#Extensibility')
  //     .attr('height', margin.top)
  //     .attr('class', 'title')
  //     .append('xhtml:div')
  //     .style('text-align', 'center')
  //     .style('word-wrap', 'break-word');
  //   // .text(pprop('sliceMetadata', 'title'));

  //   slice
  //     .attr('transform', function(d) {
  //       return 'translate(' + (d.x - textWidth / 2) + ',0)'; })
  //     .select('foreignObject')
  //     .attr('width', textWidth)
  //     .select('div');
  //   // .text(pprop('sliceMetadata', 'title'));

  //   slice.exit().remove();
  // }

  function updateGroups (svg, groups) {
    let group = svg.select('.groups').selectAll('.group')
      .data(groups)

    // EXIT
    group.exit().remove()

    // ENTER
    const enter = group.enter().append('g')
      .attr('class', 'group')
    // .on('click', selectGroup);

    enter.append('rect')
    enter.append('text')
      .attr('x', -10)
      .attr('y', -25)

    group = group.merge(enter)

    group
      .style('display', d => d.title ? 'inline' : 'none')
      .attr('transform', d => `translate(${d.rect.left},${d.rect.top})`)
      .select('rect')
      .attr('x', -10)
      .attr('y', -20)
      .attr('width', d => d.rect.right - d.rect.left + 20)
      .attr('height', d => d.rect.bottom - d.rect.top + 30)

    group.select('text')
      .text(d => d.title)
  }

  function interpolateLink (b) {
    // XXX should limit radius better
    b.points.forEach(function (p) {
      if (p.ri > 1e3) p.ri = 1e3
      if (p.ro > 1e3) p.ro = 1e3
    })
    const interp = interpolate(linkGeom(this._current), b)
    const that = this
    return function (t) {
      that._current = interp(t)
      return link(that._current)
    }
  }

  function linkGeom (l) {
    return {
      points: l.points,
      dy: l.dy
    }
  }

  function linkOrder (a, b) {
    // Selected links always on top
    if (a.id === selectedEdge) return +1
    if (b.id === selectedEdge) return -1
    // All other sorting based on group importance
    return a._groupImportance - b._groupImportance
  }

  function selectLink (d) {
    event.stopPropagation()
    const el = select(this).node()
    listeners.call('selectLink', el, d)
  }

  function selectNode (d) {
    event.stopPropagation()
    const el = select(this).node()
    listeners.call('selectNode', el, d)
  }

  // function selectGroup(d) {
  //   d3.event.stopPropagation();
  //   var el = d3.select(this)[0][0];
  //   dispatch.selectGroup.call(el, d);
  // }

  exports.margins = function (_x) {
    if (!arguments.length) return margin
    margin = {
      top: _x.top === undefined ? margin.top : _x.top,
      left: _x.left === undefined ? margin.left : _x.left,
      bottom: _x.bottom === undefined ? margin.bottom : _x.bottom,
      right: _x.right === undefined ? margin.right : _x.right
    }
    return this
  }

  exports.groups = function (_x) {
    if (!arguments.length) return groups
    groups = _x
    return this
  }

  // Node styles and title
  exports.nodeTitle = function (_x) {
    if (!arguments.length) return node.nodeTitle()
    node.nodeTitle(_x)
    linkTitle = linkTitleGenerator(_x, d => d.type, fmt)
    return this
  }

  exports.nodeValue = function (_x) {
    if (!arguments.length) return node.nodeValue()
    node.nodeValue(_x)
    return this
  }

  // Link styles and titles
  exports.linkTitle = function (_x) {
    if (!arguments.length) return linkTitle
    linkTitle = _x
    return this
  }

  exports.linkLabel = function (_x) {
    if (!arguments.length) return linkLabel
    linkLabel = _x
    return this
  }

  exports.linkColor = function (_x) {
    if (!arguments.length) return linkColor
    linkColor = _x
    return this
  }

  exports.linkMinWidth = function (_x) {
    if (!arguments.length) return link.minWidth()
    link.minWidth(_x)
    return this
  }

  exports.linkImportance = function (_x) {
    if (!arguments.length) return linkImportance
    linkImportance = _x
    return this
  }

  exports.linkImportanceAgg = function (_x) {
    if (!arguments.length) return linkImportanceAgg
    linkImportanceAgg = _x
    return this
  }

  exports.selectNode = function (_x) {
    selectedNode = _x
    return this
  }

  exports.selectLink = function (_x) {
    selectedEdge = _x
    return this
  }

  exports.on = function () {
    const value = listeners.on.apply(listeners, arguments)
    return value === listeners ? exports : value
  }

  return exports
}

function defaultLinkLabel (d) {
  return null
}

function defaultLinkImportance (d) {
  // Return negative values for special cases to put them at the bottom
  if (!d.source || (d.target && d.target.direction === 'd')) return -2
  if (!d.target || (d.source && d.source.direction === 'd')) return -1
  // Return dy (visual width) for normal links
  return d.dy
}

function defaultLinkImportanceAgg (values) {
  // Sum of importance values for all links in the group
  return values.reduce((a, b) => a + b, 0)
}