Fix pbc when exporting to ASE or Pymatgen format (#80)

README.md

@@ -135,7 +135,9 @@ volume, atoms = read_cube_data("h2o-homo.cube")
 viewer = WeasWidget()
 viewer.from_ase(atoms)
 viewer.avr.iso.volumetric_data = {"values": volume}
-viewer.avr.iso.settings = [{"isovalue": 0.0001, "mode": 0}]
+viewer.avr.iso.settings = {"positive": {"isovalue": 0.001},
+                           "negative": {"isovalue": -0.001, "color": "yellow"}
+                           }
 viewer
 ```
 <img src="docs/source/_static/images/example-isosurface.png"  width="300px"/>
@@ -166,16 +168,21 @@ Animate vibrational (phonon) modes (computed with external software).
 import numpy as np
 from ase.build import bulk
 from weas_widget import WeasWidget
-from weas_widget.utils import generate_phonon_trajectory
-
 atoms = bulk("Fe", cubic=True)
-eigenvector = np.array([[0, -0.0, 0.5], [0, 0.0, -0.5]])
-trajectory = generate_phonon_trajectory(atoms, eigenvector, repeat=[4, 4, 1])
+phonon_setting = {"eigenvectors": np.array([[[0, 0], [0, 0],[0.5, 0]],
+                                    [[0, 0], [0, 0], [-0.5, 0]]]
+                                    ),
+        "kpoint": [0, 0, 0], # optional
+        "amplitude": 5, # scale the motion of the atoms
+        "factor": 1.5, # scale the length of the arrows
+        "nframes": 20,
+        "repeat": [4, 4, 1],
+        "color": "blue",
+        "radius": 0.1,
+        }
 viewer = WeasWidget()
-viewer.from_ase(trajectory)
-# set a vector field to show the arrow
-viewer.avr.vf.settings = [{"origins": "positions", "vectors": "movement", "radius": 0.1}]
-viewer.avr.vf.show = True
+viewer.from_ase(atoms)
+viewer.avr.phonon_setting = phonon_setting
 viewer
 ```
 

docs/source/isosurface.rst

@@ -17,7 +17,9 @@ Here is an example of drawing isosurfaces for HOMO of H2O molecule.
    viewer = WeasWidget()
    viewer.from_ase(atoms)
    viewer.avr.iso.volumetric_data = {"values": volume}
-   viewer.avr.iso.settings = [{"isovalue": 0.0001, "mode": 0}]
+   viewer.avr.iso.settings = {"positive": {"isovalue": 0.001},
+                              "negative": {"isovalue": -0.001, "color": "yellow"}
+                             }
    viewer
 
 .. figure:: _static/images/example-isosurface.png

package.json

@@ -7,7 +7,7 @@
 	"dependencies": {
 		"dat.gui": "^0.7.9",
 		"three": "^0.161.0",
-		"weas": ">=0.1.32"
+		"weas": ">=0.1.33"
 	},
 	"devDependencies": {
 		"esbuild": "^0.20.0"

src/weas_widget/utils.py

@@ -40,6 +40,7 @@ def to_weas(cls, ase_atoms):
             "positions": positions.tolist(),
             "symbols": symbols,
             "attributes": attributes,
+            "pbc": ase_atoms.get_pbc().tolist(),
         }
         return weas_atoms
 
@@ -56,7 +57,9 @@ def to_ase(cls, weas_atoms):
         symbols = [weas_atoms["species"][s] for s in weas_atoms["symbols"]]
         positions = weas_atoms["positions"]
         cell = np.array(weas_atoms["cell"]).reshape(3, 3)
-        ase_atoms = Atoms(symbols=symbols, positions=positions, cell=cell)
+        ase_atoms = Atoms(
+            symbols=symbols, positions=positions, cell=cell, pbc=weas_atoms["pbc"]
+        )
         return ase_atoms
 
 
@@ -72,8 +75,10 @@ def to_weas(cls, pymatgen_structure):
         # structure is a Molecule, convert it to Structure
         if isinstance(pymatgen_structure, Molecule):
             cell = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
+            pbc = [False, False, False]
         else:
             cell = pymatgen_structure.lattice.matrix.flatten().tolist()
+            pbc = pymatgen_structure.lattice.pbc
         positions = [site.coords.tolist() for site in pymatgen_structure.sites]
         symbols = [site.species_string for site in pymatgen_structure.sites]
         for i in range(len(symbols)):
@@ -91,6 +96,7 @@ def to_weas(cls, pymatgen_structure):
             "positions": positions,
             "symbols": symbols,
             "attributes": attributes,
+            "pbc": pbc,
         }
         return weas_atoms
 
@@ -108,7 +114,7 @@ def to_pymatgen(cls, weas_atoms):
         if np.allclose(cell, np.zeros((3, 3))):
             structure = Molecule(species, sites)
         else:
-            lattice = Lattice(weas_atoms["cell"])
+            lattice = Lattice(weas_atoms["cell"], pbc=weas_atoms["pbc"])
             structure = Structure(lattice, species, sites, coords_are_cartesian=True)
         return structure
 

tests/test_weas_widget.py

@@ -1,12 +1,13 @@
 from weas_widget import WeasWidget
+import numpy as np
 
 
 def test_widget_initialization():
     widget = WeasWidget()
     assert widget is not None
 
 
-def test_set_and_get_atoms_ase(h2o):
+def test_set_and_get_atoms_ase_molecule(h2o):
     from ase.build import molecule
 
     atoms = molecule("H2O")
@@ -16,17 +17,36 @@ def test_set_and_get_atoms_ase(h2o):
     assert retrieved_atoms == atoms
 
 
+def test_set_and_get_atoms_ase_bulk():
+    from ase.build import bulk
+
+    atoms = bulk("Si", "diamond", a=5.43)
+    atoms.pbc = [True, True, False]
+    retrieved_atoms = WeasWidget(from_ase=atoms).to_ase()
+
+    assert retrieved_atoms == atoms
+    assert np.allclose(retrieved_atoms.cell, atoms.cell)
+    assert np.all(retrieved_atoms.pbc == atoms.pbc)
+
+
 def test_set_and_get_pymatgen_structure():
     from pymatgen.core import Structure, Lattice
 
-    structure = Structure.from_spacegroup(
-        "Pm-3m", Lattice.cubic(4.1437), ["Cs", "Cl"], [[0, 0, 0], [0.5, 0.5, 0.5]]
+    # Original structure with PBC in all directions
+    lattice = Lattice.from_parameters(
+        a=5, b=5, c=10, alpha=90, beta=90, gamma=90, pbc=[True, True, False]
     )
+    species = ["Si", "Si"]
+    coords = [[0, 0, 0], [0.5, 0.5, 0.5]]
+    structure = Structure(lattice, species, coords)
+
     viewer = WeasWidget(from_pymatgen=structure)
     retrieved_structure = viewer.to_pymatgen()
 
     # Assert
     assert retrieved_structure == structure
+    assert retrieved_structure.lattice.pbc == structure.lattice.pbc
+    assert np.allclose(retrieved_structure.lattice.matrix, structure.lattice.matrix)
 
 
 def test_set_and_get_pymatgen_molecule():
@@ -38,3 +58,4 @@ def test_set_and_get_pymatgen_molecule():
 
     # Assert
     assert len(retrieved_structure) == len(structure)
+    assert isinstance(retrieved_structure, Molecule)