ExpMap Javascript Port

expmap.ts

type Tangents = Array<{ tan1: Vector3; tan2: Vector3 }>;
  
let cache: {
  G: KDistances;
  tangents: Tangents;
} = {
  G: new KDistances(),
  tangents: [],
};
  
export interface ExpMapParams {
  graphNbrs: number;
  stopDist: number;
  vertex: number;
}

export const expMap = (g: BufferGeometry, params: ExpMapParams) => {
  const hasCache = cache.tangents.length !== 0;
  console.log("running expmap", hasCache);
  if (!hasCache) {
    // nearest neighbors from all vertices in the mesh. 
    cache.G = findKnn_tree(g, params.graphNbrs); 
    cache.tangents = getTan1Tan2(
      new Float32BufferAttribute(g.attributes.normal.array, 3)
    );
  }
  console.log("running dijkstras algorithm");
  const distanceThres = 2 * params.stopDist;
  
  // geodesic distance finding between source vertex and all other vertices.
  const dist = dijkstrasMaxdist(cache.G, params.vertex, distanceThres); 
  const sort_dist = Object.keys(dist)
    .map((k) => ({ key: Number(k), value: dist[Number(k)] }))
    .sort((a, b) => a.value - b.value);
  const uvs: number[] = [];
  const N = g.attributes.normal.count;
  for (let i = 0; i < N; i++) {
    uvs.push(NaN, NaN);
  }
  const si = sort_dist[0].key;
  uvs[si * 2] = 0;
  uvs[si * 2 + 1] = 0;
  const normals = g.attributes.normal.array;
  const points = g.attributes.position.array;
  for (let i = 1; i < N; i++) {
    const vtx = sort_dist[i].key;
    if (dist[vtx] > params.stopDist) break;
    // get all neighbors with distance greater than 0
    const neighbors = cache.G.kDistances[vtx].toArray
      .filter((f) => f.value !== 0)
      .filter((n) => {
        const u = uvs[n.index * 2];
        return Number.isFinite(u);
      });
    let wtSum = 0;
    let sumUV = new Vector2(0, 0);
    for (let j = 0; j < neighbors.length; j++) {
      const ui = neighbors[j].index;
      const from = new Vector3(
        points[3 * vtx],
        points[3 * vtx + 1],
        points[3 * vtx + 2]
      );
      const to = new Vector3(
        points[3 * ui],
        points[3 * ui + 1],
        points[3 * ui + 2]
      );
      const weight = 1 / from.distanceToSquared(to);
      const uv_jj = estimateUV(
        vtx,
        ui,
        si,
        points,
        normals,
        uvs,
        cache.tangents
      );
      sumUV.add(uv_jj.multiplyScalar(weight));
      wtSum += weight;
    }
    [uvs[2 * vtx], uvs[2 * vtx + 1]] = [sumUV.x / wtSum, sumUV.y / wtSum];
  }
  return uvs;
};
const estimateUV = (
  i: number,
  ui: number,
  si: number,
  points: ArrayLike<number>,
  normals: ArrayLike<number>,
  uv: number[],
  tangents: Tangents
) => {
  const ns = new Vector3(
    normals[3 * si],
    normals[3 * si + 1],
    normals[3 * si + 2]
  );
  const t1s = tangents[si].tan1;
  const nu = new Vector3(
    normals[3 * ui],
    normals[3 * ui + 1],
    normals[3 * ui + 2]
  );
  const t1u = tangents[ui].tan1;
  const t2u = tangents[ui].tan2;
  const from = new Vector3(points[3 * i], points[3 * i + 1], points[3 * i + 2]);
  const to = new Vector3(
    points[3 * ui],
    points[3 * ui + 1],
    points[3 * ui + 2]
  );
  const v = from.clone().sub(to);
  const len = Math.sqrt(v.x * v.x + v.y * v.y + v.z * v.z);
  let localUv = new Vector2(v.clone().dot(t1u), v.clone().dot(t2u));
  localUv = localUv.normalize().multiplyScalar(len);
  const nrot = vAlign(ns, nu).rotmat;
  const rt1s = matMul3(nrot, t1s);
  let { axis, theta } = vAlign(rt1s, t1u);
  if (axis.dot(nu) < 0) {
    theta = -theta;
  }
  const cosA = Math.cos(theta);
  const sinA = Math.sin(theta);
  localUv = matMul2(
    new Matrix3().fromArray([cosA, sinA, 0, -sinA, cosA, 0, 0, 0, 1]),
    localUv
  );
  const currentUV = new Vector2(uv[2 * ui], uv[2 * ui + 1]);
  return currentUV.add(localUv);
};
const matMul2 = (mat: Matrix3, vec: Vector2) => {
  const [x, y] = [
    mat.elements[0] * vec.x + mat.elements[1] * vec.y,
    mat.elements[3] * vec.x + mat.elements[4] * vec.y,
  ];
  return new Vector2(x, y);
};
const matMul3 = (mat: Matrix3, vec: Vector3) => {
  const [x, y, z] = [
    mat.elements[0] * vec.x + mat.elements[1] * vec.y + mat.elements[2] * vec.z,
    mat.elements[3] * vec.x + mat.elements[4] * vec.y + mat.elements[5] * vec.z,
    mat.elements[6] * vec.x + mat.elements[7] * vec.y + mat.elements[8] * vec.z,
  ];
  return new Vector3(x, y, z);
};
const vAlign = (from: Vector3, to: Vector3) => {
  const n1 = from.clone();
  const n2 = to.clone();
  const cosTheta = n1.clone().dot(n2) / (norm(n1) * norm(n2));
  const axis = n1.clone().cross(n2);
  if (norm(axis) > 0) {
    axis.divideScalar(norm(axis));
    const theta = Math.acos(clamp(cosTheta, -1, 1));
    const rotmat = axisrot(axis, theta);
    return { axis, theta, rotmat };
  }
  if (cosTheta < 0) {
    const tan = tangentFrame(n1);
    return {
      axis: tan.tan1,
      theta: Math.PI,
      rotmat: axisrot(tan.tan1, Math.PI),
    };
  }
  return {
    axis: new Vector3(0, 0, 1),
    theta: 0,
    rotmat: new Matrix3().identity(),
  };
};
const norm = (v: Vector3) => Math.sqrt(v.x * v.x + v.y * v.y + v.z * v.z);
const axisrot = (axis: Vector3, t: number) => {
  const fcos = Math.cos(t);
  const fsin = Math.sin(t);
  const fminuscos = 1 - fcos;
  const fx2 = axis.x * axis.x;
  const fy2 = axis.y * axis.y;
  const fz2 = axis.z * axis.z;
  const fxym = axis.x * axis.y * fminuscos;
  const fxzm = axis.x * axis.z * fminuscos;
  const fyzm = axis.y * axis.z * fminuscos;
  const fxsin = axis.x * fsin;
  const fysin = axis.y * fsin;
  const fzsin = axis.z * fsin;
  return new Matrix3().fromArray([
    fx2 * fminuscos + fcos,
    fxym - fzsin,
    fxzm + fysin,
    fxym + fzsin,
    fy2 * fminuscos + fcos,
    fyzm - fxsin,
    fxzm - fysin,
    fyzm + fxsin,
    fz2 * fminuscos + fcos,
  ]);
};