regular backup

This commit is contained in:
Jianshun Gao 2023-08-21 10:28:12 +02:00
parent 80d1de69f7
commit 602edcd6d5
3 changed files with 3008 additions and 102 deletions

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View File

@ -25,11 +25,15 @@ from scipy.special import erf, erfc
from scipy.special import gamma as gamfcn
from scipy.special import wofz
from scipy.optimize import curve_fit
from scipy.interpolate import CubicSpline
from scipy.ndimage import gaussian_filter
import xarray as xr
import copy
import matplotlib.pyplot as plt
log2 = log(2)
s2pi = sqrt(2*pi)
@ -109,6 +113,30 @@ def polylog2_2d(x, y=0.0, centerx=0.0, centery=0.0, amplitude=1.0, sigmax=1.0, s
return amplitude / 2 / np.pi / 1.20206 / max(tiny, sigmax * sigmay) * polylog(2, np.exp( -((x-centerx)**2/(2 * (sigmax)**2))-((y-centery)**2/( 2 * (sigmay)**2)) ))
def polylog_tab(pow, x, order=100):
"""Calculationg the polylog sum_i(x^i/i^pow), up to the order-th element of the sum
:param pow: power in denominator of sum
:type pow: int (can be float)
:param x: argument of Polylogarithm
:type x: float or 1D numpy array
:return: value of polylog(x)
:rtype: same as x, float or 1D numpy array
"""
sum = 0
for k in range(1,order):
sum += x ** k /k **pow
return sum
# Creating array for interpolation
x_int = np.linspace(0, 1.00001, 100000)
poly_tab = polylog_tab(2,x_int)
# Creating function interpolating between the Polylog values calculated before
polylog_int = CubicSpline(x_int, poly_tab)
def density_profile_BEC_2d(x, y=0.0, BEC_amplitude=1.0, thermal_amplitude=1.0, BEC_centerx=0.0, BEC_centery=0.0, thermal_centerx=0.0, thermal_centery=0.0,
BEC_sigmax=1.0, BEC_sigmay=1.0, thermal_sigmax=1.0, thermal_sigmay=1.0):
@ -318,115 +346,333 @@ class ThomasFermi2dModel(Model):
return update_param_vals(pars, self.prefix, **kwargs)
class DensityProfileBEC2dModel(Model):
fwhm_factor = 2*np.sqrt(2*np.log(2))
height_factor = 1./2*np.pi
def __init__(self, independent_vars=['x', 'y'], prefix='', nan_policy='raise',
**kwargs):
class DensityProfileBEC2dModel(lmfit.Model):
""" Fitting class to do 2D bimodal fit on OD of absorption images
"""
def __init__(self,
independent_vars=['x', 'y'],
prefix='',
nan_policy='raise',
atom_n_conv=144,
is_debug=False,
**kwargs
):
kwargs.update({'prefix': prefix, 'nan_policy': nan_policy,
'independent_vars': independent_vars})
self.atom_n_conv = atom_n_conv
self.is_debug=is_debug
super().__init__(density_profile_BEC_2d, **kwargs)
self._set_paramhints_prefix()
def _set_paramhints_prefix(self):
# self.set_param_hint('BEC_sigmax', min=0)
self.set_param_hint('deltax', min=0)
self.set_param_hint('BEC_sigmax', expr=f'3 * {self.prefix}thermal_sigmax - {self.prefix}deltax')
self.set_param_hint('BEC_sigmay', min=0)
self.set_param_hint('thermal_sigmax', min=0)
# self.set_param_hint('thermal_sigmay', min=0)
self.set_param_hint('BEC_amplitude', min=0)
self.set_param_hint('thermal_amplitude', min=0)
self.set_param_hint('thermalAspectRatio', min=0.8, max=1.2)
self.set_param_hint('thermal_sigmay', expr=f'{self.prefix}thermalAspectRatio * {self.prefix}thermal_sigmax')
# self.set_param_hint('betax', value=0)
# self.set_param_hint('BEC_centerx', expr=f'{self.prefix}thermal_sigmax - {self.prefix}betax')
self.set_param_hint('condensate_fraction', expr=f'{self.prefix}BEC_amplitude / ({self.prefix}BEC_amplitude + {self.prefix}thermal_amplitude)')
self.set_param_hint('amp_bec', min=0)
self.set_param_hint('amp_th', min=0)
self.set_param_hint('x0_bec', min=0)
self.set_param_hint('y0_bec', min=0)
self.set_param_hint('x0_th', min=0)
self.set_param_hint('y0_th', min=0)
self.set_param_hint('sigmax_bec', min=0)
self.set_param_hint('sigmay_bec', min=0)
self.set_param_hint('sigma_th', min=0)
def guess(self, data, x, y, negative=False, pureBECThreshold=0.5, noBECThThreshold=0.0, **kwargs):
"""Estimate initial model parameter values from data."""
fitModel = TwoGaussian2dModel()
pars = fitModel.guess(data, x=x, y=y, negative=negative)
pars['A_amplitude'].set(min=0)
pars['B_amplitude'].set(min=0)
pars['A_centerx'].set(min=pars['A_centerx'].value - 3 * pars['A_sigmax'],
max=pars['A_centerx'].value + 3 * pars['A_sigmax'],)
pars['A_centery'].set(min=pars['A_centery'].value - 3 * pars['A_sigmay'],
max=pars['A_centery'].value + 3 * pars['A_sigmay'],)
pars['B_centerx'].set(min=pars['B_centerx'].value - 3 * pars['B_sigmax'],
max=pars['B_centerx'].value + 3 * pars['B_sigmax'],)
pars['B_centery'].set(min=pars['B_centery'].value - 3 * pars['B_sigmay'],
max=pars['B_centery'].value + 3 * pars['B_sigmay'],)
fitResult = fitModel.fit(data, x=x, y=y, params=pars, **kwargs)
pars_guess = fitResult.params
BEC_amplitude = pars_guess['A_amplitude'].value
thermal_amplitude = pars_guess['B_amplitude'].value
pars = self.make_params(BEC_amplitude=BEC_amplitude,
thermal_amplitude=thermal_amplitude,
BEC_centerx=pars_guess['A_centerx'].value, BEC_centery=pars_guess['A_centery'].value,
# BEC_sigmax=(pars_guess['A_sigmax'].value / 2.355),
deltax = 3 * (pars_guess['B_sigmax'].value * s2) - (pars_guess['A_sigmax'].value / 2.355),
BEC_sigmay=(pars_guess['A_sigmay'].value / 2.355),
thermal_centerx=pars_guess['B_centerx'].value, thermal_centery=pars_guess['B_centery'].value,
thermal_sigmax=(pars_guess['B_sigmax'].value * s2),
thermalAspectRatio=(pars_guess['B_sigmax'].value * s2) / (pars_guess['B_sigmay'].value * s2)
# thermal_sigmay=(pars_guess['B_sigmay'].value * s2)
)
nBEC = pars[f'{self.prefix}BEC_amplitude'] / 2 / np.pi / 5.546 / pars[f'{self.prefix}BEC_sigmay'] / pars[f'{self.prefix}BEC_sigmax']
if (pars[f'{self.prefix}condensate_fraction']>0.95) and (np.max(data) > 1.05 * nBEC):
temp = ((np.max(data) - nBEC) * s2pi * pars[f'{self.prefix}thermal_sigmay'] / pars[f'{self.prefix}thermal_sigmax'])
if temp > pars[f'{self.prefix}BEC_amplitude']:
pars[f'{self.prefix}thermal_amplitude'].set(value=pars[f'{self.prefix}BEC_amplitude'] / 2)
def guess(self, data, x, y, **kwargs):
"""
Estimate and create initial model parameters for 2d fit
:param data: Flattened 2d array, flattened from array with (x,y) -->
:param x:
:param y:
:param kwargs:
:return:
"""
#
# global X_guess
# global bval_1d
x_width = len(np.unique(x))
y_width = len(np.unique(y))
data = np.reshape(data, (y_width, x_width))
data = data.T
shape = np.shape(data)
cut_width = np.max(shape)
thresh = self.calc_thresh(data)
center = self.calc_cen(thresh)
BEC_width_guess = self.guess_BEC_width(thresh, center)
if BEC_width_guess[0] < BEC_width_guess[1]:
if self.is_debug:
print(f'x smaller y')
s_width_ind = 0
X_guess = np.sum(data[:, round(center[1] - BEC_width_guess[1]/2) : round(center[1] + BEC_width_guess[1]/2)], 1) / len(data[0,round(center[1] - BEC_width_guess[1]/2) : round(center[1] + BEC_width_guess[1]/2)])
else:
if self.is_debug:
print(f'y smaller x')
s_width_ind = 1
X_guess = np.sum(data[round(center[0] - BEC_width_guess[0]/2) : round(center[0] + BEC_width_guess[0]/2), :], 0) / len(data[0,round(center[0] - BEC_width_guess[0]/2) : round(center[0] + BEC_width_guess[0]/2)])
if self.is_debug:
print(f'center = {center}')
print(f'BEC widths: {BEC_width_guess}')
plt.plot(X_guess)
plt.show()
x = np.linspace(0, len(X_guess), len(X_guess))
max_val = np.max(X_guess)
fitmodel_1d = lmfit.Model(density_1d, independent_vars=['x'])
params_1d = lmfit.Parameters()
params_1d.add_many(
('x0_bec', center[s_width_ind], True,0, 200),
('x0_th',center[s_width_ind], True,0, 200),
('amp_bec', 0.7 * max_val, True, 0, 1.3 * max_val),
('amp_th', 0.3 * max_val, True, 0, 1.3 * max_val),
('deltax', 3*BEC_width_guess[s_width_ind], True, 0,cut_width),
# ('sigma_bec',BEC_width_guess[i,j,0]/1.22, True, 0, 50)
('sigma_bec',BEC_width_guess[s_width_ind]/1.22, True, 0, 50)
)
params_1d.add('sigma_th', 3*BEC_width_guess[0], min=0, expr=f'0.632*sigma_bec + 0.518*deltax')
res_1d = fitmodel_1d.fit(X_guess, x=x, params=params_1d)
if self.is_debug:
params_1d.pretty_print()
self.print_bval(res_1d)
bval_1d = res_1d.best_values
amp_conv_1d_2d = np.max(gaussian_filter(data, sigma=1)) / (bval_1d['amp_bec'] + bval_1d['amp_th'])
max_val = np.max(data)
params = self.make_params()
if bval_1d['amp_th']/bval_1d['amp_bec'] > 3:
print(f'Image seems to be purely thermal (guessed from 1d fit amplitude)')
params[f'{self.prefix}amp_bec'].set(value=0, vary=False)
params[f'{self.prefix}amp_th'].set(value=amp_conv_1d_2d * bval_1d['amp_th'], max=1.3 * max_val, vary=True)
params[f'{self.prefix}x0_bec'].set(value=1, vary=False)
params[f'{self.prefix}y0_bec'].set(value=1, vary=False)
params[f'{self.prefix}x0_th'].set(value=center[0], min=center[0] -10, max=center[0] + 10, vary=True)
params[f'{self.prefix}y0_th'].set(value=center[1], min=center[1] -10, max=center[1] + 10, vary=True)
params[f'{self.prefix}sigmax_bec'].set(value=1, vary=False)
params[f'{self.prefix}sigmay_bec'].set(value=1, vary=False)
params[f'{self.prefix}sigma_th'].set(value=bval_1d['sigma_th'], max=cut_width, vary=True)
elif bval_1d['amp_bec']/bval_1d['amp_th'] > 10:
print('Image seems to be pure BEC (guessed from 1d fit amplitude)')
params[f'{self.prefix}amp_bec'].set(value=amp_conv_1d_2d * bval_1d['amp_bec'], max=1.3 * max_val, vary=True)
params[f'{self.prefix}amp_th'].set(value=0, vary=False)
params[f'{self.prefix}x0_bec'].set(value=center[0], min=center[0] -10, max=center[0] + 10, vary=True)
params[f'{self.prefix}y0_bec'].set(value=center[1], min=center[1] -10, max=center[1] + 10, vary=True)
params[f'{self.prefix}x0_th'].set(value=1, vary=False)
params[f'{self.prefix}y0_th'].set(value=1, vary=False)
params[f'{self.prefix}sigma_th'].set(value=1, vary=False)
if s_width_ind == 0:
params[f'{self.prefix}sigmax_bec'].set(value=bval_1d['sigma_bec'], max= 2*BEC_width_guess[0], vary=True)
params[f'{self.prefix}sigmay_bec'].set(value=BEC_width_guess[1]/1.22, max= 2*BEC_width_guess[1], vary=True)
elif s_width_ind == 1:
params[f'{self.prefix}sigmax_bec'].set(value=BEC_width_guess[0]/1.22, max= 2*BEC_width_guess[0], vary=True)
params[f'{self.prefix}sigmay_bec'].set(value=bval_1d['sigma_bec'], max= 2*BEC_width_guess[1], vary=True)
else:
pars[f'{self.prefix}thermal_amplitude'].set(value=temp * 10)
if BEC_amplitude / (thermal_amplitude + BEC_amplitude) > pureBECThreshold:
pars[f'{self.prefix}thermal_amplitude'].set(value=0)
pars[f'{self.prefix}BEC_amplitude'].set(value=(thermal_amplitude + BEC_amplitude))
if BEC_amplitude / (thermal_amplitude + BEC_amplitude) < noBECThThreshold:
pars[f'{self.prefix}BEC_amplitude'].set(value=0)
pars[f'{self.prefix}thermal_amplitude'].set(value=(thermal_amplitude + BEC_amplitude))
pars[f'{self.prefix}BEC_centerx'].set(
min=pars[f'{self.prefix}BEC_centerx'].value - 10 * pars[f'{self.prefix}BEC_sigmax'].value,
max=pars[f'{self.prefix}BEC_centerx'].value + 10 * pars[f'{self.prefix}BEC_sigmax'].value,
)
pars[f'{self.prefix}thermal_centerx'].set(
min=pars[f'{self.prefix}thermal_centerx'].value - 3 * pars[f'{self.prefix}thermal_sigmax'].value,
max=pars[f'{self.prefix}thermal_centerx'].value + 3 * pars[f'{self.prefix}thermal_sigmax'].value,
)
pars[f'{self.prefix}BEC_centery'].set(
min=pars[f'{self.prefix}BEC_centery'].value - 10 * pars[f'{self.prefix}BEC_sigmay'].value,
max=pars[f'{self.prefix}BEC_centery'].value + 10 * pars[f'{self.prefix}BEC_sigmay'].value,
)
pars[f'{self.prefix}thermal_centery'].set(
min=pars[f'{self.prefix}thermal_centery'].value - 3 * pars[f'{self.prefix}thermal_sigmay'].value,
max=pars[f'{self.prefix}thermal_centery'].value + 3 * pars[f'{self.prefix}thermal_sigmay'].value,
)
pars[f'{self.prefix}BEC_sigmay'].set(
max=5 * pars[f'{self.prefix}BEC_sigmay'].value,
)
pars[f'{self.prefix}thermal_sigmax'].set(
max=5 * pars[f'{self.prefix}thermal_sigmax'].value,
)
return update_param_vals(pars, self.prefix, **kwargs)
print('Error in small width BEC recogintion, s_width_ind should be 0 or 1')
else:
params[f'{self.prefix}amp_bec'].set(value=amp_conv_1d_2d * bval_1d['amp_bec'], max=1.3 * max_val, vary=True)
params[f'{self.prefix}amp_th'].set(value=amp_conv_1d_2d * bval_1d['amp_th'], max=1.3 * max_val, vary=True)
params[f'{self.prefix}x0_bec'].set(value=center[0], min=center[0] -10, max=center[0] + 10, vary=True)
params[f'{self.prefix}y0_bec'].set(value=center[1], min=center[1] -10, max=center[1] + 10, vary=True)
params[f'{self.prefix}x0_th'].set(value=center[0], min=center[0] -10, max=center[0] + 10, vary=True)
params[f'{self.prefix}y0_th'].set(value=center[1], min=center[1] -10, max=center[1] + 10, vary=True)
params[f'{self.prefix}sigma_th'].set(value=bval_1d['sigma_th'], max=cut_width, vary=True)
if s_width_ind == 0:
params[f'{self.prefix}sigmax_bec'].set(value=bval_1d['sigma_bec'], max= 2*BEC_width_guess[0], vary=True)
params[f'{self.prefix}sigmay_bec'].set(value=BEC_width_guess[1]/1.22, max= 2*BEC_width_guess[1], vary=True)
elif s_width_ind == 1:
params[f'{self.prefix}sigmax_bec'].set(value=BEC_width_guess[0]/1.22, max= 2*BEC_width_guess[0], vary=True)
params[f'{self.prefix}sigmay_bec'].set(value=bval_1d['sigma_bec'], max= 2*BEC_width_guess[1], vary=True)
else:
print('Error in small width BEC recogintion, s_width_ind should be 0 or 1')
if self.is_debug:
print('')
print('Init Params')
params.pretty_print()
return lmfit.models.update_param_vals(params, self.prefix, **kwargs)
def fit(self, data_1d, **kwargs):
res = super().fit(data_1d, **kwargs)
if self.is_debug:
print('bval first fit')
self.print_bval(res)
if res.params['amp_bec'].vary and res.params['amp_th'].vary:
bval = res.best_values
sigma_cut = max(bval['sigmay_bec'], bval['sigmax_bec'])
tf_fit = ThomasFermi_2d(kwargs['x'],kwargs['y'],centerx=bval['x0_bec'], centery=bval['y0_bec'], amplitude=bval['amp_bec'], sigmax=bval['sigmax_bec'], sigmay=bval['sigmay_bec'])
tf_fit_2 = ThomasFermi_2d(kwargs['x'],kwargs['y'],centerx=bval['x0_bec'], centery=bval['y0_bec'], amplitude=bval['amp_bec'], sigmax=1.5 * sigma_cut, sigmay=1.5*sigma_cut)
mask = np.where(tf_fit > 0, np.nan, data_1d)
#mask[i,j] = gaussian_filter(mask[i,j], sigma = 0.4)
mask = np.where(tf_fit_2 > 0, mask, np.nan)
N_c = np.nansum(mask)
# conversion N_count to Pixels
N_a = self.atom_n_conv * N_c
if N_a < 6615:
print('No thermal part detected, performing fit without thermal function')
params = res.params
params[f'{self.prefix}amp_th'].set(value=0, vary=False)
params[f'{self.prefix}x0_th'].set(value=1, vary=False)
params[f'{self.prefix}y0_th'].set(value=1, vary=False)
params[f'{self.prefix}sigma_th'].set(value=1, vary=False)
res = super().fit(data_1d, x=kwargs['x'], y=kwargs['y'], params=params)
return res
return res
def calc_thresh(self, data, thresh_val=0.3, sigma=0.4):
"""Returns thresholded binary image after blurring to guess BEC size
:param data: 2d image or 1D or 2D array containing 2d images
:type data: 2d, 3d or 4d numpy array
:param thresh_val: relative threshhold value for binarization with respect to maximum of blurred image
:param sigma: sigma of gaussian blur filter (see scipy.ndimage.gaussian_filter
:return: binary 2d image or 1D or 2D array containing 2d images
:rtype: 2d, 3d or 4d numpy array
"""
shape = np.shape(data)
thresh = np.zeros(shape)
blurred = gaussian_filter(data, sigma=sigma)
if len(shape) == 4:
for i in range(0,shape[0]):
for j in range(0, shape[1]):
thresh[i,j] = np.where(blurred[i,j] < np.max(blurred[i,j])*thresh_val, 0, 1)
elif len(shape) == 3:
for i in range(0,shape[0]):
thresh[i] = np.where(blurred[i] < np.max(blurred[i])*thresh_val, 0, 1)
elif len(shape) == 2:
thresh = np.where(blurred < np.max(blurred)*thresh_val, 0, 1)
else:
print("Shape of data is wrong, output is empty")
return thresh
def calc_cen(self, thresh1):
"""
returns array: [X_center,Y_center]
"""
cen = np.zeros(2)
(Y,X) = np.shape(thresh1)
thresh1 = thresh1 /np.sum(thresh1)
# marginal distributions
dx = np.sum(thresh1, 0)
dy = np.sum(thresh1, 1)
# expected values
cen[0] = np.sum(dx * np.arange(X))
cen[1] = np.sum(dy * np.arange(Y))
return cen
def guess_BEC_width(self, thresh, center):
"""
returns width of thresholded area along both axis through the center with shape of thresh and [X_width, Y_width] for each image
"""
shape = np.shape(thresh)
if len(shape) == 2:
BEC_width_guess = np.array([np.sum(thresh[round(center[1]), :]), np.sum(thresh[:, round(center[0])]) ])
elif len(shape) == 3:
BEC_width_guess = np.zeros((shape[0], 2))
for i in range(0, shape[0]):
BEC_width_guess[i, 0] = np.sum(thresh[i, round(center[i,j,1]), :])
BEC_width_guess[i, 1] = np.sum(thresh[i, :, round(center[i,j,0])])
elif len(shape) == 4:
BEC_width_guess = np.zeros((shape[0], shape[1], 2))
for i in range(0, shape[0]):
for j in range(0, shape[1]):
BEC_width_guess[i, j, 0] = np.sum(thresh[i, j, round(center[i,j,1]), :])
BEC_width_guess[i, j, 1] = np.sum(thresh[i, j, :, round(center[i,j,0])])
else:
print("Shape of data is wrong, output is empty")
return BEC_width_guess
def cond_frac(self, results):
"""Returns condensate fraction"""
bval = results.best_values
tf_fit = ThomasFermi_2d(X, Y, centerx=bval['x0_bec'], centery=bval['y0_bec'], amplitude=bval['amp_bec'], sigmax=bval['sigmax_bec'], sigmay=bval['sigmay_bec'])
N_bec = np.sum(tf_fit)
fit = density_profile_BEC_2d(X, Y, **bval)
N_ges = np.sum(fit)
return N_bec/N_ges
def return_atom_number(self, result, X, Y, is_print=True):
"""Printing fitted atom number in bec + thermal state"""
bval = result.best_values
tf_fit = ThomasFermi_2d(X,Y,centerx=bval['x0_bec'], centery=bval['y0_bec'], amplitude=bval['amp_bec'], sigmax=bval['sigmax_bec'], sigmay=bval['sigmay_bec'])
N_bec = self.atom_n_conv * np.sum(tf_fit)
th_fit = polylog2_2d(X, Y, centerx=bval['x0_th'], centery=bval['y0_th'], amplitude=bval['amp_th'], sigmax=bval['sigma_th'], sigmay=bval['sigma_th'])
N_th = self.atom_n_conv * np.sum(th_fit)
# fit = density_profile_BEC_2d(X,Y, **bval)
# N_ges = self.atom_n_conv * np.sum(fit)
if is_print:
print()
print('Atom numbers:')
print(f' N_bec: {N_bec :.0f}')
print(f' N_th: {N_th :.0f}')
print(f' N_ges: {N_bec + N_th :.0f}')
print(f' Cond. frac: {N_bec/(N_bec + N_th):.2f}')
print('')
def return_temperature(self, result, omg, tof, is_print=True, eff_pix=2.472e-6):
"""Returns temperature of thermal cloud"""
R_th = result.best_values['sigma_th'] * eff_pix * np.sqrt(2)
print(R_th)
T = R_th**2 * 164*const.u/const.k * (1/omg**2 + tof**2)**(-1)
if is_print:
print(f'Temperature: {T*1e9:.2f} nK')
return T
def print_bval(self, res_s):
"""nicely print best fitted values + init values + bounds """
keys = res_s.best_values.keys()
bval = res_s.best_values
init = res_s.init_params
for item in keys:
print(f'{item}: {bval[item]:.3f}, (init = {init[item].value:.3f}), bounds = [{init[item].min:.2f} : {init[item].max :.2f}] ')
print('')
class NewFitModel(Model):