Change distribution default parameters to Twiss in dashboard (#748)

* Updated UI for distributions

Twiss is now default, 'Native' changed to 'Quadratic Form' and placed as secondary option

* Refactor distribution helper function

Helper function moved to distributionFuntions.py, and simplified code structure. no functionality change

* Updated docstring/comments

* Updated export_template to handle twiss parameters

* Add twiss parameter functionality for distribution

* Updated functionality to correctly display 'Thermal' distribution parameters

Thermal is set to have a single set of parameters - 'k', 'kT', 'kT_halo', 'normali
ze', 'normalize_halo', halo'

* Add units for beta/emitt

src/python/impactx/dashboard/Input/distributionParameters/distributionFunctions.py

@@ -0,0 +1,32 @@
+from ...trame_setup import setup_server
+
+server, state, ctrl = setup_server()
+
+# -----------------------------------------------------------------------------
+# Functions
+# -----------------------------------------------------------------------------
+
+
+class DistributionFunctions:
+    """
+    Helper functions for the distribution parameters.
+    """
+
+    @staticmethod
+    def convert_distribution_parameters_to_valid_type():
+        """
+        Helper function to convert user-inputted distribution parameters
+        from string type to float type.
+
+        :return: A dictionary with parameter names as keys and their validated values.
+        """
+
+        parameter_input = {
+            param["parameter_name"]: float(param["parameter_default_value"])
+            if param_is_valid
+            else 0.0
+            for param in state.selectedDistributionParameters
+            if (param_is_valid := param["parameter_error_message"] == [])
+        }
+
+        return parameter_input

src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py

@@ -6,12 +6,16 @@
 License: BSD-3-Clause-LBNL
 """
 
+import inspect
+
+from distribution_input_helpers import twiss
 from trame.widgets import vuetify
 
 from impactx import distribution
 
 from ...trame_setup import setup_server
 from ..generalFunctions import generalFunctions
+from .distributionFunctions import DistributionFunctions
 
 server, state, ctrl = setup_server()
 
@@ -27,12 +31,13 @@
 )
 
 # -----------------------------------------------------------------------------
-# Default
+# Defaults
 # -----------------------------------------------------------------------------
 
 state.selectedDistribution = "Waterbag"
-state.selectedDistributionType = "Native"
+state.selectedDistributionType = "Twiss"
 state.selectedDistributionParameters = []
+state.distributionTypeDisabled = False
 
 # -----------------------------------------------------------------------------
 # Main Functions
@@ -43,27 +48,52 @@ def populate_distribution_parameters(selectedDistribution):
     """
     Populates distribution parameters based on the selected distribution.
     :param selectedDistribution (str): The name of the selected distribution
-        whos parameters need to be populated.
+    whose parameters need to be populated.
     """
 
-    selectedDistributionParameters = (
-        state.listOfDistributionsAndParametersAndDefault.get(selectedDistribution, [])
-    )
-
-    state.selectedDistributionParameters = [
-        {
-            "parameter_name": parameter[0],
-            "parameter_default_value": parameter[1],
-            "parameter_type": parameter[2],
-            "parameter_error_message": generalFunctions.validate_against(
-                parameter[1], parameter[2]
-            ),
-        }
-        for parameter in selectedDistributionParameters
-    ]
-
-    generalFunctions.update_simulation_validation_status()
-    return selectedDistributionParameters
+    if state.selectedDistributionType == "Twiss":
+        sig = inspect.signature(twiss)
+        state.selectedDistributionParameters = [
+            {
+                "parameter_name": param.name,
+                "parameter_default_value": param.default
+                if param.default != param.empty
+                else None,
+                "parameter_type": "float",  # Hardcoding Twiss to 'float' type.
+                "parameter_error_message": generalFunctions.validate_against(
+                    param.default if param.default != param.empty else None, "float"
+                ),
+                "parameter_units": "m"
+                if "beta" in param.name or "emitt" in param.name
+                else "",
+            }
+            for param in sig.parameters.values()
+        ]
+
+    else:  # when type == 'Quadratic Form'
+        selectedDistributionParameters = (
+            state.listOfDistributionsAndParametersAndDefault.get(
+                selectedDistribution, []
+            )
+        )
+
+        state.selectedDistributionParameters = [
+            {
+                "parameter_name": parameter[0],
+                "parameter_default_value": parameter[1],
+                "parameter_type": parameter[2],
+                "parameter_error_message": generalFunctions.validate_against(
+                    parameter[1], parameter[2]
+                ),
+                "parameter_units": "m"
+                if "beta" in parameter[0] or "emitt" in parameter[0]
+                else "",
+            }
+            for parameter in selectedDistributionParameters
+        ]
+
+        generalFunctions.update_simulation_validation_status()
+        return selectedDistributionParameters
 
 
 def update_distribution_parameters(
@@ -91,34 +121,21 @@ def update_distribution_parameters(
 # -----------------------------------------------------------------------------
 
 
-def parameter_input_checker():
-    """
-    Helper function to check if user input is valid.
-    :return: A dictionary with parameter names as keys and their validated values.
-    """
-
-    parameter_input = {}
-    for param in state.selectedDistributionParameters:
-        if param["parameter_error_message"] == []:
-            parameter_input[param["parameter_name"]] = float(
-                param["parameter_default_value"]
-            )
-        else:
-            parameter_input[param["parameter_name"]] = 0.0
-
-    return parameter_input
-
-
 def distribution_parameters():
     """
-    Writes user input for distribution parameters in suitable format for simulation code.
-    :return: An instance of the selected distribution class, initialized with user-provided parameters.
+    :return: An instance of the selected distribution class,
+    initialized with the appropriate parameters provided by the user.
     """
 
     distribution_name = state.selectedDistribution
-    parameters = parameter_input_checker()
+    parameters = DistributionFunctions.convert_distribution_parameters_to_valid_type()
+
+    if state.selectedDistributionType == "Twiss":
+        twiss_params = twiss(**parameters)
+        distr = getattr(distribution, distribution_name)(**twiss_params)
+    else:
+        distr = getattr(distribution, distribution_name)(**parameters)
 
-    distr = getattr(distribution, distribution_name)(**parameters)
     return distr
 
 
@@ -129,6 +146,12 @@ def distribution_parameters():
 
 @state.change("selectedDistribution")
 def on_distribution_name_change(selectedDistribution, **kwargs):
+    if selectedDistribution == "Thermal":
+        state.selectedDistributionType = "Quadratic Form"
+        state.distributionTypeDisabled = True
+        state.dirty("selectedDistributionType")
+    else:
+        state.distributionTypeDisabled = False
     populate_distribution_parameters(selectedDistribution)
 
 
@@ -172,21 +195,20 @@ def card():
             vuetify.VDivider()
             with vuetify.VCardText():
                 with vuetify.VRow():
-                    with vuetify.VCol(cols=8):
+                    with vuetify.VCol(cols=6):
                         vuetify.VCombobox(
                             label="Select Distribution",
                             v_model=("selectedDistribution",),
                             items=("listOfDistributions",),
                             dense=True,
                         )
-                    with vuetify.VCol(cols=4):
+                    with vuetify.VCol(cols=6):
                         vuetify.VSelect(
                             v_model=("selectedDistributionType",),
                             label="Type",
-                            items=(["Native", "Twiss"],),
-                            # change=(ctrl.kin_energy_unit_change, "[$event]"),
+                            items=(["Twiss", "Quadratic Form"],),
                             dense=True,
-                            disabled=True,
+                            disabled=("distributionTypeDisabled",),
                         )
                 with vuetify.VRow(classes="my-2"):
                     for i in range(3):
@@ -200,6 +222,7 @@ def card():
                                     vuetify.VTextField(
                                         label=("parameter.parameter_name",),
                                         v_model=("parameter.parameter_default_value",),
+                                        suffix=("parameter.parameter_units",),
                                         change=(
                                             ctrl.updateDistributionParameters,
                                             "[parameter.parameter_name, $event, parameter.parameter_type]",

src/python/impactx/dashboard/Input/inputParameters/inputFunctions.py

@@ -26,8 +26,7 @@
 
 class InputFunctions:
     """
-    Helper functions for the
-    User-Input section for beam properties.
+    Helper functions for the beam properties.
     """
 
     @staticmethod

src/python/impactx/dashboard/Toolbar/exportTemplate.py

@@ -1,4 +1,4 @@
-from ..Input.distributionParameters.distributionMain import parameter_input_checker
+from ..Input.distributionParameters.distributionFunctions import DistributionFunctions
 from ..Input.latticeConfiguration.latticeMain import parameter_input_checker_for_lattice
 from ..trame_setup import setup_server
 
@@ -15,15 +15,27 @@ def build_distribution_list():
     as a string for exporting purposes.
     """
     distribution_name = state.selectedDistribution
-    parameters = parameter_input_checker()
+    parameters = DistributionFunctions.convert_distribution_parameters_to_valid_type()
 
-    distribution_parameters = ",\n    ".join(
-        f"{key}={value}" for key, value in parameters.items()
+    indentation = " " * (8 if state.selectedDistributionType == "Twiss" else 4)
+    distribution_parameters = ",\n".join(
+        f"{indentation}{key}={value}" for key, value in parameters.items()
     )
 
-    return (
-        f"distr = distribution.{distribution_name}(\n    {distribution_parameters},\n)"
-    )
+    if state.selectedDistributionType == "Twiss":
+        return (
+            f"distr = distribution.{distribution_name}(\n"
+            f"    **twiss(\n"
+            f"{distribution_parameters},\n"
+            f"    )\n"
+            f")"
+        )
+    else:
+        return (
+            f"distr = distribution.{distribution_name}(\n"
+            f"{distribution_parameters},\n"
+            f")"
+        )
 
 
 def build_lattice_list():
@@ -56,7 +68,7 @@ def input_file():
     dashboard user inputs into a python script.
     """
     script = f"""
-from impactx import ImpactX, distribution, elements
+from impactx import ImpactX, distribution, elements, twiss
 
 sim = ImpactX()