Commit 55e48965fa90f639142288f626435ba7fe192e6b
1 parent
8cb3d71c15
Exists in
master
abs
Showing 1 changed file with 2 additions and 2 deletions Inline Diff
fit_waist.py
| 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) | |
| data_waist.append([int(f[-7:-4]), Ppopt[3]]) | 27 | 27 | data_waist.append([int(f[-7:-4]), abs(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 |