Correlation Length Calculation From Susceptibility Grid

import numpy as np
import h5py
from scipy.optimize import curve_fit

def get_max_with_coords(vals2D, lenx):
    maxcoord = (None, None)
    maxval = -999999999999999
    for x in range(lenx):
        for y in range(lenx):
            if vals2D[x, y] > maxval:
                maxcoord = (x, y)
                maxval = vals2D[x, y]
    return (maxval, maxcoord)

def oz_func_idx(q_idx, A, chi, centerxy_idx, momgrid, momdimx):
    omega = 0
    gamma = 1

    qx = np.array([momgrid[int(qidx)][0] for qidx in q_idx])
    qy = np.array([momgrid[int(qidx)][1] for qidx in q_idx])

    center = momgrid[centerxy_idx[1]*momdimx + centerxy_idx[0]]

    return A/(np.power(qx - center[0], 2) + np.power(qy - center[1], 2) + pow(chi, -2) + abs(omega)/gamma)

# for fit_window_in_percent: 0.5 is 50%
def fit_to_lorentzian(momgrid, susc, fit_window_in_percent):
    momdimx = int(np.sqrt(susc.shape[1]))
    #assuming square lattice
    susc2D = (susc[list(f["/susc_func/bgrid"]).index(0), :, 0, 0, 0, 0, 0, 0].reshape(momdimx, momdimx))
    max_susc, found_at = get_max_with_coords(susc2D, momdimx)
    suscs_to_fit = []
    ks_to_fit = []

    fit_interval_size = int(momdimx * fit_window_in_percent/2)

    kx_interval_vals = []
    ky_interval_vals = []
    for ix in np.linspace(0, 2*fit_interval_size, 2 * fit_interval_size + 1) - fit_interval_size:
        kx_interval_vals.append(int(found_at[0] + ix))
        ky_interval_vals.append(int(found_at[1] + ix))

    ks_2d_to_fit = []
    ks_to_fit = []
    #iy = int(found_at[1])
    for ix in kx_interval_vals:
        for iy in ky_interval_vals:
            ks_2d_to_fit.append((ix, iy))
            ks_to_fit.append((ix + 10*momdimx)%momdimx + ((iy + 10*momdimx)%momdimx)*momdimx)

    suscs_to_fit = np.array([susc2D[k[0], k[1]] for k in ks_2d_to_fit])

    oz_func_to_fit = lambda q_idx, A, chi: oz_func_idx(q_idx, A, chi, found_at, momgrid, momdimx)


    popt, pcov = curve_fit(oz_func_to_fit, np.array(ks_to_fit), np.array(suscs_to_fit))
    return popt, pcov


##########################################
##########################################
############## example####################
##########################################

f = h5py.File("dat_U-5_Beta7p5_TP0_Mu0_C5_KDIM16_FFT80_REFINE0_TSTART0_HUB_2D_INTFL_2LOOP_1SELOOP_SU2_ALLSYMM_LOCAL_SELFEN_FLOW_SDE_DIVERGENT.h5","r")
lam = np.array(f["Flow_obs/LAM"])[-1]
# the division by lambda^2 depends on version.. At some later point, this was added in the flowing susceptibility. But if your runs are old (e.g. from 2021, you probably need it)
susc = f["/susc_func/RE_SC"]/lam/lam
momgrid = np.array(f["/susc_func/momgrid"])

fit_window = 1.0 #percent of BZ to fit against. It seems like 1.0 (100%) gives the least error (the square root of the (1,1) component of covariance matrix) of the correlation length

ampl_corr_pair, covariance_matrix = fit_to_lorentzian(momgrid, susc, fit_window)

print("Amplitude of the lorenzian:", ampl_corr_pair[0])
print("Width of the lorenzian (the correlation length):", ampl_corr_pair[1])
print("Error in fitting the correlation length: ", np.sqrt(covariance_matrix[1][1])) # the error in fitting the correlation length