Skip to content

add PhononBandstructureAndDosComponent to jupyter notebook rendering #541

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
minhsueh wants to merge 2 commits into
base: main
Choose a base branch
from
Open

add PhononBandstructureAndDosComponent to jupyter notebook rendering #541

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5d6642b
add PhononBandstructureAndDosComponent to jupyter notebook rendering
minhsueh Jul 20, 2026
d903c73
remove old magnitude definition
minhsueh Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
176 changes: 112 additions & 64 deletions crystal_toolkit/components/phonon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}


Expand Down Expand Up @@ -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(
[
Expand All @@ -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")
)
Expand Down Expand Up @@ -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
),
Expand Down Expand Up @@ -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"]])])
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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),
Expand All @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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))
json_data, legend_layout, num_sites = self._generate_structure_scene(
bs, color_scheme, scale_x, scale_y, 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()

qpoint = 0
band_num = 0
Expand All @@ -1215,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,
Expand All @@ -1231,8 +1280,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],
Expand Down
25 changes: 21 additions & 4 deletions crystal_toolkit/core/jupyter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -29,6 +31,8 @@ class _JupyterRenderer:
MoleculeGraph: StructureMoleculeComponent,
}

ctk_registry: ClassVar[set[MPComponent]] = {PhononBandstructureAndDosComponent}

@staticmethod
def _find_available_port():
"""Find an available port.
Expand Down Expand Up @@ -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)}.")


Expand All @@ -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
Expand Down Expand Up @@ -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_
Loading