atomcloud.functions package

Submodules

Module contents

Created on Mon Mar 14 19:26:59 2022

@author: hofer

class atomcloud.functions.ConstrainedMultiFunction(funcs, constraints=None)[source]

Bases: MultiFunctionBase

Creates 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: FixedEnhancedBose1D

See 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

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

integrate_function(params)[source]

should probably check this t00

class atomcloud.functions.EnhancedBose2D[source]

Bases: Gaussian2D

Enhanced 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.

create_function(anp)[source]

Creates the 2D enhanced Bose fitting function. The polylog function which is just spence(1-z) for the 2D case is used (see jaxfuncs.spence.py). Also utilizes the gaussian2d function from the Gaussian2D class which is inherited.

default_bounds()[source]

Default bounds for the enhanced Bose function are the same as for the Gaussian except the fugacity is bounded between 0 and 1.

class atomcloud.functions.FixedEnhancedBose1D[source]

Bases: Gaussian1D

See 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

create_polylog1d(anp, n_max=50)[source]

This function needs help

integrate_function(params)[source]

should probably check this t00

class atomcloud.functions.FixedEnhancedBose2D[source]

Bases: Gaussian2D

Fixed 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

analyze_parameters(params)[source]

Analysis includes the integrated density and the sigma labwidths.

create_function(anp)[source]

Creates the 2D fixed enhanced Bose fitting function see the non-fixed version for more details of the polylog function and gaussian2d function.

class atomcloud.functions.FixedOffset[source]

Bases: SumFitBaseFunc

See 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]

Reverse integration of constant offset by dividing by number of pixels on each integration axis. Then average offsets on each axis.

class atomcloud.functions.FixedOffset1D[source]

Bases: FunctionBase

See 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: FunctionBase

Fixed offset 2D cloud function object. This is a simple offset function which can be combined with those functions objects above.

create_function(anp)[source]

Creates the 2D fixed offset fitting function.

default_bounds()[source]

This can be absolutely anything between -inf and inf.

integrate_function(params)[source]

There is no analytic integration

rescale_parameters(params, scale)[source]

We only need to rescale the offset.

class atomcloud.functions.Function1DBase[source]

Bases: FunctionBase

See 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: FunctionBase

Inherits 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: ABC

Base 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: SumFitBaseFunc

See 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

gaussian1D_amp(std)[source]

1D Gaussian equation

class atomcloud.functions.Gaussian1D[source]

Bases: Function1DBase

See 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

create_gaussian1d(anp)[source]
integrate_function(params)[source]
class atomcloud.functions.Gaussian2D[source]

Bases: Function2DBase

See FunctionBase for full documentation.

analyze_parameters(params)[source]

Calculates the lab widths and the integrated density from the fitted parameters

create_function(anp)[source]

Creates the 2D gaussian fitting function

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.

default_bounds()[source]

Returns the default bounds for the 2D Gaussian function

integrate_function(params)[source]

Integrates the 2D gaussian function analytically

rescale_analysis_params(analysis_params, scale)[source]

Rescales the analysis parameters from the func above

rescale_parameters(params, scales)[source]

Rescales the fitted parameters to the lab frame

class atomcloud.functions.MultiFunc(function_names, func_registry, constraints=None, use_jax=False)[source]

Bases: ConstrainedMultiFunction

Base 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: MultiFunc

1D 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: MultiFunc

2D 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: ABC

Class 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: SumFitBaseFunc

See 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

parabola1D_amp(rx)[source]

1D parabola equation

class atomcloud.functions.Parabola1D[source]

Bases: Function1DBase

See 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

create_parabola1d(anp)[source]
integrate_function(params)[source]
class atomcloud.functions.Parabola2D[source]

Bases: Function2DBase

Parabola 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_function(anp)[source]

Creates the 2D parabola fitting function.

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.

integrate_function(params)[source]

Integrates the 2D parabola function. This is done analytically

rescale_analysis_params(analysis_params, scale)[source]

Rescales the analysis parameters for the 2D parabola function. This is the same as the 2d gaussian function.

rescale_parameters(params, scales)[source]

Rescales the parameters for the 2D parabola function. This is the same as the 2d gaussian function.

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: SumFitBaseFunc

See 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

tf1d_amp(rx)[source]

Integrated parabola equation this needs fixing too

class atomcloud.functions.ThomasFermi1D[source]

Bases: Parabola1D

See 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

integrate_function(params)[source]
class atomcloud.functions.ThomasFermi2D[source]

Bases: Parabola2D

Thomas-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.

create_function(anp)[source]

Creates the 2D Thomas-Fermi fitting function.

integrate_function(params)[source]

Integrates the 2D Thomas-Fermi function. This is done analytically

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_polylog2d(npt)[source]
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.general_rescale(params, scales)[source]
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]

atomcloud.functions.signature(obj, *, follow_wrapped=True)[source]

Get a signature object for the passed callable.