Skip to content
Snippets Groups Projects
pyTransfocator_single.py 19.3 KiB
Newer Older
mwyman's avatar
mwyman committed
import tomllib
from transfocator_calcs import lookup_diameter, materials_to_deltas, materials_to_linear_attenuation
from transfocator_calcs import find_levels, calc_lookup_table, get_densities
mwyman's avatar
mwyman committed

MAT_MACRO = 'MAT'
NLENS_MACRO = 'NUMLENS'
RADIUS_MACRO = 'RADIUS'
LOC_MACRO = 'LOC'
THICKERR_MACRO = 'THICKERR'
        energy      : energy in keV
        L_und       : undulator length in m
        sigmaH_e    : Sigma electron source size in H direction in m
        sigmaV_e    : Sigma electron source size in V direction in m
        sigmaHp     : Sigma electron divergence in H direction in rad
        sigmaVp_e   : Sigma electron divergence in V direction in rad
        d_StoL1 : Source-to-CRL1 distance, in m
        d_min   : Minimum thickness at the apex in m
        stack_d : Stack thickness in m
DEFAULT_CONFIG = {'beam':{'energy': 15, 'L_und': 4.7, 'sigmaH_e': 14.8e-6,
                          'sigmaV_e': 3.7e-6, 'sigmaHp_e': 2.8e-6, 'sigmaVp_e': 1.5e-6},
                  'beamline': {'d_StoL1': 51.9, 'd_Stof': 66.2},
                  'crl':{'stack_d': 50.0e-3, 'd_min': 3.0e-5}}
    
    def __init__(self, crl_setup = None, beam_config = DEFAULT_CONFIG['beam'], 
                 beamline_config = DEFAULT_CONFIG['beamline'], 
mwyman's avatar
mwyman committed
                 crl_config = DEFAULT_CONFIG['crl']):
        if crl_setup is None:
            beam = beam_config
            beamline = beamline_config
            crl = crl_config
mwyman's avatar
mwyman committed
#            slits = slits_config
        else:
            with open(crl_setup, "rb") as f:
                config = tomllib.load(f)
            beam = config['beam']
            beamline = config['beamline']
            crl = config['crl']
mwyman's avatar
mwyman committed
 #           slits = config['slits'] 
mwyman's avatar
mwyman committed
        self.beam = {}
mwyman's avatar
mwyman committed
        self.bl = {}
mwyman's avatar
mwyman committed
        self.crl = {}
        # Initialize lens variables -- TODO -- this is done via a subs file -- are any of these needed prior to that loading?
        self.numlens        = np.array([1,      1,      1,      1,      1,      1,      2,      4,      8,      16])                # CRL1 number of lenses in each stack (was L1_n)
        self.radius         = np.array([2.0e-3, 1.0e-3, 5.0e-4, 3.0e-4, 2.0e-4, 1.0e-4, 1.0e-4, 1.0e-4, 1.0e-4, 1.0e-4])            # CRL1 lens radius in each stack (was L1_R)
        self.materials      = np.array(["Be",   "Be",   "Be",   "Be",   "Be",   "Be",   "Be",   "Be",   "Be",   "Be"])              # CRL1 lens material in each stack (was L1_mater)
mwyman's avatar
mwyman committed
        self.lens_loc       = np.array([4.5,    3.5,    2.5,    1.5,    0.5,    -0.5,   -1.5,   -2.5,   -3.5,   -4.5])*self.crl['stack_d']      # CRL1 lens stack location relative to center stack, positive means upstream (was L1_Loc)
        self.lens_thickerr  = np.array([1.0e-6, 1.0e-6, 1.0e-6, 1.0e-6, 1.0e-6, 1.0e-6, 1.4e-6, 2.0e-6, 2.8e-6, 4.0e-6])            # CRL1 lens RMS thickness error (was L1_HE)


        self.Lens_diameter_table = [
                                    (50, 450.0),
                                    (100, 632.0),
                                    (200, 894.0),
                                    (300, 1095.0),
                                    (500, 1414.0),
                                    (1000, 2000.0),
                                    (1500, 2450.0),
                                ]
        # Convert the lookup table to a dictionary for faster lookup        
mwyman's avatar
mwyman committed
        self.Lens_diameter_dict = {int(col1): col2 for col1, col2 in self.Lens_diameter_table}
        
        self.energy = 0  # gets value from an ao (incoming beam energy)
        self.focalSize = 0 # get value from an ao (desired focal length)
mwyman's avatar
mwyman committed
        self.lenses = 0 # sets integer (2^10) whose binary representation indicates which lenses are in or out
        self.num_stacks = 10 # Number of lenses in system
        self.thickerr_flag = True
                
    def setupSource(self, beam_properties):
        '''
        Beam properties can have entries for the following
        
        energy      : energy in keV
        L_und       : undulator length in m
        sigmaH_e    : Sigma electron source size in H direction in m
        sigmaV_e    : Sigma electron source size in V direction in m
mwyman's avatar
mwyman committed
        sigmaHp_e   : Sigma electron divergence in H direction in rad
        sigmaVp_e   : Sigma electron divergence in V direction in rad
        '''
        
mwyman's avatar
mwyman committed
        self.setEnergy(beam_properties['energy'])
        self.L_und = beam_properties['L_und']
        self.sigmaH_e = beam_properties['sigmaH_e']
        self.sigmaV_e = beam_properties['sigmaV_e']
        self.sigmaHp_e = beam_properties['sigmaHp_e']
        self.sigmaVp_e = beam_properties['sigmaVp_e']
mwyman's avatar
mwyman committed
        self.setupSourceEnergyDependent()
    
    def setupSourceEnergyDependent(self):
        '''
        Fill in later
        '''    
        self.beam['sigmaH'] =  (self.sigmaH_e**2 +  self.wl*self.L_und/2/np.pi/np.pi)**0.5
        self.beam['sigmaV'] =  (self.sigmaV_e**2 +  self.wl*self.L_und/2/np.pi/np.pi)**0.5
        self.beam['sigmaHp'] = (self.sigmaHp_e**2 + self.wl/self.L_und/2)**0.5
        self.beam['sigmaVp'] = (self.sigmaVp_e**2 + self.wl/self.L_und/2)**0.5
        
    def setupBeamline(self, beamline_properties, num=1):
        '''
        Beamline properties can contain entries for the following
        
        d_StoL1 : Source-to-CRL1 distance, in m
        self.bl['d_StoL1'] = beamline_properties['d_StoL1']
mwyman's avatar
mwyman committed
        self.bl['d_Stof'] = beamline_properties['d_Stof']
            
    def setupCRL(self, crl_properties):
        '''
        CRL properties can contiain entries for the following
        
        d_min   : Minimum thickness at the apex in m
        stack_d : Stack thickness in m
        '''
        self.crl['d_min'] = crl_properties['d_min'] 
        self.crl['stack_d'] = crl_properties['stack_d']

    def setupSlits(self, slit_properties):
        '''
        Slit properties can contain entries for the following
    def setupLookupTable(self, subs_file, n_lenses):
mwyman's avatar
mwyman committed
        '''
        lookup table created after IOC startup
        energy and slit size are updated before this is called
mwyman's avatar
mwyman committed
        '''
        print(80*'#')
        print('Setting up lens control...')
        
        self.num_stacks = n_lenses
mwyman's avatar
mwyman committed
        self.num_configs = 2**self.num_stacks
mwyman's avatar
mwyman committed
        self.configs = np.arange(self.num_configs)
mwyman's avatar
mwyman committed
        
#       self.energy = energy
mwyman's avatar
mwyman committed
        
        #read in substitutions file
        try:
            subsFile = open(subs_file,"r")
        except:
            raise RuntimeError(f"Substiution file ({subsFile}) not found.")
        subsFileContent = subsFile.readlines()
        subsFile.close()
        
        macros = subsFileContent[2].replace('{','').replace('}','').replace(',','').split()
        lens_properties = {key: [] for key in macros} # dictionary of lists
        for i in range(self.num_stacks):
mwyman's avatar
mwyman committed
            try:
                xx = subsFileContent[3+i].replace('{','').replace('}','').replace(',','').replace('"','').split()
                lens_properties[macros[0]].append(xx[0])
                lens_properties[macros[1]].append(xx[1])
                lens_properties[macros[2]].append(xx[2])
                lens_properties[macros[3]].append(xx[3])
                lens_properties[macros[4]].append(xx[4])
                lens_properties[macros[5]].append(xx[5])
                lens_properties[macros[6]].append(xx[6])
            except:
                raise RuntimeError(f"Number of lenses ({self.num_stacks}) doesn't match substitution file")
mwyman's avatar
mwyman committed
        
        self.numlens = []
        self.radius = []
        self.materials = []
        self.lens_loc = []
        self.lens_thickerr = []
            
        # get number of lens for each lens stack from lens properties dictionary-list
mwyman's avatar
mwyman committed
        print('Getting lens materials...')
        if NLENS_MACRO in macros:
mwyman's avatar
mwyman committed
            self.numlens = np.array([int(i) for i in lens_properties[NLENS_MACRO]])
mwyman's avatar
mwyman committed
            print('Number of lens read in.\n')
        else:
            raise RuntimeError(f"Number of lenses macro ({NLENS_MACRO}) not found in substituion file")
        # get radii for each lens from lens properties dictionary-list
mwyman's avatar
mwyman committed
        print('Getting lens\' radii...')
mwyman's avatar
mwyman committed
        if RADIUS_MACRO in macros:
            self.radius = np.array([float(i) for i in lens_properties[RADIUS_MACRO]])
mwyman's avatar
mwyman committed
            print('Radius of lenses read in.\n')
        else:
mwyman's avatar
mwyman committed
            raise RuntimeError(f"Radius macro ({RADIUS_MACRO}) not found in substituion file")
mwyman's avatar
mwyman committed
        # get materials from lens properties dictionary-list
        print('Getting lens materials...')
        if MAT_MACRO in macros:
            self.materials = lens_properties[MAT_MACRO]
            print('Lens material read in.\n')
        else:
            raise RuntimeError(f"Material macro ({MAT_MACRO}) not found in substituion file")
        
        # get densities from local definition (for compounds) or from xraylib (for elements)
        densities = get_densities(self.materials)
mwyman's avatar
mwyman committed
        self.densities = np.array([densities[material] for material in self.materials])
mwyman's avatar
mwyman committed

        # get location of each lens from lens properties dictionary-list
mwyman's avatar
mwyman committed
        print('Getting lens\' locations...')
        if LOC_MACRO in macros:
mwyman's avatar
mwyman committed
            self.lens_loc = np.array([float(i)*self.crl['stack_d'] for i in lens_properties[LOC_MACRO]])
mwyman's avatar
mwyman committed
            print('Location of lenses read in.\n')
        else:
mwyman's avatar
mwyman committed
            raise RuntimeError(f"Location macro ({LOC_MACRO}) not found in substituion file")
mwyman's avatar
mwyman committed

        # get thicknesses errprfrom lens properties dictionary-list
        print('Getting lens thickness error...')
mwyman's avatar
mwyman committed
        if THICKERR_MACRO in macros:
            self.lens_thickerr = np.array([float(i) for i in lens_properties[THICKERR_MACRO]])
mwyman's avatar
mwyman committed
            print('Lens thickness errors read in.\n')
        else:
mwyman's avatar
mwyman committed
            raise RuntimeError(f"Thickness errors macro ({THICKERR_MACRO}) not found in substituion file")
mwyman's avatar
mwyman committed

        print('Constructing lookup table...')
        self.construct_lookup_table()
        print('Lookup table calculation complete.\n')
        
        print('Transfocator control setup complete.')
mwyman's avatar
mwyman committed
        print(80*'#')

    def construct_lookup_table(self):
        arr_a, arr_b, arr_c = calc_lookup_table(self.num_configs, self.radius, 
                                                self.materials, self.energy, self.wl,
                                                self.numlens, 
                                                self.lens_loc, self.beam, self.bl,
                                                self.crl, self.slits[0]['H'], self.slits[0]['V'],
                                                self.lens_thickerr, flag_HE = self.thickerr_flag,
                                                verbose = self.verbose)
        self.lookupTable = arr_a
        self.sorted_invF_index = arr_b
        self.sorted_invF = arr_c                                                            
                                                                    
#        self.sort_lookup_table()
        self.updateEnergyRBV()
mwyman's avatar
mwyman committed
        self.updateSlitSizeRBV('hor')
        self.updateSlitSizeRBV('vert')
        self.updateLookupWaveform()
mwyman's avatar
mwyman committed
        self.updateLookupConfigs()
mwyman's avatar
mwyman committed

        
#    def sort_lookup_table(self):
#        '''
#        
#        '''
#        if self.verbose: print(f'Sorting lookup table of length {len(self.lookupTable)}')        
#        self.sorted_L_index = np.argsort(self.lookupTable)
mwyman's avatar
mwyman committed

    def updateConfig(self, config_BW):
        '''
        When user manually changes lenses, this gets focal size and displays it
        along with updated RBVs but it doesn't set the config PV
        '''
        # Find the configuration in the 1/f sorted list
        self.indexSorted = self.sorted_invF_index.tolist().index(self.index)

        self.setFocalSizeActual()
        self.updateLensRBV()
mwyman's avatar
mwyman committed

    def setFocalSizeActual(self):
        '''
        
        '''
mwyman's avatar
mwyman committed
#        self.focalSize_actual = self.culledTable[self.culledIndex] 
        self.focalSize_actual = self.lookupTable[self.indexSorted] 
mwyman's avatar
mwyman committed
        User selected focal size, this function finds nearest acheivable focal 
        size from the lookup table
        '''
        # Code to search lookup table for nearest focal size to desired; note the
        # lookup table is already sorted by 1/f
mwyman's avatar
mwyman committed
        if self.verbose: print(f'Searching for config closest to {self.focalSize}')
#       simple approach
#        self.indexSorted = np.argmin(np.abs(self.lookupTable - self.focalSize))

        # XS approach -- can handle nan but in pydev application don't have a good
        # way to "transmit" errors (i.e. no solution found) to user.
        indices, _ = find_levels(self.lookupTable, self.focalSize, direction='forward')[0]
        self.indexSorted = indices[0]
        

        if self.verbose: print(f'1/f-sorted config index found at {self.indexSorted}')
mwyman's avatar
mwyman committed

        self.index = self.sorted_invF_index[self.indexSorted]
        if self.verbose: print(f'Config index found at {self.index}')
mwyman's avatar
mwyman committed

        # Update PVs
        self.setFocalSizeActual()
        self.updateLensConfigPV()
        self.updateLensRBV()
mwyman's avatar
mwyman committed
    def getPreviewFocalSize(self, sortedIndex):
        '''
        
        '''
        fSize_preview = self.lookupTable[sortedIndex]
mwyman's avatar
mwyman committed
        if self.verbose: print(f'Preview focal sizes for {sortedIndex} is {fSize_preview}')
        pydev.iointr('new_preview', fSize_preview)
    
    def setSlitSize(self, size, slit):
        '''
        Update proper slit size
        '''
        
mwyman's avatar
mwyman committed
        if slit == 'hor':
#            self.slit1_H = float(size)     # H slit size before CRL 1
            self.slits[0]['H'] = float(size)     # H slit size before CRL 1
#            self.slit1_V = float(size)     # V slit size before CRL 1
            self.slits[0]['V'] = float(size)     # V slit size before CRL 1
mwyman's avatar
mwyman committed
        else: 
            raise RuntimeError(f"Slit identifier ({slit}) not recognized. Should be 'hor' or 'vert'")
           
                    
    def updateSlitSize(self, size, slit):
        '''
        Slit size updates are propagated to CRL object from EPICS.  The beam
        size lookup table is then recalculated.
        '''
mwyman's avatar
mwyman committed
        self.setSlitSize(size, slit)
mwyman's avatar
mwyman committed
        if self.verbose: 
            if slit == 'hor':
#                print(f'Horizontal slit size updated to {self.slit1_H} m')
                print(f"Horizontal slit size updated to {self.slits[0]['H']} m")
mwyman's avatar
mwyman committed
            elif slit == 'vert':
#                print(f'Vertical slit size updated to {self.slit1_V} m')
                print(f"Vertical slit size updated to {self.slits[0]['V']} m")
    def setEnergy(self, energy):
        '''
        Sets various forms of energy
        '''
        if energy > 0.0001:
            self.energy = float(energy)
            self.energy_eV = self.energy*1000.0  # Energy in keV
            self.wl = 1239.84 / (self.energy_eV * 10**9)    #Wavelength in nm(?)
            if self.verbose: print(f'Setting energy to {self.energy} keV')
            
    def updateE(self, energy):
        '''
        Beam energy updates are propagated to CRL object from EPICS. The beam
        size lookup table is then recalculated.
        '''
        if energy > 0.0001:
            # Energy variable sent from IOC as a string
            self.setEnergy(energy)
            # Update beam properties that are dependent on energy
            self.setupSourceEnergyDependent()
        else:
            if verbose: print(f'Invalid energy setting: {energy} kev; staying at {self.energy} keV')
            
    def updateFsize(self, focalSize):
        '''
        User updates desired focal size. Lookup table is traversed to find nearest
        to desired.
        '''
        # focalPoint variable sent from IOC as a string
        self.focalSize = float(focalSize)
        self.find_config()
        
    def updateIndex(self, sortedIndex):
        '''
        User has updated desired sorted index
        '''
mwyman's avatar
mwyman committed
        self.indexSorted = int(sortedIndex)
        self.index = self.sorted_invF_index[self.indexSorted]

        # Update PVs
        self.setFocalSizeActual()
        self.updateLensConfigPV()
        self.updateLensRBV()
        self.updateFocalSizeRBVs()    

    def setThickerrFlag(self, flag):
mwyman's avatar
mwyman committed
        '''
        User has updated thickness error flag so that ...
        '''
        self.thickerr_flag = int(flag)
        if self.verbose: print(f'Thickness Error Flag set to {flag}')
        self.updateThickerrFlagRBV()

    def updateThickerrFlagRBV(self):
        '''
        Thickness error flag has been updated
        '''
        if self.verbose: print(f'Thickness Error Flag RBV set to {self.thickerr_flag}')
        pydev.iointr('updated_thickerr_Flag', self.thickerr_flag)
mwyman's avatar
mwyman committed
    def updateLensConfigPV(self):
        '''
        
        '''
mwyman's avatar
mwyman committed
        self.config = self.configs[self.index]
mwyman's avatar
mwyman committed
        pydev.iointr('new_lenses', int(self.config))

    def updateLensRBV(self):
        '''
        
        '''
mwyman's avatar
mwyman committed
        pydev.iointr('new_index', int(self.indexSorted))
mwyman's avatar
mwyman committed
        pydev.iointr('updated_E', float(self.energy))
mwyman's avatar
mwyman committed
    def updateSlitSizeRBV(self, slit):
mwyman's avatar
mwyman committed
        if slit == 'hor':
#            pydev.iointr('updated_slitSize_H', float(self.slit1_H))
#            if self.verbose: print(f'Horizontal slit size RBV updated to {self.slit1_H} m')
            pydev.iointr('updated_slitSize_H', float(self.slits[0]['H']))
            if self.verbose: print(f"Horizontal slit size RBV updated to {self.slits[0]['H']} m")
#            pydev.iointr('updated_slitSize_V', float(self.slit1_V))
#            if self.verbose: print(f'Vertical slit size RBV updated to {self.slit1_V} m')
            pydev.iointr('updated_slitSize_V', float(self.slits[0]['V']))
            if self.verbose: print(f"Vertical slit size RBV updated to {self.slits[0]['V']} m")
mwyman's avatar
mwyman committed
        
    def updateFocalSizeRBVs(self):
        '''
        
        '''
        pydev.iointr('new_fSize', self.focalSize_actual)
mwyman's avatar
mwyman committed
    def updateVerbosity(self, verbosity):
        '''
        Turn on minor printing
        '''
        print(f'Verbosity set to {verbosity}')
mwyman's avatar
mwyman committed

    def updateLookupWaveform(self):
        '''
mwyman's avatar
mwyman committed
        Puts lookup table focal sizes into waveform PV
mwyman's avatar
mwyman committed
        '''
        pydev.iointr('new_lookupTable', self.lookupTable.tolist())
        pydev.iointr('new_invFind_list', self.sorted_invF_index.tolist())
        pydev.iointr('new_invF_list', self.sorted_invF.tolist())
mwyman's avatar
mwyman committed
    
    def updateLookupConfigs(self):
        '''
mwyman's avatar
mwyman committed
        Puts lookup table config integers into waveform PV
mwyman's avatar
mwyman committed
        '''
mwyman's avatar
mwyman committed
        pydev.iointr('new_configs', self.configs.tolist())