def __init__(self, length_param, rad_in, rad_out, cutoff_level = 0.01): self.length = length_param self.rad_in = rad_in self.rad_out = rad_out self.cutoff_level = cutoff_level self.normalization = 1.0 self.n_func_points = 500 z_step = length_param/self.n_func_points z_cutoff = self._findCutOff(z_step,cutoff_level) self.z_min = - z_cutoff self.z_max = + z_cutoff self.func = Function() self._normalize()
def __init__(self, length_param, acceptance_diameter_param, cutoff_level = 0.01): self.length = length_param self.acceptance_diameter = acceptance_diameter_param self.a_arr = [0.296471,4.533219,-2.270982,1.068627,-0.036391,0.022261] self.normalization= 1.0 self.n_func_points = 500 #-----find cut-off z value self.cutoff_z = self.acceptance_diameter step = self.acceptance_diameter self.cutoff_level = cutoff_level self.cutoff_z = self._findCutOff(step, cutoff_level) #------------------------------------------------------ self.func = Function() self._normalize()
def getAxisEz(self, zSimmetric = -1): """ Returns the Spline with Ez(z) on the axis of the RF. If zSimmetric > 0 the table has only half of the table, and the Function should be added points for (-Zmax) to (Zmin - step). """ stepZ = (self.Zmax - self.Zmin)/self.zSteps Ez_max = 0. for iz in range(self.zSteps+1): [z,r,Ez,Er,E,B] = self.data_arr[iz] Ez_abs = math.fabs(Ez) if(Ez_max < Ez_abs): Ez_max = Ez_abs #The z in the self.data_arr is in [cm], so to switch to [m] we use 0.01 f = Function() if(zSimmetric > 0): for iz in range(1,self.zSteps+1): [z,r,Ez,Er,E,B] = self.data_arr[iz] z = self.Zmin + stepZ*iz f.add(-z*0.01,Ez/Ez_max) for iz in range(self.zSteps+1): [z,r,Ez,Er,E,B] = self.data_arr[iz] z = self.Zmin + stepZ*iz f.add(z*0.01,Ez/Ez_max) spline = SplineCH() spline.compile(f) return spline
def __init__(self, length_param, acceptance_diameter_param, cutoff_level = 0.01): self.length = length_param self.acceptance_diameter = acceptance_diameter_param self.a_arr = [0.296471,4.533219,-2.270982,1.068627,-0.036391,0.022261] self.normalization= 1.0 self.n_func_points = 500 #-----find cut-off z value self.cutoff_z = self.acceptance_diameter step = self.acceptance_diameter self.cutoff_level = cutoff_level self.cutoff_z = self._findCutOff(step, cutoff_level) #------------------------------------------------------ self.func = Function() self._normalize()
def getNormilizedSpline(self): """ Returns the spline normilized by the integral of the absolute value. """ n = self.splineFiled.getSize() f = Function() for i in range(n): f.add(self.splineFiled.x(i),math.fabs(self.splineFiled.y(i))) integral = GaussLegendreIntegrator(500) integral.setLimits(self.splineFiled.x(0),self.splineFiled.x(n-1)) spline = SplineCH() spline.compile(f) res = integral.integral(spline) f = Function() for i in range(n): f.add(self.splineFiled.x(i),self.splineFiled.y(i)/res) spline = SplineCH() spline.compile(f) return spline
def RenormalizeFunction(func, z_min, z_max): """ It re-normalizes the Function in the new limits (z_min,z_max). We assume that region of the function definition will be cut not extended. """ spline = SplineCH() spline.compile(func) integrator = GaussLegendreIntegrator(500) integrator.setLimits(z_min, z_max) integral = integrator.integral(spline) n_points = func.getSize() step = (z_max - z_min) / (n_points - 1) new_func = Function() for i in range(n_points): x = z_min + step * i y = spline.getY(x) / integral new_func.add(x, y) new_func.setConstStep(1) return new_func
def addAxisField(cls,fl_name,dir_location = ""): """ This method add to the store the axis RF field for the RF gap node. The dir_location string variable will be added to the fl_name to get the file name. Returns the axis RF field function. """ if(cls.static_axis_field_dict.has_key(fl_name)): return cls.static_axis_field_dict[fl_name] comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD data_type = mpi_datatype.MPI_DOUBLE rank = orbit_mpi.MPI_Comm_rank(comm) main_rank = 0 x_arr = [] y_arr = [] if(rank == 0): fl_in = open(dir_location + fl_name,"r") lns = fl_in.readlines() fl_in.close() for ln in lns: res_arr = ln.split() if(len(res_arr) == 2): x = float(res_arr[0]) y = float(res_arr[1]) x_arr.append(x) y_arr.append(y) x_arr = orbit_mpi.MPI_Bcast(x_arr,data_type,main_rank,comm) y_arr = orbit_mpi.MPI_Bcast(y_arr,data_type,main_rank,comm) function = Function() for ind in range(len(x_arr)): function.add(x_arr[ind],y_arr[ind]) #---- setting the const step (if function will allow it) #---- will speed up function calculation later function.setConstStep(1) cls.static_axis_field_dict[fl_name] = function return function
def makeTransitTimeTables(self,beta_min,beta_max,n_table_points,rf_freq): """ It will calculate transit time factor tables for all RF gaps TTFs (T,S,Tp,Sp) are funcftions of the cappa variable = 2*pi*f/(c*beta) """ self.rf_freq = rf_freq c_light = 2.99792458e+8 self.beta_arr = [] self.cappa_arr = [] for i_beta in range(n_table_points): beta = beta_min + i_beta*(beta_max-beta_min)/(n_table_points-1) cappa = 2*math.pi*rf_freq/(c_light*beta) self.beta_arr.append(beta) self.cappa_arr.append(cappa) self.beta_arr.reverse() self.cappa_arr.reverse() #--calculate realtive gap amplitudes integral = GaussLegendreIntegrator(500) e0l_arr = [] e0l_sum = 0. for i in range(len(self.gap_slpline_arr)): [gap_length,x_center,splineGap] = self.gap_slpline_arr[i] n = splineGap.getSize() x_min = splineGap.x(0) x_max = splineGap.x(n-1) integral.setLimits(x_min,x_max) e0l = integral.integral(splineGap) e0l_sum += e0l e0l_arr.append(e0l) self.e0_normalized_arr = [] self.e0l_normalized_arr = [] e0_norm = e0l_arr[0]/self.gap_slpline_arr[0][0] e0l_norm = e0l_arr[0] for i in range(len(e0l_arr)): self.e0_normalized_arr.append((e0l_arr[i]/self.gap_slpline_arr[i][0])/e0_norm) self.e0l_normalized_arr.append((e0l_arr[i]/e0l_norm)) #--- calculate transit time factors self.ttp_ssp_gap_arr = [] for i in range(len(self.gap_slpline_arr)): func_T = Function() func_TP = Function() func_S = Function() func_SP = Function() self.ttp_ssp_gap_arr.append([func_T,func_TP,func_S,func_SP]) for i_gap in range(len(self.gap_slpline_arr)): [func_T,func_TP,func_S,func_SP] = self.ttp_ssp_gap_arr[i_gap] [gap_length,x0,spline] = self.gap_slpline_arr[i_gap] x_min = spline.x(0) x_max = spline.x(spline.getSize()-1) integral.setLimits(x_min,x_max) for i_beta in range(n_table_points): beta = self.beta_arr[i_beta] cappa = self.cappa_arr[i_beta] f_cos = Function() f_sin = Function() for isp in range(spline.getSize()): x = spline.x(isp) y = spline.y(isp) phase = cappa*x s = math.sin(phase) c = math.cos(phase) f_cos.add(x,c*y) f_sin.add(x,s*y) f_sp_cos = SplineCH() f_sp_sin = SplineCH() f_sp_cos.compile(f_cos) f_sp_sin.compile(f_sin) T = integral.integral(f_sp_cos) S = integral.integral(f_sp_sin) func_T.add(cappa,T/e0l_arr[i_gap]) func_S.add(cappa,S/e0l_arr[i_gap]) spline_T = SplineCH() spline_S = SplineCH() spline_T.compile(func_T) spline_S.compile(func_S) for i_beta in range(spline_T.getSize()): cappa = spline_T.x(i_beta) TP = spline_T.getYP(cappa) SP = spline_S.getYP(cappa) func_TP.add(cappa,TP) func_SP.add(cappa,SP) return self.ttp_ssp_gap_arr
# In the fitting procedure we find these parameters to reproduce data. #------------------------------------------------------------------------------- import sys import math import time from orbit_utils import Function from orbit_utils import HarmonicData import numpy as np import scipy from scipy.optimize import minimize #---- data generation. They are inside the Function f f = Function() for ind in range(0, 360, 30): x = 1.0 * ind y = 2.0 * math.cos((math.pi / 180.) * (x + 25.)) + 4.0 * math.cos( (math.pi / 180.) * (4 * x + 35.)) + 0.5 y_err = 0.01 * abs(y) f.add(x, y, y_err) #---------------------------------------------------- order = 4 harmonic_data = HarmonicData(order, f) #---- setup the parameters of 4 order harmonic fit function x_arr = [0.8, 2.1, 25.2, 0., 0., 0., 0., 4.3, 35.4] for x_ind in range(len(x_arr)):
(x, xp, y, yp, z, dE) = (0.05, 0.0, 0.05, 0.0, 0., 0.001) bunch_ini.addParticle(x, xp, y, yp, z, dE) bunch_tmp = Bunch() bunch_ini.copyBunchTo(bunch_tmp) start_ind = accLattice.getNodeIndex(quads[0]) stop_ind = accLattice.getNodeIndex(quads[len(quads) - 1]) accLattice.trackDesignBunch(bunch_tmp, None, None, index_start=start_ind, index_stop=stop_ind) traj_x_function = Function() traj_xp_function = Function() #----- if you want more beam size points along the trajectory - change pos_step here paramsDict = {"old_pos": -1., "count": 0, "pos_step": 0.1, "path_length": 0.} actionContainer = AccActionsContainer("TEAPOT Bunch Tracking") #----- start of the 1st quad pos_start = -quads[0].getLength() / 2 def action_account(paramsDict): node = paramsDict["node"] name = node.getName() bunchI = paramsDict["bunch"] pos = paramsDict["path_length"]
from orbit_utils import Function from orbit_utils import SplineCH from orbit_utils import GaussLegendreIntegrator integrator = GaussLegendreIntegrator(4) print "Number of integral points =", integrator.getnPoints() print "The integral of sin(x) from 0. to pi/2 is 1." #------------------------------------------ # integral of sin(x) from 0. to pi/2 is 1. #------------------------------------------ x_min = 0. x_max = math.pi / 2 integrator.setLimits(x_min, x_max) f = Function() n = 10 for i in range(n): x = x_min + i * (x_max - x_min) / (n - 1) f.add(x, math.sin(x)) res = integrator.integral(f) print "For Function integral =", res, " error=", math.fabs(res - 1.0) spline = SplineCH() spline.compile(f) res = integrator.integral(spline) print "For SplineCH integral =", res, " error=", math.fabs(res - 1.0) print "=============== Integration points and weights ================" point_weight_arr = integrator.getPointsAndWeights() for (x, w) in point_weight_arr:
import sys import math import random from orbit_utils import Function def FF(x): return math.sin(x) def FFP(x): return math.cos(x) f = Function() eps = 1.0e-5 n = 100 step = 2 * math.pi / n for i in range(n): x = step * i + random.uniform(-eps * step, eps * step) y = FF(x) f.add(x, y) f.dump() # let's check Const step tolerance setting procedure f.setStepEps(4 * eps) print "Const step tolerance =", f.getStepEps() res = f.setConstStep(1)
# It also has parameters of the fitting representation of this y(x) data: # y_fit(x) = A0 + A1*cos((PI/180.)*(x + phase_shift_1)) + # A2*cos((PI/180.)*(2*x + parse_shift_2)) + # ... # The fast calculation of the sumDiff2 is the main purpose of this class # This class is used for the harmonics analysis of the scan data for # the RF cavities when we measure the phase response from BPMs. #------------------------------------------------------------------------------- import sys import math from orbit_utils import Function from orbit_utils import HarmonicData f = Function() for ind in range(0, 360, 30): x = 1.0 * ind y = 2.0 * math.cos((math.pi / 180.) * (x + 25.)) + 4.0 * math.cos( (math.pi / 180.) * (4 * x + 35.)) + 0.5 y_err = 0.01 * abs(y) f.add(x, y, y_err) order = 4 harmonic_data = HarmonicData(order, f) harmonic_data.parameter(0, 0.5) harmonic_data.parameter(1, 2.0) harmonic_data.parameter(2, 25.0)
class EngeFunction: """ The Enge function with parameters from Berz's paper M.Berz, B. Erdelyn, K.Makino Fringe Field Effects in Small Rings of Large Acceptance Phys. Rev STAB, V3, 124001(2000) """ def __init__(self, length_param, acceptance_diameter_param, cutoff_level = 0.01): self.length = length_param self.acceptance_diameter = acceptance_diameter_param self.a_arr = [0.296471,4.533219,-2.270982,1.068627,-0.036391,0.022261] self.normalization= 1.0 self.n_func_points = 500 #-----find cut-off z value self.cutoff_z = self.acceptance_diameter step = self.acceptance_diameter self.cutoff_level = cutoff_level self.cutoff_z = self._findCutOff(step, cutoff_level) #------------------------------------------------------ self.func = Function() self._normalize() def setEngeCoefficients(self,a_arr): """ Sets new values for Enge function's coeffients. """ self.a_arr = a_arr step = self.length/2 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def setCutOffLevel(self,cutoff_level): """ Sets the cutoff level for quad's field """ step = self.length/2 self.cutoff_level = cutoff_level self.cutoff_z = self._findCutOff(step, cutoff_level) self._normalize() def setCutOffZ(self,cutoff_z): """ Sets the cutoff distance from the center of quad's field""" self.cutoff_z = cutoff_z self._normalize() def setLength(self,length): """ Sets the length of quad's field""" self.length = length step = self.length/2.0 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def setAcceptanceDiameter(self,acceptance_diameter): """ Sets the acceptance diameter of the quad """ self.acceptance_diameter = acceptance_diameter step = self.length/2.0 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def setNumberOfPoints(self,n_func_points): """ Sets the number of points in the field function """ self.n_func_points = n_func_points step = self.length/2.0 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def getCuttOffZ(self): """ Returns the cutoff distance from the center of quad's field""" return self.cutoff_z def getNumberOfPoints(self): """ Returns the number of points in the field function """ return self.n_func_points def _getTrueEngeFunc(self, x): """ Returns the quad's field at the distance x from the center """ # x is the distance from the center of the magnet with the iron length l """ x = (math.fabs(x) - self.length/2.0)/self.acceptance_diameter sum_exp = self.a_arr[0] x0 = x for i in range(1,len(self.a_arr)): sum_exp += self.a_arr[i]*x0 x0 *= x if(abs(sum_exp) > 30.): sum_exp = 30.0*sum_exp/abs(sum_exp) return self.normalization/(1.0+math.exp(sum_exp)) def _findCutOff(self,step, cutoff_level): """ Finds the distance from the center where the field is less than cutoff level """ self.normalization= 1.0 init_val = self._getTrueEngeFunc(0.) z = step val = self._getTrueEngeFunc(z)/init_val if(val <= cutoff_level): return z while(val > cutoff_level): z += step val = self._getTrueEngeFunc(z)/init_val z0 = z - step z1 = z step_z = step/self.n_func_points val0 = self._getTrueEngeFunc(z0)/init_val val1 = self._getTrueEngeFunc(z1)/init_val while(abs(z0-z1) > step_z): z_new = (z0+z1)/2.0 val_new = self._getTrueEngeFunc(z_new)/init_val if(val_new <= cutoff_level): z1 = z_new val1 = val_new else: z0 = z_new val0 = val_new self.cutoff_z = (z0+z1)/2.0 return self.cutoff_z def _normalize(self): """ Normalizes the quad field function to the integral of 1 """ self.normalization = 1.0 step = self.cutoff_z/(self.n_func_points - 1) self.func.clean() sum_int = 0. for ind in range(self.n_func_points): z = step*ind val = self._getTrueEngeFunc(z) self.func.add(z,val) sum_int += val sum_int -= (self._getTrueEngeFunc(0.) + self._getTrueEngeFunc(step*(self.n_func_points - 1)))/2.0 sum_int *= 2.0*step self.normalization = 1.0/sum_int self.func.setConstStep(1) def getFuncValue(self,z): """ Returns the quad's field at the distance z from the center """ if(abs(z) >= self.func.getMaxX()): return 0. return self.normalization*self.func.getY(abs(z)) def getLimitsZ(self): """ Returns the tuple with min and max Z value for this field """ z_max = self.func.getMaxX() return (-z_max,z_max) def isInside(self,z_center,z): """ Returns True if the position Z is inside the function definition region """ if(abs(z-z_center) <= self.func.getMaxX()): return True return False
#------------------------------------------------------ # This is an example of Function. # It is a containers for (x,y) (or (x,y,err) ) points #------------------------------------------------------- import sys import math from orbit_utils import Function f = Function() def FF(x): return math.sin(x) n = 10 step = 2*math.pi/n for i in range(n): x = step*i+0.1*((1.0*i)/n)**2; y = FF(x) err = y*0.001 f.add(x,y,err) f.dump() print "=======================================" for i in range(n): x = f.x(i) y = f.y(i) err = f.getYErr(x) print "i="," x,y,err= %10.7f %10.7f %10.7f "%(x,y,err*1000)
def __init__(self,splineFiled, zeroIsCenter = False): self.splineFiled = splineFiled #---------------------------------------------------- self.eps_root = 1.0e-6 self.rf_freq = 0.0 #self.e0_normalized_arr - normilized amplitudes of the gaps self.e0_normalized_arr = [] self.e0l_normalized_arr = [] # self.beta_arr - relativistic beta, cappa = 2*math.pi*rf_freq/(c_light*beta) self.beta_arr = [] self.cappa_arr = [] # self.ttp_ssp_gap_arr array with [T,Tp,S,Sp] for all gaps (T,Tp,S,Sp - Functions of cappa) self.ttp_ssp_gap_arr = [] # self.gap_polynoms_coef_arr array with polynomial coefficients for [T,Tp,S,Sp] for each gap self.gap_polynoms_coef_arr = [] # self.gap_polynoms_arr array with Polynomial instances for [T,Tp,S,Sp] for each gap self.gap_polynoms_arr = [] # self.gap_polynoms_t_tp_s_sp_err_arr - maximal relative errors for polynomial fitting self.gap_polynoms_t_tp_s_sp_err_arr = [] #----------------------------------------------------- #calculate the roots self.roots_arr = self.rootAnalysis() #find the roots of derivative - yp = y' - RFgap center positions if(zeroIsCenter): self.yp_roots_arr = [0.] else: self.yp_roots_arr = self.gapCentersAnalysis() #print "debug yp roots=",self.yp_roots_arr if(len(self.roots_arr) - 1 != len(self.yp_roots_arr)): rank = orbit_mpi.MPI_Comm_rank(mpi_comm.MPI_COMM_WORLD) if(rank == 0): print "Class RF_AxisFieldAnalysis." print "The structure of the input rf field spline is wrong!" print "roots of the filed =",self.roots_arr print "extrema positions =",self.yp_roots_arr sys.exit(1) # caluclate the position of the center of the cavity rf_center = 0 for i in range(1,len(self.yp_roots_arr)-1): rf_center += self.yp_roots_arr[i] rf_center /= (len(self.yp_roots_arr)-2) self.rf_center = rf_center # make spline for each RF gap self.gap_slpline_arr = [] #print "debug roots_arr=",self.roots_arr #---make splineGap with x in the [m] instead of [cm] for i in range(len(self.roots_arr)-1): x_center = self.yp_roots_arr[i] x0 = self.roots_arr[i] x1 = self.roots_arr[i+1] f = Function() f.add((x0-x_center),math.fabs(splineFiled.getY(x0))) for ix in range(splineFiled.getSize()-1): x = splineFiled.x(ix) if(x > x0 and x < x1): f.add((x-x_center),math.fabs(splineFiled.y(ix))) f.add((x1-x_center),math.fabs(splineFiled.getY(x1))) splineGap = SplineCH() splineGap.compile(f) n = splineGap.getSize() x_min = splineGap.x(0) x_max = splineGap.x(n-1) gap_length = x_max - x_min self.gap_slpline_arr.append([gap_length,(x_center - self.rf_center),splineGap])
#---------------------------------------------- # This is an example of a polymonial fitting #---------------------------------------------- import sys import math from orbit_utils import Function from orbit_utils import SplineCH from orbit.utils.fitting import PolynomialFit f = Function() for i in range(15): x = 0.1 * i y = -1 + 1.0 * x**2 + 2.0 * x**3 + 0.01 * x**4 f.add(x, y) print "polynomial for fitting = y = -1+1.0*x**2+2.0*x**3 + 0.01*x**4" print "================Function====================================" poly = PolynomialFit(4) poly.fitFunction(f) [coef_arr, err_arr] = poly.getCoefficientsAndErr() for i in range(len(coef_arr)): print "i=", i, " coef = %8.5f +- %8.5f" % (coef_arr[i], err_arr[i]) print "================SplineCH====================================" spline = SplineCH() spline.compile(f) poly.fitSpline(f) [coef_arr, err_arr] = poly.getCoefficientsAndErr()
def getAxisFieldFunction(fl_name): """ This function will read the .FSO file and extract the axis field Ez(z). """ fl_in = open(fl_name, "r") lns = fl_in.readlines() fl_in.close() func = Function() i_start = -1 for i in range(len(lns)): if (lns[i].find(" Z(cm) Ez(V/m)") >= 0): i_start = i + 1 break for i in range(i_start, len(lns)): res_arr = lns[i].split() if (len(res_arr) != 2): break x = 0.01 * float(res_arr[0]) y = float(res_arr[1]) func.add(x, y) #fix for the field at the 0 region func1 = Function() for i in range(1, func.getSize()): ind = func.getSize() - i x = -func.x(ind) y = func.y(ind) func1.add(x, y) for i in range(func.getSize()): x = func.x(i) y = func.y(i) func1.add(x, y) return func1
#---------------------------------------------- # This is an example of a polymonial fitting #---------------------------------------------- import sys import math from orbit_utils import Function from orbit_utils import SplineCH from orbit.utils.fitting import PolynomialFit f = Function() for i in range(15): x = 0.1*i y = -1+1.0*x**2+2.0*x**3 + 0.01*x**4 f.add(x,y) print "polynomial for fitting = y = -1+1.0*x**2+2.0*x**3 + 0.01*x**4" print "================Function====================================" poly = PolynomialFit(4) poly.fitFunction(f) [coef_arr,err_arr] = poly.getCoefficientsAndErr() for i in range(len(coef_arr)): print "i=",i," coef = %8.5f +- %8.5f"%(coef_arr[i],err_arr[i]) print "================SplineCH====================================" spline = SplineCH() spline.compile(f) poly.fitSpline(f) [coef_arr,err_arr] = poly.getCoefficientsAndErr()
class EngeFunction(AbstractQuadFieldSourceFunction): """ The Enge function with parameters from Berz's paper M.Berz, B. Erdelyn, K.Makino Fringe Field Effects in Small Rings of Large Acceptance Phys. Rev STAB, V3, 124001(2000) """ def __init__(self, length_param, acceptance_diameter_param, cutoff_level=0.01): self.length = length_param self.acceptance_diameter = acceptance_diameter_param self.a_arr = [ 0.296471, 4.533219, -2.270982, 1.068627, -0.036391, 0.022261 ] self.normalization = 1.0 self.n_func_points = 500 #-----find cut-off z value self.cutoff_z = self.acceptance_diameter step = self.acceptance_diameter self.cutoff_level = cutoff_level self.cutoff_z = self._findCutOff(step, cutoff_level) #------------------------------------------------------ self.func = Function() self._normalize() def setEngeCoefficients(self, a_arr): """ Sets new values for Enge function's coeffients. """ self.a_arr = a_arr step = self.length / 2 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def setCutOffLevel(self, cutoff_level): """ Sets the cutoff level for quad's field """ step = self.length / 2 self.cutoff_level = cutoff_level self.cutoff_z = self._findCutOff(step, cutoff_level) self._normalize() def setCutOffZ(self, cutoff_z): """ Sets the cutoff distance from the center of quad's field""" self.cutoff_z = cutoff_z self._normalize() def setLength(self, length): """ Sets the length of quad's field""" self.length = length step = self.length / 2.0 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def setAcceptanceDiameter(self, acceptance_diameter): """ Sets the acceptance diameter of the quad """ self.acceptance_diameter = acceptance_diameter step = self.length / 2.0 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def setNumberOfPoints(self, n_func_points): """ Sets the number of points in the field function """ self.n_func_points = n_func_points step = self.length / 2.0 self.cutoff_z = self._findCutOff(step, self.cutoff_level) self._normalize() def getCuttOffZ(self): """ Returns the cutoff distance from the center of quad's field""" return self.cutoff_z def getNumberOfPoints(self): """ Returns the number of points in the field function """ return self.n_func_points def _getTrueEngeFunc(self, x): """ Returns the quad's field at the distance x from the center """ # x is the distance from the center of the magnet with the iron length l """ x = (math.fabs(x) - self.length / 2.0) / self.acceptance_diameter sum_exp = self.a_arr[0] x0 = x for i in range(1, len(self.a_arr)): sum_exp += self.a_arr[i] * x0 x0 *= x if (abs(sum_exp) > 30.): sum_exp = 30.0 * sum_exp / abs(sum_exp) return self.normalization / (1.0 + math.exp(sum_exp)) def _findCutOff(self, step, cutoff_level): """ Finds the distance from the center where the field is less than cutoff level """ self.normalization = 1.0 init_val = self._getTrueEngeFunc(0.) z = step val = self._getTrueEngeFunc(z) / init_val if (val <= cutoff_level): return z while (val > cutoff_level): z += step val = self._getTrueEngeFunc(z) / init_val z0 = z - step z1 = z step_z = step / self.n_func_points val0 = self._getTrueEngeFunc(z0) / init_val val1 = self._getTrueEngeFunc(z1) / init_val while (abs(z0 - z1) > step_z): z_new = (z0 + z1) / 2.0 val_new = self._getTrueEngeFunc(z_new) / init_val if (val_new <= cutoff_level): z1 = z_new val1 = val_new else: z0 = z_new val0 = val_new self.cutoff_z = (z0 + z1) / 2.0 return self.cutoff_z def _normalize(self): """ Normalizes the quad field function to the integral of 1 """ self.normalization = 1.0 step = self.cutoff_z / (self.n_func_points - 1) self.func.clean() sum_int = 0. for ind in range(self.n_func_points): z = step * ind val = self._getTrueEngeFunc(z) self.func.add(z, val) sum_int += val sum_int -= (self._getTrueEngeFunc(0.) + self._getTrueEngeFunc(step * (self.n_func_points - 1))) / 2.0 sum_int *= 2.0 * step self.normalization = 1.0 / sum_int self.func.setConstStep(1) def getFuncValue(self, z): """ Returns the quad's field at the distance z from the center """ if (abs(z) >= self.func.getMaxX()): return 0. return self.normalization * self.func.getY(abs(z)) def getLimitsZ(self): """ Returns the tuple with min and max Z value for this field """ z_max = self.func.getMaxX() return (-z_max, z_max)
#---- Print vector function def printV(v): print "----vector--- size=",v.size() for i in xrange(v.size()): print ("v(" + str(i) + ")="+str(v.get(i)) + " "), print "" #---- Let's read the dB/dr distribution. It could be equals to zero, #---- and in this case we will treat this as drift file_name = "MEBT_fields_rf_and_quads.dat" fl_in = open(file_name,"r") lns = fl_in.readlines() fl_in.close() #---- first line in the input files is a headers line g_func = Function() for ind in range(1,len(lns)): ln = lns[ind].strip() res_arr = ln.split() if(len(res_arr) > 5): #---- here we parse the input file: 1st position - z, 4th - dB/dr(z z = float(res_arr[0]) g = float(res_arr[3]) g_func.add(z,g) #---- let's set energy and momentum mass = 0.938272 + 2*0.000511 eKin = 0.0025 momentum = math.sqrt(eKin*(eKin + 2*mass)) Brho = 3.33564*momentum # [T*m] charge = -1.0
#------------------------------------------------------ # This is an example of Function and SplineCH # They are containers for (x,y) (or (x,y,err) ) points # Function provides a linear interpolation, and SplineCH # uses 3-rd order polynomials. SplineCH can be used for # derivatives calculations. #------------------------------------------------------- import sys import math from orbit_utils import Function from orbit_utils import SplineCH f = Function() def FF(x): return math.sin(x) def FFP(x): return math.cos(x) n = 40 step = 2*math.pi/n for i in range(n): x = step*i+0.1*((1.0*i)/n)**2; y = FF(x) f.add(x,y) f.dump()
class PMQ_Trace3D_Function(AbstractQuadFieldSourceFunction): """ The PMQ Function is a represenatation of the field of permanent quad from Trace3D documantation (p 77): http://laacg.lanl.gov/laacg/services/traceman.pdf """ def __init__(self, length_param, rad_in, rad_out, cutoff_level = 0.01): self.length = length_param self.rad_in = rad_in self.rad_out = rad_out self.cutoff_level = cutoff_level self.normalization = 1.0 self.n_func_points = 500 z_step = length_param/self.n_func_points z_cutoff = self._findCutOff(z_step,cutoff_level) self.z_min = - z_cutoff self.z_max = + z_cutoff self.func = Function() self._normalize() def _findCutOff(self,step, cutoff_level): """ Finds the distance from the center where the field is less than cutoff level """ init_val = self.getPMQ_FuncValue(0.) z = step val = self.getPMQ_FuncValue(z)/init_val if(val <= cutoff_level): return z while(val > cutoff_level): z += step val = self.getPMQ_FuncValue(z)/init_val z0 = z - step z1 = z n_inner_points = 100 step_z = step/n_inner_points val0 = self.getPMQ_FuncValue(z0)/init_val val1 = self.getPMQ_FuncValue(z1)/init_val while(abs(z0-z1) > step_z): z_new = (z0+z1)/2.0 val_new = self.getPMQ_FuncValue(z_new)/init_val if(val_new <= cutoff_level): z1 = z_new val1 = val_new else: z0 = z_new val0 = val_new cutoff_z = (z0+z1)/2.0 return cutoff_z def _normalize(self): """ Normalizes the quad field function to the integral of 1 """ self.normalization = 1.0 step = self.z_max/(self.n_func_points - 1) sum_int = 0. self.func.clean() for ind in range(self.n_func_points): z = step*ind val = self.getPMQ_FuncValue(z) self.func.add(z,val) sum_int += val sum_int -= (self.getPMQ_FuncValue(0.) + self.getPMQ_FuncValue(step*(self.n_func_points - 1)))/2.0 sum_int *= 2.0*step self.normalization = 1.0/sum_int def getLimitsZ(self): """ Returns (z_min,z_max) tuple as longitudinal limits of the quad field. """ return (self.z_min,self.z_max) def pmq_func(self,z): """ This is PMQ function defined at p. 77 of the Trace3D manual. """ r1 = self.rad_in r2 = self.rad_out v1 = 1.0/math.sqrt(1.0+(z/r1)**2) v2 = 1.0/math.sqrt(1.0+(z/r2)**2) f = 0.5*(1-0.125*z*(1.0/r1+1.0/r2)*v1**2*v2**2*(v1**2+v1*v2+v1**2+4+8/v1/v2)/(v1+v2)) return f def getPMQ_FuncValue(self,z): """ Returns the total PMQ quad field distribution """ f = self.pmq_func(z - self.length/2) - self.pmq_func(z + self.length/2) return f def getFuncValue(self,z): """ Returns the quad's field at the distance z from the center """ if(abs(z) >= self.func.getMaxX()): return 0. return self.normalization*self.func.getY(abs(z)) def getFuncDerivative(self,z): """ Returns the derivative of the getFuncValue(z) """ if(abs(z) >= self.func.getMaxX()): return 0. return math.copysign(self.normalization*self.func.getYP(abs(z)),-z)