From c3f91cea6ea52fadf56637a41886d6d40b52bc4a Mon Sep 17 00:00:00 2001
From: Parthib Roy <proy@lbl.gov>
Date: Sat, 8 Feb 2025 20:19:54 -0800
Subject: [PATCH] Add trcking modes to dashboard ui and setup initial
 connection to sim

---
 src/python/impactx/dashboard/Input/defaults.py           | 6 ++++++
 .../impactx/dashboard/Input/inputParameters/inputMain.py | 8 ++++++--
 src/python/impactx/dashboard/simulation.py               | 9 +++++++--
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/python/impactx/dashboard/Input/defaults.py b/src/python/impactx/dashboard/Input/defaults.py
index 1266efd9f..e9ff80b50 100644
--- a/src/python/impactx/dashboard/Input/defaults.py
+++ b/src/python/impactx/dashboard/Input/defaults.py
@@ -18,6 +18,7 @@ class DashboardDefaults:
     }
 
     INPUT_PARAMETERS = {
+        "tracking_mode": "Particle Tracking",
         "charge_qe": -1,
         "mass_MeV": 0.51099895,
         "npart": 1000,
@@ -67,6 +68,11 @@ class DashboardDefaults:
         "poisson_solver_list": ["fft", "multigrid"],
         "particle_shape_list": [1, 2, 3],
         "max_level_list": [0, 1, 2, 3, 4],
+        "tracking_mode_list": [
+            "Particle Tracking",
+            "Envelope Tracking",
+            "Reference Tracking",
+        ],
     }
 
     # -------------------------------------------------------------------------
diff --git a/src/python/impactx/dashboard/Input/inputParameters/inputMain.py b/src/python/impactx/dashboard/Input/inputParameters/inputMain.py
index b6944b824..026eea132 100644
--- a/src/python/impactx/dashboard/Input/inputParameters/inputMain.py
+++ b/src/python/impactx/dashboard/Input/inputParameters/inputMain.py
@@ -31,8 +31,12 @@ def card(self):
         with vuetify.VCard(style="width: 340px; height: 350px"):
             CardComponents.input_header("Input Parameters")
             with vuetify.VCardText():
-                with vuetify.VRow(classes="py-2"):
-                    with vuetify.VCol(cols=6, classes="py-0"):
+                with vuetify.VRow():
+                    with vuetify.VCol(cols=5.5):
+                        InputComponents.select(
+                            label="Tracking Mode",
+                        )
+                    with vuetify.VCol(cols="auto"):
                         vuetify.VCheckbox(
                             label="Space Charge",
                             v_model=("space_charge", False),
diff --git a/src/python/impactx/dashboard/simulation.py b/src/python/impactx/dashboard/simulation.py
index 37c6b73d7..9200ed645 100644
--- a/src/python/impactx/dashboard/simulation.py
+++ b/src/python/impactx/dashboard/simulation.py
@@ -86,11 +86,16 @@ def run_simulation():
     sim.add_particles(bunch_charge_C, distribution, npart)
 
     lattice_configuration = lattice_elements()
-
     sim.lattice.extend(lattice_configuration)
 
+    tracking_modes = {
+        "Particle Tracking": sim.track_particles,
+        "Envelope Tracking": sim.track_envelope,
+        "Reference Tracking": sim.track_reference,
+    }
+
     # simulate
-    sim.evolve()
+    tracking_modes[state.tracking_mode]()
 
     fig = adjusted_settings_plot(pc)
     ctrl.matplotlib_figure_update(fig)