diff --git a/src/python/impactx/dashboard/Toolbar/exportTemplate.py b/src/python/impactx/dashboard/Toolbar/exportTemplate.py new file mode 100644 index 000000000..005663b8c --- /dev/null +++ b/src/python/impactx/dashboard/Toolbar/exportTemplate.py @@ -0,0 +1,91 @@ +from ..Input.distributionParameters.distributionMain import parameter_input_checker +from ..Input.latticeConfiguration.latticeMain import parameter_input_checker_for_lattice +from ..trame_setup import setup_server + +server, state, ctrl = setup_server() + +# ----------------------------------------------------------------------------- +# Helper Functions +# ----------------------------------------------------------------------------- + + +def build_distribution_list(): + """ + Generates an instance of distribution inputs + as a string for exporting purposes. + """ + distribution_name = state.selectedDistribution + parameters = parameter_input_checker() + + distribution_parameters = ",\n ".join( + f"{key}={value}" for key, value in parameters.items() + ) + + return ( + f"distr = distribution.{distribution_name}(\n {distribution_parameters},\n)" + ) + + +def build_lattice_list(): + """ + Generates a string representation of lattice element + inputs for export purposes. + """ + + lattice_elements = ",\n ".join( + f'elements.{element["name"]}(' + + ", ".join( + f"{key}={value}" + for key, value in parameter_input_checker_for_lattice(element).items() + ) + + ")" + for element in state.selectedLatticeList + ) + + return f"lattice_configuration = [\n {lattice_elements}\n]" + + +# ----------------------------------------------------------------------------- +# Trame setup +# ----------------------------------------------------------------------------- + + +def input_file(): + """ + This function creates the template to export + dashboard user inputs into a python script. + """ + script = f""" +from impactx import ImpactX, distribution, elements + +sim = ImpactX() + +sim.particle_shape = {state.particle_shape} +sim.space_charge = False +sim.csr = False +sim.slice_step_diagnostics = True + +sim.init_grids() + +# Initialize particle beam +kin_energy_MeV = {state.kin_energy_MeV} +bunch_charge_C = {state.bunch_charge_C} +npart = {state.npart} + +# 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) + +{build_distribution_list()} +sim.add_particles(bunch_charge_C, distr, npart) + +{build_lattice_list()} +sim.lattice.extend(lattice_configuration) + +# Simulate +sim.evolve() + +sim.finalize() +""" + + return script diff --git a/src/python/impactx/dashboard/Toolbar/toolbarMain.py b/src/python/impactx/dashboard/Toolbar/toolbarMain.py index a842e6e8b..4b71b0aca 100644 --- a/src/python/impactx/dashboard/Toolbar/toolbarMain.py +++ b/src/python/impactx/dashboard/Toolbar/toolbarMain.py @@ -9,11 +9,22 @@ from trame.widgets import vuetify from ..trame_setup import setup_server +from .exportTemplate import input_file server, state, ctrl = setup_server() state.show_dashboard_alert = True +# ----------------------------------------------------------------------------- +# Triggers +# ----------------------------------------------------------------------------- + + +@ctrl.trigger("export") +def on_export_click(): + return input_file() + + # ----------------------------------------------------------------------------- # Common toolbar elements # ----------------------------------------------------------------------------- @@ -25,6 +36,15 @@ class ToolbarElements: Vuetify UI elements for toolbar. """ + @staticmethod + def export_input_data(): + vuetify.VIcon( + "mdi-download", + style="color: #00313C;", + click="utils.download('impactx_simulation.py', trigger('export'), 'text/plain')", + disabled=("disableRunSimulationButton", True), + ) + @staticmethod def plot_options(): vuetify.VSelect( @@ -75,6 +95,8 @@ def input_toolbar(): """ (ToolbarElements.dashboard_info(),) + vuetify.VSpacer() + ToolbarElements.export_input_data() @staticmethod def run_toolbar():