Load functions via UI (#29)

* add ui feature to load function

* make the box look better

* updates

* updates

* single select

* only update selection if changed

* bugfix

* sorted

* small fix

* update version

* update documentation

* v0.1.3

Author: PythonFZ

README.md

@@ -12,18 +12,38 @@ Install via ``pip install zndraw`` or ``pip install zndraw[webview]`` to open zn
 You can use ZnDraw with the CLI ``zndraw atoms.xyz``.
 For a full list of arguments use `zndraw --help`.
 
-To interface with ``zndraw --update-function zndraw.examples.explode`` you need to be able to import via ``from module import function``.
+ZnDraw is designed to work with your Python scripts.
+To interface you can inherit from `zndraw.examples.UpdateScene` or follow this base class:
 
-The ZnDraw function expects as inputs
+```python
+import abc
+from pydantic import BaseModel
+
+class UpdateScene(BaseModel, abc.ABC):
+    @abc.abstractmethod
+    def run(self, atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
+        pass
+```
+
+The ``run`` method expects as inputs
 - atom_ids: list[int], the ids of the currently selected atoms
 - atoms: ase.Atoms, the configuration as `ase.Atoms` file where atom_ids where selected.
 
 and as an output:
 - list[ase.Atoms], a list of ase Atoms objects to display.
 
+
+You can define the parameters using `pydantic.Field` which will be displayed in the UI.
+
 ```python
-def function(atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]|Generator[ase.Atoms, None, None]:
-    ...
+class MyUpdateCls(UpdateScene):
+    steps: int = Field(100, le=1000, ge=1)
+    x: float = Field(0.5, le=5, ge=0)
+    symbol: str = Field("same")
 ```
 
+To add your method click on the `+` on the right side of the window.
+Your should be able to add your method from the working directory via `module.MyUpdateCls` as long
+as it can be imported via `from module import MyUpdateCls`.
+
 ![Alt text](https://raw.githubusercontent.com/zincware/ZnDraw/main/misc/zndraw_ui.png "ZnDraw UI")

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zndraw"
-version = "0.1.2"
+version = "0.1.3"
 description = ""
 authors = ["zincwarecode <[email protected]>"]
 license = "Apache-2.0"
@@ -15,6 +15,7 @@ flask = "^2.2.3"
 tqdm = "^4.65.0"
 pywebview = {version = "^4.0.2", optional = true}
 znh5md = "^0.1.6"
+pydantic = "^1.10.7"
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.3.0"

pyproject.toml

@@ -58,14 +58,49 @@ def bonds_step():
 @app.route("/select", methods=["POST"])
 def select() -> list[int]:
     """Update the selected atoms."""
-    return (
-        request.json
-    )  # + [x + 1 for x in request.json] + [x - 1 for x in request.json]
+    step = request.json["step"]
+    selected_ids = request.json["selected_ids"]
+    return {"selected_ids": selected_ids, "updated": False}
+
+    # atoms = globals.config.get_atoms(step)
+
+    # for id in tuple(selected_ids):
+    #     selected_symbol = atoms[id].symbol
+    #     selected_ids += [
+    #         idx for idx, atom in enumerate(atoms) if atom.symbol == selected_symbol
+    #     ]
+
+    # return {"selected_ids": list(set(selected_ids)), "updated": True}
+
+
+@app.route("/add_update_function", methods=["POST"])
+def add_update_function():
+    """Add a function to the config."""
+    globals.config.update_function = request.json
+    try:
+        signature = globals.config.get_update_signature()
+    except (ImportError, ValueError) as err:
+        return {"error": str(err)}
+    return signature
+
+
+@app.route("/update_function_values", methods=["POST"])
+def update_function_values():
+    """Update the values of the update function."""
+    globals.config.set_update_function_parameters(request.json)
+    return {}
+
+
+@app.route("/select_update_function/<name>")
+def select_update_function(name):
+    """Select a function from the config."""
+    globals.config.update_function_name = name
+    return {}
 
 
 @app.route("/update", methods=["POST"])
 def update_scene():
-    selected_ids = request.json["selected_ids"]
+    selected_ids = list(sorted(request.json["selected_ids"]))
     step = request.json["step"]
 
     function = globals.config.get_update_function()

zndraw/app.py

@@ -1,16 +1,74 @@
+import abc
+
 import ase
 import numpy as np
+from pydantic import BaseModel, Field
+
+
+class UpdateScene(BaseModel, abc.ABC):
+    @abc.abstractmethod
+    def run(self, atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
+        pass
+
+
+class Explode(UpdateScene):
+    steps: int = Field(100, le=1000, ge=1)
+    particles: int = Field(10, le=100, ge=1)
+
+    def run(self, atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
+        particles = []
+        for _atom_id in atom_ids:
+            for _ in range(self.particles):
+                particles.append(ase.Atoms("Na", positions=[atoms.positions[_atom_id]]))
+
+        for _ in range(self.steps):
+            struct = atoms.copy()
+            for particle in particles:
+                particle.positions += np.random.normal(scale=0.1, size=(1, 3))
+                struct += particle
+            yield struct
+
+
+class Delete(UpdateScene):
+    def run(self, atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
+        for idx, atom_id in enumerate(sorted(atom_ids)):
+            atoms.pop(atom_id - idx)  # we remove the atom and shift the index
+        return [atoms]
+
+
+class Move(UpdateScene):
+    x: float = Field(0.5, le=5, ge=0)
+    y: float = Field(0.5, le=5, ge=0)
+    z: float = Field(0.5, le=5, ge=0)
+
+    def run(self, atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
+        for atom_id in atom_ids:
+            atom = atoms[atom_id]
+            atom.position += np.array([self.x, self.y, self.z])
+            atoms += atom
+        return [atoms]
+
+
+class Duplicate(UpdateScene):
+    x: float = Field(0.5, le=5, ge=0)
+    y: float = Field(0.5, le=5, ge=0)
+    z: float = Field(0.5, le=5, ge=0)
+    symbol: str = Field("same")
+
+    def run(self, atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
+        for atom_id in atom_ids:
+            atom = ase.Atom(atoms[atom_id].symbol, atoms[atom_id].position)
+            atom.position += np.array([self.x, self.y, self.z])
+            atom.symbol = self.symbol if self.symbol != "same" else atom.symbol
+            atoms += atom
+        return [atoms]
 
 
-def explode(atom_id: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
-    particles = []
-    for _atom_id in atom_id:
-        for _ in range(5):
-            particles.append(ase.Atoms("Na", positions=[atoms.positions[_atom_id]]))
+class ChangeType(UpdateScene):
+    symbol: str = Field("")
 
-    for _ in range(102):
-        struct = atoms.copy()
-        for particle in particles:
-            particle.positions += np.random.normal(scale=0.1, size=(1, 3))
-            struct += particle
-        yield struct
+    def run(self, atom_ids: list[int], atoms: ase.Atoms) -> list[ase.Atoms]:
+        for atom_id in atom_ids:
+            atoms[atom_id].symbol = self.symbol
+        print(atoms)
+        return [atoms]

zndraw/examples/__init__.py

@@ -3,6 +3,7 @@
 import pathlib
 
 import ase.io
+import pydantic
 import tqdm
 import znh5md
 
@@ -20,12 +21,52 @@ class Config:
     resolution: int = 5
     repeat: tuple = (1, 1, 1)
 
+    _update_function_name: str = None
+
+    @property
+    def update_function_name(self):
+        if self._update_function_name is not None:
+            return self._update_function_name
+        return self.update_function.rsplit(".", 1)[1]
+
+    @update_function_name.setter
+    def update_function_name(self, value):
+        self._update_function_name = value
+
+    def get_update_signature(self):
+        module_name, function_name = self.update_function.rsplit(".", 1)
+        if function_name in _update_functions:
+            return _update_functions[self.update_function_name].schema()
+        module = importlib.import_module(module_name)
+        instance: pydantic.BaseModel = getattr(module, function_name)()
+        _update_functions[function_name] = instance
+        return _update_functions[self.update_function_name].schema()
+
     def get_update_function(self):
+        module_name, function_name = self.update_function.rsplit(".", 1)
         if self.update_function is None:
             return None
-        module_name, function_name = self.update_function.rsplit(".", 1)
+        if function_name in _update_functions:
+            return _update_functions[self.update_function_name].run
+
         module = importlib.import_module(module_name)
-        return getattr(module, function_name)
+        _update_functions[self.update_function_name] = getattr(module, function_name)()
+        return _update_functions[self.update_function_name].run
+
+    def set_update_function_parameters(self, value):
+        instance = _update_functions[value["function_id"]]
+        attribute = value["property"].lower()
+        value = value["value"]
+        if instance.__annotations__[attribute] == float:
+            value = float(value)
+        elif instance.__annotations__[attribute] == int:
+            value = int(value)
+        elif instance.__annotations__[attribute] == bool:
+            value = bool(value)
+        else:
+            value = value
+        print(f"Setting {attribute} of {instance} to {value}")
+        setattr(instance, attribute, value)
 
     def load_atoms(self, item=None):
         if item == 0:
@@ -58,6 +99,8 @@ def get_atoms(self, step=0) -> ase.Atoms:
 
 # TODO set defaults here and load in typer?
 
+_update_functions = {}
+
 _atoms_cache: dict = {}
 config = Config()
 

zndraw/globals.py

@@ -35,7 +35,7 @@ canvas {
     display: block;
 }
 
-#help_btn {
+#right_btn_group {
     position: absolute;
     top: 5px;
     right: 10px;
@@ -45,6 +45,12 @@ canvas {
     opacity: 0.5;
 }
 
+#add_class {
+    z-index: 2000;
+    display: none;
+
+}
+
 .atom-spinner,
 .atom-spinner * {
     box-sizing: border-box;