distortions/lib/inter_isometry.js

Helpers for visualizing distortion in nonlinear dimensionality reduction.

import * as d3 from "https://esm.sh/d3@7";
import { SVD } from "https://esm.sh/svd-js";
import { matrix, add, subtract, multiply, inv, diag, zeros, size, transpose, det } from 'https://esm.sh/mathjs';

export function isometry_update(ev, g, dataset, metrics, mappingObj, xScale, yScale, rScale, opts, margin) {
    // metric at the current mouseover location
    let f_star = [xScale.invert(ev.layerX - margin.left), yScale.invert(ev.layerY - margin.top)]
    let { h_star } = local_metric(metrics, f_star, dataset, opts.metric_bw)
    let kn_t = local_metric(metrics, f_star, dataset, opts.transformation_bw)["kn"]

    // normalize so that transformation is between 0 and 1
    kn_t = kn_t.map(k => k / d3.max(kn_t))
    let h_star_ = inv(square_root_reorient(h_star))
    let h_star_inv = inv(h_star)

    // get transformed coordinates
    let new_coords = []
    let N = Object.values(dataset).length
    for (let n = 0; n < N; n++) {
        let f_n = matrix([dataset[n][mappingObj.x], dataset[n][mappingObj.y]])
        let f_tilde_n = add(multiply(h_star_, subtract(f_n, f_star)), f_star)
        let f_n_transform = add(multiply(kn_t[n], f_tilde_n), multiply(1 - kn_t[n], f_n))

        // new SVD data, for ellipse orientation
        let h_product = multiply(h_star_inv, matrix(metrics[n]))
        let { q, v } = SVD(h_product._data)

        let sv_transform = [
            kn_t[n] * q[0] + (1 - kn_t[n]) * dataset[n]["s0"],
            kn_t[n] * q[1] + (1 - kn_t[n]) * dataset[n]["s1"]
        ]
        let v_transform = [
            kn_t[n] * v[0][0] + (1 - kn_t[n]) * dataset[n]["x0"],
            kn_t[n] * v[0][1] + (1 - kn_t[n]) * dataset[n]["y0"]
        ]

        new_coords.push({
            [mappingObj.x]: f_n_transform._data[0],
            [mappingObj.y]: f_n_transform._data[1],
            "s0": sv_transform[0], "s1": sv_transform[1],
            "x0": v_transform[0], "y0": v_transform[1],
            "new_angle": angle(v_transform[0], v_transform[1]),
            "_id": dataset[n]["_id"]
        })
    }

    for (let c in opts.otherClasses) {
        g.select(`.${opts.otherClasses[c]}`)
            .selectAll("*")
            .data(new_coords, d => d._id)
            .attr("cx", d => xScale(d[mappingObj.x]))
            .attr("cy", d => yScale(d[mappingObj.y]))
            .attr("rx", d => rScale(d[mappingObj.a]))
            .attr("ry", d => rScale(d[mappingObj.b]))
            .attr("transform", d => `rotate(${d["new_angle"]} ${xScale(d[mappingObj.x])} ${yScale(d[mappingObj.y])})`)
    }

    g.select(".isometry-links")
        .selectAll("line")
        .data(new_coords, d => d._id)
        .attr("x1", d => xScale(d[mappingObj.x]))
        .attr("y1", d => yScale(d[mappingObj.y]))
}

function similarities(f_star, dataset, gamma) {
    let result = []
    for (let n = 0; n < Object.values(dataset).length; n++) {
        let f_n = [dataset[n].embedding_0, dataset[n].embedding_1];
        result.push(similarity(f_star, f_n, gamma));
    }
    let total = d3.sum(result);
    return result.map(d => d / total)
}

function local_metric(metrics, f_star, dataset, gamma) {
    let N = Object.values(metrics).length
    let h0 = matrix(metrics[0])
    let h_star = zeros(size(h0));
    let kn = similarities(f_star, dataset, gamma)

    for (let n = 0; n < N; n++) {
        h_star = add(h_star, multiply(kn[n], matrix(metrics[n])))
    }

    let { q } = SVD(h_star._data)
    return { "h_star": h_star, "sv": q, "kn": kn }
}

function square_root_reorient(A) {
    let svd_result = SVD(A._data)
    let { v, q } = reorient(svd_result)
    return multiply(matrix(v), diag(q.map(qk => Math.sqrt(qk))))
}

function reorient(svd_result) {
    let v = matrix(svd_result["v"])
    let q = svd_result["q"]

    if (Math.abs(v._data[0][0]) < Math.abs(v._data[0][1])) {
        let P = matrix([[0, 1], [1, 0]])
        v = multiply(v, P)
        q = [q[1], q[0]]
    }
    if (det(v) < 0) {
        v = multiply(v, diag([1, -1]))
    }
    if (v._data[0][0] < 0) {
        v = multiply(v, diag([-1, -1]))
    }

    return { "v": v, "q": q }
}

function similarity(a, b, gamma = 1) {
    let d = Math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2);
    let result = Math.exp(-gamma * d);
    if (isNaN(result)) {
        return 0;
    }
    return result;
}

function angle(x1, y1) {
    return Math.atan(y1 / x1) * (180 / Math.PI) + 90;
}