distortions/lib/layering.js

Helpers for visualizing distortion in nonlinear dimensionality reduction.

// A ggplot2-like API for creating ellipse plots
//
// bundle using this command:
// esbuild --bundle --format=esm --outdir=dist layering.js 

import * as d3 from "https://esm.sh/d3@7";
import d3col from 'https://cdn.jsdelivr.net/npm/d3-svg-legend@2.25.6/+esm'
import { annotation } from "./annotation.js";
import { flatten_edges } from "./reshape.js";
import { draw_boxplot } from "./boxplot.js";
import { link_update } from "./inter_edge_link.js";
import { isometry_update } from "./inter_isometry.js";

export class DistortionPlot {
  // default dimensions
  constructor(el, options) {
    const defaults = { width: 700, height: 350, margin: { top: 40, right: 90, bottom: 60, left: 90 }, labelFontSize: 14 };
    const opts = { ...defaults, ...options };

    // initialize svg canvas
    let svg = d3.create("svg")
      .attr("width", opts.width)
      .attr("height", opts.height)
    el.appendChild(svg.node());
    this.svg = d3.select(el).select("svg")

    this.width = +this.svg.attr("width");
    this.height = +this.svg.attr("height");
    this.margin = opts.margin;
    this.plotWidth = this.width - this.margin.left - this.margin.right;
    this.plotHeight = this.height - this.margin.top - this.margin.bottom;
    this.labelFontSize = opts.labelFontSize

    // Create the base group with margins applied
    this.g = this.svg
      .append("g")
      .attr("transform", `translate(${this.margin.left},${this.margin.top})`);

    // Initialize scales
    this.xScale = d3.scaleLinear().range([0, this.plotWidth]);
    this.yScale = d3.scaleLinear().range([this.plotHeight, 0]);
    this.rScale = d3.scaleLinear();
    this.colorScale = d3.scaleOrdinal(d3.schemeCategory10);

    // Initialize data and layers
    this.dataset = null;
    this.layers = [];
    this.title = "";
    this.xLabel = "";
    this.yLabel = "";
  }

  // Set the base data for the plot
  data(dataset) {
    this.dataset = dataset;

    // create unique ID if not present
    if (Object.keys(dataset[0]).indexOf("_id") == -1) {
      for (let i = 0; i < dataset.length; i++) {
        dataset[i]._id = i;
      }
    }
    return this;
  }

  // Map data columns to x and y scales
  mapping(m) {
    this.mappingObj = m;

    // Set domains based on data
    if (this.dataset && this.mappingObj) {
      this.xScale.domain(d3.extent(this.dataset, d => +d[this.mappingObj.x])).nice();
      this.yScale.domain(d3.extent(this.dataset, d => +d[this.mappingObj.y])).nice();

      let radius_data = this.dataset.map(d => [+d[this.mappingObj.a], +d[this.mappingObj.b]])
      this.rScale.domain(d3.extent([].concat.apply([], radius_data))).nice();
    }

    return this;
  }

  // Add points/scatter layer
  geomEllipse(options = {}) {
    const defaults = {
      radiusMax: 25,
      radiusMin: 3,
      color: null,
      opacity: 1,
      className: "ellipse"
    };

    const opts = { ...defaults, ...options };

    this.layers.push({
      type: "ellipse",
      options: opts,
      render: () => {
        // update scales
        this.rScale.range([opts.radiusMin, opts.radiusMax]);
        if (this.mappingObj.color) {
          const uniqueValues = [...new Set(this.dataset.map(d => d[opts.color]))];
          this.colorScale.domain(uniqueValues);
        }

        // Draw points
        this.g.append("g")
          .attr("class", opts.className)
          .selectAll("ellipse")
          .data(this.dataset, d => d._id).enter()
          .append("ellipse")
          .attr("cx", d => this.xScale(d[this.mappingObj.x]))
          .attr("cy", d => this.yScale(d[this.mappingObj.y]))
          .attr("rx", d => this.rScale(d[this.mappingObj.a]))
          .attr("ry", d => this.rScale(d[this.mappingObj.b]))
          .attr("transform", d => `rotate(${d[this.mappingObj.angle]} ${this.xScale(d[this.mappingObj.x])} ${this.yScale(d[this.mappingObj.y])})`)
          .attr("fill", d => this.mappingObj.color ? this.colorScale(d[this.mappingObj.color]) : opts.color || "#0c0c0c")
          .attr("opacity", opts.opacity);
      }
    });

    return this;
  }

  // Add hair (line) layer, mirroring geomEllipse behavior
  geomHair(options = {}) {
    const defaults = {
      className: "hair",
      color: null,
      radiusMax: 25,
      radiusMin: 0.1,
      opacity: 1,
      strokeWidth: 1
    };
    const opts = { ...defaults, ...options };
    this.layers.push({
      type: "hair",
      options: opts,
      render: () => {
        // update radius and color scales
        this.rScale.range([opts.radiusMin, opts.radiusMax]);
        if (this.mappingObj.color) {
          const uniqueValues = [...new Set(this.dataset.map((d) => d[opts.color]))];
          this.colorScale.domain(uniqueValues);
        }
        // Draw hair as thin rectangles (using <rect> instead of <line>)
        this.g.append("g")
          .attr("class", opts.className)
          .selectAll("rect")
          .data(this.dataset).enter()
          .append("rect")
          .attr("class", opts.className)
          .attr("x", (d) => {
            const cx = this.xScale(+d[this.mappingObj.x]);
            const r = this.rScale(+d[this.mappingObj.a]);
            return cx - r;
          })
          .attr("y", (d) => {
            const cy = this.yScale(+d[this.mappingObj.y]);
            // center the rectangle vertically based on strokeWidth
            return cy - opts.strokeWidth / 2;
          })
          .attr("width", (d) => {
            const r = this.rScale(+d[this.mappingObj.a]);
            return 2 * r;
          })
          .attr("height", opts.strokeWidth)
          .attr("transform", (d) => {
            const cx = this.xScale(+d[this.mappingObj.x]);
            const cy = this.yScale(+d[this.mappingObj.y]);
            return `rotate(${+d[this.mappingObj.angle]} ${cx} ${cy})`;
          })
          .attr("fill", (d) => this.mappingObj.color
            ? this.colorScale(d[this.mappingObj.color])
            : opts.color || "#0c0c0c")
          .attr("opacity", opts.opacity);
      }
    });
    return this;
  }

  // Add scale_color for changing color mapping
  scaleColor(options = {}) {
    const defaults = {
      scheme: d3.schemeCategory10, className: null, padding: 20, x_offset: 20,
      y_offset: 0, size: 14, legendTextSize: 14, titleOffset: 15, labelOffset: 20
    }

    const opts = { ...defaults, ...options };
    this.layers.push({
      type: 'scale_color',
      options: opts,
      render: () => {
        // construct color scales
        if (typeof this.dataset[0][this.mappingObj.color] === 'string') {
          const uniqueValues = [...new Set(this.dataset.map(d => d[this.mappingObj.color]))];
          this.colorScale = d3.scaleOrdinal(uniqueValues, opts.scheme);
        } else {
          const rangeValues = d3.extent(this.dataset.map(d => d[this.mappingObj.color]));
          this.colorScale = d3.scaleLinear(rangeValues, [opts.scheme[0], opts.scheme[1]]);
        }

        // update point color
        let className = opts.className || this.layers[0].options.className;
        this.svg.select(`.${className}`)
          .selectAll("*")
          .attr("fill", d => this.colorScale(d[this.mappingObj.color]))
        this.svg.select(`.${className}-background`)
          .selectAll("*")
          .attr("fill", d => this.colorScale(d[this.mappingObj.color]))

        // d3 color legend
        const legend = d3col.legendColor()
          .shapePadding(opts.padding)
          .shapeHeight(opts.size)
          .shapeWidth(opts.size)
          // .labelOffset(opts.labelOffset)
          .title(this.mappingObj.color)
          .scale(this.colorScale)

        this.g.append("g")
          .attr("class", "legend")
          .attr("id", "colorScale")
          .attr("transform", `translate(${this.plotWidth + opts.x_offset}, ${opts.y_offset})`)
          .call(legend);

        this.g.selectAll("#colorScale .label")
          .attr("font-size", opts.legendTextSize)
          .attr("transform", `translate(${opts.labelOffset}, ${0.75 * opts.size})`)

        this.g.select("#colorScale .legendTitle")
          .attr("transform", `translate(0, -${opts.titleOffset})`)
          .attr("font-size", opts.legendTextSize)
      }
    });

    return this;
  }

 scaleSize(options = {}) {
    const defaults = {
      nCells: 4, 
      shapePadding: 20, 
      labelOffset: 20, 
      yOffset: 100, 
      legendTextSize: 14,
      xOffset: 20,
      titleOffset: 20,
      symbolColor: "#a8a8a8"
    }

    const opts = { ...defaults, ...options };
    this.layers.push({
      type: 'scale_size',
      options: opts,
      render: () => {
        const legend = d3col.legendSize()
          .scale(this.rScale)
          .shape('circle')
          .shapePadding(opts.shapePadding)
          .labelOffset(opts.labelOffset)
          .orient('vertical')
          .title("λ(Hₙ)")
          .cells(opts.nCells);

        this.g.append("g")
          .attr("class", "legend")
          .attr("id", "sizeScale")
          .attr("transform", `translate(${this.plotWidth + opts.xOffset}, ${opts.yOffset})`)
          .call(legend);

        this.g.selectAll("#sizeScale .legendCells circle")
          .attr("fill", opts.symbolColor)
        this.g.selectAll("#sizeScale .label")
          .attr("font-size", opts.legendTextSize)
        this.g.select("#sizeScale .legendTitle")
          .attr("transform", `translate(0, -${opts.titleOffset})`)
          .attr("font-size", opts.legendTextSize)
        }
    })
  }

  // Add labels (title, x-axis, y-axis)
  labs(options = {}) {
    this.layers.push(
      annotation(this.svg, options, this.width, this.height, this.margin)
    );
    return this;
  }

  geomEdgeLink(options = {}) {
    const defaults = {
      "stroke-width": 1,
      "stroke": "#363E59",
      "opacity": 1,
      className: "edge_link"
    };

    const opts = { ...defaults, ...options };
    let N = options.N
    let link_data = flatten_edges(N, this.dataset, this.mappingObj);

    this.layers.push({
      type: "edge_link",
      options: opts,
      render: () => {
        this.g.select(`.${opts.className}`)
          .selectAll("line")
          .data(link_data, d => d._id).enter()
          .insert("line")
          .attr("class", opts.className)
          .attr("x1", d => this.xScale(d.x1))
          .attr("y1", d => this.yScale(d.y1))
          .attr("x2", d => this.xScale(d.x2))
          .attr("y2", d => this.yScale(d.y2))
          .attr("stroke-width", opts["stroke-width"])
          .attr("stroke", opts.stroke)
          .attr("opacity", opts.opacity)
      }
    })
  }

  interEdgeLink(options = {}) {
    const defaults = {
      strokeWidth: 1,
      backgroundOpacity: 0.2,
      className: "inter_edge_link",
      opacity: 1,
      otherClasses: ["ellipse"],
      stroke: "#363E59",
      highlightColor: "#363E59",
      highlightStrokeWidth: 1.5,
      threshold: 1
    };

    const opts = { ...defaults, ...options };
    let N = options.N

    this.layers.push({
      type: "inter_edge_link",
      options: opts,
      render: () => {
        this.g.insert("g")
          .attr("class", opts.className)

        for (let c in opts.otherClasses) {
          // only change stroke for ellipses
          this.g.select(`.${opts.otherClasses[c]}`)
            .selectAll("ellipse")
            .attr("stroke", d => Object.keys(N).indexOf(d._id.toString()) == -1 ? "white" : opts.highlightColor)
            .attr("stroke-width", d => Object.keys(N).indexOf(d._id.toString()) == -1 ? 0 : opts.highlightStrokeWidth)
        }

        let freeze = false;
        this.svg.on("mousemove", ev => {
          if (!freeze) {
              link_update(
              ev, this.g, N, this.dataset, this.mappingObj, this.xScale, this.yScale, opts, this.margin
            );
          }
        });
        this.svg.on("dblclick", () => { freeze = !freeze; });
      }
    })
  }

  interIsometry(options = {}) {
    const defaults = {
      backgroundOpacity: 0.2,
      className: "inter_isometry",
      magnify: 1,
      metric_bw: 10,
      stroke: "#a5a5a5",
      strokeWidth: 1.5,
      otherClasses: ["ellipse"],
      transformation_bw: 2
    }
    const opts = { ...defaults, ...options }
    this.layers.push({
      type: "inter_isometry",
      options: opts,
      render: () => {
        let metrics = options.metrics

        // connections between original and background points
        this.g.append("g")
          .lower()
          .attr("class", "isometry-links")
          .selectAll("line")
          .data(this.dataset, d => d._id).enter()
          .append("line")
          .attr("x1", d => this.xScale(d[this.mappingObj.x]))
          .attr("y1", d => this.yScale(d[this.mappingObj.y]))
          .attr("x2", d => this.xScale(d[this.mappingObj.x]))
          .attr("y2", d => this.yScale(d[this.mappingObj.y]))
          .attr("stroke", opts.stroke)
          .attr("stroke-width", opts.strokeWidth)

        // original locations
        for (let c in opts.otherClasses) {
          this.g.select(`.${opts.otherClasses[c]}`)
            .clone(true)
            .lower()
            .attr("opacity", opts.backgroundOpacity)
            .attr("class", `${opts.otherClasses[c]}-background`)
            .selectAll("*")
            .data(this.dataset)
        }

        // respond to mousemove
        let freeze = false;
        this.svg.on("mousemove", ev => {
          if (!freeze) {
            isometry_update(
              ev, this.g, this.dataset, metrics, this.mappingObj, this.xScale,
              this.yScale, this.rScale, opts, this.margin
            );
          }
        });
        this.svg.on("dblclick", () => { freeze = !freeze; });
      }
    })
  }

  interBoxplot(distance_summaries, outliers, options = {}) {
    const defaults = {
      backgroundOpacity: 0.2,
      className: "inter_boxplot",
      fill: "#bcbcbc",
      highlightColor: "#363E59",
      highlightStrokeWidth: 1.5,
      legendOffset: 60,
      opacity: 1,
      otherClasses: ["ellipse"],
      outlierRadius: 2,
      relHeight: 0.3,
      relPanelMargin: 0.05,
      relWidth: 0.75,
      stroke: "#262626",
      strokeWidth: 1
    }

    // determine relative sizes of the boxplot region
    const opts = { ...defaults, ...options }
    let y_vals = outliers.map(d => d.value)
      .concat(distance_summaries.map(d => d.q1))
      .concat(distance_summaries.map(d => d.q3))

    this.xScale.range([0, this.plotWidth * (opts.relWidth - opts.relPanelMargin)])
    this.xBoxScale = d3.scaleBand()
      .domain([...new Set(distance_summaries.map(d => d.bin))])
      .range([opts.relWidth * this.plotWidth, this.plotWidth])
    this.yBoxScale = d3.scaleLinear()
      .domain([0, d3.max(y_vals)])
      .range([opts.relHeight * this.plotHeight, 0])

    this.layers.push({
      type: "inter_isometry",
      options: opts,
      render: () => {

        // avoid collision with the legend
        this.g.select(".legend")
          .attr("transform", `translate(${opts.relWidth * this.plotWidth}, ${opts.relHeight * this.plotHeight + opts.legendOffset})`)

        this.opts = opts
        draw_boxplot(this, distance_summaries, outliers)
      }
    })
  }

  // Render all layers of the plot
  render() {
    const xAxis = d3.axisBottom(this.xScale);
    const yAxis = d3.axisLeft(this.yScale);
   this.g.append("g")
      .attr("class", "x-axis")
      .attr("transform", `translate(0, ${this.plotHeight})`)
      .call(xAxis.tickFormat("").tickSize(0))
      .selectAll(".tick text")
      .attr("font-size", this.labelFontSize);

    this.g.append("g")
      .attr("class", "y-axis")
      .call(yAxis.tickFormat("").tickSize(0))
      .selectAll(".tick text")
      .attr("font-size", this.labelFontSize);

    // Render all layers in order they were added
    this.layers.forEach(l => l.render());
    return this;
  }
}