simulate import GPU from gpu.js and export forceLayout function

force-directed-gpu.js

/**
 * simulate import GPU from gpu.js and export forceLayout function
 */
import { GPU } from 'gpu.js'


// vector2 functions
function add(v1, v2) {
  return [
    v1[0] + v2[0],
    v1[1] + v2[1],
  ]
}

function sub(v1, v2) {
  return [
    v1[0] - v2[0],
    v1[1] - v2[1],
  ]
}

function scale(v, s) {
  return [
    v[0] * s,
    v[1] * s,
  ]
}

function division(v, s) {
  if (s === 0) {
    return [0, 0]
  }

  return [
    v[0] / s,
    v[1] / s,
  ]
}

function len(v) {
  return Math.sqrt(v[0] * v[0] + v[1] * v[1])
}

function vecNormalize(v) {
  const distance = len(v)
  if (distance === 0) {
    return [0, 0]
  }
  return [v[0] / distance, v[1] / distance]
}

/**
 * functions precompiled inside kernel
 * https://github.com/gpujs/gpu.js#adding-custom-functions-directly-to-kernel
 */
const vecFunctions = [
  add,
  sub,
  scale,
  division,
  len,
  vecNormalize,
]


// output (x, y)
function nodesPush(nodes, forces) {
  const idx = this.thread.x

  const force = forces[idx]
  const node = nodes[idx]

  return add(node, force)
}

// output (n1ForceX, n1ForceY, n2ForceX, n2ForceY)
function edgesForce(nodes, edges, weights, distances) {
  const friction = this.constants.friction
  const idx = this.thread.x

  const edge = edges[idx]
  const edgeDistance = distances[idx]

  const n1 = nodes[edge[0]]
  const n2 = nodes[edge[1]]
  const weight1 = weights[edge[0]]
  const weight2 = weights[edge[1]]

  let weight = 0.5

  const sumWeight = weight1 + weight2
  if (sumWeight !== 0) {
    weight = weight2 / sumWeight
  }

  let deltaVec = sub(n2, n1)
  const deltaDistance = len(deltaVec) - edgeDistance

  deltaVec = vecNormalize(deltaVec)
  const f1 = scale(deltaVec, weight * deltaDistance * friction)
  const f2 = scale(deltaVec, -(1 - weight) * deltaDistance * friction)
  return [f1[0], f1[1], f2[0], f2[1]]
}

// output (forceX, forceY)
function edgesForceOverlay(edges, forces) {
  const idx = this.thread.x
  let nodeForce = [0, 0]

  let linkedCount = 0

  for (let j = 0; j < this.constants.edgesLength; j++) {
    const edge = edges[j]
    const forceVec = forces[j]
    if (edge[0] === idx) {
      const force = [forceVec[0], forceVec[1]]
      nodeForce = add(nodeForce, force)
      linkedCount++
    }
    else if (edge[1] === idx) {
      const force = [forceVec[2], forceVec[3]]
      nodeForce = add(nodeForce, force)
      linkedCount++
    }
  }

  return division(nodeForce, linkedCount)
}


// output (forceX, forceY)
function gravityForce(nodes, center) {
  const gravity = this.constants.gravity
  const friction = this.constants.friction
  const idx = this.thread.x

  const centerVec = [center[0], center[1]]
  const node = nodes[idx]

  const deltaVec = sub(centerVec, node)
  return scale(deltaVec, gravity * friction)
}


/**
 * output: (forceX, forceY)[i, j] n1 n2 相对作用力
 * (i, j) n[j] 对 n[i] 的作用力
 * i 行即 n[i] 受到合力
 *
 *       j
 * i -> (,) (,) (,)
 *      (,) (,) (,)
 *      (,) (,) (,)
 */
function repulsiveForce(nodes, weights) {
  const i = this.thread.y
  const j = this.thread.x

  if (i >= j) {
    return [0, 0]
  }

  const n1 = nodes[i]
  const n2 = nodes[j]
  const repulsion1 = weights[i]
  const repulsion2 = weights[j]

  let deltaVec = sub(n2, n1)
  let distance = len(deltaVec)
  if (distance === 0) {
    deltaVec = [Math.random() - 0.5, Math.random() - 0.5]
    distance = 1
  }

  const repFact = (repulsion1 + repulsion2) / (distance * distance)
  return scale(deltaVec, repFact)
}


// output (forceX, forceY) node 收到排斥力合力
function repulsiveForceOverlay(forces) {
  const friction = this.constants.friction
  const idx = this.thread.x
  let overlayForce = [0, 0]

  for (let j = 0; j < this.constants.nodesLength; j++) {
    let force = [0, 0]
    if (j > idx) {
      force = scale(forces[idx][j], -1)
    } else if (j < idx) {
      force = forces[j][idx]
    }
    overlayForce = add(overlayForce, force)
  }

  return scale(overlayForce, friction)
}


function forceOverlay(forces1, forces2, forces3) {
  const idx = this.thread.x
  const force1 = forces1[idx]
  const force2 = forces2[idx]
  const force3 = forces3[idx]
  return add(add(force1, force2), force3)
}


// output: (x, y)
function toVector(array) {
  const idx = this.thread.x
  return [array[idx][0], array[idx][1]]
}


// output: x
function toWeight(array) {
  const idx = this.thread.x
  return array[idx]
}


function kernelFactory(constants) {
  const { nodesLength, edgesLength } = constants

  const gpu = new GPU({
    functions: vecFunctions,
  })

  const commonSettings = {
    loopMaxIterations: 10000,
    dynamicArguments: true,
    pipeline: true,
    constants,
  }

  const toVectorTexture = gpu.createKernel(
    toVector,
    commonSettings,
  )

  const toWeightTexture = gpu.createKernel(
    toWeight,
    commonSettings,
  )

  const nodesPushKernel = gpu.createKernel(
    nodesPush,
    {
      output: [nodesLength],
      ...commonSettings,
    },
  )

  const edgesForceKernel = gpu.createKernel(
    edgesForce,
    {
      output: [edgesLength],
      ...commonSettings,
    },
  )

  const edgesOverlayKernel = gpu.createKernel(
    edgesForceOverlay,
    {
      output: [nodesLength],
      ...commonSettings,
    },
  )

  const edgesKernel = (nodes, edges, weights, distances) =>
    edgesOverlayKernel(
      edges,
      edgesForceKernel(nodes, edges, weights, distances),
    )

  const gravityKernel = gpu.createKernel(
    gravityForce,
    {
      output: [nodesLength],
      ...commonSettings,
    },
  )

  const repulsiveForceKernel = gpu.createKernel(
    repulsiveForce,
    {
      output: [nodesLength, nodesLength],
      ...commonSettings,
    },
  )

  const repulsiveOverlayKernel = gpu.createKernel(
    repulsiveForceOverlay,
    {
      output: [nodesLength],
      ...commonSettings,
    },
  )

  const repulsiveKernel = (nodes, weights) =>
    repulsiveOverlayKernel(
      repulsiveForceKernel(nodes, weights),
    )

  const forceOverlayKernel = gpu.createKernel(
    forceOverlay,
    {
      output: [nodesLength],
      ...commonSettings,
    },
  )

  return {
    toVectorTexture,
    toWeightTexture,
    nodesPushKernel,
    edgesKernel,
    gravityKernel,
    repulsiveKernel,
    forceOverlayKernel,
  }
}

export default function forceLayout(nodes, edges, opts) {
  let friction = 0.6
  const iterateTimes = 1
  const gravity = opts.gravity == null ? 0.1 : opts.gravity

  const repulsion = 50
  const nodesLength = nodes.length
  const edgesLength = edges.length

  const rect = opts.rect
  const width = rect.width
  const height = rect.height
  const center = [rect.x + width / 2, rect.y + height / 2]

  const constants = {
    gravity,
    friction,
    repulsion,
    nodesLength,
    edgesLength,
  }

  const {
    toVectorTexture,
    toWeightTexture,
    nodesPushKernel,
    edgesKernel,
    gravityKernel,
    repulsiveKernel,
    forceOverlayKernel,
  } = kernelFactory(constants)

  for (let i = 0; i < nodesLength; i++) {
    const node = nodes[i]
    node.index = i
    node.edges = null
    if (!node.p) {
      node.p = new Float32Array([
        width / 10 * (Math.random() - 0.5) + center[0],
        height / 10 * (Math.random() - 0.5) + center[1],
      ])
    }
  }

  let vecNodes = nodes.map(node => [
    node.p[0],
    node.p[1],
  ])

  const nodesWeight = nodes.map(node => [
    node.w,
  ])

  const vecEdges = edges.map(edge => [
    edge.n1.index,
    edge.n2.index,
  ])

  const edgesDistance = edges.map(edge => [
    edge.d,
  ])

  const weightTexture = toWeightTexture.setOutput([nodesLength])(nodesWeight)
  const edgesTexture = toVectorTexture.setOutput([edgesLength])(vecEdges)
  const distanceTexture = toWeightTexture.setOutput([edgesLength])(edgesDistance)
  const centerTexture = toWeightTexture.setOutput([2])(center)

  const iterate = friction => {
    constants.friction = friction

    for (let i = 0; i < iterateTimes; i++) {
      let nodesTexture = toVectorTexture.setOutput([nodesLength])(vecNodes)

      nodesTexture = nodesPushKernel(
        nodesTexture,
        forceOverlayKernel(
          gravityKernel(nodesTexture, centerTexture),
          edgesKernel(nodesTexture, edgesTexture, weightTexture, distanceTexture),
          repulsiveKernel(nodesTexture, weightTexture),
        ),
      )

      vecNodes = nodesTexture.toArray()
    }

    for (let i = 0; i < nodesLength; i++) {
      const node = nodes[i]
      node.p[0] = vecNodes[i][0]
      node.p[1] = vecNodes[i][1]
    }

  }

  return {
    warmUp() {
      constants.friction = friction = 0.5
    },

    setFixed(idx) {
      nodes[idx].fixed = true
    },

    setUnfixed(idx) {
      nodes[idx].fixed = false
    },

    step(callback) {
      const result = iterate(friction)
      friction = friction * 0.992
      callback && callback(nodes, edges, friction < 0.01)
      return result
    },
  }
}