atomcloud.functions package
Submodules
- atomcloud.functions.func_base module
- atomcloud.functions.funcs_1d module
- atomcloud.functions.funcs_2d module
- atomcloud.functions.funcs_sumfit module
- atomcloud.functions.multi_base module
- atomcloud.functions.multi_funcs module
Module contents
Created on Mon Mar 14 19:26:59 2022
@author: hofer
- class atomcloud.functions.ConstrainedMultiFunction(funcs, constraints=None)[source]
Bases:
MultiFunctionBaseCreates a multi function which takes a list of functions and a list of constraints. The constraints are the input arguments which should have the same value when in a sub-functions parameters. The individual function parameters are then mapped to a single list of arguments which are passed to the constrained multi function.
The fit function arguments are then passed to the multi function which calls the functions with the correct parameters.
- Parameters:
funcs (list[callable]) –
constraints (list[str] | None) –
- arg_list(arg_param_map, funcs_args)[source]
Opposite to above, makes a list of lists of the same size as the parameters. Each entry is the index of that variable in the single argument list for the scipy/jaxfit compatible fit function
- Parameters:
arg_param_map (dict[str, list[tuple[int, int]]]) – See return of get_arg_map function for detailed description of this argument.
funcs_args (list[list[str]]) – all arguments of the individual functions excluding the first coordinate argument
- Returns:
List of lists of the same size as the parameters. Each entry is the index of that variable in the 1D
- Return type:
list[list[int]]
- check_contraint_strs(constraints, unique_args)[source]
Checks that the constraints are valid strings and that they are arguments of the functions.
- Parameters:
constraints (list[str]) – list of strings which are the names of the parameters which are constrained to between the individual functions comprising the multi-function.
unique_args (list[str]) – the unique input arguments across all the individual functions comprising the multi-function.
- Return type:
None
- check_parameters(aparams)[source]
Checks that the list of list of func parameters is the correct type and that each sublist is the correct length.
- Parameters:
aparams (list[list[float]]) – list of lists of function parameters for the individual functions comprising the multi-function.
- Return type:
None
- create_args_to_params()[source]
Creates a function which maps the 1D list of fit function arguments to the 2D list of function parameters. We use a creator function wrapper so that the args_to_params function can be used in a JAX (i.e. to self argument)
- fit_function(coords, *args)[source]
The fit function which takes a list of arguments and converts them to the 2D list of function parameters and then call the multi-function. This is compatible with the scipy/jaxfit curvefit function input type which requires a single argument list.
- Parameters:
coords (ndarray | Iterable[ndarray]) – coordinates of the data being fit
*args (float) – list of multi-function arguments which are converted then
comprising (converted to parameters for the individual functions) –
multi-function. (the) –
- Returns:
The value of the multi-function at the given coordinates.
- Return type:
ndarray
- get_arg_map(constraints, func_args)[source]
Creates a dictionary which maps the constrained arguments to the individual function parameters. :param constraints: list of constraints (input parameters) :param func_args: list of lists of individual function input arguments
- Returns:
Dictionary which maps the index of the single argument list for the scipy/jaxfit compatible fit function to the two indices of the individual function parameters. The first index is the index of the function in the list of functions and the second index is the index of the parameter in the list of parameters for that function.
- Parameters:
constraints (list[str]) –
func_args (list[list[str]]) –
- Return type:
dict[str, list[tuple[int, int]]]
- params_to_args(params)[source]
Converts the list of lists of individual function parameters into a 1D list of arguments for the fit function compatible with SciPy/JAX curvefit.
- Parameters:
params (list[list[float]]) – list of lists of function parameters for the individual functions comprising the multi-function.
- Returns:
Single list of arguments for the fit function which is compatible with the SciPy/JAX curvefit function.
- Return type:
list[float]
- class atomcloud.functions.EnhancedBose1D[source]
Bases:
FixedEnhancedBose1DSee FunctionBase for documentation
- create_function(anp)[source]
Creates the function using the numpy or jax object given. This method must be overridden by the child class. :param anp: The numpy or jax object to use for the function
- Returns:
The created function
- class atomcloud.functions.EnhancedBose2D[source]
Bases:
Gaussian2DEnhanced Bose 2D cloud function object. This includes the fugacity in the fitting parameters. See FunctionBase for full documentation.
- analyze_parameters(params)[source]
Analysis includes the integrated density and the sigma labwidths.
- class atomcloud.functions.FixedEnhancedBose1D[source]
Bases:
Gaussian1DSee FunctionBase for documentation
- class atomcloud.functions.FixedEnhancedBose2D[source]
Bases:
Gaussian2DFixed enhanced Bose 2D cloud function object. We assume that for this fit function we are near quantum degeneracy and the fugacity is 1. See CloudFunctionBase for full documentation
- class atomcloud.functions.FixedOffset[source]
Bases:
SumFitBaseFuncSee SumFitBaseFunc for documentation
- convert_2d_sum(XY_tuple, params)[source]
Converts 2D function parameters into two sets of 1D parameters along the x and y axes respectively. :param coords: The 2D coordinates of the data being fit :param params: The 2D parameters of the function being fit
- Returns:
The 1D function parameters along the x and y axes
- class atomcloud.functions.FixedOffset1D[source]
Bases:
FunctionBaseSee FunctionBase for documentation
- create_function(anp)[source]
Creates the function using the numpy or jax object given. This method must be overridden by the child class. :param anp: The numpy or jax object to use for the function
- Returns:
The created function
- initial_seed(x, data)[source]
Returns the initial seed parameters for the fit. The default is to return a list of ones for the parameters. This method can be overridden to return a different initial seed if something more intelligent is desired. :param coords: The coordinates to fit the function to :param data: The data to fit the function to
- Returns:
The initial seed parameters for the fit
- rescale_parameters(params, scale)[source]
Rescales the parameters of the function determined by the fit by the scales given for the x y and z axes. :param params: The parameters of the function determined by the fit :param scales: The scales for the x y and z axes
- Returns:
The rescaled fit parameters
- class atomcloud.functions.FixedOffset2D[source]
Bases:
FunctionBaseFixed offset 2D cloud function object. This is a simple offset function which can be combined with those functions objects above.
- class atomcloud.functions.Function1DBase[source]
Bases:
FunctionBaseSee FunctionBase for documentation
- analyze_parameters(params)[source]
Analyzes the fit parameters of the function and returns a dictionary of the analysis parameters. :param params: The function parameters determined by the fit
- Returns:
A dictionary of the analysis parameters
- Parameters:
params (list[float]) –
- Return type:
dict
- default_bounds()[source]
Returns the default bounds for the fit. The default is to return (-np.inf, np.inf) for all parameters. This method can be overridden to return a different set of bounds if something more intelligent is desired. :returns: The default bounds for the fit
- initial_seed(x, data)[source]
Returns the initial seed parameters for the fit. The default is to return a list of ones for the parameters. This method can be overridden to return a different initial seed if something more intelligent is desired. :param coords: The coordinates to fit the function to :param data: The data to fit the function to
- Returns:
The initial seed parameters for the fit
- rescale_analysis_params(params, scales)[source]
Rescales the analysis parameters constructed from the fit parameters by the scales given for the x y and z axes. :param params: The analysis parameters constructed from the fit parameters :param scales: The scales for the x y and z axes
- Returns:
The rescaled analysis parameters
- Parameters:
params (dict) –
scales (list) –
- Return type:
dict
- rescale_parameters(params, scale)[source]
Rescales the parameters of the function determined by the fit by the scales given for the x y and z axes. :param params: The parameters of the function determined by the fit :param scales: The scales for the x y and z axes
- Returns:
The rescaled fit parameters
- class atomcloud.functions.Function2DBase[source]
Bases:
FunctionBaseInherits from FunctionBase and then adds the 2D specific function to create the coordinate transformation functions compatible with JAX.
See FunctionBase for full documentation
- create_coord_funcs(anp)[source]
Creates the coordinate transformation functions which are compatible with JAX due to the use of the creation wrapper function. Again, we can’t use self in the function because it is not jittable
- Parameters:
anp – The numpy or jax numpy module to use for the functions which are created.
- Returns:
The overall coordinate transformation function for 2D coordinates
- class atomcloud.functions.FunctionBase[source]
Bases:
ABCBase class for function objects
- analyze_parameters(params)[source]
Analyzes the fit parameters of the function and returns a dictionary of the analysis parameters. :param params: The function parameters determined by the fit
- Returns:
A dictionary of the analysis parameters
- Parameters:
params (list[float]) –
- Return type:
dict
- abstract create_function(anp)[source]
Creates the function using the numpy or jax object given. This method must be overridden by the child class. :param anp: The numpy or jax object to use for the function
- Returns:
The created function
- Parameters:
anp (object) –
- Return type:
callable
- create_parameter_dict()[source]
Creates a dictionary of the parameters of the function
- Return type:
None
- default_bounds()[source]
Returns the default bounds for the fit. The default is to return (-np.inf, np.inf) for all parameters. This method can be overridden to return a different set of bounds if something more intelligent is desired. :returns: The default bounds for the fit
- Return type:
tuple[list[float], …]
- initial_seed(coords, data)[source]
Returns the initial seed parameters for the fit. The default is to return a list of ones for the parameters. This method can be overridden to return a different initial seed if something more intelligent is desired. :param coords: The coordinates to fit the function to :param data: The data to fit the function to
- Returns:
The initial seed parameters for the fit
- Parameters:
coords (ndarray | Iterable[ndarray]) –
data (ndarray) –
- Return type:
list[float]
- make_function(use_jax=False)[source]
Creates the class function. This function is created using the create_function method and will use jax if use_jax is True otherwise numpy will be used.
- Parameters:
use_jax (bool) – Whether or not to use jax for the function.
- Returns:
The jax or numpy function that is created
- Return type:
callable
- rescale_analysis_params(params, scales)[source]
Rescales the analysis parameters constructed from the fit parameters by the scales given for the x y and z axes. :param params: The analysis parameters constructed from the fit parameters :param scales: The scales for the x y and z axes
- Returns:
The rescaled analysis parameters
- Parameters:
params (dict) –
scales (list) –
- Return type:
dict
- rescale_parameters(params, scales)[source]
Rescales the parameters of the function determined by the fit by the scales given for the x y and z axes. :param params: The parameters of the function determined by the fit :param scales: The scales for the x y and z axes
- Returns:
The rescaled fit parameters
- Parameters:
params (list[float]) –
scales (list) –
- Return type:
list[float]
- class atomcloud.functions.Gaussian[source]
Bases:
SumFitBaseFuncSee SumFitBaseFunc for documentation
- convert_2d_sum(XY_tuple, params)[source]
Converts 2D function parameters into two sets of 1D parameters along the x and y axes respectively. :param coords: The 2D coordinates of the data being fit :param params: The 2D parameters of the function being fit
- Returns:
The 1D function parameters along the x and y axes
- convert_sum_2D(XY_tuple, xparams, yparams)[source]
Takes the 1D fit parameters of the sums along the x and y axes and converts them into 2D parameters. :param coords: The 2D coordinates of the data being fit :param xparams: The 1D sum fit parameters along the x axis :param yparams: The 1D sum fit parameters along the y axis
- Returns:
The 2D function parameters
- class atomcloud.functions.Gaussian1D[source]
Bases:
Function1DBaseSee FunctionBase for documentation
- class atomcloud.functions.Gaussian2D[source]
Bases:
Function2DBaseSee FunctionBase for full documentation.
- analyze_parameters(params)[source]
Calculates the lab widths and the integrated density from the fitted parameters
- create_gaussian2d(anp)[source]
Creates the 2D gaussian function. This will be also used in the inherited ebose classes so it’s a separate function.
- class atomcloud.functions.MultiFunc(function_names, func_registry, constraints=None, use_jax=False)[source]
Bases:
ConstrainedMultiFunctionBase class for 1D and 2D cloud multi-function classes which itself inherits from the ConstrainedMultiFunction class. This function allows the user to combine multiple functions together into a single function which can then be used in SciPy or JAXFit curve fitting functions. Can also include constraints between the parameters of the functions.
- Parameters:
function_names (list[str]) –
func_registry (object) –
constraints (list[str] | None) –
use_jax (bool) –
- class atomcloud.functions.MultiFunction1D(function_names, constraints=None, use_jax=False)[source]
Bases:
MultiFunc1D cloud multi-function class which inherits from the base class. It uses the imported dictionary of 2D function objects as it’s base dictionary of function objects, but also allows the user to add custom function objects to the dictionary of function objects.
See base class for more details.
- Parameters:
function_names (list[str]) –
constraints (list[str] | None) –
use_jax (bool) –
- class atomcloud.functions.MultiFunction2D(function_names, constraints=None, use_jax=False)[source]
Bases:
MultiFunc2D cloud multi-function class which inherits from the base class. It uses the imported dictionary of 2D function objects as it’s base dictionary of function objects, but also allows the user to add custom function objects to the dictionary of function objects.
- Parameters:
function_names (list[str]) – The keys for the function objects in the registry which will be used in the multi-function
func_registry – The registry of function objects
constraints (list[str] | None) – A list of constraints which will be applied to the the functions in the multi-function (see ConstrainedMultiFunction for more details)
use_jax (bool) – If True, the functions in the multi-function will be created using JAX. If False, the functions will be created
- Returns:
None
- class atomcloud.functions.MultiFunctionBase(funcs)[source]
Bases:
ABCClass which allows multiple functions to be added together to form a single function. This is useful for multi function fitting.
- Parameters:
funcs (list[callable]) –
- function(coords, function_parameters)[source]
Multi function which adds together the functions passed to the class initializer.
- Parameters:
coords (ndarray | Iterable[ndarray]) – coordinates of the data
function_parameters (list[list[float]]) – list of parameters for each function
- Returns:
The value of the multi function at the given coordinates
- Return type:
ndarray
- class atomcloud.functions.Parabola[source]
Bases:
SumFitBaseFuncSee SumFitBaseFunc for documentation
- convert_2d_sum(XY_tuple, params)[source]
Converts 2D function parameters into two sets of 1D parameters along the x and y axes respectively. :param coords: The 2D coordinates of the data being fit :param params: The 2D parameters of the function being fit
- Returns:
The 1D function parameters along the x and y axes
- convert_sum_2D(XY_tuple, xparams, yparams)[source]
Takes the 1D fit parameters of the sums along the x and y axes and converts them into 2D parameters. :param coords: The 2D coordinates of the data being fit :param xparams: The 1D sum fit parameters along the x axis :param yparams: The 1D sum fit parameters along the y axis
- Returns:
The 2D function parameters
- class atomcloud.functions.Parabola1D[source]
Bases:
Function1DBaseSee FunctionBase for documentation
- class atomcloud.functions.Parabola2D[source]
Bases:
Function2DBaseParabola 2D cloud function object. This is a simple parabola function and not for the integrated case, although it can be used as an approximation
- analyze_parameters(params)[source]
Analysis includes the integrated density and the radius labwidths.
- create_parabola2d(anp)[source]
Creates the 2D parabola function. This is upside down and does not go lower than 0. Using the np.where or jnp.where which makes it JAX compatible. The parabola is also used for the Thomas-Fermi object which is why it is a separate function.
- default_bounds()[source]
Default bounds for the 2D parabola function are the same as for the Gaussian.
- class atomcloud.functions.SumFitBaseFunc[source]
Bases:
ABC- abstract convert_2d_sum(coords, params)[source]
Converts 2D function parameters into two sets of 1D parameters along the x and y axes respectively. :param coords: The 2D coordinates of the data being fit :param params: The 2D parameters of the function being fit
- Returns:
The 1D function parameters along the x and y axes
- Parameters:
coords (Iterable[ndarray]) –
params (list[float]) –
- Return type:
tuple[list[float]]
- abstract convert_sum_2D(coords, xparams, yparams)[source]
Takes the 1D fit parameters of the sums along the x and y axes and converts them into 2D parameters. :param coords: The 2D coordinates of the data being fit :param xparams: The 1D sum fit parameters along the x axis :param yparams: The 1D sum fit parameters along the y axis
- Returns:
The 2D function parameters
- Parameters:
coords (Iterable[ndarray]) –
xparams (list[float]) –
yparams (list[float]) –
- Return type:
tuple[list[float]]
- class atomcloud.functions.ThomasFermi[source]
Bases:
SumFitBaseFuncSee SumFitBaseFunc for documentation
- convert_2d_sum(XY_tuple, params)[source]
Converts 2D function parameters into two sets of 1D parameters along the x and y axes respectively. :param coords: The 2D coordinates of the data being fit :param params: The 2D parameters of the function being fit
- Returns:
The 1D function parameters along the x and y axes
- convert_sum_2D(XY_tuple, xparams, yparams)[source]
Takes the 1D fit parameters of the sums along the x and y axes and converts them into 2D parameters. :param coords: The 2D coordinates of the data being fit :param xparams: The 1D sum fit parameters along the x axis :param yparams: The 1D sum fit parameters along the y axis
- Returns:
The 2D function parameters
- class atomcloud.functions.ThomasFermi1D[source]
Bases:
Parabola1DSee FunctionBase for documentation
- class atomcloud.functions.ThomasFermi2D[source]
Bases:
Parabola2DThomas-Fermi 2D cloud function object. This is a simple parabola but with a different exponent (3/2) since in absorption imaging we are integrating along the imaging beam axis.
- atomcloud.functions.amps_1D_to_2D(x_amp, y_amp, sigma_x, sigma_y, scalar_function)[source]
Converts 1D amplitudes to 2D sum amplitudes and averages them to get a single converted offset (see amps_1D_to_2D function for more details on the conversion).
- Parameters:
x_amp (float) – The 1D sum amplitude along the x axis
y_amp (float) – The 1D sum amplitude along the y axis
sigma_x (float) – The sigma along the x axis
sigma_y (float) – The sigma along the y axis
scalar_function (True) – A function that scales between sum and 2D amplitudes
- Returns:
The 2D amplitude
- Return type:
float
- atomcloud.functions.amps_2D_to_1D(amp, sigma_x, sigma_y, scalar)[source]
Converts 2D amplitudes to 1D sum amplitudes. For the 1D equations used for fitting the atom cloud integrating along one axis means the 2D peak amplitude get’s multiplied by a scalar which is a function of the functions radius. These scalars have been calculated in mathematica and are found in the ConvertAmp object.
- Parameters:
amp (float) – The 2D amplitude
sigma_x (float) – The sigma along the x axis
sigma_y (float) – The sigma along the y axis
scalar (True) – A function scales between sum and 2D amplitudes
- Returns:
The 1D amplitudes along the x and y axes
- Return type:
tuple[float, float]
- atomcloud.functions.calc_diff_elements(coords)[source]
Calculate the difference between each element in the x and y arrays. This is used to calculate the average pixel size along each axis.
- atomcloud.functions.create_spence(npt)[source]
Create a function that calculates the spence function using either scipy.special.spence or a jaxified version of the same function. I ported the Cephes library algorithm to jax, but which is equivalent for real numbers.
- Parameters:
npt – numpy or jax
- Returns:
spence function
- atomcloud.functions.general_analysis_rescale(analysis_params, scale, key=None)[source]
General rescale for clouds which have only a single lab withds parameter and an integrated density.
- atomcloud.functions.general_lab_widths(px, py, theta, key='sig')[source]
Given two orthogonal parameters which are rotated by theta with respect to the x and y axes, this function returns the widths in the lab frame, but in a dictionary format.
- atomcloud.functions.integrate_ebose(n0, sigma_x, sigma_y, fugacity)[source]
Integrates the ebose function analytically. However, currently the polylog value in the utils is a pretty good approximation. Need to make a general python polylog function that’s JAX compatible in the future.
- atomcloud.functions.rescale_1d_params(params, z_indices, x_indices, xscale, zscale)[source]
Rescale fit parameters in a 1d function
- Parameters:
params (list[float]) – list of function parameters
zinds – list of indices which are scaled by zscale
xinds – list of indices which are scaled by xscale
xscale (float) – scale to rescale x parameters by
zscale (float) – scale to rescale z parameters by
z_indices (list[int]) –
x_indices (list[int]) –
- Returns:
list of rescaled parameters
- Return type:
list[float]
- atomcloud.functions.rescale_2d_params(params, indices, scales, theta_indices=None)[source]
Rescale fit parameters in a 2d function. Must scale along x, y and z additionally some parameters may not lie along the x or y axis and so need to be rescaled differently.
- Parameters:
params (list[float]) – list of function parameters
indices (list[list]) – list of lists of indices which are scaled by the corresponding scale in scales
scales (list) – list of scales to rescale parameters by
theta_indices (list | None) – list of index of the angle parameter corresponding to
tuple (each mixed axis) –
- Returns:
list of rescaled parameters
- Return type:
list[float]