AtomCloud Multi-Functions
An AtomCloud Multi function object allow us to construct a single function out of individual subfunctions. This is needed when we use either SciPy or JAXFit to do curve fitting on data.
First, let’s install AtomCloud and JAXFit (for GPU acceleration).
[ ]:
!pip install jaxfit
!pip install atomcloud
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting jaxfit
Downloading jaxfit-0.0.5-py3-none-any.whl (36 kB)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from jaxfit) (1.22.4)
Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from jaxfit) (3.7.1)
Requirement already satisfied: scipy>=1.7.0 in /usr/local/lib/python3.10/dist-packages (from jaxfit) (1.10.1)
Requirement already satisfied: jaxlib>=0.3.10 in /usr/local/lib/python3.10/dist-packages (from jaxfit) (0.4.10+cuda11.cudnn86)
Requirement already satisfied: JAX>=0.3.10 in /usr/local/lib/python3.10/dist-packages (from jaxfit) (0.4.10)
Requirement already satisfied: ml-dtypes>=0.1.0 in /usr/local/lib/python3.10/dist-packages (from JAX>=0.3.10->jaxfit) (0.1.0)
Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from JAX>=0.3.10->jaxfit) (3.3.0)
Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (1.0.7)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (0.11.0)
Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (4.39.3)
Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (1.4.4)
Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (23.1)
Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (8.4.0)
Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (3.0.9)
Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib->jaxfit) (2.8.2)
Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib->jaxfit) (1.16.0)
Installing collected packages: jaxfit
Successfully installed jaxfit-0.0.5
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting atomcloud
Downloading atomcloud-0.0.1-py3-none-any.whl (69 kB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 69.5/69.5 kB 2.8 MB/s eta 0:00:00
Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from atomcloud) (3.7.1)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from atomcloud) (1.22.4)
Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from atomcloud) (1.5.3)
Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from atomcloud) (1.10.1)
Requirement already satisfied: tabulate in /usr/local/lib/python3.10/dist-packages (from atomcloud) (0.8.10)
Collecting uncertainties (from atomcloud)
Downloading uncertainties-3.1.7-py2.py3-none-any.whl (98 kB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 98.4/98.4 kB 13.9 MB/s eta 0:00:00
Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (1.0.7)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (0.11.0)
Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (4.39.3)
Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (1.4.4)
Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (23.1)
Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (8.4.0)
Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (3.0.9)
Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib->atomcloud) (2.8.2)
Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->atomcloud) (2022.7.1)
Requirement already satisfied: future in /usr/local/lib/python3.10/dist-packages (from uncertainties->atomcloud) (0.18.3)
Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib->atomcloud) (1.16.0)
Installing collected packages: uncertainties, atomcloud
Successfully installed atomcloud-0.0.1 uncertainties-3.1.7
1D Multi-Function
Let’s start out with the 1D case. By default atomcloud has six 1d functions to choose from (although you can add your own). The functions defined are all related to common distributions found in cold atom physics they are: “gaussian”, “ebose”, “febose”, “parabola”, “tf” and “foffset”. Ebose stands for the enhanced Bose-Einstein distribution and febose stands for fixed enhanced Bose-Einstein distribution—where the fugacity is fixed to 1. The parabola and tf functions are the parabola and Thomas-Fermi distributions respectively (where the Thomas-Fermi distribution has been integrated along two of the three dimensions). Finally, the foffset refers to a fixed offset function. This is a function that is a constant offset from zero. This is useful for fitting data that has a constant background.
We can combine these functions into a single function using the MultiFunction object. Let’s start by creating a multi function object with the “gaussian”, “parabola” and “foffset” functions. First we create our independent data.
[ ]:
import numpy as np
data_length = 256
x = np.linspace(0, data_length - 1, data_length)
Next we create a list of parameters for each function
[ ]:
gaussian_n0 = 1.2
gaussian_x0 = data_length * (1 / 4)
gaussian_std = 10
parabola_n0 = 1.5
parabola_x0 = data_length * (3 / 4)
parabola_r = 15
gaussian_params = [gaussian_n0, gaussian_std, gaussian_x0]
parabola_params = [parabola_n0, parabola_r, parabola_x0]
offset_params = [0.3]
Now we import the 1d MultiFunction object and create our function. To instantiate the MultiFunction object we need to pass it the list of functions as as string. We can optionally pass it a list of parameters to constrain, but we’ll address this later.
Once instantiated, the object can be used as a function. We can pass it our independent data and the function parameters as a list of the individual function parameters’ lists.
[ ]:
import matplotlib.pyplot as plt
from atomcloud.functions import MultiFunction1D
bimodal_equations = ['gaussian', 'parabola', 'foffset']
bimodal_parameters = [gaussian_params, parabola_params, offset_params]
bimodal_func = MultiFunction1D(bimodal_equations)
bimodal_data = bimodal_func(x, bimodal_parameters)
plt.figure()
plt.plot(x, bimodal_data)
[<matplotlib.lines.Line2D at 0x7ffa75d56ec0>]
When fitting, we often wish for two or more parameters to be constrained to be equal. For example, we may want the translation of the parabola and the offset of the gaussian to be equal.
We can do this by passing a list of parameter constraints when instantiating the multifunction object. When constraining a parameter variable, the variable will technically be defined twice (once in each function parameter list). In this case, the multifunction object will use the first definition of the variable (i.e. the first function parameter list where it is defined).
We’ll do this on the above example by constraining the “x0” parameter between the two functions.
[ ]:
bimodal_constraints = ['x0']
bimodal_func = MultiFunction1D(bimodal_equations, bimodal_constraints)
bimodal_data = bimodal_func(x, bimodal_parameters)
plt.figure()
plt.plot(x, bimodal_data)
[<matplotlib.lines.Line2D at 0x7ffa73a78760>]
Finally, we can JAX rather than Numpy to create these multifunctions by passing in a flag. This creates multifunctions that are compatible with JAXFit and allow blazing fast fitting.
[ ]:
bimodal_func = MultiFunction1D(bimodal_equations,
bimodal_constraints,
use_jax=True)
bimodal_data = bimodal_func(x, bimodal_parameters)
plt.figure()
plt.plot(x, bimodal_data)
WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[<matplotlib.lines.Line2D at 0x7ffa701bd300>]
Defining Custom Sub-Functions
By default, only ‘gaussian’, ‘parabola’, ‘tf’, ‘ebose’ ‘febose’ and ‘foffset’ are defined. However, this is a fairly limited number of functions and we might wish to add our own.
All of the function objects are defined in the 1d and 2d function registries. They all inherit from the common FunctionBase class. This means we can easily create our own function objects from this class and add them to the registry.
[ ]:
from atomcloud.functions import FunctionBase, FUNCTIONS1D
from atomcloud.analysis import rescale_1d_params
class Cos1D(FunctionBase):
"""See FunctionBase for documentation"""
def __init__(self):
super().__init__()
def create_function(self, anp):
def exp_func(x, a, b, x0):
return a * anp.cos(b * (x - x0))
return exp_func
FUNCTIONS1D.register('cos', Cos1D)
And now we can use this new function in our multifunction object.
[ ]:
from atomcloud.functions import MultiFunction1D
cos_x0 = data_length * (3 / 4)
cos_amp = .5
cos_freq = 1
cos_params = [cos_amp, cos_freq, cos_x0]
all_params = [gaussian_params, cos_params, offset_params]
cos_equations = ['gaussian', 'cos', 'foffset']
cos_func = MultiFunction1D(cos_equations)
cos_data = cos_func(x, all_params)
plt.figure()
plt.plot(x, cos_data)
plt.show()
In reality, there’s more class functions we should define if we want to use this function for fitting. Below we show a more complete example of a custom function object. Many of the additional functions are used to define default parameters for fitting including seed and bound values. Additionally, we can define rescaling functions which are used later in analysis. We refer the reader to the docs as this is a more advanced topic.
2D MultiFunctions
We can also do this for 2D functions. Below we use a fixed enhanced bose (fugacity=1), a thomas-fermi profile and a fixed offset
[ ]:
from atomcloud.functions import MultiFunction2D
coords = np.meshgrid(x, x)
#common parameters
x0 = data_length / 2
y0 = x0
theta = np.pi / 3
#subfunction parameters
gaussian_stdx = 31.2
gaussian_stdy = 28
tf_rx = 10.4
tf_ry = 13.1
offset = 0.2
gaussan_params = [1.2, x0, y0, gaussian_stdx, gaussian_stdy, theta]
tf_params = [1.5, x0, y0, tf_rx, tf_ry, theta]
offset_params = [offset]
bimodal_equations = ['gaussian', 'tf', 'foffset']
bimodal_parameters = [gaussan_params, tf_params, offset_params]
bimodal_constraints = ['x0', 'y0']
bimodal_func = MultiFunction2D(bimodal_equations,
bimodal_constraints,
use_jax=True)
bimodal_data = bimodal_func(coords, bimodal_parameters)
plt.figure()
plt.pcolormesh(*coords, bimodal_data)
plt.colorbar()
plt.show()
Now we’ll creat a custom 2D funtion and add it to the 2D function register. In contrast to the 1D function, we’ll include some of the additional object functions such as initial seed and default bounds.
[ ]:
from atomcloud.functions import FunctionBase, FUNCTIONS2D
from atomcloud.analysis import rescale_2d_params
class Sin2D(FunctionBase):
"""See FunctionBase for documentation"""
def __init__(self):
super().__init__()
def create_function(self, anp):
def a_sin_func(coords, a, b, c):
x, y = coords
return a * np.sin(b * x + c * y)
return a_sin_func
def initial_seed(self, x, data):
min_ind = np.argmin(data)
max_ind = np.argmax(data)
a = (data[max_ind] - data[min_ind]) / 2
xdelta = x[-1] - x[0]
b = xdelta / 10
c = xdelta / 12
return [a, b, x0]
def rescale_parameters(self, params, scale):
_, axis_scale, zscale = scale
zinds = [0]
xinds = [1, 2]
return rescale_2d_params(params, zinds, xinds, axis_scale, zscale)
def default_bounds(self):
min_bounds = [0, 0, -np.inf]
max_bounds = [np.inf, np.inf, np.inf]
return [min_bounds, max_bounds]
FUNCTIONS2D.register('sin', Sin2D)
Now let’s use it!
[ ]:
sin_freq_x = .1
sin_freq_y = .2
sin_amp = 1.2
sin_params = [sin_amp, sin_freq_x, sin_freq_y]
all_params = [gaussan_params, sin_params, offset_params]
sin_equations = ['gaussian', 'sin', 'foffset']
sin_func = MultiFunction2D(sin_equations)
sin_data = sin_func(coords, all_params)
plt.figure()
plt.pcolormesh(*coords, sin_data)
plt.colorbar()
plt.show()
This all seems fairly simple, but the being able to build arbitrary multifunctions from the registry objects means we can do some pretty powerful things later with both fitting and analysis of our fits.