Ramsey.py

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from scipy.optimize import basinhopping
from scipy.ndimage import gaussian_filter1d
from sklearn.model_selection import KFold


def cv_bandwidth(X, y, bandwidths=np.logspace(-1, 1, 10)):
    kf = KFold(n_splits=5)
    scores = []
    for h in bandwidths:
        score = 0
        for train_idx, val_idx in kf.split(X):
            y_smooth = gaussian_filter1d(y[train_idx], h)
            result = fit_ramsey(X[train_idx], y_smooth)
            score += np.mean((y[val_idx] - ramsey_func(X[val_idx], *result)) ** 2)
        scores.append(score)
    return bandwidths[np.argmin(scores)]


def ramsey_func(x, a, delta, t, off, phase_off):
    return a * np.cos(2 * np.pi * delta * x + phase_off) * np.exp(-x / t) + off


def fit_ramsey(X, y):
    t_init = X[len(X) // 3] / np.log(np.max(np.abs(y[:10])) / np.max(np.abs(y[-10:])))
    a_init = np.max(np.abs(y[:10]))

    # FFT on windowed data for frequency
    window = np.blackman(len(X))
    fft = np.fft.rfft(y * window)
    freqs = np.fft.rfftfreq(len(X), X[1] - X[0])
    delta_init = freqs[np.argmax(np.abs(fft[1:]))]

    initial_guess = [a_init, delta_init, t_init, np.mean(y), 0]
    minimizer_kwargs = {
        "method": "L-BFGS-B",
        "bounds": [(0, None), (0, None), (0, None), (None, None), (-np.pi, np.pi)],
    }

    result = basinhopping(
        lambda p: np.mean((y - ramsey_func(X, *p)) ** 2),
        initial_guess,
        minimizer_kwargs=minimizer_kwargs,
    )
    return result.x


def plot_ramsey_fit(X, y0):
    plt.plot(X, y0, ".-", label="data")

    a0 = y0.ptp() / 2
    off0 = y0.mean()

    y = y0 - off0
    y = y / a0
    filter_width = cv_bandwidth(X, y)
    y = gaussian_filter1d(y, filter_width)
    result = fit_ramsey(X, y)
    result[0] = result[0] * a0
    result[3] = result[3] * a0 + off0

    print(f"Amplitude: {result[0]:.2e}")
    print(f"Frequency: {result[1]:.2e}")
    print(f"T: {result[2]:.2e}")
    print(f"Offset: {result[3]:.2e}")
    print(f"Phase offset: {result[4]:.2e}")

    plt.plot(X, ramsey_func(X, *result), "g-", label="fit")
    plt.legend()
    return result
  
  
data = xr.open_zarr("./T2_Ramsey_153655/xarray.zarr/")
X = data.Delay.values
y = data["q4.resonator"].values

plot_ramsey_fit(X, y)