diff --git a/src/python/impactx/dashboard/Input/defaults.py b/src/python/impactx/dashboard/Input/defaults.py index f5b552781..af69172a8 100644 --- a/src/python/impactx/dashboard/Input/defaults.py +++ b/src/python/impactx/dashboard/Input/defaults.py @@ -34,7 +34,9 @@ class DashboardDefaults: "particle_shape": 2, "max_level": 0, "n_cell": 32, - "blocking_factor": 16, + "blocking_factor_x": 16, + "blocking_factor_y": 16, + "blocking_factor_z": 16, "prob_relative_first_value_fft": 1.1, "prob_relative_first_value_multigrid": 3.1, "mlmg_relative_tolerance": 1.0e-7, diff --git a/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py b/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py index 563522359..ec2952927 100644 --- a/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py +++ b/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py @@ -54,7 +54,6 @@ def populate_distribution_parameters(selectedDistribution): :param selectedDistribution (str): The name of the selected distribution whose parameters need to be populated. """ - if state.selectedDistributionType == "Twiss": sig = inspect.signature(twiss) state.selectedDistributionParameters = [ @@ -152,6 +151,9 @@ def distribution_parameters(): @state.change("selectedDistribution") def on_distribution_name_change(selectedDistribution, **kwargs): + if state.importing_file: + return + if selectedDistribution == "Thermal": state.selectedDistributionType = "Quadratic Form" state.distributionTypeDisabled = True @@ -163,6 +165,8 @@ def on_distribution_name_change(selectedDistribution, **kwargs): @state.change("selectedDistributionType") def on_distribution_type_change(**kwargs): + if state.importing_file: + return populate_distribution_parameters(state.selectedDistribution) diff --git a/src/python/impactx/dashboard/Toolbar/exportTemplate.py b/src/python/impactx/dashboard/Toolbar/exportTemplate.py index c7fb192e5..c67679649 100644 --- a/src/python/impactx/dashboard/Toolbar/exportTemplate.py +++ b/src/python/impactx/dashboard/Toolbar/exportTemplate.py @@ -62,9 +62,10 @@ def build_space_charge_or_csr(): Generates simulation content for space charge and csr. """ + content = "" + if state.space_charge: - content = f"""# Space Charge -sim.csr = {state.csr} + content += f"""# Space Charge sim.space_charge = {state.space_charge} sim.dynamic_size = {state.dynamic_size} sim.poisson_solver = '{state.poisson_solver}' @@ -83,15 +84,17 @@ def build_space_charge_or_csr(): 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} +""" + if state.csr: + content += f"""# Coherent Synchrotron Radiation sim.csr = {state.csr} -sim.particle_shape = {state.particle_shape} sim.csr_bins = {state.csr_bins} - """ - else: +""" + if not state.space_charge: + content += f""" +sim.particle_shape = {state.particle_shape} +""" + if not content: content = f""" sim.particle_shape = {state.particle_shape} """ diff --git a/src/python/impactx/dashboard/Toolbar/importParser.py b/src/python/impactx/dashboard/Toolbar/importParser.py new file mode 100644 index 000000000..db61f3579 --- /dev/null +++ b/src/python/impactx/dashboard/Toolbar/importParser.py @@ -0,0 +1,126 @@ +from ..Input.distributionParameters.distributionMain import ( + on_distribution_parameter_change, + populate_distribution_parameters, +) +from ..Input.latticeConfiguration.latticeMain import ( + add_lattice_element, + on_lattice_element_parameter_change, +) +from ..trame_setup import setup_server +from .importParserHelper import DashboardParserHelper + +server, state, ctrl = setup_server() + + +class DashboardParser: + """ + Provides functionality to import ImpactX simulation files + to the dashboard and auto-populate the UI with their configurations. + """ + + def file_details(file) -> None: + """ + Displays the size of the imported simulation file. + + :param file: ImpactX simulation file uploaded by the user. + """ + + file_size_in_bytes = file["size"] + if file_size_in_bytes < 1024: + size_str = f"{file_size_in_bytes} B" + elif file_size_in_bytes < 1024 * 1024: + size_str = f"{file_size_in_bytes/1024:.1f} KB" + + state.import_file_details = f"({size_str}) {file['name']}" + + def parse_impactx_simulation_file(file) -> None: + """ + Parses ImpactX simulation file contents. + + :param file: ImpactX simulation file uploaded by the user. + """ + + file_content = DashboardParserHelper.import_file_content(file) + + single_input_contents = DashboardParserHelper.parse_single_inputs(file_content) + list_input_contents = DashboardParserHelper.parse_list_inputs(file_content) + distribution_contents = DashboardParserHelper.parse_distribution(file_content) + lattice_element_contents = DashboardParserHelper.parse_lattice_elements( + file_content + ) + + parsed_values_dictionary = { + **single_input_contents, + **list_input_contents, + **distribution_contents, + **lattice_element_contents, + } + + return parsed_values_dictionary + + def populate_impactx_simulation_file_to_ui(file) -> None: + """ + Auto fills the dashboard with parsed inputs. + + :param file: ImpactX simulation file uploaded by the user. + """ + + imported_data = DashboardParser.parse_impactx_simulation_file(file) + + imported_distribution_data = imported_data["distribution"]["parameters"].items() + imported_lattice_data = imported_data["lattice_elements"] + non_state_inputs = ["distribution", "lattice_elements"] + + # Update state inputs (inputParameters, Space Charge, CSR) + for input_name, input_value in imported_data.items(): + if hasattr(state, input_name) and input_name not in non_state_inputs: + setattr(state, input_name, input_value) + + # Update distribution inputs + if imported_distribution_data: + state.selectedDistribution = imported_data["distribution"]["name"] + state.selectedDistributionType = imported_data["distribution"]["type"] + state.flush() + populate_distribution_parameters(state.selectedDistribution) + + for ( + distr_parameter_name, + distr_parameter_value, + ) in imported_distribution_data: + on_distribution_parameter_change( + distr_parameter_name, distr_parameter_value, "float" + ) + + # Update lattice elements + state.selectedLatticeList = [] + + for lattice_element_index, element in enumerate(imported_lattice_data): + parsed_element = element["element"] + parsed_parameters = element["parameters"] + + state.selectedLattice = parsed_element + add_lattice_element() + + lattice_list_parameters = state.selectedLatticeList[lattice_element_index][ + "parameters" + ] + + for ( + parsed_parameter_name, + parsed_parameter_value, + ) in parsed_parameters.items(): + parameter_type = None + + for parameter_info in lattice_list_parameters: + parameter_info_name = parameter_info["parameter_name"] + if parameter_info_name == parsed_parameter_name: + parameter_type = parameter_info["parameter_type"] + break + + if parameter_type: + on_lattice_element_parameter_change( + lattice_element_index, + parsed_parameter_name, + parsed_parameter_value, + parameter_type, + ) diff --git a/src/python/impactx/dashboard/Toolbar/importParserHelper.py b/src/python/impactx/dashboard/Toolbar/importParserHelper.py new file mode 100644 index 000000000..24f29c8b8 --- /dev/null +++ b/src/python/impactx/dashboard/Toolbar/importParserHelper.py @@ -0,0 +1,154 @@ +import ast +import re + +from ..Input.defaults import DashboardDefaults +from ..trame_setup import setup_server + +server, state, ctrl = setup_server() + + +class DashboardParserHelper: + """ + Helper functions for building dashboard parser. + """ + + @staticmethod + def import_file_content(file: str) -> dict: + """ + Retrieves and prints the content of the uploaded simulation file. + + :param content: The content of the ImpactX simulation file. + """ + if file: + content = file["content"].decode("utf-8") + return content + else: + state.file_content = "" + return "" + + @staticmethod + def parse_single_inputs(content: str) -> dict: + """ + Parses individual input parameters from the simulation file content. + + :param content: The content of the ImpactX simulation file. + """ + reference_dictionary = DashboardDefaults.DEFAULT_VALUES.copy() + + parsing_patterns = [ + r"\b{}\s*=\s*([^#\n]+)", # (param = value) + r"set_{}\(([^)]+)\)", # (set_param(value)) + ] + + for parameter_name in reference_dictionary.keys(): + if parameter_name.endswith("_list"): + continue + + for pattern in parsing_patterns: + pattern_match = re.search(pattern.format(parameter_name), content) + if pattern_match: + value = ast.literal_eval(pattern_match.group(1)) + reference_dictionary[parameter_name] = value + break + + # Handling for kin_energy + kin_energy_pattern_match = re.search( + r"\bkin_energy_MeV\s*=\s*([^#\n]+)", content + ) + if kin_energy_pattern_match: + kin_energy_value = kin_energy_pattern_match.group(1) + reference_dictionary["kin_energy"] = kin_energy_value + + return reference_dictionary + + @staticmethod + def parse_list_inputs(content: str) -> dict: + """ + Parses list-based input parameters from the simulation file content. + + :param content: The content of the ImpactX simulation file. + """ + dictionary = {} + list_inputs = ["n_cell", "prob_relative"] + list_parsing = "{} = (\\[.*?\\])" + + for input_name in list_inputs: + match = re.search(list_parsing.format(input_name), content) + if match: + values = ast.literal_eval(match.group(1).strip()) + + if input_name == "n_cell": + for i, dim in enumerate(["x", "y", "z"]): + dictionary[f"n_cell_{dim}"] = values[i] + + if input_name == "prob_relative": + dictionary["prob_relative"] = values + + return dictionary + + @staticmethod + def parse_distribution(content: str) -> dict: + """ + Parses distribution section from the simulation file content. + + :param content: The content of the ImpactX simulation file. + """ + + dictionary = {"distribution": {"name": "", "type": "", "parameters": {}}} + + distribution_name = re.search(r"distribution\.(\w+)\(", content) + distribution_type_twiss = re.search(r"twiss\((.*?)\)", content, re.DOTALL) + distribution_type_quadratic = re.search( + r"distribution\.\w+\((.*?)\)", content, re.DOTALL + ) + parameters = {} + + def extract_parameters(distribution_type, parsing_pattern): + parameter_pairs = re.findall(parsing_pattern, distribution_type.group(1)) + parsed_parameters = {} + + for param_name, param_value in parameter_pairs: + parsed_parameters[param_name] = param_value + return parsed_parameters + + if distribution_name: + dictionary["distribution"]["name"] = distribution_name.group(1) + + if distribution_type_twiss: + dictionary["distribution"]["type"] = "Twiss" + parameters = extract_parameters( + distribution_type_twiss, r"(\w+)=(\d+\.?\d*)" + ) + elif distribution_type_quadratic: + dictionary["distribution"]["type"] = "Quadratic" + parameters = extract_parameters( + distribution_type_quadratic, r"(\w+)=([^,\)]+)" + ) + + dictionary["distribution"]["parameters"] = parameters + + return dictionary + + @staticmethod + def parse_lattice_elements(content: str) -> dict: + """ + Parses lattice elements from the simulation file content. + + :param content: The content of the ImpactX simulation file. + """ + + dictionary = {"lattice_elements": []} + + lattice_elements = re.findall(r"elements\.(\w+)\((.*?)\)", content) + + for element_name, element_parameter in lattice_elements: + element = {"element": element_name, "parameters": {}} + + parameter_pairs = re.findall(r"(\w+)=([^,\)]+)", element_parameter) + for parameter_name, parameter_value in parameter_pairs: + parameter_value_cleaned = parameter_value.strip("'\"") + element["parameters"][parameter_name] = parameter_value_cleaned + + dictionary["lattice_elements"].append(element) + + return dictionary diff --git a/src/python/impactx/dashboard/Toolbar/toolbarMain.py b/src/python/impactx/dashboard/Toolbar/toolbarMain.py index 4b71b0aca..ab28abd3a 100644 --- a/src/python/impactx/dashboard/Toolbar/toolbarMain.py +++ b/src/python/impactx/dashboard/Toolbar/toolbarMain.py @@ -6,25 +6,53 @@ License: BSD-3-Clause-LBNL """ -from trame.widgets import vuetify +from trame.widgets import html, vuetify from ..trame_setup import setup_server from .exportTemplate import input_file +from .importParser import DashboardParser server, state, ctrl = setup_server() state.show_dashboard_alert = True +state.import_file = False +state.import_file_details = None +state.import_file_error = False +state.importing_file = False # ----------------------------------------------------------------------------- -# Triggers +# Triggers/Controllers # ----------------------------------------------------------------------------- +@ctrl.add("reset_import_file") +def reset_import_file(): + # later replaced by calling reset_function + state.import_file_error = None + state.import_file_details = None + state.import_file = None + state.importing_file = False + + @ctrl.trigger("export") def on_export_click(): return input_file() +@state.change("import_file") +def on_import_file_change(import_file, **kwargs): + if import_file: + try: + state.importing_file = True + DashboardParser.file_details(import_file) + DashboardParser.populate_impactx_simulation_file_to_ui(import_file) + except Exception: + state.import_file_error = True + state.import_file_error_message = "Unable to parse" + finally: + state.importing_file = False + + # ----------------------------------------------------------------------------- # Common toolbar elements # ----------------------------------------------------------------------------- @@ -37,13 +65,15 @@ class ToolbarElements: """ @staticmethod - def export_input_data(): - vuetify.VIcon( - "mdi-download", - style="color: #00313C;", + def export_button(): + with vuetify.VBtn( click="utils.download('impactx_simulation.py', trigger('export'), 'text/plain')", + outlined=True, + small=True, disabled=("disableRunSimulationButton", True), - ) + ): + vuetify.VIcon("mdi-download", left=True, small=True) + html.Span("Export") @staticmethod def plot_options(): @@ -66,6 +96,50 @@ def run_simulation_button(): disabled=("disableRunSimulationButton", True), ) + @staticmethod + def import_button(): + vuetify.VFileInput( + v_model=("import_file",), + accept=".py", + __properties=["accept"], + style="display: none;", + ref="fileInput", + ) + with html.Div( + style="position: relative; margin-right: 8px;", + ): + with vuetify.VBtn( + click="$refs.fileInput.$refs.input.click()", + outlined=True, + small=True, + disabled=("(!import_file_error && import_file_details)",), + color=("import_file_error ? 'error' : ''",), + ): + vuetify.VIcon( + "mdi-upload", + left=True, + small=True, + ) + html.Span("Import") + with html.Div( + style="position: absolute; font-size: 10px; width: 100%; padding-top: 2px; display: flex; justify-content: center; white-space: nowrap;" + ): + html.Span( + "{{ import_file_error ? import_file_error_message : import_file_details }}", + style="text-overflow: ellipsis; overflow: hidden;", + classes=( + "import_file_error ? 'error--text' : 'grey--text text--darken-1'", + ), + ) + vuetify.VIcon( + "mdi-close", + x_small=True, + style="cursor: pointer;", + click=ctrl.reset_import_file, + v_if="import_file_details || import_file_error", + color=("import_file_error ? 'error' : 'grey darken-1'",), + ) + @staticmethod def dashboard_info(): """ @@ -96,7 +170,8 @@ def input_toolbar(): (ToolbarElements.dashboard_info(),) vuetify.VSpacer() - ToolbarElements.export_input_data() + ToolbarElements.import_button() + ToolbarElements.export_button() @staticmethod def run_toolbar():