@antv/layout 2.0.0-alpha.1 → 2.0.0-alpha.3

package/src/algorithm/d3-force/index.ts

@@ -0,0 +1,592 @@
+import { deepMix } from '@antv/util';
+import type { ForceLink, Simulation } from 'd3-force';
+import {
+  forceCenter,
+  forceCollide,
+  forceLink,
+  forceManyBody,
+  forceRadial,
+  forceSimulation,
+  forceX,
+  forceY,
+} from 'd3-force';
+import type { ID, Position } from '../../types';
+import { assignDefined, normalizeViewport } from '../../util';
+import { formatNodeSizeFn } from '../../util/format';
+import { BaseLayoutWithIterations } from '../base-layout';
+import forceInABox from './force-in-a-box';
+import type {
+  D3ForceCommonOptions,
+  D3ForceLayoutOptions,
+  EdgeDatum,
+  NodeDatum,
+} from './types';
+
+export type { D3ForceLayoutOptions };
+
+const DEFAULTS_LAYOUT_OPTIONS: Partial<D3ForceLayoutOptions> = {
+  link: {
+    id: (d) => String(d.id),
+  },
+
+  manyBody: {
+    strength: -30,
+  },
+
+  preventOverlap: false,
+  nodeSize: 10,
+  nodeSpacing: 0,
+
+  x: false,
+  y: false,
+
+  clustering: false,
+  clusterNodeStrength: -1,
+  clusterEdgeStrength: 0.1,
+  clusterEdgeDistance: 100,
+  clusterFociStrength: 0.8,
+  clusterNodeSize: 10,
+};
+
+export class D3ForceLayout<
+  T extends D3ForceCommonOptions = D3ForceLayoutOptions,
+  N extends NodeDatum = NodeDatum,
+  E extends EdgeDatum<N> = EdgeDatum<N>,
+> extends BaseLayoutWithIterations<T> {
+  public id = 'd3-force';
+
+  public simulation!: Simulation<N, E>;
+
+  private d3Nodes: N[] = [];
+  private d3Edges: E[] = [];
+
+  protected config = {
+    simulationAttrs: [
+      'alpha',
+      'alphaMin',
+      'alphaDecay',
+      'alphaTarget',
+      'velocityDecay',
+      'randomSource',
+    ],
+  };
+
+  protected getDefaultOptions(): T {
+    return DEFAULTS_LAYOUT_OPTIONS as T;
+  }
+
+  protected mergeOptions<T>(base: T, patch?: Partial<T>): T {
+    return deepMix({}, base, patch) as T;
+  }
+
+  constructor(options?: Partial<T>) {
+    super(options);
+
+    if (this.options.forceSimulation) {
+      this.simulation = this.options.forceSimulation as unknown as Simulation<
+        N,
+        E
+      >;
+    }
+  }
+
+  public stop(): this {
+    if (this.simulation) {
+      this.simulation.stop();
+    }
+    return this;
+  }
+
+  public tick(iterations: number = 1): this {
+    if (this.simulation) {
+      for (let i = 0; i < iterations; i++) {
+        this.simulation.tick();
+      }
+      this.syncPositionsFromD3();
+      this.options.onTick?.(this);
+    }
+    return this;
+  }
+
+  public restart(alpha?: number): this {
+    if (this.simulation) {
+      if (alpha !== undefined) {
+        this.simulation.alpha(alpha);
+      }
+      this.simulation.restart();
+    }
+    return this;
+  }
+
+  public reheat(): this {
+    return this.restart(1);
+  }
+
+  public getAlpha(): number {
+    return this.simulation?.alpha() ?? 0;
+  }
+
+  public setAlpha(alpha: number): this {
+    if (this.simulation) {
+      this.simulation.alpha(alpha);
+    }
+    return this;
+  }
+
+  public getForce(name: string): any {
+    return this.simulation?.force(name);
+  }
+
+  public force(name: string, force: any): this {
+    if (this.simulation) {
+      this.simulation.force(name, force);
+    }
+    return this;
+  }
+
+  public nodes(): N[] {
+    return this.simulation?.nodes() ?? [];
+  }
+
+  public find(x: number, y: number, radius?: number): N | undefined {
+    if (!this.simulation) return undefined;
+    return this.simulation.find(x, y, radius);
+  }
+
+  public setFixedPosition(id: ID, position: Position | null[] | null): void {
+    const d3Node = this.d3Nodes.find((n) => (n as any).id === id);
+    const node = this.model.node(id);
+    if (!node || !d3Node) return;
+
+    const keys = ['fx', 'fy', 'fz'] as const;
+
+    if (position === null) {
+      // Unset fixed position
+      keys.forEach((key) => {
+        delete (node as any)[key];
+        delete (d3Node as any)[key];
+      });
+      return;
+    }
+
+    position.forEach((value, index) => {
+      if (
+        index < keys.length &&
+        (typeof value === 'number' || value === null)
+      ) {
+        (node as any)[keys[index]] = value;
+        (d3Node as any)[keys[index]] = value;
+      }
+    });
+  }
+
+  protected parseOptions(options: Partial<T>): T {
+    const _ = options;
+    // process nodeSize
+    if (_.collide && _.collide?.radius === undefined) {
+      _.collide = _.collide || {};
+      // @ts-ignore
+      _.collide.radius = _.nodeSize ?? 10;
+    }
+    // process iterations
+    if (_.iterations === undefined) {
+      if (_.link && _.link.iterations === undefined) {
+        _.iterations = _.link.iterations;
+      }
+      if (_.collide && _.collide.iterations === undefined) {
+        _.iterations = _.collide.iterations;
+      }
+    }
+
+    return _ as T;
+  }
+
+  protected async layout(): Promise<void> {
+    const options = this.parseOptions(this.options || {});
+
+    this.createD3Copies();
+
+    const simulation = this.setSimulation(options);
+
+    simulation.nodes(this.d3Nodes);
+    simulation.force<ForceLink<N, E>>('link')?.links(this.d3Edges);
+
+    return new Promise<void>((resolve) => {
+      simulation.on('end', () => {
+        this.syncPositionsFromD3();
+        resolve();
+      });
+    });
+  }
+
+  private createD3Copies() {
+    this.d3Nodes = [];
+    this.d3Edges = [];
+
+    this.model.forEachNode((node) => {
+      this.d3Nodes.push({ ...(node as any) });
+    });
+    this.model.forEachEdge((edge) => {
+      this.d3Edges.push({ ...(edge as any) });
+    });
+  }
+
+  protected syncPositionsFromD3() {
+    this.d3Nodes.forEach((d3Node) => {
+      const node = this.model.node(d3Node.id);
+      if (node) {
+        node.x = d3Node.x!;
+        node.y = d3Node.y!;
+        if (d3Node.z !== undefined) node.z = d3Node.z;
+        // 同步固定位置属性
+        if (d3Node.fx !== undefined) node.fx = d3Node.fx;
+        if (d3Node.fy !== undefined) node.fy = d3Node.fy;
+        if (d3Node.fz !== undefined) node.fz = d3Node.fz;
+        // 同步速度属性
+        if (d3Node.vx !== undefined) node.vx = d3Node.vx;
+        if (d3Node.vy !== undefined) node.vy = d3Node.vy;
+        if (d3Node.vz !== undefined) node.vz = d3Node.vz;
+      }
+    });
+  }
+
+  protected initSimulation() {
+    return forceSimulation<N, E>();
+  }
+
+  protected setSimulation(options: T) {
+    const simulation =
+      this.simulation ||
+      (this.options.forceSimulation as unknown as Simulation<N, E>) ||
+      this.initSimulation();
+
+    if (!this.simulation) {
+      this.simulation = simulation.on('tick', () => {
+        this.syncPositionsFromD3();
+        options.onTick?.(this);
+      });
+    }
+
+    apply(
+      simulation,
+      this.config.simulationAttrs.map((name) => [
+        name,
+        options[name as keyof T],
+      ]),
+    );
+
+    this.setupForces(simulation, options);
+
+    return simulation;
+  }
+
+  protected setupForces(simulation: Simulation<N, E>, options: T) {
+    this.setupLinkForce(simulation, options);
+    this.setupManyBodyForce(simulation, options);
+    this.setupCenterForce(simulation, options);
+    this.setupCollisionForce(simulation, options);
+    this.setupXForce(simulation, options);
+    this.setupYForce(simulation, options);
+    this.setupRadialForce(simulation, options);
+    this.setupClusterForce(simulation, options);
+  }
+
+  private getCenterOptions(options: T): T['center'] | undefined {
+    if (options.center === false) return undefined;
+
+    const viewport = normalizeViewport({
+      width: options.width,
+      height: options.height,
+    });
+    return assignDefined({}, options.center || {}, {
+      x: viewport.width / 2,
+      y: viewport.height / 2,
+      strength: options.centerStrength,
+    }) as T['center'];
+  }
+
+  protected setupCenterForce(simulation: Simulation<N, E>, options: T) {
+    const center = this.getCenterOptions(options);
+
+    if (center) {
+      let force = simulation.force('center');
+      if (!force) {
+        force = forceCenter(center.x, center.y);
+        simulation.force('center', force as any);
+      }
+
+      const params: [string, any][] = [];
+      if (center.x !== undefined) params.push(['x', center.x]);
+      if (center.y !== undefined) params.push(['y', center.y]);
+      if (center.strength !== undefined)
+        params.push(['strength', center.strength]);
+
+      apply(force, params);
+    } else {
+      simulation.force('center', null);
+    }
+  }
+
+  private getManyBodyOptions(options: T): D3ForceLayoutOptions['manyBody'] {
+    if (options.manyBody === false) return undefined;
+
+    return assignDefined({}, options.manyBody || {}, {
+      strength: options.nodeStrength,
+      distanceMin: options.distanceMin,
+      distanceMax: options.distanceMax,
+      theta: options.theta,
+    });
+  }
+
+  protected setupManyBodyForce(simulation: Simulation<N, E>, options: T) {
+    const manyBody = this.getManyBodyOptions(options);
+
+    if (manyBody) {
+      let force = simulation.force('charge');
+      if (!force) {
+        force = forceManyBody();
+        simulation.force('charge', force as any);
+      }
+
+      const params: [string, any][] = [];
+
+      if (manyBody.strength !== undefined)
+        params.push(['strength', manyBody.strength]);
+      if (manyBody.distanceMin !== undefined)
+        params.push(['distanceMin', manyBody.distanceMin]);
+      if (manyBody.distanceMax !== undefined)
+        params.push(['distanceMax', manyBody.distanceMax]);
+      if (manyBody.theta !== undefined) params.push(['theta', manyBody.theta]);
+
+      apply(force, params);
+    } else {
+      simulation.force('charge', null);
+    }
+  }
+
+  private getLinkOptions(options: T): D3ForceLayoutOptions['link'] {
+    if (options.link === false) return undefined;
+
+    return assignDefined({}, options.link || {}, {
+      id: options.edgeId,
+      distance: options.linkDistance,
+      strength: options.edgeStrength,
+      iterations: options.edgeIterations,
+    });
+  }
+
+  protected setupLinkForce(simulation: Simulation<N, E>, options: T) {
+    const edges = this.model.edges();
+
+    const link = this.getLinkOptions(options);
+
+    if (edges.length > 0 && link) {
+      let force = simulation.force<ForceLink<N, E>>('link');
+      if (!force) {
+        force = forceLink<N, E>();
+        simulation.force('link', force);
+      }
+
+      const params: [string, any][] = [];
+      if (link.id !== undefined) params.push(['id', link.id]);
+      if (link.distance !== undefined) params.push(['distance', link.distance]);
+      if (link.strength !== undefined) params.push(['strength', link.strength]);
+      if (link.iterations !== undefined)
+        params.push(['iterations', link.iterations]);
+
+      apply(force, params);
+    } else {
+      simulation.force('link', null);
+    }
+  }
+
+  private getCollisionOptions(options: T): D3ForceLayoutOptions['collide'] {
+    if (options.preventOverlap === false || options.collide === false)
+      return undefined;
+
+    const radius =
+      options.nodeSize || options.nodeSpacing
+        ? (d: NodeDatum) =>
+            formatNodeSizeFn(
+              options.nodeSize,
+              options.nodeSpacing,
+            )(d._original) / 2
+        : undefined;
+
+    return assignDefined({}, options.collide || {}, {
+      radius: options.collide || radius,
+      strength: options.collideStrength,
+      iterations: options.collideIterations,
+    });
+  }
+
+  protected setupCollisionForce(simulation: Simulation<N, E>, options: T) {
+    const collide = this.getCollisionOptions(options);
+    if (collide) {
+      let force = simulation.force('collide');
+      if (!force) {
+        force = forceCollide();
+        simulation.force('collide', force as any);
+      }
+
+      const params: [string, any][] = [];
+
+      if (collide.radius !== undefined) params.push(['radius', collide.radius]);
+      if (collide.strength !== undefined)
+        params.push(['strength', collide.strength]);
+      if (collide.iterations !== undefined)
+        params.push(['iterations', collide.iterations]);
+
+      apply(force, params);
+    } else {
+      simulation.force('collide', null);
+    }
+  }
+
+  private getXForceOptions(options: T): D3ForceLayoutOptions['x'] {
+    if (options.x === false) return undefined;
+
+    const center = this.getCenterOptions(options);
+    return assignDefined({}, options.x || {}, {
+      x: options.forceXPosition ?? (center && center.x),
+      strength: options.forceXStrength,
+    });
+  }
+
+  protected setupXForce(simulation: Simulation<N, E>, options: T) {
+    const x = this.getXForceOptions(options);
+
+    if (x) {
+      let force = simulation.force('x');
+      if (!force) {
+        force = forceX();
+        simulation.force('x', force as any);
+      }
+
+      const params: [string, any][] = [];
+      if (x.x !== undefined) params.push(['x', x.x]);
+      if (x.strength !== undefined) params.push(['strength', x.strength]);
+
+      apply(force, params);
+    } else {
+      simulation.force('x', null);
+    }
+  }
+
+  private getYForceOptions(options: T): D3ForceLayoutOptions['y'] {
+    if (options.y === false) return undefined;
+
+    const center = this.getCenterOptions(options);
+    return assignDefined({}, options.y || {}, {
+      y: options.forceYPosition ?? (center && center.y),
+      strength: options.forceYStrength,
+    });
+  }
+
+  protected setupYForce(simulation: Simulation<N, E>, options: T) {
+    const y = this.getYForceOptions(options);
+    if (y) {
+      let force = simulation.force('y');
+      if (!force) {
+        force = forceY();
+        simulation.force('y', force as any);
+      }
+
+      const params: [string, any][] = [];
+      if (y.y !== undefined) params.push(['y', y.y]);
+      if (y.strength !== undefined) params.push(['strength', y.strength]);
+
+      apply(force, params);
+    } else {
+      simulation.force('y', null);
+    }
+  }
+
+  private getRadialOptions(options: T): D3ForceLayoutOptions['radial'] {
+    if (
+      options.radial !== undefined ||
+      options.radialStrength !== undefined ||
+      options.radialRadius !== undefined ||
+      options.radialX !== undefined ||
+      options.radialY !== undefined
+    ) {
+      const center = this.getCenterOptions(options);
+      return assignDefined({}, options.radial || {}, {
+        strength: options.radialStrength,
+        radius: options.radialRadius ?? 100,
+        x: options.radialX ?? (center && center.x),
+        y: options.radialY ?? (center && center.y),
+      });
+    }
+    return undefined;
+  }
+
+  protected setupRadialForce(simulation: Simulation<N, E>, options: T) {
+    const radial = this.getRadialOptions(options);
+
+    if (radial) {
+      let force = simulation.force('radial');
+      if (!force) {
+        force = forceRadial(radial.radius || 100, radial.x, radial.y);
+        simulation.force('radial', force as any);
+      }
+
+      const params: [string, any][] = [];
+      if (radial.radius !== undefined) params.push(['radius', radial.radius]);
+      if (radial.strength !== undefined)
+        params.push(['strength', radial.strength]);
+      if (radial.x !== undefined) params.push(['x', radial.x]);
+      if (radial.y !== undefined) params.push(['y', radial.y]);
+
+      apply(force, params);
+    } else {
+      simulation.force('radial', null);
+    }
+  }
+
+  protected setupClusterForce(simulation: Simulation<N, E>, options: T) {
+    const { clustering } = options;
+
+    if (clustering) {
+      const {
+        clusterFociStrength,
+        clusterEdgeDistance,
+        clusterEdgeStrength,
+        clusterNodeStrength,
+        clusterNodeSize,
+        clusterBy,
+      } = options;
+
+      const center = this.getCenterOptions(options);
+
+      let force = simulation.force('group');
+      if (!force) {
+        force = forceInABox();
+        simulation.force('group', force as any);
+      }
+
+      apply(force, [
+        ['centerX', center && center.x],
+        ['centerY', center && center.y],
+        ['template', 'force'],
+        ['strength', clusterFociStrength],
+        ['groupBy', clusterBy],
+        ['nodes', this.model.nodes()],
+        ['links', this.model.edges()],
+        ['forceLinkDistance', clusterEdgeDistance],
+        ['forceLinkStrength', clusterEdgeStrength],
+        ['forceCharge', clusterNodeStrength],
+        ['forceNodeSize', clusterNodeSize],
+      ]);
+    } else {
+      simulation.force('group', null);
+    }
+  }
+}
+
+export const apply = (target: any, params: [string, any][]) => {
+  return params.reduce((acc, [method, param]) => {
+    if (!acc[method] || param === undefined) return acc;
+    return acc[method].call(target, param);
+  }, target);
+};