bmarechal / PM100_waist

Download as
Browse Code ยป

Commit f808f4392a5680ed0808eeae115f3708c4f6745f

Authored by bmarechal
1 parent c81fde478c
Exists in master

use more robust trf method instead lms

Showing 1 changed file with 2 additions and 2 deletions Inline Diff

from scipy.optimize import curve_fit 1 1 from scipy.optimize import curve_fit
import csv, numpy, glob 2 2 import csv, numpy, glob
from scipy.special import erf 3 3 from scipy.special import erf
import matplotlib.pyplot as plt 4 4 import matplotlib.pyplot as plt
5 5
'''power function to optimize''' 6 6 '''power function to optimize'''
def P(x, Po, Pmax, xo, w): 7 7 def P(x, Po, Pmax, xo, w):
return Po+0.5*Pmax*(1.-erf(2.**0.5*(x-xo)/w)) 8 8 return Po+0.5*Pmax*(1.-erf(2.**0.5*(x-xo)/w))
9 9
10 10
'''load and fit beam section''' 11 11 '''load and fit beam section'''
files = glob.glob('*.dat') 12 12 files = glob.glob('*.dat')
files.sort() 13 13 files.sort()
data_waist = [] 14 14 data_waist = []
plt.close() 15 15 plt.close()
fig, p = plt.subplots(2, 1) 16 16 fig, p = plt.subplots(2, 1)
for f in files: 17 17 for f in files:
with open(f, 'r') as dest_f: 18 18 with open(f, 'r') as dest_f:
raw = csv.reader(dest_f, delimiter = '\t', quotechar = '"') 19 19 raw = csv.reader(dest_f, delimiter = '\t', quotechar = '"')
data = [value for value in raw] 20 20 data = [value for value in raw]
data = numpy.asarray(data, dtype = float) 21 21 data = numpy.asarray(data, dtype = float)
xmes = data[:,0] 22 22 xmes = data[:,0]
Pmes = data[:,1] 23 23 Pmes = data[:,1]
24 24
'''optimization with non-linear least squares method''' 25 25 '''optimization with non-linear least squares method'''
Ppopt, Pcov = curve_fit(P, xmes, Pmes) 26 26 Ppopt, Pcov = curve_fit(P, xmes, Pmes, method = 'trf')
data_waist.append([int(f[-7:-4]), abs(Ppopt[3])]) 27 27 data_waist.append([int(f[-7:-4]), Ppopt[3]])
28 28
'''plot''' 29 29 '''plot'''
p[0].plot(xmes, Pmes, 'o') 30 30 p[0].plot(xmes, Pmes, 'o')
p[0].plot(numpy.linspace(xmes[0], xmes[-1], 100), P(numpy.linspace(xmes[0], xmes[-1], 100), *Ppopt)) 31 31 p[0].plot(numpy.linspace(xmes[0], xmes[-1], 100), P(numpy.linspace(xmes[0], xmes[-1], 100), *Ppopt))
32 32
p[0].grid() 33 33 p[0].grid()
34 34
'''return waist(z) table''' 35 35 '''return waist(z) table'''
data_waist = numpy.asarray(data_waist, dtype = float) 36 36 data_waist = numpy.asarray(data_waist, dtype = float)
print(data_waist) 37 37 print(data_waist)
38 38
'''waist function to optimize''' 39 39 '''waist function to optimize'''
def W(z, w0, z0): 40 40 def W(z, w0, z0):
return w0*(1.+((z-z0)*1542e-6/(numpy.pi*w0**2))**2)**0.5 41 41 return w0*(1.+((z-z0)*1542e-6/(numpy.pi*w0**2))**2)**0.5
42 42
popt, cov = curve_fit(W, data_waist[:,0], data_waist[:,1]) 43 43 popt, cov = curve_fit(W, data_waist[:,0], data_waist[:,1])