refactor Phonon plugin (#14)

Add a Phonon module to generate the trajectory from a phonon mode.

demo/demo.js

@@ -140,21 +140,32 @@ async function updateAtoms(filename, fileContent = null) {
       filename = "graphene.cif";
       structureData = fileContent || (await fetchFile(filename));
       atoms = parseCIF(structureData);
-      editor.avr.fromPhononTrajectory(
-        atoms,
-        [
-          [0, 1, 0],
-          [0, -1, 0],
+      editor.avr.fromPhononMode({
+        atoms: atoms,
+        eigenvectors: [
+          [
+            [0, 0.5],
+            [1, 0.5],
+            [0, 0.5],
+          ],
+          [
+            [0, 1],
+            [-1, 0.5],
+            [0, 1],
+          ],
         ],
-        1,
-        50,
-      );
+        amplitude: 1,
+        nframes: 50,
+        kpoint: [0.5, 0.5, 0.5],
+        repeat: [4, 4, 1],
+      });
+
+      // control the speed of the animation
       editor.avr.boundary = [
-        [-2, 3],
-        [-2, 3],
-        [0, 1],
+        [-0.01, 1.01],
+        [-0.01, 1.01],
+        [-0.01, 1.01],
       ];
-      // control the speed of the animation
       editor.avr.frameDuration = 20;
       editor.avr.VFManager.addSetting({ origins: "positions", vectors: "movement", color: "#ff0000", radius: 0.1 });
       editor.avr.modelStyle = 1;

docs/source/atoms/vector_field.rst

@@ -53,7 +53,26 @@ In order to update the vector fields when the atom is moved in a animation, the
     // update the atoms viewer (avr) from phonon trajectory
     // inputs: atoms, egenvectors, amplitude and nframes
     // the egenvectors should have the same length as the atoms
-    editor.avr.fromPhononTrajectory(atoms, [[0, 1, 0], [0, -1, 0]], 1, 15);
+    // the egenvector consists of the real and imaginary part
+    editor.avr.fromPhononMode({
+        atoms: atoms,
+        eigenvectors: [
+          [
+            [0, 0.5],
+            [1, 0.5],
+            [0, 0.5],
+          ],
+          [
+            [0, 1],
+            [-1, 0.5],
+            [0, 1],
+          ],
+        ],
+        amplitude: 1,
+        nframes: 50,
+        kpoint: [0.5, 0.5, 0.5],
+        repeat: [4, 4, 1],
+      });
     // Then add a vector field (VF) to show the arrows of the movement
     editor.avr.VFManager.addSetting({ origins: "positions", vectors: "movement", color: "#ff0000", radius: 0.1 });
     editor.avr.drawModels();

src/atoms/AtomsViewer.js

@@ -14,6 +14,7 @@ import { Measurement } from "./plugins/measurement.js";
 import { getAtomColors } from "./color.js";
 import { AtomsGUI } from "./atomsGui.js";
 import { defaultViewerSettings } from "../config.js";
+import { Phonon } from "./plugins/phonon.js";
 
 class AtomsViewer {
   constructor({ weas, atoms = [new Atoms()], viewerConfig = {} }) {
@@ -55,6 +56,7 @@ class AtomsViewer {
     this.Measurement = new Measurement(this);
     this.VFManager = new VectorField(this);
     this.animate = this.animate.bind(this); // Bind once in the constructor
+    this._atoms = null;
     this.init(atoms);
   }
 
@@ -149,6 +151,13 @@ class AtomsViewer {
     this.tjs.render();
   }
 
+  get originalAtoms() {
+    if (this._atoms) {
+      return this._atoms;
+    }
+    return this.atoms;
+  }
+
   get atoms() {
     const atoms = this.trajectory[this.currentFrame];
     atoms.uuid = this.uuid;
@@ -170,6 +179,7 @@ class AtomsViewer {
     } else {
       this.trajectory = [atoms];
     }
+    this._atoms = null;
     this._currentFrame = 0;
     // set cell
     this.cellManager.atoms = this.atoms;
@@ -193,19 +203,16 @@ class AtomsViewer {
   }
 
   // set atoms from phonon trajectory
-  fromPhononTrajectory(atoms, eigenvectors, amplitude, nframes) {
-    const trajectory = [];
-    const times = Array.from({ length: nframes }, (_, i) => 2 * Math.PI * (i / nframes));
-    times.forEach((t) => {
-      const vectors = eigenvectors.map((vec) => vec.map((val) => val * amplitude * Math.sin(t)));
-      const newAtoms = atoms.copy();
-      for (let i = 0; i < newAtoms.positions.length; i++) {
-        newAtoms.positions[i] = newAtoms.positions[i].map((pos, j) => pos + vectors[i][j] / 5);
-      }
-      newAtoms.newAttribute("movement", vectors);
-      trajectory.push(newAtoms);
-    });
+  fromPhononMode({ atoms, eigenvectors, amplitude, nframes, kpoint = [0, 0, 0], repeat = [1, 1, 1] }) {
+    console.log("--------------------------------------From Phonon Mode--------------------------------------");
+    const phonon = new Phonon(atoms, kpoint, eigenvectors, true);
+    const trajectory = phonon.getTrajectory(amplitude, nframes, null, null, null, repeat);
     this.atoms = trajectory;
+    this._atoms = atoms.multiply(...repeat);
+    this._atoms.uuid = this.uuid;
+    this.drawModels();
+    this.play();
+    // console.log("this._atoms: ", this._atoms);
   }
 
   get ready() {
@@ -463,7 +470,7 @@ class AtomsViewer {
     // Map the symbols to their radii
     this.cutoffs = this.bondManager.buildBondDict();
     // find neighbor atoms in the original cell
-    this.neighbors = findNeighbors(this.atoms, this.cutoffs);
+    this.neighbors = findNeighbors(this.originalAtoms, this.cutoffs);
     if (this.debug) {
       console.log("neighbors: ", this.neighbors);
     }

src/atoms/atoms.js

@@ -278,7 +278,22 @@ class Atoms {
         }
       }
     }
-    console.timeEnd("multiply");
+    // repeat attributes for each atom
+    for (const name in this.attributes["atom"]) {
+      const values = this.attributes["atom"][name];
+      const newValues = [];
+      for (let ix = 0; ix < mx; ix++) {
+        for (let iy = 0; iy < my; iy++) {
+          for (let iz = 0; iz < mz; iz++) {
+            for (let i = 0; i < values.length; i++) {
+              newValues.push(values[i]);
+            }
+          }
+        }
+      }
+      newAtoms.newAttribute(name, newValues, "atom");
+    }
+    // console.timeEnd("multiply");
     // Return the new Atoms object
     return newAtoms;
   }

src/atoms/cell.js

@@ -102,7 +102,7 @@ export class CellManager {
   }
 
   drawUnitCellVectors() {
-    console.log("drawUnitCellVectors");
+    // console.log("drawUnitCellVectors");
     const cell = this.atoms.cell;
     if (!cell || cell.length !== 3) {
       console.warn("Invalid or missing unit cell data for vectors");

src/atoms/plugins/bond.js

@@ -46,7 +46,7 @@ export class BondManager {
     */
     console.log("init bond settings");
     this.settings = [];
-    const atoms = this.viewer.atoms;
+    const atoms = this.viewer.originalAtoms;
     const symbols = atoms.symbols;
     const speciesSet = new Set(symbols);
     const speciesList = Array.from(speciesSet);
@@ -118,7 +118,7 @@ export class BondManager {
     }
     // I don't add bonded atoms to offsets, because the bondlist will add them through the bondedAtoms["bonds"]
     // console.log("offsets: ", offsets);
-    this.bondList = buildBonds(this.viewer.atoms, offsets, this.viewer.neighbors["map"], this.viewer._boundary, this.viewer.modelSticks);
+    this.bondList = buildBonds(this.viewer.originalAtoms, offsets, this.viewer.neighbors["map"], this.viewer._boundary, this.viewer.modelSticks);
     // merge the bondList and the bondedAtoms["bonds"]
     this.bondList = this.bondList.concat(this.viewer.bondedAtoms["bonds"]);
     if (this.viewer.debug) {
@@ -130,7 +130,7 @@ export class BondManager {
     if (this.viewer.colorBy !== "Element") {
       atomColors = this.viewer.atomColors;
     }
-    this.bondMesh = drawStick(this.viewer.atoms, this.bondList, this.buildBondDict(), this.viewer.bondRadius, this.viewer._materialType, atomColors);
+    this.bondMesh = drawStick(this.viewer.originalAtoms, this.bondList, this.buildBondDict(), this.viewer.bondRadius, this.viewer._materialType, atomColors);
     return this.bondMesh;
   }
 
@@ -141,7 +141,7 @@ export class BondManager {
     */
     // console.log("updateBondMesh: ", atomIndex);
     if (atoms === null) {
-      atoms = this.viewer.atoms;
+      atoms = this.viewer.originalAtoms;
     }
     let bondIndices = [];
     if (atomIndex) {

src/atoms/plugins/phonon.js

@@ -0,0 +1,90 @@
+// Helper function to perform vector dot product
+function vec_dot(a, b) {
+  return a.reduce((sum, value, index) => sum + value * b[index], 0);
+}
+
+// Helper function to create a complex number in polar form
+function complexPolar(magnitude, angle) {
+  return {
+    real: magnitude * Math.cos(angle),
+    imag: magnitude * Math.sin(angle),
+    mult(other) {
+      return {
+        real: this.real * other.real - this.imag * other.imag,
+        imag: this.real * other.imag + this.imag * other.real,
+      };
+    },
+  };
+}
+
+// Helper function to multiply complex numbers
+function complexMult(complex, phase) {
+  return phase.mult({ real: complex[0], imag: complex[1] });
+}
+
+export class Phonon {
+  constructor(atoms, kpoint = null, eigenvectors = null, addatomphase = true) {
+    this.atoms = atoms;
+    this.kpoint = kpoint;
+    this.eigenvectors = eigenvectors;
+    this.addatomphase = addatomphase;
+    this.vibrations = [];
+  }
+
+  // Compute initial phases and vibrations
+  calculateVibrations() {
+    const fractional_positions = this.atoms.calculateFractionalCoordinates();
+    const natoms = this.atoms.positions.length;
+    let atom_phase = [];
+
+    if (this.addatomphase) {
+      atom_phase = fractional_positions.map((position) => vec_dot(this.kpoint, position));
+    } else {
+      atom_phase = new Array(natoms).fill(0);
+    }
+    console.log("atom_phase: ", atom_phase);
+
+    for (let i = 0; i < natoms; i++) {
+      let sprod = atom_phase[i];
+      let phase = complexPolar(1.0, sprod * 2.0 * Math.PI);
+      this.vibrations.push(this.eigenvectors[i].map((vector) => complexMult(vector, phase)));
+    }
+    console.log("vibrations: ", this.vibrations);
+  }
+
+  // Get the trajectory of the phonon mode
+  getTrajectory(amplitude, nframes, kpoint = null, eigenvectors = null, atoms = null, repeat = [1, 1, 1], addatomphase = null) {
+    if (atoms) {
+      this.atoms = atoms;
+    }
+    if (kpoint) {
+      this.kpoint = kpoint;
+    }
+    if (eigenvectors) {
+      this.eigenvectors = eigenvectors;
+    }
+    if (this.kpoint === null || this.eigenvectors === null) {
+      throw new Error("kpoint and eigenvectors must be provided");
+    }
+    if (addatomphase !== null) {
+      this.addatomphase = addatomphase;
+    }
+
+    this.calculateVibrations();
+    const trajectory = [];
+    const times = Array.from({ length: nframes }, (_, i) => 2 * Math.PI * (i / nframes));
+    times.forEach((t) => {
+      const newAtoms = this.atoms.copy();
+      let phase = complexPolar(amplitude, t);
+      const movement = [];
+      for (let i = 0; i < this.atoms.positions.length; i++) {
+        let displacement = this.vibrations[i].map((v) => phase.real * v.real);
+        newAtoms.positions[i] = this.atoms.positions[i].map((pos, index) => pos + displacement[index] / 5);
+        movement.push(displacement);
+      }
+      newAtoms.newAttribute("movement", movement);
+      trajectory.push(newAtoms.multiply(...repeat));
+    });
+    return trajectory;
+  }
+}

tests/e2e/gui.spec.js

@@ -362,7 +362,18 @@ test.describe("Phonon", () => {
     await page.mouse.move(centerX, centerY);
   });
 
-  test("Play", async ({ page }) => {
+  test("Frame", async ({ page }) => {
+    await page.evaluate(() => {
+      window.editor.avr.pause();
+      const timeline = document.getElementById("timeline");
+      timeline.value = 0;
+      // Creating and dispatching the event must happen within the page context
+      const event = new Event("input", {
+        bubbles: true,
+        cancelable: true,
+      });
+      timeline.dispatchEvent(event);
+    });
     await expect(page).toHaveScreenshot("Phonon-frame-0.png");
     // set frame 10
     await page.evaluate(() => {

tests/e2e/testPhonon.html

@@ -25,26 +25,39 @@
       fetchFile(filename).then((fileContent) => {
         const atoms = weas.parseCIF(fileContent);
         let editor = new weas.WEAS({ domElement });
-        editor.avr.fromPhononTrajectory(
-          atoms,
-          [
-            [0, 1, 0],
-            [0, -1, 0],
+        editor.avr.fromPhononMode({
+          atoms: atoms,
+          eigenvectors: [
+            [
+              [0, 0],
+              [1, 0],
+              [0, 0],
+            ],
+            [
+              [0, 0],
+              [-1, 0],
+              [0, 0],
+            ],
           ],
-          1,
-          50,
-        );
+          amplitude: 1,
+          nframes: 50,
+          kpoint: [0, 0, 0],
+          repeat: [4, 4, 1],
+        });
+
+        // control the speed of the animation
         editor.avr.boundary = [
-          [-2, 3],
-          [-2, 3],
-          [0, 1],
+          [-0.01, 1.01],
+          [-0.01, 1.01],
+          [-0.01, 1.01],
         ];
         // control the speed of the animation
         editor.avr.frameDuration = 20;
         editor.avr.VFManager.addSetting({ origins: "positions", vectors: "movement", color: "#ff0000", radius: 0.1 });
         editor.avr.modelStyle = 1;
         editor.avr.bondManager.hideLongBonds = false;
         editor.render();
+        window.editor = editor;
       });
     </script>
   </body>