2023-05-04 13:47:33 +02:00
import numpy as np
2023-05-07 00:38:52 +02:00
from uncertainties import ufloat
2023-05-04 13:47:33 +02:00
import lmfit
from lmfit . models import ( ConstantModel , ComplexConstantModel , LinearModel , QuadraticModel ,
PolynomialModel , SineModel , GaussianModel , Gaussian2dModel , LorentzianModel ,
SplitLorentzianModel , VoigtModel , PseudoVoigtModel , MoffatModel ,
Pearson7Model , StudentsTModel , BreitWignerModel , LognormalModel ,
DampedOscillatorModel , ExponentialGaussianModel , SkewedGaussianModel ,
SkewedVoigtModel , ThermalDistributionModel , DoniachModel , PowerLawModel ,
ExponentialModel , StepModel , RectangleModel , ExpressionModel , DampedHarmonicOscillatorModel )
from lmfit . models import ( guess_from_peak , guess_from_peak2d , fwhm_expr , height_expr ,
update_param_vals )
from lmfit . lineshapes import ( not_zero , breit_wigner , damped_oscillator , dho , doniach ,
expgaussian , exponential , gaussian , gaussian2d ,
linear , lognormal , lorentzian , moffat , parabolic ,
pearson7 , powerlaw , pvoigt , rectangle , sine ,
skewed_gaussian , skewed_voigt , split_lorentzian , step ,
students_t , thermal_distribution , tiny , voigt )
from lmfit import Model
import numpy as np
from numpy import ( arctan , copysign , cos , exp , isclose , isnan , log , pi , real ,
sin , sqrt , where )
from scipy . special import erf , erfc
from scipy . special import gamma as gamfcn
from scipy . special import wofz
from scipy . optimize import curve_fit
2023-08-21 10:28:12 +02:00
from scipy . interpolate import CubicSpline
2023-09-08 14:13:27 +02:00
from scipy . ndimage import gaussian_filter , rotate
2023-09-14 14:13:49 +02:00
import scipy . constants as const
2023-05-04 13:47:33 +02:00
import xarray as xr
2023-06-02 18:42:18 +02:00
import copy
2023-08-21 10:28:12 +02:00
import matplotlib . pyplot as plt
2023-05-04 13:47:33 +02:00
log2 = log ( 2 )
s2pi = sqrt ( 2 * pi )
s2 = sqrt ( 2.0 )
def gaussianWithOffset ( x , amplitude = 1.0 , center = 0.0 , sigma = 1.0 , offset = 0.0 ) :
""" Return a 1-dimensional Gaussian function with an offset.
gaussian ( x , amplitude , center , sigma ) =
( amplitude / ( s2pi * sigma ) ) * exp ( - ( 1.0 * x - center ) * * 2 / ( 2 * sigma * * 2 ) )
"""
return ( ( amplitude / ( max ( tiny , s2pi * sigma ) ) )
* exp ( - ( 1.0 * x - center ) * * 2 / max ( tiny , ( 2 * sigma * * 2 ) ) ) + offset )
def lorentzianWithOffset ( x , amplitude = 1.0 , center = 0.0 , sigma = 1.0 , offset = 0.0 ) :
return ( ( amplitude / ( 1 + ( ( 1.0 * x - center ) / max ( tiny , sigma ) ) * * 2 ) )
/ max ( tiny , ( pi * sigma ) ) + offset )
def exponentialWithOffset ( x , amplitude = 1.0 , decay = 1.0 , offset = 0.0 ) :
decay = not_zero ( decay )
return amplitude * exp ( - x / decay ) + offset
def expansion ( x , amplitude = 1.0 , offset = 0.0 ) :
return np . sqrt ( amplitude * x * x + offset )
def dampingOscillation ( x , center = 0 , amplitude = 1.0 , frequency = 1.0 , decay = 1.0 , offset = 0.0 ) :
return amplitude * np . exp ( - decay * x ) * np . sin ( 2 * np . pi * frequency * ( x - center ) ) + offset
2023-05-04 18:32:17 +02:00
def two_gaussian2d ( x , y = 0.0 , A_amplitude = 1.0 , A_centerx = 0.0 , A_centery = 0.0 , A_sigmax = 1.0 ,
A_sigmay = 1.0 , B_amplitude = 1.0 , B_centerx = 0.0 , B_centery = 0.0 , B_sigmax = 1.0 ,
B_sigmay = 1.0 ) :
""" Return a 2-dimensional Gaussian function.
gaussian2d ( x , y , amplitude , centerx , centery , sigmax , sigmay ) =
amplitude / ( 2 * pi * sigmax * sigmay ) * exp ( - ( x - centerx ) * * 2 / ( 2 * sigmax * * 2 )
- ( y - centery ) * * 2 / ( 2 * sigmay * * 2 ) )
"""
z = A_amplitude * ( gaussian ( x , amplitude = 1 , center = A_centerx , sigma = A_sigmax ) *
gaussian ( y , amplitude = 1 , center = A_centery , sigma = A_sigmay ) )
z + = B_amplitude * ( gaussian ( x , amplitude = 1 , center = B_centerx , sigma = B_sigmax ) *
gaussian ( y , amplitude = 1 , center = B_centery , sigma = B_sigmay ) )
return z
2023-05-04 13:47:33 +02:00
2023-08-21 10:28:12 +02:00
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 1 D numpy array
: return : value of polylog ( x )
: rtype : same as x , float or 1 D 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 )
2023-08-21 11:26:02 +02:00
def thermal ( x , x0 , amp , sigma ) :
2023-08-22 09:54:04 +02:00
""" Calculating thermal density distribution in 1D (scaled such that if amp=1, return = 1)
: param x : axis
: type x : float or 1 d array
: param x0 : position of peak along axis
: type x0 : float
: param amp : amplitude of function
: type amp : float
: param sigma : width of function
: type sigma : float
: return : calculated function value
: rtype : float or 1 D array
"""
2023-08-21 11:26:02 +02:00
res = np . exp ( - 0.5 * ( x - x0 ) * * 2 / sigma * * 2 )
return amp / 1.643 * polylog_int ( res )
def Thomas_Fermi_1d ( x , x0 , amp , sigma ) :
res = ( 1 - ( ( x - x0 ) / sigma ) * * 2 )
res = np . where ( res > 0 , res , 0 )
res = res * * ( 3 / 2 )
return amp * res
def density_1d ( x , x0_bec , x0_th , amp_bec , amp_th , sigma_bec , sigma_th ) :
return thermal ( x , x0_th , amp_th , sigma_th ) + Thomas_Fermi_1d ( x , x0_bec , amp_bec , sigma_bec )
def polylog ( pow , x ) :
order = 15
sum = 0
for k in range ( 1 , order ) :
sum + = x * * k / k * * pow
return sum
def ThomasFermi_2d ( x , y = 0.0 , centerx = 0.0 , centery = 0.0 , amplitude = 1.0 , sigmax = 1.0 , sigmay = 1.0 ) :
res = ( 1 - ( ( x - centerx ) / ( sigmax ) ) * * 2 - ( ( y - centery ) / ( sigmay ) ) * * 2 )
res = np . where ( res > 0 , res , 0 )
res = res * * ( 3 / 2 )
return amplitude * res
def polylog2_2d ( x , y = 0.0 , centerx = 0.0 , centery = 0.0 , amplitude = 1.0 , sigmax = 1.0 , sigmay = 1.0 ) :
## Approximation of the polylog function with 2D gaussian as argument. -> discribes the thermal part of the cloud
return amplitude / 1.643 * polylog_int ( np . exp ( - ( ( x - centerx ) * * 2 / ( 2 * sigmax * * 2 ) ) - ( ( y - centery ) * * 2 / ( 2 * sigmay * * 2 ) ) ) )
2023-09-08 14:13:27 +02:00
def rotate_coord ( x , y , rot_angle ) :
rot_angle * = 2 * np . pi / 360
x_og = x
x = x_og * np . cos ( rot_angle ) + y * np . sin ( rot_angle )
y = - x_og * np . sin ( rot_angle ) + y * np . cos ( rot_angle )
return x , y
2023-08-21 11:26:02 +02:00
def density_profile_BEC_2d ( x , y = 0.0 , amp_bec = 1.0 , amp_th = 1.0 , x0_bec = 0.0 , y0_bec = 0.0 , x0_th = 0.0 , y0_th = 0.0 ,
2023-09-08 14:13:27 +02:00
sigmax_bec = 1.0 , sigmay_bec = 1.0 , sigma_th = 1.0 , rot_angle = None ) :
if rot_angle is not None :
x , y = rotate_coord ( x , y , rot_angle )
x0_bec , y0_bec = rotate_coord ( x0_bec , y0_bec , rot_angle )
x0_th , y0_th = rotate_coord ( x0_th , y0_th , rot_angle )
2023-08-21 11:26:02 +02:00
return ThomasFermi_2d ( x = x , y = y , centerx = x0_bec , centery = y0_bec ,
amplitude = amp_bec , sigmax = sigmax_bec , sigmay = sigmay_bec
) + polylog2_2d ( x = x , y = y , centerx = x0_th , centery = y0_th ,
amplitude = amp_th , sigmax = sigma_th , sigmay = sigma_th )
2023-05-23 22:14:03 +02:00
2023-05-04 13:47:33 +02:00
class GaussianWithOffsetModel ( Model ) :
fwhm_factor = 2 * np . sqrt ( 2 * np . log ( 2 ) )
height_factor = 1. / np . sqrt ( 2 * np . pi )
def __init__ ( self , independent_vars = [ ' x ' ] , nan_policy = ' raise ' , prefix = ' ' , name = None , * * kwargs ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( gaussianWithOffset , * * kwargs )
self . _set_paramhints_prefix ( )
def _set_paramhints_prefix ( self ) :
self . set_param_hint ( ' sigma ' , min = 0 )
self . set_param_hint ( ' fwhm ' , expr = fwhm_expr ( self ) )
self . set_param_hint ( ' height ' , expr = height_expr ( self ) )
def guess ( self , data , x , negative = False , * * kwargs ) :
offset = np . min ( data )
data = data - offset
pars = guess_from_peak ( self , data , x , negative )
pars . add ( ' offset ' , value = offset )
return update_param_vals ( pars , self . prefix , * * kwargs )
class LorentzianWithOffsetModel ( Model ) :
fwhm_factor = 2.0
height_factor = 1. / np . pi
def __init__ ( self , independent_vars = [ ' x ' ] , prefix = ' ' , nan_policy = ' raise ' ,
* * kwargs ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( lorentzianWithOffset , * * kwargs )
self . _set_paramhints_prefix ( )
def _set_paramhints_prefix ( self ) :
self . set_param_hint ( ' sigma ' , min = 0 )
self . set_param_hint ( ' fwhm ' , expr = fwhm_expr ( self ) )
self . set_param_hint ( ' height ' , expr = height_expr ( self ) )
def guess ( self , data , x , negative = False , * * kwargs ) :
""" Estimate initial model parameter values from data. """
offset = np . min ( data )
data = data - offset
pars = guess_from_peak ( self , data , x , negative , ampscale = 1.25 )
pars . add ( ' offset ' , value = offset )
return update_param_vals ( pars , self . prefix , * * kwargs )
class ExponentialWithOffsetModel ( Model ) :
def __init__ ( self , independent_vars = [ ' x ' ] , prefix = ' ' , nan_policy = ' raise ' ,
* * kwargs ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( exponentialWithOffset , * * kwargs )
def guess ( self , data , x , * * kwargs ) :
""" Estimate initial model parameter values from data. """
offset = np . min ( data )
data = data - offset
try :
sval , oval = np . polyfit ( x , np . log ( abs ( data ) + 1.e-15 ) , 1 )
except TypeError :
sval , oval = 1. , np . log ( abs ( max ( data ) + 1.e-9 ) )
pars = self . make_params ( amplitude = np . exp ( oval ) , decay = - 1.0 / sval )
pars . add ( ' offset ' , value = offset )
return update_param_vals ( pars , self . prefix , * * kwargs )
class ExpansionModel ( Model ) :
def __init__ ( self , independent_vars = [ ' x ' ] , prefix = ' ' , nan_policy = ' raise ' ,
* * kwargs ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( expansion , * * kwargs )
def guess ( self , data , x , * * kwargs ) :
""" Estimate initial model parameter values from data. """
popt1 , pcov1 = curve_fit ( expansion , x , data )
pars = self . make_params ( amplitude = popt1 [ 0 ] , offset = popt1 [ 1 ] )
return update_param_vals ( pars , self . prefix , * * kwargs )
class DampingOscillationModel ( Model ) :
def __init__ ( self , independent_vars = [ ' x ' ] , prefix = ' ' , nan_policy = ' raise ' ,
* * kwargs ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( dampingOscillation , * * kwargs )
def guess ( self , data , x , * * kwargs ) :
""" Estimate initial model parameter values from data. """
try :
popt1 , pcov1 = curve_fit ( dampingOscillation , x , data , np . array ( 0 , 5 , 5e2 , 1e3 , 16 ) )
pars = self . make_params ( center = popt1 [ 0 ] , amplitude = popt1 [ 1 ] , frequency = popt1 [ 2 ] , decay = popt1 [ 3 ] , offset = popt1 [ 4 ] )
except :
pars = self . make_params ( center = 0 , amplitude = 5.0 , frequency = 5e2 , decay = 1.0e3 , offset = 16.0 )
return update_param_vals ( pars , self . prefix , * * kwargs )
2023-05-04 18:32:17 +02:00
class TwoGaussian2dModel ( 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 ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
self . helperModel = Gaussian2dModel ( )
super ( ) . __init__ ( two_gaussian2d , * * kwargs )
2023-05-07 00:38:52 +02:00
self . _set_paramhints_prefix ( )
def _set_paramhints_prefix ( self ) :
self . set_param_hint ( ' delta ' , value = - 1 , max = 0 )
self . set_param_hint ( ' A_sigmax ' , expr = f ' { self . prefix } delta + { self . prefix } B_sigmax ' )
2023-05-04 18:32:17 +02:00
def guess ( self , data , x , y , negative = False , * * kwargs ) :
pars_guess = guess_from_peak2d ( self . helperModel , data , x , y , negative )
2023-05-07 00:38:52 +02:00
pars = self . make_params ( A_amplitude = pars_guess [ ' amplitude ' ] . value , A_centerx = pars_guess [ ' centerx ' ] . value , A_centery = pars_guess [ ' centery ' ] . value ,
A_sigmax = pars_guess [ ' sigmax ' ] . value , A_sigmay = pars_guess [ ' sigmay ' ] . value ,
B_amplitude = pars_guess [ ' amplitude ' ] . value , B_centerx = pars_guess [ ' centerx ' ] . value , B_centery = pars_guess [ ' centery ' ] . value ,
B_sigmax = pars_guess [ ' sigmax ' ] . value , B_sigmay = pars_guess [ ' sigmay ' ] . value )
2023-05-04 18:32:17 +02:00
pars [ f ' { self . prefix } A_sigmax ' ] . set ( expr = f ' delta + { self . prefix } B_sigmax ' )
2023-05-07 00:38:52 +02:00
pars . add ( f ' { self . prefix } delta ' , value = - 1 , max = 0 , min = - np . inf , vary = True )
2023-05-04 18:32:17 +02:00
pars [ f ' { self . prefix } A_sigmay ' ] . set ( min = 0.0 )
pars [ f ' { self . prefix } B_sigmax ' ] . set ( min = 0.0 )
pars [ f ' { self . prefix } B_sigmay ' ] . set ( min = 0.0 )
return pars
2023-05-23 22:14:03 +02:00
class Polylog22dModel ( 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 ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( polylog2_2d , * * kwargs )
self . _set_paramhints_prefix ( )
def _set_paramhints_prefix ( self ) :
self . set_param_hint ( ' Rx ' , min = 0 )
self . set_param_hint ( ' Ry ' , min = 0 )
def guess ( self , data , x , y , negative = False , * * kwargs ) :
""" Estimate initial model parameter values from data. """
pars = guess_from_peak2d ( self , data , x , y , negative )
return update_param_vals ( pars , self . prefix , * * kwargs )
class ThomasFermi2dModel ( Model ) :
fwhm_factor = 1
height_factor = 0.5
def __init__ ( self , independent_vars = [ ' x ' , ' y ' ] , prefix = ' ' , nan_policy = ' raise ' ,
* * kwargs ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( ThomasFermi_2d , * * kwargs )
self . _set_paramhints_prefix ( )
def _set_paramhints_prefix ( self ) :
self . set_param_hint ( ' Rx ' , min = 0 )
self . set_param_hint ( ' Ry ' , min = 0 )
def guess ( self , data , x , y , negative = False , * * kwargs ) :
""" Estimate initial model parameter values from data. """
pars = guess_from_peak2d ( self , data , x , y , negative )
# amplitude = pars['amplitude'].value
# simgax = pars['sigmax'].value
# sigmay = pars['sigmay'].value
# pars['amplitude'].set(value=amplitude/s2pi/simgax/sigmay)
simgax = pars [ ' sigmax ' ] . value
sigmay = pars [ ' sigmay ' ] . value
pars [ ' simgax ' ] . set ( value = simgax / 2.355 )
pars [ ' sigmay ' ] . set ( value = sigmay / 2.355 )
return update_param_vals ( pars , self . prefix , * * kwargs )
2023-08-21 10:28:12 +02:00
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 ,
2023-08-22 09:54:04 +02:00
pre_check = False ,
post_check = False ,
2023-08-21 10:28:12 +02:00
is_debug = False ,
* * kwargs
) :
2023-05-23 22:14:03 +02:00
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
2023-08-21 10:28:12 +02:00
self . atom_n_conv = atom_n_conv
2023-08-22 09:54:04 +02:00
self . pre_check = pre_check
self . post_check = post_check
2023-08-21 10:28:12 +02:00
self . is_debug = is_debug
2023-05-23 22:14:03 +02:00
super ( ) . __init__ ( density_profile_BEC_2d , * * kwargs )
self . _set_paramhints_prefix ( )
def _set_paramhints_prefix ( self ) :
2023-06-29 11:54:10 +02:00
# self.set_param_hint('BEC_sigmax', min=0)
2023-08-21 10:28:12 +02:00
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 )
2023-09-08 14:13:27 +02:00
2023-09-14 14:13:49 +02:00
self . set_param_hint ( ' rot_angle ' , min = - 90 , max = 90 )
2023-08-25 15:36:43 +02:00
self . set_param_hint ( ' atom_number_bec ' , expr = f ' { self . prefix } amp_bec / 5 * 2 * 3.14159265359 * { self . prefix } sigmax_bec * { self . prefix } sigmay_bec ' )
self . set_param_hint ( ' atom_number_th ' , expr = f ' { self . prefix } amp_th * 2 * 3.14159265359 * 1.20206 / 1.643 * { self . prefix } sigma_th * { self . prefix } sigma_th ' )
self . set_param_hint ( ' condensate_fraction ' , expr = f ' { self . prefix } atom_number_bec / ( { self . prefix } atom_number_bec + { self . prefix } atom_number_th) ' )
2023-08-21 10:28:12 +02:00
2023-09-08 14:13:27 +02:00
def guess ( self , data , x , y , rot_angle = 0 , vary_rot = False , pre_check = False , post_check = False , * * kwargs ) :
2023-08-22 09:54:04 +02:00
""" Estimate and create initial model parameters for 2d bimodal fit, by doing a 1d bimodal fit along an integrated slice of the image
: param data : Flattened 2 d array , in form [ a_00 , a_10 , a_20 , . . . , a_01 , a_02 , . . , a_XY ] with a_xy , x_dim = X , y_dim = Y
: type data : 1 d numpy array
: param x : flattened X output of np . meshgrid ( x_axis , y_axis ) in form : [ x1 , x2 , . . , xX , x1 , x2 , . . , xX , . . Y times . . ]
: type x : 1 d numpy array
: param y : flattened Y output of np . meshgrid ( x_axis , y_axis ) in form : [ y1 , y1 , . . , y1 ( X times ) , y2 , y2 , . . , y2 ( X times ) , . . Y times . . ]
: type y : 1 d numpy array
: param pre_check : if True the amplitude of the 1 d fit is used to guess if the image is purely BEC or thermal and
the corresponding amplitude of the 2 d fit is set to zero and not varied to speed up the fitting , defaults to False
: type pre_check : bool , optional
: param post_check : if True , after doing a 2 d bimodal fit the number of atoms surrounding the fitted BEC is counted and if the value is
below a certain threshhold the fit is done again with the thermal amplitude set to zero , defaults to False
: type post_check : bool , optional
: return : initial parameters for 2 d fit
: rtype : params object ( lmfit )
2023-08-21 10:28:12 +02:00
"""
2023-08-22 09:54:04 +02:00
self . pre_check = pre_check
self . post_check = post_check
# reshaping the image to 2D in the form [[a_00, a_01, .., a_0Y], [a_10,.., a_1Y], .., [a_X0, .., a_XY]], with a_xy
2023-08-21 10:28:12 +02:00
x_width = len ( np . unique ( x ) )
y_width = len ( np . unique ( y ) )
2023-08-23 18:23:33 +02:00
x_1d = np . linspace ( x [ 0 ] , x [ - 1 ] , x_width )
y_1d = np . linspace ( y [ 0 ] , y [ - 1 ] , y_width )
2023-08-22 09:54:04 +02:00
2023-08-21 10:28:12 +02:00
data = np . reshape ( data , ( y_width , x_width ) )
data = data . T
2023-09-08 14:13:27 +02:00
if rot_angle != 0 :
data = rotate ( data , rot_angle , reshape = False )
2023-08-21 10:28:12 +02:00
shape = np . shape ( data )
2023-08-22 09:54:04 +02:00
if self . is_debug :
print ( f ' shape: { shape } ' )
max_width = np . max ( shape )
2023-08-21 10:28:12 +02:00
2023-08-22 09:54:04 +02:00
# binarizing image to guess BEC width and calculate center
thresh = self . calc_thresh ( data , thresh_val = 0.5 )
# calculating center of cloud by statistical distribution of binarized image
2023-08-23 18:23:33 +02:00
center_pix = self . calc_cen_pix ( thresh )
center = self . center_pix_conv ( center_pix , x_1d , y_1d )
2023-08-22 09:54:04 +02:00
# guessing BEC width, or better of width of center blob if no BEC is present
2023-08-23 18:23:33 +02:00
BEC_width_guess = self . guess_BEC_width ( thresh , center_pix )
2023-08-21 10:28:12 +02:00
2023-08-22 09:54:04 +02:00
# plot binarized image and center position for debugging
if self . is_debug :
2023-08-23 18:23:33 +02:00
X , Y = np . meshgrid ( x_1d , y_1d )
plt . pcolormesh ( X , Y , thresh . T , cmap = ' jet ' )
2023-08-22 09:54:04 +02:00
plt . plot ( center [ 0 ] , center [ 1 ] , marker = ' x ' , markersize = 25 , color = ' green ' )
plt . gca ( ) . set_aspect ( ' equal ' )
plt . title ( f ' Binarized image for guessing BEC width + center position (BEC_width: x= { BEC_width_guess [ 0 ] : .0f } , y= { BEC_width_guess [ 1 ] : .0f } pix) ' )
plt . xlabel ( ' x_axis ' )
plt . ylabel ( ' y_axis ' )
plt . show ( )
# The 1d fit is done along the short axis of the BEC (decided via the BEC_width guess)
2023-08-21 10:28:12 +02:00
if BEC_width_guess [ 0 ] < BEC_width_guess [ 1 ] :
if self . is_debug :
2023-08-22 09:54:04 +02:00
print ( f ' x smaller y, 1d fit along x ' )
2023-08-21 10:28:12 +02:00
s_width_ind = 0
2023-08-23 18:23:33 +02:00
x_fit = x_1d
2023-08-22 09:54:04 +02:00
# slice of the image along the short BEC axis with width of BEC width is taken
2023-08-23 18:23:33 +02:00
X_guess = np . sum ( data [ : , round ( center_pix [ 1 ] - BEC_width_guess [ 1 ] / 2 ) : round ( center_pix [ 1 ] + BEC_width_guess [ 1 ] / 2 ) ] , 1 ) / len ( data [ 0 , round ( center_pix [ 1 ] - BEC_width_guess [ 1 ] / 2 ) : round ( center_pix [ 1 ] + BEC_width_guess [ 1 ] / 2 ) ] )
2023-08-21 10:28:12 +02:00
else :
if self . is_debug :
2023-08-22 09:54:04 +02:00
print ( f ' y smaller x, 1d fit along y ' )
2023-08-21 10:28:12 +02:00
s_width_ind = 1
2023-08-23 18:23:33 +02:00
x_fit = y_1d
X_guess = np . sum ( data [ round ( center_pix [ 0 ] - BEC_width_guess [ 0 ] / 2 ) : round ( center_pix [ 0 ] + BEC_width_guess [ 0 ] / 2 ) , : ] , 0 ) / len ( data [ 0 , round ( center_pix [ 0 ] - BEC_width_guess [ 0 ] / 2 ) : round ( center_pix [ 0 ] + BEC_width_guess [ 0 ] / 2 ) ] )
2023-08-21 10:28:12 +02:00
2023-08-22 09:54:04 +02:00
# Creating 1d fit init params + Performing fit
2023-08-23 18:23:33 +02:00
2023-08-21 10:28:12 +02:00
max_val = np . max ( X_guess )
fitmodel_1d = lmfit . Model ( density_1d , independent_vars = [ ' x ' ] )
params_1d = lmfit . Parameters ( )
params_1d . add_many (
2023-08-22 09:54:04 +02:00
( ' x0_bec ' , center [ s_width_ind ] , True , center [ s_width_ind ] - 10 , center [ s_width_ind ] + 10 ) ,
( ' x0_th ' , center [ s_width_ind ] , True , center [ s_width_ind ] - 10 , center [ s_width_ind ] + 10 ) ,
( ' amp_bec ' , 0.5 * max_val , True , 0 , 1.3 * max_val ) ,
( ' amp_th ' , 0.5 * max_val , True , 0 , 1.3 * max_val ) ,
( ' deltax ' , 3 * BEC_width_guess [ s_width_ind ] , True , 0 , max_width ) ,
2023-08-21 10:28:12 +02:00
# ('sigma_bec',BEC_width_guess[i,j,0]/1.22, True, 0, 50)
2023-08-22 09:54:04 +02:00
( ' sigma_bec ' , BEC_width_guess [ s_width_ind ] / 1.22 , True , 0 , BEC_width_guess [ s_width_ind ] * 2 )
2023-08-21 10:28:12 +02:00
)
params_1d . add ( ' sigma_th ' , 3 * BEC_width_guess [ 0 ] , min = 0 , expr = f ' 0.632*sigma_bec + 0.518*deltax ' )
2023-05-23 22:14:03 +02:00
2023-08-21 10:28:12 +02:00
2023-08-23 18:23:33 +02:00
res_1d = fitmodel_1d . fit ( X_guess , x = x_fit , params = params_1d )
2023-08-22 09:54:04 +02:00
bval_1d = res_1d . best_values
2023-08-21 10:28:12 +02:00
if self . is_debug :
2023-08-22 09:54:04 +02:00
print ( ' ' )
print ( ' 1d fit initialization ' )
print ( f ' center = { center } ' )
print ( f ' BEC widths: { BEC_width_guess } ' )
print ( ' ' )
print ( ' 1d init fit values ' )
2023-08-21 10:28:12 +02:00
params_1d . pretty_print ( )
2023-08-22 09:54:04 +02:00
print ( ' 1d fitted values ' )
2023-08-21 10:28:12 +02:00
self . print_bval ( res_1d )
2023-08-23 18:23:33 +02:00
plt . plot ( x_fit , X_guess , label = ' 1d int. data ' )
plt . plot ( x_fit , density_1d ( x_fit , * * bval_1d ) , label = ' bimodal fit ' )
plt . plot ( x_fit , thermal ( x_fit , x0 = bval_1d [ ' x0_th ' ] , amp = bval_1d [ ' amp_th ' ] , sigma = bval_1d [ ' sigma_th ' ] ) , label = ' thermal part ' )
2023-08-22 09:54:04 +02:00
plt . legend ( )
if s_width_ind == 0 :
plt . title ( ' 1d fit of data along x-axis ' )
2023-08-23 18:23:33 +02:00
plt . xlabel ( ' x_ ' )
2023-08-22 09:54:04 +02:00
else :
plt . title ( ' 1d fit of data along y-axis ' )
2023-08-23 18:23:33 +02:00
plt . xlabel ( ' y_axis ' )
2023-08-22 09:54:04 +02:00
plt . show ( )
2023-08-21 10:28:12 +02:00
2023-08-22 09:54:04 +02:00
# scaling amplitudes of 1d fit with the maximum value of blurred 2d data
2023-08-21 10:28:12 +02:00
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 ( )
2023-08-22 09:54:04 +02:00
# if precheck enabled and amp_th is 7x higher than amp_bec (value might be changed), amplitude of BEC in 2d fit is set to zero
if bval_1d [ ' amp_th ' ] / bval_1d [ ' amp_bec ' ] > 7 and self . pre_check :
2023-08-21 10:28:12 +02:00
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 )
2023-08-22 09:54:04 +02:00
params [ f ' { self . prefix } sigma_th ' ] . set ( value = bval_1d [ ' sigma_th ' ] , max = max_width , vary = True )
2023-08-21 10:28:12 +02:00
2023-08-22 09:54:04 +02:00
# if precheck enabled and amp_bec is 10x higher than amp_th (value might be changed), amplitude of thermal part in 2d fit is set to zero
elif bval_1d [ ' amp_bec ' ] / bval_1d [ ' amp_th ' ] > 10 and self . pre_check :
2023-08-21 10:28:12 +02:00
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 )
2023-05-24 16:54:29 +02:00
else :
2023-08-21 10:28:12 +02:00
print ( ' Error in small width BEC recogintion, s_width_ind should be 0 or 1 ' )
2023-08-22 09:54:04 +02:00
# params for normal 2d bimodal fit are initialized
2023-08-21 10:28:12 +02:00
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 )
2023-08-22 09:54:04 +02:00
params [ f ' { self . prefix } sigma_th ' ] . set ( value = bval_1d [ ' sigma_th ' ] , max = max_width , vary = True )
2023-08-21 10:28:12 +02:00
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 ' )
2023-09-08 14:13:27 +02:00
params [ f ' { self . prefix } rot_angle ' ] . set ( value = - rot_angle , min = - rot_angle - 30 , max = - rot_angle + 30 , vary = vary_rot )
2023-08-21 10:28:12 +02:00
if self . is_debug :
print ( ' ' )
print ( ' Init Params ' )
params . pretty_print ( )
2023-08-22 09:54:04 +02:00
print ( ' ' )
2023-08-21 10:28:12 +02:00
return lmfit . models . update_param_vals ( params , self . prefix , * * kwargs )
2023-08-21 11:26:02 +02:00
def fit ( self , data , * * kwargs ) :
2023-08-22 09:54:04 +02:00
""" fitting function overwrites parent class fitting function of lmfit, in order to check (if post_check is enabled)
if thermal fit completely lies in BEC fit by counting sourrounding number of atoms and comparing it to threshold value
: param data : Flattened 2 d array , in form [ a_00 , a_10 , a_20 , . . . , a_01 , a_02 , . . , a_XY ] with a_xy , x_dim = X , y_dim = Y
: type data : 1 d numpy array
: return : result of 2 d fit
: rtype : result object ( lmfit )
"""
2023-08-21 10:28:12 +02:00
2023-08-22 09:54:04 +02:00
res = super ( ) . fit ( data , * * kwargs )
2023-08-21 10:28:12 +02:00
if self . is_debug :
print ( ' bval first fit ' )
self . print_bval ( res )
2023-08-22 09:54:04 +02:00
bval = res . best_values
# Do described post_check if enabled
2023-08-23 18:23:33 +02:00
2023-08-22 09:54:04 +02:00
if res . params [ ' amp_bec ' ] . vary and res . params [ ' amp_th ' ] . vary and bval [ ' amp_bec ' ] > 0.5 * bval [ ' amp_th ' ] and self . post_check :
# creating image by cutting out region around BEC and counting number of atoms
2023-08-21 10:28:12 +02:00
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 )
2023-08-22 09:54:04 +02:00
mask = np . where ( tf_fit > 0 , np . nan , data )
2023-08-21 10:28:12 +02:00
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
2023-08-22 09:54:04 +02:00
#TODO change fixed threshhold to variable
# If number of atoms around BEC is small the image is guessed to be purely BEC and another 2d fit is performed with setting the thermal amplitude to zero
2023-08-21 10:28:12 +02:00
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 )
2023-08-22 09:54:04 +02:00
res = super ( ) . fit ( data , x = kwargs [ ' x ' ] , y = kwargs [ ' y ' ] , params = params )
2023-08-21 10:28:12 +02:00
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
2023-08-22 09:54:04 +02:00
: param data : 2 d image
: type data : 2 d numpy array
2023-08-21 10:28:12 +02:00
: param thresh_val : relative threshhold value for binarization with respect to maximum of blurred image
2023-08-22 09:54:04 +02:00
: param sigma : sigma of gaussian blur filter ( see scipy . ndimage . gaussian_filter )
: return : binary 2 d image
: rtype : 2 d numpy array
2023-08-21 10:28:12 +02:00
"""
shape = np . shape ( data )
thresh = np . zeros ( shape )
blurred = gaussian_filter ( data , sigma = sigma )
2023-08-22 09:54:04 +02:00
thresh = np . where ( blurred < np . max ( blurred ) * thresh_val , 0 , 1 )
2023-08-21 10:28:12 +02:00
return thresh
2023-08-23 18:23:33 +02:00
def calc_cen_pix ( self , thresh1 ) :
""" Calculating the center (in pixel) of a blob (atom cloud) in a binarized image by first calculating the probability distribution along both axes and afterwards the expectation value
2023-08-22 09:54:04 +02:00
: param thresh1 : Binary 2 D image in the form [ [ a_00 , a_01 , . . , a_0Y ] , [ a_10 , . . , a_1Y ] , . . , [ a_X0 , . . , a_XY ] ] , with a_xy , x_dim = X , y_dim = Y
: type thresh1 : 2 D numpy array
: return : center coordinates of blob in form [ x_center , y_center ]
: rtype : 1 d numpy array ( shape = ( 1 , 2 ) )
2023-08-21 10:28:12 +02:00
"""
cen = np . zeros ( 2 )
2023-08-22 09:54:04 +02:00
( X , Y ) = np . shape ( thresh1 )
2023-08-21 10:28:12 +02:00
thresh1 = thresh1 / np . sum ( thresh1 )
# marginal distributions
2023-08-22 09:54:04 +02:00
dx = np . sum ( thresh1 , 1 )
dy = np . sum ( thresh1 , 0 )
2023-08-21 10:28:12 +02:00
# expected values
cen [ 0 ] = np . sum ( dx * np . arange ( X ) )
cen [ 1 ] = np . sum ( dy * np . arange ( Y ) )
return cen
2023-08-23 18:23:33 +02:00
def center_pix_conv ( self , center_pix , x , y ) :
""" Converts center in pixel to center in values of x and y
: param center_pix : pixel values of center
: type center_pix : numpy array ( shape = ( 1 , 2 ) )
: param x : x - axis
: type x : 1 d numpy array
: param y : y - axis
: type y : 1 d numpy array
: return : center coordinates in form [ x_center , y_center ] with respect to the axes
: rtype : numpy array ( shap = ( 1 , 2 ) )
"""
center = np . empty ( 2 )
center [ 0 ] = x [ round ( center_pix [ 0 ] ) ]
center [ 1 ] = y [ round ( center_pix [ 1 ] ) ]
return center
2023-08-21 10:28:12 +02:00
def guess_BEC_width ( self , thresh , center ) :
2023-08-22 09:54:04 +02:00
""" returns width of blob in binarized image along both axis through the center
: param thresh : Binary 2 D image in the form [ [ a_00 , a_01 , . . , a_0Y ] , [ a_10 , . . , a_1Y ] , . . , [ a_X0 , . . , a_XY ] ] , with a_xy , x_dim = X , y_dim = Y
: type thresh : 2 d numpy array
2023-08-23 18:23:33 +02:00
: param center : center of blob in image in form [ x_center , y_center ] in pixel
2023-08-22 09:54:04 +02:00
: type center : 1 d numpy array ( shape = ( 1 , 2 ) )
: return : width of blob in image as [ x_width , y_width ]
: rtype : 1 d numpy array ( shape = ( 1 , 2 ) )
2023-08-21 10:28:12 +02:00
"""
shape = np . shape ( thresh )
if len ( shape ) == 2 :
2023-08-22 09:54:04 +02:00
BEC_width_guess = np . array ( [ np . sum ( thresh [ : , round ( center [ 1 ] ) ] ) , np . sum ( thresh [ round ( center [ 0 ] ) , : ] ) ] )
for i in range ( 2 ) :
if BEC_width_guess [ i ] < = 0 :
BEC_width_guess [ i ] = 1
2023-08-21 10:28:12 +02:00
else :
print ( " Shape of data is wrong, output is empty " )
return BEC_width_guess
2023-08-22 09:54:04 +02:00
def cond_frac ( self , results , X , Y ) :
""" Returns condensate fraction of 2d fitting result
: param results : result of 2 d bimodal fit
: type results : result object ( lmfit )
: param X : X output of np . meshgrid ( x_axis , y_axis ) in form : [ [ x1 , x2 , . . , xX ] , [ x1 , x2 , . . , xX ] . . Y times . . ]
: type X : 2 d numpy array
: param Y : Y output of np . meshgrid ( x_axis , y_axis ) in form : [ [ y1 , y1 , . . , y1 ( X times ) ] , [ y2 , y2 , . . , y2 ( X times ) ] , . . Y times . . ]
: type Y : 2 d numpy array
: return : condensate fraction
: rtype : float between 0 and 1
"""
2023-08-21 10:28:12 +02:00
bval = results . best_values
2023-08-22 09:54:04 +02:00
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 ' ] )
2023-08-21 10:28:12 +02:00
N_bec = np . sum ( tf_fit )
2023-08-22 09:54:04 +02:00
fit = density_profile_BEC_2d ( X , Y , * * bval )
2023-08-21 10:28:12 +02:00
N_ges = np . sum ( fit )
return N_bec / N_ges
def return_atom_number ( self , result , X , Y , is_print = True ) :
2023-08-22 09:54:04 +02:00
""" Calculating (and printing if enabled) fitted atom number in BEC + thermal state, and condensate fraction
: param result : result of 2 d bimodal fit
: type result : result object ( lmfit )
: param X : X output of np . meshgrid ( x_axis , y_axis ) in form : [ [ x1 , x2 , . . , xX ] , [ x1 , x2 , . . , xX ] . . Y times . . ]
: type X : 2 d numpy array
: param Y : Y output of np . meshgrid ( x_axis , y_axis ) in form : [ [ y1 , y1 , . . , y1 ( X times ) ] , [ y2 , y2 , . . , y2 ( X times ) ] , . . Y times . . ]
: type Y : 2 d numpy array
: param is_print : if true results are printed , defaults to True
: type is_print : bool , optional
: return : dictionary with total atom number N , BEC N_bec , thermal N_th and condensate fraction cond_f
: rtype : dictionary
"""
2023-08-21 10:28:12 +02:00
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 )
2023-08-22 09:54:04 +02:00
N = N_bec + N_th
frac = N_bec / N
2023-08-21 10:28:12 +02:00
# 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 } ' )
2023-08-22 09:54:04 +02:00
print ( f ' N_ges: { N : .0f } ' )
print ( f ' Cond. frac: { frac * 1e2 : .2f } % ' )
2023-08-21 10:28:12 +02:00
print ( ' ' )
2023-08-22 09:54:04 +02:00
atom_n = { ' N ' : N , ' N_bec ' : N_bec , ' N_th ' : N_th , ' cond_f ' : frac }
return atom_n
def return_temperature ( self , result , tof , omg = None , is_print = True , eff_pix = 2.472e-6 ) :
""" Returns temperature of thermal cloud
: param result : result of 2 d bimodal fit
: type result : result object ( lmfit )
: param tof : time of flight
: type tof : float
: param omg : geometric average of trapping frequencies , defaults to None
: type omg : float , if NONE initial cloud size is neglected optional
: param is_print : if True temperature is printed , defaults to True
: type is_print : bool , optional
: param eff_pix : effective pixel size of imaging setup , defaults to 2.472e-6
: type eff_pix : float , optional
: return : temperature of atom cloud
: rtype : float
"""
2023-08-21 10:28:12 +02:00
R_th = result . best_values [ ' sigma_th ' ] * eff_pix * np . sqrt ( 2 )
2023-08-22 09:54:04 +02:00
# print(R_th)
if omg is None :
T = R_th * * 2 * 164 * const . u / const . k * ( tof * * 2 ) * * ( - 1 )
else :
T = R_th * * 2 * 164 * const . u / const . k * ( 1 / omg * * 2 + tof * * 2 ) * * ( - 1 )
2023-08-21 10:28:12 +02:00
if is_print :
print ( f ' Temperature: { T * 1e9 : .2f } nK ' )
return T
def print_bval ( self , res_s ) :
2023-08-22 09:54:04 +02:00
""" nicely print best fitted values + init values + bounds
: param res_s : result of 2 d bimodal fit
: type res_s : result object ( lmfit )
"""
2023-08-21 10:28:12 +02:00
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 ( ' ' )
2023-05-23 22:14:03 +02:00
2023-05-17 11:03:34 +02:00
class NewFitModel ( Model ) :
def __init__ ( self , func , independent_vars = [ ' x ' ] , prefix = ' ' , nan_policy = ' raise ' ,
* * kwargs ) :
kwargs . update ( { ' prefix ' : prefix , ' nan_policy ' : nan_policy ,
' independent_vars ' : independent_vars } )
super ( ) . __init__ ( func , * * kwargs )
def guess ( self , * args , * * kwargs ) :
return self . make_params ( )
2023-05-04 13:47:33 +02:00
lmfit_models = { ' Constant ' : ConstantModel ,
' Complex Constant ' : ComplexConstantModel ,
' Linear ' : LinearModel ,
' Quadratic ' : QuadraticModel ,
' Polynomial ' : PolynomialModel ,
' Gaussian ' : GaussianModel ,
' Gaussian-2D ' : Gaussian2dModel ,
' Lorentzian ' : LorentzianModel ,
' Split-Lorentzian ' : SplitLorentzianModel ,
' Voigt ' : VoigtModel ,
' PseudoVoigt ' : PseudoVoigtModel ,
' Moffat ' : MoffatModel ,
' Pearson7 ' : Pearson7Model ,
' StudentsT ' : StudentsTModel ,
' Breit-Wigner ' : BreitWignerModel ,
' Log-Normal ' : LognormalModel ,
' Damped Oscillator ' : DampedOscillatorModel ,
' Damped Harmonic Oscillator ' : DampedHarmonicOscillatorModel ,
' Exponential Gaussian ' : ExponentialGaussianModel ,
' Skewed Gaussian ' : SkewedGaussianModel ,
' Skewed Voigt ' : SkewedVoigtModel ,
' Thermal Distribution ' : ThermalDistributionModel ,
' Doniach ' : DoniachModel ,
' Power Law ' : PowerLawModel ,
' Exponential ' : ExponentialModel ,
' Step ' : StepModel ,
' Rectangle ' : RectangleModel ,
' Expression ' : ExpressionModel ,
' Gaussian With Offset ' : GaussianWithOffsetModel ,
' Lorentzian With Offset ' : LorentzianWithOffsetModel ,
' Expansion ' : ExpansionModel ,
2023-05-04 18:32:17 +02:00
' Damping Oscillation Model ' : DampingOscillationModel ,
' Two Gaussian-2D ' : TwoGaussian2dModel ,
2023-07-12 17:23:18 +02:00
' Thomas Fermi-2D ' : ThomasFermi2dModel ,
' Density Profile of BEC-2D ' : DensityProfileBEC2dModel ,
' Polylog2-2D ' : polylog2_2d ,
2023-05-04 13:47:33 +02:00
}
class FitAnalyser ( ) :
2023-07-03 15:19:12 +02:00
""" This is a class integrated all the functions to do a fit.
"""
2023-05-04 13:47:33 +02:00
def __init__ ( self , fitModel , fitDim = 1 , * * kwargs ) - > None :
2023-07-03 15:19:12 +02:00
""" Initialize function
: param fitModel : The fitting model of fit function
: type fitModel : lmfit Model
: param fitDim : The dimension of the fit , defaults to 1
: type fitDim : int , optional
"""
2023-05-04 13:47:33 +02:00
if isinstance ( fitModel , str ) :
self . fitModel = lmfit_models [ fitModel ] ( * * kwargs )
else :
self . fitModel = fitModel
self . fitDim = fitDim
2023-05-07 00:38:52 +02:00
2023-05-22 19:35:09 +02:00
def print_params_set_template ( self , params = None ) :
2023-07-03 15:19:12 +02:00
""" Print a template to manually set the initial parameters of the fit
: param params : An object of Parameters class to print , defaults to None
: type params : lmfit Paraemters , optional
"""
2023-05-07 00:38:52 +02:00
if params is None :
params = self . fitModel . make_params ( )
for key in params :
res = " params.add( "
res + = " name= \" " + key + " \" , "
if not params [ key ] . expr is None :
2023-06-02 18:42:18 +02:00
res + = " expr= \" " + params [ key ] . expr + " \" ) "
2023-05-07 00:38:52 +02:00
else :
res + = " value= " + f ' { params [ key ] . value : 3g } ' + " , "
if str ( params [ key ] . max ) == " inf " :
res + = " max=np.inf, "
else :
res + = " max= " + f ' { params [ key ] . max : 3g } ' + " , "
if str ( params [ key ] . min ) == " -inf " :
res + = " min=-np.inf, "
else :
res + = " min= " + f ' { params [ key ] . min : 3g } ' + " , "
res + = " vary= " + str ( params [ key ] . vary ) + " ) "
print ( res )
2023-05-04 13:47:33 +02:00
2023-05-04 18:32:17 +02:00
def _guess_1D ( self , data , x , * * kwargs ) :
2023-07-03 15:19:12 +02:00
""" Call the guess function of the 1D fit model to guess the initial value.
: param data : The data to fit
: type data : 1 D numpy array
: param x : The data of x axis
: type x : 1 D numpy array
: return : The guessed initial parameters for the fit
: rtype : lmfit Parameters
"""
2023-05-04 18:32:17 +02:00
return self . fitModel . guess ( data = data , x = x , * * kwargs )
2023-05-04 13:47:33 +02:00
2023-05-04 18:32:17 +02:00
def _guess_2D ( self , data , x , y , * * kwargs ) :
2023-07-03 15:19:12 +02:00
""" Call the guess function of the 2D fit model to guess the initial value.
: param data : The flattened data to fit
: type data : 1 D numpy array
: param x : The flattened data of x axis
: type x : 1 D numpy array
: param y : The flattened data of y axis
: type y : 1 D numpy array
: return : The guessed initial parameters for the fit
: rtype : lmfit Parameters
"""
2023-05-15 17:05:16 +02:00
data = data . flatten ( order = ' F ' )
2023-05-04 18:32:17 +02:00
return self . fitModel . guess ( data = data , x = x , y = y , * * kwargs )
2023-05-04 13:47:33 +02:00
2023-05-10 19:03:03 +02:00
def guess ( self , dataArray , x = None , y = None , guess_kwargs = { } , input_core_dims = None , dask = ' parallelized ' , vectorize = True , keep_attrs = True , daskKwargs = None , * * kwargs ) :
2023-07-03 18:37:44 +02:00
""" Call the guess function of the fit model to guess the initial value.
2023-07-03 15:19:12 +02:00
: param dataArray : The data for the fit
: type dataArray : xarray DataArray
: param x : The name of x axis in data or the data of x axis , defaults to None
: type x : str or numpy array , optional
: param y : The name of y axis in data or the data of y axis , defaults to None
: type y : str or numpy array , optional
: param guess_kwargs : the keyworded arguments to send to the guess function , defaults to { }
: type guess_kwargs : dict , optional
: param input_core_dims : over write of the same argument in xarray . apply_ufunc , defaults to None
2023-07-03 18:37:44 +02:00
: type input_core_dims : list or array like , optional
2023-07-03 15:19:12 +02:00
: param dask : over write of the same argument in xarray . apply_ufunc , , defaults to ' parallelized '
: type dask : str , optional
2023-07-05 12:15:03 +02:00
: param vectorize : over write of the same argument in xarray . apply_ufunc , defaults to True
2023-07-03 15:19:12 +02:00
: type vectorize : bool , optional
2023-07-05 12:15:03 +02:00
: param keep_attrs : over write of the same argument in xarray . apply_ufunc , defaults to True
2023-07-03 15:19:12 +02:00
: type keep_attrs : bool , optional
2023-07-05 12:15:03 +02:00
: param daskKwargs : over write of the same argument in xarray . apply_ufunc , defaults to None
2023-07-03 15:19:12 +02:00
: type daskKwargs : dict , optional
2023-07-03 18:37:44 +02:00
: return : The guessed initial parameters for the fit
: rtype : xarray DataArray
2023-07-03 15:19:12 +02:00
"""
2023-05-04 13:47:33 +02:00
kwargs . update (
{
" dask " : dask ,
" vectorize " : vectorize ,
2023-05-04 18:32:17 +02:00
" input_core_dims " : input_core_dims ,
' keep_attrs ' : keep_attrs ,
2023-05-10 19:03:03 +02:00
2023-05-04 13:47:33 +02:00
}
)
2023-05-10 19:03:03 +02:00
if not daskKwargs is None :
kwargs . update ( { " dask_gufunc_kwargs " : daskKwargs } )
2023-05-04 13:47:33 +02:00
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ ' x ' ] ] ,
}
)
if x is None :
if ' x ' in dataArray . dims :
x = dataArray [ ' x ' ] . to_numpy ( )
else :
if isinstance ( x , str ) :
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ x ] ] ,
}
)
2023-05-05 18:25:03 +02:00
x = dataArray [ x ] . to_numpy ( )
2023-05-04 13:47:33 +02:00
if self . fitDim == 1 :
2023-05-04 18:32:17 +02:00
guess_kwargs . update (
{
' x ' : x ,
}
)
return xr . apply_ufunc ( self . _guess_1D , dataArray , kwargs = guess_kwargs ,
2023-05-04 13:47:33 +02:00
output_dtypes = [ type ( self . fitModel . make_params ( ) ) ] ,
* * kwargs
)
if self . fitDim == 2 :
if y is None :
if ' y ' in dataArray . dims :
y = dataArray [ ' y ' ] . to_numpy ( )
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ ' x ' , ' y ' ] ] ,
}
)
else :
if isinstance ( y , str ) :
kwargs [ " input_core_dims " ] [ 0 ] = np . append ( kwargs [ " input_core_dims " ] [ 0 ] , y )
2023-05-05 18:25:03 +02:00
y = dataArray [ y ] . to_numpy ( )
2023-05-04 13:47:33 +02:00
elif input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ ' x ' , ' y ' ] ] ,
}
)
_x , _y = np . meshgrid ( x , y )
_x = _x . flatten ( )
_y = _y . flatten ( )
2023-05-15 17:05:16 +02:00
# dataArray = dataArray.stack(_z=(kwargs["input_core_dims"][0][0], kwargs["input_core_dims"][0][1]))
2023-05-04 13:47:33 +02:00
2023-05-15 17:05:16 +02:00
# kwargs["input_core_dims"][0] = ['_z']
2023-05-04 13:47:33 +02:00
2023-05-04 18:32:17 +02:00
guess_kwargs . update (
{
' x ' : _x ,
' y ' : _y ,
}
)
return xr . apply_ufunc ( self . _guess_2D , dataArray , kwargs = guess_kwargs ,
2023-05-04 13:47:33 +02:00
output_dtypes = [ type ( self . fitModel . make_params ( ) ) ] ,
* * kwargs
)
def _fit_1D ( self , data , params , x ) :
2023-07-03 18:37:44 +02:00
""" Call the fit function of the 1D fit model to do the fit.
: param data : The data to fit
: type data : 1 D numpy array
: param params : The initial paramters of the fit
: type params : lmfit Parameters
: param x : The data of x axis
: type x : 1 D numpy array
: return : The result of the fit
: rtype : lmfit FitResult
"""
2023-05-16 15:51:13 +02:00
return self . fitModel . fit ( data = data , x = x , params = params , nan_policy = ' omit ' )
2023-05-04 13:47:33 +02:00
def _fit_2D ( self , data , params , x , y ) :
2023-07-03 18:37:44 +02:00
""" Call the fit function of the 2D fit model to do the fit.
: param data : The flattened data to fit
: type data : 1 D numpy array
: param params : The flattened initial paramters of the fit
: type params : lmfit Parameters
: param x : The flattened data of x axis
: type x : 1 D numpy array
: param y : The flattened data of y axis
: type y : 1 D numpy array
: return : The result of the fit
: rtype : lmfit FitResult
"""
2023-05-15 17:05:16 +02:00
data = data . flatten ( order = ' F ' )
2023-05-16 15:51:13 +02:00
return self . fitModel . fit ( data = data , x = x , y = y , params = params , nan_policy = ' omit ' )
2023-05-04 13:47:33 +02:00
2023-05-10 19:03:03 +02:00
def fit ( self , dataArray , paramsArray , x = None , y = None , input_core_dims = None , dask = ' parallelized ' , vectorize = True , keep_attrs = True , daskKwargs = None , * * kwargs ) :
2023-07-03 18:37:44 +02:00
""" Call the fit function of the fit model to do the fit.
: param dataArray : The data for the fit
: type dataArray : xarray DataArray
: param paramsArray : The flattened initial paramters of the fit
: type paramsArray : xarray DataArray or lmfit Parameters
: param x : The name of x axis in data or the data of x axis , defaults to None
: type x : str or numpy array , optional
: param y : The name of y axis in data or the data of y axis , defaults to None
: type y : str or numpy array , optional
: param input_core_dims : over write of the same argument in xarray . apply_ufunc , defaults to None , defaults to None
: type input_core_dims : list or array like , optional
: param dask : over write of the same argument in xarray . apply_ufunc , defaults to None , defaults to ' parallelized '
: type dask : str , optional
: param vectorize : over write of the same argument in xarray . apply_ufunc , defaults to None , defaults to True
: type vectorize : bool , optional
: param keep_attrs : over write of the same argument in xarray . apply_ufunc , defaults to None , defaults to True
: type keep_attrs : bool , optional
: param daskKwargs : over write of the same argument in xarray . apply_ufunc , defaults to None , defaults to None
: type daskKwargs : dict , optional
: return : The fit result
: rtype : xarray DataArray
"""
2023-05-04 13:47:33 +02:00
kwargs . update (
{
" dask " : dask ,
" vectorize " : vectorize ,
" input_core_dims " : input_core_dims ,
2023-05-04 18:32:17 +02:00
' keep_attrs ' : keep_attrs ,
2023-05-04 13:47:33 +02:00
}
)
2023-05-10 19:03:03 +02:00
if not daskKwargs is None :
kwargs . update ( { " dask_gufunc_kwargs " : daskKwargs } )
2023-05-04 13:47:33 +02:00
if isinstance ( paramsArray , type ( self . fitModel . make_params ( ) ) ) :
2023-05-07 00:38:52 +02:00
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ ' x ' ] ] ,
}
)
if x is None :
if ' x ' in dataArray . dims :
x = dataArray [ ' x ' ] . to_numpy ( )
else :
if isinstance ( x , str ) :
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ x ] ] ,
}
)
x = dataArray [ x ] . to_numpy ( )
2023-05-04 13:47:33 +02:00
if self . fitDim == 1 :
2023-06-02 18:42:18 +02:00
res = xr . apply_ufunc ( self . _fit_1D , dataArray , kwargs = { ' params ' : paramsArray , ' x ' : x } ,
2023-05-04 13:47:33 +02:00
output_dtypes = [ type ( lmfit . model . ModelResult ( self . fitModel , self . fitModel . make_params ( ) ) ) ] ,
* * kwargs )
2023-06-02 18:42:18 +02:00
if keep_attrs :
res . attrs = copy . copy ( dataArray . attrs )
return res
2023-05-04 13:47:33 +02:00
if self . fitDim == 2 :
if y is None :
if ' y ' in dataArray . dims :
y = dataArray [ ' y ' ] . to_numpy ( )
if input_core_dims is None :
kwargs . update (
{
2023-05-07 00:38:52 +02:00
" input_core_dims " : [ [ ' x ' , ' y ' ] ] ,
2023-05-04 13:47:33 +02:00
}
)
else :
if isinstance ( y , str ) :
kwargs [ " input_core_dims " ] [ 0 ] = np . append ( kwargs [ " input_core_dims " ] [ 0 ] , y )
2023-05-05 18:25:03 +02:00
y = dataArray [ y ] . to_numpy ( )
2023-05-04 13:47:33 +02:00
elif input_core_dims is None :
kwargs . update (
{
2023-05-07 00:38:52 +02:00
" input_core_dims " : [ [ ' x ' , ' y ' ] ] ,
2023-05-04 13:47:33 +02:00
}
)
_x , _y = np . meshgrid ( x , y )
_x = _x . flatten ( )
_y = _y . flatten ( )
2023-05-15 17:05:16 +02:00
# dataArray = dataArray.stack(_z=(kwargs["input_core_dims"][0][0], kwargs["input_core_dims"][0][1]))
2023-05-04 13:47:33 +02:00
2023-05-15 17:05:16 +02:00
# kwargs["input_core_dims"][0] = ['_z']
2023-05-04 13:47:33 +02:00
2023-06-02 18:42:18 +02:00
res = xr . apply_ufunc ( self . _fit_2D , dataArray , kwargs = { ' params ' : paramsArray , ' x ' : _x , ' y ' : _y } ,
2023-05-04 13:47:33 +02:00
output_dtypes = [ type ( lmfit . model . ModelResult ( self . fitModel , self . fitModel . make_params ( ) ) ) ] ,
* * kwargs )
2023-06-02 18:42:18 +02:00
if keep_attrs :
res . attrs = copy . copy ( dataArray . attrs )
return res
2023-05-04 13:47:33 +02:00
else :
2023-05-07 00:38:52 +02:00
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ ' x ' ] , [ ] ] ,
}
)
if x is None :
if ' x ' in dataArray . dims :
x = dataArray [ ' x ' ] . to_numpy ( )
else :
if isinstance ( x , str ) :
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ x ] , [ ] ] ,
}
)
x = dataArray [ x ] . to_numpy ( )
2023-05-04 13:47:33 +02:00
if self . fitDim == 1 :
return xr . apply_ufunc ( self . _fit_1D , dataArray , paramsArray , kwargs = { ' x ' : x } ,
output_dtypes = [ type ( lmfit . model . ModelResult ( self . fitModel , self . fitModel . make_params ( ) ) ) ] ,
* * kwargs )
if self . fitDim == 2 :
if input_core_dims is None :
kwargs . update (
{
" input_core_dims " : [ [ ' x ' , ' y ' ] , [ ] ] ,
}
)
if y is None :
if ' y ' in dataArray . dims :
y = dataArray [ ' y ' ] . to_numpy ( )
else :
if isinstance ( y , str ) :
y = dataArray [ y ] . to_numpy ( )
kwargs [ " input_core_dims " ] [ 0 ] = np . append ( kwargs [ " input_core_dims " ] [ 0 ] , y )
_x , _y = np . meshgrid ( x , y )
_x = _x . flatten ( )
_y = _y . flatten ( )
2023-05-15 17:05:16 +02:00
# dataArray = dataArray.stack(_z=(kwargs["input_core_dims"][0][0], kwargs["input_core_dims"][0][1]))
2023-05-04 13:47:33 +02:00
2023-05-15 17:05:16 +02:00
# kwargs["input_core_dims"][0] = ['_z']
2023-05-04 13:47:33 +02:00
return xr . apply_ufunc ( self . _fit_2D , dataArray , paramsArray , kwargs = { ' x ' : _x , ' y ' : _y } ,
output_dtypes = [ type ( lmfit . model . ModelResult ( self . fitModel , self . fitModel . make_params ( ) ) ) ] ,
* * kwargs )
def _eval_1D ( self , fitResult , x ) :
2023-07-03 18:37:44 +02:00
""" Call the eval function of the 1D fit model to calculate the curve.
: param fitResult : The result of fit
: type fitResult : lmfit FitResult
: param x : The data of x axis
: type x : 1 D numpy array
: return : The curve based on the fit result
: rtype : 1 D numpy array
"""
2023-05-04 13:47:33 +02:00
return self . fitModel . eval ( x = x , * * fitResult . best_values )
def _eval_2D ( self , fitResult , x , y , shape ) :
2023-07-03 18:37:44 +02:00
""" Call the eval function of the 2D fit model to calculate the curve.
: param fitResult : The result of fit
: type fitResult : lmfit FitResult
: param x : The data of x axis
: type x : 1 D numpy array
: param y : The flattened data of y axis
: type y : 1 D numpy array
: param shape : The desired shape for output
: type shape : list , tuple or array like
: return : The curve based on the fit result
: rtype : 2 D numpy array
"""
2023-05-04 13:47:33 +02:00
res = self . fitModel . eval ( x = x , y = y , * * fitResult . best_values )
2023-05-15 17:05:16 +02:00
return res . reshape ( shape , order = ' F ' )
2023-05-04 13:47:33 +02:00
2023-05-10 19:03:03 +02:00
def eval ( self , fitResultArray , x = None , y = None , output_core_dims = None , prefix = " " , dask = ' parallelized ' , vectorize = True , daskKwargs = None , * * kwargs ) :
2023-07-05 12:15:03 +02:00
""" Call the eval function of the fit model to calculate the curve.
2023-07-03 18:37:44 +02:00
: param fitResultArray : The result of fit
: type fitResultArray : xarray DataArray
: param x : The name of x axis in data or the data of x axis , defaults to None
: type x : str or numpy array , optional
: param y : The name of y axis in data or the data of y axis , defaults to None
: type y : str or numpy array , optional
: param output_core_dims : over write of the same argument in xarray . apply_ufunc , defaults to None
: type output_core_dims : list , optional
: param prefix : prefix str of the fit model , defaults to " "
: type prefix : str , optional
: param dask : over write of the same argument in xarray . apply_ufunc , defaults to ' parallelized '
: type dask : str , optional
: param vectorize : over write of the same argument in xarray . apply_ufunc , defaults to True
: type vectorize : bool , optional
: param daskKwargs : over write of the same argument in xarray . apply_ufunc , defaults to None
: type daskKwargs : dict , optional
: return : The curve based on the fit result
: rtype : xarray
"""
2023-05-04 13:47:33 +02:00
kwargs . update (
{
" dask " : dask ,
" vectorize " : vectorize ,
" output_core_dims " : output_core_dims ,
}
)
2023-05-10 19:03:03 +02:00
if daskKwargs is None :
daskKwargs = { }
2023-05-04 13:47:33 +02:00
if self . fitDim == 1 :
if output_core_dims is None :
kwargs . update (
{
2023-05-05 18:25:03 +02:00
" output_core_dims " : prefix + ' x ' ,
2023-05-04 13:47:33 +02:00
}
)
output_core_dims = [ prefix + ' x ' ]
2023-05-10 19:03:03 +02:00
daskKwargs . update (
{
' output_sizes ' : {
output_core_dims [ 0 ] : np . size ( x ) ,
} ,
' meta ' : np . ndarray ( ( 0 , 0 ) , dtype = float )
}
)
2023-05-04 13:47:33 +02:00
kwargs . update (
{
2023-05-10 19:03:03 +02:00
" dask_gufunc_kwargs " : daskKwargs ,
2023-05-04 13:47:33 +02:00
}
)
2023-05-08 16:57:58 +02:00
res = xr . apply_ufunc ( self . _eval_1D , fitResultArray , kwargs = { " x " : x } , * * kwargs )
return res . assign_coords ( { prefix + ' x ' : np . array ( x ) } )
2023-05-04 13:47:33 +02:00
if self . fitDim == 2 :
if output_core_dims is None :
kwargs . update (
{
" output_core_dims " : [ [ prefix + ' x ' , prefix + ' y ' ] ] ,
}
)
output_core_dims = [ prefix + ' x ' , prefix + ' y ' ]
2023-05-10 19:03:03 +02:00
daskKwargs . update (
{
' output_sizes ' : {
output_core_dims [ 0 ] : np . size ( x ) ,
output_core_dims [ 1 ] : np . size ( y ) ,
} ,
' meta ' : np . ndarray ( ( 0 , 0 ) , dtype = float )
} ,
)
2023-05-04 13:47:33 +02:00
kwargs . update (
{
2023-05-10 19:03:03 +02:00
" dask_gufunc_kwargs " : daskKwargs ,
2023-05-04 13:47:33 +02:00
}
)
_x , _y = np . meshgrid ( x , y )
_x = _x . flatten ( )
_y = _y . flatten ( )
2023-05-08 16:57:58 +02:00
res = xr . apply_ufunc ( self . _eval_2D , fitResultArray , kwargs = { " x " : _x , " y " : _y , " shape " : ( len ( x ) , len ( y ) ) } , * * kwargs )
return res . assign_coords ( { prefix + ' x ' : np . array ( x ) , prefix + ' y ' : np . array ( y ) } )
2023-05-04 13:47:33 +02:00
2023-05-07 00:38:52 +02:00
def _get_fit_value_single ( self , fitResult , key ) :
2023-07-03 18:37:44 +02:00
""" get value of one single parameter from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param key : The name of the parameter
: type key : str
: return : The vaule of the parameter
: rtype : float
"""
2023-05-07 00:38:52 +02:00
return fitResult . params [ key ] . value
def _get_fit_value ( self , fitResult , params ) :
2023-07-03 18:37:44 +02:00
""" get value of parameters from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param params : A list of the name of paramerters to read
: type params : list or array like
: return : The vaule of the parameter
: rtype : 1 D numpy array
"""
2023-05-07 00:38:52 +02:00
func = np . vectorize ( self . _get_fit_value_single )
res = tuple (
func ( fitResult , key )
for key in params
)
return res
def get_fit_value ( self , fitResult , dask = ' parallelized ' , * * kwargs ) :
2023-07-03 18:37:44 +02:00
""" get value of parameters from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param dask : over write of the same argument in xarray . apply_ufunc , defaults to ' parallelized '
: type dask : str , optional
: return : The vaule of the parameter
: rtype : xarray DataArray
"""
2023-05-07 00:38:52 +02:00
firstIndex = {
key : fitResult [ key ] [ 0 ]
for key in fitResult . dims
}
firstFitResult = fitResult . sel ( firstIndex ) . item ( )
params = list ( firstFitResult . params . keys ( ) )
output_core_dims = [ [ ] for _ in range ( len ( params ) ) ]
kwargs . update (
{
" dask " : dask ,
" output_core_dims " : output_core_dims ,
}
)
value = xr . apply_ufunc ( self . _get_fit_value , fitResult , kwargs = dict ( params = params ) , * * kwargs )
value = xr . Dataset (
data_vars = {
params [ i ] : value [ i ]
for i in range ( len ( params ) )
} ,
attrs = fitResult . attrs
)
return value
def _get_fit_std_single ( self , fitResult , key ) :
2023-07-03 18:37:44 +02:00
""" get standard deviation of one single parameter from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param key : The name of the parameter
: type key : str
: return : The vaule of the parameter
: rtype : float
"""
2023-05-07 00:38:52 +02:00
return fitResult . params [ key ] . stderr
def _get_fit_std ( self , fitResult , params ) :
2023-07-03 18:37:44 +02:00
""" get standard deviation of parameters from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param params : A list of the name of paramerters to read
: type params : list or array like
: return : The vaule of the parameter
: rtype : 1 D numpy array
"""
2023-05-07 00:38:52 +02:00
func = np . vectorize ( self . _get_fit_std_single )
res = tuple (
func ( fitResult , key )
for key in params
)
return res
def get_fit_std ( self , fitResult , dask = ' parallelized ' , * * kwargs ) :
2023-07-03 18:37:44 +02:00
""" get standard deviation of parameters from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param dask : over write of the same argument in xarray . apply_ufunc , defaults to ' parallelized '
: type dask : str , optional
: return : The vaule of the parameter
: rtype : xarray DataArray
"""
2023-05-07 00:38:52 +02:00
firstIndex = {
key : fitResult [ key ] [ 0 ]
for key in fitResult . dims
}
firstFitResult = fitResult . sel ( firstIndex ) . item ( )
params = list ( firstFitResult . params . keys ( ) )
output_core_dims = [ [ ] for _ in range ( len ( params ) ) ]
kwargs . update (
{
" dask " : dask ,
" output_core_dims " : output_core_dims ,
}
)
value = xr . apply_ufunc ( self . _get_fit_std , fitResult , kwargs = dict ( params = params ) , * * kwargs )
value = xr . Dataset (
data_vars = {
params [ i ] : value [ i ]
for i in range ( len ( params ) )
} ,
attrs = fitResult . attrs
)
return value
def _get_fit_full_result_single ( self , fitResult , key ) :
2023-07-03 18:37:44 +02:00
""" get the full result of one single parameter from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param key : The name of the parameter
: type key : str
: return : The vaule of the parameter
: rtype : float
"""
2023-05-08 11:47:35 +02:00
if not fitResult . params [ key ] . value is None :
value = fitResult . params [ key ] . value
else :
value = np . nan
if not fitResult . params [ key ] . stderr is None :
std = fitResult . params [ key ] . stderr
else :
std = np . nan
return ufloat ( value , std )
2023-05-07 00:38:52 +02:00
def _get_fit_full_result ( self , fitResult , params ) :
2023-07-03 18:37:44 +02:00
""" get the full result of parameters from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param params : A list of the name of paramerters to read
: type params : list or array like
: return : The vaule of the parameter
: rtype : 1 D numpy array
"""
2023-05-07 00:38:52 +02:00
func = np . vectorize ( self . _get_fit_full_result_single )
res = tuple (
func ( fitResult , key )
for key in params
)
return res
def get_fit_full_result ( self , fitResult , dask = ' parallelized ' , * * kwargs ) :
2023-07-03 18:37:44 +02:00
""" get the full result of parameters from fit result
: param fitResult : The result of the fit
: type fitResult : lmfit FitResult
: param dask : over write of the same argument in xarray . apply_ufunc , defaults to ' parallelized '
: type dask : str , optional
: return : The vaule of the parameter
: rtype : xarray DataArray
"""
2023-05-07 00:38:52 +02:00
firstIndex = {
key : fitResult [ key ] [ 0 ]
for key in fitResult . dims
}
firstFitResult = fitResult . sel ( firstIndex ) . item ( )
params = list ( firstFitResult . params . keys ( ) )
output_core_dims = [ [ ] for _ in range ( len ( params ) ) ]
kwargs . update (
{
" dask " : dask ,
" output_core_dims " : output_core_dims ,
}
)
value = xr . apply_ufunc ( self . _get_fit_full_result , fitResult , kwargs = dict ( params = params ) , * * kwargs )
value = xr . Dataset (
data_vars = {
params [ i ] : value [ i ]
for i in range ( len ( params ) )
} ,
attrs = fitResult . attrs
)
return value