Wrap atoms on move

docs/source/atoms/edit.rst

@@ -50,6 +50,18 @@ to apply the transform and click to confirm.
 - ``r`` rotate
 - ``s`` scale
 
+Wrap On Move (PBC)
+-------------------
+By default, translations do not wrap atoms back into the unit cell. To enable
+PBC wrapping after moving atoms, set ``wrapOnMove`` to ``true`` (a valid cell
+and ``pbc`` flags are required).
+
+.. code-block:: javascript
+
+   editor.avr.wrapOnMove = true;
+
+You can also toggle this in the GUI under **Boundary**.
+
 Rotation defaults to the camera axis through the selection center.
 To rotate around a custom axis, press ``r`` to enter rotate mode, then press ``a`` and click one, two, or three atoms, then press ``a`` again to exit axis picking.
 One atom sets the rotation center (camera axis), two atoms define the bond axis, and three atoms define the plane normal through their centroid.

src/atoms/AtomsViewer.js

@@ -2,7 +2,7 @@ import * as THREE from "three";
 import { CellManager } from "./cell.js";
 import { AtomManager } from "./plugins/atom.js";
 import { BondManager, defaultBondRadius, searchBondedAtoms } from "./plugins/bond.js";
-import { clearObjects, clearObject, toIndexArray, toVector3 } from "../utils.js";
+import { clearObjects, clearObject, toIndexArray, toVector3, calculateCartesianCoordinates, calculateInverseMatrix, multiplyMatrixVector } from "../utils.js";
 import { PolyhedraManager } from "./plugins/polyhedra.js";
 import { BoundaryManager } from "./plugins/boundary.js";
 import { AtomLabelManager } from "./plugins/atomLabel.js";
@@ -38,6 +38,7 @@ class AtomsViewer {
     this._showBondedAtoms = viewerSettings.showBondedAtoms;
     this._boundary = viewerSettings.boundary;
     this._atomScale = viewerSettings.atomScale;
+    this._wrapOnMove = viewerSettings.wrapOnMove;
     this._backgroundColor = viewerSettings.backgroundColor;
     this.tjs.scene.background = new THREE.Color(this._backgroundColor);
     this._selectedAtomsIndices = new Array(); // Store selected atoms
@@ -636,6 +637,19 @@ class AtomsViewer {
     this.applyState({ atomScale: newValue }, { redraw: "render" });
   }
 
+  get wrapOnMove() {
+    return this._wrapOnMove;
+  }
+
+  set wrapOnMove(newValue) {
+    if (this._syncingState) {
+      this._wrapOnMove = newValue;
+      this.weas.eventHandlers.dispatchViewerUpdated({ wrapOnMove: newValue });
+      return;
+    }
+    this.applyState({ wrapOnMove: newValue }, { redraw: "none" });
+  }
+
   get atomScales() {
     return this._atomScales;
   }
@@ -1036,18 +1050,60 @@ class AtomsViewer {
   }
 
   setAtomPosition({ index, position }) {
+    const nextPosition = this.wrapPositionIfNeeded(position);
     // Update the atom position
     const matrix = new THREE.Matrix4();
     this.atomManager.meshes["atom"].getMatrixAt(index, matrix);
-    matrix.setPosition(position);
+    matrix.setPosition(nextPosition);
     this.atomManager.meshes["atom"].setMatrixAt(index, matrix);
-    this.atoms.positions[index] = [position.x, position.y, position.z];
+    this.atoms.positions[index] = [nextPosition.x, nextPosition.y, nextPosition.z];
     // update the other meshes
     this.atomManager.updateImageAtomsMesh(index);
     this.bondManager.updateBondMesh(index);
     this.polyhedraManager.updatePolyhedraMesh(index);
   }
 
+  wrapPositionIfNeeded(position) {
+    if (!this._wrapOnMove) {
+      return position;
+    }
+    if (!this.atoms || !Array.isArray(this.atoms.pbc) || !this.atoms.pbc.some(Boolean)) {
+      return position;
+    }
+    if (typeof this.atoms.isUndefinedCell === "function" && this.atoms.isUndefinedCell()) {
+      return position;
+    }
+    const cell = this.atoms.cell;
+    if (!Array.isArray(cell) || cell.length !== 3) {
+      return position;
+    }
+    try {
+      const cellT = cell[0].map((_, i) => cell.map((row) => row[i]));
+      const invCellT = calculateInverseMatrix(cellT);
+      const fractional = multiplyMatrixVector(invCellT, [position.x, position.y, position.z]);
+      let changed = false;
+      for (let i = 0; i < 3; i++) {
+        if (!this.atoms.pbc[i]) {
+          continue;
+        }
+        const f = fractional[i];
+        const wrapped = f - Math.floor(f);
+        if (wrapped !== f) {
+          changed = true;
+        }
+        fractional[i] = wrapped;
+      }
+      if (!changed) {
+        return position;
+      }
+      const wrappedCartesian = calculateCartesianCoordinates(cell, fractional);
+      return new THREE.Vector3(wrappedCartesian[0], wrappedCartesian[1], wrappedCartesian[2]);
+    } catch (err) {
+      this.logger.debug("wrapPositionIfNeeded failed:", err);
+      return position;
+    }
+  }
+
   resetSelectedAtomsPositions(initialAtomPositionsOrOptions, indices = null) {
     let initialAtomPositions = initialAtomPositionsOrOptions;
     if (initialAtomPositionsOrOptions && typeof initialAtomPositionsOrOptions === "object" && Object.prototype.hasOwnProperty.call(initialAtomPositionsOrOptions, "initialAtomPositions")) {

src/atoms/atomsGui.js

@@ -224,6 +224,13 @@ class AtomsGUI {
       );
     });
     boundaryFolder.add({ apply: () => this.applyBoundaryChanges() }, "apply").name("Apply Changes");
+    this.wrapOnMoveController = boundaryFolder
+      .add({ wrapOnMove: this.viewer.wrapOnMove }, "wrapOnMove")
+      .name("Wrap On Move")
+      .onChange((value) => {
+        if (this.isSyncing || this.viewer.weas.ops.isRestoring) return;
+        this.viewer.setState({ wrapOnMove: value }, { record: true, redraw: "none" });
+      });
   }
 
   addColorControl() {
@@ -448,6 +455,9 @@ class AtomsGUI {
         case "boundary":
           this.updateBoundary(value);
           break;
+        case "wrapOnMove":
+          this.updateWrapOnMove(value);
+          break;
         default:
           break;
       }
@@ -534,6 +544,12 @@ class AtomsGUI {
       }
     }
   }
+
+  updateWrapOnMove(newValue) {
+    if (this.wrapOnMoveController && this.wrapOnMoveController.getValue() !== newValue) {
+      this.wrapOnMoveController.setValue(newValue);
+    }
+  }
 }
 
 export { AtomsGUI };

src/config.js

@@ -30,6 +30,7 @@ const defaultViewerSettings = {
     [0, 1],
   ],
   atomScale: 0.4, // Default atom scale
+  wrapOnMove: false, // Wrap atoms into the unit cell after moving
   backgroundColor: "#ffffff", // Default background color (white)
   logLevel: "warn", // Default log level
   continuousUpdate: true, // Default continuous update

src/index.d.ts

@@ -163,6 +163,8 @@ export declare class AtomsViewer {
   boundary: any;
   /** Atom scale factor */
   atomScale: number;
+  /** Whether to wrap atoms into the unit cell after moving */
+  wrapOnMove: boolean;
   /** Per-atom scale overrides */
   atomScales: number[];
   /** Per-atom stick style overrides */

src/state/defaultState.js

@@ -15,6 +15,7 @@ function createDefaultState() {
       showBondedAtoms: viewerDefaults.showBondedAtoms,
       boundary: viewerDefaults.boundary,
       atomScale: viewerDefaults.atomScale,
+      wrapOnMove: viewerDefaults.wrapOnMove,
       atomScales: [],
       modelSticks: [],
       modelPolyhedras: [],

tests/e2e/gui.spec.js

@@ -231,6 +231,34 @@ test("Transform Translate Axis", async ({ page }) => {
   await expect(await page.evaluate(() => window.editor.eventHandlers.transformControls.translatePlanePending)).toBe(false);
 });
 
+test("Wrap on move", async ({ page }) => {
+  await page.goto("http://127.0.0.1:8080/tests/e2e/testCrystal.html");
+  await page.waitForFunction(() => window.editor);
+  // focus the element
+  const element = await page.$("#viewer");
+  await element.focus();
+  const boundingBox = await element.boundingBox();
+  // Calculate the center of the element
+  const centerX = boundingBox.x + boundingBox.width / 2;
+  const centerY = boundingBox.y + boundingBox.height / 2;
+  page.centerX = centerX;
+  page.centerY = centerY;
+  // Move the mouse to the center of the element
+  await page.mouse.move(centerX, centerY);
+
+  await page.evaluate(() => {
+    const editor = window.editor;
+    editor.avr.selectedAtomsIndices = [1];
+    editor.avr.wrapOnMove = true;
+    editor.eventHandlers.currentMousePosition.set(300, 200);
+    editor.eventHandlers.transformControls.enterMode("translate", editor.eventHandlers.currentMousePosition);
+    editor.tjs.render();
+  });
+  // mouse move to the center of the canvas element
+  await page.mouse.move(page.centerX - 100, page.centerY);
+  await expect.soft(page).toHaveScreenshot("Wrap-on-move.png");
+});
+
 test("Text Manager", async ({ page }) => {
   await page.goto("http://127.0.0.1:8080/tests/e2e/testTextManager.html");
   await expect.soft(page).toHaveScreenshot("TextManager.png");