Skip to main content

A widget to visualize and interact with atomistic structures in Jupyter Notebook.

Project description

Welcome to WEAS Widget!

PyPI version Docs status Unit test

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: Colab

🌟 Gallery: projects using weas-widget

📢 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.

Operation Shortcut
Move g
Rotate r
Duplicate d

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

License

MIT

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

weas_widget-0.2.6.tar.gz (329.2 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

weas_widget-0.2.6-py3-none-any.whl (345.3 kB view details)

Uploaded Python 3

File details

Details for the file weas_widget-0.2.6.tar.gz.

File metadata

  • Download URL: weas_widget-0.2.6.tar.gz
  • Upload date:
  • Size: 329.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.9.25

File hashes

Hashes for weas_widget-0.2.6.tar.gz
Algorithm Hash digest
SHA256 9f4524a57ffbbf6060eaa1d22f9c6072155c219f3c81a39da36d9c7ac9dff28b
MD5 dcbed6dd44be83e27444c00e0d05d1a3
BLAKE2b-256 352339fb0478928f4e7f78ec9e0eb1c8b71a11d7a2fafde9831531c404af128d

See more details on using hashes here.

File details

Details for the file weas_widget-0.2.6-py3-none-any.whl.

File metadata

  • Download URL: weas_widget-0.2.6-py3-none-any.whl
  • Upload date:
  • Size: 345.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.9.25

File hashes

Hashes for weas_widget-0.2.6-py3-none-any.whl
Algorithm Hash digest
SHA256 06bec3bfca85f27ae9089da1ff6ab33a97e9604ee70eb009f5ea5b8a2c38bcf4
MD5 8632e9ef6d00fa49f1e61b46e78b01bd
BLAKE2b-256 03ce089287b3a09f1c552d54e27cc07211f2e1fd5301e89e321cabcdc9296d2d

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page