<commit_before>
from syned.beamline.beamline import Beamline

from wofrysrw.srw_object import SRWObject
from wofrysrw.storage_ring.srw_light_source import SRWLightSource
from wofrysrw.propagator.wavefront2D.srw_wavefront import WavefrontPropagationParameters, WavefrontPropagationOptionalParameters

from srwlib import *

class Where:
    DRIFT_BEFORE = "before"
    DRIFT_AFTER = "after"
    OE = "oe"

    @classmethod
    def tuple(cls):
        return cls.DRIFT_BEFORE, cls.OE, cls.DRIFT_AFTER

class SRWBeamline(Beamline, SRWObject):

    def __init__(self,
                 light_source=SRWLightSource(),
                 beamline_elements_list=[]):
        super().__init__(light_source=light_source, beamline_elements_list=beamline_elements_list)

        self._wavefront_propagation_parameters_list = {}
        self._wavefront_propagation_parameters_list[Where.DRIFT_BEFORE] = []
        self._wavefront_propagation_parameters_list[Where.OE]           = []
        self._wavefront_propagation_parameters_list[Where.DRIFT_AFTER]  = []

    # overwrites the SynedObject method for dealing with list
    def to_dictionary(self): # TODO: to be adjusted....
        dict_to_save = super().to_dictionary()
        dict_to_save["srw_wavefront_propagation_parameters"] = [ el[0].to_dictionary() for el in self._wavefront_propagation_parameters_list[Where.OE]]

        return dict_to_save

    def append_wavefront_propagation_parameters(self, wavefront_propagation_parameters=WavefrontPropagationParameters(), wavefront_propagation_optional_parameters=WavefrontPropagationOptionalParameters(), where=Where.OE):
        self._wavefront_propagation_parameters_list[where].append([wavefront_propagation_parameters, wavefront_propagation_optional_parameters])

    def get_wavefront_propagation_parameters(self, where=Where.OE):
        return self._wavefront_propagation_parameters_list[where]

    def get_wavefront_propagation_parameters_at(self, index, where=Where.OE):
        if index >= len(self._wavefront_propagation_parameters_list[where]):
            raise IndexError("Index " + str(index) + " out of bounds")

        return self._wavefront_propagation_parameters_list[where][index]

    def duplicate(self):
        beamline_elements_list = []
        for beamline_element in self._beamline_elements_list:
            beamline_elements_list.append(beamline_element)

        new_beamline = SRWBeamline(light_source=self._light_source, beamline_elements_list = beamline_elements_list)

        for where in Where.tuple():
            for element in self.get_wavefront_propagation_parameters(where):
                new_beamline.append_wavefront_propagation_parameters(element[0], element[1], where)

        return new_beamline

    #TODO: to be completed
    def to_python_code(self, data=None):
        wavefront = data[0]
        is_multi_electron = data[1] == True

        text_code  = "from srwlib import *\nfrom uti_plot import *\nimport numpy\n\n"

        if is_multi_electron: text_code += "#if not srwl_uti_proc_is_master(): exit()\n"
        else: text_code += "if not srwl_uti_proc_is_master(): exit()\n"

        text_code += "\n####################################################\n# LIGHT SOURCE\n\n"
        text_code += self.get_light_source().to_python_code(is_multi_electron)

        if not is_multi_electron:
            text_code += "\n"
            text_code += "mesh0 = deepcopy(wfr.mesh)" + "\n"
            text_code += "arI = array('f', [0]*mesh0.nx*mesh0.ny)" + "\n"
            text_code += "srwl.CalcIntFromElecField(arI, wfr, 6, 0, 3, mesh0.eStart, 0, 0)" + "\n"
            text_code += "arIx = array('f', [0]*mesh0.nx)" + "\n"
            text_code += "srwl.CalcIntFromElecField(arIx, wfr, 6, 0, 1, mesh0.eStart, 0, 0)" + "\n"
            text_code += "arIy = array('f', [0]*mesh0.ny)" + "\n"
            text_code += "srwl.CalcIntFromElecField(arIy, wfr, 6, 0, 2, mesh0.eStart, 0, 0)" + "\n"
            text_code += "#save ascii file with intensity" + "\n"
            text_code += "#srwl_uti_save_intens_ascii(arI, mesh0, <file_path>)" + "\n"
            text_code += "plotMesh0x = [1000*mesh0.xStart, 1000*mesh0.xFin, mesh0.nx]" + "\n"
            text_code += "plotMesh0y = [1000*mesh0.yStart, 1000*mesh0.yFin, mesh0.ny]" + "\n"
            text_code += "uti_plot2d1d (arI, plotMesh0x, plotMesh0y, labels=['Horizontal Position [mm]', 'Vertical Position [mm]', 'Intensity Before Propagation'])" + "\n"

        #TODO: HERE CODE FOR PRPAGATION
        if self.get_beamline_elements_number() > 0:
            text_code += "\n####################################################\n# BEAMLINE\n\n"
            text_code += "srw_oe_array = []" + "\n"
            text_code += "srw_pp_array = []" + "\n"
            text_code += "\n"

            for index in range(self.get_beamline_elements_number()):
                oe_name = "oe_" + str(index)
                data_oe = [oe_name, wavefront]

                beamline_element = self.get_beamline_element_at(index)

                optical_element = beamline_element.get_optical_element()
                coordinates = beamline_element.get_coordinates()

                if coordinates.p() != 0.0:
                    text_code += "drift_before_" + oe_name + " = " + "SRWLOptD(" + str(coordinates.p()) + ")" + "\n"

                    wp, wop = self.get_wavefront_propagation_parameters_at(index, Where.DRIFT_BEFORE)
                    text_array = wp.to_python_code()
                    if not wop is None: text_array = wop.append_to_python_code(text_array)

                    text_code += "pp_drift_before_" + oe_name + " = " + text_array  + "\n"
                    text_code += "\n"
                    text_code += "srw_oe_array.append(drift_before_" + oe_name + ")" + "\n"
                    text_code += "srw_pp_array.append(pp_drift_before_" + oe_name + ")" + "\n"
                    text_code += "\n"

                wp, wop = self.get_wavefront_propagation_parameters_at(index, Where.OE)
<commit_msg>Added Specific Propagation parameters<commit_after>
from syned.beamline.beamline import Beamline

from wofrysrw.srw_object import SRWObject
from wofrysrw.storage_ring.srw_light_source import SRWLightSource
from wofrysrw.propagator.wavefront2D.srw_wavefront import WavefrontPropagationParameters, WavefrontPropagationOptionalParameters

from srwlib import *

class Where:
    DRIFT_BEFORE = "before"
    DRIFT_AFTER = "after"
    OE = "oe"

    @classmethod
    def tuple(cls):
        return cls.DRIFT_BEFORE, cls.OE, cls.DRIFT_AFTER

class SRWBeamline(Beamline, SRWObject):

    def __init__(self,
                 light_source=SRWLightSource(),
                 beamline_elements_list=[]):
        super().__init__(light_source=light_source, beamline_elements_list=beamline_elements_list)

        self._wavefront_propagation_parameters_list = {}
        self._wavefront_propagation_parameters_list[Where.DRIFT_BEFORE] = []
        self._wavefront_propagation_parameters_list[Where.OE]           = []
        self._wavefront_propagation_parameters_list[Where.DRIFT_AFTER]  = []

    # overwrites the SynedObject method for dealing with list
    def to_dictionary(self): # TODO: to be adjusted....
        dict_to_save = super().to_dictionary()
        dict_to_save["srw_wavefront_propagation_parameters"] = [ el[0].to_dictionary() for el in self._wavefront_propagation_parameters_list[Where.OE]]

        return dict_to_save

    def append_wavefront_propagation_parameters(self, wavefront_propagation_parameters=WavefrontPropagationParameters(), wavefront_propagation_optional_parameters=WavefrontPropagationOptionalParameters(), where=Where.OE):
        self._wavefront_propagation_parameters_list[where].append([wavefront_propagation_parameters, wavefront_propagation_optional_parameters])

    def get_wavefront_propagation_parameters(self, where=Where.OE):
        return self._wavefront_propagation_parameters_list[where]

    def get_wavefront_propagation_parameters_at(self, index, where=Where.OE):
        if index >= len(self._wavefront_propagation_parameters_list[where]):
            raise IndexError("Index " + str(index) + " out of bounds")

        return self._wavefront_propagation_parameters_list[where][index]

    def duplicate(self):
        beamline_elements_list = []
        for beamline_element in self._beamline_elements_list:
            beamline_elements_list.append(beamline_element)

        new_beamline = SRWBeamline(light_source=self._light_source, beamline_elements_list = beamline_elements_list)

        for where in Where.tuple():
            for element in self.get_wavefront_propagation_parameters(where):
                new_beamline.append_wavefront_propagation_parameters(element[0], element[1], where)

        return new_beamline

    def to_python_code(self, data=None):
        wavefront = data[0]
        is_multi_electron = data[1] == True

        text_code  = "from srwlib import *\nfrom uti_plot import *\nimport numpy\n\n"

        if is_multi_electron: text_code += "#if not srwl_uti_proc_is_master(): exit()\n"
        else: text_code += "if not srwl_uti_proc_is_master(): exit()\n"

        text_code += "\n####################################################\n# LIGHT SOURCE\n\n"
        text_code += self.get_light_source().to_python_code(is_multi_electron)

        if not is_multi_electron:
            text_code += "\n"
            text_code += "mesh0 = deepcopy(wfr.mesh)" + "\n"
            text_code += "arI = array('f', [0]*mesh0.nx*mesh0.ny)" + "\n"
            text_code += "srwl.CalcIntFromElecField(arI, wfr, 6, 0, 3, mesh0.eStart, 0, 0)" + "\n"
            text_code += "arIx = array('f', [0]*mesh0.nx)" + "\n"
            text_code += "srwl.CalcIntFromElecField(arIx, wfr, 6, 0, 1, mesh0.eStart, 0, 0)" + "\n"
            text_code += "arIy = array('f', [0]*mesh0.ny)" + "\n"
            text_code += "srwl.CalcIntFromElecField(arIy, wfr, 6, 0, 2, mesh0.eStart, 0, 0)" + "\n"
            text_code += "#save ascii file with intensity" + "\n"
            text_code += "#srwl_uti_save_intens_ascii(arI, mesh0, <file_path>)" + "\n"
            text_code += "plotMesh0x = [1000*mesh0.xStart, 1000*mesh0.xFin, mesh0.nx]" + "\n"
            text_code += "plotMesh0y = [1000*mesh0.yStart, 1000*mesh0.yFin, mesh0.ny]" + "\n"
            text_code += "uti_plot2d1d (arI, plotMesh0x, plotMesh0y, labels=['Horizontal Position [mm]', 'Vertical Position [mm]', 'Intensity Before Propagation'])" + "\n"

        if self.get_beamline_elements_number() > 0:
            text_code += "\n####################################################\n# BEAMLINE\n\n"
            text_code += "srw_oe_array = []" + "\n"
            text_code += "srw_pp_array = []" + "\n"
            text_code += "\n"

            for index in range(self.get_beamline_elements_number()):
                oe_name = "oe_" + str(index)
                data_oe = [oe_name, wavefront]

                beamline_element = self.get_beamline_element_at(index)

                optical_element = beamline_element.get_optical_element()
                coordinates = beamline_element.get_coordinates()

                if coordinates.p() != 0.0:
                    text_code += "drift_before_" + oe_name + " = " + "SRWLOptD(" + str(coordinates.p()) + ")" + "\n"

                    wp, wop = self.get_wavefront_propagation_parameters_at(index, Where.DRIFT_BEFORE)
                    text_array = wp.to_python_code()
                    if not wop is None: text_array = wop.append_to_python_code(text_array)

                    text_code += "pp_drift_before_" + oe_name + " = " + text_array  + "\n"
                    text_code += "\n"
                    text_code += "srw_oe_array.append(drift_before_" + oe_name + ")" + "\n"
                    text_code += "srw_pp_array.append(pp_drift_before_" + oe_name + ")" + "\n"
                    text_code += "\n"

                wp, wop = self.get_wavefront_propagation_parameters_at(index, Where.OE)

{'fraction_non_alphanumeric': 0.12949520233625364, 'fraction_numerical': 0.010680016687526075, 'mean_word_length': 3.064428619871143, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1537130', 'n_tokens_mistral': 4031, 'n_tokens_neox': 3799, 'n_words': 666}