Add atom legend

docs/source/gui.rst

@@ -137,3 +137,11 @@ Set viewer width and height
    from weas_widget import WeasWidget
    viewer = WeasWidget(viewerStyle = {"width": "800px", "height": "600px"})
    viewer
+
+Show atoms legend
+------------------
+
+.. code-block:: python
+
+   from weas_widget import WeasWidget
+   viewer.avr.show_atoms_legend = True

docs/source/quick_start.ipynb

@@ -140,7 +140,7 @@
    "outputs": [],
    "source": [
     "# only draw stick (bond) for first four atoms\n",
-    "viewer.model_sticks = [1, 1, 1, 1, 0, 0, 0, 0]\n",
+    "viewer.avr.model_sticks = [1, 1, 1, 1, 0, 0, 0, 0]\n",
     "viewer.avr.draw()"
    ]
   },
@@ -308,9 +308,8 @@
     "viewer = WeasWidget()\n",
     "viewer.from_ase(atoms)\n",
     "viewer.avr.iso.volumetric_data = {\"values\": volume}\n",
-    "# mode 0: plot both positive and negative isosurface for the isovalue\n",
-    "# mode !=0, only plot one isosurface\n",
-    "viewer.avr.iso.settings = [{\"isovalue\": 0.0001, \"mode\": 0}]\n",
+    "viewer.avr.iso.settings = {\"positive\": {\"isovalue\": 0.001},\n",
+    "                           \"negative\": {\"isovalue\": -0.001, \"color\": \"yellow\"},}\n",
     "viewer"
    ]
   },
@@ -363,7 +362,7 @@
     "origins = atoms.positions\n",
     "# the vertor\n",
     "vectors = [[0, 0, 1]]*len(atoms)\n",
-    "viewer.avr.vf.settings = [{\"origins\": origins, \"vectors\": vectors, \"color\": \"red\"}]"
+    "viewer.avr.vf.settings[\"test\"]= {\"origins\": origins, \"vectors\": vectors, \"color\": \"red\"}"
    ]
   },
   {
@@ -387,15 +386,22 @@
     "import numpy as np\n",
     "from ase.build import bulk\n",
     "from weas_widget import WeasWidget\n",
-    "from weas_widget.utils import generate_phonon_trajectory\n",
-    "\n",
     "atoms = bulk(\"Fe\", cubic=True)\n",
-    "eigenvector = np.array([[0, -0.0, 0.5], [0, 0.0, -0.5]])\n",
-    "trajectory = generate_phonon_trajectory(atoms, eigenvector, repeat=[4, 4, 1])\n",
+    "phonon_setting = {\"eigenvectors\": np.array([[[0, 0], [0, 0],[0.5, 0]],\n",
+    "                                    [[0, 0], [0, 0], [-0.5, 0]]]\n",
+    "                                   ),\n",
+    "        \"kpoint\": [0, 0, 0], # optional\n",
+    "        \"amplitude\": 5, # scale the motion of the atoms\n",
+    "        \"factor\": 1.5, # scale the length of the arrows\n",
+    "        \"nframes\": 20,\n",
+    "        \"repeat\": [4, 4, 1],\n",
+    "        \"color\": \"blue\",\n",
+    "        \"radius\": 0.1,\n",
+    "        }\n",
     "viewer = WeasWidget()\n",
-    "viewer.from_ase(trajectory)\n",
-    "# set a vector field to show the arrow\n",
-    "viewer.avr.vf.settings = [{\"origins\": \"positions\", \"vectors\": \"movement\", \"radius\": 0.1}]\n",
+    "viewer.from_ase(atoms)\n",
+    "viewer.avr.phonon_setting = phonon_setting\n",
+    "viewer.camera.setting = {\"direction\": [0, -1, 0]}\n",
     "viewer"
    ]
   },

js/widget.js

@@ -35,6 +35,14 @@ function render({ model, el }) {
         }
         // console.log("atoms: ", atoms);
         const guiConfig = model.get("guiConfig");
+        if (guiConfig.legend) {
+            guiConfig.legend.enabled = model.get("showAtomLegend");
+        } else {
+            guiConfig.legend = {enabled: model.get("showAtomLegend"),
+                position: "bottom-right",
+            };
+        }
+
         const viewerConfig = {
              logLevel: model.get("logLevel"),
             _modelStyle: model.get("modelStyle"),
@@ -232,6 +240,11 @@ function render({ model, el }) {
         editor.tjs.controls.update();
         editor.tjs.render();
     });
+    // frame
+    model.on("change:showAtomLegend", () => {
+        editor.avr.guiManager.guiConfig.legend.enabled = model.get("showAtomLegend");
+        editor.avr.guiManager.updateLegend();
+    });
 }
 function createVolumeData(data, cell=[[1, 0, 0], [0, 1, 0], [0, 0, 1]]) {
     // get the dimensions

package.json

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

src/weas_widget/atoms_viewer.py

@@ -21,6 +21,7 @@ class AtomsViewer(WidgetWrapper):
         "show_cell": "showCell",
         "show_bonded_atoms": "showBondedAtoms",
         "hide_long_bonds": "hideLongBonds",
+        "show_atom_legend": "showAtomLegend",
         "show_hydrogen_bonds": "showHydrogenBonds",
         "atom_label_type": "atomLabelType",
         "material_type": "materialType",

src/weas_widget/base_widget.py

@@ -69,6 +69,8 @@ class BaseWidget(anywidget.AnyWidget):
     phonon = tl.Dict({}).tag(sync=True)
     # highlight
     highlightSettings = tl.Dict({}).tag(sync=True)
+    #
+    showAtomLegend = tl.Bool(False).tag(sync=True)
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)

src/weas_widget/plugins/vector_field.py

@@ -1,20 +1,29 @@
-from ..base_class import WidgetWrapper
+from ..base_class import WidgetWrapper, ChangeTrackingDict
 
 
 class VectorField(WidgetWrapper):
 
     catalog = "vector_field"
 
     _attribute_map = {
-        "settings": "vectorField",
         "show": "showVectorField",
     }
 
-    _extra_allowed_attrs = []
+    _extra_allowed_attrs = ["_settings", "settings"]
 
     def __init__(self, _widget):
         super().__init__(_widget)
 
+    @property
+    def settings(self):
+        return self._settings
+
+    @settings.setter
+    def settings(self, value):
+        self._settings = ChangeTrackingDict(
+            value, widget=self._widget, key="vectorField"
+        )
+
     def update_atoms(self):
         self.settings = self.set_moment()
 

tests/notebooks/tests/notebooks/widgets.ipynb

@@ -167,6 +167,7 @@
     "del viewer.avr.bond.settings['S-O']\n",
     "viewer.avr.bond.settings['C-H'][\"max\"]=1.2\n",
     "viewer.avr.bond.settings['C-H'][\"color1\"]=\"red\"\n",
+    "viewer.avr.show_atom_legend = True\n",
     "viewer"
    ]
   },