bands-visualiser/src/bandvisualiser.js

A plotly.js renderer for Bands, DOS, combined Bands/Dos and fatBands.

import Plotly from "plotly.js-basic-dist-min";
import deepmerge from "deepmerge";

import { plotCombinedLayout } from "./utils/MCXDPlotTemplates.js";
import {
  mergeLegendTracesByName,
  getColor,
  addfermiLines,
} from "./utils/plotUtils.js";
import { getBandStructure } from "./getBandStructure.js";
import { getDos } from "./getDos";
import { getProjectionTraces } from "./getProjectionTraces";

import { ensureArray, sanitizeBandsInput } from "./utils/dataUtils.js";

function handleSettings(settings, layout) {
  for (const key in settings) {
    if (
      settings[key] !== null &&
      typeof settings[key] === "object" &&
      !Array.isArray(settings[key]) &&
      layout[key] &&
      typeof layout[key] === "object" &&
      !Array.isArray(layout[key])
    ) {
      // Deep merge nested objects (like layout.xaxis, layout.title, etc)
      layout[key] = deepmerge(layout[key], settings[key]);
    } else {
      // Replace primitive values or arrays directly
      layout[key] = settings[key];
    }
  }
}

export function BandsVisualiser(
  container,
  {
    bandsDataArray = [],
    dosDataArray = [],
    fermiLevels = [],
    customTraces = [],
    settings = {},
  }
) {
  let bandTraces = [],
    bandLayout = {},
    projTraces = [],
    dosTraces = [],
    dosLayout = {};

  // normalize inputs as arrays, in the case of a single bandsDataObject being passed.
  bandsDataArray = ensureArray(bandsDataArray);
  dosDataArray = ensureArray(dosDataArray);
  fermiLevels = ensureArray(fermiLevels);
  customTraces = ensureArray(customTraces);

  // santize array if it doesnt adhere to the internal api.
  bandsDataArray = sanitizeBandsInput(bandsDataArray);
  console.log(bandsDataArray);

  const hasBands = bandsDataArray.length > 0;
  const hasDOS = dosDataArray.length > 0;

  if (hasBands) {
    const result = getBandStructure(bandsDataArray);
    bandTraces = result.bandTraces;
    bandLayout = result.bandLayout;

    for (const [index, bandEntry] of bandsDataArray.entries()) {
      if (!bandEntry.hasOwnProperty("projections")) continue;
      for (const [i, pr] of bandEntry.projections.entries()) {
        const formatting =
          pr.traceFormat && Object.keys(pr.traceFormat).length > 0
            ? pr.traceFormat
            : {
                name: `Projection ${i}`,
                line: { color: getColor(i), width: 0 },
              };
        projTraces.push(
          ...getProjectionTraces(
            bandEntry.bandsData.x, // this may break if old aiida format
            bandEntry.bandsData.y, // this may break if old aiida format
            pr.projData.weights,
            formatting
          )
        );
      }
    }
  }

  if (hasDOS) {
    const orientation = hasBands ? "v" : "h";
    const result = getDos(dosDataArray, orientation);
    dosTraces = result.dosTraces;
    dosLayout = result.dosLayout;
  }

  let layout = {},
    traces = [];

  if (hasDOS && !hasBands) {
    traces = dosTraces;
    layout = { ...dosLayout, showlegend: dosLayout?.showlegend ?? true };
  } else if (hasBands && !hasDOS) {
    traces = [...bandTraces, ...projTraces];
    layout = { ...bandLayout, showlegend: bandLayout?.showlegend ?? true };
  } else if (hasDOS && hasBands) {
    dosTraces = dosTraces.map((trace) => ({
      ...trace,
      xaxis: "x2",
      yaxis: "y",
    }));

    // merge projection legend entries onto dos legend entries
    const mergedTraces = mergeLegendTracesByName([...projTraces, ...dosTraces]);
    traces = [...bandTraces, ...mergedTraces];

    layout = {
      ...plotCombinedLayout,
      xaxis: {
        ...(plotCombinedLayout.xaxis || {}),
        ...(bandLayout?.xaxis || {}),
      },
      xaxis2: {
        ...(plotCombinedLayout.xaxis2 || {}),
        ...(dosLayout?.xaxis2 || {}),
      },
      yaxis: {
        ...(plotCombinedLayout.yaxis || {}),
        ...(bandLayout?.yaxis || {}),
        ...(dosLayout?.yaxis || {}),
      },
    };
  }

  // Add custom traces
  for (const item of customTraces) {
    console.log(item);
    if (item.target === "dos") {
      traces.push({ ...item, xaxis: "x2", yaxis: "y" });
    } else if (item.target === "bands") {
      traces.push({ ...item, xaxis: "x", yaxis: "y" });
    } else {
      traces.push({ ...item });
    }
  }
  // add fermi lines
  const traces_ = addfermiLines(fermiLevels, traces, hasDOS, hasBands);
  handleSettings(settings, layout);

  Plotly.react(container, traces_, layout, {
    displaylogo: false,
    responsive: true,
  });
}

export default BandsVisualiser;