### Install JAXFit Library

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_2D_Gaussian_Demo.ipynb

This snippet provides the command to install the JAXFit library using pip. It is a necessary first step to use JAXFit in your environment.

```python
!pip install jaxfit
```

--------------------------------

### Install JAXFIT Library with pip

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This snippet demonstrates the command to install the JAXFIT library using pip. It's the essential first step to set up your Python environment for using JAXFIT's functionalities.

```python
!pip install jaxfit
```

--------------------------------

### Install JAXFit with pip

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet demonstrates how to install the JAXFit library using pip. It's the first step to setting up your environment for curve fitting with JAX.

```Python
!pip install jaxfit
```

--------------------------------

### Installing JAXFit with pip

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit 2D Gaussian Demo.ipynb

This snippet demonstrates how to install the JAXFit library using pip. It's recommended to set the runtime type to GPU for optimal performance in environments like Google Colab.

```python
!pip install jaxfit
```

--------------------------------

### Optimize JAXFit for Drastically Different Data Sizes

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This example shows how to handle analyses with drastically different data sizes by instantiating multiple `CurveFit` objects, each with an appropriate `flength`. This allows for optimized fit speeds across different data ranges, preventing a single large `flength` from slowing down fits for smaller datasets, by tailoring the fixed length to specific data ranges.

```python
lmin = 10**3
lmax = 10**6
nlengths = 20
lengths1 = np.linspace(10**3, 5 * 10**4, nlengths, dtype=int)
lengths2 = np.linspace(10**5, 10**6, nlengths, dtype=int)

fixed_length1 = np.amax(lengths1)
fixed_length2 = np.amax(lengths2)

jcf1 = CurveFit(flength=fixed_length1)
jcf2 = CurveFit(flength=fixed_length2)

jax_fit_times1 = []
jax_fit_times2 = []

for length1, length2 in zip(lengths1, lengths2):
  xdata1, ydata1 = get_random_data(length1)
  xdata2, ydata2 = get_random_data(length2)

  start_time = time.time()
  popt1, pcov1 = jcf1.curve_fit(linear, xdata1, ydata1, p0=(1,1))
  jax_fit_times1.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = jcf2.curve_fit(linear, xdata2, ydata2, p0=(1,1))
  jax_fit_times2.append(time.time() - start_time)

plt.figure()
plt.title('Fit Speeds')
plt.plot(jax_fit_times1, label='Small Data')
plt.plot(jax_fit_times2, label='Large Data')
plt.legend()

plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Import JAXFIT and JAX Modules

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This code imports the necessary `CurveFit` class from JAXFIT and `jax.numpy` for numerical operations. It is crucial to import JAXFIT before JAX to ensure JAX is configured for 64-bit array computation, which is important for precision in fitting.

```python
from jaxfit import CurveFit
import jax.numpy as jnp
```

--------------------------------

### Perform Curve Fitting with JAXFIT on Synthetic Data

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This snippet demonstrates the core functionality of JAXFIT by using the `CurveFit` class to fit the synthetic data to the linear function. It initializes the fitter, executes the fit, and outputs the actual versus fitted parameters, showcasing JAXFIT's ease of use.

```python
jcf = CurveFit()
popt, pcov = jcf.curve_fit(linear, x, y, p0=(1,1))
y_fit = linear(x, *popt) 

print('Actual Parameters', params)
print('Fit Parameters', popt)
```

--------------------------------

### Install JAX Core Library

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/README.md

This command installs the core JAX library, ensuring its version corresponds to the previously installed jaxlib library for compatibility and proper functioning.

```shell
pip install jax==0.3.14
```

--------------------------------

### Benchmark JAXFIT Performance Against SciPy

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This code block sets up a performance benchmark to compare JAXFIT's speed against SciPy's `curve_fit` for multiple large datasets. It includes a helper function to generate random parameters and measures the time taken for each fit, illustrating JAXFIT's significant speed advantages after the initial JAX tracing.

```python
from scipy.optimize import curve_fit
import time

def get_random_parameters(mmin=1, mmax=10, bmin=0, bmax=10):
  deltam = mmax - mmin
  deltab = bmax - bmin
  m = mmin + deltam * np.random.random()
  b = bmin + deltab * np.random.random()
  return m, b
  
  
length = 3 * 10**5
x = np.linspace(0, 10, length)

jcf = CurveFit()
jax_fit_times = []
scipy_fit_times = []
nsamples = 21
for i in range(nsamples):
  params = get_random_parameters()
  y = linear(x, *params) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, x, y, p0=(1,1))
  jax_fit_times.append(time.time() - start_time)

plt.figure()
plt.title('Fit Speeds')
plt.plot(jax_fit_times, label='JAXFit')
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Demonstrating JAX Retracing with a Single CurveFit Object

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This snippet illustrates the performance overhead when a single CurveFit object is used to fit multiple different functions (linear and quadratic exponential). It shows that JAX retracing occurs for every fit, leading to increased execution times.

```python
import jax.numpy as jnp

def quad_exp(x, a, b, c, d):
  return a * x**2 + b * x + c + jnp.exp(d)

length = 3 * 10**5
x = np.linspace(0, 10, length)

jcf = CurveFit()
nsamples = 21

all_linear_params = np.random.random(size=(nsamples, 2))
all_quad_params = np.random.random(size=(nsamples, 4))

linear_fit_times = []
quad_fit_times = []
for i in range(nsamples):
  y_linear = linear(x, *all_linear_params[i]) + np.random.normal(0, 0.2, size=length)
  y_quad = quad_exp(x, *all_quad_params[i]) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, x, y_linear, p0=(.5,.5,))
  linear_fit_times.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = jcf.curve_fit(quad_exp, x, y_quad, p0=(.5,.5,.5,.5))
  quad_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(linear_fit_times + quad_fit_times))

plt.figure()
plt.plot(linear_fit_times, label='Linear Function')
plt.plot(quad_fit_times, label='Quad Function')
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
plt.legend()
plt.show()
```

--------------------------------

### Install JAX Core Library Matching JAXlib Version

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This command installs the core JAX library, ensuring its version matches the previously installed jaxlib. It's crucial to align these versions for proper functionality and compatibility.

```Bash
pip install jax==0.3.14
```

--------------------------------

### Install JAXFit via Pip

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/main.rst

This command installs the JAXFit library using pip. It should be executed after JAX has been successfully installed, as JAXFit is built on top of JAX.

```Bash
pip install jaxfit
```

--------------------------------

### Generate and Visualize Synthetic Data

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This code generates synthetic data based on the defined linear function, adds random noise, and then visualizes it using `matplotlib`. This step simulates typical experimental data that would be subjected to curve fitting, providing a clear context for the fitting process.

```python
import numpy as np
import matplotlib.pyplot as plt

# make the synthetic data
length = 1000
x = np.linspace(0, 10, length)
params = (3, 5)
y = linear(x, *params) 
# add a little noise to the data to make things interesting
y += np.random.normal(0, 0.2, size=length)

plt.figure()
plt.title('Some Noisy Data')
plt.plot(x, y)
plt.show()
```

--------------------------------

### Pip Install JAXlib Pre-built Wheel for Windows GPU

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/main.rst

This command installs a specific pre-built JAXlib wheel for Windows, optimized for CUDA 11.1 and Python 3.9. It directly fetches the wheel from a provided URL, which is a common method for Windows JAX installations.

```Bash
pip install https://whls.blob.core.windows.net/unstable/cuda111/jaxlib-0.3.14+cuda11.cudnn82-cp39-none-win_amd64.whl
```

--------------------------------

### Comparing JAXFit and SciPy Curve Fitting Speeds

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This snippet compares the execution speed of JAXFit's curve_fit against SciPy's curve_fit for a quadratic exponential function. It highlights JAXFit's significant speed advantage, especially after the initial tracing, due to its optimization capabilities and GPU utilization.

```python
import jax.numpy as jnp
  
length = 3 * 10**5
x = np.linspace(0, 10, length)

jcf = CurveFit()
jax_fit_times = []
scipy_fit_times = []
nsamples = 21

all_params = np.random.random(size=(nsamples, 4))

for i in range(nsamples):
  params = get_random_parameters()
  y = quad_exp(x, *all_params[i]) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(quad_exp, x, y, p0=(.5,.5,.5,.5))
  jax_fit_times.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = curve_fit(quad_exp, x, y, p0=(.5,.5,.5,.5))
  scipy_fit_times.append(time.time() - start_time)

plt.figure()
plt.title('Fit Speeds')
plt.plot(jax_fit_times, label='JAXFit')
plt.plot(scipy_fit_times, label='SciPy')
plt.legend()
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### JAXFit CurveFit Initialization Example

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_2D_Gaussian_Demo.ipynb

This commented-out line provides an example of initializing the JAXFit CurveFit object with an explicit `flength` parameter. This can be useful for optimizing performance with large datasets.

```python
# jcf = CurveFit(flength=length**2)
```

--------------------------------

### Optimizing JAXFit Performance with Multiple CurveFit Objects

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This snippet demonstrates how to improve performance by instantiating separate CurveFit objects for each distinct function being fitted. This approach minimizes JAX retracing, ensuring it only occurs for the first fit of each object, resulting in faster subsequent fits.

```python
jcf_linear = CurveFit()
jcf_quad = CurveFit()


linear_fit_times = []
quad_fit_times = []
for i in range(nsamples):

  y_linear = linear(x, *all_linear_params[i]) + np.random.normal(0, 0.2, size=length)
  y_quad = quad_exp(x, *all_quad_params[i]) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf_linear.curve_fit(linear, x, y_linear, p0=(.5,.5,))
  linear_fit_times.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = jcf_quad.curve_fit(quad_exp, x, y_quad, p0=(.5,.5,.5,.5))
  quad_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(linear_fit_times + quad_fit_times))

plt.figure()
plt.plot(linear_fit_times, label='Linear Function')
plt.plot(quad_fit_times, label='Quad Function')
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
plt.legend()
plt.show()
```

--------------------------------

### Install Pre-built JAXlib Wheel for CUDA 11.1 (Windows)

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/README.md

This command installs a specific pre-built JAXlib wheel for Windows, compatible with CUDA 11.1 and Python 3.9. Using pre-built wheels simplifies the JAX GPU installation process on Windows.

```shell
pip install https://whls.blob.core.windows.net/unstable/cuda111/jaxlib-0.3.14+cuda11.cudnn82-cp39-none-win_amd64.whl
```

--------------------------------

### Install JAXFit Python Package

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/README.md

This command installs the JAXFit library using pip, assuming JAX has been pre-installed. JAXFit is a pure Python package based on the JAX framework.

```shell
pip install jaxfit
```

--------------------------------

### Import JAXFit and JAX for numerical operations

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet imports the necessary modules: `CurveFit` from JAXFit for fitting functionalities and `jax.numpy` for numerical operations. It's crucial to import JAXFit first to ensure JAX uses 64-bit arrays.

```Python
from jaxfit import CurveFit
import jax.numpy as jnp
```

--------------------------------

### Install JAXFit Python Package via Pip

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This command installs the JAXFit library using pip after JAX has been successfully installed. JAXFit is a pure Python package built on JAX.

```Bash
pip install jaxfit
```

--------------------------------

### Analyzing Fit Performance and Parameters

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit 2D Gaussian Demo.ipynb

This code prints the average fit time for JAXFit (excluding the first run due to JIT compilation overhead) and compares the fitted parameters obtained from JAXFit and SciPy. It also generates a plot showing the fit speed over multiple iterations for JAXFit, visually demonstrating its efficiency.

```python
print('Average fit time', np.mean(times[1:]))
print('JAXFit parameters', popt)
print('SciPy parameters', popt2)

plt.figure()
plt.plot(times[1:])
plt.xlabel('Fit Number')
plt.ylabel('Fit Speed (seconds)')
plt.show()
```

--------------------------------

### Install Pre-built JAXlib Wheel for Windows GPU

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This command installs a specific pre-built JAXlib wheel from CloudHan's repository. This wheel is optimized for Windows systems with CUDA 11.1, CUDnn 8.2, and Python 3.9, enabling GPU acceleration for JAX.

```Bash
pip install https://whls.blob.core.windows.net/unstable/cuda111/jaxlib-0.3.14+cuda11.cudnn82-cp39-none-win_amd64.whl
```

--------------------------------

### Measure JAXFit Performance with Varying Data Size

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This code defines helper functions to generate random data and then measures the time taken by JAXFit's `curve_fit` function to fit data arrays of increasing sizes. It demonstrates the performance slowdown due to JAX's retrace mechanism when input array shapes change, as JAX must retrace the function for each new size.

```python
def get_coordinates(length, xmin=0, xmax=10):
  return np.linspace(xmin, xmax, length)


def get_random_data(length):
  xdata = get_coordinates(length)
  params = get_random_parameters()
  ydata = linear(xdata, *params) + np.random.normal(0, 0.2)
  return xdata, ydata


lmin = 10**3
lmax = 10**6
nlengths = 20
lengths = np.linspace(lmin, lmax, nlengths, dtype=int)

jcf = CurveFit()
jax_fit_times = []
for length in lengths:
  xdata, ydata = get_random_data(length)

  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, xdata, ydata, p0=(1,1))
  jax_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(jax_fit_times))

plt.figure()
plt.title('Fit Speeds')
plt.plot(lengths, jax_fit_times, label='JAXFit')
plt.xlabel('Data Length')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Pip Install JAX Version

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/main.rst

This command installs a specific version of the JAX library that corresponds to the previously installed jaxlib version. It's crucial to match JAX and jaxlib versions for compatibility.

```Bash
pip install jax==0.3.14
```

--------------------------------

### Install CUDA Toolkit Developer Tools for JAX (Windows GPU)

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This command installs additional developer tools required by JAX, which are not included in the standard CUDA Toolkit package installed previously. This is necessary for full JAX functionality with GPU support.

```Bash
conda install -c conda-forge cudatoolkit-dev
```

--------------------------------

### Perform curve fitting using JAXFit

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet demonstrates how to use `JAXFit.curve_fit` to fit the synthetic data to the `linear` function. It initializes `CurveFit`, performs the fit, and prints the actual versus fitted parameters, showcasing JAXFit's core functionality.

```Python
jcf = CurveFit()
popt, pcov = jcf.curve_fit(linear, x, y, p0=(1,1))
y_fit = linear(x, *popt) 

print('Actual Parameters', params)
print('Fit Parameters', popt)
```

--------------------------------

### Define a Linear Function for Curve Fitting

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This snippet defines a simple linear function `linear(x, m, b)` using `jax.numpy`. This function serves as the model that will be fitted to data, illustrating how to create custom functions compatible with JAXFIT's curve fitting capabilities.

```python
def linear(x, m, b):
  return m * x + b
```

--------------------------------

### Import JAXFit and JAX Libraries

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_2D_Gaussian_Demo.ipynb

This code imports the necessary CurveFit class from JAXFit and the jax.numpy module. It is crucial to import JAXFit before JAX to ensure JAX computations use 64-bit arrays.

```python
from jaxfit import CurveFit
import jax.numpy as jnp
```

--------------------------------

### Analyze and Compare Fit Performance

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_2D_Gaussian_Demo.ipynb

This code prints the average fit time for JAXFit and displays the fitted parameters from both JAXFit and SciPy. It also generates a plot to visualize the fit speed over multiple iterations, demonstrating JAXFit's efficiency.

```python
print('Average fit time', np.mean(times[1:]))
print('JAXFit parameters', popt)
print('SciPy parameters', popt2)

plt.figure()
plt.plot(times[1:])
plt.xlabel('Fit Number')
plt.ylabel('Fit Speed (seconds)')
plt.show()
```

--------------------------------

### Optimize JAXFit Performance with Fixed Data Array Size

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_Quickstart.ipynb

This code demonstrates how to significantly improve JAXFit performance by instantiating the `CurveFit` object with a `flength` parameter. This fixed-length approach prevents JAX from re-tracing the function for different input array sizes, leading to much faster fits, especially for varying data lengths, by using dummy data internally to maintain a consistent array shape.

```python
fixed_length = np.amax(lengths)
jcf = CurveFit(flength=fixed_length)

jax_fit_times = []
for length in lengths:
  xdata, ydata = get_random_data(length)

  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, xdata, ydata, p0=(1,1))
  jax_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(jax_fit_times))

plt.figure()
plt.title('Fit Speeds')
plt.plot(lengths, jax_fit_times, label='JAXFit')
plt.xlabel('Data Length')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Comparing JAXFit and SciPy Curve Fitting Speed

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet compares the performance of JAXFit's `curve_fit` against SciPy's `curve_fit` for a quadratic exponential function. It highlights JAXFit's significant speed advantage, particularly after the initial tracing phase, due to its ability to utilize GPUs and avoid repeated retracing.

```python
import jax.numpy as jnp
  
length = 3 * 10**5
x = np.linspace(0, 10, length)

jcf = CurveFit()
jax_fit_times = []
scipy_fit_times = []
nsamples = 21

all_params = np.random.random(size=(nsamples, 4))

for i in range(nsamples):
  params = get_random_parameters()
  y = quad_exp(x, *all_params[i]) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(quad_exp, x, y, p0=(.5,.5,.5,.5))
  jax_fit_times.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = curve_fit(quad_exp, x, y, p0=(.5,.5,.5,.5))
  scipy_fit_times.append(time.time() - start_time)

plt.figure()
plt.title('Fit Speeds')
plt.plot(jax_fit_times, label='JAXFit')
plt.plot(scipy_fit_times, label='SciPy')
plt.legend()
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Install CUDA Toolkit Developer Tools via Conda

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/README.md

This command installs additional developer tools for the CUDA Toolkit, which are required by JAX for proper GPU functionality and compilation.

```shell
conda install -c conda-forge cudatoolkit-dev
```

--------------------------------

### Activate JAX Conda Environment

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This command activates the previously created 'jaxenv' Conda environment. All subsequent installations for JAX and JAXFit should be performed within this activated environment.

```Bash
conda activate jaxenv
```

--------------------------------

### Benchmark JAXFit performance for multiple fits

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet measures and plots the execution time of multiple JAXFit curve fitting operations on large datasets. It highlights JAX's 'first-run slow, subsequent-runs fast' behavior due to function tracing.

```Python
from scipy.optimize import curve_fit
import time

def get_random_parameters(mmin=1, mmax=10, bmin=0, bmax=10):
  deltam = mmax - mmin
  deltab = bmax - bmin
  m = mmin + deltam * np.random.random()
  b = bmin + deltab * np.random.random()
  return m, b
  
  
length = 3 * 10**5
x = np.linspace(0, 10, length)

jcf = CurveFit()
jax_fit_times = []
scipy_fit_times = []
nsamples = 21
for i in range(nsamples):
  params = get_random_parameters()
  y = linear(x, *params) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, x, y, p0=(1,1))
  jax_fit_times.append(time.time() - start_time)

plt.figure()
plt.title('Fit Speeds')
plt.plot(jax_fit_times, label='JAXFit')
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Activate Conda Environment for JAX

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/README.md

This command activates the previously created 'jaxenv' Anaconda environment, making it the active environment for subsequent package installations.

```shell
conda activate jaxenv
```

--------------------------------

### Install CUDA Toolkit Developer Tools with Conda

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/main.rst

This command installs additional developer tools for the CUDA Toolkit into the Conda environment. These tools are required by JAX for full functionality and proper integration with CUDA.

```Bash
conda install -c conda-forge cudatoolkit-dev
```

--------------------------------

### Perform Curve Fitting with JAXFit and SciPy

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_2D_Gaussian_Demo.ipynb

This snippet initializes the JAXFit CurveFit object and performs 100 iterations of curve fitting on the synthetic data, recording the time for each fit. It also includes a single fit using SciPy's `curve_fit` for direct comparison.

```python
from scipy.optimize import curve_fit

def get_random_float(low, high):
    delta = high - low
    return low + delta * np.random.random()

flat_data = zdata.flatten()
flat_XY_tuple = [coord.flatten() for coord in XY_tuple]
jcf = CurveFit()

loop = 100
times = []
stimes = []
for i in range(loop):
    seed = [val * get_random_float(.9, 1.2) for val in params]
    st = time.time()
    popt, pcov = jcf.curve_fit(gaussian2d, flat_XY_tuple, flat_data, p0=seed)
    times.append(time.time() - st)

popt2, pcov2 = curve_fit(gaussian2d, flat_XY_tuple, flat_data, p0=seed)
```

--------------------------------

### Importing JAXFit and JAX NumPy

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit 2D Gaussian Demo.ipynb

This code imports `CurveFit` from JAXFit and `jax.numpy` as `jnp`. It's crucial to import JAXFit before JAX to ensure that JAX computations are configured to use 64-bit arrays, which is important for precision in scientific computing.

```python
from jaxfit import CurveFit
import jax.numpy as jnp
```

--------------------------------

### Install CUDA Toolkit 11.1 and CUDnn 8.2 via Conda

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/README.md

This command installs the specified versions of CUDA Toolkit and CUDnn into the active Conda environment. These are essential prerequisites for enabling GPU acceleration with JAX on Windows.

```shell
conda install -c conda-forge cudatoolkit=11.1 cudnn=8.2.0
```

--------------------------------

### Simulate and plot noisy synthetic data

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet generates synthetic data based on the `linear` function, adds random noise, and then visualizes it using Matplotlib. This data will serve as the input for the curve fitting process.

```Python
import numpy as np
import matplotlib.pyplot as plt

# make the synthetic data
length = 1000
x = np.linspace(0, 10, length)
params = (3, 5)
y = linear(x, *params) 
# add a little noise to the data to make things interesting
y += np.random.normal(0, 0.2, size=length)

plt.figure()
plt.title('Some Noisy Data')
plt.plot(x, y)
plt.show()
```

--------------------------------

### Efficient JAXFit Usage: Separate CurveFit Objects for Multiple Functions

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet illustrates the recommended approach for fitting multiple functions with JAXFit. By instantiating a separate `CurveFit` object for each distinct function, retracing is minimized, occurring only for the first fit of each object, significantly improving performance for subsequent fits.

```python
jcf_linear = CurveFit()
jcf_quad = CurveFit()


linear_fit_times = []
quad_fit_times = []
for i in range(nsamples):

  y_linear = linear(x, *all_linear_params[i]) + np.random.normal(0, 0.2, size=length)
  y_quad = quad_exp(x, *all_quad_params[i]) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf_linear.curve_fit(linear, x, y_linear, p0=(.5,.5,))
  linear_fit_times.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = jcf_quad.curve_fit(quad_exp, x, y_quad, p0=(.5,.5,.5,.5))
  quad_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(linear_fit_times + quad_fit_times))

plt.figure()
plt.plot(linear_fit_times, label='Linear Function')
plt.plot(quad_fit_times, label='Quad Function')
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
plt.legend()
plt.show()
```

--------------------------------

### Install CUDA Toolkit and CUDnn for JAX (Windows GPU)

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This command installs specific versions of the CUDA Toolkit (11.1) and CUDnn (8.2.0) into the active Conda environment. These are crucial for running JAX on a CUDA-compatible GPU on Windows systems.

```Bash
conda install -c conda-forge cudatoolkit=11.1 cudnn=8.2.0
```

--------------------------------

### Install CUDA Toolkit and CUDnn with Conda

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/main.rst

This command installs specific versions of the CUDA Toolkit (11.1) and CUDnn (8.2.0) into the active Conda environment. These are essential for enabling GPU acceleration with JAX on compatible hardware.

```Bash
conda install -c conda-forge cudatoolkit=11.1 cudnn=8.2.0
```

--------------------------------

### Define a linear function for curve fitting

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet defines a simple linear function `y = mx + b` using `jax.numpy`. This function will be used as the model for curve fitting with JAXFit.

```Python
def linear(x, m, b):
  return m * x + b
```

--------------------------------

### Initializing JAXFit CurveFit with flength (Commented)

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit 2D Gaussian Demo.ipynb

This commented-out line shows an alternative way to initialize the `CurveFit` object, potentially specifying the flattened length of the data (`flength`). This might be used for optimizing internal memory allocation or processing for specific data shapes.

```python
# jcf = CurveFit(flength=length**2)
```

--------------------------------

### Activate JAX Conda Environment

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/main.rst

This command activates the previously created 'jaxenv' Anaconda environment. Activating the environment makes its Python interpreter and installed packages available for use.

```Bash
conda activate jaxenv
```

--------------------------------

### Create Conda Environment for JAX Installation

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This command creates a new Anaconda environment named 'jaxenv' with Python version 3.9. It's recommended for managing dependencies for JAX and JAXFit.

```Bash
conda create -n jaxenv python=3.9
```

--------------------------------

### Create Python 3.9 Conda Environment for JAX

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/README.md

This command creates a new Anaconda environment named 'jaxenv' with Python version 3.9, which is recommended for JAX installations on Windows systems.

```shell
conda create -n jaxenv python=3.9
```

--------------------------------

### Measure JAXFit Performance with Varying Data Sizes

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This Python snippet demonstrates the default behavior of JAXFit's `curve_fit` when processing input data of varying lengths. It measures and plots the fit times, illustrating that JAX's re-tracing mechanism for different array sizes leads to slower performance as data length increases. The code includes helper functions for generating random data and parameters.

```python
def get_coordinates(length, xmin=0, xmax=10):
  return np.linspace(xmin, xmax, length)


def get_random_data(length):
  xdata = get_coordinates(length)
  params = get_random_parameters()
  ydata = linear(xdata, *params) + np.random.normal(0, 0.2)
  return xdata, ydata


lmin = 10**3
lmax = 10**6
nlengths = 20
lengths = np.linspace(lmin, lmax, nlengths, dtype=int)

jcf = CurveFit()
jax_fit_times = []
for length in lengths:
  xdata, ydata = get_random_data(length)

  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, xdata, ydata, p0=(1,1))
  jax_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(jax_fit_times))

plt.figure()
plt.title('Fit Speeds')
plt.plot(lengths, jax_fit_times, label='JAXFit')
plt.xlabel('Data Length')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Optimize JAXFit for Disparate Data Sizes Using Multiple CurveFit Objects

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This Python snippet presents an advanced optimization strategy for JAXFit when dealing with datasets of drastically different sizes. It demonstrates instantiating multiple `CurveFit` objects, each configured with a `flength` appropriate for a specific range of data sizes. This approach allows for optimized performance across varied data scales by leveraging JAX's fixed-size optimization for each distinct data range.

```python
lmin = 10**3
lmax = 10**6
nlengths = 20
lengths1 = np.linspace(10**3, 5 * 10**4, nlengths, dtype=int)
lengths2 = np.linspace(10**5, 10**6, nlengths, dtype=int)

fixed_length1 = np.amax(lengths1)
fixed_length2 = np.amax(lengths2)

jcf1 = CurveFit(flength=fixed_length1)
jcf2 = CurveFit(flength=fixed_length2)

jax_fit_times1 = []
jax_fit_times2 = []

for length1, length2 in zip(lengths1, lengths2):
  xdata1, ydata1 = get_random_data(length1)
  xdata2, ydata2 = get_random_data(length2)

  start_time = time.time()
  popt1, pcov1 = jcf1.curve_fit(linear, xdata1, ydata1, p0=(1,1))
  jax_fit_times1.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = jcf2.curve_fit(linear, xdata2, ydata2, p0=(1,1))
  jax_fit_times2.append(time.time() - start_time)

plt.figure()
plt.title('Fit Speeds')
plt.plot(jax_fit_times1, label='Small Data')
plt.plot(jax_fit_times2, label='Large Data')
plt.legend()

plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Simulating Synthetic 2D Gaussian Data

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit 2D Gaussian Demo.ipynb

This code simulates 2D Gaussian data with added noise, which will be used as synthetic data for fitting. It includes helper functions `get_coordinates` to generate meshgrid coordinates and `get_gaussian_parameters` to define the parameters for the 2D Gaussian. The simulated data is then visualized using Matplotlib's `imshow`.

```python
import numpy as np
import matplotlib.pyplot as plt
import time

def get_coordinates(width, height):
    x = np.linspace(0, width - 1, width)
    y = np.linspace(0, height - 1, height)
    X, Y = np.meshgrid(x, y)
    return X, Y


def get_gaussian_parameters(length):
  n0 = 1
  x0 = length / 2
  y0 = length / 2
  sigx = length / 6
  sigy = length / 8
  theta = np.pi / 3

  offset = .1 * n0
  params = [n0, x0, y0, sigx, sigy, theta, offset]
  return params

length = 500
XY_tuple = get_coordinates(length, length)

params = get_gaussian_parameters(length)
zdata = gaussian2d(XY_tuple, *params)
zdata += np.random.normal(0, .1, size=(length, length))

plt.imshow(zdata)
plt.show()
```

--------------------------------

### Simulate Synthetic 2D Gaussian Data

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_2D_Gaussian_Demo.ipynb

This code generates synthetic 2D Gaussian data with added noise, which will be used as input for the curve fitting process. It includes helper functions to create coordinate grids and define Gaussian parameters, and visualizes the generated data using Matplotlib.

```python
import numpy as np
import matplotlib.pyplot as plt
import time

def get_coordinates(width, height):
    x = np.linspace(0, width - 1, width)
    y = np.linspace(0, height - 1, height)
    X, Y = np.meshgrid(x, y)
    return X, Y


def get_gaussian_parameters(length):
  n0 = 1
  x0 = length / 2
  y0 = length / 2
  sigx = length / 6
  sigy = length / 8
  theta = np.pi / 3

  offset = .1 * n0
  params = [n0, x0, y0, sigx, sigy, theta, offset]
  return params

length = 500
XY_tuple = get_coordinates(length, length)

params = get_gaussian_parameters(length)
zdata = gaussian2d(XY_tuple, *params)
zdata += np.random.normal(0, .1, size=(length, length))

plt.imshow(zdata)
plt.show()
```

--------------------------------

### Fit a Linear Function with JAXFit

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/README.rst

This example demonstrates how to use JAXFit as a drop-in replacement for SciPy's curve_fit. It defines a simple linear fit function, prepares sample data, and then uses the CurveFit class to perform the nonlinear least squares fitting, returning optimal parameters and covariance.

```python
import numpy as np
from jaxfit import CurveFit

def linear(x, m, b): # fit function
	return m * x + b

x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
y = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

cf = CurveFit()
popt, pcov = cf.curve_fit(linear, x, y)
```

--------------------------------

### Optimize JAXFit Performance with Fixed Data Array Size

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This Python snippet shows how to significantly improve JAXFit's performance by initializing the `CurveFit` object with a `flength` parameter. By setting a fixed maximum data size, JAX avoids re-tracing for subsequent fits with smaller or equal data lengths, resulting in much faster and consistent fit times. A caveat is that the fit speed is always that of the fixed data size.

```python
fixed_length = np.amax(lengths)
jcf = CurveFit(flength=fixed_length)

jax_fit_times = []
for length in lengths:
  xdata, ydata = get_random_data(length)

  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, xdata, ydata, p0=(1,1))
  jax_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(jax_fit_times))

plt.figure()
plt.title('Fit Speeds')
plt.plot(lengths, jax_fit_times, label='JAXFit')
plt.xlabel('Data Length')
plt.ylabel('Fit Time (seconds)')
```

--------------------------------

### Perform Linear Curve Fitting with JAXFit

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/main.rst

This example demonstrates how to use JAXFit's `CurveFit` class to perform a nonlinear least squares fit on a simple linear dataset. It defines a custom linear function, prepares input and output data, and then applies `curve_fit` to obtain the optimal parameters and their covariance matrix.

```python
import numpy as np
from jaxfit import CurveFit

def linear(x, m, b): # fit function
	return m * x + b

x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
y = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

cf = CurveFit()
popt, pcov = cf.curve_fit(linear, x, y)
```

--------------------------------

### Define 2D Gaussian Function with JAX

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/docs/notebooks/JAXFit_2D_Gaussian_Demo.ipynb

This snippet defines two functions: `rotate_coordinates2D` for rotating 2D coordinates and `gaussian2d` for calculating the density of a 2D Gaussian. Both functions are constructed using `jax.numpy` to ensure compatibility with JAX's computation graph.

```python
def rotate_coordinates2D(coords, theta):
    R = jnp.array([[jnp.cos(theta), -jnp.sin(theta)],
                  [jnp.sin(theta), jnp.cos(theta)]])
  
    shape = coords[0].shape
    coords = jnp.stack([coord.flatten() for coord in coords])
    rcoords = R @ coords
    return [jnp.reshape(coord, shape) for coord in rcoords]


def gaussian2d(coords, n0, x0, y0, sigma_x, sigma_y, theta, offset):
    coords = [coords[0] - x0, coords[1] - y0] #translate first
    X, Y = rotate_coordinates2D(coords, theta)
    density = n0 * jnp.exp(-.5 * (X**2 / sigma_x**2 + Y**2 / sigma_y**2))
    return density + offset
```

--------------------------------

### Inefficient JAXFit Usage: Single CurveFit Object for Multiple Functions

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit Quickstart.ipynb

This snippet demonstrates the performance impact of using a single `CurveFit` object to fit multiple different functions. It shows that JAX retracing occurs for every fit, leading to increased computation times, especially when fitting different function types repeatedly.

```python
import jax.numpy as jnp

def quad_exp(x, a, b, c, d):
  return a * x**2 + b * x + c + jnp.exp(d)

length = 3 * 10**5
x = np.linspace(0, 10, length)

jcf = CurveFit()
nsamples = 21

all_linear_params = np.random.random(size=(nsamples, 2))
all_quad_params = np.random.random(size=(nsamples, 4))

linear_fit_times = []
quad_fit_times = []
for i in range(nsamples):
  y_linear = linear(x, *all_linear_params[i]) + np.random.normal(0, 0.2, size=length)
  y_quad = quad_exp(x, *all_quad_params[i]) + np.random.normal(0, 0.2, size=length)

  # fit the data
  start_time = time.time()
  popt1, pcov1 = jcf.curve_fit(linear, x, y_linear, p0=(.5,.5,))
  linear_fit_times.append(time.time() - start_time)

  start_time = time.time()
  popt2, pcov2 = jcf.curve_fit(quad_exp, x, y_quad, p0=(.5,.5,.5,.5))
  quad_fit_times.append(time.time() - start_time)

print('Summed Fit Times', np.sum(linear_fit_times + quad_fit_times))

plt.figure()
plt.plot(linear_fit_times, label='Linear Function')
plt.plot(quad_fit_times, label='Quad Function')
plt.xlabel('Fit Iteration')
plt.ylabel('Fit Time (seconds)')
plt.legend()
plt.show()
```

--------------------------------

### Fitting Synthetic Data with JAXFit and SciPy

Source: https://github.com/dipolar-quantum-gases/jaxfit/blob/main/examples/JAXFit 2D Gaussian Demo.ipynb

This snippet initializes the JAXFit `CurveFit` object and performs 100 fits on the synthetic data, each with a different random seed for initial parameters to test robustness. It also performs a single fit using SciPy's `curve_fit` for direct comparison. The `get_random_float` helper function generates random floats within a specified range for seeding.

```python
from scipy.optimize import curve_fit

def get_random_float(low, high):
    delta = high - low
    return low + delta * np.random.random()

flat_data = zdata.flatten()
flat_XY_tuple = [coord.flatten() for coord in XY_tuple]
jcf = CurveFit()

loop = 100
times = []
stimes = []
for i in range(loop):
    print(i, 'of', loop)
    seed = [val * get_random_float(.9, 1.2) for val in params]
    st = time.time()
    popt, pcov = jcf.curve_fit(gaussian2d, flat_XY_tuple, flat_data, p0=seed)
    times.append(time.time() - st)

popt2, pcov2 = curve_fit(gaussian2d, flat_XY_tuple, flat_data, p0=seed)
```