diff --git a/src/python/impactx/dashboard/Input/csrConfiguration/csrMain.py b/src/python/impactx/dashboard/Input/csrConfiguration/csrMain.py new file mode 100644 index 000000000..a9f664155 --- /dev/null +++ b/src/python/impactx/dashboard/Input/csrConfiguration/csrMain.py @@ -0,0 +1,66 @@ +from trame.widgets import vuetify + +from ...trame_setup import setup_server +from ..generalFunctions import generalFunctions + +server, state, ctrl = setup_server() + +# ----------------------------------------------------------------------------- +# Default State Variables +# ----------------------------------------------------------------------------- + +state.csr_bins = 150 +state.csr_bins_error_message = "" + +# ----------------------------------------------------------------------------- +# +# ----------------------------------------------------------------------------- + + +@state.change("csr_bins") +def on_csr_bins_change(csr_bins, **kwargs): + error_message = generalFunctions.validate_against(csr_bins, "int", ["positive"]) + state.csr_bins_error_message = error_message + generalFunctions.update_simulation_validation_status() + + +# ----------------------------------------------------------------------------- +# UI +# ----------------------------------------------------------------------------- + + +class csrConfiguration: + @staticmethod + def card(): + """ + Creates UI content for CSR. + """ + + with vuetify.VCard(v_show="csr", style="width: 170px;"): + with vuetify.VCardTitle("CSR"): + vuetify.VSpacer() + vuetify.VIcon( + "mdi-information", + classes="ml-2", + click=lambda: generalFunctions.documentation("CSR"), + style="color: #00313C;", + ) + vuetify.VDivider() + with vuetify.VCardText(): + with vuetify.VRow(classes="my-0"): + with vuetify.VCol(classes="py-0"): + vuetify.VSelect( + label="Particle Shape", + v_model=("particle_shape",), + items=([1, 2, 3],), + dense=True, + ) + with vuetify.VRow(classes="my-0"): + with vuetify.VCol(classes="py-0"): + vuetify.VTextField( + label="CSR Bins", + v_model=("csr_bins",), + error_messages=("csr_bins_error_message",), + type="number", + dense=True, + ) diff --git a/src/python/impactx/dashboard/Input/generalFunctions.py b/src/python/impactx/dashboard/Input/generalFunctions.py index 42e7f66cd..94eabcbc7 100644 --- a/src/python/impactx/dashboard/Input/generalFunctions.py +++ b/src/python/impactx/dashboard/Input/generalFunctions.py @@ -28,14 +28,16 @@ def documentation(section_name): Opens a tab to the specified section link in the documentation. :param section_name (str): The name of the documentation section to open. """ - - if section_name == "LatticeElements": - url = "https://impactx.readthedocs.io/en/latest/usage/python.html#lattice-elements" - elif section_name == "BeamDistributions": - url = "https://impactx.readthedocs.io/en/latest/usage/python.html#initial-beam-distributions" - elif section_name == "pythonParameters": - url = "https://impactx.readthedocs.io/en/latest/usage/python.html#general" - else: + url_dict = { + "LatticeElements": "https://impactx.readthedocs.io/en/latest/usage/python.html#lattice-elements", + "BeamDistributions": "https://impactx.readthedocs.io/en/latest/usage/python.html#initial-beam-distributions", + "pythonParameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#general", + "space_charge_documentation": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#space-charge", + "CSR": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#coherent-synchrotron-radiation-csr", + } + + url = url_dict.get(section_name) + if url is None: raise ValueError(f"Invalid section name: {section_name}") if "WSL_DISTRO_NAME" in os.environ: @@ -104,11 +106,11 @@ def validate_against(input_value, value_type, additional_conditions=None): if errors == [] and additional_conditions: for condition in additional_conditions: if condition == "non_zero" and value == 0: - errors.append("Must be non-zero") - if condition == "positive" and value <= 0: - errors.append("Must be positive") - if condition == "negative" and value >= 0: - errors.append("Must be negative") + errors.append("Must be non-zero.") + if condition == "positive" and value < 0: + errors.append("Must be positive.") + if condition == "negative" and value > 0: + errors.append("Must be negative.") return errors @@ -151,6 +153,35 @@ def update_simulation_validation_status(): if state.selectedLatticeList == []: error_details.append("LatticeListIsEmpty") + # Check for errors in CSR parameters + if state.csr_bins_error_message: + error_details.append(f"CSR Bins: {state.csr_bins_error_message}") + + # Check for errors in Space Charge parameters + if state.space_charge: + # n_cell parameters + for direction in ["x", "y", "z"]: + n_cell_error = getattr(state, f"error_message_n_cell_{direction}") + if n_cell_error: + error_details.append(f"n_cell_{direction}: {n_cell_error}") + + # Blocking factor parameters + for direction in ["x", "y", "z"]: + blocking_factor_error = getattr( + state, f"error_message_blocking_factor_{direction}" + ) + if blocking_factor_error: + error_details.append( + f"blocking_factor_{direction}: {blocking_factor_error}" + ) + + # Prob Relative Fields + for index, field in enumerate(state.prob_relative_fields): + if field["error_message"]: + error_details.append( + f"prob_relative[{index}]: {field['error_message']}" + ) + state.disableRunSimulationButton = bool(error_details) # ----------------------------------------------------------------------------- diff --git a/src/python/impactx/dashboard/Input/inputParameters/inputMain.py b/src/python/impactx/dashboard/Input/inputParameters/inputMain.py index 1f9699564..baa5a6386 100644 --- a/src/python/impactx/dashboard/Input/inputParameters/inputMain.py +++ b/src/python/impactx/dashboard/Input/inputParameters/inputMain.py @@ -96,12 +96,19 @@ def card(self): ) vuetify.VDivider() with vuetify.VCardText(): - vuetify.VCombobox( - v_model=("particle_shape",), - label="Particle Shape", - items=([1, 2, 3],), - dense=True, - ) + with vuetify.VRow(classes="py-2"): + with vuetify.VCol(cols=6, classes="py-0"): + vuetify.VCheckbox( + label="Space Charge", + v_model=("space_charge", False), + dense=True, + ) + with vuetify.VCol(cols=6, classes="py-0"): + vuetify.VCheckbox( + label="CSR", + v_model=("csr", False), + dense=True, + ) with vuetify.VRow(classes="my-2"): with vuetify.VCol(cols=6, classes="py-0"): vuetify.VTextField( diff --git a/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeFunctions.py b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeFunctions.py new file mode 100644 index 000000000..514978b85 --- /dev/null +++ b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeFunctions.py @@ -0,0 +1,76 @@ +from ...trame_setup import setup_server +from ..generalFunctions import generalFunctions + +server, state, ctrl = setup_server() + +# ----------------------------------------------------------------------------- +# Functions +# ----------------------------------------------------------------------------- + + +class SpaceChargeFunctions: + @staticmethod + def validate_prob_relative_fields(index, prob_relative_value): + """ + This function checks specific validation requirements + for prob_relative_fields. + :param index: The index of the prob_relative_field modified. + :param prob_relative_value: The numerical value entered by the user. + :return: An error message. An empty string if there is no error. + """ + error_message = "" + + try: + prob_relative_value = float(prob_relative_value) + poisson_solver = state.poisson_solver + + if index == 0: + if poisson_solver == "multigrid": + if prob_relative_value < 3: + error_message = "Must be greater than 3." + elif poisson_solver == "fft": + if prob_relative_value <= 1: + error_message = "Must be greater than 1." + else: + previous_value = float(state.prob_relative[index - 1]) + if prob_relative_value >= previous_value: + error_message = ( + f"Must be less than previous value ({previous_value})." + ) + else: + if prob_relative_value <= 1: + error_message = "Must be greater than 1." + except ValueError: + error_message = "Must be a float." + + return error_message + + @staticmethod + def validate_n_cell_and_blocking_factor(direction): + """ + Validation function for n_cell and blocking_factor parameters. + """ + n_cell_value = getattr(state, f"n_cell_{direction}", None) + blocking_factor_value = getattr(state, f"blocking_factor_{direction}", None) + + n_cell_errors = generalFunctions.validate_against(n_cell_value, "int") + blocking_factor_errors = generalFunctions.validate_against( + blocking_factor_value, "int", ["non_zero", "positive"] + ) + + setattr(state, f"error_message_n_cell_{direction}", "; ".join(n_cell_errors)) + setattr( + state, + f"error_message_blocking_factor_{direction}", + "; ".join(blocking_factor_errors), + ) + + if not n_cell_errors and not blocking_factor_errors: + n_cell_value = int(n_cell_value) + blocking_factor_value = int(blocking_factor_value) + if n_cell_value % blocking_factor_value != 0: + setattr( + state, + f"error_message_n_cell_{direction}", + "Must be a multiple of blocking factor.", + ) diff --git a/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeMain.py b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeMain.py new file mode 100644 index 000000000..1f0d1f4f3 --- /dev/null +++ b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeMain.py @@ -0,0 +1,328 @@ +from trame.widgets import vuetify + +from ...trame_setup import setup_server +from ..generalFunctions import generalFunctions +from .spaceChargeFunctions import SpaceChargeFunctions + +server, state, ctrl = setup_server() + +# ----------------------------------------------------------------------------- +# Default +# ----------------------------------------------------------------------------- + +state.dynamic_size = False +state.max_level = 0 +state.n_cell = [] +state.prob_relative = [] +state.particle_shape = 2 +state.poisson_solver = "fft" + +state.prob_relative_fields = [] +state.n_cell_x = 32 +state.n_cell_y = 32 +state.n_cell_z = 32 + +state.blocking_factor_x = 16 +state.blocking_factor_y = 16 +state.blocking_factor_z = 16 + +state.mlmg_relative_tolerance = 1.0e-7 +state.mlmg_absolute_tolerance = 0.0 +state.mlmg_max_iters = 100 +state.mlmg_verbosity = 1 + +state.error_message_mlmg_relative_tolerance = "" +state.error_message_mlmg_absolute_tolerance = "" +state.error_message_mlmg_max_iters = "" +state.error_message_mlmg_verbosity = "" + +# ----------------------------------------------------------------------------- +# Helper functions +# ----------------------------------------------------------------------------- + + +def populate_prob_relative_fields(max_level): + num_prob_relative_fields = int(max_level) + 1 + + if state.poisson_solver == "fft": + state.prob_relative = [1.1] + [0.0] * (num_prob_relative_fields - 1) + elif state.poisson_solver == "multigrid": + state.prob_relative = [3.1] + [0.0] * (num_prob_relative_fields - 1) + else: + state.prob_relative = [0.0] * num_prob_relative_fields + + state.prob_relative_fields = [ + { + "value": state.prob_relative[i], + "error_message": SpaceChargeFunctions.validate_prob_relative_fields( + i, state.prob_relative[i] + ), + } + for i in range(num_prob_relative_fields) + ] + + +def update_blocking_factor_and_n_cell(category, kwargs): + directions = ["x", "y", "z"] + + for state_name, value in kwargs.items(): + if any(state_name == f"{category}_{dir}" for dir in directions): + direction = state_name.split("_")[-1] + SpaceChargeFunctions.validate_n_cell_and_blocking_factor(direction) + + n_cell_error = getattr(state, f"error_message_n_cell_{direction}") + blocking_factor_error = getattr( + state, f"error_message_blocking_factor_{direction}" + ) + + if not n_cell_error: + n_cell_value = getattr(state, f"n_cell_{direction}") + setattr(state, f"n_cell_{direction}", int(n_cell_value)) + + if not blocking_factor_error: + blocking_factor_value = getattr(state, f"blocking_factor_{direction}") + setattr( + state, f"blocking_factor_{direction}", int(blocking_factor_value) + ) + + state.n_cell = [getattr(state, f"n_cell_{dir}", 0) for dir in directions] + state.blocking_factor = [ + getattr(state, f"blocking_factor_{dir}", 0) for dir in directions + ] + + +# ----------------------------------------------------------------------------- +# Decorators +# ----------------------------------------------------------------------------- +@state.change("poisson_solver") +def on_poisson_solver_change(poisson_solver, **kwargs): + populate_prob_relative_fields(state.max_level) + state.dirty("prob_relative_fields") + generalFunctions.update_simulation_validation_status() + + +@state.change("space_charge") +def on_space_charge_change(space_charge, **kwargs): + state.dynamic_size = space_charge + generalFunctions.update_simulation_validation_status() + + +@state.change("max_level") +def on_max_level_change(max_level, **kwargs): + populate_prob_relative_fields(max_level) + generalFunctions.update_simulation_validation_status() + + +@state.change("blocking_factor_x", "blocking_factor_y", "blocking_factor_z") +def on_blocking_factor_change(**kwargs): + update_blocking_factor_and_n_cell("blocking_factor", kwargs) + generalFunctions.update_simulation_validation_status() + + +@state.change("n_cell_x", "n_cell_y", "n_cell_z") +def on_n_cell_change(**kwargs): + update_blocking_factor_and_n_cell("n_cell", kwargs) + generalFunctions.update_simulation_validation_status() + + +@ctrl.add("update_prob_relative") +def on_update_prob_relative_call(index, value): + index = int(index) + + try: + prob_relative_value = float(value) + state.prob_relative[index] = prob_relative_value + except ValueError: + prob_relative_value = 0.0 + state.prob_relative[index] = prob_relative_value + + # Validate the updated value + error_message = SpaceChargeFunctions.validate_prob_relative_fields( + index, prob_relative_value + ) + + state.prob_relative_fields[index]["value"] = value + state.prob_relative_fields[index]["error_message"] = error_message + + # Validate next index if it exists + if index + 1 < len(state.prob_relative): + next_value = state.prob_relative[index + 1] + + next_error_message = SpaceChargeFunctions.validate_prob_relative_fields( + index + 1, next_value + ) + state.prob_relative_fields[index + 1]["error_message"] = next_error_message + + state.dirty("prob_relative_fields") + generalFunctions.update_simulation_validation_status() + + +# ----------------------------------------------------------------------------- +# UI +# ----------------------------------------------------------------------------- + + +class SpaceChargeConfiguration: + @staticmethod + def card(): + """ + Creates UI content for space charge configuration + """ + + with vuetify.VDialog(v_model=("showSpaceChargeDialog", False), width="500px"): + SpaceChargeConfiguration.dialog_space_charge_settings() + + with vuetify.VCard(v_show="space_charge", style="width: 340px;"): + with vuetify.VCardTitle("Space Charge"): + vuetify.VSpacer() + vuetify.VIcon( + "mdi-cog", + classes="ml-2", + v_if="poisson_solver == 'multigrid'", + click="showSpaceChargeDialog = true", + style="cursor: pointer;", + ) + vuetify.VIcon( + "mdi-information", + classes="ml-2", + click=lambda: generalFunctions.documentation( + "space_charge_documentation" + ), + style="color: #00313C;", + ) + vuetify.VDivider() + with vuetify.VCardText(): + with vuetify.VRow(classes="my-0"): + with vuetify.VCol(cols=5, classes="py-0"): + vuetify.VSelect( + label="Poisson Solver", + v_model=("poisson_solver",), + items=(["multigrid", "fft"],), + dense=True, + hide_details=True, + ) + with vuetify.VCol(cols=4, classes="py-0"): + vuetify.VSelect( + label="Particle Shape", + v_model=("particle_shape",), + items=([1, 2, 3],), + dense=True, + ) + with vuetify.VCol(cols=3, classes="py-0"): + vuetify.VSelect( + label="Max Level", + v_model=("max_level",), + items=([0, 1, 2, 3, 4],), + dense=True, + ) + with vuetify.VCol(classes="pa-0"): + vuetify.VListItemSubtitle( + "nCell", + classes="font-weight-bold black--text", + ) + with vuetify.VRow(classes="my-0"): + for direction in ["x", "y", "z"]: + with vuetify.VCol(cols=4, classes="py-0"): + vuetify.VTextField( + placeholder=direction, + v_model=(f"n_cell_{direction}",), + error_messages=(f"error_message_n_cell_{direction}",), + type="number", + dense=True, + style="margin-top: -5px", + ) + with vuetify.VCol(classes="pa-0"): + vuetify.VListItemSubtitle( + "Blocking Factor", + classes="font-weight-bold black--text mt-2", + ) + with vuetify.VRow(classes="my-0"): + for direction in ["x", "y", "z"]: + with vuetify.VCol(cols=4, classes="py-0"): + vuetify.VTextField( + placeholder=direction, + v_model=(f"blocking_factor_{direction}",), + error_messages=( + f"error_message_blocking_factor_{direction}", + ), + type="number", + dense=True, + style="margin-top: -5px", + ) + with vuetify.VCol(classes="pa-0"): + vuetify.VListItemSubtitle( + "prob_relative", + classes="font-weight-bold black--text mt-2", + ) + with vuetify.VRow(classes="my-0"): + with vuetify.VCol( + v_for=("(field, index) in prob_relative_fields",), + classes="py-0", + ): + vuetify.VTextField( + placeholder=("val."), + v_model=("field.value",), + input=(ctrl.update_prob_relative, "[index, $event]"), + error_messages=("field.error_message",), + type="number", + dense=True, + style="margin-top: -5px", + ) + + @staticmethod + def dialog_space_charge_settings(): + """ + Creates UI content for space charge configuration + settings. + """ + with vuetify.VCard(): + with vuetify.VTabs( + v_model=("space_charge_tab", "Advanced Multigrid Settings") + ): + vuetify.VTab("Settings") + vuetify.VDivider() + with vuetify.VTabsItems(v_model="space_charge_tab"): + with vuetify.VTabItem(): + with vuetify.VContainer(fluid=True): + with vuetify.VRow( + classes="my-2", v_if="poisson_solver == 'multigrid'" + ): + with vuetify.VCol(cols=6, classes="py-0"): + vuetify.VTextField( + label="MLMG Relative Tolerance", + v_model=("mlmg_relative_tolerance",), + error_messages=( + "error_message_mlmg_relative_tolerance", + ), + type="number", + dense=True, + ) + with vuetify.VCol(cols=6, classes="py-0"): + vuetify.VTextField( + label="MLMG Absolute Tolerance", + v_model=("mlmg_absolute_tolerance",), + error_messages=( + "error_message_mlmg_absolute_tolerance", + ), + type="number", + dense=True, + ) + with vuetify.VRow( + classes="my-0", v_if="poisson_solver == 'multigrid'" + ): + with vuetify.VCol(cols=6, classes="py-0"): + vuetify.VTextField( + label="MLMG Max Iterations", + v_model=("mlmg_max_iters",), + error_messages=("error_message_mlmg_max_iters",), + type="number", + dense=True, + ) + with vuetify.VCol(cols=6, classes="py-0"): + vuetify.VTextField( + label="MLMG Verbosity", + v_model=("mlmg_verbosity",), + error_messages=("error_message_mlmg_verbosity",), + type="number", + dense=True, + ) diff --git a/src/python/impactx/dashboard/Toolbar/exportTemplate.py b/src/python/impactx/dashboard/Toolbar/exportTemplate.py index 4effd2de6..c7fb192e5 100644 --- a/src/python/impactx/dashboard/Toolbar/exportTemplate.py +++ b/src/python/impactx/dashboard/Toolbar/exportTemplate.py @@ -57,6 +57,48 @@ def build_lattice_list(): return f"lattice_configuration = [\n {lattice_elements}\n]" +def build_space_charge_or_csr(): + """ + Generates simulation content for space charge + and csr. + """ + if state.space_charge: + content = f"""# Space Charge +sim.csr = {state.csr} +sim.space_charge = {state.space_charge} +sim.dynamic_size = {state.dynamic_size} +sim.poisson_solver = '{state.poisson_solver}' +sim.particle_shape = {state.particle_shape} +sim.max_level = {state.max_level} +sim.n_cell = {state.n_cell} +sim.blocking_factor_x = {state.blocking_factor_x} +sim.blocking_factor_y = {state.blocking_factor_y} +sim.blocking_factor_z = {state.blocking_factor_z} +sim.prob_relative = {state.prob_relative} +""" + if state.poisson_solver == "multigrid": + content += f""" +# Space Charge - Multigrid-Specific Numerical Options +sim.mlmg_relative_tolerance = {state.mlmg_relative_tolerance} +sim.mlmg_absolute_tolerance = {state.mlmg_absolute_tolerance} +sim.mlmg_max_iters = {state.mlmg_max_iters} +sim.mlmg_verbosity = {state.mlmg_verbosity} + """ + elif state.csr: + content = f"""# Coherent Synchrotron Radiation +sim.space_charge = {state.space_charge} +sim.csr = {state.csr} +sim.particle_shape = {state.particle_shape} +sim.csr_bins = {state.csr_bins} + """ + else: + content = f""" +sim.particle_shape = {state.particle_shape} +""" + + return content + + # ----------------------------------------------------------------------------- # Trame setup # ----------------------------------------------------------------------------- @@ -72,9 +114,7 @@ def input_file(): sim = ImpactX() -sim.particle_shape = {state.particle_shape} -sim.space_charge = False -sim.csr = False +{build_space_charge_or_csr()} sim.slice_step_diagnostics = True sim.init_grids() @@ -97,6 +137,7 @@ def input_file(): # Simulate sim.track_particles() +# Clean Shutdown sim.finalize() """ diff --git a/src/python/impactx/dashboard/__main__.py b/src/python/impactx/dashboard/__main__.py index b7a48d8e6..4f15fa468 100644 --- a/src/python/impactx/dashboard/__main__.py +++ b/src/python/impactx/dashboard/__main__.py @@ -13,9 +13,11 @@ from trame.widgets import router, vuetify, xterm from .Analyze.plotsMain import AnalyzeSimulation +from .Input.csrConfiguration.csrMain import csrConfiguration from .Input.distributionParameters.distributionMain import DistributionParameters from .Input.inputParameters.inputMain import InputParameters from .Input.latticeConfiguration.latticeMain import LatticeConfiguration +from .Input.space_charge_configuration.spaceChargeMain import SpaceChargeConfiguration from .Input.trameFunctions import TrameFunctions from .start import main from .Toolbar.toolbarMain import Toolbars @@ -36,6 +38,11 @@ with vuetify.VRow(no_gutters=True): with vuetify.VCol(cols="auto", classes="pa-2"): inputParameters.card() + with vuetify.VCol(cols="auto", classes="pa-2"): + SpaceChargeConfiguration.card() + with vuetify.VCol(cols="auto", classes="pa-2"): + csrConfiguration.card() + with vuetify.VRow(no_gutters=True): with vuetify.VCol(cols="auto", classes="pa-2"): DistributionParameters.card() with vuetify.VRow(no_gutters=True): diff --git a/src/python/impactx/dashboard/simulation.py b/src/python/impactx/dashboard/simulation.py index fbf696518..b76b48845 100644 --- a/src/python/impactx/dashboard/simulation.py +++ b/src/python/impactx/dashboard/simulation.py @@ -1,81 +1,104 @@ -""" -This file is part of ImpactX - -Copyright 2024 ImpactX contributors -Authors: Parthib Roy, Axel Huebl -License: BSD-3-Clause-LBNL -""" - -from .trame_setup import setup_server - -server, state, ctrl = setup_server() - -import base64 -import io - -from impactx import Config, ImpactX - -from .Analyze.plot_PhaseSpaceProjections.phaseSpaceSettings import ( - adjusted_settings_plot, -) -from .Input.distributionParameters.distributionMain import distribution_parameters -from .Input.latticeConfiguration.latticeMain import lattice_elements - -# Call MPI_Init and MPI_Finalize only once: -if Config.have_mpi: - from mpi4py import MPI # noqa - - -def fig_to_base64(fig): - """ - Puts png in trame-compatible form - """ - buf = io.BytesIO() - fig.savefig(buf, format="png") - buf.seek(0) - return base64.b64encode(buf.read()).decode("utf-8") - - -def run_simulation(): - """ - This tests using ImpactX and Pandas Dataframes - """ - sim = ImpactX() - - sim.particle_shape = state.particle_shape - sim.space_charge = False - sim.slice_step_diagnostics = True - sim.init_grids() - - # init particle beam - kin_energy_MeV = state.kin_energy_MeV - bunch_charge_C = state.bunch_charge_C - npart = state.npart - - # reference particle - pc = sim.particle_container() - ref = pc.ref_particle() - ref.set_charge_qe(state.charge_qe).set_mass_MeV(state.mass_MeV).set_kin_energy_MeV( - kin_energy_MeV - ) - - distribution = distribution_parameters() - sim.add_particles(bunch_charge_C, distribution, npart) - - lattice_configuration = lattice_elements() - - sim.lattice.extend(lattice_configuration) - - # simulate - sim.track_particles() - - fig = adjusted_settings_plot(pc) - fig_original = pc.plot_phasespace() - - if fig_original is not None: - image_base64 = fig_to_base64(fig_original) - state.image_data = f"data:image/png;base64, {image_base64}" - - sim.finalize() - - return fig +""" +This file is part of ImpactX + +Copyright 2024 ImpactX contributors +Authors: Parthib Roy, Axel Huebl +License: BSD-3-Clause-LBNL +""" + +from .trame_setup import setup_server + +server, state, ctrl = setup_server() + +import base64 +import io + +from impactx import Config, ImpactX + +from .Analyze.plot_PhaseSpaceProjections.phaseSpaceSettings import ( + adjusted_settings_plot, +) +from .Input.distributionParameters.distributionMain import distribution_parameters +from .Input.latticeConfiguration.latticeMain import lattice_elements + +# Call MPI_Init and MPI_Finalize only once: +if Config.have_mpi: + from mpi4py import MPI # noqa + + +def fig_to_base64(fig): + """ + Puts png in trame-compatible form + """ + buf = io.BytesIO() + fig.savefig(buf, format="png") + buf.seek(0) + return base64.b64encode(buf.read()).decode("utf-8") + + +def run_simulation(): + """ + This tests runs a simulation on ImpactX + based on user inputs from the dashboard. + """ + sim = ImpactX() + + # space charge selections + if state.space_charge: + sim.max_level = state.max_level + sim.n_cell = state.n_cell + sim.particle_shape = state.particle_shape + sim.poisson_solver = state.poisson_solver + sim.space_charge = state.space_charge + sim.dynamic_size = state.dynamic_size + sim.prob_relative = state.prob_relative + sim.blocking_factor_x = [state.blocking_factor_x] + sim.blocking_factor_y = [state.blocking_factor_y] + sim.blocking_factor_z = [state.blocking_factor_z] + + if state.poisson_solver == "multigrid": + sim.mlmg_relative_tolerance = state.mlmg_relative_tolerance + sim.mlmg_absolute_tolerance = state.mlmg_absolute_tolerance + sim.mlmg_max_iters = state.mlmg_max_iters + sim.mlmg_verbosity = state.mlmg_verbosity + # csr + if state.csr: + sim.csr = state.csr + sim.csr_bins = state.csr_bins + + sim.particle_shape = state.particle_shape + sim.slice_step_diagnostics = True + sim.init_grids() + + # init particle beam + kin_energy_MeV = state.kin_energy_MeV + bunch_charge_C = state.bunch_charge_C + npart = state.npart + + # reference particle + pc = sim.particle_container() + ref = pc.ref_particle() + ref.set_charge_qe(state.charge_qe).set_mass_MeV(state.mass_MeV).set_kin_energy_MeV( + kin_energy_MeV + ) + + distribution = distribution_parameters() + sim.add_particles(bunch_charge_C, distribution, npart) + + lattice_configuration = lattice_elements() + + sim.lattice.extend(lattice_configuration) + + # simulate + sim.evolve() + + fig = adjusted_settings_plot(pc) + fig_original = pc.plot_phasespace() + + if fig_original is not None: + image_base64 = fig_to_base64(fig_original) + state.image_data = f"data:image/png;base64, {image_base64}" + + sim.finalize() + + return fig