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