{ "cells": [ { "cell_type": "markdown", "id": "8665bce9-d57c-445a-821a-a596efa3f4b3", "metadata": {}, "source": [ "## Utility Functions" ] }, { "cell_type": "markdown", "id": "5a9fc409-a0e1-4c6b-9544-2c9494f40e22", "metadata": {}, "source": [ "SimFrame provides various utility functions for inspecting, modifying and saving lattices and lattice elements. This notebook will demonstrate some of this functionality.\n", "\n", "For general tips on setting up SimFrame, see [getting_started](./getting_started.ipynb)." ] }, { "cell_type": "code", "execution_count": 2, "id": "ba2a922d-f9fd-4b80-b6a7-7248b20d1601", "metadata": {}, "outputs": [], "source": [ "import sys\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "\n", "import SimulationFramework.Framework as fw # noqa E402\n", "\n", "# Define a new framework instance, in directory 'getting_started'.\n", "# \"clean\" will empty (delete everything!) in the directory if true\n", "# \"verbose\" will print a progressbar if true\n", "framework = fw.Framework(\n", " directory=\"./getting_started\",\n", " clean=False,\n", " verbose=False\n", " )\n", "\n", "scaling = 4" ] }, { "cell_type": "markdown", "id": "9dd5ffc7-c315-4cb2-b5c1-14fafe6f655c", "metadata": {}, "source": [ "Load the settings from a definitions file and check which codes correspond to each line." ] }, { "cell_type": "code", "execution_count": 2, "id": "1c3e2a3f-bc8e-4ddf-a890-fd30188423df", "metadata": {}, "outputs": [], "source": [ "framework.loadSettings(\"Lattices/clara400_v13_SP3.def\")" ] }, { "cell_type": "code", "execution_count": 8, "id": "35af4ed8-0a7a-41b0-931b-9d091c5bd3af", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "['generator: ASTRA', 'injector400: astra', 'S02: elegant', 'L02: elegant', 'S03: elegant', 'L03: elegant', 'S04: elegant', 'L4H: elegant', 'S05: elegant', 'VBC: elegant', 'S06: elegant', 'L04: elegant', 'S07: elegant', 'SP3: elegant']\n" ] } ], "source": [ "print([f\"{line}: {framework[line].code}\" for line in framework.lines])" ] }, { "cell_type": "markdown", "id": "3d394e42-af4b-419e-a706-ad1cdc21aef6", "metadata": {}, "source": [ "Check the global parameters accessible to this instance of SimFrame and all of the `frameworkLattice` objects." ] }, { "cell_type": "code", "execution_count": 9, "id": "d58a7024-55fb-4ca6-9ce3-8d34b2d490eb", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'beam': {'filename': None, 'code': None},\n", " 'GPTLICENSE': '',\n", " 'delete_tracking_files': False,\n", " 'astra_use_wsl': 1,\n", " 'master_subdir': '/home/xkc85723/Documents/simframe/examples/notebooks/getting_started',\n", " 'master_lattice_location': '/home/xkc85723/Documents/masterlattice/MasterLattice/./',\n", " 'simcodes_location': '/home/xkc85723/Documents/simcodes/SimCodes/./'}" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "print(framework.global_parameters)" ] }, { "cell_type": "code", "execution_count": 10, "id": "22d9e304-ef7e-4f29-982e-c3f74b006aae", "metadata": {}, "outputs": [], "source": [ "framework.generator.load_defaults(\"clara_400_2ps_Gaussian\")" ] }, { "cell_type": "code", "execution_count": 15, "id": "31b3d59d-571f-4162-a972-c099dcc5f20c", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "FrameworkSettings({'settingsFilename': '/home/xkc85723/Documents/masterlattice/MasterLattice/./Lattices/clara400_v13_SP3.def', 'files': {'injector400': {'code': 'ASTRA', 'charge': {'cathode': True, 'space_charge_mode': '2D', 'mirror_charge': True}, 'input': {'particle_definition': 'initial_distribution'}, 'output': {'zstart': 0, 'end_element': 'CLA-S02-SIM-APER-01'}}, 'S02': {'code': 'elegant', 'output': {'start_element': 'CLA-S02-SIM-APER-01', 'end_element': 'CLA-L02-SIM-APER-01'}}, 'L02': {'code': 'elegant', 'output': {'start_element': 'CLA-L02-SIM-APER-01', 'end_element': 'CLA-S03-SIM-APER-01'}, 'input': {}}, 'S03': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S03-SIM-APER-01', 'end_element': 'CLA-L03-SIM-APER-01'}}, 'L03': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-L03-SIM-APER-01', 'end_element': 'CLA-S04-SIM-APER-01'}}, 'S04': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S04-SIM-APER-01', 'end_element': 'CLA-L4H-SIM-APER-01'}, 'groups': {'laser-heater': {'dipoleangle': 0.0001, 'width': 0.2, 'gap': 0.02}}}, 'L4H': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-L4H-SIM-APER-01', 'end_element': 'CLA-S05-SIM-APER-01'}}, 'S05': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S05-SIM-APER-01', 'end_element': 'CLA-S05-SIM-MARK-01'}}, 'VBC': {'code': 'elegant', 'model': 'test', 'input': {}, 'output': {'start_element': 'CLA-S05-SIM-MARK-01', 'end_element': 'CLA-S06-SIM-MARK-01'}, 'groups': {'bunch_compressor': {'dipoleangle': 0.1185, 'width': 0.2, 'gap': 0.02}}}, 'S06': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S06-SIM-MARK-01', 'end_element': 'CLA-L04-SIM-APER-01'}}, 'L04': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-L04-SIM-APER-01', 'end_element': 'CLA-S07-SIM-APER-01'}}, 'S07': {'code': 'elegant', 'input': {}, 'charge': {'cathode': False, 'space_charge_mode': '3D'}, 'output': {'start_element': 'CLA-S07-SIM-APER-01', 'end_element': 'CLA-SP3-SIM-MARK-01'}}, 'SP3': {'code': 'elegant', 'input': {}, 'charge': {'cathode': False, 'space_charge_mode': '3D'}, 'output': {'start_element': 'CLA-SP3-SIM-MARK-01', 'end_element': 'CLA-SP3-DIA-FCUP-01'}}}, 'generator': {'code': 'astra', 'default': 'clara_400_3ps'}, 'elements': {'filename': ['YAML/Injector400.yaml', 'YAML/S02.yaml', 'YAML/L02.yaml', 'YAML/S03.yaml', 'YAML/L03.yaml', 'YAML/S04.yaml', 'YAML/L4H.yaml', 'YAML/S05.yaml', 'YAML/VBC.yaml', 'YAML/S06.yaml', 'YAML/L04.yaml', 'YAML/S07_SP3.yaml']}, 'groups': {'bunch_compressor': {'type': 'chicane', 'elements': ['CLA-VBC-MAG-DIP-01', 'CLA-VBC-MAG-DIP-02', 'CLA-VBC-MAG-DIP-03', 'CLA-VBC-MAG-DIP-04']}}, 'global': {'ASTRAsettings': {'High_res': True, 'N_red': 1, 'EmitS': True, 'PhaseS': True, 'TrackS': True, 'RefS': True, 'Track_All': True, 'Phase_Scan': False, 'Auto_Phase': True, 'check_ref_part': False, 'Zphase': 1, 'Zemit': 1050, 'H_max': 0.0007, 'H_min': 0.0007, 'title': 'trial', 'run_no': 1, 'Smooth_x': 2, 'Smooth_y': 2, 'Smooth_z': 2, 'Max_scale': 0.1, 'Max_count': 100, 'cell_var': 2, 'min_grid': 3.424657e-13, 'Lmagnetized': False}}})" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "print(framework.settings)" ] }, { "cell_type": "code", "execution_count": 16, "id": "e727cf61-40e4-4046-9929-2770923c1d2e", "metadata": {}, "outputs": [], "source": [ "framework.change_subdirectory('./utility_functions')" ] }, { "cell_type": "markdown", "id": "a2ae0777-71a6-4f5a-a451-f01069af06b2", "metadata": {}, "source": [ "### Accessing and modifying elements" ] }, { "cell_type": "markdown", "id": "707a8cbe-9cea-4fbd-b013-7f9f9097e35f", "metadata": {}, "source": [ "Elements and their parameters can be retrieved either using `getElement` or as a key in the `framework` object itself." ] }, { "cell_type": "code", "execution_count": 5, "id": "3bac6e0d-b3da-4a78-8269-91eb0a79ba16", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "objectname='CLA-S02-MAG-QUAD-01' objecttype='quadrupole' objectdefaults={} allowedkeywords=['subelement', 'sub_elements', 'length', 'online_model_name', 'controller_name', 'pv_root', 'pv_suffixes', 'field_integral_coefficients', 'global_rotation', 'centre', 'datum', 'start', 'end', 'buffer_start', 'buffer_end', 'buffer_start_length', 'buffer_end_length', 'position_errors', 'rotation_errors', 'tilt_angle', 'array_names', 'rotation', 'smooth', 'field_reference_position', 'parent', 'k1l', 'k1', 'field_definition', 'gradient', 'fringe_field_coefficient', 'field_type', 'scale_field', 'multipoles'] length=0.1007 position_errors=[0, 0, 0] rotation_errors=[0, 0, 0] global_rotation=[0.0, 0.0, 0.0] rotation=[0, 0, 0] starting_rotation=0.0 conversion_rules_elegant={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_file_name': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} conversion_rules_ocelot={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_definition_sdds': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} starting_offset=[0, 0, 0] subelement=False field_definition=None wakefield_definition=None k1l=-0.128366 n_kicks=4 field_reference_position='middle' fringe_field_coefficient=0.0 strength_errors=[0] gradient=0.0 scale_field=False multipoles=[0] smooth=2 centre=[0.0, 0.0, 3.52715] datum=[0.0, 0.0, 3.5775] global_parameters={'beam': {'filename': None, 'code': None}, 'GPTLICENSE': '', 'delete_tracking_files': False, 'astra_use_wsl': 1, 'master_subdir': '/home/xkc85723/Documents/simframe/examples/notebooks/getting_started', 'master_lattice_location': '/home/xkc85723/Documents/masterlattice/MasterLattice/./', 'simcodes_location': '/home/xkc85723/Documents/simcodes/SimCodes/./'} PV='CLA-S02-MAG-QUAD-01' isr_enable=True sr_enable=True type='quadrupole' bore=0.037\n", "objectname='CLA-S02-MAG-QUAD-01' objecttype='quadrupole' objectdefaults={} allowedkeywords=['subelement', 'sub_elements', 'length', 'online_model_name', 'controller_name', 'pv_root', 'pv_suffixes', 'field_integral_coefficients', 'global_rotation', 'centre', 'datum', 'start', 'end', 'buffer_start', 'buffer_end', 'buffer_start_length', 'buffer_end_length', 'position_errors', 'rotation_errors', 'tilt_angle', 'array_names', 'rotation', 'smooth', 'field_reference_position', 'parent', 'k1l', 'k1', 'field_definition', 'gradient', 'fringe_field_coefficient', 'field_type', 'scale_field', 'multipoles'] length=0.1007 position_errors=[0, 0, 0] rotation_errors=[0, 0, 0] global_rotation=[0.0, 0.0, 0.0] rotation=[0, 0, 0] starting_rotation=0.0 conversion_rules_elegant={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_file_name': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} conversion_rules_ocelot={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_definition_sdds': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} starting_offset=[0, 0, 0] subelement=False field_definition=None wakefield_definition=None k1l=-0.128366 n_kicks=4 field_reference_position='middle' fringe_field_coefficient=0.0 strength_errors=[0] gradient=0.0 scale_field=False multipoles=[0] smooth=2 centre=[0.0, 0.0, 3.52715] datum=[0.0, 0.0, 3.5775] global_parameters={'beam': {'filename': None, 'code': None}, 'GPTLICENSE': '', 'delete_tracking_files': False, 'astra_use_wsl': 1, 'master_subdir': '/home/xkc85723/Documents/simframe/examples/notebooks/getting_started', 'master_lattice_location': '/home/xkc85723/Documents/masterlattice/MasterLattice/./', 'simcodes_location': '/home/xkc85723/Documents/simcodes/SimCodes/./'} PV='CLA-S02-MAG-QUAD-01' isr_enable=True sr_enable=True type='quadrupole' bore=0.037\n" ] } ], "source": [ "print(framework.getElement(\"CLA-S02-MAG-QUAD-01\"))\n", "print(framework[\"CLA-S02-MAG-QUAD-01\"])" ] }, { "cell_type": "code", "execution_count": 5, "id": "fc044a1a-0ecc-4abb-9fcb-2e596926213d", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "-0.128366\n", "-0.128366\n" ] } ], "source": [ "print(framework.getElement(\"CLA-S02-MAG-QUAD-01\", \"k1l\"))\n", "print(framework[\"CLA-S02-MAG-QUAD-01\"].k1l)" ] }, { "cell_type": "markdown", "id": "f8074391-a6fc-4a9d-87ae-215bb8dceac8", "metadata": {}, "source": [ "All objects of a given type defined in the entire lattice can be accessed as follows.\n", "\n", "(NB here cavity on-crest phase is zero and ASTRA convention is followed for the phase sign)" ] }, { "cell_type": "code", "execution_count": 17, "id": "c7554ba3-8c96-46d9-9ddf-3abab3894d3c", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "{'CLA-HRG1-GUN-CAV-01': 9, 'CLA-L01-LIN-CAV-01': 16, 'CLA-L02-LIN-CAV-01': 23, 'CLA-L03-LIN-CAV-01': 8, 'CLA-L4H-LIN-CAV-01': -184, 'CLA-L04-LIN-CAV-01': -45}\n" ] } ], "source": [ "cavnames = [c[\"name\"] for c in framework.getElementType(\"cavity\")]\n", "print({k: v for k, v in zip(cavnames, framework.getElementType(\"cavity\", \"phase\"))})" ] }, { "cell_type": "markdown", "id": "2b33487a-a78b-4e6c-aac2-aeaf0316dff4", "metadata": {}, "source": [ "Various options are available for changing a lattice element, or multiple elements of a given type simultaneously." ] }, { "cell_type": "code", "execution_count": 20, "id": "02bb21f7-b70e-42db-ade7-fcb91a9ea7bb", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "{'CLA-VBC-MAG-DIP-01': 0.0, 'CLA-VBC-MAG-DIP-02': 0.0, 'CLA-VBC-MAG-DIP-03': 0.0, 'CLA-VBC-MAG-DIP-04': 0.0, 'CLA-SP3-MAG-DIP-01': 0.0}\n", "{'CLA-VBC-MAG-DIP-01': 0.1, 'CLA-VBC-MAG-DIP-02': 0.1, 'CLA-VBC-MAG-DIP-03': 0.1, 'CLA-VBC-MAG-DIP-04': 0.1, 'CLA-SP3-MAG-DIP-01': 0.1}\n" ] } ], "source": [ "# Set all dipole angles in the lattice to zero and use `setElementType`\n", "values = [0 for _ in framework.getElementType(\"dipole\")]\n", "dipnames = [d[\"name\"] for d in framework.getElementType(\"dipole\")]\n", "framework.setElementType(\"dipole\", \"angle\", values)\n", "print({k: v for k, v in zip(dipnames, framework.getElementType(\"dipole\", \"angle\"))})\n", "\n", "# Set all dipole angles to 0.1 and use modifyElementType\n", "value = 0.1\n", "framework.modifyElementType(\"dipole\", \"angle\", 0.1)\n", "print({k: v for k, v in zip(dipnames, framework.getElementType(\"dipole\", \"angle\"))})" ] }, { "cell_type": "code", "execution_count": 49, "id": "ee6d7c47-d443-45f4-b93a-0d191563c23c", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CLA-S02-MAG-QUAD-01 k1l = 1.0\n", "CLA-S02-MAG-QUAD-01 k1l = 2.0\n", "CLA-S02-MAG-QUAD-01 k1l = 3.0\n", "CLA-S02-MAG-QUAD-01 k1l = 3.0, n_kicks = 5\n", "CLA-S02-MAG-QUAD-01 k1l = 4.0, n_kicks = 6\n", "CLA-S02-MAG-QUAD-01 k1l = 5.0, n_kicks = 7, CLA-S02-MAG-QUAD-02 k1l = 5.0, n_kicks = 7\n" ] } ], "source": [ "s02q1 = \"CLA-S02-MAG-QUAD-01\"\n", "s02q2 = \"CLA-S02-MAG-QUAD-02\"\n", "setattr(framework.elementObjects[s02q1], 'k1l', 1)\n", "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}\")\n", "framework[s02q1].k1l = 2\n", "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}\")\n", "framework.modifyElement(s02q1, \"k1l\", 3)\n", "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}\")\n", "framework.modifyElement(s02q1, [\"k1l\", \"n_kicks\"], [3, 5])\n", "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}\")\n", "framework.modifyElement(s02q1, {\"k1l\": 4, \"n_kicks\": 6})\n", "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}\")\n", "framework.modifyElements([s02q1, s02q2], {\"k1l\": 5, \"n_kicks\": 7})\n", "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}, \\\n", "{s02q2} k1l = {framework.getElement(s02q2, 'k1l')}, n_kicks = {framework[s02q2].n_kicks})" ] }, { "cell_type": "code", "execution_count": 6, "id": "0221e2f2-cf64-491e-b04a-ea5b83b4af1e", "metadata": {}, "outputs": [], "source": [ "framework.saveParametersFile(\"params.yaml\", framework.getElementType(\"dipole\"))" ] }, { "cell_type": "code", "execution_count": null, "id": "bd6f5236-69a3-49c2-9ac3-692345207f2d", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "711f8980-2f15-436b-8b12-4d0b2f971f5a", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "ab02d5cd-d699-4da5-a8c7-a70a1017e38a", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "0135558a-8ad0-49b2-a95e-5681d8dc1a12", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "1184abda-df45-4895-ac97-6fb68f8a89bf", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "3dd64637-4261-4370-976e-1c68a2c1f2bc", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 5 }