pyTransfocator_single.py

import numpy as np
import tomllib
import xraylib
from transfocator_calcs import lookup_diameter, materials_to_deltas, materials_to_linear_attenuation, calc_lookup_table, get_densities

MAT_MACRO = 'MAT'
NLENS_MACRO = 'NUMLENS'
RADIUS_MACRO = 'RADIUS'
LOC_MACRO = 'LOC'
THICKERR_MACRO = 'THICKERR'

'''
Config variables

Beam Properties
        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
Beamline properties
        d_StoL1 : Source-to-CRL1 distance, in m
        d_Stof  : Source-to-focus distance, in m
CRL properties
        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_StoL': 51.9, 'd_Stof': 66.2},
                  'crl':{'stack_d': 50.0e-3, 'd_min': 3.0e-5}}
"""
pyDevice TO DO: 

WHAT inputs change the focal size arrays? Energy, what else?
WHAT inputs change the search through the arrays? desired focal size, what else?

IOC init functions
-get lens stack parameters (# of lenses in each stack, radius, location, thickness, thickness error) -- from substitution file but put into PVs? Update with autosave?
-get source info
    -energy from from ID IOC
    -hor/vert sizes and divergence (also energy dependent)
-lens diameter table? What is it doing?

-desired focal size is changed --> what needs updating? --> nothing, just need to search focal size array again
    -multiple flags: is focal size achievable? is it achievable at sample?

recalc function -- should probably be same as init function
-energy is updated --> what needs updating?

-what else could user/staff change? sample position?
"""


'''
Update the following to accommodate XS code
'''

class singleTF():
    
    def __init__(self, crl_setup = None, beam_config = DEFAULT_CONFIG['beam'], 
                 beamline_config = DEFAULT_CONFIG['beamline'], 
                 crl_config = DEFAULT_CONFIG['crl']):
        if crl_setup is None:
            beam = beam_config
            beamline = beamline_config
            crl = crl_config
#            slits = slits_config
        else:
            with open(crl_setup, "rb") as f:
                config = tomllib.load(f)
            beam = config['beam']
            beamline = config['beamline']
            crl = config['crl']
 #           slits = config['slits'] 

        self.beam = {}
        self.setupSource(beam)
        
        self.bl = {}
        self.setupBeamline(beamline)
        
        self.crl = {}
        self.setupCRL(crl)
        
        # TODO is setupSlits necessary?
        #self.setupSlits(slits)
        
        # 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)
        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)


        #<----- lens diameter stuff done in transfocator calcs, should this be removed from here?
        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        
        self.Lens_diameter_dict = {int(col1): col2 for col1, col2 in self.Lens_diameter_table}
        # end lens diameter stuff ------->
        
        self.energy = 0  # gets value from an ao (incoming beam energy)
        self.focalSize = 0 # get value from an ao (desired focal length)
        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.verbose = True

        self.lookupTable = []
        
        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
        sigmaHp_e   : Sigma electron divergence in H direction in rad
        sigmaVp_e   : Sigma electron divergence in V direction in rad
        '''
        
        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']
        
        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):
        '''
        Beamline properties can contain entries for the following
        
        d_StoL1 : Source-to-CRL1 distance, in m
        d_Stof  : Source-to-focus distance, in m
        '''            
        
        self.bl['d_StoL'] = beamline_properties['d_StoL']
        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
        
        '''
        pass 
                
    def setupLookupTable(self, subs_file, n_lenses):
        '''
        lookup table created after IOC startup
        energy and slit size are updated before this is called
        '''
        print(80*'#')
        print('Setting up lens control...')
        
        self.num_stacks = n_lenses
        self.num_configs = 2**self.num_stacks
        self.configs = np.arange(self.num_configs)
        
#       self.energy = energy
        
        #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):
            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")
        
        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
        print('Getting lens materials...')
        if NLENS_MACRO in macros:
            self.numlens = np.array([int(i) for i in lens_properties[NLENS_MACRO]])
            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
        print('Getting lens\' radii...')
        if RADIUS_MACRO in macros:
            self.radius = np.array([float(i) for i in lens_properties[RADIUS_MACRO]])
            print('Radius of lenses read in.\n')
        else:
            raise RuntimeError(f"Radius macro ({RADIUS_MACRO}) not found in substituion file")

        # 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)
        self.densities = np.array([densities[material] for material in self.materials])

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

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

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

    def construct_lookup_table(self):
        self.lookupTable = 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.slit1_H, self.slit1_V,
                                             self.lens_thickerr, flag_HE = self.thickerr_flag,
                                             verbose = self.verbose)
#        self.cull_lookup_table()
        self.sort_lookup_table()
        self.updateEnergyRBV()
        self.updateSlitSizeRBV('hor')
        self.updateSlitSizeRBV('vert')
        self.updateLookupWaveform()
        self.updateLookupConfigs()

#    def cull_lookup_table(self):
#        '''
#        Culls the lookup table based on lenses that are locked and or disabled
#        '''
#        self.culledSize = 2**(self.num_stacks - (self.outMask | self.inMask).bit_count())
#        if self.verbose: print(f'Operating spaced now at {self.culledSize} configurations')
#        if self.verbose: print(f'Culling table with in mask {self.inMask} and out mask {self.outMask}')
#        
#        self.culledConfigs = np.empty(self.culledSize, dtype=int)
#        self.culledTable = np.empty(self.culledSize)
#        j = 0
#        for i in range(2**self.num_stacks):
#            if ((i & self.outMask == 0) and (i & self.inMask == self.inMask)):
#                self.culledConfigs[j]=i
#                self.culledTable[j] = self.lookupTable[i]
#                j += 1
#                
#        self.sort_lookup_table()
        
    def sort_lookup_table(self):
        '''
        
        '''
#        if self.verbose: print(f'Sorting culled lookup table of length {len(self.culledTable)}')        
#        self.sorted_index = np.argsort(self.culledTable)
        if self.verbose: print(f'Sorting lookup table of length {len(self.lookupTable)}')        
        self.sorted_index = np.argsort(self.lookupTable)

    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
        '''
        self.index = int(config_BW)
        self.indexSorted = self.sorted_index.tolist().index(self.index)

        self.setFocalSizeActual()
        self.updateLensRBV()
        self.updateFocalSizeRBVs()      

#    def setInMask(self, inMask):
#        '''
#        update mask for lenses that are locked in
#        '''
#        self.inMask = int(inMask)
#        self.cull_lookup_table()
#        if self.verbose: print(f'Converting culled index via in Mask')
#        self.convertCulledIndex()
#        if self.verbose: print(f'Updating transfocator RBV via in Mask')
#        self.updateLensRBV()
#        if self.verbose: print(f'Setting in mask RBV to {self.inMask}')
#        pydev.iointr('new_inMask', int(self.inMask))
#
#    def setOutMask(self, outMask):
#        '''
#        update mask for lenses that must remain out (either disabled or locked)
#        '''
#        self.outMask = int(outMask)
#        self.cull_lookup_table()
#        if self.verbose: print(f'Converting culled index via out Mask')
#        self.convertCulledIndex()
#        if self.verbose: print(f'Updating transfocator RBV via out Mask')
#        self.updateLensRBV()
#        if self.verbose: print(f'Setting out mask RBV to {self.outMask}')
#        pydev.iointr('new_outMask', int(self.outMask))

#   def convertCulledIndex(self):
#       '''
#       When available configs change, need to update index so that tweaks 
#       continue to work
#       '''
#       if self.verbose: print('Converting ...')
#       self.culledIndex = (np.where(self.culledConfigs == self.config))[0][0]
#       if self.verbose: print(f'Culled index is {self.culledIndex}')
#       self.culledIndexSorted = self.sorted_index.tolist().index(self.culledIndex)
#       if self.verbose: print(f'Sorted culled index is {self.culledIndexSorted}')

    def setFocalSizeActual(self):
        '''
        
        '''
#        self.focalSize_actual = self.culledTable[self.culledIndex] 
        self.focalSize_actual = self.lookupTable[self.index] 

    def find_config(self):
        ''' 
        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
        if self.verbose: print(f'Searching for config closest to {self.focalSize}')
#       self.culledIndex = np.argmin(np.abs(self.culledTable - self.focalSize))
#       if self.verbose: print(f'Config index found at {self.culledIndex}')
        self.index = np.argmin(np.abs(self.lookupTable - self.focalSize))
        if self.verbose: print(f'Config index found at {self.index}')

#       self.culledIndexSorted = self.sorted_index.tolist().index(self.culledIndex)
#       if self.verbose: print(f'Sorted config index found at {self.culledIndexSorted}')
        self.indexSorted = self.sorted_index.tolist().index(self.index)
        if self.verbose: print(f'Sorted config index found at {self.indexSorted}')

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

    def getPreviewFocalSize(self, sortedIndex):
        '''
        
        '''
#       fSize_preview = self.culledTable[self.sorted_index[int(sortedIndex)]]
        fSize_preview = self.lookupTable[self.sorted_index[int(sortedIndex)]]
        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
        '''
        
        if slit == 'hor':
            self.slit1_H = float(size)     # H slit size before CRL 1
        elif slit == 'vert':
            self.slit1_V = float(size)     # V slit size before CRL 1
        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.
        '''
        self.setSlitSize(size, slit)
        
        if self.verbose: 
            if slit == 'hor':
                print(f'Horizontal slit size updated to {self.slit1_H} m')
            elif slit == 'vert':
                print(f'Vertical slit size updated to {self.slit1_V} m')
        
        
        # self.updateSlitSizeRBV(slit)
        # Testing calling the lookup table reconstruction in EPICS instead of python
        #! self.construct_lookup_table()

    def setEnergy(self, energy):
        '''
        Sets various forms of energy
        '''
        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(?)
        
    def updateE(self, energy):
        '''
        Beam energy updates are propagated to CRL object from EPICS. The beam
        size lookup table is then recalculated.
        '''
        # Energy variable sent from IOC as a string
        self.setEnergy(energy)
            
        if self.verbose: print(f'Setting energy to {self.energy} keV')
            
        # Update beam properties that are dependent on energy
        self.setupSourceEnergyDependent()
        
        # Testing calling the lookup table reconstruction in EPICS instead of python
        #! self.construct_lookup_table()

        # Do I need to find what the current config would produce as far as focal size and location?
        # self.focalSizeRBV = 

    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
        '''
#       self.culledIndexSorted = int(sortedIndex)
#       self.culledIndex = self.sorted_index[self.culledIndexSorted]
        self.indexSorted = int(sortedIndex)
        self.index = self.sorted_index[self.indexSorted]

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

    def setThickerrFlag(self, flag):
        '''
        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)

        
    def updateLensConfigPV(self):
        '''
        
        '''
#       self.config = self.culledConfigs[self.culledIndex]
        self.config = self.configs[self.index]
        pydev.iointr('new_lenses', int(self.config))

    def updateLensRBV(self):
        '''
        
        '''
#        pydev.iointr('new_index', int(self.culledIndexSorted))
        pydev.iointr('new_index', int(self.indexSorted))

    def updateEnergyRBV(self):
        '''
        
        '''
        pydev.iointr('updated_E', float(self.energy))

    def updateSlitSizeRBV(self, slit):
        '''
        Update proper slit size
        '''
        
        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')
        elif slit == 'vert':
            pydev.iointr('updated_slitSize_V', float(self.slit1_V))
            if self.verbose: print(f'Vertical slit size RBV updated to {self.slit1_V} m')
            

        
    def updateFocalSizeRBVs(self):
        '''
        
        '''
        pydev.iointr('new_fSize', self.focalSize_actual)
        
    def updateVerbosity(self, verbosity):
        '''
        Turn on minor printing
        '''
        print(f'Verbosity set to {verbosity}')
        self.verbose = int(verbosity)

    def updateLookupWaveform(self):
        '''
        Puts lookup table focal sizes into waveform PV
        '''
        pydev.iointr('new_lookupTable', self.lookupTable.tolist())
#       pydev.iointr('new_culledTable', self.culledTable.tolist())
    
    def updateLookupConfigs(self):
        '''
        Puts lookup table config integers into waveform PV
        '''
#       pydev.iointr('new_culledConfigs', self.culledConfigs.tolist())
        pydev.iointr('new_configs', self.configs.tolist())