Welcome to WEAS Widget!

A widget to visualize and edit atomic structures in Jupyter Notebooks. It uses WEAS (Web Environment For Atomistic Structure) in the backend. Interact with the widget using LLM agents (e.g., LangChain) to guide edits and analysis from your notebook.

Features:

• Model: space-filling, ball-stick, polyhedral.
• Supports importing data from ASE and Pymatgen.
• Edit structure: move, rotate, delete and replace atoms.
• Supports periodic boundary conditions
• Animations
• Isosurfaces
• Vector fields, e.g., magnetic moments, phonons, ...
• Fermi surface
• Lattice planes
• Integrate with LLMs via LangChain Agent

Demo

Please try the widget in the following links:

🌟 Gallery: projects using weas-widget

• AiiDAlab Plugins:
  • aiidalab-qe-vibroscopy
  • aiidalab-qe-pp
  • aiidalab-qe-wannier90
• pack-mm
• ipyvasp

📢 Want to Feature Your Project?

If your project uses weas-widget, we’d love to showcase it here!

• Submit a pull request (PR) with a link to your project.
• Need a project removed? Open a PR or issue, and we’ll take care of it.

🚀 Let's build a growing community of weas-widget users!

Installation

With pip:

pip install weas-widget

To install the latest version from source, first clone the repository and then install using pip:

git clone https://github.com/superstar54/weas-widget
cd weas-widget
npm install
npm run build
pip install -e .

How to use

from ase.build import molecule
from weas_widget import WeasWidget
atoms = molecule("C2H6SO")
viewer = WeasWidget()
viewer.from_ase(atoms)
viewer

Full documentation at: https://weas-widget.readthedocs.io/en/latest/index.html

Reporting issues

If you encounter any problems, please first update the widget to the latest version.

    pip install weas-widget  --upgrade

If the problem persists, please open a GitHub issue

Features

Select Atoms

• Pick Selection: Click directly on an atom to select it.
• Box Selection: Hold the Shift key and drag to select a group of atoms.
• Lasso Selection: Hold Shift + Alt and drag to select a group of atoms.

Shortcuts for editing selected atoms

Press the keyboard shortcut, and move your mouse.

Rotate selected atoms

To rotate around a custom axis, press r then a and click one or two atoms, then press a again to exit axis picking. Rotation defaults to the camera axis through the selection center. One atom sets the rotation center (camera axis), two atoms define the bond axis. The axis is shown with orange crosses and a long orange line (for two atoms). Press a again to exit axis picking and rotate; click an axis atom again to deselect it. Press r then x, y, or z to lock rotation to a world axis (press the same key again to unlock).

Translate along one axis

To translate along one axis, press g then x, y, or z to lock movement. To translate along an atom-defined axis, press g then a and click two atoms, then press a again to exit axis picking. To translate in an atom-defined plane, press g then a and click three atoms, press a again to exit axis picking, then press p for plane or n for normal.

Delete selected atoms

Press the Delete key

Export

• Export the modified atomic structure to ASE or Pymatgen

atoms = viewer.to_ase()

• Save image to a path by:

viewer.save_image("/home/xing/filename.png")

• Download image by:

viewer.download_image("filename.png")

Different styles for the atoms.

Visualizing crystal structures

For a nice visualization of a crystal, show

• unit cell
• bonded atoms outside the cell
• polyhedra

from weas_widget import WeasWidget
viewer1 = WeasWidget()
viewer1.load_example("tio2.cif")
viewer1.avr.model_style = 2
viewer1.avr.boundary = [[-0.1, 1.1], [-0.1, 1.1], [-0.1, 1.1]]
viewer1.avr.show_bonded_atoms = True
viewer1.avr.color_type = "VESTA"
viewer1

Isosurfaces

from ase.build import molecule
from weas_widget import WeasWidget
from ase.io.cube import read_cube_data
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 = {"positive": {"isovalue": 0.001},
                           "negative": {"isovalue": -0.001, "color": "yellow"}
                           }
viewer

Fermi surface

Requires the optional fermi-surface dependencies (includes seekpath):

pip install "weas-widget[fermi-surface]"

from weas_widget import WeasWidget
viewer = WeasWidget()
viewer.add_fermi_surface_from_bxsf("copper.bxsf", clip_bz=True)
viewer

Magnetic moments

Show the magnetic moments as a vector field.

from ase.build import bulk
from weas_widget import WeasWidget
import numpy as np
atoms = bulk("Fe", cubic=True)
atoms*=[2, 2, 1]
atoms.set_array("moment", np.ones(len(atoms)))
viewer = WeasWidget()
viewer.from_ase(atoms)
viewer.avr.model_style = 1
viewer

Phonons

Animate vibrational (phonon) modes (computed with external software).

import numpy as np
from ase.build import bulk
from weas_widget import WeasWidget
atoms = bulk("Fe", cubic=True)
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(atoms)
viewer.avr.phonon_setting = phonon_setting
viewer

Lattice plane

Draw a plane that is defined by the miller indices and distance from the origin or by selecting the atoms.

viewer.avr.lp.add_plane_from_indices(name = "111",
                                     indices = [1, 1, 1],
                                     distance = 4,
                                     scale = 1.0,
                                     color = [0, 1, 1, 0.5])
viewer.avr.lp.build_plane()

Slice 2D

from ase.build import molecule
from weas_widget import WeasWidget
from ase.io.cube import read_cube_data
volume, atoms = read_cube_data("h2o-homo.cube")
viewer = WeasWidget()
viewer.from_ase(atoms)
viewer.avr.model_style = 1
viewer.avr.volume_slice.volumetric_data = {"values": volume}
viewer.avr.volume_slice.settings = {"Slice 1": {"h": 0, "k": 1, "l": 0, "distance": 5.5, "samplingDistance": 0.1 },
                                    "Slice 2": {"h": 1, "k": 1, "l": 0, "distance": 5.5, "samplingDistance": 0.1 },
                                   }
viewer.camera.setting = {"direction": [0.5, 1, 2], "zoom": 1.5}
viewer

Save and restore state

state = viewer.export_state()
viewer.save_state("snapshot.json")

# later
viewer.load_state("snapshot.json")

# or create a viewer directly from a saved state
viewer = WeasWidget.from_state_file("snapshot.json")

Integrate with LangChain Agent

WeasToolkit provides agent-ready tools for inspection and editing (style controls, selection, structure loading, atom edits, camera, measurements, meshes, and more). You can also extend it with your own tools.

from weas_widget import WeasWidget, WeasToolkit
from langchain_openai import ChatOpenAI
from langchain_widget import LangChainWidget
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import ipywidgets as ipw

load_dotenv()

viewer = WeasWidget()

chat_model = ChatOpenAI(model="gpt-4o-mini")

chat = LangChainWidget(
    chat_model=chat_model,
    tools=WeasToolkit(viewer=viewer),
    title="WEAS Agent Chat",
    system_prompt=(
        "You are a scientific assistant. "
        "Use the available tools to inspect and manipulate the 3D structure."
    ),
)
ipw.VBox([viewer, chat])

See the agent tools reference in the docs: https://weas-widget.readthedocs.io/en/latest/agent_tools.html

Test

Unit test

pytest

End-to-end test

The e2e test is similar to ipywidgets.

For the first time, one needs to install the dependence.

cd tests/notebooks/
yarn install

Then run in a terminal:

yarn start

In another terminal:

yarn test

If the snapshots need to be updated:

yarn test:update

Contact

• Xing Wang xingwang1991@gmail.com

License

MIT