From 71cad49e9e084d1b9ca08fa67dad9e6d63af0476 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Tue, 3 Nov 2026 18:20:39 -0600 Subject: [PATCH 1/2] Replace py4DSTEM TEM diffraction with the quantEM browser simulator TEMDiffractionComponent now shows DiffractionSim from the quantem.diffsim package. The server converts the structure to CIF text of its conventional standard cell, and the browser computes the reflections, structure factors and Bloch wave intensities, so the component no longer imports py4DSTEM or NumPy. The temdiff extra installs quantem.diffsim>=0.1, and the example app switches between Si, SrTiO3 and GaN. Fixes #547. Co-Authored-By: Claude Opus 5.5 --- .../apps/examples/diffraction_tem.py | 63 ++- crystal_toolkit/components/diffraction_tem.py | 405 ++---------------- pyproject.toml | 2 +- 3 files changed, 89 insertions(+), 381 deletions(-) diff --git a/crystal_toolkit/apps/examples/diffraction_tem.py b/crystal_toolkit/apps/examples/diffraction_tem.py index 6e029c65..47e9bebd 100644 --- a/crystal_toolkit/apps/examples/diffraction_tem.py +++ b/crystal_toolkit/apps/examples/diffraction_tem.py @@ -1,8 +1,10 @@ from __future__ import annotations import dash +from dash import dcc +from dash.dependencies import Input, Output -# create our crystal structure using pymatgen +# create our crystal structures using pymatgen from pymatgen.core import Lattice, Structure import crystal_toolkit.components as ctc @@ -12,28 +14,46 @@ app = dash.Dash(assets_folder=SETTINGS.ASSETS_PATH) -structure = Structure( - Lattice.cubic(5.44), - ["Si"] * 8, - [ - [0.25, 0.75, 0.25], - [0.0, 0.0, 0.5], - [0.25, 0.25, 0.75], - [0.0, 0.5, 0.0], - [0.75, 0.75, 0.75], - [0.5, 0.0, 0.0], - [0.75, 0.25, 0.25], - [0.5, 0.5, 0.5], - ], -) +structures = { + "Si": Structure( + Lattice.cubic(5.44), + ["Si"] * 8, + [ + [0.25, 0.75, 0.25], + [0.0, 0.0, 0.5], + [0.25, 0.25, 0.75], + [0.0, 0.5, 0.0], + [0.75, 0.75, 0.75], + [0.5, 0.0, 0.0], + [0.75, 0.25, 0.25], + [0.5, 0.5, 0.5], + ], + ), + "SrTiO3": Structure( + Lattice.cubic(3.905), + ["Sr", "Ti", "O", "O", "O"], + [[0, 0, 0], [0.5, 0.5, 0.5], [0.5, 0.5, 0], [0.5, 0, 0.5], [0, 0.5, 0.5]], + ), + "GaN": Structure.from_spacegroup( + "P6_3mc", + Lattice.hexagonal(3.189, 5.185), + ["Ga", "N"], + [[1 / 3, 2 / 3, 0], [1 / 3, 2 / 3, 0.377]], + ), +} -tem_component = ctc.TEMDiffractionComponent(initial_structure=structure) +tem_component = ctc.TEMDiffractionComponent(initial_structure=structures["Si"]) + +structure_choice = dcc.Dropdown( + id="structure-choice", options=list(structures), value="Si", clearable=False +) # example layout to demonstrate capabilities of component layout = Container( [ H1("TEMDiffractionComponent Example"), H3("Generated from Structure object"), + structure_choice, tem_component.layout(), ] ) @@ -41,6 +61,17 @@ # as explained in "preamble" section in documentation ctc.register_crystal_toolkit(app=app, layout=layout) + +@app.callback( + Output(tem_component.id("structure"), "data"), + Input(structure_choice, "value"), + prevent_initial_call=True, +) +def load_structure(name: str) -> dict: + """Show the chosen structure in the simulator.""" + return structures[name].as_dict() + + # run this app with "python path/to/this/file.py" # in production, deploy behind gunicorn or similar # see Dash docs for more info diff --git a/crystal_toolkit/components/diffraction_tem.py b/crystal_toolkit/components/diffraction_tem.py index 8c9638c7..2ec3c09a 100644 --- a/crystal_toolkit/components/diffraction_tem.py +++ b/crystal_toolkit/components/diffraction_tem.py @@ -1,397 +1,74 @@ from __future__ import annotations -from time import time from typing import TYPE_CHECKING from warnings import warn -import numpy as np -import plotly.graph_objects as go -from dash import dcc, html from dash.dependencies import Input, Output +from dash.exceptions import PreventUpdate from crystal_toolkit.core.mpcomponent import MPComponent -from crystal_toolkit.helpers.layouts import Box, Column, Columns, Loading, Reveal +from crystal_toolkit.helpers.layouts import Column, Columns if TYPE_CHECKING: + from dash import html from pymatgen.core import Structure try: - import py4DSTEM + import quantem.diffsim as diffsim + import quantem.diffsim.dash as diffsim_dash except ImportError: - no_py4dstem_msg = "requires the py4DSTEM package. Please pip install py4DSTEM." - py4DSTEM = None + no_diffsim_msg = ( + "requires the quantem.diffsim package. Please pip install quantem.diffsim." + ) + diffsim = None -# Author: Steven Zeltmann -# Contact: steven.zeltmann@lbl.gov +# Author: Colin Ophus class TEMDiffractionComponent(MPComponent): + """Electron diffraction simulator from quantEM (the quantem.diffsim package). + + The server converts the structure to CIF text of its conventional standard + cell, and the simulator computes the reflections, structure factors and + Bloch wave intensities in the browser. + """ + def __init__( self, *args, initial_structure: Structure | None = None, **kwargs ) -> None: super().__init__(*args, **kwargs) self.create_store("structure", initial_data=initial_structure) - self.calculator = TEMDiffractionCalculator() - def layout(self) -> Columns: - if not py4DSTEM: - warn(f"{type(self).__name__} {no_py4dstem_msg}") + def layout(self) -> html.Div | Columns: + """The simulator's settings store and the element it draws into.""" + if not diffsim: + warn(f"{type(self).__name__} {no_diffsim_msg}") col = Column( - "This feature will not work unless py4DSTEM is installed on the server." + "This feature will not work unless quantem.diffsim is installed on the server." ) return Columns([col]) - voltage = self.get_numerical_input( - kwarg_label="voltage", - default=200, - step=10.0, - label="Voltage / kV", - help_str="Accelerating voltage for electron beam.", - ) - - beam_direction = self.get_numerical_input( - kwarg_label="beam_direction", - default=[0, 0, 1], - label="Beam Direction", - help_str="The direction of the electron beam fired onto the sample.", - shape=(3,), - is_int=True, - ) - - k_max = self.get_numerical_input( - kwarg_label="k_max", - default=1.5, - step=0.25, - label="Maximum Scattering Angle [Å⁻¹]", - help_str="Maximum scattering angle to compute reciprocal lattice.", - max=10, - ) - - use_dynamical = self.get_bool_input( - kwarg_label="use_dynamical", - default=False, - label="Use Bloch Wave Dynamical Calculator", - help_str="The Bloch wave calculator gives accurate diffraction " - "intensities for thick crystals by including multiple scattering" - " of the electron beam, but requires substantially more computation time.", - ) - - thickness = self.get_numerical_input( - kwarg_label="thickness", - default=500.0, - step=1.0, - label="Thickness [Å]", - help_str="Sample thickness in Ångströms, for dynamical simulation.", - ) - - # Advanced options - - excitation_tol = self.get_numerical_input( - kwarg_label="sigma_excitation_error", - default=0.02, - step=0.001, - label="Excitation error tolerance [Å⁻¹]", - help_str="Standard deviation of Gaussian function for damping reciprocal lattice points.", - max=0.2, - ) - - Fhkl_tol = self.get_numerical_input( - kwarg_label="tol_structure_factor", - default=0.0, - step=0.001, - label="|Fkhl| tolerance", - help_str="Minimum structure factor intensity to include a reflection. Setting" - " this value to zero allows kinematically forbidden reflections to be excited" - " in Bloch wave calculations, but increases computation time.", - ) - - absorption_method_names = { - "Lobato (Elastic)": "Lobato", - "Lobato (Hashimoto absorptive)": "Lobato-absorptive", - "Weickenmeier-Kohl (Elastic)": "WK", - "Weickenmeier-Kohl (Core only)": "WK-C", - "Weickenmeier-Kohl (Phonon only)": "WK-P", - "Weickenmeier-Kohl (Core + Phonon)": "WK-CP", - } - - absorption_methods = self.get_choice_input( - kwarg_label="dynamical_method", - label="Scattering Factor Parameterization", - default="WK-CP", - help_str="Parameterization of absorptive scattering factors, used only" - " for dynamical calculations. Kinematic calculations always use Lobato.", - options=[ - { - "label": name, - "value": shortname, - } - for name, shortname in absorption_method_names.items() - ], - ) - - DWF = self.get_numerical_input( - kwarg_label="DWF", - label="RMS Atomic Displacements [Å]", - default=0.08, - step=0.01, - help_str="RMS atomic displacements used to include thermal smearing of" - "the electrostatic potential when a Weickenmeier-Kohl scattering factor is chosen", - ) - - gamma = self.get_numerical_input( - kwarg_label="gamma", - label="Display Gamma", - default=0.5, - step=0.1, - help_str="Power for scaling intensities in the displayed pattern", - ) - - advanced_options = Reveal( - title="Advanced Options", - children=[ - excitation_tol, - html.Br(), - Fhkl_tol, - html.Br(), - absorption_methods, - html.Br(), - DWF, - html.Br(), - gamma, - ], - id="tem-advanced-options", - ) - - return Columns( - [ - Column([Box(Loading(id=self.id("tem-plot")))], size=8), - Column( - [ - voltage, - html.Br(), - beam_direction, - html.Br(), - k_max, - html.Br(), - use_dynamical, - html.Br(), - thickness, - html.Br(), - advanced_options, - ], - size=4, - ), - ], - ) + return diffsim_dash.layout(self.id("quantem-diffsim")) def generate_callbacks(self, app, cache) -> None: - @app.callback( - Output(self.id("tem-plot"), "children"), - [ - Input(self.id("structure"), "data"), - Input(self.get_all_kwargs_id(), "value"), - ], - ) - def generate_diffraction_pattern(structure, *args): - structure = self.from_data(structure) - kwargs = self.reconstruct_kwargs_from_state() - - # calculator = TEMCalculator(**kwargs) - - print("kwargs", kwargs) - - return dcc.Graph( - figure=self.calculator.get_plot_2d(structure, **kwargs), - responsive=False, - config=dict(displayModeBar=False, displaylogo=False), - ) + """Send the simulator settings for each new structure.""" + if not diffsim: + return + # serves the simulator bundle and adds the clientside callback that draws it + diffsim_dash.register(app) -class TEMDiffractionCalculator: - """Docstring.""" - - def __init__(self) -> None: - # Initialize parameter caches to invalid so that on the first run, - # everything gets computed from scratch. - self.crystal = None - self.voltage = np.nan - self.k_max = np.nan - self.tol_structure_factor = np.nan - self.sigma_excitation_error = np.nan - self.DWF = np.nan - self.dynamical_method = "" - - def get_plot_2d( - self, - structure, - beam_direction, - voltage: float, - k_max: float, - thickness: float, - tol_structure_factor: float, - sigma_excitation_error: float, - use_dynamical: bool, - dynamical_method: str, - DWF: float, - gamma, - # **kwargs, - ) -> go: - """Generate diffraction pattern using py4DSTEM and return as a plotly Figure object.""" - if not py4DSTEM: - raise ImportError(f"{type(self).__name__} {no_py4dstem_msg}") - t0 = time() - # figure out what needs to be recomputed: - new_crystal = py4DSTEM.process.diffraction.Crystal.from_pymatgen_structure( - structure - ) - needs_structure = not self.crystal or not ( - self.crystal.positions.shape[0] == new_crystal.positions.shape[0] - and np.allclose(self.crystal.numbers, new_crystal.numbers) - and np.allclose(self.crystal.cell, new_crystal.cell) - and np.allclose(self.crystal.positions, new_crystal.positions) - ) - - needs_kinematic_SFs = ( - needs_structure - or (self.voltage != voltage) - or (self.k_max != k_max) - or (self.tol_structure_factor != tol_structure_factor) - ) - - needs_dynamic_SFs = use_dynamical and ( - needs_structure - or self.DWF != DWF - or self.dynamical_method != dynamical_method - ) - - # # Check if the cache logic is working - # print( - # f"Needs structure?\t{needs_structure}\nNeeds SFs?:\t{needs_kinematic_SFs}\nNeeds Ug?:\t{needs_dynamic_SFs}" - # ) - - if needs_structure: - self.crystal = py4DSTEM.process.diffraction.Crystal.from_pymatgen_structure( - structure=structure, - ) - - if needs_kinematic_SFs: - self.update_structure_factors(voltage, k_max, tol_structure_factor) - - if needs_dynamic_SFs: - self.update_dynamic_structure_factors(dynamical_method, DWF) - - # generate diffraction pattern - pattern = self.crystal.generate_diffraction_pattern( - zone_axis_lattice=beam_direction, tol_intensity=0.0 - ) - - # rescale intensities - pattern.data["intensity"] /= pattern.data["intensity"].max() - - # perform dynamical simulation, if Bloch is selected - if use_dynamical: - pattern = self.crystal.generate_dynamical_diffraction_pattern( - pattern, thickness=thickness, zone_axis_lattice=beam_direction + @app.callback( + Output(diffsim_dash.settings_id(self.id("quantem-diffsim")), "data"), + Input(self.id("structure"), "data"), + ) + def update_settings(structure): + if not structure: + raise PreventUpdate + return diffsim.settings( + self.from_data(structure), + zone_axis=[0, 0, 1], + mode="nanobeam", + energy_keV=200, + title="Electron diffraction", ) - - print(f"Generated pattern in {time() - t0:.3f} seconds") - - # generate plotly Figure - return self.pointlist_to_spots(pattern, beam_direction, gamma) - - def update_structure_factors( - self, - voltage, - k_max, - tol_structure_factor, - ): - self.crystal.setup_diffraction(accelerating_voltage=voltage * 1e3) - self.crystal.calculate_structure_factors( - k_max=k_max, tol_structure_factor=tol_structure_factor - ) - - self.voltage = voltage - self.k_max = k_max - self.tol_structure_factor = tol_structure_factor - - def update_dynamic_structure_factors( - self, - dynamical_method, - DWF, - ): - self.crystal.calculate_dynamical_structure_factors( - accelerating_voltage=self.voltage * 1e3, - method=dynamical_method, - k_max=self.k_max, - thermal_sigma=DWF, - recompute_kinematic_structure_factors=False, - ) - - self.dynamical_method = dynamical_method - self.DWF = DWF - - def pointlist_to_spots(self, pattern, beam_direction, gamma): - hkl_strings = [ - f"({r['h']} {r['k']} {r['l']})
I: {r['intensity']:.3e}" - for r in pattern.data - ] - - scaled_intensity = pattern.data["intensity"] ** gamma - scaled_intensity /= scaled_intensity.max() - - data = go.Scatter( - x=np.round(pattern.data["qx"], 3), - y=np.round(pattern.data["qy"], 3), - hovertemplate="%{text}
qx: %{x:.2f} Å⁻¹
qy: %{y:.2f}Å⁻¹", - text=hkl_strings, - mode="markers", - marker=dict( - size=12, - cmax=1, - cmin=0, - color=scaled_intensity, - colorscale="gray_r", - ), - showlegend=False, - ) - - plot_max = self.k_max * 1.2 - - layout = go.Layout( - title="2D Diffraction Pattern
Beam Direction: (" - + "".join(str(int(e)) for e in beam_direction) - + ")", - font=dict(size=14, color="#7f7f7f"), - hovermode="closest", - xaxis=dict( - title="qx [Å-1]", - range=[-plot_max, plot_max], - showgrid=False, - zeroline=False, - tickmode="linear", - dtick=0.5, - showticklabels=True, - mirror=True, - ticks="outside", - showline=True, - linecolor="#444", - ), - yaxis=dict( - title="qy [Å-1]", - range=[-plot_max, plot_max], - showgrid=False, - zeroline=False, - tickmode="linear", - dtick=0.5, - showticklabels=True, - mirror=True, - ticks="outside", - showline=True, - linecolor="#444", - ), - width=550, - height=550, - paper_bgcolor="white", - plot_bgcolor="white", - ) - return go.Figure(data=data, layout=layout) diff --git a/pyproject.toml b/pyproject.toml index e61a5e98..d6796b40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dependencies = [ [project.optional-dependencies] server = ["dash-extensions", "gunicorn[gevent]", "habanero", "hiphive", "redis"] robocrys = ["robocrys"] -temdiff = ["py4DSTEM>=0.13.11"] +temdiff = ["quantem.diffsim>=0.1"] fermi = [ "ifermi", "pyfftw", From 2a842f4e0a5f4b00c776d270146d6001b60fe654 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Tue, 3 Nov 2026 18:21:49 -0600 Subject: [PATCH 2/2] Test the TEM diffraction settings callback The test checks the settings that TEMDiffractionComponent sends to the simulator and the bundle route, and skips when quantem.diffsim is not installed. Co-Authored-By: Claude Opus 5.5 --- tests/test_diffraction_tem.py | 39 +++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 tests/test_diffraction_tem.py diff --git a/tests/test_diffraction_tem.py b/tests/test_diffraction_tem.py new file mode 100644 index 00000000..e5e0ba95 --- /dev/null +++ b/tests/test_diffraction_tem.py @@ -0,0 +1,39 @@ +import pytest +from dash import Dash +from dash.exceptions import PreventUpdate +from pymatgen.core import Lattice, Structure + +from crystal_toolkit.components.diffraction_tem import TEMDiffractionComponent + +diffsim_dash = pytest.importorskip("quantem.diffsim.dash") + +Si = Structure.from_spacegroup("Fd-3m", Lattice.cubic(5.431), ["Si"], [[0, 0, 0]]) + + +def test_settings_callback(): + component = TEMDiffractionComponent(initial_structure=Si) + app = Dash() + component.generate_callbacks(app, None) + + # the server callback that writes the settings, and the bundle route from register() + name = component.id("quantem-diffsim") + (key,) = [ + key + for key in app.callback_map + if name in key and diffsim_dash.SETTINGS_TYPE in key + ] + assert any( + diffsim_dash.BUNDLE_ROUTE in rule.rule + for rule in app.server.url_map.iter_rules() + ) + + update_settings = app.callback_map[key]["callback"].__wrapped__ + settings = update_settings(Si.as_dict()) + assert settings["zone_axis"] == [0, 0, 1] + assert settings["mode"] == "nanobeam" + assert settings["energy_keV"] == 200 + assert settings["structure_name"] == "Si" + assert "Fd-3m" in settings["cif"] + + with pytest.raises(PreventUpdate): + update_settings(None)