From 5d6642b2ea922aee17ac686b854a6a0d77efb08a Mon Sep 17 00:00:00 2001 From: Min-Hsueh Chiu Date: Mon, 20 Jul 2026 16:22:27 -0700 Subject: [PATCH 1/2] add PhononBandstructureAndDosComponent to jupyter notebook rendering --- crystal_toolkit/components/phonon.py | 177 +++++++++++++++++---------- crystal_toolkit/core/jupyter.py | 25 +++- 2 files changed, 136 insertions(+), 66 deletions(-) diff --git a/crystal_toolkit/components/phonon.py b/crystal_toolkit/components/phonon.py index c9ee6f99..a7c20d22 100644 --- a/crystal_toolkit/components/phonon.py +++ b/crystal_toolkit/components/phonon.py @@ -47,8 +47,15 @@ MIN_MAGNITUDE = 0 MAX_SUPERCELL_SITES = 500 -DEFAULTS: dict[str, str | bool] = { +DEFAULTS: dict[str, str | bool | int] = { "color_scheme": "VESTA", + "scale_x": 1, + "scale_y": 1, + "scale_z": 1, + "band_num": 0, + "qpoint": 0, + "magnitude_fraction": 1, + "velocity": 1.0, } @@ -83,12 +90,20 @@ def __init__( **kwargs, ) + if bandstructure_symm_line and density_of_states: + # initialize this for jupyter notebook rendering + self.create_store("ph_bs", bandstructure_symm_line) + self.create_store("ph_dos", density_of_states) + @property def _sub_layouts(self) -> dict[str, Component]: # defaults state = {"label-select": "sc", "dos-select": "ap"} - fig = PhononBandstructureAndDosComponent.get_figure(None, None) + ph_bs = self._initial_data.get("ph_bs", None) + ph_dos = self._initial_data.get("ph_dos", None) + + fig = PhononBandstructureAndDosComponent.get_figure(ph_bs, ph_dos) # Main plot graph = html.Div( [ @@ -103,7 +118,9 @@ def _sub_layouts(self) -> dict[str, Component]: ) # Brillouin zone - zone_scene = self.get_brillouin_zone_scene(None) + zone_scene = self.get_brillouin_zone_scene( + self._initial_data.get("ph_bs", None) + ) zone = CrystalToolkitScene( data=zone_scene.to_json(), sceneSize="500px", id=self.id("zone") ) @@ -217,13 +234,33 @@ def _get_animation_panel(self): ) # crystal visualization + tf_data = None + if ph_bs := self._initial_data.get("ph_bs", None): + # if there is _initial_data + bs = PhononBS.from_pmg(ph_bs) + json_data, _, _ = self._generate_structure_scene( + bs, + DEFAULTS["color_scheme"], + DEFAULTS["scale_x"], + DEFAULTS["scale_y"], + DEFAULTS["scale_z"], + ) + tf_data = PhononBandstructureAndDosComponent._get_time_function_json( + ph_bs=bs, + json_data=json_data, + band=DEFAULTS["band_num"], + qpoint=DEFAULTS["qpoint"], + magnitude=DEFAULTS["magnitude_fraction"], + velocity=DEFAULTS["velocity"], + ) + crystal_animation = html.Div( # CrystalToolkitAnimationScene( PhononAnimationScene( - data={"app": "phonon"}, + data=tf_data if tf_data else {"app": "phonon"}, sceneSize="400px", id=self.id("crystal-animation"), - settings={"defaultZoom": 1.2}, + settings={"defaultZoom": 1.2, "extractAxis": True}, axisView="SW", showControls=False, # disable download for now ), @@ -429,7 +466,7 @@ def _get_animation_panel(self): ), ] - def layout(self) -> html.Div: + def layout(self, jupyter=False) -> html.Div: sub_layouts = self._sub_layouts graph = Columns([Column([sub_layouts["graph"]])]) hints = Columns([Column([sub_layouts["hints"]])]) @@ -457,6 +494,16 @@ def layout(self) -> html.Div: ] ) + if jupyter: + return html.Div( + [ + graph, + hints, + crystal_animation_button_container, + crystal_animation_container, + ] + ) + return html.Div( [ graph, @@ -495,9 +542,7 @@ def _get_time_function_json( json_data: dict, band: int = 0, qpoint: int = 0, - precision: int = 15, magnitude: int = MAX_MAGNITUDE / 2, - total_repeat_cell_cnt: int = 1, velocity: float = 1.0, ) -> dict: if not ph_bs or not json_data: @@ -586,7 +631,7 @@ def _get_time_function_json( rdata["phases"] = phases[qpoint].tolist() # amplitude (A) - rdata["amplitude"] = 1 / np.linalg.norm( + rdata["amplitude"] = magnitude / np.linalg.norm( ph_bs.eigendisplacements[0][0] ) # magnitude @@ -846,7 +891,6 @@ def get_figure( "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)", } - return go.Figure(layout=empty_plot_style) if freq_range[0] is None: @@ -1012,6 +1056,57 @@ def get_font_color(hex_code): className="buttons", ) + def _generate_structure_scene( + self, bs, color_scheme, scale_x=1, scale_y=1, scale_z=1 + ): + struct = bs.structure + total_repeat_cell_cnt = 1 + + # + num_sites = struct.num_sites + + # update structure if the controls got triggered + total_repeat_cell_cnt = scale_x * scale_y * scale_z + + # create supercell + trans = SupercellTransformation( + ((scale_x, 0, 0), (0, scale_y, 0), (0, 0, scale_z)) + ) + struct = trans.apply_transformation(struct) + + struc_graph = StructureGraph.from_local_env_strategy(struct, CrystalNN()) + + # legend + legend = Legend( + struc_graph.structure, + color_scheme=color_scheme, + # radius_scheme=radius_strategy, + cmap_range=None, + ) + self._legend = legend + legend_layout = html.Div(self._make_legend(legend.get_legend())) + + # scene + scene = struc_graph.get_scene( + draw_image_atoms=False, + bonded_sites_outside_unit_cell=False, + site_get_scene_kwargs={ + "retain_atom_idx": True, + "total_repeat_cell_cnt": total_repeat_cell_cnt, + }, + legend=legend, + ) + + # axis + axes = struct.lattice._axes_from_lattice() + axes.visible = True + scene.contents.append(axes) + + # + json_data = scene.to_json() + + return json_data, legend_layout, num_sites + def generate_callbacks(self, app, cache) -> None: @app.callback( Output(self.id("ph-bsdos-graph"), "figure", allow_duplicate=True), @@ -1022,9 +1117,10 @@ def generate_callbacks(self, app, cache) -> None: ), Input(self.id("ph_bs"), "data"), Input(self.id("ph_dos"), "data"), - # prevent_intial_call=True, + prevent_initial_call=True, ) def update_graph(bs, dos): + # this is triggered from the web when updating ph_bs and ph_dos if isinstance(bs, dict): # bs = PhononBS.from_pmg(bs) bs = PhononBandStructureSymmLine.from_dict(bs) @@ -1047,7 +1143,7 @@ def update_graph(bs, dos): State(self.id("ph-bsdos-graph"), "figure"), Input(self.id("ph-bsdos-graph"), "clickData"), Input(self.id("animation-button"), "n_clicks"), - prevent_intial_call=True, + prevent_initial_call=True, ) def update_pointer_graph(figure, nclick, animation_click): if not animation_click: @@ -1105,7 +1201,7 @@ def update_pointer_graph(figure, nclick, animation_click): self.id("animation-button-container"), "style", allow_duplicate=True ), Input(self.id("animation-button"), "n_clicks"), - prevent_intial_call=True, + prevent_initial_call=True, ) def create_animation(nclick): if not nclick: @@ -1160,51 +1256,9 @@ def update_crystal_animation( bs = PhononBS.from_pmg(bs) # bs = PhononBandStructureSymmLine.from_dict(bs) - struct = bs.structure - total_repeat_cell_cnt = 1 - - # - num_sites = struct.num_sites - - # update structure if the controls got triggered - total_repeat_cell_cnt = scale_x * scale_y * scale_z - - # create supercell - trans = SupercellTransformation( - ((scale_x, 0, 0), (0, scale_y, 0), (0, 0, scale_z)) - ) - struct = trans.apply_transformation(struct) - - struc_graph = StructureGraph.from_local_env_strategy(struct, CrystalNN()) - - # legend - legend = Legend( - struc_graph.structure, - color_scheme=color_scheme, - # radius_scheme=radius_strategy, - cmap_range=None, + json_data, legend_layout, num_sites = self._generate_structure_scene( + bs, color_scheme, scale_x, scale_y, scale_z ) - self._legend = legend - legend_layout = html.Div(self._make_legend(legend.get_legend())) - - # scene - scene = struc_graph.get_scene( - draw_image_atoms=False, - bonded_sites_outside_unit_cell=False, - site_get_scene_kwargs={ - "retain_atom_idx": True, - "total_repeat_cell_cnt": total_repeat_cell_cnt, - }, - legend=legend, - ) - - # axis - axes = struct.lattice._axes_from_lattice() - axes.visible = True - scene.contents.append(axes) - - # - json_data = scene.to_json() qpoint = 0 band_num = 0 @@ -1216,9 +1270,9 @@ def update_crystal_animation( raise ValueError("qpoint and band_num are invalid") # magnitude - magnitude = ( - MAX_MAGNITUDE - MIN_MAGNITUDE - ) * magnitude_fraction + MIN_MAGNITUDE + # magnitude = ( + # MAX_MAGNITUDE - MIN_MAGNITUDE + # ) * magnitude_fraction + MIN_MAGNITUDE # set maximum scale for supercell to limit size max_sc_scale = max( @@ -1231,8 +1285,7 @@ def update_crystal_animation( json_data=json_data, band=band_num, qpoint=qpoint, - total_repeat_cell_cnt=total_repeat_cell_cnt, - magnitude=magnitude, + magnitude=magnitude_fraction, velocity=velocity, ), [None, legend_layout], diff --git a/crystal_toolkit/core/jupyter.py b/crystal_toolkit/core/jupyter.py index 3f65187d..887b848b 100644 --- a/crystal_toolkit/core/jupyter.py +++ b/crystal_toolkit/core/jupyter.py @@ -3,7 +3,7 @@ from __future__ import annotations import socketserver -from typing import TYPE_CHECKING, ClassVar +from typing import ClassVar from warnings import warn from dash import Dash @@ -13,12 +13,14 @@ from pymatgen.core.structure import SiteCollection import crystal_toolkit.helpers.layouts as ctl +from crystal_toolkit.components.phonon import PhononBandstructureAndDosComponent from crystal_toolkit.components.structure import StructureMoleculeComponent +from crystal_toolkit.core.mpcomponent import MPComponent from crystal_toolkit.core.plugin import CrystalToolkitPlugin from crystal_toolkit.settings import SETTINGS -if TYPE_CHECKING: - from crystal_toolkit.core.mpcomponent import MPComponent +# if TYPE_CHECKING: +# from crystal_toolkit.core.mpcomponent import MPComponent class _JupyterRenderer: @@ -29,6 +31,8 @@ class _JupyterRenderer: MoleculeGraph: StructureMoleculeComponent, } + ctk_registry: ClassVar[set[MPComponent]] = {PhononBandstructureAndDosComponent} + @staticmethod def _find_available_port(): """Find an available port. @@ -65,6 +69,13 @@ def display(self, obj): ) return self.run(layout) + if any(isinstance(obj, kls) for kls in self.ctk_registry): + layout = ctl.Block( + [obj.layout(jupyter=True)], + style={"margin-top": "1rem", "margin-left": "1rem"}, + ) + return self.run(layout) + raise ValueError(f"No component defined for object of type {type(obj)}.") @@ -88,8 +99,12 @@ def _repr_mimebundle_(self, include=None, exclude=None): def _ipython_display_(self): """Display MSONable objects using a Crystal Toolkit component, if available.""" - if any(isinstance(self, x) for x in _JupyterRenderer.registry): + if any(isinstance(self, x) for x in _JupyterRenderer.registry) or any( + isinstance(self, x) for x in _JupyterRenderer.ctk_registry + ): return _JupyterRenderer().display(self) + # if any(isinstance(self, x) for x in _JupyterRenderer.registry): + # return _JupyterRenderer().display(self) # To be strict here, we could use inspect.signature # and .return_annotation is either a Scene or a go.Figure respectively @@ -125,3 +140,5 @@ def patch_msonable(): MSONable._repr_mimebundle_ = _repr_mimebundle_ MSONable.display_json = _display_json MSONable._ipython_display_ = _ipython_display_ + + MPComponent._ipython_display_ = _ipython_display_ From d903c73954cacb537ad736d78296ecdcc584c970 Mon Sep 17 00:00:00 2001 From: Min-Hsueh Chiu Date: Mon, 20 Jul 2026 16:25:52 -0700 Subject: [PATCH 2/2] remove old magnitude definition --- crystal_toolkit/components/phonon.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/crystal_toolkit/components/phonon.py b/crystal_toolkit/components/phonon.py index a7c20d22..36b63d5e 100644 --- a/crystal_toolkit/components/phonon.py +++ b/crystal_toolkit/components/phonon.py @@ -1269,11 +1269,6 @@ def update_crystal_animation( if qpoint == -1 or band_num == -1: raise ValueError("qpoint and band_num are invalid") - # magnitude - # magnitude = ( - # MAX_MAGNITUDE - MIN_MAGNITUDE - # ) * magnitude_fraction + MIN_MAGNITUDE - # set maximum scale for supercell to limit size max_sc_scale = max( 1,