Integrate stk as a structure provider, and generate shapes for custom bonds

.github/workflows/website.yml

@@ -21,7 +21,7 @@ jobs:
           cache: 'npm'
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
       - name: configure access to git repositories in package-lock.json
         run: git config --global url."https://github.com/".insteadOf ssh://[email protected]
       - name: install npm dependencies

docs/requirements.txt

@@ -8,6 +8,12 @@ numpy
 ipywidgets
 matplotlib
 ase==3.22.1
+# These three pins are to avoid rdkit version 2024.3.5. Once a later version
+# is available, the latest version (assuming Python >= 3.11) of all three
+# should be ok.
+rdkit==2023.9.5
+stk==2024.9.23.1
+stko==2024.8.29.1
 
 # required for the examples of chemiscope.explore
 mace-torch

docs/src/python/reference.rst

@@ -28,3 +28,5 @@
 .. autofunction:: chemiscope.explore
 
 .. autofunction:: chemiscope.metatensor_featurizer
+
+.. autofunction:: chemiscope.convert_stk_bonds_as_shapes

python/README.md

@@ -24,6 +24,8 @@ import chemiscope
 import ase.io
 
 # read frames using ase
+# frames can also be stk objets, e.g.
+# frames = [stk.BuildingBlock(smiles="NCCN")]
 frames = ase.io.read("structures.xyz", ":")
 
 # add additional properties to display
@@ -44,6 +46,8 @@ import chemiscope
 import ase.io
 
 # read frames using ase
+# frames can also be stk objets, e.g.
+# frames = [stk.BuildingBlock(smiles="NCCN")]
 frames = ase.io.read("structures.xyz", ":")
 
 # add additional properties to display

python/chemiscope/__init__.py

@@ -10,6 +10,7 @@
     ellipsoid_from_tensor,
     extract_properties,
     librascal_atomic_environments,
+    convert_stk_bonds_as_shapes,
 )
 from .explore import explore, metatensor_featurizer  # noqa: F401
 from .version import __version__  # noqa: F401

python/chemiscope/input.py

@@ -2,6 +2,7 @@
 """
 Generate JSON input files for the default chemiscope visualizer.
 """
+
 import gzip
 import json
 import os
@@ -282,6 +283,7 @@ def create_input(
         data["structures"] = frames_to_json(frames)
         n_structures = len(data["structures"])
         n_atoms = sum(s["size"] for s in data["structures"])
+
     else:
         n_atoms = 0
 
@@ -378,6 +380,7 @@ def create_input(
     # should be removed in version 0.6 of chemiscope.
     if frames is not None:
         found_one_from_frame = False
+
         atom_properties = _list_atom_properties(frames)
         for name in atom_properties:
             if name in data["properties"]:

python/chemiscope/jupyter.py

@@ -155,7 +155,7 @@ def show_input(path, mode="default"):
     try:
         meta = dict_input["meta"]
 
-        if meta is {"name": " "}:
+        if meta == {"name": " "}:
             has_metadata = False
         else:
             has_metadata = True

python/chemiscope/structures/__init__.py

@@ -20,6 +20,15 @@
     ase_tensors_to_ellipsoids,
     ase_vectors_to_arrows,
 )
+from ._stk import (  # noqa: F401
+    _stk_valid_structures,
+    _stk_to_json,
+    convert_stk_bonds_as_shapes,
+    _stk_all_atomic_environments,
+    _stk_composition_properties,
+    _stk_list_atom_properties,
+    _stk_list_structure_properties,
+)
 
 
 def _guess_adapter(frames):
@@ -32,6 +41,10 @@ def _guess_adapter(frames):
     if use_ase:
         return ase_frames, "ASE"
 
+    stk_frames, use_stk = _stk_valid_structures(frames)
+    if use_stk:
+        return stk_frames, "stk"
+
     raise Exception(f"unknown frame type: '{frames[0].__class__.__name__}'")
 
 
@@ -48,6 +61,8 @@ def frames_to_json(frames):
 
     if adapter == "ASE":
         return [_ase_to_json(frame) for frame in frames]
+    elif adapter == "stk":
+        return [_stk_to_json(frame) for frame in frames]
     else:
         raise Exception("reached unreachable code")
 
@@ -62,6 +77,9 @@ def _list_atom_properties(frames):
 
     if adapter == "ASE":
         return _ase_list_atom_properties(frames)
+    elif adapter == "stk":
+        return _stk_list_atom_properties(frames)
+
     else:
         raise Exception("reached unreachable code")
 
@@ -76,6 +94,8 @@ def _list_structure_properties(frames):
 
     if adapter == "ASE":
         return _ase_list_structure_properties(frames)
+    elif adapter == "stk":
+        return _stk_list_structure_properties(frames)
     else:
         raise Exception("reached unreachable code")
 
@@ -96,6 +116,12 @@ def extract_properties(frames, only=None, environments=None):
 
     if adapter == "ASE":
         return _ase_extract_properties(frames, only, environments)
+
+    elif adapter == "stk":
+        raise RuntimeError(
+            "stk molecules do not contain properties, you must manually provide them"
+        )
+
     else:
         raise Exception("reached unreachable code")
 
@@ -119,6 +145,10 @@ def composition_properties(frames, environments=None):
 
     if adapter == "ASE":
         return _ase_composition_properties(frames, environments)
+
+    elif adapter == "stk":
+        return _stk_composition_properties(frames, environments)
+
     else:
         raise Exception("reached unreachable code")
 
@@ -137,6 +167,8 @@ def all_atomic_environments(frames, cutoff=3.5):
 
     if adapter == "ASE":
         return _ase_all_atomic_environments(frames, cutoff)
+    elif adapter == "stk":
+        return _stk_all_atomic_environments(frames, cutoff)
     else:
         raise Exception("reached unreachable code")
 

python/chemiscope/structures/_ase.py

@@ -15,7 +15,11 @@
 
 
 def _ase_valid_structures(frames):
-    frames_list = list(frames)
+    try:
+        frames_list = list(frames)
+    except TypeError:
+        return [], False
+
     if HAVE_ASE and isinstance(frames_list[0], ase.Atoms):
         for frame in frames_list:
             assert isinstance(frame, ase.Atoms)