<commit_before>import numpy

from syned.beamline.optical_elements.mirrors.mirror import Mirror
from syned.beamline.shape import Rectangle, Ellipse, Circle

from wofrysrw.beamline.optical_elements.srw_optical_element import SRWOpticalElement, Orientation
from wofrysrw.propagator.wavefront2D.srw_wavefront import WavefrontPropagationParameters

from srwlib import SRWLOptC, SRWLOptMir
from srwlib import srwl, srwl_opt_setup_surf_height_1d, srwl_opt_setup_surf_height_2d, srwl_uti_read_data_cols


class ApertureShape:
    RECTANGULAR = 'r'
    ELLIPTIC = 'e'

class SimulationMethod:
    THIN = 1
    THICK = 2

class TreatInputOutput:
    WAVEFRONT_INPUT_PLANE_BEFORE_OUTPUT_PLANE_AFTER  = 0
    WAVEFRONT_INPUT_CENTER_OUTPUT_CENTER  = 1
    WAVEFRONT_INPUT_CENTER_OUTPUT_CENTER_DRIFT_BACK  = 2

class ScaleType:
    LINEAR = 'lin'
    LOGARITHMIC = 'log'

class SRWMirror(Mirror, SRWOpticalElement):
    def __init__(self,
                 name                               = "Undefined",
                 tangential_size                    = 1.2,
                 sagittal_size                      = 0.01,
                 vertical_position_of_mirror_center = 0.0,
                 horizontal_position_of_mirror_center = 0.0,
                 grazing_angle                      = 0.003,
                 orientation_of_reflection_plane    = Orientation.UP,
                 invert_tangent_component           = False,
                 height_profile_data_file           = "mirror.dat",
                 height_profile_data_file_dimension = 1,
                 height_amplification_coefficient   = 1.0):


        Mirror.__init__(self,
                        name=name,
                        boundary_shape=Rectangle(x_left=horizontal_position_of_mirror_center - 0.5*sagittal_size,
                                                 x_right=horizontal_position_of_mirror_center + 0.5*sagittal_size,
                                                 y_bottom=vertical_position_of_mirror_center - 0.5*tangential_size,
                                                 y_top=vertical_position_of_mirror_center + 0.5*tangential_size),
                        surface_shape=self.get_shape())

        self.tangential_size                                  = tangential_size
        self.sagittal_size                                    = sagittal_size
        self.grazing_angle                                    = grazing_angle
        self.orientation_of_reflection_plane                  = orientation_of_reflection_plane
        self.invert_tangent_component                         = invert_tangent_component

        self.height_profile_data_file = height_profile_data_file
        self.height_profile_data_file_dimension = height_profile_data_file_dimension
        self.height_amplification_coefficient = height_amplification_coefficient

    def get_shape(self):
        raise NotImplementedError()

    def applyOpticalElement(self, wavefront=None, parameters=None):
        optical_elements = [self.toSRWLOpt()]
        propagation_parameters = [self.get_srw_wavefront_propagation_parameter(parameters)]

        if not self.height_profile_data_file is None:
            optical_elements.append(self.get_optTrEr())
            propagation_parameters.append(WavefrontPropagationParameters().to_SRW_array())

        optBL = SRWLOptC(optical_elements, propagation_parameters)

        srwl.PropagElecField(wavefront, optBL)

        return wavefront

    def add_to_srw_native_array(self, oe_array = [], pp_array=[], parameters=None, wavefront=None):
        super(SRWMirror, self).add_to_srw_native_array(oe_array, pp_array, parameters)

        if not self.height_profile_data_file is None:
            oe_array.append(self.get_optTrEr())
            pp_array.append(WavefrontPropagationParameters().to_SRW_array())

    def get_optTrEr(self):
        if self.orientation_of_reflection_plane == Orientation.LEFT or self.orientation_of_reflection_plane == Orientation.RIGHT:
            dim = 'x'
        elif self.orientation_of_reflection_plane == Orientation.UP or self.orientation_of_reflection_plane == Orientation.DOWN:
            dim = 'y'

        if self.height_profile_data_file_dimension == 1:
            height_profile_data = srwl_uti_read_data_cols(self.height_profile_data_file,
                                                          _str_sep='\t',
                                                          _i_col_start=0,
                                                          _i_col_end=1)

            optTrEr = srwl_opt_setup_surf_height_1d(_height_prof_data=height_profile_data,
                                                    _ang=self.grazing_angle,
                                                    _dim=dim,
                                                    _amp_coef=self.height_amplification_coefficient)
        elif self.height_profile_data_file_dimension == 2:
            height_profile_data = srwl_uti_read_data_cols(self.height_profile_data_file,
                                                          _str_sep='\t')

            optTrEr = srwl_opt_setup_surf_height_2d(_height_prof_data=height_profile_data,
                                                    _ang=self.grazing_angle,
                                                    _dim=dim,
                                                    _amp_coef=self.height_amplification_coefficient)
        return optTrEr

    def toSRWLOpt(self):
        nvx, nvy, nvz, tvx, tvy = self.get_orientation_vectors()
        x, y = self.getXY()

        if isinstance(self.get_boundary_shape(), Rectangle):
            ap_shape = ApertureShape.RECTANGULAR
        elif isinstance(self.get_boundary_shape(), Ellipse) or isinstance(self.get_boundary_shape(), Circle):
            ap_shape = ApertureShape.ELLIPTIC

        return self.get_SRWLOptMir(nvx, nvy, nvz, tvx, tvy, x, y, ap_shape)

    def get_SRWLOptMir(self, nvx, nvy, nvz, tvx, tvy, x, y, ap_shape):
        mirror = SRWLOptMir()

        mirror.set_dim_sim_meth(_size_tang=self.tangential_size,
                                _size_sag=self.sagittal_size,
                                _ap_shape=ap_shape,
                                _sim_meth=SimulationMethod.THICK,
                                _treat_in_out=TreatInputOutput.WAVEFRONT_INPUT_CENTER_OUTPUT_CENTER)
        mirror.set_orient(_nvx=nvx,
                          _nvy=nvy,
                          _nvz=nvz,
                          _tvx=tvx,
                          _tvy=tvy,
                          _x = x,<commit_msg>Added Reflectivity to SRW<commit_after>import numpy

from syned.beamline.optical_elements.mirrors.mirror import Mirror
from syned.beamline.shape import Rectangle, Ellipse, Circle

from wofrysrw.beamline.optical_elements.srw_optical_element import SRWOpticalElement, Orientation
from wofrysrw.propagator.wavefront2D.srw_wavefront import WavefrontPropagationParameters

from srwlib import SRWLOptC, SRWLOptMir
from srwlib import srwl, srwl_opt_setup_surf_height_1d, srwl_opt_setup_surf_height_2d, srwl_uti_read_data_cols


class ApertureShape:
    RECTANGULAR = 'r'
    ELLIPTIC = 'e'

class SimulationMethod:
    THIN = 1
    THICK = 2

class TreatInputOutput:
    WAVEFRONT_INPUT_PLANE_BEFORE_OUTPUT_PLANE_AFTER  = 0
    WAVEFRONT_INPUT_CENTER_OUTPUT_CENTER  = 1
    WAVEFRONT_INPUT_CENTER_OUTPUT_CENTER_DRIFT_BACK  = 2

class ScaleType:
    LINEAR = 'lin'
    LOGARITHMIC = 'log'

class SRWMirror(Mirror, SRWOpticalElement):
    def __init__(self,
                 name                               = "Undefined",
                 tangential_size                    = 1.2,
                 sagittal_size                      = 0.01,
                 vertical_position_of_mirror_center = 0.0,
                 horizontal_position_of_mirror_center = 0.0,
                 grazing_angle                      = 0.003,
                 orientation_of_reflection_plane    = Orientation.UP,
                 invert_tangent_component           = False,
                 height_profile_data_file           = "mirror.dat",
                 height_profile_data_file_dimension = 1,
                 height_amplification_coefficient   = 1.0):


        Mirror.__init__(self,
                        name=name,
                        boundary_shape=Rectangle(x_left=horizontal_position_of_mirror_center - 0.5*sagittal_size,
                                                 x_right=horizontal_position_of_mirror_center + 0.5*sagittal_size,
                                                 y_bottom=vertical_position_of_mirror_center - 0.5*tangential_size,
                                                 y_top=vertical_position_of_mirror_center + 0.5*tangential_size),
                        surface_shape=self.get_shape())

        self.tangential_size                                  = tangential_size
        self.sagittal_size                                    = sagittal_size
        self.grazing_angle                                    = grazing_angle
        self.orientation_of_reflection_plane                  = orientation_of_reflection_plane
        self.invert_tangent_component                         = invert_tangent_component

        self.height_profile_data_file = height_profile_data_file
        self.height_profile_data_file_dimension = height_profile_data_file_dimension
        self.height_amplification_coefficient = height_amplification_coefficient

    # this method remembers the original SRW object
    def set_reflectivity(self,
                         reflectivity_data=0.95,
                         energies_number=1,
                         angles_number=1,
                         components_number=1,
                         energy_start=0.0,
                         energy_end=0.0,
                         energy_scale_type=ScaleType.LINEAR,
                         angle_start=0.0,
                         angle_end=0.0,
                         angle_scale_type=ScaleType.LINEAR):
        self.reflectivity_data = reflectivity_data
        self.energies_number   = energies_number
        self.angles_number     = angles_number
        self.components_number = components_number
        self.energy_start      = energy_start
        self.energy_end        = energy_end
        self.energy_scale_type = energy_scale_type
        self.angle_start       = angle_start
        self.angle_end         = angle_end
        self.angle_scale_type  = angle_scale_type

    def get_shape(self):
        raise NotImplementedError()

    def applyOpticalElement(self, wavefront=None, parameters=None):
        optical_elements = [self.toSRWLOpt()]
        propagation_parameters = [self.get_srw_wavefront_propagation_parameter(parameters)]

        if not self.height_profile_data_file is None:
            optical_elements.append(self.get_optTrEr())
            propagation_parameters.append(WavefrontPropagationParameters().to_SRW_array())

        optBL = SRWLOptC(optical_elements, propagation_parameters)

        srwl.PropagElecField(wavefront, optBL)

        return wavefront

    def add_to_srw_native_array(self, oe_array = [], pp_array=[], parameters=None, wavefront=None):
        super(SRWMirror, self).add_to_srw_native_array(oe_array, pp_array, parameters)

        if not self.height_profile_data_file is None:
            oe_array.append(self.get_optTrEr())
            pp_array.append(WavefrontPropagationParameters().to_SRW_array())

    def get_optTrEr(self):
        if self.orientation_of_reflection_plane == Orientation.LEFT or self.orientation_of_reflection_plane == Orientation.RIGHT:
            dim = 'x'
        elif self.orientation_of_reflection_plane == Orientation.UP or self.orientation_of_reflection_plane == Orientation.DOWN:
            dim = 'y'

        if self.height_profile_data_file_dimension == 1:
            height_profile_data = srwl_uti_read_data_cols(self.height_profile_data_file,
                                                          _str_sep='\t',
                                                          _i_col_start=0,
                                                          _i_col_end=1)

            optTrEr = srwl_opt_setup_surf_height_1d(_height_prof_data=height_profile_data,
                                                    _ang=self.grazing_angle,
                                                    _dim=dim,
                                                    _amp_coef=self.height_amplification_coefficient)
        elif self.height_profile_data_file_dimension == 2:
            height_profile_data = srwl_uti_read_data_cols(self.height_profile_data_file,
                                                          _str_sep='\t')

            optTrEr = srwl_opt_setup_surf_height_2d(_height_prof_data=height_profile_data,
                                                    _ang=self.grazing_angle,
                                                    _dim=dim,
                                                    _amp_coef=self.height_amplification_coefficient)
        return optTrEr

    def toSRWLOpt(self):
        nvx, nvy, nvz, tvx, tvy = self.get_orientation_vectors()
        x, y = self.getXY()

        if isinstance(self.get_boundary_shape(), Rectangle):
            ap_shape = ApertureShape.RECTANGULAR
        elif isinstance(self.get_boundary_shape(), Ellipse) or isinstance(self.get_boundary_shape(), Circle):
            ap_shape = ApertureShape.ELLIPTIC

        mirror = self.get_SRWLOptMir(nvx, nvy, nvz, tvx, tvy, x, y, ap_shape)

        if hasattr(self, "reflectivity_data"):
            mirror.set_reflect(_refl=self.reflectivity_data,
                               _n_ph_en=self.energies_number,
                               _n_ang=self.angles_number,
                               _n_comp=self.components_number,
                               _ph_en_start=self.energy_start,
                               _ph_en_fin=self.energy_end,
                               _ph_en_scale_type=self.energy_scale_type,
                               _ang_start=self.angle_start,
                               _ang_fin=self.angle_end,
                               _ang_scale_type=self.angle_scale_type)

        return mirror


    def get_SRWLOptMir(self, nvx, nvy, nvz, tvx, tvy, x, y, ap_shape):
        mirror = SRWLOptMir()

        mirror.set_dim_sim_meth(_size_tang=self.tangential_size,
                                _size_sag=self.sagittal_size,
                                _ap_shape=ap_shape,
                                _sim_meth=SimulationMethod.THICK,
                                _treat_in_out=TreatInputOutput.WAVEFRONT_INPUT_CENTER_OUTPUT_CENTER)
        mirror.set_orient(_nvx=nvx,
                          _nvy=nvy,
                          _nvz=nvz,
                          _tvx=tvx,
                          _tvy=tvy,
                          _x = x,
{'fraction_non_alphanumeric': 0.05915966386554622, 'fraction_numerical': 0.006050420168067227, 'mean_word_length': 1.7145985401459853, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 62, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1214466', 'n_tokens_mistral': 4337, 'n_tokens_neox': 4085, 'n_words': 570}