[Draft] Dashboard

Early draft of a control and status dashboard written in trame.

setup.py

@@ -230,6 +230,20 @@ def build_extension(self, ext):
     if ImpactX_MPI == "ON":
         install_requires.append("mpi4py>=2.1.0")
 
+# Extra dependencies
+tests_require = ["matplotlib", "numpy", "pandas", "pytest", "scipy"]
+extras = {
+    "default": install_requires,
+    "test": tests_require,
+    "dashboard": ["trame"],
+    # trame-components trame-vuetify
+    # trame-plotly
+    # trame-matplotlib mpld3
+    # trame-vtk vtk
+}
+extras["all"] = list(set(sum(extras.values(), [])))
+
+
 # keyword reference:
 #   https://packaging.python.org/guides/distributing-packages-using-setuptools
 setup(
@@ -266,8 +280,9 @@ def build_extension(self, ext):
     cmdclass=cmdclass,
     zip_safe=False,
     python_requires=">=3.8",
-    tests_require=["numpy", "pandas", "pytest", "scipy"],
+    tests_require=tests_require,
     install_requires=install_requires,
+    extras_require=extras,
     # cmdclass={'test': PyTest},
     # platforms='any',
     classifiers=[

src/python/ImpactX.cpp

@@ -53,7 +53,7 @@ namespace detail
 
 void init_ImpactX (py::module& m)
 {
-    py::class_<ImpactX> impactx(m, "ImpactX");
+    py::class_<ImpactX> impactx(m, "ImpactX", py::dynamic_attr());
     impactx
         .def(py::init<>())
 

src/python/impactx/Dashboard.py

@@ -0,0 +1,175 @@
+"""
+This file is part of ImpactX
+
+Copyright 2023 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+import asyncio
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def figure_size():
+    if state.figure_size is None:
+        return {}
+
+    dpi = state.figure_size.get("dpi")
+    rect = state.figure_size.get("size")
+    w_inch = rect.get("width") / dpi
+    h_inch = rect.get("height") / dpi
+
+    return {
+        "figsize": (w_inch, h_inch),
+        "dpi": dpi,
+    }
+
+
+def FirstDemo():
+    plt.close("all")
+    fig, ax = plt.subplots(**figure_size())
+    np.random.seed(0)
+    ax.plot(
+        np.random.normal(size=100), np.random.normal(size=100), "or", ms=10, alpha=0.3
+    )
+    ax.plot(
+        np.random.normal(size=100), np.random.normal(size=100), "ob", ms=20, alpha=0.1
+    )
+
+    ax.set_xlabel("this is x")
+    ax.set_ylabel("this is y")
+    ax.set_title("Matplotlib Plot Rendered in D3!", size=14)
+    ax.grid(color="lightgray", alpha=0.7)
+
+    return fig
+
+
+# Function to update histogram in a given Matplotlib widget
+def update_histogram(widget, data):
+    widget.figure.clear()
+    ax = widget.figure.add_subplot(111)
+    ax.hist(data.flatten(), bins=50)
+    widget.draw_idle()
+    # widget.update(fig1)
+
+
+def create_app(server):
+    from trame.ui.vuetify import SinglePageLayout
+    from trame.widgets import matplotlib
+    from trame.widgets.vuetify import VContainer, VSelect, VSpacer
+
+    layout = SinglePageLayout(server)
+    layout.title.set_text("Hello trame")
+    # plot = matplotlib.Figure()
+
+    return layout
+
+    from trame.widgets.trame import SizeObserver
+
+    ctrl = server.controller
+
+    # create GUI layout
+    with SinglePageLayout(server) as layout:
+        layout.title.set_text("trame ❤️ matplotlib")
+
+        with layout.toolbar:
+            VSpacer()
+            VSelect(
+                v_model=("active_figure", "FirstDemo"),
+                items=(
+                    "figures",
+                    [
+                        {"text": "First Demo", "value": "FirstDemo"},
+                        #                        {"text": "Multi Lines", "value": "MultiLines"},
+                        #                        {"text": "Dots and Points", "value": "DotsandPoints"},
+                        #                        {"text": "Moving Window Average", "value": "MovingWindowAverage"},
+                        #                        {"text": "Subplots", "value": "Subplots"},
+                    ],
+                ),
+                hide_details=True,
+                dense=True,
+            )
+
+        with layout.content:
+            with VContainer(fluid=True, classes="fill-height pa-0 ma-0"):
+                with SizeObserver("figure_size"):
+                    html_figure = matplotlib.Figure(style="position: absolute")
+                    ctrl.update_figure = html_figure.update
+
+
+def update_dashboard(sim):
+    update_histogram(widget, data)
+
+
+# -----------------------------------------------------------------------------
+# Life Cycle events
+# -----------------------------------------------------------------------------
+
+
+def server_ready(**state):
+    import json
+
+    print("on_server_ready")
+    print("  => current state:")
+    print(json.dumps(state, indent=2))
+    print("-" * 60)
+
+
+def client_connected():
+    print("on_client_connected")
+
+
+def client_unmounted():
+    print("on_client_unmounted")
+
+
+def client_exited():
+    print("on_client_exited")
+
+
+def server_exited(**state):
+    import json
+
+    print("on_server_exited")
+    print("  => current state:")
+    print(json.dumps(state, indent=2))
+    print("-" * 60)
+
+
+async def start(server):
+    await server.start(
+        exec_mode="coroutine",
+        open_browser=True,
+    )
+
+
+async def init_dashboard(sim):
+    from trame.app import get_server
+
+    sim.trame_server = get_server()
+    sim.trame_server.client_type = "vue2"  # Until Matplotlib is ported to vue3
+
+    # -----------------------------------------------------------------------------
+    # Life Cycle registration
+    # -----------------------------------------------------------------------------
+    ctrl = sim.trame_server.controller
+
+    ctrl.on_server_ready.add(server_ready)
+    ctrl.on_client_connected.add(client_connected)
+    ctrl.on_client_unmounted.add(client_unmounted)
+    ctrl.on_client_exited.add(client_exited)
+    ctrl.on_server_exited.add(server_exited)
+
+    layout = create_app(sim.trame_server)
+
+    await asyncio.create_task(start(sim.trame_server))
+
+
+def register_dashboard(sim):
+    """Simulation helper methods for the Dashboard"""
+
+    # register member functions for ImpactX simulation class
+    sim.dashboard = init_dashboard
+    sim.update_dashboard = update_dashboard

src/python/impactx/__init__.py

@@ -17,6 +17,7 @@
 
 # import core bindings to C++
 from . import impactx_pybind as cxx
+from .Dashboard import register_dashboard
 from .ImpactXParticleContainer import (
     register_ImpactXParticleContainer_extension,
 )
@@ -41,3 +42,4 @@
 
 # Pure Python extensions to ImpactX types
 register_ImpactXParticleContainer_extension(ImpactXParticleContainer)
+register_dashboard(ImpactX)

tests/python/test_dashboard.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 The ImpactX Community
+#
+# Authors: Axel Huebl
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import asyncio
+
+from impactx import ImpactX, distribution, elements, push
+
+
+def test_dashboard():
+    """
+    This tests using ImpactX with an interactive dashboard.
+    """
+    sim = ImpactX()
+
+    sim.particle_shape = 2
+    sim.slice_step_diagnostics = True
+    sim.init_grids()
+
+    # init particle beam
+    kin_energy_MeV = 2.0e3
+    bunch_charge_C = 1.0e-9
+    npart = 10000
+
+    #   reference particle
+    ref = sim.particle_container().ref_particle()
+    ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+    #   particle bunch
+    distr = distribution.Waterbag(
+        sigmaX=3.9984884770e-5,
+        sigmaY=3.9984884770e-5,
+        sigmaT=1.0e-3,
+        sigmaPx=2.6623538760e-5,
+        sigmaPy=2.6623538760e-5,
+        sigmaPt=2.0e-3,
+        muxpx=-0.846574929020762,
+        muypy=0.846574929020762,
+        mutpt=0.0,
+    )
+    sim.add_particles(bunch_charge_C, distr, npart)
+
+    pc = sim.particle_container()
+    assert pc.TotalNumberOfParticles() == npart
+
+    # init accelerator lattice
+    fodo = [
+        elements.Drift(0.25),
+        elements.Quad(1.0, 1.0),
+        elements.Drift(0.5),
+        elements.Quad(1.0, -1.0),
+        elements.Drift(0.25),
+    ]
+    #  assign a fodo segment
+    # sim.lattice = fodo
+
+    #  add 4 more FODO segments
+    for i in range(4):
+        sim.lattice.extend(fodo)
+
+    # add 2 more drifts
+    for i in range(4):
+        sim.lattice.append(elements.Drift(0.25))
+
+    async def run():
+        # start interactive dashboard
+        dashboard = sim.dashboard()
+
+        # run simulation
+        sim.evolve()
+        # TODO / Idea:
+        # - add callbacks to a python "async def" function that calls
+        #     await asyncio.sleep(0)
+        #   to yield control to the event loop
+        # - await sim.evolve_async()
+        # - OR make evolve return an Awaitable object that has a
+        #   __await__ method for each iteration
+        # - add an option of the sorts of "pause" to evolve,
+        #   fall into event loop yield busy loop and listen to step
+        #   or evolve events / state changes
+        # - await sim.evolve_async(pause=True)
+        #   before calling the blocking dashboard is an option, too
+        #   -> but, using the dashboard blocking might not be ideal for
+        #      Jupyter usage...?
+
+        await dashboard
+
+    asyncio.run(run())
+
+
+if __name__ == "__main__":
+    test_dashboard()