visjs-network/lib/network/modules/components/physics/HierarchicalSpringSolver.js

A dynamic, browser-based network visualization library.

/**
 * Hierarchical Spring Solver
 */
class HierarchicalSpringSolver {
  /**
   * @param {Object} body
   * @param {{physicsNodeIndices: Array, physicsEdgeIndices: Array, forces: {}, velocities: {}}} physicsBody
   * @param {Object} options
   */
  constructor(body, physicsBody, options) {
    this.body = body
    this.physicsBody = physicsBody
    this.setOptions(options)
  }

  /**
   *
   * @param {Object} options
   */
  setOptions(options) {
    this.options = options
  }

  /**
   * This function calculates the springforces on the nodes, accounting for the support nodes.
   *
   * @private
   */
  solve() {
    var edgeLength, edge
    var dx, dy, fx, fy, springForce, distance
    var edges = this.body.edges
    var factor = 0.5

    var edgeIndices = this.physicsBody.physicsEdgeIndices
    var nodeIndices = this.physicsBody.physicsNodeIndices
    var forces = this.physicsBody.forces

    // initialize the spring force counters
    for (let i = 0; i < nodeIndices.length; i++) {
      let nodeId = nodeIndices[i]
      forces[nodeId].springFx = 0
      forces[nodeId].springFy = 0
    }

    // forces caused by the edges, modelled as springs
    for (let i = 0; i < edgeIndices.length; i++) {
      edge = edges[edgeIndices[i]]
      if (edge.connected === true) {
        edgeLength =
          edge.options.length === undefined
            ? this.options.springLength
            : edge.options.length

        dx = edge.from.x - edge.to.x
        dy = edge.from.y - edge.to.y
        distance = Math.sqrt(dx * dx + dy * dy)
        distance = distance === 0 ? 0.01 : distance

        // the 1/distance is so the fx and fy can be calculated without sine or cosine.
        springForce =
          (this.options.springConstant * (edgeLength - distance)) / distance

        fx = dx * springForce
        fy = dy * springForce

        if (edge.to.level != edge.from.level) {
          if (forces[edge.toId] !== undefined) {
            forces[edge.toId].springFx -= fx
            forces[edge.toId].springFy -= fy
          }
          if (forces[edge.fromId] !== undefined) {
            forces[edge.fromId].springFx += fx
            forces[edge.fromId].springFy += fy
          }
        } else {
          if (forces[edge.toId] !== undefined) {
            forces[edge.toId].x -= factor * fx
            forces[edge.toId].y -= factor * fy
          }
          if (forces[edge.fromId] !== undefined) {
            forces[edge.fromId].x += factor * fx
            forces[edge.fromId].y += factor * fy
          }
        }
      }
    }

    // normalize spring forces
    springForce = 1
    var springFx, springFy
    for (let i = 0; i < nodeIndices.length; i++) {
      let nodeId = nodeIndices[i]
      springFx = Math.min(
        springForce,
        Math.max(-springForce, forces[nodeId].springFx)
      )
      springFy = Math.min(
        springForce,
        Math.max(-springForce, forces[nodeId].springFy)
      )

      forces[nodeId].x += springFx
      forces[nodeId].y += springFy
    }

    // retain energy balance
    var totalFx = 0
    var totalFy = 0
    for (let i = 0; i < nodeIndices.length; i++) {
      let nodeId = nodeIndices[i]
      totalFx += forces[nodeId].x
      totalFy += forces[nodeId].y
    }
    var correctionFx = totalFx / nodeIndices.length
    var correctionFy = totalFy / nodeIndices.length

    for (let i = 0; i < nodeIndices.length; i++) {
      let nodeId = nodeIndices[i]
      forces[nodeId].x -= correctionFx
      forces[nodeId].y -= correctionFy
    }
  }
}

export default HierarchicalSpringSolver